20309855018

Committed 17 Dec 2025 04:27PM UTC coverage: 41.938% (-21.7%) from 63.664%

Build # 20309855018

Build Type

Pull #276

github

Committed by

web-flow

Commit Message

Merge 5453ac0d2 into 98d7c449a

Pull Request Pull Request #276: updated black version in pre-commit

Run Details

30 of 1579 new or added lines in 20 files covered. (1.9%)

7 existing lines in 4 files now uncovered.

1883 of 4490 relevant lines covered (41.94%)

0.42 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/src/clm/module_library/sequence_module.py

from torch import nn
import functools


class SequenceModule(nn.Module):
    """Abstract sequence model class. All models must adhere to this interface

    A SequenceModule is generally a model that transforms an input of shape
    (n_batch, l_sequence, d_model) to (n_batch, l_sequence, d_output)

    REQUIRED methods and attributes
    forward, d_model, d_output: controls standard forward pass, a sequence-to-sequence transformation
    __init__ should also satisfy the following interface; see SequenceIdentity for an example
        def __init__(self, d_model, transposed=False, **kwargs)

    OPTIONAL methods
    default_state, step: allows stepping the model recurrently with a hidden state
    state_to_tensor, d_state: allows decoding from hidden state
    """

    @property
    def d_model(self):
        """Model dimension (generally same as input dimension).

        This attribute is required for all SequenceModule instantiations.
        It is used by the rest of the pipeline (e.g. model backbone, encoder) to track the internal shapes of the full model.
        """
        if getattr(self, "_d_model", None) is None:
            raise NotImplementedError("SequenceModule instantiation must set d_model")
        return self._d_model

    @d_model.setter
    def d_model(self, d):
        self._d_model = d

    @property
    def d_output(self):
        """Output dimension of model.

        This attribute is required for all SequenceModule instantiations.
        It is used by the rest of the pipeline (e.g. model backbone, decoder) to track the internal shapes of the full model.
        """
        if getattr(self, "_d_output", None) is None:
            raise NotImplementedError(
                "SequenceModule instantiation must specify d_output for decoder"
            )
        return self._d_output

    @d_output.setter
    def d_output(self, d):
        self._d_output = d

    def forward(self, x, state=None, **kwargs):
        """Forward pass of sequence model, a sequence-to-sequence transformation with an optional state.

        Generally, this should map a tensor of shape (batch, length, self.d_model) to (batch, length, self.d_output)

        Additionally, it returns a "state" which can be any additional information
        For example, RNN and SSM layers may return their hidden state,
        while some types of transformer layers (e.g. Transformer-XL) may want to pass a state as well
        """
        return x, None

    @property
    def state_to_tensor(self):
        """Returns a function mapping a state to a single tensor.

        This method should be implemented if one wants to use the hidden state instead of the output sequence for final prediction.
        Currently only used with the StateDecoder.
        """
        return lambda _: None

    @property
    def d_state(self):
        """Returns dimension of output of self.state_to_tensor"""
        return None

    def default_state(self, *batch_shape, device=None):
        """Create initial state for a batch of inputs."""

        return None

    def step(self, x, state=None, **kwargs):
        """Step the model recurrently for one step of the input sequence.

        For example, this should correspond to unrolling an RNN for one step.
        If the forward pass has signature (B, L, H1) -> (B, L, H2),
        this method should generally have signature (B, H1) -> (B, H2) with an optional recurrent state.
        """
        raise NotImplementedError


def TransposedModule(module):
    """Wrap a SequenceModule class to accept transposed parameter, handle state, absorb kwargs"""

    # https://stackoverflow.com/a/65470430/1980685
    @functools.wraps(module, updated=())
    class TransposedModule(module):
        def __init__(self, *args, transposed=False, **kwargs):
            super().__init__(*args, **kwargs)
            self.transposed = transposed

        def forward(self, x, state=None, **kwargs):
            if self.transposed:
                x = x.transpose(-1, -2)
            x, next_state = super().forward(x, state)  # Don't use kwarg because nn.LSTM
            next_state = None if state is None else next_state
            if self.transposed:
                x = x.transpose(-1, -2)
            return x, next_state

    # https://stackoverflow.com/questions/5352781/how-to-set-class-names-dynamically
    # TransposedModule.__name__ = module.__name__ # functools wraps is better solution
    return TransposedModule


@TransposedModule
class SequenceIdentity(SequenceModule):
    """Simple SequenceModule for testing purposes"""

    def __init__(self, d_model, dropout=0.0, **kwargs):
        """Default interface for SequenceModule

        d_model: input dimension (sometimes denoted H for hidden dimension)
        transposed: if True, inputs have axis ordering (B, H, L) instead of (B, H, L)
        """
        super().__init__()
        self.d_model = d_model
        self.d_output = d_model

    def forward(self, x, state=None):
        return x, state

    def default_state(self, *batch_shape, device=None):
        return None

    def step(self, x, state=None, **kwargs):
        return x, state

NEW 1	from torch import nn	×
NEW 2	import functools	×
3
4
NEW 5	class SequenceModule(nn.Module):	×
6	"""Abstract sequence model class. All models must adhere to this interface
7
8	A SequenceModule is generally a model that transforms an input of shape
9	(n_batch, l_sequence, d_model) to (n_batch, l_sequence, d_output)
10
11	REQUIRED methods and attributes
12	forward, d_model, d_output: controls standard forward pass, a sequence-to-sequence transformation
13	__init__ should also satisfy the following interface; see SequenceIdentity for an example
14	def __init__(self, d_model, transposed=False, **kwargs)
15
16	OPTIONAL methods
17	default_state, step: allows stepping the model recurrently with a hidden state
18	state_to_tensor, d_state: allows decoding from hidden state
19	"""
20
NEW 21	@property	×
NEW 22	def d_model(self):	×
23	"""Model dimension (generally same as input dimension).
24
25	This attribute is required for all SequenceModule instantiations.
26	It is used by the rest of the pipeline (e.g. model backbone, encoder) to track the internal shapes of the full model.
27	"""
NEW 28	if getattr(self, "_d_model", None) is None:	×
NEW 29	raise NotImplementedError("SequenceModule instantiation must set d_model")	×
NEW 30	return self._d_model	×
31
NEW 32	@d_model.setter	×
NEW 33	def d_model(self, d):	×
NEW 34	self._d_model = d	×
35
NEW 36	@property	×
NEW 37	def d_output(self):	×
38	"""Output dimension of model.
39
40	This attribute is required for all SequenceModule instantiations.
41	It is used by the rest of the pipeline (e.g. model backbone, decoder) to track the internal shapes of the full model.
42	"""
NEW 43	if getattr(self, "_d_output", None) is None:	×
NEW 44	raise NotImplementedError(	×
45	"SequenceModule instantiation must specify d_output for decoder"
46	)
NEW 47	return self._d_output	×
48
NEW 49	@d_output.setter	×
NEW 50	def d_output(self, d):	×
NEW 51	self._d_output = d	×
52
NEW 53	def forward(self, x, state=None, **kwargs):	×
54	"""Forward pass of sequence model, a sequence-to-sequence transformation with an optional state.
55
56	Generally, this should map a tensor of shape (batch, length, self.d_model) to (batch, length, self.d_output)
57
58	Additionally, it returns a "state" which can be any additional information
59	For example, RNN and SSM layers may return their hidden state,
60	while some types of transformer layers (e.g. Transformer-XL) may want to pass a state as well
61	"""
NEW 62	return x, None	×
63
NEW 64	@property	×
NEW 65	def state_to_tensor(self):	×
66	"""Returns a function mapping a state to a single tensor.
67
68	This method should be implemented if one wants to use the hidden state instead of the output sequence for final prediction.
69	Currently only used with the StateDecoder.
70	"""
NEW 71	return lambda _: None	×
72
NEW 73	@property	×
NEW 74	def d_state(self):	×
75	"""Returns dimension of output of self.state_to_tensor"""
NEW 76	return None	×
77
NEW 78	def default_state(self, *batch_shape, device=None):	×
79	"""Create initial state for a batch of inputs."""
80
NEW 81	return None	×
82
NEW 83	def step(self, x, state=None, **kwargs):	×
84	"""Step the model recurrently for one step of the input sequence.
85
86	For example, this should correspond to unrolling an RNN for one step.
87	If the forward pass has signature (B, L, H1) -> (B, L, H2),
88	this method should generally have signature (B, H1) -> (B, H2) with an optional recurrent state.
89	"""
NEW 90	raise NotImplementedError	×
91
92
NEW 93	def TransposedModule(module):	×
94	"""Wrap a SequenceModule class to accept transposed parameter, handle state, absorb kwargs"""
95
96	# https://stackoverflow.com/a/65470430/1980685
NEW 97	@functools.wraps(module, updated=())	×
NEW 98	class TransposedModule(module):	×
NEW 99	def __init__(self, args, transposed=False, *kwargs):	×
NEW 100	super().__init__(args, *kwargs)	×
NEW 101	self.transposed = transposed	×
102
NEW 103	def forward(self, x, state=None, **kwargs):	×
NEW 104	if self.transposed:	×
NEW 105	x = x.transpose(-1, -2)	×
NEW 106	x, next_state = super().forward(x, state) # Don't use kwarg because nn.LSTM	×
NEW 107	next_state = None if state is None else next_state	×
NEW 108	if self.transposed:	×
NEW 109	x = x.transpose(-1, -2)	×
NEW 110	return x, next_state	×
111
112	# https://stackoverflow.com/questions/5352781/how-to-set-class-names-dynamically
113	# TransposedModule.__name__ = module.__name__ # functools wraps is better solution
NEW 114	return TransposedModule	×
115
116
NEW 117	@TransposedModule	×
NEW 118	class SequenceIdentity(SequenceModule):	×
119	"""Simple SequenceModule for testing purposes"""
120
NEW 121	def __init__(self, d_model, dropout=0.0, **kwargs):	×
122	"""Default interface for SequenceModule
123
124	d_model: input dimension (sometimes denoted H for hidden dimension)
125	transposed: if True, inputs have axis ordering (B, H, L) instead of (B, H, L)
126	"""
NEW 127	super().__init__()	×
NEW 128	self.d_model = d_model	×
NEW 129	self.d_output = d_model	×
130
NEW 131	def forward(self, x, state=None):	×
NEW 132	return x, state	×
133
NEW 134	def default_state(self, *batch_shape, device=None):	×
NEW 135	return None	×
136
NEW 137	def step(self, x, state=None, **kwargs):	×
NEW 138	return x, state	×

skinniderlab / CLM / 20309855018

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous