Batching Strategies

Online learning algorithms need to handle batched data efficiently. In braintrace, there are two main batching strategies:

• Map-based batching (recommended): Wrap single-sample model logic with brainstate.nn.Map, then compile from one complete batched time step.

• Single-sample mode: Process one sample at a time, without any batching.

The choice of strategy affects how model states are initialized and how the online learning algorithm is called.

This tutorial walks through each strategy with concrete examples and shows how to build a full training loop using Map-based batching.

Map-Based Batching (Recommended)

Keep the model’s update logic single-sample and let brainstate.nn.Map manage independent state copies across the batch. Set the mapped learner up explicitly:

1. Create exactly one brainstate.nn.Map(model, init_map_size=B).

2. Initialize it with mapped_model.init_all_states().

3. Construct the online-learning algorithm with the mapped model.

4. Compile the ETP graph from one complete batched time step with shape (batch_size, n_in).

Do not pass mapped_model to braintrace.compile(..., vmap=True). That setup path owns its batching transformation and would map an already mapped model a second time. When using an explicit Map, construct and compile the algorithm directly as shown below.

import jax
import jax.numpy as jnp
import brainstate
import braintools
import braintrace

class SimpleGRU(brainstate.nn.Module):
    def __init__(self, n_in, n_rec, n_out):
        super().__init__()
        self.rnn = braintrace.nn.GRUCell(n_in, n_rec)
        self.out = braintrace.nn.Linear(n_rec, n_out)

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

model = SimpleGRU(10, 64, 5)
batch_size = 16
example_input = jnp.zeros((batch_size, 10))

# Create and initialize exactly one mapped model.
mapped_model = brainstate.nn.Map(model, init_map_size=batch_size)
mapped_model.init_all_states()

# Compile the algorithm directly from the complete batched example input.
mapped_algo = braintrace.D_RTRL(mapped_model)
mapped_algo.compile_graph(example_input)

x_batch = jnp.ones((batch_size, 10))
out = mapped_algo(x_batch)
print("Output shape:", out.shape)  # (16, 5)

How it works:

• brainstate.nn.Map(model, init_map_size=B) creates the mapped state owner; mapped_model.init_all_states() initializes its independent recurrent states.

• The algorithm receives that mapped model and compiles against the complete batched example input, keeping the batch axis inside the graph.

• Each learner call maps the wrapped model over axis 0 while sharing parameter states and maintaining independent recurrent states.

• The Map is created once. It is not passed to another API that would wrap it again.

Single-Sample Mode

For debugging or situations where batch processing is unnecessary, you can compile and run the algorithm on individual samples directly. No vmap or state replication is needed.

model2 = SimpleGRU(10, 64, 5)

# Single-sample mode: omit batch_size so states are created unbatched.
algo2 = braintrace.compile(model2, braintrace.D_RTRL, jnp.zeros(10))

# Process one sample at a time
x_single = jnp.ones(10)
out = algo2(x_single)
print("Single sample output shape:", out.shape)  # (5,)

This mode is straightforward: initialize the model, compile the graph, and call the algorithm. It is useful for step-by-step debugging or when processing a single stream of data.

Multi-Step Data

braintrace provides SingleStepData and MultiStepData wrappers to control how the algorithm processes input along the time dimension.

• SingleStepData: Wraps data for a single time step. The algorithm processes it as one forward pass.

• MultiStepData: Wraps a sequence of time steps. The algorithm internally scans over all steps in the sequence.

This is useful when you want to pass an entire sequence to the algorithm and have it handle the temporal loop internally, rather than manually iterating over time steps.

# Single-step: process one time step at a time
x_single = braintrace.SingleStepData(jnp.ones(10))

# Multi-step: process a sequence
sequence = jnp.ones((20, 10))  # 20 time steps, 10 features
x_multi = braintrace.MultiStepData(sequence)

When a MultiStepData object is passed to the algorithm, it will iterate over the first axis (time steps) internally. When a SingleStepData object (or a plain array) is passed, the algorithm processes it as a single forward step.

Full Training Loop with Map Batching

Below is a complete example that combines Map-based batching with a temporal training loop. The pattern is:

1. Map and initialize independent model states across the batch.

2. Compile the algorithm from one batched time step.

3. Drive the sequence with etrace_grad, which walks the time axis and accumulates the per-step online gradients for you — no hand-written scan.

4. Update parameters with the accumulated gradients.

@brainstate.transform.jit
def train_step(inputs, targets):
    """inputs: (n_steps, batch_size, n_in), targets: (batch_size,)"""
    def step_loss(inp):
        out = mapped_algo(inp)
        return jnp.mean((out - targets) ** 2)

    # etrace_grad drives the sequence and accumulates per-step online
    # gradients. The explicitly mapped model keeps the batch axis inside the
    # compiled graph.
    return mapped_algo.etrace_grad(
        inputs, step_fn=step_loss, reduction='sum'
    )

# Example usage
model = SimpleGRU(10, 64, 5)
inputs = jnp.ones((20, 16, 10))  # 20 steps, batch 16, 10 features
targets = jnp.zeros((16, 5))

mapped_model = brainstate.nn.Map(model, init_map_size=inputs.shape[1])
mapped_model.init_all_states()
mapped_algo = braintrace.D_RTRL(mapped_model)
mapped_algo.compile_graph(inputs[0])

grads = train_step(inputs, targets)
print("Gradient keys:", list(grads.keys()))

What happens in train_step:

1. The setup creates one brainstate.nn.Map, initializes it, constructs D_RTRL(mapped_model), and compiles from the complete batched time step.

2. mapped_algo.etrace_grad iterates over time, calls step_fn, and accumulates online gradients; reduction='sum' accumulates without dividing.

3. The returned gradient keys match mapped_algo.param_states. Register those learner parameter states with the optimizer when applying the gradients.

Summary

• Create one brainstate.nn.Map(model, init_map_size=B) for batched online learning and call mapped_model.init_all_states().

• Construct the algorithm with that mapped model and call compile_graph(example_input) using one complete batched time step.

• Never pass an already mapped model to compile(..., vmap=True); doing so would apply batching twice.

• After compilation, drive the time axis with etrace_grad (gradients) or etrace_evolve (state/trace only) rather than writing the scan yourself.

• For one stream, initialize and compile the original model without Map.