brainstate documentation

brainstate implements a State-based transformation system for programming compilation.

---

Features

State-based Transformation

BrainState provides a concise interface to write State-based programs with composable transformation capabilities.

Neural Network Support

BrainState implements a neural network module system for building and training ANNs/SNNs.

---

Installation

CPU

pip install -U brainstate[cpu]

GPU

pip install -U brainstate[cuda12]
pip install -U brainstate[cuda13]

TPU

pip install -U brainstate[tpu]

---

See also the ecosystem

brainstate is one part of our brain simulation ecosystem.

---

Quick Start

Wrap mutable arrays in a State, build models from brainstate.nn.Module, and compose JAX transformations that track state automatically:

import brainstate
import jax.numpy as jnp


class Linear(brainstate.nn.Module):
    def __init__(self, din, dout):
        super().__init__()
        self.w = brainstate.ParamState(brainstate.random.randn(din, dout) * 0.1)
        self.b = brainstate.ParamState(jnp.zeros(dout))

    def __call__(self, x):
        return x @ self.w.value + self.b.value


model = Linear(3, 2)
params = model.states(brainstate.ParamState)

def loss_fn(x, y):
    return jnp.mean((model(x) - y) ** 2)

@brainstate.transform.jit
def train_step(x, y):
    grads = brainstate.transform.grad(loss_fn, grad_states=params)(x, y)
    for key, grad in grads.items():
        params[key].value -= 0.1 * grad
    return loss_fn(x, y)

x = brainstate.random.randn(16, 3)
y = brainstate.random.randn(16, 2)
for step in range(100):
    loss = train_step(x, y)

Parameters are the only states differentiated; jit, grad, and vmap thread state through automatically, with no manual bookkeeping. For a guided walkthrough, start with Quickstart.