Returns a linear nn.Module with control over axes order, initialization, and activation

d_state: the dimension of the state, also denoted by N l_max: the maximum kernel length, also denoted by L. Set l_max=None to always

Register a tensor with a configurable learning rate and 0 weight decay

d_state: the dimension of the state, also denoted by N l_max: the maximum sequence length, also denoted by L if this is not

Construct dA and dB for discretized state equation

Version of the step function that has time O(N) instead of O(N^2) per step, which takes advantage of the DPLR form and bilinear discretizatio

Args. complex_spec: torch.Tensor, shape = [B, D, T] length: int Return: wav: torch.Tensor, shape

Compute A^L and the scan sum_i A^i v_i A: (..., N, N) v: (..., N, L)

Args. wav: torch.Tensor, shape = [B, T] Return: complex_spec: torch.Tensor, shape = [B, D, T]

Construct C~ from C Two modes are supported: go directly to length L if self.L is 1, or length is doubled

Return w, p, q, V, B such that (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V i.e. A = V[w - p q^*]V^*, B

Calculate (and cache) FFT nodes and their "unprocessed" version with the bilinear transform This should be called everytime the internal leng

Create parameters that allow fast linear stepping of state

Must be called after self.default_state() is used to construct an initial state!

Args. outputs: torch.Tensor [B, C, D, T] noisy_features: Dict Return. torch.Tensor [B, C, D, T]

Pykeops version

v, w: (..., N) z: (..., L) returns: (..., L)

Forward the state through a sequence, i.e. computes the state after passing chunk through SSM state: (B, H, N) u: (B, H, L) R

Initializes multi-GPU-friendly python logger.

u: ... H L v: ... H N x: ... H N Returns: ... H N V = Vandermonde(a, L) : (H N L) contract_L(V * u *

x: (..., L1, L2, .., Lk) L: list of length k [l1, l2, .., lk] returns: x[..., :l1, :l2, .., :lk]

Return low-rank matrix L such that A + L is normal

Args. complex_spec: torch.Tensor, shape = [B, D, T] Return: complex_spec: torch.Tensor, shape = [B, D, T]

Args. complex_spec: torch.Tensor, shape = [B, D, T] Return: complex_spec: torch.Tensor, shape = [B, D, T]

Dispatcher to create single SSM initialization N: state size R: rank (for DPLR parameterization) H: number of independent SSM copies

Step one time step as a recurrent model. Intended to be used during validation. u: (B H) state: (B H N) Returns: output (B H)

Implements the Stochastic Depth from `"Deep Networks with Stochastic Depth" <https://arxiv.org/abs/1603.09382>`_ used for randomly dropping r

A, B transition matrices for different measures

L: Maximum length; this module computes an SSM kernel of length L A is represented by diag(w) - PP^* w: (S, N) diagonal part

State Space Kernel which computes the convolution kernel $\\bar{K}$ H: Number of independent SSM copies; controls the size of the model. Also

Feed-forward block. Args: d_model: dimension of input expand: expansion factor for inverted bottleneck

Residual S4 block. Args: d_model: dimension of the model layer: a layer config dropout: dropout r

Feed-forward block. Args: d_model: dimension of input ff_expand_factor: expansion factor for inverted bottlen

Residual S4 block. Args: d_model: dimension of the model layer: a layer config dropout: dropout r

tie: tie dropout mask across sequence lengths (Dropout1d/2d/3d)

Sets time kernel to custom value

Set up dA, dB, dC discretized parameters for stepping

input: (batch, d_model, freq, length) output: (batch, d_model, freq, length)

state: (B, H, N) initial state rate: sampling rate factor L: target length returns: (C, H, L) convolution ker

state: (B, H, N) initial state rate: sampling rate factor L: target length returns: (C, H, L) convolution ker

u: (B H L) if self.transposed else (B L H) state: (H N) never needed unless you know what you're doing Returns: same shape as