lrnnx.models.ltv.base module

Base class for LTV (Linear Time-Varying) models.

LTV models have time-varying dynamics, meaning the state transition matrices (A, B, C, etc.) can change at each timestep based on the input. This makes them more expressive than LTI models but prevents the use of FFT-based convolution for training.

Key differences from LTI models: - Cannot use convolution mode (FFT) since the kernel varies per timestep - Support async/event-driven discretization with variable timesteps - Must use scan for both training and inference

class LTV_LRNN[source]

Bases: LRNN

Base class for all LTV (Linear Time-Varying) LRNN models.

Note

LTV models support async discretization for event-driven processing where timesteps between events may vary. This is specified via the integration_timesteps parameter in forward().

Example

>>> from lrnnx.models.ltv import LTV_LRNN
>>> my_lrnn = LTV_LRNN("zoh")
>>> # create dummy input tensor and perform forward pass
>>> # in subclass
__init__(discretization: Literal['zoh', 'bilinear', 'dirac', 'async', 'no_discretization'] | None)[source]

Initialize the LTV LRNN base class.

Parameters:

discretization (Literal["zoh", "bilinear", "dirac", "async", "no_discretization"] | None) – Discretization method to use. Can be one of "zoh", "bilinear", "dirac", "async", "no_discretization", or None for models that handle discretization internally.

abstractmethod forward(x: torch.Tensor, integration_timesteps: torch.Tensor | None = None, lengths: torch.Tensor | None = None, inference_cache: Dict[str, Any] | None = None) → torch.Tensor[source]

Forward pass through the LTV model.

Parameters:

  • x (torch.Tensor) – Input tensor, shape (B, L, H).

  • integration_timesteps (torch.Tensor, optional) – Timesteps for async/event-driven discretization (Reference: https://arxiv.org/abs/2404.18508), shape (B, L). If None, uniform timesteps are assumed. Defaults to None.

  • lengths (torch.Tensor, optional) – Lengths of sequences, shape (B,), required for variable-length sequences or bidirectional models. Defaults to None.

  • inference_cache (dict, optional) – Cache containing states and pre-computed values for efficient autoregressive generation. If provided during inference, enables incremental processing. Defaults to None.

Returns:

Output tensor, same shape as input (x), i.e., (B, L, H).

Return type:

torch.Tensor

abstractmethod step(x: torch.Tensor, inference_cache: Dict[str, Any]) → Tuple[torch.Tensor, Dict[str, Any]][source]

Performs a single recurrent step of the LTV model.

This method is used for autoregressive inference, where inputs are processed one timestep at a time. Unlike LTI models, the dynamics may vary at each step based on the input.

Parameters:

  • x (torch.Tensor) – Input at current timestep, shape (B, 1, H).

  • inference_cache (Dict[str, Any]) – Cache dictionary containing model states. This is the same format returned by allocate_inference_cache(). The cache is updated in-place and also returned for convenience.

Returns:

A tuple containing:

  • y : Output at current timestep, shape (B, 1, H).

  • inference_cache : Updated cache dictionary.

Return type:

tuple[torch.Tensor, Dict[str, Any]]

abstractmethod allocate_inference_cache(batch_size: int, max_seqlen: int, dtype: torch.dtype | None = None) → Dict[str, Any][source]

Allocates cache for efficient autoregressive inference.

For LTV models, this typically includes:

  • Initial hidden state(s)

  • Any auxiliary states (e.g., convolution state for Mamba)

  • Metadata for tracking sequence position

Parameters:

  • batch_size (int) – The batch size for inference.

  • max_seqlen (int) – Maximum sequence length to support.

  • dtype (torch.dtype, optional) – Data type for allocated tensors. If None, uses the model’s default dtype. Defaults to None.

Returns:

Cache dictionary that can be passed to forward().

Should contain at minimum: - Model state tensors (e.g., “lrnn_state”, “conv_state”) - “seqlen_offset”: Current position in the sequence

Return type:

Dict[str, Any]