Neural Network Modules

braintrace.nn provides neural network layers that use ETP primitives internally. These layers are drop-in replacements for standard layers but automatically participate in online learning.

For example, braintrace.nn.Linear uses braintrace.matmul internally, so its weight is automatically included in eligibility trace computation. Similarly, braintrace.nn.GRUCell uses braintrace.matmul for its recurrent weight and braintrace.element_wise for gate operations.

Linear Layers

Linear	Linear transformation layer.
SignedWLinear	Linear layer with signed absolute weights.
SparseLinear	Linear layer with sparse weight matrix.
LoRA	A standalone LoRA layer.

Convolutional Layers

Conv1d	One-dimensional convolution layer.
Conv2d	Two-dimensional convolution layer.
Conv3d	Three-dimensional convolution layer.

Recurrent Layers

ValinaRNNCell	Vanilla RNN cell.
GRUCell	Gated Recurrent Unit (GRU) cell.
MGUCell	Minimal Gated Recurrent Unit (MGU) cell.
LSTMCell	Long short-term memory (LSTM) RNN core.
URLSTMCell	Update-Reset LSTM (URLSTM) cell.
MinimalRNNCell	Minimal RNN Cell.
MiniGRU	Minimal GRU cell.
MiniLSTM	Minimal LSTM cell.
LRUCell	Linear Recurrent Unit (LRU) layer.

Normalization Layers

BatchNorm0d	0-D batch normalization.
BatchNorm1d	1-D batch normalization.
BatchNorm2d	2-D batch normalization.
BatchNorm3d	3-D batch normalization.
LayerNorm	Layer normalization layer [1]_.
RMSNorm	Root Mean Square Layer Normalization [1]_.
GroupNorm	Group Normalization layer [1]_.

Readout Layers

LeakyRateReadout

Leaky dynamics for the read-out module used in Real-Time Recurrent Learning.