Neural-Network Layers

braintrace.nn mirrors a subset of brainstate.nn, but routes each layer’s trainable forward pass through ETP primitives. As a result, the layers are drop-in replacements whose parameters automatically participate in online learning — no manual wiring required.

For example, Linear uses braintrace.matmul() internally, so its weight is included in eligibility-trace computation; GRUCell uses braintrace.matmul() for its recurrent maps and braintrace.element_wise() for gate operations.

Note

Activation, normalization, and pooling layers are intentionally not re-implemented here. Accessing them through braintrace.nn (e.g. braintrace.nn.LayerNorm) emits a DeprecationWarning and forwards to brainstate.nn / brainpy.state; use those packages directly.

Linear Layers

Linear	Linear transformation layer.
GroupedLinear	Block-diagonal (grouped) linear layer backed by ETP braintrace.grouped_matmul().
SignedWLinear	Linear layer with signed absolute weights.
ScaledWSLinear	Linear layer with weight standardization.
SparseLinear	Linear layer with sparse weight matrix.
LoRA	A standalone LoRA layer.

Embedding Layers

Embedding

A lookup table whose gather is routed through the ETP embedding op.

Convolutional Layers

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

Recurrent Layers

Single-step recurrent cells. Each updates its hidden state in place and returns the new hidden state (or, for LRUCell, the projected output).

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.
CFNCell	Chaos Free Networks (CFN) cell.

Readout Layers

LeakyRateReadout

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