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)

Backward function.

Utility method: register a tensor as a buffer or trainable parameter

Matrix mult A @ u with special case to save memory if u has additional batch dim The batch dimension is assumed to be the second dimension A

Interleave two tensors of same shape

Bottom-up divide-and-conquer version of variable_unroll.

Unroll with variable (in time/length) transitions A. A : ([L], ..., N, N) dimension L should exist iff variable is True u : (L, [B], ..., N)

Fix `stride()` information of input tensor. This addresses some inconsistency in stride information in PyTorch. For a tensor t, if t.size(0) =

Calls the compiled function after checking the input format. This function is called in three different modes. t1 x t2 = r ==> t1 and t2 are n

Bottom-up divide-and-conquer version of unroll.

A : (..., N, N) u : (L, ..., N) output : x (..., N) x[i, ...] = A^{i} @ u[0, ...] + ... + A @ u[i-1, ...] + u[i, ...]

Unroll with variable (in time/length) transitions A with general associative operation A : ([L], ..., N, N) dimension L should exist iff variabl

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

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

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

Bottom-up divide-and-conquer version of variable_unroll. compose is an optional function that defines how to compose A without multiplying by a

Construct C~ from C double_length: current C is for length L, convert it to length 2L

Construct dA and dB for discretized state equation

Computes (I + delta A) u A: (n, n) u: (..., n) delta: (...) or scalar output: (..., n)

Get the attention mask of PAM-Naive

Computes (I - d A)^-1 u

Bottom-up divide-and-conquer version of unroll, implemented iteratively

Same as parallel_unroll_recursive but uses bit reversal for locality.

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

Gather features from PAM's pyramid sequences

Test the time and CUDA memory consumption of normal self attention.

SpeedUp version of Autocorrelation (a batch-normalization style design) This is for the inference phase.

SpeedUp version of Autocorrelation (a batch-normalization style design) This is for the training phase.

A, B transition matrices for different measures. measure: the type of measure legt - Legendre (translated) legs - Legendre (scaled)

Unroll with variable (in time/length) transitions A with general associative operation A : ([L], ..., N) dimension L should exist iff variable i

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