lrnnx.models.lti.base module

Base class for LTI models.

class LTI_LRNN

Bases: LRNN

Base class for all LTI LRNN models.

Note

LTI models do not support async discretization as that requires time-varying dynamics. For async/event-driven models, use LTV models.

Example

>>> from lrnnx.models.lti import LTI_LRNN
>>> my_lrnn = LTI_LRNN("zoh")
>>> # create dummy input tensor and perform forward pass
>>> # in subclass

__init__(discretization: Literal['zoh', 'bilinear', 'dirac', 'no_discretization'])

Initialize the LTI LRNN base class.

Parameters:

discretization (Literal["zoh", "bilinear", "dirac", "no_discretization"]) – Discretization method to use.

abstractmethod discretize() → tuple[torch.Tensor, torch.Tensor | float, torch.Tensor]

This function discretizes the A, B and C matrices, with a learned step-size delta. This could be done inside the compute_kernel method itself, but doing this explicitly outside allows for more flexibility later.

Returns:

A tuple of tensors representing the

discretized A, B, C matrices, ideally of shapes (B, N), (B, N, H) or float, and (B, H, N) respectively.

Return type:

tuple[torch.Tensor, torch.Tensor | float, torch.Tensor]

abstractmethod compute_kernel() → tuple[torch.Tensor, torch.Tensor]

Computes the convolution kernel for efficient parallel processing.

This function is only relevant for LTI models; for LTV models this will materialize a huge vector in-memory at every timestep, which is not efficient. Reference: https://github.com/kunibald413/aTENNuate/blob/15a27dab00d3bf2c27cbbbc3bd41a3d9196dca1e/attenuate/model.py#L30

Parameters:

*args – Model-specific arguments (e.g., sequence length, discretized matrices). See subclass implementations for details.

Returns:

A tuple containing:

• K : Powers of A matrix (A^0, A^1, …, A^{L-1}), shape (N, L)

• B_norm : Normalized input projection matrix, shape (N, H)

Return type:

tuple[torch.Tensor, torch.Tensor]

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

Performs a single recurrent step of the LTI model.

This method is used for autoregressive inference, where inputs are processed one timestep at a time.

Parameters:

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

• inference_cache (Dict[str, Any]) – Cache dictionary from allocate_inference_cache() containing recurrent state and pre-computed matrices. Updated in-place and returned.

Returns:

A tuple containing:

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

• inference_cache : Updated cache dictionary.

Return type:

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

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

Allocates initial state and caches matrices for efficient inference.

For LTI models, the system matrices (A, B, C) are time-invariant, so they can be pre-computed once and reused for all timesteps during autoregressive generation.

Parameters:

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

• max_seqlen (int, optional) – Maximum sequence length (unused for LTI, kept for interface consistency with LTV models). Defaults to 1.

• dtype (torch.dtype, optional) – Data type for allocated tensors. Defaults to None.

Returns:

Cache dictionary containing initial state and

pre-computed matrices for use in step().

Return type:

Dict[str, Any]