Layers

Reusable model layers and wrappers.

SeqLinear

SeqLinear(input_size: int, output_size: int, hidden_size: int = 100, hidden_layer: int = 1, act=Mish, batch_first: bool = True)

Bases: Module

Pointwise MLP applied independently at each sequence position.

Maps the last (feature) dimension through an MLP while preserving all leading dimensions, so [batch, seq, features] in gives [batch, seq, output_size] out (2-D and higher-rank inputs work too). Implemented with nn.Linear, which broadcasts over leading dims since PyTorch 0.4; earlier versions used Conv1d(kernel_size=1) and those checkpoints still load (see _load_from_state_dict).

Parameters:

Name	Type	Description	Default
input_size	int	number of input features	required
output_size	int	number of output features	required
hidden_size	int	number of hidden units per layer	100
hidden_layer	int	number of hidden layers (0 = a single linear map, no activation)	1
act		activation function class	Mish
batch_first	bool	has no effect (nn.Linear preserves leading-dim order); retained for API compatibility.	True

Source code in tsfast/models/_core/layers.py

def __init__(
    self,
    input_size: int,
    output_size: int,
    hidden_size: int = 100,
    hidden_layer: int = 1,
    act=Mish,
    batch_first: bool = True,
):
    super().__init__()
    self.batch_first = batch_first

    def lin_act(inp, out):
        return nn.Sequential(nn.Linear(inp, out), act())

    if hidden_layer < 1:
        self.lin = nn.Linear(input_size, output_size)
    else:
        self.lin = nn.Sequential(
            lin_act(input_size, hidden_size),
            *[lin_act(hidden_size, hidden_size) for _ in range(hidden_layer - 1)],
            nn.Linear(hidden_size, output_size),
        )

AR_Model

AR_Model(model: Module, ar: bool = True, model_has_state: bool = False, return_state: bool = False, out_sz: int | None = None)

Bases: Module

Autoregressive model container.

Runs autoregressively when the output sequence is not provided, otherwise uses teacher forcing. Normalization should be handled externally via ScaledModel wrapping.

Parameters:

Name	Type	Description	Default
model	Module	inner model to wrap	required
ar	bool	if True, default to autoregressive mode in forward	True
model_has_state	bool	if True, the inner model accepts and returns hidden state	False
return_state	bool	if True, return (output, hidden_state) tuple	False
out_sz	int | None	output feature size, used to initialize autoregressive seed	None

Source code in tsfast/models/_core/layers.py

def __init__(
    self,
    model: nn.Module,
    ar: bool = True,
    model_has_state: bool = False,
    return_state: bool = False,
    out_sz: int | None = None,
):
    super().__init__()
    self.model = model
    self.ar = ar
    self.model_has_state = model_has_state
    self.return_state = return_state
    self.out_sz = out_sz
    if return_state and not model_has_state:
        raise ValueError("return_state=True requires model_has_state=True")

SeqAggregation

SeqAggregation(func: Callable = lambda x, dim: x.select(dim, -1), dim: int = 1)

Bases: Module

Aggregation layer that reduces the sequence dimension.

Parameters:

Name	Type	Description	Default
func	Callable	aggregation function taking (tensor, dim) and returning reduced tensor	lambda x, dim: select(dim, -1)
dim	int	sequence dimension to aggregate over	1

Source code in tsfast/models/_core/layers.py

def __init__(
    self,
    func: Callable = lambda x, dim: x.select(dim, -1),
    dim: int = 1,
):
    super().__init__()
    self.func = func
    self.dim = dim

forward

forward(x: Tensor) -> torch.Tensor

Apply the aggregation function to the input tensor.

Source code in tsfast/models/_core/layers.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    "Apply the aggregation function to the input tensor."
    return self.func(x, dim=self.dim)