popgym.baselines.models.lmu

Module Contents

Classes

LMUCell	LMU Cell.
LMU	LMU layer.
LMUFFT	Parallelized LMU Layer.

Functions

leCunUniform(tensor)

LeCun Uniform Initializer

popgym.baselines.models.lmu.leCunUniform(tensor)

LeCun Uniform Initializer

class popgym.baselines.models.lmu.LMUCell(input_size, hidden_size, memory_size, theta, learn_a=False, learn_b=False)

Bases: torch.nn.Module

LMU Cell.

Parameters:

input_size (int) :

Size of the input vector (x_t)

hidden_size (int) :

Size of the hidden vector (h_t)

memory_size (int) :

Size of the memory vector (m_t)

theta (int) :

The number of timesteps in the sliding window that is represented using the LTI system

learn_a (boolean) :

Whether to learn the matrix A (default = False)

learn_b (boolean) :

Whether to learn the matrix B (default = False)

initParameters()

Initialize the cell’s parameters.

stateSpaceMatrices(memory_size, theta)

Returns the discretized state space matrices A and B.

forward(x, state)

Parameters:

x (torch.tensor):

Input of size [batch_size, input_size]

state (tuple):

h (torch.tensor) : [batch_size, hidden_size] m (torch.tensor) : [batch_size, memory_size]

class popgym.baselines.models.lmu.LMU(input_size, hidden_size, memory_size, theta, learn_a=False, learn_b=False)

Bases: torch.nn.Module

LMU layer.

Parameters:

input_size (int) :

Size of the input vector (x_t)

hidden_size (int) :

Size of the hidden vector (h_t)

memory_size (int) :

Size of the memory vector (m_t)

theta (int) :

The number of timesteps in the sliding window that is represented using the LTI system

learn_a (boolean) :

Whether to learn the matrix A (default = False)

learn_b (boolean) :

Whether to learn the matrix B (default = False)

forward(x, state=None)

Parameters:

x (torch.tensor):

Input of size [batch_size, seq_len, input_size]

state (tuple)(default = None)

h (torch.tensor) : [batch_size, hidden_size] m (torch.tensor) : [batch_size, memory_size]

class popgym.baselines.models.lmu.LMUFFT(input_size, hidden_size, memory_size, seq_len, theta)

Bases: torch.nn.Module

Parallelized LMU Layer.

Parameters:

input_size (int) :

Size of the input vector (x_t)

hidden_size (int) :

Size of the hidden vector (h_t)

memory_size (int) :

Size of the memory vector (m_t)

seq_len (int) :

Size of the sequence length (n)

theta (int) :

The number of timesteps in the sliding window that is represented using the LTI system

stateSpaceMatrices()

Returns the discretized state space matrices A and B.

impulse()

Returns the matrices H and the 1D Fourier transform of H (Equations 23, 26 of the paper)

forward(x)

Parameters:

x (torch.tensor):

Input of size [batch_size, seq_len, input_size]