Neural Network Layers for Online Learning

After choosing the parameter operations that participate in online learning, compose them into reusable network blocks. braintrace.nn provides layers whose trainable forwards already use ETP operators. This tutorial explains operation-based selection, surveys the layer families, and builds a recurrent model whose parameter-to-hidden paths remain visible to the compiler.

See the complete Neural Network Layers API for constructor signatures and generated class pages.

Selection is operation-based

A parameter participates in online learning only when the forward path uses an ETP operation. braintrace.nn.Linear, for example, uses braintrace.matmul(); the corresponding brainstate.nn.Linear uses an ordinary JAX operation. The class namespace is convenient, but the operation in the traced graph is what the compiler recognizes.

import brainstate
import jax.numpy as jnp

import braintrace

brainstate.random.seed(23)
linear = braintrace.nn.Linear(3, 2)
linear_output = linear(jnp.ones(3))
print("Linear output shape:", linear_output.shape)

Layer families

Family	Representative API classes	ETP operation
Linear maps	braintrace.nn.Linear, braintrace.nn.GroupedLinear, braintrace.nn.SparseLinear, braintrace.nn.LoRA	braintrace.matmul(), braintrace.sparse_matmul(), braintrace.lora_matmul()
Embeddings	braintrace.nn.Embedding	indexed ETP-aware weight access
Convolutions	braintrace.nn.Conv1d, braintrace.nn.Conv2d, braintrace.nn.Conv3d	braintrace.conv()
Recurrent cells	braintrace.nn.GRUCell, braintrace.nn.LSTMCell, braintrace.nn.MiniGRU, braintrace.nn.MiniLSTM	ETP dense and element-wise operations
Readouts	braintrace.nn.LeakyRateReadout	ETP-aware projection

Activation, normalization, and pooling layers should be imported directly from brainstate.nn; compatibility forwarding through braintrace.nn is deprecated. The Neural Network Layers API is the authoritative list of BrainTrace-owned layers.

Compose and compile a recurrent model

The recurrent layer writes hidden state, while the final Linear layer maps that state to an observable output. braintrace.compile() discovers the ETP relationships and returns the online learner.

class TinySequenceModel(brainstate.nn.Module):
    def __init__(self):
        super().__init__()
        self.rnn = braintrace.nn.MiniGRU(in_size=1, out_size=4)
        self.readout = braintrace.nn.Linear(4, 1)

    def update(self, x):
        return self.readout(self.rnn(x))


model = TinySequenceModel()
sample = jnp.ones(1)
learner = braintrace.compile(
    model,
    braintrace.D_RTRL,
    sample,
)

Run repeated steps with a compiled transform

Use a brainstate.transform loop for repeated model execution. State is carried automatically and the outputs are stacked along the leading time axis.

sequence = jnp.linspace(-1.0, 1.0, 6).reshape(6, 1)
brainstate.nn.reset_all_states(model)
learner.reset_state()
outputs = learner.etrace_evolve(sequence, return_outputs=True)
print("Sequence output shape:", outputs.shape)

Temporal and non-temporal parameters

The recurrent weights influence hidden state and therefore need eligibility traces. The readout weight is still trainable, but its output does not feed a hidden state, so the compiler classifies it as non-temporal and computes its instantaneous gradient directly.

Relation boundaries

Trace every parameter path to hidden state before composing custom blocks. A weight -> weight -> hidden path crosses two trainable ETP operations. The compiler stops at the downstream operation, so the upstream weight is not recorded as an independent ETP relation; recording both would double-count a contribution that per-operation rules cannot represent jointly.

Next steps

Continue with Hidden States for Online Learning, then compile and compare single-layer and stacked models in the RNN Compiler Walkthrough. Keep the Neural Network Layers API open when selecting a concrete layer.