Imported from chaobrain/BrainX-skill (brainx-display-cases/creative-experiment-verification/05-alpha-rhythm/run0/.agents/skills/brainstate/SKILL.md). Install upstream with npx skills add chaobrain/BrainX-skill --skill brainstate. Copyright stays with the author.

Brainstate

Purpose and boundary

Use this skill for BrainState's general stateful workflow: mutable State, registered brainstate.nn.Module graphs, environment-scoped simulation, state initialization and collection, randomness, state-aware transforms, and simulation or training.

Canonical path:

classify State roles -> construct Modules -> register State and children -> initialize -> set the run environment -> transform the whole operation -> validate State and outputs

Open the routed reference for specialized variants; route specialized neuronal or network dynamics to the matching BrainCell or BrainPy-State skill.

Underlying principle of BrainState

BrainState reconciles mutable models with JAX by making mutation explicit: store every value that changes during transformed execution in State. Its State roles match distinct model lifecycles; for example, HiddenState holds voltages, firing rates, and other dynamics updated each simulation step.

Wrap the complete stateful operation in brainstate.transform. Unlike raw JAX transforms, its JAX-like jit, grad, and vmap discover State reads and writes, differentiate selected State collections, and map or share State during vectorization.

API structure

Choose the namespace that owns the operation:

1. State is the mutation boundary

State is a typed, mutable container for an array or stable PyTree; read and replace it through .value, including inside BrainState transformations. Treat it as a fixed-structure slot: keep static configuration in ordinary attributes and preserve value type, shape, dtype, and PyTree structure across writes.

Create scalar, array, and pytree state

import brainstate
import brainstate.nn as nn
import brainunit as u
import jax.numpy as jnp

counter = brainstate.State(jnp.array(0))
vector = brainstate.State(jnp.zeros(10))
neuron = brainstate.State({
    "V": jnp.zeros(5),
    "u": jnp.ones(5),
})

value = neuron.value
neuron.value = {
    "V": value["V"] + 0.1,
    "u": value["u"],
}

with brainstate.check_state_value_tree():
    neuron.value = {
        "V": jnp.zeros(5),
        "u": jnp.ones(5),
    }

Substate

State subclasses remain State containers but act as semantic markers for filtering and model organization. Use these non-parameter roles for hidden or runtime values:

h = brainstate.HiddenState(jnp.zeros(5))
last_spike = brainstate.ShortTermState(jnp.full(5, -1e7))
running_mean = brainstate.LongTermState(jnp.zeros(5))

2. Modules form registered state graphs

Subclass brainstate.nn.Module, assign each State and child Module to an attribute, and implement computation in update(); assignment registers a traversable graph whose State leaves can be filtered by role:

params = model.states(brainstate.ParamState)

Add state to a module

class Counter(brainstate.nn.Module):
    def __init__(self):
        super().__init__()
        self.count = brainstate.ShortTermState(jnp.array(0))

    def update(self, x):
        self.count.value = self.count.value + 1
        return x * self.count.value


counter_module = Counter()
for _ in range(5):
    print(counter_module(jnp.array(10.0)), counter_module.count.value)

The output advances through (10, 1), (20, 2), ..., (50, 5) because the registered State is updated explicitly.

Use basic prebuilt layers

brainstate.random.seed(42)

linear = nn.Linear(in_size=(10,), out_size=(5,))
relu = nn.ReLU()
conv = nn.Conv2d(
    in_size=(28, 28, 3),
    out_channels=32,
    kernel_size=3,
    padding="SAME",
)

x = brainstate.random.randn(8, 10)
y = relu(linear(x))
assert y.shape == (8, 5)

Open the layer and activation catalogs instead of guessing an uncommon class name or signature.

Parameter

nn.Param wraps an underlying ParamState with an optional constraint transform and regularizer; model code reads .value(), while optimizers update .val.

Use nn.SoftplusT(lower=L) for a value constrained to (L, infinity); .value() stays in range while .val remains unconstrained. Do not confuse parameter transforms with execution transforms such as brainstate.transform.jit.

w = brainstate.ParamState(brainstate.random.randn(10, 5) * 0.1)
b = brainstate.ParamState(jnp.zeros(5))

gain = nn.Param(jnp.array(1.0))
positive_tau = nn.Param(jnp.array(2.0), t=nn.SoftplusT(lower=0.1))
fixed_scale = nn.Const(jnp.array(10.0))

print(gain.value())          # IdentityT by default; usable and stored values coincide
print(positive_tau.value())  # constrained model value, always greater than 0.1
print(positive_tau.val)      # underlying unconstrained ParamState
print(fixed_scale.value())   # fixed value, excluded from ParamState collection

3. Size inference drives composition

Size-aware Modules carry feature-shape metadata without the batch dimension, so composition code can construct and validate each next layer before execution.

model = nn.Sequential(
    nn.Linear(in_size=(10,), out_size=(8,)),
    nn.ReLU(),
    nn.Linear.desc(out_size=(2,)),
)

x = brainstate.random.randn(4, 10)
y = model(x)

assert model.layers[2].in_size == (8,)
assert model.out_size == (2,)
assert y.shape == (4, 2)

Open references/size-inference-variations.md to compose ComplexNet with Sequential and .desc(), or when convolution, pooling, and flattening make size propagation non-obvious.

4. Environment context drives simulations

Use brainstate.environ to share run settings across a computation, where scoped contexts override persistent defaults and model code reads the active value.

with brainstate.environ.context(dt=0.1 * u.ms, fit=False):
    brainstate.nn.init_all_states(net)
    times = u.math.arange(
        0.0 * u.ms,
        100.0 * u.ms,
        brainstate.environ.get_dt(),
    )
    step_indices = jnp.arange(times.shape[0])

    def step(t, i):
        with brainstate.environ.context(t=t, i=i):
            return net.update(input_current)

    outputs = brainstate.transform.for_loop(
        step,
        times,
        step_indices,
    )

Open references/simulation-environment.md for persistent defaults, generic setting access, nested or isolated environments, precision and platform controls, exp_euler_step(), and environment-specific failures.

5. State-aware transforms

Wrap the complete forward, simulation, or training step in brainstate.transform; its JAX-like jit, grad, and vmap track State reads and writes, while raw JAX transforms can lose State mutations.

Canonical transformation setup

brainstate.random.seed(0)
model = nn.Linear(in_size=(3,), out_size=(1,))
x_train = brainstate.random.randn(64, 3)
y_train = brainstate.random.randn(64, 1)
params = model.states(brainstate.ParamState)

State-aware jit example

forward = brainstate.transform.jit(model)
prediction = forward(x_train)
assert prediction.shape == (64, 1)

The first compatible call traces and compiles the complete forward pass; later compatible calls reuse it while BrainState handles State effects.

Gradient and parameter update example

grad differentiates with respect to a State collection and returns gradients keyed by the same State paths. return_value=True returns the loss from the same pass.

def loss_fn():
    return jnp.mean((model(x_train) - y_train) ** 2)


grads, loss = brainstate.transform.grad(
    loss_fn,
    params,
    return_value=True,
)()

for key in params:
    params[key].value -= 0.1 * grads[key]

Composed training-step transform

@brainstate.transform.jit
def train_step():
    grads, loss = brainstate.transform.grad(
        loss_fn,
        params,
        return_value=True,
    )()
    for key in params:
        params[key].value -= 0.1 * grads[key]
    return loss


loss = train_step()

Use jit(grad(...)) as the default compiled training-step backbone. Open the optimizer reference when manual updates should become optimizer-managed updates.

State-aware vmap

def predict_one(x_row):
    return model(x_row[None, :])[0]


predict_batch = brainstate.transform.vmap(predict_one)
batched_prediction = predict_batch(x_train)
assert batched_prediction.shape == (64, 1)

This maps a function written for one example over a batch. Open the vmap expansion for mapped State axes, ensembles, or parameter sweeps. Use vmap(..., in_states=..., out_states=...) to declare specific mapped State instances; use vmap2(..., state_in_axes=..., state_out_axes=...) to select State by filters with configurable output-axis inference. Do not add either transform when the owning package already exposes the required batch axis.

6. Randomness

In BrainX modeling, seed BrainState's stateful generator before model construction or simulation so initialization, stochastic inputs, noise, dropout, and data order follow the same automatically split key sequence.

Generate reproducible stochastic input

brainstate.random.seed(42)

num_steps, batch_size, num_inputs = 100, 128, 100
firing_rate = 5.0 * u.Hz

with brainstate.environ.context(dt=1.0 * u.ms):
    spike_probability = firing_rate * brainstate.environ.get_dt()
    input_spikes = (
        brainstate.random.rand(num_steps, batch_size, num_inputs)
        < spike_probability
    )

assert input_spikes.shape == (num_steps, batch_size, num_inputs)

Open references/brainstate/randomness-and-reproducibility.md for independent RandomState streams, direct key control, stochastic mapping, exact replay, or checkpointed RNG State.

Script references

• scripts/integrator_rnn.py: full stateful sequence-training workflow with a custom RNN cell, trainable initial state, optimization, compiled steps, and evaluation.
• scripts/lif_neuron_model.py: extended combination of HiddenState, ShortTermState, and ParamState with explicit .value updates.
• scripts/modern_cnn.py: full convolution, normalization, activation, pooling, dropout, and dense Module composition. Select it through the layer or activation branch.
• scripts/resnet.py: residual Modules and dynamically registered child blocks.

Reference routing

Route by the outcome the task needs, then open only the smallest reference that owns that variant.

State graphs, collections, and lifecycle

Simulation environment

Model composition, extension, and interoperation

Parameters, optimization, and randomness

The remaining nested reference has one inbound route:

• Only parameter-constraints-regularization.md may open parameter-transforms-regularizers-catalog.md.

Layer libraries

Stateful transformations

Do not route to dynamics or solver references from this skill; they are outside the architecture supplied for this BrainState skill.