lrnnx.core.discretization module

Discretization methods for continuous-time systems.

zoh(A: torch.Tensor, delta: torch.Tensor, integration_timesteps: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor]

Zero-Order Hold (ZOH) discretization method, used across most models.

\[\begin{split}\bar{A} &= \exp(\Delta A) \\ \bar{\gamma} &= A^{-1} (\bar{A} - I)\end{split}\]

Reference: https://hazyresearch.stanford.edu/blog/2022-01-14-s4-3

Parameters:

• A (torch.Tensor) – The continuous-time state matrix.

• delta (torch.Tensor) – The discretization step size.

• integration_timesteps (torch.Tensor, optional) – Not used in ZOH discretization. Defaults to None.

Returns:

A tuple containing:

• A_bar : The discretized system matrix.

• gamma_bar : The input normalizer.

Return type:

tuple[torch.Tensor, torch.Tensor]

bilinear(A: torch.Tensor, delta: torch.Tensor, integration_timesteps: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor]

Bilinear method first used in S4.

\[\begin{split}\bar{A} &= (I + 0.5 \Delta A)^{-1} (I - 0.5 \Delta A) \\ \bar{\gamma} &= (I + 0.5 \Delta A)^{-1} \Delta\end{split}\]

Reference: https://hazyresearch.stanford.edu/blog/2022-01-14-s4-3

Parameters:

• A (torch.Tensor) – Continuous-time system matrix, shape: (N,), i.e., only diagonal elements.

• delta (torch.Tensor) – Time step for discretization.

• integration_timesteps (torch.Tensor, optional) – Not used in bilinear discretization. Defaults to None.

Returns:

A tuple containing:

• A_bar : The discretized system matrix.

• gamma_bar : The input normalizer.

Return type:

tuple[torch.Tensor, torch.Tensor]

dirac(A: torch.Tensor, delta: torch.Tensor, integration_timesteps: torch.Tensor | None = None) → tuple[torch.Tensor, float]

Dirac discretization method.

\[\begin{split}\bar{A} &= \exp(\Delta A) \\ \bar{\gamma} &= 1.0\end{split}\]

Reference: https://github.com/Efficient-Scalable-Machine-Learning/event-ssm/blob/main/event_ssm/ssm.py

Parameters:

• A (torch.Tensor) – Continuous-time system matrix.

• delta (torch.Tensor) – Time step for discretization.

• integration_timesteps (torch.Tensor, optional) – Not used in dirac discretization. Defaults to None.

Returns:

A tuple containing:

• A_bar : The discretized system matrix.

• gamma_bar : The input normalizer (1.0).

Return type:

tuple[torch.Tensor, float]

async_(A: torch.Tensor, delta: torch.Tensor, integration_timesteps: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor]

Asynchronous discretization method, introduced in https://arxiv.org/abs/2404.18508. This helps provide a strong inductive bias for asynchronous event-streams.

\[\begin{split}\bar{A} &= \exp(\Delta \cdot \text{integration\_timesteps} \cdot A) \\ \bar{\gamma} &= A^{-1} (\exp(\Delta A) - I)\end{split}\]

This method is only for legacy reasons; it is not possible to use this method (or any other discretization with async timesteps) with LTI models.

Parameters:

• A (torch.Tensor) – Continuous-time system matrix.

• delta (torch.Tensor) – Time step for discretization.

• integration_timesteps (torch.Tensor) – Timesteps for async discretization, ideally of shape (B, L), i.e., difference in timesteps of events.

Returns:

A tuple containing:

• A_bar : The discretized system matrix.

• gamma_bar : The input normalizer.

Return type:

tuple[torch.Tensor, torch.Tensor]

no_discretization(A: torch.Tensor, delta: torch.Tensor, integration_timesteps: torch.Tensor | None = None) → tuple[torch.Tensor, float]

No discretization method, identity operation.

\[\begin{split}\bar{A} &= A \\ \bar{\gamma} &= 1.0\end{split}\]

Parameters:

• A (torch.Tensor) – Continuous-time system matrix.

• delta (torch.Tensor) – Time step for discretization (unused).

• integration_timesteps (torch.Tensor, optional) – Not used in no_discretization. Defaults to None.

Returns:

A tuple containing:

• A_bar : Same as A.

• gamma_bar : 1.0, as B_bar = B.

Return type:

tuple[torch.Tensor, float]