Migrate examples, docs and docstrings onto compile / etrace_grad / etrace_evolve

braintrace/__init__.py

@@ -41,15 +41,24 @@
    hidden->weight / hidden->hidden Jacobian computations.
 4. **Algorithms** — online-learning orchestrators: the exact algorithms
    :class:`D_RTRL` / :func:`pp_prop` / :class:`ES_D_RTRL`, and the SNN family
-   :class:`EProp`, :class:`OSTLRecurrent`, :class:`OSTLFeedforward`.
+   :class:`EProp`, :class:`OSTLRecurrent`, :class:`OSTLFeedforward`. Every one
+   of them carries the two sequence drivers of
+   :class:`SequenceDriverMixin` — ``etrace_grad`` and ``etrace_evolve`` — so a
+   caller never writes the per-step gradient-accumulation loop by hand.
 
 The :mod:`braintrace.nn` subpackage provides ready-made ETP-wired layers
 (linear maps, convolutions, recurrent cells, read-outs).
 
 Notes
 -----
 The convenience entry point :func:`compile` wraps a model together with an
-algorithm into a single trainable object and is the recommended starting point.
+algorithm into a single trainable object and is the recommended starting point:
+one call replaces ``init_all_states`` + algorithm construction +
+``compile_graph`` (+ the vmap wrapper, with ``vmap=True``). The object it
+returns drives a sequence with :meth:`~SequenceDriverMixin.etrace_grad`
+(accumulate online gradients under a loss) or
+:meth:`~SequenceDriverMixin.etrace_evolve` (advance hidden state and the
+eligibility trace with no loss).
 The ``braintrace.MatMulOp`` / ``ETraceParam`` style names from the v0.1.x API
 are deprecated shims served lazily with a :class:`DeprecationWarning`; new code
 should mark parameters by routing them through ETP ops instead.

braintrace/_algorithm/d_rtrl.py

@@ -80,11 +80,19 @@ class D_RTRL(ParamDimVjpAlgorithm):
 
         >>> import brainstate
         >>> import braintrace
+        >>> import jax.numpy as jnp
         >>>
         >>> model = braintrace.nn.ValinaRNNCell(2, 4, activation='tanh')
         >>> x0 = brainstate.random.randn(2)
         >>> learner = braintrace.compile(model, braintrace.D_RTRL, x0)
         >>> y = learner.update(x0)
+        >>>
+        >>> # etrace_grad drives the sequence and accumulates the online gradients
+        >>> xs = brainstate.random.randn(10, 2)   # (T, ...)
+        >>> ys = brainstate.random.randn(10, 4)
+        >>> def step_loss(x, y):
+        ...     return jnp.mean((learner(x) - y) ** 2)
+        >>> grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
 
     References
     ----------

braintrace/_algorithm/e_prop.py

@@ -151,6 +151,7 @@ class EProp(ParamDimVjpAlgorithm):
 
         >>> import brainstate
         >>> import braintrace
+        >>> import jax.numpy as jnp
         >>>
         >>> class RSNN(brainstate.nn.Module):
         ...     def __init__(self):
@@ -165,6 +166,13 @@ class EProp(ParamDimVjpAlgorithm):
         >>> # one call: initialise states, build the trace graph, return a learner
         >>> learner = braintrace.compile(model, braintrace.EProp, x0, kappa_filter_decay=0.9)
         >>> y = learner(x0)             # forward pass + eligibility-trace update
+        >>>
+        >>> # etrace_grad drives the sequence and accumulates the online gradients
+        >>> xs = brainstate.random.randn(10, 1)   # (T, ...)
+        >>> ys = brainstate.random.randn(10, 1)
+        >>> def step_loss(x, y):
+        ...     return jnp.mean((learner(x) - y) ** 2)
+        >>> grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
 
     References
     ----------

braintrace/_algorithm/io_dim_vjp.py

@@ -753,6 +753,7 @@ class IODimVjpAlgorithm(ETraceVjpAlgorithm):
 
         >>> import brainstate
         >>> import braintrace
+        >>> import jax.numpy as jnp
         >>>
         >>> class RNN(brainstate.nn.Module):
         ...     def __init__(self):
@@ -767,6 +768,13 @@ class IODimVjpAlgorithm(ETraceVjpAlgorithm):
         >>> # one call: initialise states, build the trace graph, return a learner
         >>> learner = braintrace.compile(model, braintrace.pp_prop, x0, decay_or_rank=0.9)  # or rank: decay_or_rank=19
         >>> y = learner(x0)             # forward pass + eligibility-trace update
+        >>>
+        >>> # etrace_grad drives the sequence and accumulates the online gradients
+        >>> xs = brainstate.random.randn(10, 1)   # (T, ...)
+        >>> ys = brainstate.random.randn(10, 1)
+        >>> def step_loss(x, y):
+        ...     return jnp.mean((learner(x) - y) ** 2)
+        >>> grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
 
     References
     ----------

braintrace/_algorithm/ostl.py

@@ -118,6 +118,7 @@ class OSTLRecurrent(ParamDimVjpAlgorithm):
 
         >>> import brainstate
         >>> import braintrace
+        >>> import jax.numpy as jnp
         >>>
         >>> class Net(brainstate.nn.Module):
         ...     def __init__(self):
@@ -132,6 +133,13 @@ class OSTLRecurrent(ParamDimVjpAlgorithm):
         >>> # one call: initialise states, build the trace graph, return a learner
         >>> learner = braintrace.compile(model, braintrace.OSTLRecurrent, x0)
         >>> y = learner(x0)
+        >>>
+        >>> # etrace_grad drives the sequence and accumulates the online gradients
+        >>> xs = brainstate.random.randn(10, 1)   # (T, ...)
+        >>> ys = brainstate.random.randn(10, 1)
+        >>> def step_loss(x, y):
+        ...     return jnp.mean((learner(x) - y) ** 2)
+        >>> grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
 
     References
     ----------
@@ -204,6 +212,7 @@ class OSTLFeedforward(pp_prop):
 
         >>> import brainstate
         >>> import braintrace
+        >>> import jax.numpy as jnp
         >>>
         >>> class Net(brainstate.nn.Module):
         ...     def __init__(self):
@@ -218,6 +227,13 @@ class OSTLFeedforward(pp_prop):
         >>> # one call: initialise states, build the trace graph, return a learner
         >>> learner = braintrace.compile(model, braintrace.OSTLFeedforward, x0)
         >>> y = learner(x0)
+        >>>
+        >>> # etrace_grad drives the sequence and accumulates the online gradients
+        >>> xs = brainstate.random.randn(10, 1)   # (T, ...)
+        >>> ys = brainstate.random.randn(10, 1)
+        >>> def step_loss(x, y):
+        ...     return jnp.mean((learner(x) - y) ** 2)
+        >>> grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
 
     References
     ----------

braintrace/_algorithm/param_dim_vjp.py

@@ -1073,6 +1073,7 @@ class ParamDimVjpAlgorithm(ETraceVjpAlgorithm):
 
         >>> import brainstate
         >>> import braintrace
+        >>> import jax.numpy as jnp
         >>>
         >>> class RNN(brainstate.nn.Module):
         ...     def __init__(self):
@@ -1088,6 +1089,13 @@ class ParamDimVjpAlgorithm(ETraceVjpAlgorithm):
         >>> # initialises states, builds the trace graph, and returns a learner.
         >>> learner = braintrace.compile(model, braintrace.D_RTRL, x0)
         >>> y = learner(x0)             # forward pass + eligibility-trace update
+        >>>
+        >>> # etrace_grad drives the sequence and accumulates the online gradients
+        >>> xs = brainstate.random.randn(10, 1)   # (T, ...)
+        >>> ys = brainstate.random.randn(10, 1)
+        >>> def step_loss(x, y):
+        ...     return jnp.mean((learner(x) - y) ** 2)
+        >>> grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
 
     References
     ----------

braintrace/_algorithm/pp_prop.py

@@ -82,6 +82,7 @@ class pp_prop(IODimVjpAlgorithm):
 
         >>> import brainstate
         >>> import braintrace
+        >>> import jax.numpy as jnp
         >>>
         >>> class RNN(brainstate.nn.Module):
         ...     def __init__(self):
@@ -96,6 +97,13 @@ class pp_prop(IODimVjpAlgorithm):
         >>> # one call: initialise states, build the trace graph, return a learner
         >>> learner = braintrace.compile(model, braintrace.pp_prop, x0, decay_or_rank=0.9)  # or rank: decay_or_rank=19
         >>> y = learner(x0)             # forward pass + eligibility-trace update
+        >>>
+        >>> # etrace_grad drives the sequence and accumulates the online gradients
+        >>> xs = brainstate.random.randn(10, 1)   # (T, ...)
+        >>> ys = brainstate.random.randn(10, 1)
+        >>> def step_loss(x, y):
+        ...     return jnp.mean((learner(x) - y) ** 2)
+        >>> grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
 
     References
     ----------

braintrace/_algorithm/snap_n.py

@@ -160,6 +160,7 @@ class SnAp(ParamDimVjpAlgorithm):
 
         >>> import brainstate
         >>> import braintrace
+        >>> import jax.numpy as jnp
         >>>
         >>> class Net(brainstate.nn.Module):
         ...     def __init__(self):
@@ -173,6 +174,13 @@ class SnAp(ParamDimVjpAlgorithm):
         >>> x0 = brainstate.random.randn(1)
         >>> learner = braintrace.compile(model, braintrace.SnAp, x0, n=2)
         >>> y = learner(x0)
+        >>>
+        >>> # etrace_grad drives the sequence and accumulates the online gradients
+        >>> xs = brainstate.random.randn(10, 1)   # (T, ...)
+        >>> ys = brainstate.random.randn(10, 1)
+        >>> def step_loss(x, y):
+        ...     return jnp.mean((learner(x) - y) ** 2)
+        >>> grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
     """
 
     __module__ = 'braintrace'

braintrace/_algorithm/three_factor.py

@@ -130,6 +130,21 @@ class ThreeFactor(ParamDimVjpAlgorithm):
 
     The keyword takes precedence for the call it appears on.
 
+    Over a sequence, the per-step reward is simply a second sequence:
+    ``etrace_grad`` slices every sequence in lockstep and hands the slices to
+    ``step_fn`` positionally, and ``step_fn`` -- not the driver -- owns the
+    model call, so there is nowhere the modulator has to be threaded through.
+
+    .. code-block:: python
+
+        >>> xs = jnp.zeros((10, 1, 4))
+        >>> rewards = jnp.linspace(-1.0, 1.0, 10)
+        >>> ys = jnp.zeros((10, 1, 4))
+        >>> def step_loss(x, reward, y):
+        ...     return jnp.mean((learner.update(x, modulator=reward) - y) ** 2)
+        >>> grads, losses = learner.etrace_grad(
+        ...     xs, rewards, ys, step_fn=step_loss, return_value=True)
+
     Notes
     -----
     Under single-step, every **plain** (non-ETP) parameter's gradient is exactly

braintrace/_algorithm/uoro.py

@@ -118,6 +118,21 @@ class UORO(RandomProjectionVjpAlgorithm):
         >>> learner.compile_graph(braintrace.MultiStepData(jnp.zeros((1, 1, 4))))
         >>> learner.init_etrace_state()
 
+    UORO is multi-step by construction, so ``etrace_grad`` drives it in **window
+    mode**: pass ``chunk_size=k`` with ``k >= 2``, and ``step_fn`` receives a
+    ``(k, ...)`` slice, wraps its model input in :class:`MultiStepData`, and
+    returns a ``(k,)`` vector of per-step losses rather than a scalar.
+
+    .. code-block:: python
+
+        >>> xs = jnp.zeros((10, 1, 4))
+        >>> ys = jnp.zeros((10, 1, 4))
+        >>> def window_loss(x, y):          # x, y are (k, 1, 4)
+        ...     out = learner(braintrace.MultiStepData(x))
+        ...     return jnp.mean((out - y) ** 2, axis=(1, 2))    # (k,)
+        >>> grads, losses = learner.etrace_grad(
+        ...     xs, ys, step_fn=window_loss, chunk_size=5, return_value=True)
+
     Notes
     -----
     **Variance grows with the number of window boundaries.** The estimate is

braintrace/_compile.py

@@ -253,6 +253,7 @@ def compile(
 
         >>> import brainstate
         >>> import braintrace
+        >>> import jax.numpy as jnp
         >>>
         >>> class RNN(brainstate.nn.Module):
         ...     def __init__(self):
@@ -276,6 +277,16 @@ def compile(
         >>> config = braintrace.ETraceConfig()
         >>> by_config = braintrace.compile(
         ...     model, config, x0, batch_size=1)
+        >>>
+        >>> # Every learner compile returns carries the two sequence drivers.
+        >>> xs = brainstate.random.randn(10, 1, 3)   # (T, batch, features)
+        >>> ys = brainstate.random.randn(10, 1, 1)
+        >>> def step_loss(x, y):
+        ...     return jnp.mean((by_name(x) - y) ** 2)
+        >>> grads, losses = by_name.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
+        >>> # etrace_evolve is the same drive with no loss: it advances hidden
+        >>> # state and the eligibility trace, optionally stacking the outputs.
+        >>> outs = by_name.etrace_evolve(xs, return_outputs=True)
     """
     cls = _resolve_algorithm(algorithm)
     if isinstance(algorithm, ETraceConfig):

docs/advanced/batching.ipynb

@@ -194,7 +194,13 @@
    "cell_type": "markdown",
    "id": "a1b2c3d4e5f60016",
    "metadata": {},
-   "source": "**What happens in `train_step`:**\n\n1. `weights` collects all `ParamState` objects from the model.\n2. `braintrace.compile(model, braintrace.D_RTRL, inputs[0], batch_size=B, vmap=True)` initialises per-sample hidden states, compiles the computation graph using a single time step's batched input (`inputs[0]`, shape `(batch_size, n_in)`), and wraps the result in `brainstate.nn.Vmap` for batch-parallel execution — all in one call.\n3. `brainstate.transform.scan` iterates over the time dimension (`inputs` has shape `(n_steps, batch_size, n_in)`). At each step, `step_fn` computes the loss and its gradients with respect to `weights`, then accumulates them.\n4. The returned `grads` dictionary can be passed to an optimizer (e.g., `braintools.optim.Adam`) for a parameter update."
+   "source": [
+    "**What happens in `train_step`:**\n",
+    "\n",
+    "1. `braintrace.compile(model, braintrace.D_RTRL, inputs[0], batch_size=B, vmap=True)` initialises per-sample hidden states, compiles the computation graph using a single time step's batched input (`inputs[0]`, shape `(batch_size, n_in)`), and wraps the result in a `braintrace.ETraceVmap` (a `brainstate.nn.Vmap` that also carries the sequence drivers) for batch-parallel execution — all in one call.\n",
+    "2. `vmapped_algo.etrace_grad` iterates over the time dimension (`inputs` has shape `(n_steps, batch_size, n_in)`). At each step it calls `step_fn` for the loss and takes its online gradient, then accumulates; `reduction='sum'` accumulates without dividing. The vmapped learner carries the same drivers as an unbatched one, so this line is identical in both batching strategies.\n",
+    "3. The returned `grads` dictionary can be passed to an optimizer (e.g., `braintools.optim.Adam`) for a parameter update."
+   ]
   },
   {
    "cell_type": "markdown",