memorax.networks.sequence_models

Sequence models for temporal processing.

RNN Models

RNN - Wrapper for Flax RNN cells (LSTM, GRU, etc.).

sLSTMCell - Scalar LSTM cell with enhanced gating.

SHMCell - Stable Hadamard Memory cell.

Memoroid Models

Memoroid - Wrapper for parallel-scannable sequence models.

MemoroidCellBase - Base class for memoroid cells.

MambaCell - Selective State Space Model cell.

S5Cell - Simplified Structured State Space cell.

LRUCell - Linear Recurrent Unit cell.

MinGRUCell - Minimal GRU cell.

mLSTMCell - Matrix LSTM cell.

FFMCell - Fast and Forgetful Memory cell.

LinearAttentionCell - Linear attention cell.

Attention

SelfAttention - Multi-head self-attention.

Wrappers

SequenceModelWrapper - Wraps non-recurrent models.

MetaMaskWrapper - Wrapper for meta-learning with episode boundaries.