braintrace documentation

braintrace implements scalable online learning for recurrent neural networks (RNNs) and spiking neural networks (SNNs) using eligibility trace propagation (ETP).

The key idea: mark weight operations with ETP primitives (braintrace.matmul, braintrace.conv, etc.) to include them in online learning. Regular JAX operations are automatically excluded — no special parameter classes needed.

---

Basic Usage

import braintrace
import brainstate

class MyRNN(brainstate.nn.Module):
    def __init__(self):
        super().__init__()
        self.rnn = braintrace.nn.GRUCell(10, 64)
        self.out = braintrace.nn.Linear(64, 10)

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

model = MyRNN()
example_input = brainstate.random.randn(1, 10)   # (batch_size, n_in)

# One call initialises states, compiles the graph, and returns a ready learner.
# The example input carries the batch axis to match ``batch_size=1``.
learner = braintrace.compile(model, braintrace.D_RTRL, example_input, batch_size=1)
y = learner(example_input)

# Now use brainstate.transform.grad as usual — gradients are
# computed online via eligibility traces, not BPTT.

---

Installation

CPU

pip install -U braintrace[cpu]

GPU (CUDA 12.0)

pip install -U braintrace[cuda12]

TPU

pip install -U braintrace[tpu]

---

See also the ecosystem

braintrace is part of the brain simulation ecosystem.

---