From 9fd2bedd5c239842cf42f3b42ec8cf6c7a49cf6e Mon Sep 17 00:00:00 2001 From: poilsosart <128177087+poilsosart@users.noreply.github.com> Date: Tue, 28 Jul 2026 17:08:24 +0800 Subject: [PATCH 1/4] refactor: expose warnings and use Map initialization --- braintrace/_algorithm/base.py | 12 +- .../tests/conv_vmap_correctness_test.py | 64 ++++------- .../tests/diagnostic_exploration_test.py | 21 ++-- .../_algorithm/tests/while_support_test.py | 14 +-- braintrace/_compile.py | 65 ++++------- braintrace/_compile_test.py | 27 ++--- braintrace/_compiler/canonicalize_test.py | 22 ++-- braintrace/_compiler/hidden_group_test.py | 17 +-- braintrace/_compiler/module_info.py | 33 +++++- braintrace/_compiler/module_info_test.py | 14 +++ braintrace/_compiler/scenario_catalog_test.py | 36 ++---- .../tests/cell_relation_guardrail_test.py | 8 +- .../_compiler/tests/compiler_oracle_test.py | 20 ++-- .../_compiler/tests/compiler_property_test.py | 23 ++-- braintrace/_legacy/_ops_test.py | 19 +--- braintrace/_legacy/_params_test.py | 11 +- braintrace/_op/_primitive.py | 3 + braintrace/_op/conv.py | 48 +++++++- docs/advanced/batching.ipynb | 107 ++++++++++++++---- docs/quickstart/concepts.ipynb | 5 +- docs/quickstart/quickstart.ipynb | 22 ++-- ...6-07-28-warnings-and-map-initialization.md | 62 ++++++++++ docs/tutorials/drtrl.ipynb | 21 ++-- docs/tutorials/hidden_states.ipynb | 86 +++++++++----- docs/tutorials/neural_network_layers.ipynb | 17 +-- docs/tutorials/pp_prop.ipynb | 23 ++-- docs/tutorials/rnn_online_learning.ipynb | 29 ++--- docs/tutorials/snn_online_learning.ipynb | 17 ++- examples/002-coba-ei-rsnn.py | 8 +- ...03-snn-memory-and-speed-evaluation-vmap.py | 38 ++++--- examples/004-feedforward-conv-snn.py | 16 +-- examples/100-gru-on-copying-task.py | 13 +-- examples/drtrl/02-batching-vmap.py | 8 +- examples/pp_prop/05-batching-vmap.py | 16 +-- .../pp_prop/12-classification-neuromorphic.py | 13 +-- .../pp_prop/14-knob-vjp-method-contrast.py | 13 +-- examples/pp_prop/README.md | 2 +- examples/pp_prop/_shared.py | 14 +-- examples/snn_models.py | 29 ++--- 39 files changed, 536 insertions(+), 480 deletions(-) create mode 100644 docs/specs/2026-07-28-warnings-and-map-initialization.md diff --git a/braintrace/_algorithm/base.py b/braintrace/_algorithm/base.py index b403164e..ebd7dafb 100644 --- a/braintrace/_algorithm/base.py +++ b/braintrace/_algorithm/base.py @@ -279,14 +279,10 @@ def compile_graph(self, *args: Any) -> None: The input arguments. """ - # ``vmap_new_states`` / ``vmap2_new_states`` run an eager *discovery - # probe* that executes the surrounding ``init`` (including this call) - # once against throwaway, un-batched states before the real mapped - # pass. Compiling there would bind the executor to those probe states - # (which are discarded and left untagged), so the subsequent - # ``brainstate.nn.Vmap(..., vmap_states='new')`` would not cover them - # and writing a batched value raises ``BatchAxisError``. Defer to the - # real mapped pass, which creates the 'new'-tagged batched states. + # Legacy mapped-state transforms run an eager discovery probe before + # the real mapped pass. Compiling against the throwaway probe states + # would bind the executor to states that are discarded immediately, so + # defer compilation until the real mapped pass. _in_probe = getattr(brainstate.transform, 'in_new_state_probe', None) if _in_probe is not None and _in_probe(): return diff --git a/braintrace/_algorithm/tests/conv_vmap_correctness_test.py b/braintrace/_algorithm/tests/conv_vmap_correctness_test.py index c2f318f0..494e05e7 100644 --- a/braintrace/_algorithm/tests/conv_vmap_correctness_test.py +++ b/braintrace/_algorithm/tests/conv_vmap_correctness_test.py @@ -13,17 +13,17 @@ # limitations under the License. # ============================================================================== -"""Conv / mixed ETP under ``brainstate.nn.Vmap(vmap_states='new')`` correctness. +"""Conv and mixed ETP correctness under ``brainstate.nn.Map``. Regression coverage for the eligibility-trace path through the *batched* online -executor wrapped by ``brainstate.nn.Vmap`` (the ``OnlineVmapTrainer`` flow used +executor backed by ``brainstate.nn.Map`` (the mapped batching flow used by ``examples/004``). This path was previously uncovered — conv was exercised only at the rule level (``_op/conv_test.py``) and the "conv" model in ``transform_correctness_test`` is actually a matmul — which let two regressions through: 1. *Pure conv.* A conv forward forces a leading batch axis on its input, but - under ``vmap_states='new'`` the hidden-state traces are per-lane and carry no + under mapped state initialization the hidden-state traces are per-lane and carry no batch axis, so the instantaneous, recurrent and solve terms saw a singleton batch on the input but none on the cotangent. @@ -36,13 +36,13 @@ 3. *Norm in the transition.* ``conv -> LayerNorm -> IF`` makes ``dh/dy`` non-diagonal; the all-ones jvp returns its row sums, exactly zero for the shift-invariant norm. A ``use_fast_variance=True`` norm leaves a float32 - residual instead, which under ``vmap_states='new'`` the recurrent trace and + residual instead, which under mapped execution the recurrent trace and ``rsqrt(var+eps)`` amplify into an overflow that diverges from the eager reference — the ``examples/004`` ``loss=ln(10)`` stall. -**Oracle (exact, transform-invariance).** ``brainstate.nn.Vmap`` is a transform; +**Oracle (exact, transform-invariance).** ``brainstate.nn.Map`` is a transform; for parameters shared across lanes its grad sums the per-lane gradients. So the -gradient from the ``vmap_new_states`` + ``Vmap`` path on a batch of ``B`` samples +gradient from the mapped path on a batch of ``B`` samples must equal the sum over ``b`` of the *eager, batch=1* gradient on sample ``b``. The eager batch=1 path is independently healthy for conv (states are initialised *with* a size-1 batch, so input and trace batch axes agree), which makes it a @@ -59,17 +59,6 @@ import braintools import brainpy.state -# `etp_conv` has no registered batched counterpart, so every model here that -# routes a sample through `braintrace.nn.Conv2d` under `brainstate.nn.Vmap` -# hits the identity-preserving batching rule's decomposition fallback in -# `braintrace/_op/_primitive.py`, which emits a `UserWarning`. That warning is -# expected-but-not-under-test in this module (the module tests gradient -# correctness, not the vmap-decomposition warning itself — that is covered by -# `braintrace/_op/_primitive_test.py`), so it is filtered narrowly by message. -pytestmark = pytest.mark.filterwarnings( - "ignore:ETP primitive 'etp_conv' was decomposed:UserWarning" -) - H = W = 6 C_IN = 2 C_OUT = 3 @@ -154,24 +143,19 @@ def loss_fn(x): return grads -def _vmap_grad(data, targets, make_net): - """The ``OnlineVmapTrainer`` flow: vmap_new_states init + Vmap(vmap_states='new').""" +def _map_grad(data, targets, make_net): + """Initialize mapped states explicitly and accumulate mapped gradients.""" net = make_net() - model = braintrace.D_RTRL(net) - - @brainstate.transform.vmap_new_states(state_tag='new', axis_size=data.shape[1]) - def init(): - brainstate.nn.init_all_states(net) - with brainstate.environ.context(fit=True): - model.compile_graph(data[0, 0]) - - init() - vmodel = brainstate.nn.Vmap(model, vmap_states='new') + mapped_net = brainstate.nn.Map(net, init_map_size=data.shape[1]) + mapped_net.init_all_states() + model = braintrace.D_RTRL(mapped_net) + with brainstate.environ.context(fit=True): + model.compile_graph(data[0]) weights = net.states().subset(brainstate.ParamState) def _grad(inp): with brainstate.environ.context(fit=True): - return _loss(vmodel(inp), targets) + return _loss(model(inp), targets) def _step(prev, x): g = brainstate.transform.grad(_grad, weights)(x) @@ -189,7 +173,7 @@ def _step(prev, x): def _make_mixed_net(): """conv -> IF -> flatten -> Linear -> IF: a *mixed* batched/unbatched model. - Under ``vmap_states='new'`` the graph is compiled per-lane, so the conv stays + Under ``brainstate.nn.Map`` the graph is compiled across mapped lanes, so the conv stays a *batched* primitive (its parent layer forces a leading batch axis) while the flattened ``Linear`` input is 1-D and dispatches to the *unbatched* ``etp_mv``. The solve's trailing batch-sum must collapse only the conv gradient's batch @@ -232,7 +216,7 @@ def _make_conv_ln_net(use_fast_variance): upstream conv gets no eligibility gradient through the norm — a documented approximation, matching the eager path). That exactness is numerical: with ``use_fast_variance=True`` the ``E[x^2]-E[x]^2`` variance leaves a float32 - residual instead of zero, and under ``Vmap(vmap_states='new')`` the recurrent + residual instead of zero, and under mapped execution the recurrent trace and the large ``rsqrt(var+eps)`` factor amplify it into an overflow that diverges from the eager reference (the ``examples/004`` ``loss=ln(10)`` stall). """ @@ -272,8 +256,8 @@ def _assert_grads_match(ref, got): @pytest.mark.parametrize('neuron', ['IF', 'ALIF'], ids=['num_state1_IF', 'num_state2_ALIF']) -def test_conv_vmap_grad_equals_sum_of_eager_single_sample(neuron): - """vmap_new_states+Vmap conv D_RTRL grad == sum over samples of eager batch=1 grad.""" +def test_conv_map_grad_equals_sum_of_eager_single_sample(neuron): + """Mapped conv D-RTRL gradient equals the sum of eager sample gradients.""" rng = np.random.RandomState(42) data = jnp.asarray(rng.rand(N_STEP, B, H, W, C_IN).astype('float32')) targets = jnp.asarray(rng.rand(B, H, W, C_OUT).astype('float32')) @@ -284,11 +268,11 @@ def test_conv_vmap_grad_equals_sum_of_eager_single_sample(neuron): g = _eager_grad_one(data[:, b], targets[b], make_net) ref = g if ref is None else jax.tree.map(lambda a, c: a + c, ref, g) - got = _vmap_grad(data, targets, make_net) + got = _map_grad(data, targets, make_net) _assert_grads_match(ref, got) -def test_mixed_conv_dense_vmap_grad_equals_sum_of_eager_single_sample(): +def test_mixed_conv_dense_map_grad_equals_sum_of_eager_single_sample(): """Mixed batched(conv)+unbatched(dense-mv) model: vmap grad == sum of eager batch=1. Regression for the ``examples/004`` layer4 failure — the unbatched ``etp_mv`` @@ -304,11 +288,11 @@ def test_mixed_conv_dense_vmap_grad_equals_sum_of_eager_single_sample(): g = _eager_grad_one(data[:, b], targets[b], _make_mixed_net) ref = g if ref is None else jax.tree.map(lambda a, c: a + c, ref, g) - got = _vmap_grad(data, targets, _make_mixed_net) + got = _map_grad(data, targets, _make_mixed_net) _assert_grads_match(ref, got) -def test_conv_layernorm_vmap_grad_matches_eager_and_stays_finite(): +def test_conv_layernorm_map_grad_matches_eager_and_stays_finite(): """conv -> LayerNorm -> IF: vmap grad == sum of eager batch=1, and stays finite. Regression for the ``examples/004`` ``loss=ln(10)`` stall. A mean-subtracting @@ -316,7 +300,7 @@ def test_conv_layernorm_vmap_grad_matches_eager_and_stays_finite(): its row sums, which for shift-invariance are exactly zero, so the conv weight gets no eligibility gradient through the norm (both paths agree on ~0). With a numerically stable variance (``use_fast_variance=False``) that exact zero holds - under ``Vmap(vmap_states='new')``; the transform-invariance oracle then makes + under mapped execution; the transform-invariance oracle then makes vmap == sum-of-eager, and neither explodes. """ rng = np.random.RandomState(42) @@ -329,7 +313,7 @@ def test_conv_layernorm_vmap_grad_matches_eager_and_stays_finite(): g = _eager_grad_one(data[:, b], targets[b], make_net) ref = g if ref is None else jax.tree.map(lambda a, c: a + c, ref, g) - got = _vmap_grad(data, targets, make_net) + got = _map_grad(data, targets, make_net) # No overflow/NaN in either path (the bug produced ~1e14 -> NaN under vmap). for leaf in jax.tree.leaves(got): assert np.all(np.isfinite(np.asarray(leaf))), 'vmap grad is non-finite' diff --git a/braintrace/_algorithm/tests/diagnostic_exploration_test.py b/braintrace/_algorithm/tests/diagnostic_exploration_test.py index 085e2c4c..204f68bc 100644 --- a/braintrace/_algorithm/tests/diagnostic_exploration_test.py +++ b/braintrace/_algorithm/tests/diagnostic_exploration_test.py @@ -36,7 +36,6 @@ """ import importlib.util -import warnings import pytest @@ -82,10 +81,8 @@ def _drtrl(model): def _assert_exact_equals_bptt(spec, inputs): """D_RTRL multi-step gradient == BPTT gradient for every ParamState.""" - with warnings.catch_warnings(): - warnings.simplefilter('ignore') - expected = bptt_param_gradients(spec.factory, inputs) - actual = online_param_gradients(spec.factory, inputs, algo_factory=_drtrl) + expected = bptt_param_gradients(spec.factory, inputs) + actual = online_param_gradients(spec.factory, inputs, algo_factory=_drtrl) assert_param_gradients_close(actual, expected, atol=ATOL) @@ -163,12 +160,10 @@ def test_batch_invariance_over_dims(n_in, n_rec, batch, seq_len, seed): summed per-step SSE loss over the batch axis.""" seq = jnp.asarray( np.random.RandomState(seed).randn(seq_len, batch, n_in).astype('float32')) - with warnings.catch_warnings(): - warnings.simplefilter('ignore') - batched = _batched_multistep_grad(n_in, n_rec, batch, seq, seed) - summed = None - for b in range(batch): - sub = seq[:, b:b + 1, :] - g = _batched_multistep_grad(n_in, n_rec, 1, sub, seed) - summed = g if summed is None else {k: summed[k] + g[k] for k in g} + batched = _batched_multistep_grad(n_in, n_rec, batch, seq, seed) + summed = None + for b in range(batch): + sub = seq[:, b:b + 1, :] + g = _batched_multistep_grad(n_in, n_rec, 1, sub, seed) + summed = g if summed is None else {k: summed[k] + g[k] for k in g} assert_param_gradients_close(batched, summed, atol=ATOL) diff --git a/braintrace/_algorithm/tests/while_support_test.py b/braintrace/_algorithm/tests/while_support_test.py index 1ae022c3..f7fb998a 100644 --- a/braintrace/_algorithm/tests/while_support_test.py +++ b/braintrace/_algorithm/tests/while_support_test.py @@ -49,8 +49,6 @@ WARNING-level ``CONTROL_FLOW_OPAQUE_FWD`` diagnostic for each detach. """ -import warnings - import brainstate import jax import jax.numpy as jnp @@ -239,13 +237,11 @@ def test_upstream_layer_gradient_is_zero_behind_while_DOCUMENTED_LIMITATION(): inputs = _inputs(6, 3) def grads(while_layer): - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - return online_param_gradients_singlestep_naive( - lambda: _StackedWhileNet(while_layer=while_layer), - inputs, - algo_factory=braintrace.D_RTRL, - ) + return online_param_gradients_singlestep_naive( + lambda: _StackedWhileNet(while_layer=while_layer), + inputs, + algo_factory=braintrace.D_RTRL, + ) g_while = grads(True) g_twin = grads(False) diff --git a/braintrace/_compile.py b/braintrace/_compile.py index f214c89a..94f3556f 100644 --- a/braintrace/_compile.py +++ b/braintrace/_compile.py @@ -17,14 +17,12 @@ from typing import Any, Type, Union -import jax import brainstate from ._misc import CompilationError from ._algorithm import ( ETraceAlgorithm, ETraceConfig, - ETraceVmap, IODimVjpAlgorithm, ParamDimVjpAlgorithm, RandomProjectionVjpAlgorithm, @@ -128,7 +126,7 @@ def compile( verbose: int = 0, vmap: bool = False, **options: Any, -) -> ETraceAlgorithm | brainstate.nn.Vmap: +) -> ETraceAlgorithm: """Define an eligibility-trace online-learning model in one call. This is the unified entry point. It initializes the model's states, builds @@ -174,29 +172,22 @@ def compile( vmap : bool, optional When ``False`` (default) states are initialized with ``init_all_states(model, batch_size=batch_size)``. When ``True``, states - are created under - ``brainstate.transform.vmap_new_states(state_tag='new', axis_size=batch_size)`` - and the learner is wrapped in :class:`ETraceVmap`. In vmap mode: - ``example_inputs`` carry the batch axis (axis 0); ``batch_size`` is - **required** and used as the vmap ``axis_size``; the return value is a - :class:`ETraceVmap` whose ``.module`` is the unbatched learner (use - ``result.module.report`` for its report). Drive sequences through the - returned wrapper, never through ``result.module``. Requires a model - whose hidden states are all (re)created in ``init_all_states``; models - holding construction-time states may raise - ``brainstate.transform.BatchAxisError``. + initialized by wrapping the model in + ``brainstate.nn.Map(model, init_map_size=batch_size)`` and calling the + mapped model's ``init_all_states()`` method. In vmap mode, + ``example_inputs`` carry the batch axis (axis 0), ``batch_size`` is + **required** and sets the map size, and the returned learner exposes + ``report``, ``etrace_grad``, and ``etrace_evolve`` directly. **options : Any Forwarded to the algorithm constructor. See *Algorithm options* below. Returns ------- - ETraceAlgorithm or ETraceVmap - When ``vmap=False``, the compiled learner carries a - :attr:`~ETraceAlgorithm.report`; call ``.update(*inputs)`` to train. - When ``vmap=True``, returns an :class:`ETraceVmap` wrapper (also a - ``brainstate.nn.Vmap``); access the underlying learner's report as - ``.module.report``. Call ``etrace_grad`` and ``etrace_evolve`` on the - wrapper itself, not on ``.module``. + ETraceAlgorithm + The compiled learner, carrying a :attr:`~ETraceAlgorithm.report`. Call + ``.update(*inputs)`` for one step, or use ``etrace_grad`` and + ``etrace_evolve`` for sequences. This return contract is identical in + mapped and directly batched modes. Raises ------ @@ -315,34 +306,22 @@ def compile( raise ValueError(f'verbose must be 0, 1, or 2, got {verbose!r}.') if vmap and batch_size is None: raise ValueError( - 'compile(..., vmap=True) requires batch_size, used as the per-sample ' - 'vmap axis size. Pass batch_size= matching the batch axis ' - '(axis 0) of example_inputs.' + 'compile(..., vmap=True) requires batch_size, used as the ' + 'brainstate.nn.Map size. Pass batch_size= matching axis 0 ' + 'of example_inputs.' ) if vmap: - # Per-sample vmap scheme: example_inputs carry the batch axis (axis 0); - # the eligibility-trace graph is built per-lane on an unbatched sample, - # while hidden + trace states are created with the new per-sample axis. - learner = cls(model, **options) - unbatched = jax.tree.map(lambda a: a[0], example_inputs) - - @brainstate.transform.vmap_new_states(state_tag='new', axis_size=batch_size) - def _init() -> None: - brainstate.nn.init_all_states(model) - learner.compile_graph(*unbatched) - + # Per-sample map scheme: example_inputs carry the batch axis (axis 0). + model = brainstate.nn.Map(model, init_map_size=batch_size) if seed is not None: with brainstate.random.seed_context(seed): - _init() + model.init_all_states() else: - _init() - # ETraceVmap, not brainstate.nn.Vmap: the wrapper must carry - # etrace_grad / etrace_evolve so the call site is identical in batched - # and unbatched mode. Reaching into `.module` instead would drive the - # *unbatched* learner and silently give per-lane-wrong results. It is - # still a brainstate.nn.Vmap, so existing users are unaffected. - result: ETraceAlgorithm | brainstate.nn.Vmap = ETraceVmap(learner, vmap_states='new') + model.init_all_states() + learner = cls(model, **options) + learner.compile_graph(*example_inputs) + result = learner else: # --- state initialization (always) --- # if seed is not None: diff --git a/braintrace/_compile_test.py b/braintrace/_compile_test.py index 14ae3bbc..18b0afdb 100644 --- a/braintrace/_compile_test.py +++ b/braintrace/_compile_test.py @@ -259,7 +259,10 @@ def test_compile_vmap_builds_forwards_and_grads(): B = 4 xb = jnp.ones((B, 3), dtype='float32') learner = braintrace.compile(model, 'D_RTRL', xb, batch_size=B, vmap=True) - assert isinstance(learner, brainstate.nn.Vmap) + assert isinstance(learner, braintrace.D_RTRL) + assert isinstance(learner.graph_executor.model, brainstate.nn.Map) + assert learner.graph_executor.model.init_map_size == B + assert learner.graph_executor.model._init out = learner(xb) assert out.shape[0] == B @@ -279,14 +282,14 @@ def test_compile_vmap_requires_batch_size(): assert 'batch_size' in str(exc.value) -def test_compile_vmap_returns_wrapper_exposing_report(): +def test_compile_vmap_returns_algorithm_exposing_report(): model = _VmapRNN() B = 4 xb = jnp.ones((B, 3), dtype='float32') learner = braintrace.compile(model, 'D_RTRL', xb, batch_size=B, vmap=True) - assert isinstance(learner.module, braintrace.D_RTRL) - assert learner.module.report is not None - assert learner.module.is_compiled + assert isinstance(learner, braintrace.D_RTRL) + assert learner.report is not None + assert learner.is_compiled # --- both-modes coverage across RNN architectures + algorithms --------------- @@ -402,24 +405,14 @@ def update(self, x): @pytest.mark.parametrize('name,builder,algo,kw,feat', _BOTH_MODE_CASES, ids=[c[0] for c in _BOTH_MODE_CASES]) @pytest.mark.parametrize('vmap', [False, True], ids=['no_vmap', 'vmap']) -# `conv1d_minigru_d_rtrl` (etp_conv) has no registered batched counterpart, -# so under `vmap=True` compilation it hits the identity-preserving batching -# rule's decomposition fallback and warns (see `braintrace/_op/_primitive.py`). -# This test asserts gradient finiteness/non-zero-ness, not the -# vmap-decomposition warning (covered by -# `braintrace/_op/_primitive_test.py`), so the expected warning is filtered -# narrowly by message rather than left uncaptured. `lora_d_rtrl` (etp_lora_mv) -# now has a registered batched counterpart (`etp_lora_mm`) and is promoted -# instead of decomposed, so no filter is needed for it. -@pytest.mark.filterwarnings( - "ignore:ETP primitive 'etp_conv' was decomposed:UserWarning") def test_compile_both_modes_finite_nonzero_grad(name, builder, algo, kw, feat, vmap): B, T = 4, 5 xs = brainstate.random.randn(T, B, *feat) model = builder() learner = braintrace.compile(model, algo, xs[0], batch_size=B, vmap=vmap, **kw) if vmap: - assert isinstance(learner, brainstate.nn.Vmap) + assert isinstance(learner, braintrace.ETraceAlgorithm) + assert isinstance(learner.graph_executor.model, brainstate.nn.Map) weights = model.states(brainstate.ParamState) def total_loss(xs): diff --git a/braintrace/_compiler/canonicalize_test.py b/braintrace/_compiler/canonicalize_test.py index a08b3697..70eda7c9 100644 --- a/braintrace/_compiler/canonicalize_test.py +++ b/braintrace/_compiler/canonicalize_test.py @@ -13,7 +13,6 @@ # limitations under the License. # ============================================================================== -import warnings import brainstate import jax @@ -808,21 +807,20 @@ def test_skip_length_exceeds_limit(self): kinds = [r.kind for r in reporter.records()] assert kinds.count(DiagnosticKind.SCAN_UNROLL_SKIPPED) == 1 - def test_skip_length_exceeds_limit_info_under_descent_auto(self): + def test_skip_length_exceeds_limit_info_under_descent_auto(self, recwarn): # Phase 4: with scan_descent='auto' an over-limit scan is no longer a # dead end, so the skip record downgrades to INFO (no UserWarning) # and points at the descent path. f, closed, w, h0, xs = self._etp_scan_jaxpr() with diagnostic_context() as reporter: - with warnings.catch_warnings(): - warnings.simplefilter('error') - conv = _unroll( - closed, - weights=[closed.jaxpr.invars[0]], - policy=ControlFlowPolicy(scan_unroll_limit=self.L - 1, - scan_descent='auto'), - ) + conv = _unroll( + closed, + weights=[closed.jaxpr.invars[0]], + policy=ControlFlowPolicy(scan_unroll_limit=self.L - 1, + scan_descent='auto'), + ) assert 'scan' in _primitive_names(conv) + assert not any(issubclass(w.category, UserWarning) for w in recwarn) recs = [r for r in reporter.records() if r.kind is DiagnosticKind.SCAN_UNROLL_SKIPPED] assert len(recs) == 1 @@ -1229,9 +1227,7 @@ def body_fn(carry): return h with pytest.raises(NotImplementedError, match='while'): - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - self._graph_for(WhileCell) + self._graph_for(WhileCell) def test_drtrl_gradient_parity_with_unrolled_model(self): def build_and_grads(cell_cls): diff --git a/braintrace/_compiler/hidden_group_test.py b/braintrace/_compiler/hidden_group_test.py index 352cb878..dadcd817 100644 --- a/braintrace/_compiler/hidden_group_test.py +++ b/braintrace/_compiler/hidden_group_test.py @@ -15,7 +15,6 @@ import unittest -import warnings from pprint import pprint import brainstate @@ -1497,12 +1496,10 @@ def _compile_mixing(self, include_recurrent_mixing=False, n=4): cell = WhileMixingCell(n) brainstate.nn.init_all_states(cell) x = brainstate.random.rand(n) - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - with diagnostic_context() as reporter: - groups, path_to_group = find_hidden_groups_from_module( - cell, x, include_recurrent_mixing=include_recurrent_mixing, - ) + with diagnostic_context() as reporter: + groups, path_to_group = find_hidden_groups_from_module( + cell, x, include_recurrent_mixing=include_recurrent_mixing, + ) return cell, x, groups, reporter def test_default_mode_falls_back_to_zero_recurrence(self): @@ -1582,10 +1579,8 @@ def test_jit_wrapped_mixing_in_body_is_still_a_boundary(self): cell = WhileJitMixingCell(4) brainstate.nn.init_all_states(cell) x = brainstate.random.rand(4) - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - with diagnostic_context() as reporter: - groups, _pg = find_hidden_groups_from_module(cell, x) + with diagnostic_context() as reporter: + groups, _pg = find_hidden_groups_from_module(cell, x) assert len(groups) == 1 group = groups[0] # zero-recurrence fallback, exactly like the un-jitted mixing cell diff --git a/braintrace/_compiler/module_info.py b/braintrace/_compiler/module_info.py index cf9f42d9..334ba762 100644 --- a/braintrace/_compiler/module_info.py +++ b/braintrace/_compiler/module_info.py @@ -74,12 +74,32 @@ def _check_consistent_states_between_model_and_compiler( id(st): st for st in compiled_model_states } - id_to_path = { - id(st): path + id_to_path = {} + for path, st in retrieved_model_states.items(): + state_id = id(st) + previous = id_to_path.get(state_id) + if previous is None: + id_to_path[state_id] = path + continue + + # Map exposes mapped states through both its internal registry and the + # wrapped module. Prefer the module-facing path: numeric registry keys + # look like layer boundaries to hidden-state grouping. + internal_depth = sum(part == 'dict_vmap_states' for part in path) + previous_internal_depth = sum( + part == 'dict_vmap_states' for part in previous + ) + if internal_depth < previous_internal_depth: + id_to_path[state_id] = path + + # Graph traversal may expose the same state through more than one path. + # Keep the canonical path selected above so each compiled state has exactly + # one model path. + paths_to_remove = [ + path for path, st in retrieved_model_states.items() - } - - paths_to_remove = [] + if id_to_path[id(st)] != path + ] for id_ in id_to_path: if id_ not in id_to_compiled_state: path = id_to_path[id_] @@ -175,6 +195,9 @@ def abstractify_model( "Since it allows the explicit definition of the model structure." ) model_retrieved_states = brainstate.graph.states(model) + if isinstance(model, brainstate.nn.Map): + for path, state in brainstate.graph.states(model.module).items(): + model_retrieved_states[('module', *path)] = state # --- stateful model, for extracting states, weights, and variables --- # # diff --git a/braintrace/_compiler/module_info_test.py b/braintrace/_compiler/module_info_test.py index d37bef8f..c0d06739 100644 --- a/braintrace/_compiler/module_info_test.py +++ b/braintrace/_compiler/module_info_test.py @@ -64,6 +64,20 @@ def test_add_jaxpr_outs_preserves_policy(self): class Test_extract_model_info: + def test_map_hidden_state_aliases_are_deduplicated(self): + batch_size = 3 + rnn = braintrace.nn.GRUCell(2, 4) + mapped = brainstate.nn.Map(rnn, init_map_size=batch_size) + mapped.init_all_states() + + minfo = braintrace.extract_module_info( + mapped, brainstate.random.rand(batch_size, 2) + ) + states = minfo.retrieved_model_states + + assert len({id(state) for state in states.values()}) == len(states) + assert all('dict_vmap_states' not in path for path in states) + @pytest.mark.parametrize( "cls", [ diff --git a/braintrace/_compiler/scenario_catalog_test.py b/braintrace/_compiler/scenario_catalog_test.py index 0b10aa9c..48960263 100644 --- a/braintrace/_compiler/scenario_catalog_test.py +++ b/braintrace/_compiler/scenario_catalog_test.py @@ -34,8 +34,6 @@ - ``W -> non-gradient-enabled W -> h`` excludes the preceding weight. """ -import warnings - import brainstate import jax import jax.numpy as jnp @@ -56,14 +54,8 @@ def _compile(model, *inputs): - """Compile, suppressing expected weight-exclusion UserWarnings. - - Tests still assert on the structured ``DiagnosticKind`` records, so we - silence the warning-stream duplicate for readability. - """ - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - return compile_etrace_graph(model, *inputs, include_hidden_perturb=False) + """Compile a model and retain both warnings and structured diagnostics.""" + return compile_etrace_graph(model, *inputs, include_hidden_perturb=False) def _relation_set(graph): @@ -969,11 +961,7 @@ def test_cond_branches_full_pipeline_converts(self): def test_scan_body_full_pipeline_unrolls(self): model = ScanBodyRNN(4, loops=3) brainstate.nn.init_all_states(model) - with warnings.catch_warnings(): - # Earlier sub-steps' weights are excluded per the - # weight->weight->hidden invariant and warn about it. - warnings.simplefilter('ignore', UserWarning) - graph = compile_etrace_graph(model, jnp.ones(4)) + graph = compile_etrace_graph(model, jnp.ones(4)) names = [eqn.primitive.name for eqn in graph.module_info.jaxpr.eqns] assert 'scan' not in names # Only the final sub-step's two ETP ops are relations; the earlier @@ -1009,10 +997,8 @@ def test_scan_body_etp_exclude_policy_warns_and_drops(self): jaxpr = make_scan_body_etp_jaxpr(3, 4) policy = braintrace.ControlFlowPolicy(etp_in_control_flow='exclude') - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - with diagnostic_context() as reporter: - top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) + with diagnostic_context() as reporter: + top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) assert top == [], 'ETP inside scan body must NOT bubble up' records = [ @@ -1026,10 +1012,8 @@ def test_cond_branches_etp_exclude_policy_diagnostic_per_branch(self): jaxpr = make_cond_branches_etp_jaxpr(3, 4) policy = braintrace.ControlFlowPolicy(etp_in_control_flow='exclude') - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - with diagnostic_context() as reporter: - top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) + with diagnostic_context() as reporter: + top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) assert top == [] n_cf = sum( @@ -1069,10 +1053,8 @@ def test_while_body_etp_exclude_policy_warns_and_drops(self): jaxpr = make_while_body_etp_jaxpr(4, 4) policy = braintrace.ControlFlowPolicy(etp_in_control_flow='exclude') - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - with diagnostic_context() as reporter: - top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) + with diagnostic_context() as reporter: + top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) assert top == [] kinds = [r.kind for r in reporter.records()] diff --git a/braintrace/_compiler/tests/cell_relation_guardrail_test.py b/braintrace/_compiler/tests/cell_relation_guardrail_test.py index ae578498..06e26413 100644 --- a/braintrace/_compiler/tests/cell_relation_guardrail_test.py +++ b/braintrace/_compiler/tests/cell_relation_guardrail_test.py @@ -24,8 +24,6 @@ compiler_property_test.py, graph_test.py) and are not duplicated here. """ -import warnings - import brainstate import pytest @@ -51,11 +49,7 @@ def _compile_cell(name, n_in=3, n_out=4): cell = cls(n_in, n_out) brainstate.nn.init_all_states(cell) inp = brainstate.random.rand(n_in) - with warnings.catch_warnings(): - # GRUCell legitimately warns when it excludes Wr (W->W->h); the guardrail - # checks the diagnostic records, not the warning. - warnings.simplefilter('ignore') - return braintrace.compile_etrace_graph(cell, inp, include_hidden_perturb=False) + return braintrace.compile_etrace_graph(cell, inp, include_hidden_perturb=False) @pytest.mark.parametrize('cell_name', list(_CELL_GUARDRAILS)) diff --git a/braintrace/_compiler/tests/compiler_oracle_test.py b/braintrace/_compiler/tests/compiler_oracle_test.py index ec572379..d58553f2 100644 --- a/braintrace/_compiler/tests/compiler_oracle_test.py +++ b/braintrace/_compiler/tests/compiler_oracle_test.py @@ -39,8 +39,6 @@ -import warnings - import brainstate import jax import jax.numpy as jnp @@ -54,10 +52,8 @@ ) -def _silent_compile(model, *args): - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - return compile_etrace_graph(model, *args, include_hidden_perturb=False) +def _compile(model, *args): + return compile_etrace_graph(model, *args, include_hidden_perturb=False) def _transition_callable(rel, group, const_vals): @@ -90,7 +86,7 @@ def test_unbatched_mv_rnn_dhdy_matches_analytic(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.3, -0.7, 1.1]) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) rel = graph.hidden_param_op_relations[0] @@ -122,7 +118,7 @@ def test_elemwise_only_rnn_dhdy_matches_analytic(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.5, -0.2, 0.9, -1.0]) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) rel = graph.hidden_param_op_relations[0] @@ -165,7 +161,7 @@ def test_w1_dhdy_is_direct_only(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.4, -0.6, 1.2]) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) by_path = {r.path: r for r in graph.hidden_param_op_relations} @@ -190,7 +186,7 @@ def test_w2_dhdy_is_correct(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.4, -0.6, 1.2]) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) by_path = {r.path: r for r in graph.hidden_param_op_relations} @@ -237,7 +233,7 @@ def test_unbatched_mv_rnn_fd_matches_ad(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.1, 0.2, -0.3]) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) rel = graph.hidden_param_op_relations[0] self._check_fd(rel, rel.hidden_groups[0], temps) @@ -247,7 +243,7 @@ def test_partial_path_w1_fd_matches_ad(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.05, -0.1, 0.2]) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) rel = next( r for r in graph.hidden_param_op_relations diff --git a/braintrace/_compiler/tests/compiler_property_test.py b/braintrace/_compiler/tests/compiler_property_test.py index 0aeb81d2..35920121 100644 --- a/braintrace/_compiler/tests/compiler_property_test.py +++ b/braintrace/_compiler/tests/compiler_property_test.py @@ -42,7 +42,6 @@ import importlib.util -import warnings import pytest @@ -73,10 +72,8 @@ ) -def _silent_compile(model, *args): - with warnings.catch_warnings(): - warnings.simplefilter('ignore', UserWarning) - return compile_etrace_graph(model, *args, include_hidden_perturb=False) +def _compile(model, *args): + return compile_etrace_graph(model, *args, include_hidden_perturb=False) def _summary(graph): @@ -122,8 +119,8 @@ def test_unbatched_mv_rnn_is_idempotent(self, n_in, n_out): m2 = UnbatchedMvRNN(n_in, n_out) brainstate.nn.init_all_states(m2) - s1 = _summary(_silent_compile(m1, inp)) - s2 = _summary(_silent_compile(m2, inp)) + s1 = _summary(_compile(m1, inp)) + s2 = _summary(_compile(m2, inp)) assert s1 == s2 @given( @@ -138,8 +135,8 @@ def test_partial_path_rnn_is_idempotent(self, n): m2 = PartialPathRNN(n, n) brainstate.nn.init_all_states(m2) - s1 = _summary(_silent_compile(m1, inp)) - s2 = _summary(_silent_compile(m2, inp)) + s1 = _summary(_compile(m1, inp)) + s2 = _summary(_compile(m2, inp)) assert s1 == s2 @@ -163,7 +160,7 @@ def test_one_relation_per_layer_scoped_to_own_h(self, depth, n_in, n_out): brainstate.nn.init_all_states(model) inp = jnp.zeros(n_in) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) rels = graph.hidden_param_op_relations assert len(rels) == depth @@ -190,7 +187,7 @@ def test_two_call_sites_yield_two_relations(self, n): brainstate.nn.init_all_states(model) inp = jnp.zeros(n) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) rels = graph.hidden_param_op_relations assert len(rels) == 2 @@ -258,7 +255,7 @@ def test_only_last_weight_registers(self, chain_len, n): brainstate.nn.init_all_states(model) inp = jnp.zeros(n) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) rels = graph.hidden_param_op_relations included_paths = {r.path for r in rels} last = (f'w{chain_len - 1}',) @@ -294,7 +291,7 @@ def test_classification_is_shape_invariant(self, n_in, n_out): brainstate.nn.init_all_states(model) inp = jnp.zeros(n_in) - graph = _silent_compile(model, inp) + graph = _compile(model, inp) by_path = {r.path: r for r in graph.hidden_param_op_relations} assert by_path[('w1',)].path_classification == { diff --git a/braintrace/_legacy/_ops_test.py b/braintrace/_legacy/_ops_test.py index 2470ec3d..1b81e096 100644 --- a/braintrace/_legacy/_ops_test.py +++ b/braintrace/_legacy/_ops_test.py @@ -25,11 +25,10 @@ -import warnings - import jax import jax.numpy as jnp import numpy as np +import pytest import braintrace from braintrace._legacy import ( @@ -151,22 +150,14 @@ class TestDeprecationWarnings: # and not when importing from the private ``braintrace._legacy`` submodule. def test_matmul_op_access_warns(self): - with warnings.catch_warnings(record=True) as captured: - warnings.simplefilter('always') + with pytest.warns(DeprecationWarning, match='MatMulOp'): _ = braintrace.MatMulOp - assert any( - issubclass(w.category, DeprecationWarning) - and 'MatMulOp' in str(w.message) - for w in captured - ) - def test_construction_does_not_warn(self): + def test_construction_does_not_warn(self, recwarn): # The shim classes themselves no longer warn; construction is silent. - with warnings.catch_warnings(record=True) as captured: - warnings.simplefilter('always') - MatMulOp() + MatMulOp() assert not any( - issubclass(w.category, DeprecationWarning) for w in captured + issubclass(w.category, DeprecationWarning) for w in recwarn ) diff --git a/braintrace/_legacy/_params_test.py b/braintrace/_legacy/_params_test.py index 52408cdf..799e52e6 100644 --- a/braintrace/_legacy/_params_test.py +++ b/braintrace/_legacy/_params_test.py @@ -26,11 +26,10 @@ -import warnings - import brainstate import jax.numpy as jnp import numpy as np +import pytest import braintrace from braintrace._legacy import ( @@ -172,14 +171,8 @@ class TestDeprecationWarnings: # and not when importing from the private ``braintrace._legacy`` submodule. def test_etrace_param_access_warns(self): - with warnings.catch_warnings(record=True) as captured: - warnings.simplefilter('always') + with pytest.warns(DeprecationWarning, match='ETraceParam'): _ = braintrace.ETraceParam - assert any( - issubclass(w.category, DeprecationWarning) - and 'ETraceParam' in str(w.message) - for w in captured - ) # --------------------------------------------------------------------------- diff --git a/braintrace/_op/_primitive.py b/braintrace/_op/_primitive.py index faadc7ad..92e52b4c 100644 --- a/braintrace/_op/_primitive.py +++ b/braintrace/_op/_primitive.py @@ -58,6 +58,8 @@ 'register_primitive', ] +_ETP_BATCHING_RULES: dict[Primitive, Callable[..., Any]] = {} + class ETPPrimitive(Primitive): """A JAX ``Primitive`` with ETP rule registration helpers. @@ -336,6 +338,7 @@ def _batching(args: Any, dims: Any, **params: Any) -> Any: ) return jax.vmap(partial(impl_fn, **params), in_axes=dims)(*args), 0 + _ETP_BATCHING_RULES[p] = _batching batching.primitive_batchers[p] = _batching return p diff --git a/braintrace/_op/conv.py b/braintrace/_op/conv.py index f53dbe7b..a8777352 100644 --- a/braintrace/_op/conv.py +++ b/braintrace/_op/conv.py @@ -121,8 +121,9 @@ import jax import jax.numpy as jnp import brainunit as u +from jax.interpreters import batching -from ._primitive import register_primitive +from ._primitive import _ETP_BATCHING_RULES, register_primitive from ._registries import ETP_RULES_INSTANT_DRTRL, ETP_RULES_SOLVE_DRTRL from braintrace._typing import ArrayLike, WeightFn @@ -758,6 +759,51 @@ def _conv_init_pp(x_var: Any, y_var: Any, weight_vars: dict[str, Any], trainable_invars_fn=_conv_trainable_invars, x_invar_index=0, ) + +_default_conv_batcher = _ETP_BATCHING_RULES[etp_conv_p] + + +def _conv_lhs_batch_axis(params: dict[str, Any]) -> int: + """Return the input batch-axis position encoded by dimension numbers.""" + dn = params.get('dimension_numbers') + if dn is None: + return 0 + if isinstance(dn, tuple) and len(dn) == 3 and isinstance(dn[0], str): + return dn[0].index('N') + return dn.lhs_spec[0] + + +def _conv_batcher(args: Any, dims: Any, **params: Any) -> Any: + """Preserve ``etp_conv_p`` when mapping only the convolution input.""" + x_idx = 0 + if ( + dims[x_idx] is None + or any(d is not None for i, d in enumerate(dims) if i != x_idx) + ): + return _default_conv_batcher(args, dims, **params) + + x = jnp.moveaxis(args[x_idx], dims[x_idx], 0) + lhs_batch_axis = _conv_lhs_batch_axis(params) + x = jnp.moveaxis(x, lhs_batch_axis + 1, 1) + map_size, inner_batch_size = x.shape[:2] + merged_x = x.reshape(map_size * inner_batch_size, *x.shape[2:]) + merged_x = jnp.moveaxis(merged_x, 0, lhs_batch_axis) + + merged_args = tuple( + merged_x if i == x_idx else arg for i, arg in enumerate(args) + ) + y = etp_conv_p.bind(*merged_args, **params) + + _, _, output_batch_axis, _ = _conv_layout(params) + y = jnp.moveaxis(y, output_batch_axis, 0) + y = y.reshape(map_size, inner_batch_size, *y.shape[1:]) + y = jnp.moveaxis(y, 1, output_batch_axis + 1) + return y, 0 + + +batching.primitive_batchers[etp_conv_p] = _conv_batcher + + def _conv_snap_anchor(eqn_params: dict) -> bool: """Declare the SnAp-n trace anchor for ``etp_conv``. diff --git a/docs/advanced/batching.ipynb b/docs/advanced/batching.ipynb index 845b760b..7b97702f 100644 --- a/docs/advanced/batching.ipynb +++ b/docs/advanced/batching.ipynb @@ -9,19 +9,34 @@ "\n", "Online learning algorithms need to handle batched data efficiently. In braintrace, there are two main batching strategies:\n", "\n", - "- **Vmap-based batching** (recommended): Compile the computation graph for a single sample, then use `vmap` to automatically vectorize across the batch dimension.\n", + "- **Map-based batching** (recommended): Wrap single-sample model logic with `brainstate.nn.Map`, then compile from one complete batched time step.\n", "- **Single-sample mode**: Process one sample at a time, without any batching.\n", "\n", "The choice of strategy affects how model states are initialized and how the online learning algorithm is called.\n", "\n", - "This tutorial walks through each strategy with concrete examples and shows how to build a full training loop using vmap-based batching." + "This tutorial walks through each strategy with concrete examples and shows how to build a full training loop using Map-based batching." ] }, { "cell_type": "markdown", "id": "a1b2c3d4e5f60002", "metadata": {}, - "source": "## Vmap-Based Batching (Recommended)\n\nThe recommended approach is to compile the online learning graph for a **single sample**, then leverage JAX's `vmap` to parallelize across the batch. `braintrace.compile(..., batch_size=B, vmap=True)` does this in one call:\n\n1. It initialises per-sample states for `batch_size` samples.\n2. It compiles the ETP graph from a single **batched** time step (shape `(batch_size, n_in)`); the batch axis (axis 0) is stripped internally to recover the per-sample example.\n3. It wraps the algorithm with `brainstate.nn.Vmap` for parallel execution.\n4. It returns the vmapped learner, ready to call on batched inputs.\n\nThis pattern keeps the model definition simple (single-sample logic) while gaining efficient batch parallelism automatically." + "source": [ + "## Map-Based Batching (Recommended)\n", + "\n", + "The recommended approach is to keep the model's update logic single-sample\n", + "and let `brainstate.nn.Map` manage independent state copies across the batch.\n", + "`braintrace.compile(..., batch_size=B, vmap=True)` does this in one call:\n", + "\n", + "1. It wraps the model with `brainstate.nn.Map(model, init_map_size=B)`.\n", + "2. It initializes the mapped states through `mapped_model.init_all_states()`.\n", + "3. It compiles the ETP graph from one batched time step with shape\n", + " `(batch_size, n_in)`.\n", + "4. It returns the concrete online-learning algorithm, ready for batched calls.\n", + "\n", + "The returned learner exposes `report` and the rest of the algorithm API\n", + "directly." + ] }, { "cell_type": "code", @@ -60,13 +75,41 @@ "id": "a1b2c3d4e5f60005", "metadata": {}, "outputs": [], - "source": "model = SimpleGRU(10, 64, 5)\nbatch_size = 16\n\n# braintrace.compile with vmap=True:\n# - initialises per-sample hidden states (batch_size independent copies)\n# - compiles the ETP graph from a single batched time step: axis 0 is the batch\n# axis, which compile strips internally to recover the per-sample example\n# - wraps the result in brainstate.nn.Vmap for parallel execution\nalgo_vmapped = braintrace.compile(\n model, braintrace.D_RTRL, jnp.zeros((batch_size, 10)),\n batch_size=batch_size, vmap=True,\n)\n\n# Run on batched input — the returned learner handles the batch axis transparently\nx_batch = jnp.ones((batch_size, 10))\nout = algo_vmapped(x_batch)\nprint(\"Output shape:\", out.shape) # (16, 5)" + "source": [ + "model = SimpleGRU(10, 64, 5)\n", + "batch_size = 16\n", + "\n", + "# braintrace.compile with vmap=True:\n", + "# - initialises per-sample hidden states (batch_size independent copies)\n", + "# - wraps the model in brainstate.nn.Map and initializes mapped states\n", + "# - compiles the ETP graph from one batched time step\n", + "# - returns the concrete algorithm for parallel mapped execution\n", + "mapped_algo = braintrace.compile(\n", + " model, braintrace.D_RTRL, jnp.zeros((batch_size, 10)),\n", + " batch_size=batch_size, vmap=True,\n", + ")\n", + "\n", + "# Run on batched input — the returned learner handles the batch axis transparently\n", + "x_batch = jnp.ones((batch_size, 10))\n", + "out = mapped_algo(x_batch)\n", + "print(\"Output shape:\", out.shape) # (16, 5)" + ] }, { "cell_type": "markdown", "id": "a1b2c3d4e5f60006", "metadata": {}, - "source": "**How it works:**\n\n- `braintrace.compile(..., batch_size=B, vmap=True)` internally runs `brainstate.transform.vmap_new_states` to create `B` independent copies of all model states (tagged `'new'`), then compiles the ETP graph on a single-sample input, and finally wraps the learner in `brainstate.nn.Vmap(algo, vmap_states='new')` for parallel execution.\n- Each call to `algo_vmapped(x_batch)` automatically splits the batch input across the per-sample states, runs the forward pass independently for each sample, and stacks the outputs.\n- The model itself only ever sees single-sample inputs — all batch handling is transparent." + "source": [ + "**How it works:**\n", + "\n", + "- `braintrace.compile(..., batch_size=B, vmap=True)` creates\n", + " `brainstate.nn.Map(model, init_map_size=B)` and calls\n", + " `mapped_model.init_all_states()`.\n", + "- The algorithm compiles against the batched example input and keeps the ETP\n", + " primitives visible to the compiler.\n", + "- Each learner call maps the wrapped model over axis 0 while sharing parameter\n", + " states and maintaining independent recurrent states." + ] }, { "cell_type": "markdown", @@ -84,7 +127,17 @@ "id": "a1b2c3d4e5f60008", "metadata": {}, "outputs": [], - "source": "model2 = SimpleGRU(10, 64, 5)\n\n# Single-sample mode: omit batch_size so states are created unbatched.\nalgo2 = braintrace.compile(model2, braintrace.D_RTRL, jnp.zeros(10))\n\n# Process one sample at a time\nx_single = jnp.ones(10)\nout = algo2(x_single)\nprint(\"Single sample output shape:\", out.shape) # (5,)" + "source": [ + "model2 = SimpleGRU(10, 64, 5)\n", + "\n", + "# Single-sample mode: omit batch_size so states are created unbatched.\n", + "algo2 = braintrace.compile(model2, braintrace.D_RTRL, jnp.zeros(10))\n", + "\n", + "# Process one sample at a time\n", + "x_single = jnp.ones(10)\n", + "out = algo2(x_single)\n", + "print(\"Single sample output shape:\", out.shape) # (5,)" + ] }, { "cell_type": "markdown", @@ -137,12 +190,12 @@ "id": "a1b2c3d4e5f60013", "metadata": {}, "source": [ - "## Full Training Loop with Vmap Batching\n", + "## Full Training Loop with Map Batching\n", "\n", - "Below is a complete example that combines vmap-based batching with a temporal training loop. The pattern is:\n", + "Below is a complete example that combines Map-based batching with a temporal training loop. The pattern is:\n", "\n", - "1. **Initialize** model states and compile the graph for a single sample.\n", - "2. **Vmap** the algorithm across the batch dimension.\n", + "1. **Map and initialize** independent model states across the batch.\n", + "2. **Compile** the algorithm from one batched time step.\n", "3. **Scan** over time steps, accumulating gradients at each step.\n", "4. **Update** parameters with the accumulated gradients." ] @@ -157,22 +210,22 @@ "@brainstate.transform.jit\n", "def train_step(inputs, targets):\n", " \"\"\"inputs: (n_steps, batch_size, n_in), targets: (batch_size,)\"\"\"\n", - " # braintrace.compile with vmap=True replaces the manual init + compile_graph + Vmap pattern.\n", - " # Pass the batched single time step inputs[0] (shape (batch_size, n_in)); compile strips\n", - " # axis 0 internally to recover the per-sample example — do NOT pass inputs[0, 0].\n", - " vmapped_algo = braintrace.compile(\n", + " # braintrace.compile with vmap=True replaces manual Map initialization and compilation.\n", + " # Pass inputs[0] with shape (batch_size, n_in); the compiler traces the mapped model\n", + " # against the complete batched time step, so do not pass inputs[0, 0].\n", + " mapped_algo = braintrace.compile(\n", " model, braintrace.D_RTRL, inputs[0],\n", " batch_size=inputs.shape[1], vmap=True,\n", " )\n", "\n", " def step_loss(inp):\n", - " out = vmapped_algo(inp)\n", + " out = mapped_algo(inp)\n", " return jnp.mean((out - targets) ** 2)\n", "\n", " # etrace_grad drives the whole sequence and accumulates the per-step online\n", - " # gradients. The vmapped learner carries the same driver methods as an\n", - " # unbatched one, so nothing here changes when you switch batching strategy.\n", - " return vmapped_algo.etrace_grad(inputs, step_fn=step_loss, reduction='sum')" + " # gradients. Map keeps the batch axis inside the compiled graph, while the\n", + " # learner exposes the same driver methods as the unbatched path.\n", + " return mapped_algo.etrace_grad(inputs, step_fn=step_loss, reduction='sum')" ] }, { @@ -197,16 +250,26 @@ "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." + "1. `braintrace.compile(model, braintrace.D_RTRL, inputs[0], batch_size=B, vmap=True)` wraps the model with `brainstate.nn.Map`, initializes independent per-sample states, compiles from the batched time step, and returns the algorithm directly.\n", + "2. `vmapped_algo.etrace_grad` iterates over time, calls `step_fn`, and accumulates online gradients; `reduction='sum'` accumulates without dividing. The call is identical for mapped and directly batched learners.\n", + "3. The returned `grads` dictionary keeps the original model parameter paths and can be passed to an optimizer such as `braintools.optim.Adam`." ] }, { "cell_type": "markdown", "id": "a1b2c3d4e5f60017", "metadata": {}, - "source": "## Summary\n\n- **`braintrace.compile(..., batch_size=B, vmap=True)` is the recommended way** to set up batched online learning. It replaces the three-step manual pattern (`init_all_states` + `compile_graph` + `Vmap`) with a single call.\n- The workflow is: **`braintrace.compile` (with `vmap=True`) → scan over time steps → accumulate gradients → update parameters**.\n- `SingleStepData` and `MultiStepData` control whether the algorithm processes one time step or scans over an entire sequence internally.\n- For non-batched (single-sample) use, call `braintrace.compile` without `vmap=True`; states are initialised for a single sample." + "source": [ + "## Summary\n", + "\n", + "- `braintrace.compile(..., batch_size=B, vmap=True)` is the recommended setup\n", + " for batched online learning.\n", + "- Internally it uses `brainstate.nn.Map(model, init_map_size=B)` followed by\n", + " `mapped_model.init_all_states()`.\n", + "- The workflow is: compile with `vmap=True`, scan over time, accumulate\n", + " gradients, and update parameters.\n", + "- For one stream, compile without `vmap=True`; states remain unbatched." + ] } ], "metadata": { diff --git a/docs/quickstart/concepts.ipynb b/docs/quickstart/concepts.ipynb index 1bbf0586..434ef8c0 100644 --- a/docs/quickstart/concepts.ipynb +++ b/docs/quickstart/concepts.ipynb @@ -115,10 +115,7 @@ "import jax\n", "import jax.numpy as jnp\n", "import brainstate\n", - "import braintrace\n", - "import warnings\n", - "\n", - "warnings.filterwarnings(\"ignore\", message=r\"ETP primitive .*\")" + "import braintrace" ] }, { diff --git a/docs/quickstart/quickstart.ipynb b/docs/quickstart/quickstart.ipynb index 64d860df..9fe44f23 100644 --- a/docs/quickstart/quickstart.ipynb +++ b/docs/quickstart/quickstart.ipynb @@ -80,8 +80,6 @@ } ], "source": [ - "import warnings\n", - "\n", "import brainstate\n", "import braintools\n", "import braintrace\n", @@ -107,18 +105,14 @@ "targets = 0.7 * inputs + 0.2\n", "\n", "# Compile once. inputs[0] is one batched time step with shape (1, 1).\n", - "with warnings.catch_warnings():\n", - " # The readout does not feed recurrent state, so it is non-temporal.\n", - " warnings.filterwarnings(\n", - " \"ignore\",\n", - " message=r\"ETP primitive etp_mm.*has no connected hidden states.*\",\n", - " )\n", - " learner = braintrace.compile(\n", - " model,\n", - " braintrace.D_RTRL,\n", - " inputs[0],\n", - " batch_size=1,\n", - " )\n", + "# The compiler reports that the readout is non-temporal because it does not\n", + "# feed a recurrent state.\n", + "learner = braintrace.compile(\n", + " model,\n", + " braintrace.D_RTRL,\n", + " inputs[0],\n", + " batch_size=1,\n", + ")\n", "weights = model.states(brainstate.ParamState)\n", "optimizer = braintools.optim.SGD(lr=0.08)\n", "optimizer.register_trainable_weights(weights)" diff --git a/docs/specs/2026-07-28-warnings-and-map-initialization.md b/docs/specs/2026-07-28-warnings-and-map-initialization.md new file mode 100644 index 00000000..90d743e4 --- /dev/null +++ b/docs/specs/2026-07-28-warnings-and-map-initialization.md @@ -0,0 +1,62 @@ +# Warning Visibility and Map Initialization + +## Status + +Approved for implementation. + +## Motivation + +BrainTrace warnings are part of the public diagnostic surface. Library code, +tests, examples, and documentation must not suppress them with +`warnings.catch_warnings` or `warnings.filterwarnings`. + +Mapped state initialization must use the state-management abstraction provided +by BrainState. The supported flow is: + +```python +model = brainstate.nn.Map(model, init_map_size=batch_size) +model.init_all_states() +``` + +This replaces executable uses of +`brainstate.transform.vmap_new_states`. Historical changelog entries may retain +the old symbol when they describe behavior from an earlier release. + +## Requirements + +1. Remove every executable use of `warnings.catch_warnings`, + `warnings.filterwarnings`, and bare `filterwarnings`. +2. Do not add replacement warning filters or warning-suppression helpers. +3. Let BrainTrace warnings reach users and test output unchanged. +4. Replace mapped-state discovery and initialization with + `brainstate.nn.Map(model, init_map_size=...)` followed by + `model.init_all_states()`. +5. Compile mapped algorithms against the complete batched example input and + return the algorithm object directly. +6. Keep mapped model states discoverable by the compiler without duplicating + state paths. +7. Preserve ETP primitives and ETP-specific batching behavior under + `brainstate.nn.Map`. +8. Update affected tests, examples, and tutorials to use the supported Map + workflow. + +## Non-goals + +- Do not change algorithm equations, optimization behavior, or public callable + signatures unrelated to mapped initialization. +- Do not suppress third-party compatibility warnings. +- Do not address the separate `braintools`/`saiunit` quantity compatibility + failures. +- Do not add custom documentation CSS, JavaScript, Sphinx hooks, or static API + pages. + +## Verification + +1. Search the repository for prohibited warning filters and executable + `vmap_new_states` calls. +2. Run focused compiler, mapped-state, convolution batching, and public API + tests. +3. Build the documentation with Sphinx warnings treated as errors. +4. Run the complete test suite and allow it to finish naturally. +5. Report dependency-related failures separately from regressions caused by + this change. diff --git a/docs/tutorials/drtrl.ipynb b/docs/tutorials/drtrl.ipynb index cf1712b4..5239af03 100644 --- a/docs/tutorials/drtrl.ipynb +++ b/docs/tutorials/drtrl.ipynb @@ -59,7 +59,6 @@ "import braintrace\n", "import jax.numpy as jnp\n", "import matplotlib.pyplot as plt\n", - "import warnings\n", "\n", "brainstate.random.seed(7)" ] @@ -102,18 +101,14 @@ "inputs = jnp.linspace(-1.0, 1.0, 12).reshape(12, 1, 1)\n", "targets = 0.7 * inputs + 0.2\n", "\n", - "with warnings.catch_warnings():\n", - " # The Linear readout is intentionally non-temporal; Section 5 explains why.\n", - " warnings.filterwarnings(\n", - " \"ignore\",\n", - " message=r\"ETP primitive etp_mm.*has no connected hidden states.*\",\n", - " )\n", - " learner = braintrace.compile(\n", - " model,\n", - " braintrace.D_RTRL,\n", - " inputs[0],\n", - " batch_size=1,\n", - " )\n", + "# The Linear readout is intentionally non-temporal; Section 5 explains\n", + "# the compiler diagnostic emitted for it.\n", + "learner = braintrace.compile(\n", + " model,\n", + " braintrace.D_RTRL,\n", + " inputs[0],\n", + " batch_size=1,\n", + ")\n", "weights = model.states(brainstate.ParamState)\n", "optimizer = braintools.optim.SGD(lr=0.08)\n", "optimizer.register_trainable_weights(weights);" diff --git a/docs/tutorials/hidden_states.ipynb b/docs/tutorials/hidden_states.ipynb index c97f3110..aefdc61d 100644 --- a/docs/tutorials/hidden_states.ipynb +++ b/docs/tutorials/hidden_states.ipynb @@ -364,7 +364,7 @@ "- **Single-sample initialization**: `brainstate.nn.init_all_states(model)` -- state tensors have shape `(M,)`.\n", "- **Batched initialization**: `brainstate.nn.init_all_states(model, batch_size=N)` -- state tensors have shape `(N, M)`, where `N` is the batch size. This is used for manual batching.\n", "\n", - "For automatic batching with `vmap`, you can use `brainstate.transform.vmap_new_states` to initialize per-sample states while keeping the model definition simple." + "For automatic batching, wrap the model with `brainstate.nn.Map(model, init_map_size=N)` and call `mapped_model.init_all_states()`. The wrapper keeps the model definition single-sample while managing independent per-sample states." ] }, { @@ -403,7 +403,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": null, "id": "raeddpwnz5o", "metadata": { "execution": { @@ -413,29 +413,16 @@ "shell.execute_reply": "2026-06-26T13:41:44.376744Z" } }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "After vmap_new_states, model is ready for automatic batching.\n" - ] - } - ], + "outputs": [], "source": [ - "# --- Automatic batching with vmap_new_states ---\n", + "# --- Automatic batching with brainstate.nn.Map ---\n", "model = SimpleNeuron(32)\n", + "mapped_model = brainstate.nn.Map(model, init_map_size=16)\n", + "mapped_model.init_all_states()\n", "\n", - "@brainstate.transform.vmap_new_states(state_tag='new', axis_size=16)\n", - "def init():\n", - " brainstate.nn.init_all_states(model)\n", - "\n", - "init()\n", - "\n", - "# After vmap initialization, hidden states are managed per-sample internally.\n", - "# The model still \"thinks\" it processes a single sample, but vmap replicates\n", - "# the computation across the batch dimension automatically.\n", - "print(\"After vmap_new_states, model is ready for automatic batching.\")" + "# Map manages independent hidden states for each sample and accepts inputs\n", + "# whose leading axis has size 16.\n", + "print(\"Mapped model is ready for automatic batching.\")" ] }, { @@ -468,7 +455,20 @@ } }, "outputs": [], - "source": "# Complete example: model -> compile -> inspect graph structure\n\nmodel = SimpleNeuron(8)\n\n# braintrace.compile initialises states, compiles the ETP graph, and returns a ready learner.\n# Access the compiled graph via learner.graph and the report via learner.report.\nalgo = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros(8), batch_size=1)\n\n# Display the discovered graph structure:\n# - Which hidden groups were found\n# - Which weight parameters are associated with each group\nalgo.show_graph()" + "source": [ + "# Complete example: model -> compile -> inspect graph structure\n", + "\n", + "model = SimpleNeuron(8)\n", + "\n", + "# braintrace.compile initialises states, compiles the ETP graph, and returns a ready learner.\n", + "# Access the compiled graph via learner.graph and the report via learner.report.\n", + "algo = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros(8), batch_size=1)\n", + "\n", + "# Display the discovered graph structure:\n", + "# - Which hidden groups were found\n", + "# - Which weight parameters are associated with each group\n", + "algo.show_graph()" + ] }, { "cell_type": "markdown", @@ -496,7 +496,41 @@ } }, "outputs": [], - "source": "# A two-layer recurrent network to demonstrate multi-group discovery\n\nclass TwoLayerRNN(brainstate.nn.Module):\n \"\"\"Two stacked recurrent layers, each with its own hidden state.\"\"\"\n\n def __init__(self, in_size, hidden_size, out_size):\n super().__init__()\n # Layer 1\n self.w1_in = brainstate.ParamState(brainstate.random.randn(in_size, hidden_size) * 0.01)\n self.w1_rec = brainstate.ParamState(brainstate.random.randn(hidden_size, hidden_size) * 0.01)\n self.h1 = brainstate.HiddenState(jnp.zeros(hidden_size))\n\n # Layer 2\n self.w2_in = brainstate.ParamState(brainstate.random.randn(hidden_size, out_size) * 0.01)\n self.w2_rec = brainstate.ParamState(brainstate.random.randn(out_size, out_size) * 0.01)\n self.h2 = brainstate.HiddenState(jnp.zeros(out_size))\n\n def update(self, x):\n # Layer 1: x feeds in, h1 recurs\n self.h1.value = jax.nn.tanh(\n x @ self.w1_in.value + braintrace.matmul(self.h1.value, self.w1_rec.value)\n )\n # Layer 2: h1 feeds in, h2 recurs\n self.h2.value = jax.nn.tanh(\n self.h1.value @ self.w2_in.value + braintrace.matmul(self.h2.value, self.w2_rec.value)\n )\n return self.h2.value\n\n\nmodel_2layer = TwoLayerRNN(in_size=10, hidden_size=16, out_size=8)\n\nalgo_2layer = braintrace.compile(model_2layer, braintrace.D_RTRL, jnp.zeros(10), batch_size=1)\nalgo_2layer.show_graph()" + "source": [ + "# A two-layer recurrent network to demonstrate multi-group discovery\n", + "\n", + "class TwoLayerRNN(brainstate.nn.Module):\n", + " \"\"\"Two stacked recurrent layers, each with its own hidden state.\"\"\"\n", + "\n", + " def __init__(self, in_size, hidden_size, out_size):\n", + " super().__init__()\n", + " # Layer 1\n", + " self.w1_in = brainstate.ParamState(brainstate.random.randn(in_size, hidden_size) * 0.01)\n", + " self.w1_rec = brainstate.ParamState(brainstate.random.randn(hidden_size, hidden_size) * 0.01)\n", + " self.h1 = brainstate.HiddenState(jnp.zeros(hidden_size))\n", + "\n", + " # Layer 2\n", + " self.w2_in = brainstate.ParamState(brainstate.random.randn(hidden_size, out_size) * 0.01)\n", + " self.w2_rec = brainstate.ParamState(brainstate.random.randn(out_size, out_size) * 0.01)\n", + " self.h2 = brainstate.HiddenState(jnp.zeros(out_size))\n", + "\n", + " def update(self, x):\n", + " # Layer 1: x feeds in, h1 recurs\n", + " self.h1.value = jax.nn.tanh(\n", + " x @ self.w1_in.value + braintrace.matmul(self.h1.value, self.w1_rec.value)\n", + " )\n", + " # Layer 2: h1 feeds in, h2 recurs\n", + " self.h2.value = jax.nn.tanh(\n", + " self.h1.value @ self.w2_in.value + braintrace.matmul(self.h2.value, self.w2_rec.value)\n", + " )\n", + " return self.h2.value\n", + "\n", + "\n", + "model_2layer = TwoLayerRNN(in_size=10, hidden_size=16, out_size=8)\n", + "\n", + "algo_2layer = braintrace.compile(model_2layer, braintrace.D_RTRL, jnp.zeros(10), batch_size=1)\n", + "algo_2layer.show_graph()" + ] }, { "cell_type": "markdown", @@ -526,7 +560,7 @@ "1. **Automatic discovery**: The compiler traces the model's Jaxpr and automatically identifies which states are recurrent. No manual annotation of hidden states is needed -- just use `brainstate`'s state classes.\n", "2. **Grouping**: Related hidden states are grouped together for efficient Jacobian computation. `HiddenGroupState` and `HiddenTreeState` explicitly declare a group; separate `HiddenState` variables are grouped by data flow analysis.\n", "3. **Operation-based selection**: Whether a weight participates in online learning depends on the operation used (`braintrace.matmul` vs. regular `@`), not on the parameter class.\n", - "4. **Flexible initialization**: Use `init_all_states` for single-sample or manual batching, and `vmap_new_states` for automatic batching." + "4. **Flexible initialization**: Use `init_all_states` for single-sample or manual batching, and `brainstate.nn.Map` plus `mapped_model.init_all_states()` for automatic batching." ] } ], @@ -551,4 +585,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/docs/tutorials/neural_network_layers.ipynb b/docs/tutorials/neural_network_layers.ipynb index cbfbb009..ebab4679 100644 --- a/docs/tutorials/neural_network_layers.ipynb +++ b/docs/tutorials/neural_network_layers.ipynb @@ -33,8 +33,6 @@ "metadata": {}, "outputs": [], "source": [ - "import warnings\n", - "\n", "import brainstate\n", "import jax.numpy as jnp\n", "\n", @@ -93,16 +91,11 @@ "\n", "model = TinySequenceModel()\n", "sample = jnp.ones(1)\n", - "with warnings.catch_warnings():\n", - " warnings.filterwarnings(\n", - " \"ignore\",\n", - " message=r\"ETP primitive etp_mv.*has no connected hidden states.*\",\n", - " )\n", - " learner = braintrace.compile(\n", - " model,\n", - " braintrace.D_RTRL,\n", - " sample,\n", - " )\n" + "learner = braintrace.compile(\n", + " model,\n", + " braintrace.D_RTRL,\n", + " sample,\n", + ")\n" ] }, { diff --git a/docs/tutorials/pp_prop.ipynb b/docs/tutorials/pp_prop.ipynb index 35b7f07e..54738c13 100644 --- a/docs/tutorials/pp_prop.ipynb +++ b/docs/tutorials/pp_prop.ipynb @@ -59,7 +59,6 @@ "import braintrace\n", "import jax.numpy as jnp\n", "import matplotlib.pyplot as plt\n", - "import warnings\n", "\n", "brainstate.random.seed(7)" ] @@ -102,19 +101,15 @@ "inputs = jnp.linspace(-1.0, 1.0, 12).reshape(12, 1, 1)\n", "targets = 0.7 * inputs + 0.2\n", "\n", - "with warnings.catch_warnings():\n", - " # The Linear readout is intentionally non-temporal; Section 5 explains why.\n", - " warnings.filterwarnings(\n", - " \"ignore\",\n", - " message=r\"ETP primitive etp_mm.*has no connected hidden states.*\",\n", - " )\n", - " learner = braintrace.compile(\n", - " model,\n", - " braintrace.pp_prop,\n", - " inputs[0],\n", - " batch_size=1,\n", - " decay_or_rank=0.9,\n", - " )\n", + "# The Linear readout is intentionally non-temporal; Section 5 explains\n", + "# the compiler diagnostic emitted for it.\n", + "learner = braintrace.compile(\n", + " model,\n", + " braintrace.pp_prop,\n", + " inputs[0],\n", + " batch_size=1,\n", + " decay_or_rank=0.9,\n", + ")\n", "weights = model.states(brainstate.ParamState)\n", "optimizer = braintools.optim.SGD(lr=0.08)\n", "optimizer.register_trainable_weights(weights);" diff --git a/docs/tutorials/rnn_online_learning.ipynb b/docs/tutorials/rnn_online_learning.ipynb index f90bbc50..505c0ca8 100644 --- a/docs/tutorials/rnn_online_learning.ipynb +++ b/docs/tutorials/rnn_online_learning.ipynb @@ -55,7 +55,6 @@ "import braintools\n", "import braintrace\n", "import matplotlib.pyplot as plt\n", - "import warnings\n", "\n", "brainstate.random.seed(17)" ] @@ -174,7 +173,7 @@ "\n", "The key steps in the online training loop are:\n", "\n", - "1. **Compile the model**: `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initialises hidden states, compiles the eligibility-trace graph, and returns a vmapped learner — all in one call.\n", + "1. **Compile the model**: `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initializes mapped hidden states, compiles the eligibility-trace graph, and returns the concrete learner - all in one call.\n", "2. **Warm-up phase**: Use `learner.etrace_evolve(...)` to advance hidden states and eligibility traces without computing a loss gradient.\n", "3. **Learning phase**: Use `learner.etrace_grad(..., step_fn=step_loss)` to drive the remaining sequence and accumulate online gradients.\n", "4. **Parameter update**: After processing the full sequence, apply the accumulated gradients to update the parameters.\n", @@ -205,12 +204,10 @@ " opt.register_trainable_weights(weights)\n", "\n", " batch_size = input_batches.shape[2]\n", - " with warnings.catch_warnings():\n", - " warnings.filterwarnings(\"ignore\", message=r\"ETP primitive .*\")\n", - " learner = braintrace.compile(\n", - " model, braintrace.D_RTRL, input_batches[0, 0],\n", - " batch_size=batch_size, vmap=True,\n", - " )\n", + " learner = braintrace.compile(\n", + " model, braintrace.D_RTRL, input_batches[0],\n", + " batch_size=batch_size, vmap=True,\n", + " )\n", "\n", " @brainstate.transform.jit\n", " def train_step(inputs, targets):\n", @@ -279,23 +276,19 @@ "\n", " @brainstate.transform.jit\n", " def train_step(inputs, targets):\n", - " @brainstate.transform.vmap_new_states(\n", - " state_tag=\"new\", axis_size=inputs.shape[1]\n", + " mapped_model = brainstate.nn.Map(\n", + " model, init_map_size=inputs.shape[1]\n", " )\n", - " def init():\n", - " brainstate.nn.init_all_states(model)\n", - "\n", - " init()\n", - " vmapped_model = brainstate.nn.Vmap(model, vmap_states=\"new\")\n", + " mapped_model.init_all_states()\n", "\n", " def run_step(inp, tar):\n", - " out = vmapped_model(inp)\n", + " out = mapped_model(inp)\n", " loss = braintools.metric.softmax_cross_entropy_with_integer_labels(out, tar).mean()\n", " return out, loss\n", "\n", " def bptt_forward():\n", " n_sim = time_lag + 10\n", - " brainstate.transform.for_loop(vmapped_model, inputs[:n_sim])\n", + " brainstate.transform.for_loop(mapped_model, inputs[:n_sim])\n", " outs, losses = brainstate.transform.for_loop(run_step, inputs[n_sim:], targets)\n", " return losses.mean(), outs\n", "\n", @@ -447,7 +440,7 @@ "**Key takeaways:**\n", "\n", "- **D-RTRL** provides approximate online gradients with `O(B * theta)` complexity, where `B` is the batch size and `theta` is the number of parameters. Unlike BPTT, it does not need to store the full unrolled computation graph.\n", - "- The online training loop uses `braintrace.compile` to set up the algorithm. A single call to `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initialises hidden states, compiles the eligibility-trace graph, and returns a vmapped learner ready for batched training.\n", + "- A single call to `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initializes `brainstate.nn.Map` states, compiles the eligibility-trace graph, and returns the concrete learner ready for batched training.\n", "- Online learning uses `learner.etrace_evolve` for gradient-free sequence prefixes and `learner.etrace_grad` for sequence objectives; both compose with `brainstate.transform.jit`.\n", "- `braintrace` is particularly effective for RNN models with gating mechanisms (GRU, LSTM), where the internal dynamics naturally support eligibility trace propagation.\n", "\n", diff --git a/docs/tutorials/snn_online_learning.ipynb b/docs/tutorials/snn_online_learning.ipynb index 21586a95..4459a398 100644 --- a/docs/tutorials/snn_online_learning.ipynb +++ b/docs/tutorials/snn_online_learning.ipynb @@ -62,7 +62,6 @@ "import brainunit as u\n", "import brainpy.state\n", "import matplotlib.pyplot as plt\n", - "import warnings\n", "\n", "brainstate.random.seed(31)" ] @@ -205,7 +204,7 @@ "\n", "Use the float form to set the decay directly, or the integer form to select the corresponding decay through the documented conversion. Neither form creates an independent rank-versus-memory trade-off.\n", "\n", - "`braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.5)` initialises per-sample states, builds the ETP graph, and returns a vmapped learner — no separate `init_all_states`, `compile_graph`, or `Vmap` calls are needed.\n", + "`braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.5)` initializes mapped per-sample states, builds the ETP graph, and returns the concrete algorithm — no separate `Map`, `init_all_states`, or `compile_graph` calls are needed.\n", "\n", "`braintrace.D_RTRL` is the alternative parameter-dimensional estimator. It stores parameter-shaped eligibility traces and can use substantially more memory; neither estimator is generally gradient-equivalent to BPTT outside its documented assumptions." ] @@ -253,12 +252,10 @@ " weights = model.states(brainstate.ParamState)\n", " opt.register_trainable_weights(weights)\n", "\n", - " with warnings.catch_warnings():\n", - " warnings.filterwarnings(\"ignore\", message=r\"ETP primitive .*\")\n", - " learner = braintrace.compile(\n", - " model, braintrace.pp_prop, input_batches[0, 0],\n", - " batch_size=batch_size, vmap=True, decay_or_rank=0.5,\n", - " )\n", + " learner = braintrace.compile(\n", + " model, braintrace.pp_prop, input_batches[0],\n", + " batch_size=batch_size, vmap=True, decay_or_rank=0.5,\n", + " )\n", "\n", " @brainstate.transform.jit\n", " def train_step(inputs, targets):\n", @@ -428,11 +425,11 @@ "\n", "1. **Model Construction**: Use `braintrace.nn.Linear` and `braintrace.nn.LeakyRateReadout` for layers that should participate in online learning (ETP-aware). Combine them with spiking neuron models from `brainpy.state` (e.g., `LIF`).\n", "\n", - "2. **Online Learning Setup**: Use `learner = braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.5)` to initialise states, compile the ETP graph, and return a vmapped learner in one call. Then use `learner.etrace_grad(...)` to drive the sequence and accumulate online gradients.\n", + "2. **Online Learning Setup**: Use `learner = braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.5)` to initialize mapped states, compile the ETP graph, and return the concrete learner in one call. Then use `learner.etrace_grad(...)` to drive the sequence and accumulate online gradients.\n", "\n", "3. **Scalability**: ES-D-RTRL achieves O(B(I+O)) memory complexity, making it practical for large spiking networks. The `decay_or_rank` parameter controls the trace approximation quality.\n", "\n", - "4. **Batching**: `braintrace.compile(..., batch_size=B, vmap=True)` handles per-sample state initialisation and vmapped execution automatically.\n", + "4. **Batching**: `braintrace.compile(..., batch_size=B, vmap=True)` handles `brainstate.nn.Map` initialization and mapped execution automatically.\n", "\n", "For more advanced topics, including training on real neuromorphic datasets (N-MNIST) and comparing online learning with BPTT, see:\n", "- [pp_prop algorithm tutorial](pp_prop.ipynb)\n", diff --git a/examples/002-coba-ei-rsnn.py b/examples/002-coba-ei-rsnn.py index 72d182c3..81511c45 100644 --- a/examples/002-coba-ei-rsnn.py +++ b/examples/002-coba-ei-rsnn.py @@ -297,12 +297,8 @@ def visualize(self, inputs, n2show: int = 5): n_seq = inputs.shape[0] batch_size = inputs.shape[1] - @brainstate.transform.vmap_new_states(state_tag='new', axis_size=batch_size) - def init(): - brainstate.nn.init_all_states(self) - - init() - model = brainstate.nn.Vmap(self, vmap_states='new') + model = brainstate.nn.Map(self, init_map_size=batch_size) + model.init_all_states() def step(inp): out = model(inp) diff --git a/examples/003-snn-memory-and-speed-evaluation-vmap.py b/examples/003-snn-memory-and-speed-evaluation-vmap.py index 4bbf3706..63c8eb0c 100644 --- a/examples/003-snn-memory-and-speed-evaluation-vmap.py +++ b/examples/003-snn-memory-and-speed-evaluation-vmap.py @@ -458,32 +458,34 @@ def _step(i, inp): return losses.mean(), acc def _compile_etrace_function(self, input_info): - # kept manual: this *is* compile(..., vmap=True)'s scheme -- - # vmap_new_states(state_tag='new') + init_all_states + compile_graph on - # the unbatched sample + a Vmap wrapper -- but compile cannot take this - # example input. It strips the batch axis with `a[0]`, and `input_info` - # is an unbatched jax.ShapeDtypeStruct, which is not subscriptable. - # (A benchmark builds the graph from a shape, never from real data.) + # Kept explicit because this benchmark compiles from an unbatched + # ShapeDtypeStruct rather than a concrete batched example input. + mapped_target = brainstate.nn.Map( + self.target, init_map_size=self.args.batch_size + ) + mapped_target.init_all_states() + if self.args.method == 'expsm_diag': - model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay, ) + model = braintrace.ES_D_RTRL( + mapped_target, self.args.etrace_decay, + ) elif self.args.method == 'diag': - model = braintrace.D_RTRL(self.target, ) + model = braintrace.D_RTRL(mapped_target) else: raise ValueError(f'Unknown online learning methods: {self.args.method}.') - # initialize the states - @brainstate.transform.vmap_new_states(state_tag='new', axis_size=self.args.batch_size) - def init(): - brainstate.nn.init_all_states(self.target) - model.compile_graph(input_info) - - init() - run_model = brainstate.nn.Vmap(model, vmap_states='new') + batched_input_info = jax.ShapeDtypeStruct( + (self.args.batch_size, *input_info.shape), input_info.dtype + ) + model.compile_graph(batched_input_info) + run_model = model @brainstate.transform.jit - @brainstate.transform.vmap(in_states=run_model.states('new')) def reset_state(): - brainstate.nn.reset_all_states(run_model) + brainstate.nn.reset_all_states( + self.target, batch_size=self.args.batch_size + ) + run_model.reset_state(batch_size=self.args.batch_size) @brainstate.transform.jit def _etrace_single_run(i, batch_inp): diff --git a/examples/004-feedforward-conv-snn.py b/examples/004-feedforward-conv-snn.py index 4dcf6845..22898993 100644 --- a/examples/004-feedforward-conv-snn.py +++ b/examples/004-feedforward-conv-snn.py @@ -246,25 +246,15 @@ def batch_train(self, inputs, targets): # inputs: [n_step, n_batch, ...] # targets: [n_batch, n_out] - # One call replaces init_all_states + compile_graph + Vmap. Pass the - # batched single step inputs[0]; compile strips axis 0 to recover the - # per-sample example, so this is the same graph the manual expansion in - # examples/drtrl/02-batching-vmap.py builds by hand. - # model = braintrace.compile(self.target, braintrace.ES_D_RTRL, inputs[0], - # batch_size=inputs.shape[1], vmap=True, - # decay_or_rank=self.decay_or_rank) + # compile wraps the model with Map, initializes per-sample states, and + # builds the graph from the batched single-step input. with brainstate.environ.context(fit=True): model = braintrace.compile( self.target, braintrace.D_RTRL, inputs[0], batch_size=inputs.shape[1], vmap=True, ) - # show_graph() is a post-compile diagnostic and lives on the learner, not - # on the vmap wrapper -- ETraceVmap forwards the drivers, not the - # introspection surface. Reading through .module is fine here; only - # *driving* through it would be wrong (it would drive the unbatched - # learner and give per-lane-wrong results). - model.module.show_graph() + model.show_graph() def _etrace_grad(inp): with brainstate.environ.context(fit=True): diff --git a/examples/100-gru-on-copying-task.py b/examples/100-gru-on-copying-task.py index 7cd3f5b7..9699647c 100644 --- a/examples/100-gru-on-copying-task.py +++ b/examples/100-gru-on-copying-task.py @@ -179,14 +179,11 @@ def batch_train(self, inputs, targets): # 需要求解梯度的参数 weights = self.target.states(brainstate.ParamState) - # kept manual: BPTT baseline — no online algorithm to migrate - # initialize the states - @brainstate.transform.vmap_new_states(state_tag='new', axis_size=inputs.shape[1]) - def init(): - brainstate.nn.init_all_states(self.target) - - init() - model = brainstate.nn.Vmap(self.target, vmap_states='new') + # kept manual: BPTT baseline, with mapped per-sample states + model = brainstate.nn.Map( + self.target, init_map_size=inputs.shape[1] + ) + model.init_all_states() def _run_step_train(inp, tar): out = model(inp) diff --git a/examples/drtrl/02-batching-vmap.py b/examples/drtrl/02-batching-vmap.py index 76047b4c..51fe06c3 100644 --- a/examples/drtrl/02-batching-vmap.py +++ b/examples/drtrl/02-batching-vmap.py @@ -1,9 +1,9 @@ # Copyright 2026 BrainX Ecosystem Limited. Licensed under the Apache License, 2.0. """02 Batching with the canonical public learner workflow. -``braintrace.compile(..., vmap=True)`` creates one eligibility-trace learner -per batch lane. Drive the returned learner directly with ``etrace_grad`` so -each sample keeps its own hidden and eligibility-trace state. +``braintrace.compile(..., vmap=True)`` wraps the model in +``brainstate.nn.Map``, initializes independent per-sample states, and returns +the compiled learner. Drive that learner directly with ``etrace_grad``. """ import pathlib @@ -63,7 +63,7 @@ def step_loss(inp, tar): plt.plot(losses); plt.xlabel('epoch'); plt.ylabel('MSE') - plt.title('02 Batching via public learner workflow'); + plt.title('02 Batching via public Map-backed learner workflow'); plt.show() return {"losses": losses} diff --git a/examples/pp_prop/05-batching-vmap.py b/examples/pp_prop/05-batching-vmap.py index ad91bf38..e10b94c7 100644 --- a/examples/pp_prop/05-batching-vmap.py +++ b/examples/pp_prop/05-batching-vmap.py @@ -1,16 +1,10 @@ # Copyright 2026 BrainX Ecosystem Limited. Licensed under the Apache License, 2.0. -"""05 · Batching via ``braintrace.compile(..., vmap=True)``. +# ``online_train_epoch`` uses braintrace.compile(..., vmap=True). +"""05 - Batching via ``braintrace.compile(..., vmap=True)``. -The network and the pp_prop algorithm are defined unbatched; ``compile`` with -``vmap=True`` replicates them across the batch dimension (it initializes the -states inside a ``vmap_new_states`` scope, builds the eligibility-trace graph -on one unbatched sample, and returns an ``ETraceVmap``). pp_prop's per-rule -init is aware of batching and allocates batched eligibility traces -automatically. This is the default batching path used by examples 01-04, and -it lives in ``_shared.online_train_epoch``, which this file calls. - -For the same three steps written out by hand, see -``examples/drtrl/02-batching-vmap.py``. +The model keeps single-sample update logic. ``braintrace.compile`` wraps it in +``brainstate.nn.Map``, initializes per-sample states, and compiles pp_prop from +one batched time step. """ import pathlib diff --git a/examples/pp_prop/12-classification-neuromorphic.py b/examples/pp_prop/12-classification-neuromorphic.py index 5d5c2730..5e094d7d 100644 --- a/examples/pp_prop/12-classification-neuromorphic.py +++ b/examples/pp_prop/12-classification-neuromorphic.py @@ -36,14 +36,11 @@ def _accuracy(outputs_seq, labels): def _eval(model, inputs, labels): - # kept manual: eval re-init, no online construction - @brainstate.transform.vmap_new_states(state_tag="new", axis_size=inputs.shape[1]) - def init(): - brainstate.nn.init_all_states(model) - - init() - vmap_model = brainstate.nn.Vmap(model, vmap_states="new") - outs = brainstate.transform.for_loop(lambda x: vmap_model(x), inputs) + mapped_model = brainstate.nn.Map( + model, init_map_size=inputs.shape[1] + ) + mapped_model.init_all_states() + outs = brainstate.transform.for_loop(lambda x: mapped_model(x), inputs) return _accuracy(outs, labels) diff --git a/examples/pp_prop/14-knob-vjp-method-contrast.py b/examples/pp_prop/14-knob-vjp-method-contrast.py index b3470578..02c9997c 100644 --- a/examples/pp_prop/14-knob-vjp-method-contrast.py +++ b/examples/pp_prop/14-knob-vjp-method-contrast.py @@ -35,14 +35,11 @@ def _accuracy(outputs_seq, labels): def _eval(model, inputs, labels): - # kept manual: eval re-init, no online construction - @brainstate.transform.vmap_new_states(state_tag="new", axis_size=inputs.shape[1]) - def init(): - brainstate.nn.init_all_states(model) - - init() - vmap_model = brainstate.nn.Vmap(model, vmap_states="new") - outs = brainstate.transform.for_loop(lambda x: vmap_model(x), inputs) + mapped_model = brainstate.nn.Map( + model, init_map_size=inputs.shape[1] + ) + mapped_model.init_all_states() + outs = brainstate.transform.for_loop(lambda x: mapped_model(x), inputs) return _accuracy(outs, labels) diff --git a/examples/pp_prop/README.md b/examples/pp_prop/README.md index dd5ae14c..f45611de 100644 --- a/examples/pp_prop/README.md +++ b/examples/pp_prop/README.md @@ -34,7 +34,7 @@ if sklearn is missing). | 02 | `02-neurons-alif-dms.py` | ALIF (adaptive threshold) on delayed-match-to-sample | | 03 | `03-neurons-gif-working-memory.py` | GIF with heterogeneous tau_I2 on working-memory recall | | 04 | `04-neurons-coba-ei-rsnn.py` | Dale-law E/I RSNN on small Poisson-MNIST | -| 05 | `05-batching-vmap.py` | Batching via `brainstate.nn.Vmap(vmap_states='new')` | +| 05 | `05-batching-vmap.py` | Batching via `brainstate.nn.Map` | | 06 | `06-batching-batched.py` | Batching via the batched ETP primitive path | | 07 | `07-vjp-single-step.py` | `vjp_method='single-step'` (default) | | 08 | `08-vjp-multi-step.py` | `vjp_method='multi-step'` for temporal credit | diff --git a/examples/pp_prop/_shared.py b/examples/pp_prop/_shared.py index c03ee6b4..a42c2fe2 100644 --- a/examples/pp_prop/_shared.py +++ b/examples/pp_prop/_shared.py @@ -389,16 +389,14 @@ def bptt_train_epoch_fixed_target( """BPTT baseline with per-step softmax-cross-entropy over a fixed label.""" weights = model.states(brainstate.ParamState) - # kept manual: BPTT re-init — no algorithm construction, no compile_graph - @brainstate.transform.vmap_new_states(state_tag="new", axis_size=inputs.shape[1]) - def init(): - brainstate.nn.init_all_states(model) - - init() - vmap_model = brainstate.nn.Vmap(model, vmap_states="new") + # kept manual: BPTT baseline, with mapped per-sample states + mapped_model = brainstate.nn.Map( + model, init_map_size=inputs.shape[1] + ) + mapped_model.init_all_states() def run_step(inp): - out = vmap_model(inp) + out = mapped_model(inp) loss = braintools.metric.softmax_cross_entropy_with_integer_labels( out, target_labels ).mean() diff --git a/examples/snn_models.py b/examples/snn_models.py index c2fefdfb..a6652bae 100644 --- a/examples/snn_models.py +++ b/examples/snn_models.py @@ -143,12 +143,8 @@ def verify(self, input_spikes, num_show=5, sps_inc=10.): xs = np.transpose(input_spikes, (1, 0, 2)) # [n_steps, n_samples, n_in] # 运行仿真模型 - @brainstate.transform.vmap_new_states(state_tag='new', axis_size=xs.shape[1]) - def init(): - brainstate.nn.init_all_states(self) - - init() - model = brainstate.nn.Vmap(self, vmap_states='new') + model = brainstate.nn.Map(self, init_map_size=xs.shape[1]) + model.init_all_states() outs, sps, vs = brainstate.transform.for_loop( lambda x: (model(x), self.r.get_spike(), self.r.V.value), @@ -448,14 +444,11 @@ class BPTTTrainer(Trainer): def batch_train(self, inputs, targets): weights = self.target.states().subset(brainstate.ParamState) - # kept manual: BPTT baseline — no online algorithm to migrate - # initialize the states - @brainstate.transform.vmap_new_states(state_tag='new', axis_size=inputs.shape[1]) - def init(): - brainstate.nn.init_all_states(self.target) - - init() - model = brainstate.nn.Vmap(self.target, vmap_states='new') + # kept manual: BPTT baseline, with mapped per-sample states + model = brainstate.nn.Map( + self.target, init_map_size=inputs.shape[1] + ) + model.init_all_states() # the model for a single step def _run_step_train(inp): @@ -521,12 +514,8 @@ def update(self, spk): @brainstate.transform.jit(static_argnums=0) def eval(self, xs): - @brainstate.transform.vmap_new_states(state_tag='new', axis_size=xs.shape[1]) - def init(): - brainstate.nn.init_all_states(self) - - init() - model = brainstate.nn.Vmap(self, vmap_states='new') + model = brainstate.nn.Map(self, init_map_size=xs.shape[1]) + model.init_all_states() outs, sps, vs = brainstate.transform.for_loop( lambda x: (model(x), self.neu.get_spike(), self.neu.V.value), xs From 4001dba872a04d52f4f55b43e826a981c43b5cc4 Mon Sep 17 00:00:00 2001 From: poilsosart <128177087+poilsosart@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:34:35 +0800 Subject: [PATCH 2/4] fix: align sequence drivers with Map semantics --- braintrace/_algorithm/base.py | 3 + braintrace/_algorithm/sequence.py | 40 +++--- braintrace/_algorithm/sequence_test.py | 188 ++++++++++++++++++------- 3 files changed, 163 insertions(+), 68 deletions(-) diff --git a/braintrace/_algorithm/base.py b/braintrace/_algorithm/base.py index ebd7dafb..5e74b6a0 100644 --- a/braintrace/_algorithm/base.py +++ b/braintrace/_algorithm/base.py @@ -233,6 +233,9 @@ def other_states(self) -> brainstate.util.FlattedDict: @property def _seq_param_states(self) -> brainstate.util.FlattedDict: """Sequence-driver hook: the default set of weights to differentiate.""" + model = self.graph_executor.model + if isinstance(model, brainstate.nn.Map): + return model.module.states(brainstate.ParamState) return self.param_states @property diff --git a/braintrace/_algorithm/sequence.py b/braintrace/_algorithm/sequence.py index 1b1d2fbc..174f9dc5 100644 --- a/braintrace/_algorithm/sequence.py +++ b/braintrace/_algorithm/sequence.py @@ -168,9 +168,10 @@ class SequenceDriverMixin: bypass the window-mode validation entirely. """ - #: Set by :class:`ETraceVmap`. Window mode is refused when true, because - #: ``compile(vmap=True)`` maps ``in_axes=0`` and a ``(k, B, ...)`` window - #: slice would map *time* as the batch axis. + #: Set by :class:`ETraceVmap`. Window mode is refused for that compatibility + #: wrapper because its ``in_axes=0`` would map a window's time axis. + #: ``compile(vmap=True)`` now uses ``brainstate.nn.Map``, keeps the + #: batch axis inside the graph, and therefore leaves this false. _seq_is_vmapped: bool = False @property @@ -196,12 +197,12 @@ def _seq_check_window(self, chunk_size: Optional[int], length: int, if self._seq_is_vmapped: raise ValueError( f'chunk_size={chunk_size} (window mode) is not supported under ' - f'a vmapped learner. compile(vmap=True) maps in_axes=0, so a ' - f'(chunk_size, batch, ...) window slice would map the *time* ' - f'axis as the batch axis -- which is silently wrong whenever ' - f'chunk_size equals the batch size. Use the batched (non-vmap) ' - f'mode, which carries the batch axis inside the compiled graph, ' - f'or drive with chunk_size=None.' + f'an ETraceVmap learner. Its in_axes=0 mapping would treat the ' + f'*time* axis of a (chunk_size, batch, ...) window slice as the ' + f'batch axis -- which is silently wrong whenever ' + f'chunk_size equals the batch size. Use compile(vmap=True), ' + f'which carries the batch axis inside a brainstate.nn.Map, or ' + f'drive this compatibility wrapper with chunk_size=None.' ) if for_grad and self._seq_vjp_method != 'multi-step': @@ -266,7 +267,7 @@ def etrace_grad( ``step_fn`` must return a ``(k,)`` vector of per-step losses and wrap its model inputs in :class:`MultiStepData`. Window mode requires ``vjp_method='multi-step'`` and ``T % k == 0``, and is not - available under a vmapped learner. + available under an :class:`ETraceVmap` compatibility wrapper. weights : dict, optional The :class:`brainstate.ParamState` to differentiate. Defaults to the learner's own ``param_states``. @@ -302,8 +303,9 @@ def etrace_grad( ``T == 0`` -- there is nothing to slice. If *chunk_size* is below ``1``; if ``k >= 2`` but the learner's ``vjp_method`` is not ``'multi-step'`` (the executor would raise three frames down), the - learner is vmapped (``in_axes=0`` would map time as the batch - axis), or ``T % k != 0``. If *mask* is not shape ``(T,)``. If + learner is an :class:`ETraceVmap` (``in_axes=0`` would map time as + the batch axis), or ``T % k != 0``. If *mask* is not shape + ``(T,)``. If *reduction* or *loss_output* is not one of its legal values. If ``step_fn`` returns a non-scalar in plain mode, or anything but shape ``(k,)`` in window mode. If the learner has not been @@ -498,7 +500,8 @@ def etrace_evolve( As in :meth:`etrace_grad`, **except** that a window is *not* refused for being on a single-step learner -- no loss VJP is taken here, so the restriction does not apply. Windows are still refused - under a vmapped learner, and ``T % chunk_size == 0`` still holds. + under an :class:`ETraceVmap` compatibility wrapper, and + ``T % chunk_size == 0`` still holds. Examples -------- @@ -537,11 +540,12 @@ def body(*slices): class ETraceVmap(SequenceDriverMixin, brainstate.nn.Vmap): - """Provide sequence drivers on a ``brainstate.nn.Vmap`` wrapper. + """Provide sequence drivers for explicitly constructed legacy Vmap models. - Returned by ``braintrace.compile(..., vmap=True)`` so the call site is - identical in batched and unbatched mode. Because it *is* a - ``brainstate.nn.Vmap``, calling it, its attributes and every + ``braintrace.compile(..., vmap=True)`` uses ``brainstate.nn.Map`` and + returns the algorithm directly. This public compatibility type remains for + callers that explicitly construct a ``brainstate.nn.Vmap`` learner. + Because it *is* a ``brainstate.nn.Vmap``, calling it and every ``isinstance(x, brainstate.nn.Vmap)`` check keep working; only the added methods are new. @@ -555,7 +559,7 @@ class ETraceVmap(SequenceDriverMixin, brainstate.nn.Vmap): would drive the **unbatched** learner and silently produce per-lane-wrong results. - Window mode is refused here -- see + Window mode is refused here; see :meth:`SequenceDriverMixin.etrace_grad`. """ __module__ = 'braintrace' diff --git a/braintrace/_algorithm/sequence_test.py b/braintrace/_algorithm/sequence_test.py index aa52b743..f9a887e8 100644 --- a/braintrace/_algorithm/sequence_test.py +++ b/braintrace/_algorithm/sequence_test.py @@ -34,6 +34,7 @@ import inspect import brainstate +import braintools import brainunit as u import jax import jax.numpy as jnp @@ -609,20 +610,18 @@ def test_a_weighted_mask_reweights_the_objective(self): # --------------------------------------------------------------------------- -# 18--20. vmap +# 18--20. Map # --------------------------------------------------------------------------- class _VmapNet(brainstate.nn.Module): """A model whose hidden state is created in ``init_state``, not ``__init__``. - ``om.tanh_rnn`` allocates its ``HiddenState`` in ``__init__``, so - ``vmap_new_states`` has nothing to batch and ``compile(vmap=True)`` raises - ``BatchAxisError`` on the first call -- the whole vmap section was - previously written against a fixture that could not run. ``ValinaRNNCell`` - defers its state to ``init_state``, which is the property that matters. + ``ValinaRNNCell`` defers its state to ``init_state``, allowing + ``brainstate.nn.Map.init_all_states`` to create one independent state per + mapped lane. ``wout`` is a plain (non-ETP) parameter, so it is exactly zero under - ``vjp_method='single-step'`` (F-33); the vmap fixture therefore runs + ``vjp_method='single-step'`` (F-33); the Map fixture therefore runs ``'multi-step'``, which keeps every key live and the comparisons honest. """ @@ -637,7 +636,7 @@ def update(self, x): return self.cell(x) @ self.wout.value -def _vmap_learner(batch, vjp_method='multi-step', **opts): +def _map_learner(batch, vjp_method='multi-step', **opts): return braintrace.compile(_VmapNet(), 'D_RTRL', jnp.zeros((batch, N_IN)), batch_size=batch, vmap=True, vjp_method=vjp_method, **opts) @@ -662,20 +661,46 @@ def _lane_data(batch, *, seed=7): return xs, ys -class TestVmap: - def test_the_vmapped_learner_carries_the_driver_methods(self): +class TestMap: + def test_the_mapped_learner_carries_the_driver_methods(self): """Spec test 18. - ``compile(vmap=True)`` must return something that *has* ``etrace_grad``; - before this change it returned a bare ``brainstate.nn.Vmap``, which does - not. Reaching into ``.module`` instead would drive the unbatched learner - and silently give per-lane-wrong results. + ``compile(vmap=True)`` returns the algorithm itself with a mapped model, + so sequence drivers and compilation reports remain directly available. """ - learner = _vmap_learner(3) + learner = _map_learner(3) + assert isinstance(learner, braintrace.ETraceAlgorithm) + assert isinstance(learner.graph_executor.model, brainstate.nn.Map) assert hasattr(learner, 'etrace_grad') assert hasattr(learner, 'etrace_evolve') - def test_the_vmapped_gradient_is_the_sum_over_independent_lanes(self): + def test_default_gradients_keep_the_original_model_parameter_paths(self): + """Map internals must not leak into the public optimizer contract.""" + batch = 3 + xs, ys = _lane_data(batch) + model = _VmapNet() + weights = model.states(brainstate.ParamState) + optimizer = braintools.optim.Adam(lr=1e-3) + optimizer.register_trainable_weights(weights) + learner = braintrace.compile( + model, + 'D_RTRL', + jnp.zeros((batch, N_IN)), + batch_size=batch, + vmap=True, + vjp_method='multi-step', + ) + + grads = learner.etrace_grad( + xs, + ys, + step_fn=lambda inp, tar: jnp.sum((learner(inp) - tar) ** 2), + ) + + assert set(grads) == set(weights) + optimizer.update(grads) + + def test_the_mapped_gradient_is_the_sum_over_independent_lanes(self): """Spec test 18, the part that has content. The parameters are shared across lanes, so the batched gradient must be @@ -688,7 +713,7 @@ def test_the_vmapped_gradient_is_the_sum_over_independent_lanes(self): batch = 3 xs, ys = _lane_data(batch) - batched = _vmap_learner(batch) + batched = _map_learner(batch) def step_fn(inp, tar): return jnp.sum((batched(inp) - tar) ** 2) @@ -708,13 +733,16 @@ def lane_step(inp, tar, learner=lane): lambda a, b: a + b, lane_total, g_lane) flat = _arrays(g_batched) - assert set(flat) == set(_arrays(lane_total)), 'gradient keys diverged' + lane_flat = _arrays(lane_total) + assert set(flat) == set(lane_flat), 'gradient keys diverged' for k, v in flat.items(): assert np.max(np.abs(v)) > 0.0, f'{k} is identically zero -- vacuous' - _assert_trees_equal(g_batched, lane_total, rtol=2e-6, atol=1e-6, - msg='batched vs sum over independent lanes') + np.testing.assert_allclose( + v, lane_flat[k], rtol=2e-6, atol=1e-6, + err_msg=f'mapped vs sum over independent lanes: at {k}', + ) - def test_permuting_the_lanes_changes_the_vmapped_gradient(self): + def test_permuting_the_lanes_changes_the_mapped_gradient(self): """Spec test 18, the negative control. Without this, the sum-over-lanes identity could hold for a driver that @@ -725,11 +753,11 @@ def test_permuting_the_lanes_changes_the_vmapped_gradient(self): xs, ys = _lane_data(batch) perm = jnp.asarray([1, 0, 2]) - straight = _vmap_learner(batch) + straight = _map_learner(batch) g_straight = straight.etrace_grad( xs, ys, step_fn=lambda i, t: jnp.sum((straight(i) - t) ** 2)) - swapped = _vmap_learner(batch) + swapped = _map_learner(batch) g_swapped = swapped.etrace_grad( xs, ys[:, perm], step_fn=lambda i, t: jnp.sum((swapped(i) - t) ** 2)) @@ -737,50 +765,110 @@ def test_permuting_the_lanes_changes_the_vmapped_gradient(self): oracle.assert_gradients_differ(_arrays(g_straight), _arrays(g_swapped), min_rel=1e-3) - @pytest.mark.parametrize('batch', [3, K]) # B != k, and the silent B == k case - def test_window_mode_is_refused_under_vmap(self, batch): + @pytest.mark.parametrize('batch', [3, K]) + def test_window_mode_matches_independent_lane_windows_under_map(self, batch): """Spec test 19. - ``compile(vmap=True)`` maps ``in_axes=0``, so a ``(k, B, ...)`` window - slice would map *time* as the batch axis. At ``B != k`` that is a loud - shape error; at ``B == k`` the shapes line up and it would train on - transposed data. The ``B == k`` parametrization is the one that matters. + Map keeps the batch axis inside the compiled graph, so a + ``(K, batch, ...)`` window must equal the sum of independently driven + ``(K, 1, ...)`` lanes. Comparing with the plain path is not a valid + oracle: a multi-step update and K single-step updates are different + approximation regimes for this recurrent fixture. """ - learner = _vmap_learner(batch) - xs = jnp.zeros((T, batch, N_IN)) + xs, ys = _lane_data(batch) + mapped = _map_learner(batch) - with pytest.raises(ValueError, match='vmap'): - learner.etrace_grad(xs, step_fn=lambda x: jnp.zeros(K), - chunk_size=K) - with pytest.raises(ValueError, match='vmap'): - learner.etrace_evolve(xs, chunk_size=K) + def mapped_window_loss(inp, tar): + out = mapped(braintrace.MultiStepData(inp)) + return jnp.sum((out - tar) ** 2, axis=(1, 2)) + + g_mapped = mapped.etrace_grad( + xs, + ys, + step_fn=mapped_window_loss, + chunk_size=K, + reduction='sum', + ) + + lane_total = None + for lane_index in range(batch): + lane = _lane_learner() + + def lane_window_loss(inp, tar, learner=lane): + out = learner(braintrace.MultiStepData(inp)) + return jnp.sum((out - tar) ** 2, axis=(1, 2)) + + g_lane = lane.etrace_grad( + xs[:, lane_index:lane_index + 1], + ys[:, lane_index:lane_index + 1], + step_fn=lane_window_loss, + chunk_size=K, + reduction='sum', + ) + lane_total = g_lane if lane_total is None else jax.tree.map( + lambda a, b: a + b, lane_total, g_lane) + + mapped_arrays = _arrays(g_mapped) + lane_arrays = _arrays(lane_total) + assert set(mapped_arrays) == set(lane_arrays), 'gradient keys diverged' + for key, mapped_value in mapped_arrays.items(): + np.testing.assert_allclose( + mapped_value, + lane_arrays[key], + rtol=2e-6, + atol=1e-6, + err_msg=f'Map window vs independent lane windows: at {key}', + ) + + @pytest.mark.parametrize('batch', [3, K]) + def test_window_evolve_matches_independent_lanes_under_map(self, batch): + """Window evolution preserves separate time and mapped-lane axes.""" + xs, _ = _lane_data(batch) + mapped = _map_learner(batch) + mapped_outputs = mapped.etrace_evolve( + xs, chunk_size=K, return_outputs=True) + + lane_outputs = [] + for lane_index in range(batch): + lane = _lane_learner() + lane_outputs.append( + lane.etrace_evolve( + xs[:, lane_index:lane_index + 1], + chunk_size=K, + return_outputs=True, + ) + ) + expected = jnp.concatenate(lane_outputs, axis=2) + np.testing.assert_allclose( + np.asarray(mapped_outputs), + np.asarray(expected), + rtol=2e-6, + atol=1e-6, + ) @pytest.mark.parametrize('batch', [3, K]) - def test_chunk_size_one_is_admitted_under_vmap(self, batch): - """Spec test 19, the other half -- the refusal must not overreach. + def test_chunk_size_one_is_admitted_under_map(self, batch): + """Spec test 19, the plain path remains available under Map. - ``chunk_size=1`` is the plain path, so it carries none of the axis - collision that makes ``k >= 2`` unsafe. A guard written as - ``if chunk_size is not None`` would refuse it, which is why both - methods are exercised rather than just ``etrace_evolve``. + ``chunk_size=1`` is normalized to the ordinary step-by-step path. """ xs, ys = _lane_data(batch) - learner = _vmap_learner(batch) + learner = _map_learner(batch) learner.etrace_evolve(xs, chunk_size=1) - grad_learner = _vmap_learner(batch) + grad_learner = _map_learner(batch) grads = grad_learner.etrace_grad( xs, ys, chunk_size=1, step_fn=lambda i, t: jnp.sum((grad_learner(i) - t) ** 2)) for k, v in _arrays(grads).items(): assert np.all(np.isfinite(v)), f'{k} is not finite' - def test_the_vmap_return_value_is_still_a_brainstate_vmap(self): - """Spec test 20 -- existing ``vmap=True`` users must be unaffected.""" - learner = _vmap_learner(3) - assert isinstance(learner, brainstate.nn.Vmap) - assert isinstance(learner, braintrace.ETraceVmap) - assert isinstance(learner.module, braintrace.ETraceAlgorithm) + def test_vmap_option_returns_an_algorithm_with_a_mapped_model(self): + """Spec test 20 -- the option selects Map-based state initialization.""" + learner = _map_learner(3) + assert isinstance(learner, braintrace.ETraceAlgorithm) + assert not isinstance(learner, brainstate.nn.Vmap) + assert isinstance(learner.graph_executor.model, brainstate.nn.Map) # --------------------------------------------------------------------------- From 14eb4592f7fa00479a7c1b23721ce0dddbc88d2c Mon Sep 17 00:00:00 2001 From: poilsosart <128177087+poilsosart@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:20:34 +0800 Subject: [PATCH 3/4] fix: canonicalize Map state paths --- braintrace/_algorithm/base.py | 3 --- braintrace/_algorithm/sequence_test.py | 8 ++++---- braintrace/_compile_test.py | 2 +- braintrace/_compiler/module_info.py | 9 ++++++--- braintrace/_compiler/module_info_test.py | 7 +++++++ .../003-snn-memory-and-speed-evaluation-all.py | 13 ++++--------- ...-snn-memory-and-speed-evaluation-batched.py | 13 ++++--------- examples/snn_models.py | 2 +- examples/tests/test_compile_modes.py | 18 +++++++++++++++++- 9 files changed, 44 insertions(+), 31 deletions(-) diff --git a/braintrace/_algorithm/base.py b/braintrace/_algorithm/base.py index 5e74b6a0..ebd7dafb 100644 --- a/braintrace/_algorithm/base.py +++ b/braintrace/_algorithm/base.py @@ -233,9 +233,6 @@ def other_states(self) -> brainstate.util.FlattedDict: @property def _seq_param_states(self) -> brainstate.util.FlattedDict: """Sequence-driver hook: the default set of weights to differentiate.""" - model = self.graph_executor.model - if isinstance(model, brainstate.nn.Map): - return model.module.states(brainstate.ParamState) return self.param_states @property diff --git a/braintrace/_algorithm/sequence_test.py b/braintrace/_algorithm/sequence_test.py index f9a887e8..ffdb4af6 100644 --- a/braintrace/_algorithm/sequence_test.py +++ b/braintrace/_algorithm/sequence_test.py @@ -679,9 +679,6 @@ def test_default_gradients_keep_the_original_model_parameter_paths(self): batch = 3 xs, ys = _lane_data(batch) model = _VmapNet() - weights = model.states(brainstate.ParamState) - optimizer = braintools.optim.Adam(lr=1e-3) - optimizer.register_trainable_weights(weights) learner = braintrace.compile( model, 'D_RTRL', @@ -690,6 +687,8 @@ def test_default_gradients_keep_the_original_model_parameter_paths(self): vmap=True, vjp_method='multi-step', ) + optimizer = braintools.optim.Adam(lr=1e-3) + optimizer.register_trainable_weights(learner.param_states) grads = learner.etrace_grad( xs, @@ -697,7 +696,8 @@ def test_default_gradients_keep_the_original_model_parameter_paths(self): step_fn=lambda inp, tar: jnp.sum((learner(inp) - tar) ** 2), ) - assert set(grads) == set(weights) + assert set(learner.param_states) == set(model.states(brainstate.ParamState)) + assert set(grads) == set(learner.param_states) optimizer.update(grads) def test_the_mapped_gradient_is_the_sum_over_independent_lanes(self): diff --git a/braintrace/_compile_test.py b/braintrace/_compile_test.py index 18b0afdb..9dfd6080 100644 --- a/braintrace/_compile_test.py +++ b/braintrace/_compile_test.py @@ -295,7 +295,7 @@ def test_compile_vmap_returns_algorithm_exposing_report(): # --- both-modes coverage across RNN architectures + algorithms --------------- # Each architecture/algorithm must build, forward, and back-prop a finite, # non-zero gradient under BOTH compile(vmap=False) (internal batch primitive) -# and compile(vmap=True) (per-sample vmap lanes). A multi-step scan exercises +# and compile(vmap=True) (Map-owned per-sample states). A multi-step scan exercises # the eligibility trace (single-step would never engage it). _NI, _NR = 3, 4 diff --git a/braintrace/_compiler/module_info.py b/braintrace/_compiler/module_info.py index 334ba762..290db31a 100644 --- a/braintrace/_compiler/module_info.py +++ b/braintrace/_compiler/module_info.py @@ -194,10 +194,13 @@ def abstractify_model( "The model should be an instance of brainstate.nn.Module. " "Since it allows the explicit definition of the model structure." ) - model_retrieved_states = brainstate.graph.states(model) if isinstance(model, brainstate.nn.Map): - for path, state in brainstate.graph.states(model.module).items(): - model_retrieved_states[('module', *path)] = state + # ``Map`` exposes the same states through implementation paths such as + # ``module`` and ``dict_vmap_states``. Compiler paths are public model + # paths, so retrieve them from the wrapped module directly. + model_retrieved_states = brainstate.graph.states(model.module) + else: + model_retrieved_states = brainstate.graph.states(model) # --- stateful model, for extracting states, weights, and variables --- # # diff --git a/braintrace/_compiler/module_info_test.py b/braintrace/_compiler/module_info_test.py index c0d06739..ace0c8d7 100644 --- a/braintrace/_compiler/module_info_test.py +++ b/braintrace/_compiler/module_info_test.py @@ -70,12 +70,19 @@ def test_map_hidden_state_aliases_are_deduplicated(self): mapped = brainstate.nn.Map(rnn, init_map_size=batch_size) mapped.init_all_states() + expected_states = brainstate.graph.states(rnn) + expected_params = rnn.states(brainstate.ParamState) + minfo = braintrace.extract_module_info( mapped, brainstate.random.rand(batch_size, 2) ) states = minfo.retrieved_model_states + assert set(states) == set(expected_states) + assert set(minfo.weight_path_to_invars) == set(expected_params) + assert all(states[path] is state for path, state in expected_states.items()) assert len({id(state) for state in states.values()}) == len(states) + assert all('module' not in path for path in states) assert all('dict_vmap_states' not in path for path in states) @pytest.mark.parametrize( diff --git a/examples/003-snn-memory-and-speed-evaluation-all.py b/examples/003-snn-memory-and-speed-evaluation-all.py index 171cb2da..a102d1bc 100644 --- a/examples/003-snn-memory-and-speed-evaluation-all.py +++ b/examples/003-snn-memory-and-speed-evaluation-all.py @@ -353,15 +353,10 @@ def _step(i, inp): def _compile_etrace_function(self, input_info): # kept manual: braintrace.compile has no path for this state scheme. - # It offers two: init_all_states(batch_size=B) (vmap=False), or - # vmap_new_states(state_tag='new') + compile_graph on an *unbatched* - # sample + an ETraceVmap wrapper (vmap=True). This benchmark uses a - # third -- vmap_init_all_states(state_tag='new') for the per-sample - # states, compile_graph on the *batched* example, no wrapper, and an - # explicit brainstate.transform.vmap(in_states=...) only for the reset. - # compile's vmap branch would also reject `input_info`: it strips the - # batch axis with `a[0]`, and a jax.ShapeDtypeStruct is not - # subscriptable. + # It owns explicitly tagged per-sample states and resets those states + # through ``transform.vmap``. ``compile(vmap=True)`` instead wraps the + # target in ``brainstate.nn.Map``, which would change this benchmark's + # state ownership and reset contract. if self.args.method == 'expsm_diag': model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay) elif self.args.method == 'diag': diff --git a/examples/003-snn-memory-and-speed-evaluation-batched.py b/examples/003-snn-memory-and-speed-evaluation-batched.py index fb23afc6..22c8077d 100644 --- a/examples/003-snn-memory-and-speed-evaluation-batched.py +++ b/examples/003-snn-memory-and-speed-evaluation-batched.py @@ -477,15 +477,10 @@ def _step(i, inp): def _compile_etrace_function(self, input_info): # kept manual: braintrace.compile has no path for this state scheme. - # It offers two: init_all_states(batch_size=B) (vmap=False), or - # vmap_new_states(state_tag='new') + compile_graph on an *unbatched* - # sample + an ETraceVmap wrapper (vmap=True). This benchmark uses a - # third -- vmap_init_all_states(state_tag='new') for the per-sample - # states, compile_graph on the *batched* example, no wrapper, and an - # explicit brainstate.transform.vmap(in_states=...) only for the reset. - # compile's vmap branch would also reject `input_info`: it strips the - # batch axis with `a[0]`, and a jax.ShapeDtypeStruct is not - # subscriptable. + # It owns explicitly tagged per-sample states and resets those states + # through ``transform.vmap``. ``compile(vmap=True)`` instead wraps the + # target in ``brainstate.nn.Map``, which would change this benchmark's + # state ownership and reset contract. if self.args.method == 'expsm_diag': model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay) elif self.args.method == 'diag': diff --git a/examples/snn_models.py b/examples/snn_models.py index a6652bae..d33e5d48 100644 --- a/examples/snn_models.py +++ b/examples/snn_models.py @@ -404,7 +404,7 @@ def batch_train(self, inputs, targets): model = braintrace.compile(self.target, braintrace.pp_prop, inputs[0], batch_size=inputs.shape[1], vmap=True, decay_or_rank=self.decay_or_rank) - model.module.show_graph() + model.show_graph() def _etrace_grad(inp): # call the model diff --git a/examples/tests/test_compile_modes.py b/examples/tests/test_compile_modes.py index ebd0c690..417f80ec 100644 --- a/examples/tests/test_compile_modes.py +++ b/examples/tests/test_compile_modes.py @@ -1,7 +1,7 @@ # Copyright 2026 BrainX Ecosystem Limited. Licensed under the Apache License, 2.0. """Verify the example SNN cells compile and run under BOTH ``braintrace.compile(vmap=False)`` (batched, internal batch primitive) and -``braintrace.compile(vmap=True)`` (per-sample vmap lanes). +``braintrace.compile(vmap=True)`` (``brainstate.nn.Map`` state ownership). The custom ``GIF`` neuron in ``snn_models.py`` originally defined ``init_state(self)`` without ``batch_size``, so the non-vmap path @@ -88,3 +88,19 @@ def test_gif_neuron_init_state_accepts_batch_size(): brainstate.nn.init_all_states(neu, batch_size=B) assert neu.V.value.shape == (B, N_REC) assert neu.I2.value.shape == (B, N_REC) + + +def test_mapped_compile_exposes_show_graph_directly(): + """Mapped compile results expose reports without a wrapper ``.module``.""" + xs = jnp.zeros((B, N_IN)) + learner = braintrace.compile( + braintrace.nn.GRUCell(N_IN, N_REC), braintrace.D_RTRL, xs, + batch_size=B, vmap=True, + ) + + report = learner.show_graph(verbose=False, return_msg=True) + + assert isinstance(report, str) + assert "model.module.show_graph()" not in ( + EXAMPLES_DIR / "snn_models.py" + ).read_text(encoding="utf-8") From fa00494c2620e082aad5389622619c39e6b8137b Mon Sep 17 00:00:00 2001 From: poilsosart <128177087+poilsosart@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:26:56 +0800 Subject: [PATCH 4/4] docs: use explicit Map initialization workflow --- braintrace/_algorithm/base.py | 12 +- braintrace/_algorithm/sequence.py | 40 ++-- braintrace/_algorithm/sequence_test.py | 188 +++++------------- .../tests/conv_vmap_correctness_test.py | 64 +++--- .../tests/diagnostic_exploration_test.py | 21 +- .../_algorithm/tests/while_support_test.py | 14 +- braintrace/_compile.py | 65 ++++-- braintrace/_compile_test.py | 29 ++- braintrace/_compiler/canonicalize_test.py | 22 +- braintrace/_compiler/hidden_group_test.py | 17 +- braintrace/_compiler/module_info.py | 38 +--- braintrace/_compiler/module_info_test.py | 21 -- braintrace/_compiler/scenario_catalog_test.py | 36 +++- .../tests/cell_relation_guardrail_test.py | 8 +- .../_compiler/tests/compiler_oracle_test.py | 20 +- .../_compiler/tests/compiler_property_test.py | 23 ++- braintrace/_legacy/_ops_test.py | 19 +- braintrace/_legacy/_params_test.py | 11 +- braintrace/_op/_primitive.py | 3 - braintrace/_op/conv.py | 48 +---- docs/advanced/batching.ipynb | 99 ++++----- ...6-07-28-warnings-and-map-initialization.md | 62 ------ docs/tutorials/rnn_online_learning.ipynb | 62 +++--- docs/tutorials/snn_online_learning.ipynb | 74 ++++--- examples/002-coba-ei-rsnn.py | 8 +- ...003-snn-memory-and-speed-evaluation-all.py | 13 +- ...snn-memory-and-speed-evaluation-batched.py | 13 +- ...03-snn-memory-and-speed-evaluation-vmap.py | 38 ++-- examples/004-feedforward-conv-snn.py | 16 +- examples/100-gru-on-copying-task.py | 13 +- examples/drtrl/02-batching-vmap.py | 8 +- examples/pp_prop/05-batching-vmap.py | 16 +- .../pp_prop/12-classification-neuromorphic.py | 13 +- .../pp_prop/14-knob-vjp-method-contrast.py | 13 +- examples/pp_prop/README.md | 2 +- examples/pp_prop/_shared.py | 14 +- examples/snn_models.py | 31 ++- examples/tests/test_compile_modes.py | 18 +- 38 files changed, 569 insertions(+), 643 deletions(-) delete mode 100644 docs/specs/2026-07-28-warnings-and-map-initialization.md diff --git a/braintrace/_algorithm/base.py b/braintrace/_algorithm/base.py index ebd7dafb..b403164e 100644 --- a/braintrace/_algorithm/base.py +++ b/braintrace/_algorithm/base.py @@ -279,10 +279,14 @@ def compile_graph(self, *args: Any) -> None: The input arguments. """ - # Legacy mapped-state transforms run an eager discovery probe before - # the real mapped pass. Compiling against the throwaway probe states - # would bind the executor to states that are discarded immediately, so - # defer compilation until the real mapped pass. + # ``vmap_new_states`` / ``vmap2_new_states`` run an eager *discovery + # probe* that executes the surrounding ``init`` (including this call) + # once against throwaway, un-batched states before the real mapped + # pass. Compiling there would bind the executor to those probe states + # (which are discarded and left untagged), so the subsequent + # ``brainstate.nn.Vmap(..., vmap_states='new')`` would not cover them + # and writing a batched value raises ``BatchAxisError``. Defer to the + # real mapped pass, which creates the 'new'-tagged batched states. _in_probe = getattr(brainstate.transform, 'in_new_state_probe', None) if _in_probe is not None and _in_probe(): return diff --git a/braintrace/_algorithm/sequence.py b/braintrace/_algorithm/sequence.py index 174f9dc5..1b1d2fbc 100644 --- a/braintrace/_algorithm/sequence.py +++ b/braintrace/_algorithm/sequence.py @@ -168,10 +168,9 @@ class SequenceDriverMixin: bypass the window-mode validation entirely. """ - #: Set by :class:`ETraceVmap`. Window mode is refused for that compatibility - #: wrapper because its ``in_axes=0`` would map a window's time axis. - #: ``compile(vmap=True)`` now uses ``brainstate.nn.Map``, keeps the - #: batch axis inside the graph, and therefore leaves this false. + #: Set by :class:`ETraceVmap`. Window mode is refused when true, because + #: ``compile(vmap=True)`` maps ``in_axes=0`` and a ``(k, B, ...)`` window + #: slice would map *time* as the batch axis. _seq_is_vmapped: bool = False @property @@ -197,12 +196,12 @@ def _seq_check_window(self, chunk_size: Optional[int], length: int, if self._seq_is_vmapped: raise ValueError( f'chunk_size={chunk_size} (window mode) is not supported under ' - f'an ETraceVmap learner. Its in_axes=0 mapping would treat the ' - f'*time* axis of a (chunk_size, batch, ...) window slice as the ' - f'batch axis -- which is silently wrong whenever ' - f'chunk_size equals the batch size. Use compile(vmap=True), ' - f'which carries the batch axis inside a brainstate.nn.Map, or ' - f'drive this compatibility wrapper with chunk_size=None.' + f'a vmapped learner. compile(vmap=True) maps in_axes=0, so a ' + f'(chunk_size, batch, ...) window slice would map the *time* ' + f'axis as the batch axis -- which is silently wrong whenever ' + f'chunk_size equals the batch size. Use the batched (non-vmap) ' + f'mode, which carries the batch axis inside the compiled graph, ' + f'or drive with chunk_size=None.' ) if for_grad and self._seq_vjp_method != 'multi-step': @@ -267,7 +266,7 @@ def etrace_grad( ``step_fn`` must return a ``(k,)`` vector of per-step losses and wrap its model inputs in :class:`MultiStepData`. Window mode requires ``vjp_method='multi-step'`` and ``T % k == 0``, and is not - available under an :class:`ETraceVmap` compatibility wrapper. + available under a vmapped learner. weights : dict, optional The :class:`brainstate.ParamState` to differentiate. Defaults to the learner's own ``param_states``. @@ -303,9 +302,8 @@ def etrace_grad( ``T == 0`` -- there is nothing to slice. If *chunk_size* is below ``1``; if ``k >= 2`` but the learner's ``vjp_method`` is not ``'multi-step'`` (the executor would raise three frames down), the - learner is an :class:`ETraceVmap` (``in_axes=0`` would map time as - the batch axis), or ``T % k != 0``. If *mask* is not shape - ``(T,)``. If + learner is vmapped (``in_axes=0`` would map time as the batch + axis), or ``T % k != 0``. If *mask* is not shape ``(T,)``. If *reduction* or *loss_output* is not one of its legal values. If ``step_fn`` returns a non-scalar in plain mode, or anything but shape ``(k,)`` in window mode. If the learner has not been @@ -500,8 +498,7 @@ def etrace_evolve( As in :meth:`etrace_grad`, **except** that a window is *not* refused for being on a single-step learner -- no loss VJP is taken here, so the restriction does not apply. Windows are still refused - under an :class:`ETraceVmap` compatibility wrapper, and - ``T % chunk_size == 0`` still holds. + under a vmapped learner, and ``T % chunk_size == 0`` still holds. Examples -------- @@ -540,12 +537,11 @@ def body(*slices): class ETraceVmap(SequenceDriverMixin, brainstate.nn.Vmap): - """Provide sequence drivers for explicitly constructed legacy Vmap models. + """Provide sequence drivers on a ``brainstate.nn.Vmap`` wrapper. - ``braintrace.compile(..., vmap=True)`` uses ``brainstate.nn.Map`` and - returns the algorithm directly. This public compatibility type remains for - callers that explicitly construct a ``brainstate.nn.Vmap`` learner. - Because it *is* a ``brainstate.nn.Vmap``, calling it and every + Returned by ``braintrace.compile(..., vmap=True)`` so the call site is + identical in batched and unbatched mode. Because it *is* a + ``brainstate.nn.Vmap``, calling it, its attributes and every ``isinstance(x, brainstate.nn.Vmap)`` check keep working; only the added methods are new. @@ -559,7 +555,7 @@ class ETraceVmap(SequenceDriverMixin, brainstate.nn.Vmap): would drive the **unbatched** learner and silently produce per-lane-wrong results. - Window mode is refused here; see + Window mode is refused here -- see :meth:`SequenceDriverMixin.etrace_grad`. """ __module__ = 'braintrace' diff --git a/braintrace/_algorithm/sequence_test.py b/braintrace/_algorithm/sequence_test.py index ffdb4af6..aa52b743 100644 --- a/braintrace/_algorithm/sequence_test.py +++ b/braintrace/_algorithm/sequence_test.py @@ -34,7 +34,6 @@ import inspect import brainstate -import braintools import brainunit as u import jax import jax.numpy as jnp @@ -610,18 +609,20 @@ def test_a_weighted_mask_reweights_the_objective(self): # --------------------------------------------------------------------------- -# 18--20. Map +# 18--20. vmap # --------------------------------------------------------------------------- class _VmapNet(brainstate.nn.Module): """A model whose hidden state is created in ``init_state``, not ``__init__``. - ``ValinaRNNCell`` defers its state to ``init_state``, allowing - ``brainstate.nn.Map.init_all_states`` to create one independent state per - mapped lane. + ``om.tanh_rnn`` allocates its ``HiddenState`` in ``__init__``, so + ``vmap_new_states`` has nothing to batch and ``compile(vmap=True)`` raises + ``BatchAxisError`` on the first call -- the whole vmap section was + previously written against a fixture that could not run. ``ValinaRNNCell`` + defers its state to ``init_state``, which is the property that matters. ``wout`` is a plain (non-ETP) parameter, so it is exactly zero under - ``vjp_method='single-step'`` (F-33); the Map fixture therefore runs + ``vjp_method='single-step'`` (F-33); the vmap fixture therefore runs ``'multi-step'``, which keeps every key live and the comparisons honest. """ @@ -636,7 +637,7 @@ def update(self, x): return self.cell(x) @ self.wout.value -def _map_learner(batch, vjp_method='multi-step', **opts): +def _vmap_learner(batch, vjp_method='multi-step', **opts): return braintrace.compile(_VmapNet(), 'D_RTRL', jnp.zeros((batch, N_IN)), batch_size=batch, vmap=True, vjp_method=vjp_method, **opts) @@ -661,46 +662,20 @@ def _lane_data(batch, *, seed=7): return xs, ys -class TestMap: - def test_the_mapped_learner_carries_the_driver_methods(self): +class TestVmap: + def test_the_vmapped_learner_carries_the_driver_methods(self): """Spec test 18. - ``compile(vmap=True)`` returns the algorithm itself with a mapped model, - so sequence drivers and compilation reports remain directly available. + ``compile(vmap=True)`` must return something that *has* ``etrace_grad``; + before this change it returned a bare ``brainstate.nn.Vmap``, which does + not. Reaching into ``.module`` instead would drive the unbatched learner + and silently give per-lane-wrong results. """ - learner = _map_learner(3) - assert isinstance(learner, braintrace.ETraceAlgorithm) - assert isinstance(learner.graph_executor.model, brainstate.nn.Map) + learner = _vmap_learner(3) assert hasattr(learner, 'etrace_grad') assert hasattr(learner, 'etrace_evolve') - def test_default_gradients_keep_the_original_model_parameter_paths(self): - """Map internals must not leak into the public optimizer contract.""" - batch = 3 - xs, ys = _lane_data(batch) - model = _VmapNet() - learner = braintrace.compile( - model, - 'D_RTRL', - jnp.zeros((batch, N_IN)), - batch_size=batch, - vmap=True, - vjp_method='multi-step', - ) - optimizer = braintools.optim.Adam(lr=1e-3) - optimizer.register_trainable_weights(learner.param_states) - - grads = learner.etrace_grad( - xs, - ys, - step_fn=lambda inp, tar: jnp.sum((learner(inp) - tar) ** 2), - ) - - assert set(learner.param_states) == set(model.states(brainstate.ParamState)) - assert set(grads) == set(learner.param_states) - optimizer.update(grads) - - def test_the_mapped_gradient_is_the_sum_over_independent_lanes(self): + def test_the_vmapped_gradient_is_the_sum_over_independent_lanes(self): """Spec test 18, the part that has content. The parameters are shared across lanes, so the batched gradient must be @@ -713,7 +688,7 @@ def test_the_mapped_gradient_is_the_sum_over_independent_lanes(self): batch = 3 xs, ys = _lane_data(batch) - batched = _map_learner(batch) + batched = _vmap_learner(batch) def step_fn(inp, tar): return jnp.sum((batched(inp) - tar) ** 2) @@ -733,16 +708,13 @@ def lane_step(inp, tar, learner=lane): lambda a, b: a + b, lane_total, g_lane) flat = _arrays(g_batched) - lane_flat = _arrays(lane_total) - assert set(flat) == set(lane_flat), 'gradient keys diverged' + assert set(flat) == set(_arrays(lane_total)), 'gradient keys diverged' for k, v in flat.items(): assert np.max(np.abs(v)) > 0.0, f'{k} is identically zero -- vacuous' - np.testing.assert_allclose( - v, lane_flat[k], rtol=2e-6, atol=1e-6, - err_msg=f'mapped vs sum over independent lanes: at {k}', - ) + _assert_trees_equal(g_batched, lane_total, rtol=2e-6, atol=1e-6, + msg='batched vs sum over independent lanes') - def test_permuting_the_lanes_changes_the_mapped_gradient(self): + def test_permuting_the_lanes_changes_the_vmapped_gradient(self): """Spec test 18, the negative control. Without this, the sum-over-lanes identity could hold for a driver that @@ -753,11 +725,11 @@ def test_permuting_the_lanes_changes_the_mapped_gradient(self): xs, ys = _lane_data(batch) perm = jnp.asarray([1, 0, 2]) - straight = _map_learner(batch) + straight = _vmap_learner(batch) g_straight = straight.etrace_grad( xs, ys, step_fn=lambda i, t: jnp.sum((straight(i) - t) ** 2)) - swapped = _map_learner(batch) + swapped = _vmap_learner(batch) g_swapped = swapped.etrace_grad( xs, ys[:, perm], step_fn=lambda i, t: jnp.sum((swapped(i) - t) ** 2)) @@ -765,110 +737,50 @@ def test_permuting_the_lanes_changes_the_mapped_gradient(self): oracle.assert_gradients_differ(_arrays(g_straight), _arrays(g_swapped), min_rel=1e-3) - @pytest.mark.parametrize('batch', [3, K]) - def test_window_mode_matches_independent_lane_windows_under_map(self, batch): + @pytest.mark.parametrize('batch', [3, K]) # B != k, and the silent B == k case + def test_window_mode_is_refused_under_vmap(self, batch): """Spec test 19. - Map keeps the batch axis inside the compiled graph, so a - ``(K, batch, ...)`` window must equal the sum of independently driven - ``(K, 1, ...)`` lanes. Comparing with the plain path is not a valid - oracle: a multi-step update and K single-step updates are different - approximation regimes for this recurrent fixture. + ``compile(vmap=True)`` maps ``in_axes=0``, so a ``(k, B, ...)`` window + slice would map *time* as the batch axis. At ``B != k`` that is a loud + shape error; at ``B == k`` the shapes line up and it would train on + transposed data. The ``B == k`` parametrization is the one that matters. """ - xs, ys = _lane_data(batch) - mapped = _map_learner(batch) - - def mapped_window_loss(inp, tar): - out = mapped(braintrace.MultiStepData(inp)) - return jnp.sum((out - tar) ** 2, axis=(1, 2)) - - g_mapped = mapped.etrace_grad( - xs, - ys, - step_fn=mapped_window_loss, - chunk_size=K, - reduction='sum', - ) + learner = _vmap_learner(batch) + xs = jnp.zeros((T, batch, N_IN)) - lane_total = None - for lane_index in range(batch): - lane = _lane_learner() - - def lane_window_loss(inp, tar, learner=lane): - out = learner(braintrace.MultiStepData(inp)) - return jnp.sum((out - tar) ** 2, axis=(1, 2)) - - g_lane = lane.etrace_grad( - xs[:, lane_index:lane_index + 1], - ys[:, lane_index:lane_index + 1], - step_fn=lane_window_loss, - chunk_size=K, - reduction='sum', - ) - lane_total = g_lane if lane_total is None else jax.tree.map( - lambda a, b: a + b, lane_total, g_lane) - - mapped_arrays = _arrays(g_mapped) - lane_arrays = _arrays(lane_total) - assert set(mapped_arrays) == set(lane_arrays), 'gradient keys diverged' - for key, mapped_value in mapped_arrays.items(): - np.testing.assert_allclose( - mapped_value, - lane_arrays[key], - rtol=2e-6, - atol=1e-6, - err_msg=f'Map window vs independent lane windows: at {key}', - ) - - @pytest.mark.parametrize('batch', [3, K]) - def test_window_evolve_matches_independent_lanes_under_map(self, batch): - """Window evolution preserves separate time and mapped-lane axes.""" - xs, _ = _lane_data(batch) - mapped = _map_learner(batch) - mapped_outputs = mapped.etrace_evolve( - xs, chunk_size=K, return_outputs=True) - - lane_outputs = [] - for lane_index in range(batch): - lane = _lane_learner() - lane_outputs.append( - lane.etrace_evolve( - xs[:, lane_index:lane_index + 1], - chunk_size=K, - return_outputs=True, - ) - ) - expected = jnp.concatenate(lane_outputs, axis=2) - np.testing.assert_allclose( - np.asarray(mapped_outputs), - np.asarray(expected), - rtol=2e-6, - atol=1e-6, - ) + with pytest.raises(ValueError, match='vmap'): + learner.etrace_grad(xs, step_fn=lambda x: jnp.zeros(K), + chunk_size=K) + with pytest.raises(ValueError, match='vmap'): + learner.etrace_evolve(xs, chunk_size=K) @pytest.mark.parametrize('batch', [3, K]) - def test_chunk_size_one_is_admitted_under_map(self, batch): - """Spec test 19, the plain path remains available under Map. + def test_chunk_size_one_is_admitted_under_vmap(self, batch): + """Spec test 19, the other half -- the refusal must not overreach. - ``chunk_size=1`` is normalized to the ordinary step-by-step path. + ``chunk_size=1`` is the plain path, so it carries none of the axis + collision that makes ``k >= 2`` unsafe. A guard written as + ``if chunk_size is not None`` would refuse it, which is why both + methods are exercised rather than just ``etrace_evolve``. """ xs, ys = _lane_data(batch) - learner = _map_learner(batch) + learner = _vmap_learner(batch) learner.etrace_evolve(xs, chunk_size=1) - grad_learner = _map_learner(batch) + grad_learner = _vmap_learner(batch) grads = grad_learner.etrace_grad( xs, ys, chunk_size=1, step_fn=lambda i, t: jnp.sum((grad_learner(i) - t) ** 2)) for k, v in _arrays(grads).items(): assert np.all(np.isfinite(v)), f'{k} is not finite' - def test_vmap_option_returns_an_algorithm_with_a_mapped_model(self): - """Spec test 20 -- the option selects Map-based state initialization.""" - learner = _map_learner(3) - assert isinstance(learner, braintrace.ETraceAlgorithm) - assert not isinstance(learner, brainstate.nn.Vmap) - assert isinstance(learner.graph_executor.model, brainstate.nn.Map) + def test_the_vmap_return_value_is_still_a_brainstate_vmap(self): + """Spec test 20 -- existing ``vmap=True`` users must be unaffected.""" + learner = _vmap_learner(3) + assert isinstance(learner, brainstate.nn.Vmap) + assert isinstance(learner, braintrace.ETraceVmap) + assert isinstance(learner.module, braintrace.ETraceAlgorithm) # --------------------------------------------------------------------------- diff --git a/braintrace/_algorithm/tests/conv_vmap_correctness_test.py b/braintrace/_algorithm/tests/conv_vmap_correctness_test.py index 494e05e7..c2f318f0 100644 --- a/braintrace/_algorithm/tests/conv_vmap_correctness_test.py +++ b/braintrace/_algorithm/tests/conv_vmap_correctness_test.py @@ -13,17 +13,17 @@ # limitations under the License. # ============================================================================== -"""Conv and mixed ETP correctness under ``brainstate.nn.Map``. +"""Conv / mixed ETP under ``brainstate.nn.Vmap(vmap_states='new')`` correctness. Regression coverage for the eligibility-trace path through the *batched* online -executor backed by ``brainstate.nn.Map`` (the mapped batching flow used +executor wrapped by ``brainstate.nn.Vmap`` (the ``OnlineVmapTrainer`` flow used by ``examples/004``). This path was previously uncovered — conv was exercised only at the rule level (``_op/conv_test.py``) and the "conv" model in ``transform_correctness_test`` is actually a matmul — which let two regressions through: 1. *Pure conv.* A conv forward forces a leading batch axis on its input, but - under mapped state initialization the hidden-state traces are per-lane and carry no + under ``vmap_states='new'`` the hidden-state traces are per-lane and carry no batch axis, so the instantaneous, recurrent and solve terms saw a singleton batch on the input but none on the cotangent. @@ -36,13 +36,13 @@ 3. *Norm in the transition.* ``conv -> LayerNorm -> IF`` makes ``dh/dy`` non-diagonal; the all-ones jvp returns its row sums, exactly zero for the shift-invariant norm. A ``use_fast_variance=True`` norm leaves a float32 - residual instead, which under mapped execution the recurrent trace and + residual instead, which under ``vmap_states='new'`` the recurrent trace and ``rsqrt(var+eps)`` amplify into an overflow that diverges from the eager reference — the ``examples/004`` ``loss=ln(10)`` stall. -**Oracle (exact, transform-invariance).** ``brainstate.nn.Map`` is a transform; +**Oracle (exact, transform-invariance).** ``brainstate.nn.Vmap`` is a transform; for parameters shared across lanes its grad sums the per-lane gradients. So the -gradient from the mapped path on a batch of ``B`` samples +gradient from the ``vmap_new_states`` + ``Vmap`` path on a batch of ``B`` samples must equal the sum over ``b`` of the *eager, batch=1* gradient on sample ``b``. The eager batch=1 path is independently healthy for conv (states are initialised *with* a size-1 batch, so input and trace batch axes agree), which makes it a @@ -59,6 +59,17 @@ import braintools import brainpy.state +# `etp_conv` has no registered batched counterpart, so every model here that +# routes a sample through `braintrace.nn.Conv2d` under `brainstate.nn.Vmap` +# hits the identity-preserving batching rule's decomposition fallback in +# `braintrace/_op/_primitive.py`, which emits a `UserWarning`. That warning is +# expected-but-not-under-test in this module (the module tests gradient +# correctness, not the vmap-decomposition warning itself — that is covered by +# `braintrace/_op/_primitive_test.py`), so it is filtered narrowly by message. +pytestmark = pytest.mark.filterwarnings( + "ignore:ETP primitive 'etp_conv' was decomposed:UserWarning" +) + H = W = 6 C_IN = 2 C_OUT = 3 @@ -143,19 +154,24 @@ def loss_fn(x): return grads -def _map_grad(data, targets, make_net): - """Initialize mapped states explicitly and accumulate mapped gradients.""" +def _vmap_grad(data, targets, make_net): + """The ``OnlineVmapTrainer`` flow: vmap_new_states init + Vmap(vmap_states='new').""" net = make_net() - mapped_net = brainstate.nn.Map(net, init_map_size=data.shape[1]) - mapped_net.init_all_states() - model = braintrace.D_RTRL(mapped_net) - with brainstate.environ.context(fit=True): - model.compile_graph(data[0]) + model = braintrace.D_RTRL(net) + + @brainstate.transform.vmap_new_states(state_tag='new', axis_size=data.shape[1]) + def init(): + brainstate.nn.init_all_states(net) + with brainstate.environ.context(fit=True): + model.compile_graph(data[0, 0]) + + init() + vmodel = brainstate.nn.Vmap(model, vmap_states='new') weights = net.states().subset(brainstate.ParamState) def _grad(inp): with brainstate.environ.context(fit=True): - return _loss(model(inp), targets) + return _loss(vmodel(inp), targets) def _step(prev, x): g = brainstate.transform.grad(_grad, weights)(x) @@ -173,7 +189,7 @@ def _step(prev, x): def _make_mixed_net(): """conv -> IF -> flatten -> Linear -> IF: a *mixed* batched/unbatched model. - Under ``brainstate.nn.Map`` the graph is compiled across mapped lanes, so the conv stays + Under ``vmap_states='new'`` the graph is compiled per-lane, so the conv stays a *batched* primitive (its parent layer forces a leading batch axis) while the flattened ``Linear`` input is 1-D and dispatches to the *unbatched* ``etp_mv``. The solve's trailing batch-sum must collapse only the conv gradient's batch @@ -216,7 +232,7 @@ def _make_conv_ln_net(use_fast_variance): upstream conv gets no eligibility gradient through the norm — a documented approximation, matching the eager path). That exactness is numerical: with ``use_fast_variance=True`` the ``E[x^2]-E[x]^2`` variance leaves a float32 - residual instead of zero, and under mapped execution the recurrent + residual instead of zero, and under ``Vmap(vmap_states='new')`` the recurrent trace and the large ``rsqrt(var+eps)`` factor amplify it into an overflow that diverges from the eager reference (the ``examples/004`` ``loss=ln(10)`` stall). """ @@ -256,8 +272,8 @@ def _assert_grads_match(ref, got): @pytest.mark.parametrize('neuron', ['IF', 'ALIF'], ids=['num_state1_IF', 'num_state2_ALIF']) -def test_conv_map_grad_equals_sum_of_eager_single_sample(neuron): - """Mapped conv D-RTRL gradient equals the sum of eager sample gradients.""" +def test_conv_vmap_grad_equals_sum_of_eager_single_sample(neuron): + """vmap_new_states+Vmap conv D_RTRL grad == sum over samples of eager batch=1 grad.""" rng = np.random.RandomState(42) data = jnp.asarray(rng.rand(N_STEP, B, H, W, C_IN).astype('float32')) targets = jnp.asarray(rng.rand(B, H, W, C_OUT).astype('float32')) @@ -268,11 +284,11 @@ def test_conv_map_grad_equals_sum_of_eager_single_sample(neuron): g = _eager_grad_one(data[:, b], targets[b], make_net) ref = g if ref is None else jax.tree.map(lambda a, c: a + c, ref, g) - got = _map_grad(data, targets, make_net) + got = _vmap_grad(data, targets, make_net) _assert_grads_match(ref, got) -def test_mixed_conv_dense_map_grad_equals_sum_of_eager_single_sample(): +def test_mixed_conv_dense_vmap_grad_equals_sum_of_eager_single_sample(): """Mixed batched(conv)+unbatched(dense-mv) model: vmap grad == sum of eager batch=1. Regression for the ``examples/004`` layer4 failure — the unbatched ``etp_mv`` @@ -288,11 +304,11 @@ def test_mixed_conv_dense_map_grad_equals_sum_of_eager_single_sample(): g = _eager_grad_one(data[:, b], targets[b], _make_mixed_net) ref = g if ref is None else jax.tree.map(lambda a, c: a + c, ref, g) - got = _map_grad(data, targets, _make_mixed_net) + got = _vmap_grad(data, targets, _make_mixed_net) _assert_grads_match(ref, got) -def test_conv_layernorm_map_grad_matches_eager_and_stays_finite(): +def test_conv_layernorm_vmap_grad_matches_eager_and_stays_finite(): """conv -> LayerNorm -> IF: vmap grad == sum of eager batch=1, and stays finite. Regression for the ``examples/004`` ``loss=ln(10)`` stall. A mean-subtracting @@ -300,7 +316,7 @@ def test_conv_layernorm_map_grad_matches_eager_and_stays_finite(): its row sums, which for shift-invariance are exactly zero, so the conv weight gets no eligibility gradient through the norm (both paths agree on ~0). With a numerically stable variance (``use_fast_variance=False``) that exact zero holds - under mapped execution; the transform-invariance oracle then makes + under ``Vmap(vmap_states='new')``; the transform-invariance oracle then makes vmap == sum-of-eager, and neither explodes. """ rng = np.random.RandomState(42) @@ -313,7 +329,7 @@ def test_conv_layernorm_map_grad_matches_eager_and_stays_finite(): g = _eager_grad_one(data[:, b], targets[b], make_net) ref = g if ref is None else jax.tree.map(lambda a, c: a + c, ref, g) - got = _map_grad(data, targets, make_net) + got = _vmap_grad(data, targets, make_net) # No overflow/NaN in either path (the bug produced ~1e14 -> NaN under vmap). for leaf in jax.tree.leaves(got): assert np.all(np.isfinite(np.asarray(leaf))), 'vmap grad is non-finite' diff --git a/braintrace/_algorithm/tests/diagnostic_exploration_test.py b/braintrace/_algorithm/tests/diagnostic_exploration_test.py index 204f68bc..085e2c4c 100644 --- a/braintrace/_algorithm/tests/diagnostic_exploration_test.py +++ b/braintrace/_algorithm/tests/diagnostic_exploration_test.py @@ -36,6 +36,7 @@ """ import importlib.util +import warnings import pytest @@ -81,8 +82,10 @@ def _drtrl(model): def _assert_exact_equals_bptt(spec, inputs): """D_RTRL multi-step gradient == BPTT gradient for every ParamState.""" - expected = bptt_param_gradients(spec.factory, inputs) - actual = online_param_gradients(spec.factory, inputs, algo_factory=_drtrl) + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + expected = bptt_param_gradients(spec.factory, inputs) + actual = online_param_gradients(spec.factory, inputs, algo_factory=_drtrl) assert_param_gradients_close(actual, expected, atol=ATOL) @@ -160,10 +163,12 @@ def test_batch_invariance_over_dims(n_in, n_rec, batch, seq_len, seed): summed per-step SSE loss over the batch axis.""" seq = jnp.asarray( np.random.RandomState(seed).randn(seq_len, batch, n_in).astype('float32')) - batched = _batched_multistep_grad(n_in, n_rec, batch, seq, seed) - summed = None - for b in range(batch): - sub = seq[:, b:b + 1, :] - g = _batched_multistep_grad(n_in, n_rec, 1, sub, seed) - summed = g if summed is None else {k: summed[k] + g[k] for k in g} + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + batched = _batched_multistep_grad(n_in, n_rec, batch, seq, seed) + summed = None + for b in range(batch): + sub = seq[:, b:b + 1, :] + g = _batched_multistep_grad(n_in, n_rec, 1, sub, seed) + summed = g if summed is None else {k: summed[k] + g[k] for k in g} assert_param_gradients_close(batched, summed, atol=ATOL) diff --git a/braintrace/_algorithm/tests/while_support_test.py b/braintrace/_algorithm/tests/while_support_test.py index f7fb998a..1ae022c3 100644 --- a/braintrace/_algorithm/tests/while_support_test.py +++ b/braintrace/_algorithm/tests/while_support_test.py @@ -49,6 +49,8 @@ WARNING-level ``CONTROL_FLOW_OPAQUE_FWD`` diagnostic for each detach. """ +import warnings + import brainstate import jax import jax.numpy as jnp @@ -237,11 +239,13 @@ def test_upstream_layer_gradient_is_zero_behind_while_DOCUMENTED_LIMITATION(): inputs = _inputs(6, 3) def grads(while_layer): - return online_param_gradients_singlestep_naive( - lambda: _StackedWhileNet(while_layer=while_layer), - inputs, - algo_factory=braintrace.D_RTRL, - ) + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + return online_param_gradients_singlestep_naive( + lambda: _StackedWhileNet(while_layer=while_layer), + inputs, + algo_factory=braintrace.D_RTRL, + ) g_while = grads(True) g_twin = grads(False) diff --git a/braintrace/_compile.py b/braintrace/_compile.py index 94f3556f..f214c89a 100644 --- a/braintrace/_compile.py +++ b/braintrace/_compile.py @@ -17,12 +17,14 @@ from typing import Any, Type, Union +import jax import brainstate from ._misc import CompilationError from ._algorithm import ( ETraceAlgorithm, ETraceConfig, + ETraceVmap, IODimVjpAlgorithm, ParamDimVjpAlgorithm, RandomProjectionVjpAlgorithm, @@ -126,7 +128,7 @@ def compile( verbose: int = 0, vmap: bool = False, **options: Any, -) -> ETraceAlgorithm: +) -> ETraceAlgorithm | brainstate.nn.Vmap: """Define an eligibility-trace online-learning model in one call. This is the unified entry point. It initializes the model's states, builds @@ -172,22 +174,29 @@ def compile( vmap : bool, optional When ``False`` (default) states are initialized with ``init_all_states(model, batch_size=batch_size)``. When ``True``, states - initialized by wrapping the model in - ``brainstate.nn.Map(model, init_map_size=batch_size)`` and calling the - mapped model's ``init_all_states()`` method. In vmap mode, - ``example_inputs`` carry the batch axis (axis 0), ``batch_size`` is - **required** and sets the map size, and the returned learner exposes - ``report``, ``etrace_grad``, and ``etrace_evolve`` directly. + are created under + ``brainstate.transform.vmap_new_states(state_tag='new', axis_size=batch_size)`` + and the learner is wrapped in :class:`ETraceVmap`. In vmap mode: + ``example_inputs`` carry the batch axis (axis 0); ``batch_size`` is + **required** and used as the vmap ``axis_size``; the return value is a + :class:`ETraceVmap` whose ``.module`` is the unbatched learner (use + ``result.module.report`` for its report). Drive sequences through the + returned wrapper, never through ``result.module``. Requires a model + whose hidden states are all (re)created in ``init_all_states``; models + holding construction-time states may raise + ``brainstate.transform.BatchAxisError``. **options : Any Forwarded to the algorithm constructor. See *Algorithm options* below. Returns ------- - ETraceAlgorithm - The compiled learner, carrying a :attr:`~ETraceAlgorithm.report`. Call - ``.update(*inputs)`` for one step, or use ``etrace_grad`` and - ``etrace_evolve`` for sequences. This return contract is identical in - mapped and directly batched modes. + ETraceAlgorithm or ETraceVmap + When ``vmap=False``, the compiled learner carries a + :attr:`~ETraceAlgorithm.report`; call ``.update(*inputs)`` to train. + When ``vmap=True``, returns an :class:`ETraceVmap` wrapper (also a + ``brainstate.nn.Vmap``); access the underlying learner's report as + ``.module.report``. Call ``etrace_grad`` and ``etrace_evolve`` on the + wrapper itself, not on ``.module``. Raises ------ @@ -306,22 +315,34 @@ def compile( raise ValueError(f'verbose must be 0, 1, or 2, got {verbose!r}.') if vmap and batch_size is None: raise ValueError( - 'compile(..., vmap=True) requires batch_size, used as the ' - 'brainstate.nn.Map size. Pass batch_size= matching axis 0 ' - 'of example_inputs.' + 'compile(..., vmap=True) requires batch_size, used as the per-sample ' + 'vmap axis size. Pass batch_size= matching the batch axis ' + '(axis 0) of example_inputs.' ) if vmap: - # Per-sample map scheme: example_inputs carry the batch axis (axis 0). - model = brainstate.nn.Map(model, init_map_size=batch_size) + # Per-sample vmap scheme: example_inputs carry the batch axis (axis 0); + # the eligibility-trace graph is built per-lane on an unbatched sample, + # while hidden + trace states are created with the new per-sample axis. + learner = cls(model, **options) + unbatched = jax.tree.map(lambda a: a[0], example_inputs) + + @brainstate.transform.vmap_new_states(state_tag='new', axis_size=batch_size) + def _init() -> None: + brainstate.nn.init_all_states(model) + learner.compile_graph(*unbatched) + if seed is not None: with brainstate.random.seed_context(seed): - model.init_all_states() + _init() else: - model.init_all_states() - learner = cls(model, **options) - learner.compile_graph(*example_inputs) - result = learner + _init() + # ETraceVmap, not brainstate.nn.Vmap: the wrapper must carry + # etrace_grad / etrace_evolve so the call site is identical in batched + # and unbatched mode. Reaching into `.module` instead would drive the + # *unbatched* learner and silently give per-lane-wrong results. It is + # still a brainstate.nn.Vmap, so existing users are unaffected. + result: ETraceAlgorithm | brainstate.nn.Vmap = ETraceVmap(learner, vmap_states='new') else: # --- state initialization (always) --- # if seed is not None: diff --git a/braintrace/_compile_test.py b/braintrace/_compile_test.py index 9dfd6080..14ae3bbc 100644 --- a/braintrace/_compile_test.py +++ b/braintrace/_compile_test.py @@ -259,10 +259,7 @@ def test_compile_vmap_builds_forwards_and_grads(): B = 4 xb = jnp.ones((B, 3), dtype='float32') learner = braintrace.compile(model, 'D_RTRL', xb, batch_size=B, vmap=True) - assert isinstance(learner, braintrace.D_RTRL) - assert isinstance(learner.graph_executor.model, brainstate.nn.Map) - assert learner.graph_executor.model.init_map_size == B - assert learner.graph_executor.model._init + assert isinstance(learner, brainstate.nn.Vmap) out = learner(xb) assert out.shape[0] == B @@ -282,20 +279,20 @@ def test_compile_vmap_requires_batch_size(): assert 'batch_size' in str(exc.value) -def test_compile_vmap_returns_algorithm_exposing_report(): +def test_compile_vmap_returns_wrapper_exposing_report(): model = _VmapRNN() B = 4 xb = jnp.ones((B, 3), dtype='float32') learner = braintrace.compile(model, 'D_RTRL', xb, batch_size=B, vmap=True) - assert isinstance(learner, braintrace.D_RTRL) - assert learner.report is not None - assert learner.is_compiled + assert isinstance(learner.module, braintrace.D_RTRL) + assert learner.module.report is not None + assert learner.module.is_compiled # --- both-modes coverage across RNN architectures + algorithms --------------- # Each architecture/algorithm must build, forward, and back-prop a finite, # non-zero gradient under BOTH compile(vmap=False) (internal batch primitive) -# and compile(vmap=True) (Map-owned per-sample states). A multi-step scan exercises +# and compile(vmap=True) (per-sample vmap lanes). A multi-step scan exercises # the eligibility trace (single-step would never engage it). _NI, _NR = 3, 4 @@ -405,14 +402,24 @@ def update(self, x): @pytest.mark.parametrize('name,builder,algo,kw,feat', _BOTH_MODE_CASES, ids=[c[0] for c in _BOTH_MODE_CASES]) @pytest.mark.parametrize('vmap', [False, True], ids=['no_vmap', 'vmap']) +# `conv1d_minigru_d_rtrl` (etp_conv) has no registered batched counterpart, +# so under `vmap=True` compilation it hits the identity-preserving batching +# rule's decomposition fallback and warns (see `braintrace/_op/_primitive.py`). +# This test asserts gradient finiteness/non-zero-ness, not the +# vmap-decomposition warning (covered by +# `braintrace/_op/_primitive_test.py`), so the expected warning is filtered +# narrowly by message rather than left uncaptured. `lora_d_rtrl` (etp_lora_mv) +# now has a registered batched counterpart (`etp_lora_mm`) and is promoted +# instead of decomposed, so no filter is needed for it. +@pytest.mark.filterwarnings( + "ignore:ETP primitive 'etp_conv' was decomposed:UserWarning") def test_compile_both_modes_finite_nonzero_grad(name, builder, algo, kw, feat, vmap): B, T = 4, 5 xs = brainstate.random.randn(T, B, *feat) model = builder() learner = braintrace.compile(model, algo, xs[0], batch_size=B, vmap=vmap, **kw) if vmap: - assert isinstance(learner, braintrace.ETraceAlgorithm) - assert isinstance(learner.graph_executor.model, brainstate.nn.Map) + assert isinstance(learner, brainstate.nn.Vmap) weights = model.states(brainstate.ParamState) def total_loss(xs): diff --git a/braintrace/_compiler/canonicalize_test.py b/braintrace/_compiler/canonicalize_test.py index 70eda7c9..a08b3697 100644 --- a/braintrace/_compiler/canonicalize_test.py +++ b/braintrace/_compiler/canonicalize_test.py @@ -13,6 +13,7 @@ # limitations under the License. # ============================================================================== +import warnings import brainstate import jax @@ -807,20 +808,21 @@ def test_skip_length_exceeds_limit(self): kinds = [r.kind for r in reporter.records()] assert kinds.count(DiagnosticKind.SCAN_UNROLL_SKIPPED) == 1 - def test_skip_length_exceeds_limit_info_under_descent_auto(self, recwarn): + def test_skip_length_exceeds_limit_info_under_descent_auto(self): # Phase 4: with scan_descent='auto' an over-limit scan is no longer a # dead end, so the skip record downgrades to INFO (no UserWarning) # and points at the descent path. f, closed, w, h0, xs = self._etp_scan_jaxpr() with diagnostic_context() as reporter: - conv = _unroll( - closed, - weights=[closed.jaxpr.invars[0]], - policy=ControlFlowPolicy(scan_unroll_limit=self.L - 1, - scan_descent='auto'), - ) + with warnings.catch_warnings(): + warnings.simplefilter('error') + conv = _unroll( + closed, + weights=[closed.jaxpr.invars[0]], + policy=ControlFlowPolicy(scan_unroll_limit=self.L - 1, + scan_descent='auto'), + ) assert 'scan' in _primitive_names(conv) - assert not any(issubclass(w.category, UserWarning) for w in recwarn) recs = [r for r in reporter.records() if r.kind is DiagnosticKind.SCAN_UNROLL_SKIPPED] assert len(recs) == 1 @@ -1227,7 +1229,9 @@ def body_fn(carry): return h with pytest.raises(NotImplementedError, match='while'): - self._graph_for(WhileCell) + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + self._graph_for(WhileCell) def test_drtrl_gradient_parity_with_unrolled_model(self): def build_and_grads(cell_cls): diff --git a/braintrace/_compiler/hidden_group_test.py b/braintrace/_compiler/hidden_group_test.py index dadcd817..352cb878 100644 --- a/braintrace/_compiler/hidden_group_test.py +++ b/braintrace/_compiler/hidden_group_test.py @@ -15,6 +15,7 @@ import unittest +import warnings from pprint import pprint import brainstate @@ -1496,10 +1497,12 @@ def _compile_mixing(self, include_recurrent_mixing=False, n=4): cell = WhileMixingCell(n) brainstate.nn.init_all_states(cell) x = brainstate.random.rand(n) - with diagnostic_context() as reporter: - groups, path_to_group = find_hidden_groups_from_module( - cell, x, include_recurrent_mixing=include_recurrent_mixing, - ) + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + with diagnostic_context() as reporter: + groups, path_to_group = find_hidden_groups_from_module( + cell, x, include_recurrent_mixing=include_recurrent_mixing, + ) return cell, x, groups, reporter def test_default_mode_falls_back_to_zero_recurrence(self): @@ -1579,8 +1582,10 @@ def test_jit_wrapped_mixing_in_body_is_still_a_boundary(self): cell = WhileJitMixingCell(4) brainstate.nn.init_all_states(cell) x = brainstate.random.rand(4) - with diagnostic_context() as reporter: - groups, _pg = find_hidden_groups_from_module(cell, x) + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + with diagnostic_context() as reporter: + groups, _pg = find_hidden_groups_from_module(cell, x) assert len(groups) == 1 group = groups[0] # zero-recurrence fallback, exactly like the un-jitted mixing cell diff --git a/braintrace/_compiler/module_info.py b/braintrace/_compiler/module_info.py index 290db31a..cf9f42d9 100644 --- a/braintrace/_compiler/module_info.py +++ b/braintrace/_compiler/module_info.py @@ -74,32 +74,12 @@ def _check_consistent_states_between_model_and_compiler( id(st): st for st in compiled_model_states } - id_to_path = {} - for path, st in retrieved_model_states.items(): - state_id = id(st) - previous = id_to_path.get(state_id) - if previous is None: - id_to_path[state_id] = path - continue - - # Map exposes mapped states through both its internal registry and the - # wrapped module. Prefer the module-facing path: numeric registry keys - # look like layer boundaries to hidden-state grouping. - internal_depth = sum(part == 'dict_vmap_states' for part in path) - previous_internal_depth = sum( - part == 'dict_vmap_states' for part in previous - ) - if internal_depth < previous_internal_depth: - id_to_path[state_id] = path - - # Graph traversal may expose the same state through more than one path. - # Keep the canonical path selected above so each compiled state has exactly - # one model path. - paths_to_remove = [ - path + id_to_path = { + id(st): path for path, st in retrieved_model_states.items() - if id_to_path[id(st)] != path - ] + } + + paths_to_remove = [] for id_ in id_to_path: if id_ not in id_to_compiled_state: path = id_to_path[id_] @@ -194,13 +174,7 @@ def abstractify_model( "The model should be an instance of brainstate.nn.Module. " "Since it allows the explicit definition of the model structure." ) - if isinstance(model, brainstate.nn.Map): - # ``Map`` exposes the same states through implementation paths such as - # ``module`` and ``dict_vmap_states``. Compiler paths are public model - # paths, so retrieve them from the wrapped module directly. - model_retrieved_states = brainstate.graph.states(model.module) - else: - model_retrieved_states = brainstate.graph.states(model) + model_retrieved_states = brainstate.graph.states(model) # --- stateful model, for extracting states, weights, and variables --- # # diff --git a/braintrace/_compiler/module_info_test.py b/braintrace/_compiler/module_info_test.py index ace0c8d7..d37bef8f 100644 --- a/braintrace/_compiler/module_info_test.py +++ b/braintrace/_compiler/module_info_test.py @@ -64,27 +64,6 @@ def test_add_jaxpr_outs_preserves_policy(self): class Test_extract_model_info: - def test_map_hidden_state_aliases_are_deduplicated(self): - batch_size = 3 - rnn = braintrace.nn.GRUCell(2, 4) - mapped = brainstate.nn.Map(rnn, init_map_size=batch_size) - mapped.init_all_states() - - expected_states = brainstate.graph.states(rnn) - expected_params = rnn.states(brainstate.ParamState) - - minfo = braintrace.extract_module_info( - mapped, brainstate.random.rand(batch_size, 2) - ) - states = minfo.retrieved_model_states - - assert set(states) == set(expected_states) - assert set(minfo.weight_path_to_invars) == set(expected_params) - assert all(states[path] is state for path, state in expected_states.items()) - assert len({id(state) for state in states.values()}) == len(states) - assert all('module' not in path for path in states) - assert all('dict_vmap_states' not in path for path in states) - @pytest.mark.parametrize( "cls", [ diff --git a/braintrace/_compiler/scenario_catalog_test.py b/braintrace/_compiler/scenario_catalog_test.py index 48960263..0b10aa9c 100644 --- a/braintrace/_compiler/scenario_catalog_test.py +++ b/braintrace/_compiler/scenario_catalog_test.py @@ -34,6 +34,8 @@ - ``W -> non-gradient-enabled W -> h`` excludes the preceding weight. """ +import warnings + import brainstate import jax import jax.numpy as jnp @@ -54,8 +56,14 @@ def _compile(model, *inputs): - """Compile a model and retain both warnings and structured diagnostics.""" - return compile_etrace_graph(model, *inputs, include_hidden_perturb=False) + """Compile, suppressing expected weight-exclusion UserWarnings. + + Tests still assert on the structured ``DiagnosticKind`` records, so we + silence the warning-stream duplicate for readability. + """ + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + return compile_etrace_graph(model, *inputs, include_hidden_perturb=False) def _relation_set(graph): @@ -961,7 +969,11 @@ def test_cond_branches_full_pipeline_converts(self): def test_scan_body_full_pipeline_unrolls(self): model = ScanBodyRNN(4, loops=3) brainstate.nn.init_all_states(model) - graph = compile_etrace_graph(model, jnp.ones(4)) + with warnings.catch_warnings(): + # Earlier sub-steps' weights are excluded per the + # weight->weight->hidden invariant and warn about it. + warnings.simplefilter('ignore', UserWarning) + graph = compile_etrace_graph(model, jnp.ones(4)) names = [eqn.primitive.name for eqn in graph.module_info.jaxpr.eqns] assert 'scan' not in names # Only the final sub-step's two ETP ops are relations; the earlier @@ -997,8 +1009,10 @@ def test_scan_body_etp_exclude_policy_warns_and_drops(self): jaxpr = make_scan_body_etp_jaxpr(3, 4) policy = braintrace.ControlFlowPolicy(etp_in_control_flow='exclude') - with diagnostic_context() as reporter: - top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + with diagnostic_context() as reporter: + top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) assert top == [], 'ETP inside scan body must NOT bubble up' records = [ @@ -1012,8 +1026,10 @@ def test_cond_branches_etp_exclude_policy_diagnostic_per_branch(self): jaxpr = make_cond_branches_etp_jaxpr(3, 4) policy = braintrace.ControlFlowPolicy(etp_in_control_flow='exclude') - with diagnostic_context() as reporter: - top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + with diagnostic_context() as reporter: + top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) assert top == [] n_cf = sum( @@ -1053,8 +1069,10 @@ def test_while_body_etp_exclude_policy_warns_and_drops(self): jaxpr = make_while_body_etp_jaxpr(4, 4) policy = braintrace.ControlFlowPolicy(etp_in_control_flow='exclude') - with diagnostic_context() as reporter: - top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + with diagnostic_context() as reporter: + top = _scan_jaxpr_for_etp_eqns(jaxpr, policy=policy) assert top == [] kinds = [r.kind for r in reporter.records()] diff --git a/braintrace/_compiler/tests/cell_relation_guardrail_test.py b/braintrace/_compiler/tests/cell_relation_guardrail_test.py index 06e26413..ae578498 100644 --- a/braintrace/_compiler/tests/cell_relation_guardrail_test.py +++ b/braintrace/_compiler/tests/cell_relation_guardrail_test.py @@ -24,6 +24,8 @@ compiler_property_test.py, graph_test.py) and are not duplicated here. """ +import warnings + import brainstate import pytest @@ -49,7 +51,11 @@ def _compile_cell(name, n_in=3, n_out=4): cell = cls(n_in, n_out) brainstate.nn.init_all_states(cell) inp = brainstate.random.rand(n_in) - return braintrace.compile_etrace_graph(cell, inp, include_hidden_perturb=False) + with warnings.catch_warnings(): + # GRUCell legitimately warns when it excludes Wr (W->W->h); the guardrail + # checks the diagnostic records, not the warning. + warnings.simplefilter('ignore') + return braintrace.compile_etrace_graph(cell, inp, include_hidden_perturb=False) @pytest.mark.parametrize('cell_name', list(_CELL_GUARDRAILS)) diff --git a/braintrace/_compiler/tests/compiler_oracle_test.py b/braintrace/_compiler/tests/compiler_oracle_test.py index d58553f2..ec572379 100644 --- a/braintrace/_compiler/tests/compiler_oracle_test.py +++ b/braintrace/_compiler/tests/compiler_oracle_test.py @@ -39,6 +39,8 @@ +import warnings + import brainstate import jax import jax.numpy as jnp @@ -52,8 +54,10 @@ ) -def _compile(model, *args): - return compile_etrace_graph(model, *args, include_hidden_perturb=False) +def _silent_compile(model, *args): + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + return compile_etrace_graph(model, *args, include_hidden_perturb=False) def _transition_callable(rel, group, const_vals): @@ -86,7 +90,7 @@ def test_unbatched_mv_rnn_dhdy_matches_analytic(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.3, -0.7, 1.1]) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) rel = graph.hidden_param_op_relations[0] @@ -118,7 +122,7 @@ def test_elemwise_only_rnn_dhdy_matches_analytic(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.5, -0.2, 0.9, -1.0]) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) rel = graph.hidden_param_op_relations[0] @@ -161,7 +165,7 @@ def test_w1_dhdy_is_direct_only(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.4, -0.6, 1.2]) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) by_path = {r.path: r for r in graph.hidden_param_op_relations} @@ -186,7 +190,7 @@ def test_w2_dhdy_is_correct(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.4, -0.6, 1.2]) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) by_path = {r.path: r for r in graph.hidden_param_op_relations} @@ -233,7 +237,7 @@ def test_unbatched_mv_rnn_fd_matches_ad(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.1, 0.2, -0.3]) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) rel = graph.hidden_param_op_relations[0] self._check_fd(rel, rel.hidden_groups[0], temps) @@ -243,7 +247,7 @@ def test_partial_path_w1_fd_matches_ad(self): brainstate.nn.init_all_states(model) inp = jnp.array([0.05, -0.1, 0.2]) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) _, _, _, temps = graph.module_info.jaxpr_call(inp) rel = next( r for r in graph.hidden_param_op_relations diff --git a/braintrace/_compiler/tests/compiler_property_test.py b/braintrace/_compiler/tests/compiler_property_test.py index 35920121..0aeb81d2 100644 --- a/braintrace/_compiler/tests/compiler_property_test.py +++ b/braintrace/_compiler/tests/compiler_property_test.py @@ -42,6 +42,7 @@ import importlib.util +import warnings import pytest @@ -72,8 +73,10 @@ ) -def _compile(model, *args): - return compile_etrace_graph(model, *args, include_hidden_perturb=False) +def _silent_compile(model, *args): + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + return compile_etrace_graph(model, *args, include_hidden_perturb=False) def _summary(graph): @@ -119,8 +122,8 @@ def test_unbatched_mv_rnn_is_idempotent(self, n_in, n_out): m2 = UnbatchedMvRNN(n_in, n_out) brainstate.nn.init_all_states(m2) - s1 = _summary(_compile(m1, inp)) - s2 = _summary(_compile(m2, inp)) + s1 = _summary(_silent_compile(m1, inp)) + s2 = _summary(_silent_compile(m2, inp)) assert s1 == s2 @given( @@ -135,8 +138,8 @@ def test_partial_path_rnn_is_idempotent(self, n): m2 = PartialPathRNN(n, n) brainstate.nn.init_all_states(m2) - s1 = _summary(_compile(m1, inp)) - s2 = _summary(_compile(m2, inp)) + s1 = _summary(_silent_compile(m1, inp)) + s2 = _summary(_silent_compile(m2, inp)) assert s1 == s2 @@ -160,7 +163,7 @@ def test_one_relation_per_layer_scoped_to_own_h(self, depth, n_in, n_out): brainstate.nn.init_all_states(model) inp = jnp.zeros(n_in) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) rels = graph.hidden_param_op_relations assert len(rels) == depth @@ -187,7 +190,7 @@ def test_two_call_sites_yield_two_relations(self, n): brainstate.nn.init_all_states(model) inp = jnp.zeros(n) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) rels = graph.hidden_param_op_relations assert len(rels) == 2 @@ -255,7 +258,7 @@ def test_only_last_weight_registers(self, chain_len, n): brainstate.nn.init_all_states(model) inp = jnp.zeros(n) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) rels = graph.hidden_param_op_relations included_paths = {r.path for r in rels} last = (f'w{chain_len - 1}',) @@ -291,7 +294,7 @@ def test_classification_is_shape_invariant(self, n_in, n_out): brainstate.nn.init_all_states(model) inp = jnp.zeros(n_in) - graph = _compile(model, inp) + graph = _silent_compile(model, inp) by_path = {r.path: r for r in graph.hidden_param_op_relations} assert by_path[('w1',)].path_classification == { diff --git a/braintrace/_legacy/_ops_test.py b/braintrace/_legacy/_ops_test.py index 1b81e096..2470ec3d 100644 --- a/braintrace/_legacy/_ops_test.py +++ b/braintrace/_legacy/_ops_test.py @@ -25,10 +25,11 @@ +import warnings + import jax import jax.numpy as jnp import numpy as np -import pytest import braintrace from braintrace._legacy import ( @@ -150,14 +151,22 @@ class TestDeprecationWarnings: # and not when importing from the private ``braintrace._legacy`` submodule. def test_matmul_op_access_warns(self): - with pytest.warns(DeprecationWarning, match='MatMulOp'): + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter('always') _ = braintrace.MatMulOp + assert any( + issubclass(w.category, DeprecationWarning) + and 'MatMulOp' in str(w.message) + for w in captured + ) - def test_construction_does_not_warn(self, recwarn): + def test_construction_does_not_warn(self): # The shim classes themselves no longer warn; construction is silent. - MatMulOp() + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter('always') + MatMulOp() assert not any( - issubclass(w.category, DeprecationWarning) for w in recwarn + issubclass(w.category, DeprecationWarning) for w in captured ) diff --git a/braintrace/_legacy/_params_test.py b/braintrace/_legacy/_params_test.py index 799e52e6..52408cdf 100644 --- a/braintrace/_legacy/_params_test.py +++ b/braintrace/_legacy/_params_test.py @@ -26,10 +26,11 @@ +import warnings + import brainstate import jax.numpy as jnp import numpy as np -import pytest import braintrace from braintrace._legacy import ( @@ -171,8 +172,14 @@ class TestDeprecationWarnings: # and not when importing from the private ``braintrace._legacy`` submodule. def test_etrace_param_access_warns(self): - with pytest.warns(DeprecationWarning, match='ETraceParam'): + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter('always') _ = braintrace.ETraceParam + assert any( + issubclass(w.category, DeprecationWarning) + and 'ETraceParam' in str(w.message) + for w in captured + ) # --------------------------------------------------------------------------- diff --git a/braintrace/_op/_primitive.py b/braintrace/_op/_primitive.py index 92e52b4c..faadc7ad 100644 --- a/braintrace/_op/_primitive.py +++ b/braintrace/_op/_primitive.py @@ -58,8 +58,6 @@ 'register_primitive', ] -_ETP_BATCHING_RULES: dict[Primitive, Callable[..., Any]] = {} - class ETPPrimitive(Primitive): """A JAX ``Primitive`` with ETP rule registration helpers. @@ -338,7 +336,6 @@ def _batching(args: Any, dims: Any, **params: Any) -> Any: ) return jax.vmap(partial(impl_fn, **params), in_axes=dims)(*args), 0 - _ETP_BATCHING_RULES[p] = _batching batching.primitive_batchers[p] = _batching return p diff --git a/braintrace/_op/conv.py b/braintrace/_op/conv.py index a8777352..f53dbe7b 100644 --- a/braintrace/_op/conv.py +++ b/braintrace/_op/conv.py @@ -121,9 +121,8 @@ import jax import jax.numpy as jnp import brainunit as u -from jax.interpreters import batching -from ._primitive import _ETP_BATCHING_RULES, register_primitive +from ._primitive import register_primitive from ._registries import ETP_RULES_INSTANT_DRTRL, ETP_RULES_SOLVE_DRTRL from braintrace._typing import ArrayLike, WeightFn @@ -759,51 +758,6 @@ def _conv_init_pp(x_var: Any, y_var: Any, weight_vars: dict[str, Any], trainable_invars_fn=_conv_trainable_invars, x_invar_index=0, ) - -_default_conv_batcher = _ETP_BATCHING_RULES[etp_conv_p] - - -def _conv_lhs_batch_axis(params: dict[str, Any]) -> int: - """Return the input batch-axis position encoded by dimension numbers.""" - dn = params.get('dimension_numbers') - if dn is None: - return 0 - if isinstance(dn, tuple) and len(dn) == 3 and isinstance(dn[0], str): - return dn[0].index('N') - return dn.lhs_spec[0] - - -def _conv_batcher(args: Any, dims: Any, **params: Any) -> Any: - """Preserve ``etp_conv_p`` when mapping only the convolution input.""" - x_idx = 0 - if ( - dims[x_idx] is None - or any(d is not None for i, d in enumerate(dims) if i != x_idx) - ): - return _default_conv_batcher(args, dims, **params) - - x = jnp.moveaxis(args[x_idx], dims[x_idx], 0) - lhs_batch_axis = _conv_lhs_batch_axis(params) - x = jnp.moveaxis(x, lhs_batch_axis + 1, 1) - map_size, inner_batch_size = x.shape[:2] - merged_x = x.reshape(map_size * inner_batch_size, *x.shape[2:]) - merged_x = jnp.moveaxis(merged_x, 0, lhs_batch_axis) - - merged_args = tuple( - merged_x if i == x_idx else arg for i, arg in enumerate(args) - ) - y = etp_conv_p.bind(*merged_args, **params) - - _, _, output_batch_axis, _ = _conv_layout(params) - y = jnp.moveaxis(y, output_batch_axis, 0) - y = y.reshape(map_size, inner_batch_size, *y.shape[1:]) - y = jnp.moveaxis(y, 1, output_batch_axis + 1) - return y, 0 - - -batching.primitive_batchers[etp_conv_p] = _conv_batcher - - def _conv_snap_anchor(eqn_params: dict) -> bool: """Declare the SnAp-n trace anchor for ``etp_conv``. diff --git a/docs/advanced/batching.ipynb b/docs/advanced/batching.ipynb index 7b97702f..195c5f12 100644 --- a/docs/advanced/batching.ipynb +++ b/docs/advanced/batching.ipynb @@ -24,18 +24,19 @@ "source": [ "## Map-Based Batching (Recommended)\n", "\n", - "The recommended approach is to keep the model's update logic single-sample\n", - "and let `brainstate.nn.Map` manage independent state copies across the batch.\n", - "`braintrace.compile(..., batch_size=B, vmap=True)` does this in one call:\n", + "Keep the model's update logic single-sample and let `brainstate.nn.Map` manage\n", + "independent state copies across the batch. Set the mapped learner up explicitly:\n", "\n", - "1. It wraps the model with `brainstate.nn.Map(model, init_map_size=B)`.\n", - "2. It initializes the mapped states through `mapped_model.init_all_states()`.\n", - "3. It compiles the ETP graph from one batched time step with shape\n", + "1. Create exactly one `brainstate.nn.Map(model, init_map_size=B)`.\n", + "2. Initialize it with `mapped_model.init_all_states()`.\n", + "3. Construct the online-learning algorithm with the mapped model.\n", + "4. Compile the ETP graph from one complete batched time step with shape\n", " `(batch_size, n_in)`.\n", - "4. It returns the concrete online-learning algorithm, ready for batched calls.\n", "\n", - "The returned learner exposes `report` and the rest of the algorithm API\n", - "directly." + "Do not pass `mapped_model` to `braintrace.compile(..., vmap=True)`. That setup\n", + "path owns its batching transformation and would map an already mapped model a\n", + "second time. When using an explicit `Map`, construct and compile the algorithm\n", + "directly as shown below." ] }, { @@ -78,18 +79,16 @@ "source": [ "model = SimpleGRU(10, 64, 5)\n", "batch_size = 16\n", + "example_input = jnp.zeros((batch_size, 10))\n", "\n", - "# braintrace.compile with vmap=True:\n", - "# - initialises per-sample hidden states (batch_size independent copies)\n", - "# - wraps the model in brainstate.nn.Map and initializes mapped states\n", - "# - compiles the ETP graph from one batched time step\n", - "# - returns the concrete algorithm for parallel mapped execution\n", - "mapped_algo = braintrace.compile(\n", - " model, braintrace.D_RTRL, jnp.zeros((batch_size, 10)),\n", - " batch_size=batch_size, vmap=True,\n", - ")\n", + "# Create and initialize exactly one mapped model.\n", + "mapped_model = brainstate.nn.Map(model, init_map_size=batch_size)\n", + "mapped_model.init_all_states()\n", + "\n", + "# Compile the algorithm directly from the complete batched example input.\n", + "mapped_algo = braintrace.D_RTRL(mapped_model)\n", + "mapped_algo.compile_graph(example_input)\n", "\n", - "# Run on batched input — the returned learner handles the batch axis transparently\n", "x_batch = jnp.ones((batch_size, 10))\n", "out = mapped_algo(x_batch)\n", "print(\"Output shape:\", out.shape) # (16, 5)" @@ -102,13 +101,14 @@ "source": [ "**How it works:**\n", "\n", - "- `braintrace.compile(..., batch_size=B, vmap=True)` creates\n", - " `brainstate.nn.Map(model, init_map_size=B)` and calls\n", - " `mapped_model.init_all_states()`.\n", - "- The algorithm compiles against the batched example input and keeps the ETP\n", - " primitives visible to the compiler.\n", + "- `brainstate.nn.Map(model, init_map_size=B)` creates the mapped state owner;\n", + " `mapped_model.init_all_states()` initializes its independent recurrent states.\n", + "- The algorithm receives that mapped model and compiles against the complete\n", + " batched example input, keeping the batch axis inside the graph.\n", "- Each learner call maps the wrapped model over axis 0 while sharing parameter\n", - " states and maintaining independent recurrent states." + " states and maintaining independent recurrent states.\n", + "- The Map is created once. It is not passed to another API that would wrap it\n", + " again." ] }, { @@ -210,22 +210,16 @@ "@brainstate.transform.jit\n", "def train_step(inputs, targets):\n", " \"\"\"inputs: (n_steps, batch_size, n_in), targets: (batch_size,)\"\"\"\n", - " # braintrace.compile with vmap=True replaces manual Map initialization and compilation.\n", - " # Pass inputs[0] with shape (batch_size, n_in); the compiler traces the mapped model\n", - " # against the complete batched time step, so do not pass inputs[0, 0].\n", - " mapped_algo = braintrace.compile(\n", - " model, braintrace.D_RTRL, inputs[0],\n", - " batch_size=inputs.shape[1], vmap=True,\n", - " )\n", - "\n", " def step_loss(inp):\n", " out = mapped_algo(inp)\n", " return jnp.mean((out - targets) ** 2)\n", "\n", - " # etrace_grad drives the whole sequence and accumulates the per-step online\n", - " # gradients. Map keeps the batch axis inside the compiled graph, while the\n", - " # learner exposes the same driver methods as the unbatched path.\n", - " return mapped_algo.etrace_grad(inputs, step_fn=step_loss, reduction='sum')" + " # etrace_grad drives the sequence and accumulates per-step online\n", + " # gradients. The explicitly mapped model keeps the batch axis inside the\n", + " # compiled graph.\n", + " return mapped_algo.etrace_grad(\n", + " inputs, step_fn=step_loss, reduction='sum'\n", + " )" ] }, { @@ -239,6 +233,12 @@ "model = SimpleGRU(10, 64, 5)\n", "inputs = jnp.ones((20, 16, 10)) # 20 steps, batch 16, 10 features\n", "targets = jnp.zeros((16, 5))\n", + "\n", + "mapped_model = brainstate.nn.Map(model, init_map_size=inputs.shape[1])\n", + "mapped_model.init_all_states()\n", + "mapped_algo = braintrace.D_RTRL(mapped_model)\n", + "mapped_algo.compile_graph(inputs[0])\n", + "\n", "grads = train_step(inputs, targets)\n", "print(\"Gradient keys:\", list(grads.keys()))" ] @@ -250,9 +250,13 @@ "source": [ "**What happens in `train_step`:**\n", "\n", - "1. `braintrace.compile(model, braintrace.D_RTRL, inputs[0], batch_size=B, vmap=True)` wraps the model with `brainstate.nn.Map`, initializes independent per-sample states, compiles from the batched time step, and returns the algorithm directly.\n", - "2. `vmapped_algo.etrace_grad` iterates over time, calls `step_fn`, and accumulates online gradients; `reduction='sum'` accumulates without dividing. The call is identical for mapped and directly batched learners.\n", - "3. The returned `grads` dictionary keeps the original model parameter paths and can be passed to an optimizer such as `braintools.optim.Adam`." + "1. The setup creates one `brainstate.nn.Map`, initializes it, constructs\n", + " `D_RTRL(mapped_model)`, and compiles from the complete batched time step.\n", + "2. `mapped_algo.etrace_grad` iterates over time, calls `step_fn`, and\n", + " accumulates online gradients; `reduction='sum'` accumulates without\n", + " dividing.\n", + "3. The returned gradient keys match `mapped_algo.param_states`. Register those\n", + " learner parameter states with the optimizer when applying the gradients." ] }, { @@ -262,13 +266,14 @@ "source": [ "## Summary\n", "\n", - "- `braintrace.compile(..., batch_size=B, vmap=True)` is the recommended setup\n", - " for batched online learning.\n", - "- Internally it uses `brainstate.nn.Map(model, init_map_size=B)` followed by\n", - " `mapped_model.init_all_states()`.\n", - "- The workflow is: compile with `vmap=True`, scan over time, accumulate\n", - " gradients, and update parameters.\n", - "- For one stream, compile without `vmap=True`; states remain unbatched." + "- Create one `brainstate.nn.Map(model, init_map_size=B)` for batched online\n", + " learning and call `mapped_model.init_all_states()`.\n", + "- Construct the algorithm with that mapped model and call\n", + " `compile_graph(example_input)` using one complete batched time step.\n", + "- Never pass an already mapped model to `compile(..., vmap=True)`; doing so\n", + " would apply batching twice.\n", + "- Scan over time with `etrace_grad` or `etrace_evolve` after compilation.\n", + "- For one stream, initialize and compile the original model without `Map`." ] } ], diff --git a/docs/specs/2026-07-28-warnings-and-map-initialization.md b/docs/specs/2026-07-28-warnings-and-map-initialization.md deleted file mode 100644 index 90d743e4..00000000 --- a/docs/specs/2026-07-28-warnings-and-map-initialization.md +++ /dev/null @@ -1,62 +0,0 @@ -# Warning Visibility and Map Initialization - -## Status - -Approved for implementation. - -## Motivation - -BrainTrace warnings are part of the public diagnostic surface. Library code, -tests, examples, and documentation must not suppress them with -`warnings.catch_warnings` or `warnings.filterwarnings`. - -Mapped state initialization must use the state-management abstraction provided -by BrainState. The supported flow is: - -```python -model = brainstate.nn.Map(model, init_map_size=batch_size) -model.init_all_states() -``` - -This replaces executable uses of -`brainstate.transform.vmap_new_states`. Historical changelog entries may retain -the old symbol when they describe behavior from an earlier release. - -## Requirements - -1. Remove every executable use of `warnings.catch_warnings`, - `warnings.filterwarnings`, and bare `filterwarnings`. -2. Do not add replacement warning filters or warning-suppression helpers. -3. Let BrainTrace warnings reach users and test output unchanged. -4. Replace mapped-state discovery and initialization with - `brainstate.nn.Map(model, init_map_size=...)` followed by - `model.init_all_states()`. -5. Compile mapped algorithms against the complete batched example input and - return the algorithm object directly. -6. Keep mapped model states discoverable by the compiler without duplicating - state paths. -7. Preserve ETP primitives and ETP-specific batching behavior under - `brainstate.nn.Map`. -8. Update affected tests, examples, and tutorials to use the supported Map - workflow. - -## Non-goals - -- Do not change algorithm equations, optimization behavior, or public callable - signatures unrelated to mapped initialization. -- Do not suppress third-party compatibility warnings. -- Do not address the separate `braintools`/`saiunit` quantity compatibility - failures. -- Do not add custom documentation CSS, JavaScript, Sphinx hooks, or static API - pages. - -## Verification - -1. Search the repository for prohibited warning filters and executable - `vmap_new_states` calls. -2. Run focused compiler, mapped-state, convolution batching, and public API - tests. -3. Build the documentation with Sphinx warnings treated as errors. -4. Run the complete test suite and allow it to finish naturally. -5. Report dependency-related failures separately from regressions caused by - this change. diff --git a/docs/tutorials/rnn_online_learning.ipynb b/docs/tutorials/rnn_online_learning.ipynb index 505c0ca8..a3192a9f 100644 --- a/docs/tutorials/rnn_online_learning.ipynb +++ b/docs/tutorials/rnn_online_learning.ipynb @@ -169,16 +169,27 @@ "source": [ "## 4. Online Training with D-RTRL\n", "\n", - "D-RTRL (Diagonal Real-Time Recurrent Learning) is an online learning algorithm provided by `braintrace`. Unlike BPTT, which requires storing the entire computation graph across all time steps, D-RTRL computes approximate gradients incrementally at each time step using **eligibility traces**. It is not generally gradient-equivalent to BPTT outside the assumptions of its diagonal Jacobian approximation.\n", + "D-RTRL (Diagonal Real-Time Recurrent Learning) is an online learning algorithm\n", + "provided by `braintrace`. Unlike BPTT, which requires storing the entire\n", + "computation graph across all time steps, D-RTRL computes approximate gradients\n", + "incrementally using **eligibility traces**. It is not generally\n", + "gradient-equivalent to BPTT outside the assumptions of its diagonal Jacobian\n", + "approximation.\n", "\n", "The key steps in the online training loop are:\n", "\n", - "1. **Compile the model**: `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initializes mapped hidden states, compiles the eligibility-trace graph, and returns the concrete learner - all in one call.\n", - "2. **Warm-up phase**: Use `learner.etrace_evolve(...)` to advance hidden states and eligibility traces without computing a loss gradient.\n", - "3. **Learning phase**: Use `learner.etrace_grad(..., step_fn=step_loss)` to drive the remaining sequence and accumulate online gradients.\n", - "4. **Parameter update**: After processing the full sequence, apply the accumulated gradients to update the parameters.\n", - "\n", - "The `D_RTRL` class wraps the model and handles the eligibility trace bookkeeping automatically." + "1. **Map once**: Create `brainstate.nn.Map(model, init_map_size=B)` and call\n", + " `mapped_model.init_all_states()`.\n", + "2. **Compile directly**: Construct `braintrace.D_RTRL(mapped_model)` and compile\n", + " from one complete batched time step.\n", + "3. **Warm up**: Use `learner.etrace_evolve(...)` to advance hidden states and\n", + " eligibility traces without computing a loss gradient.\n", + "4. **Learn**: Use `learner.etrace_grad(..., step_fn=step_loss)` to accumulate\n", + " online gradients, then update the parameters.\n", + "\n", + "An already mapped model must not be passed to\n", + "`braintrace.compile(..., vmap=True)`, because that would apply a second mapping\n", + "layer." ] }, { @@ -199,20 +210,20 @@ " \"\"\"Train one GRU with D-RTRL over precomputed copying batches.\"\"\"\n", " brainstate.random.seed(21)\n", " model = GRUNet(10, 64, 10)\n", - " opt = braintools.optim.Adam(lr)\n", - " weights = model.states().subset(brainstate.ParamState)\n", - " opt.register_trainable_weights(weights)\n", - "\n", " batch_size = input_batches.shape[2]\n", - " learner = braintrace.compile(\n", - " model, braintrace.D_RTRL, input_batches[0],\n", - " batch_size=batch_size, vmap=True,\n", - " )\n", + "\n", + " mapped_model = brainstate.nn.Map(model, init_map_size=batch_size)\n", + " mapped_model.init_all_states()\n", + " learner = braintrace.D_RTRL(mapped_model)\n", + " learner.compile_graph(input_batches[0])\n", + "\n", + " opt = braintools.optim.Adam(lr)\n", + " opt.register_trainable_weights(learner.param_states)\n", "\n", " @brainstate.transform.jit\n", " def train_step(inputs, targets):\n", - " brainstate.nn.reset_all_states(model, batch_size=batch_size)\n", - " learner.reset_state(batch_size=batch_size)\n", + " brainstate.nn.reset_all_states(mapped_model)\n", + " learner.reset_state()\n", "\n", " # The loss for ONE step. `etrace_grad` owns the loop; this owns the\n", " # model call, so multi-head models and regularizers need no special\n", @@ -435,18 +446,23 @@ "source": [ "## 8. Summary\n", "\n", - "In this tutorial, we demonstrated how to use `braintrace` for online learning of a GRU network on the copying task.\n", + "In this tutorial, we demonstrated online learning of a GRU on the copying task.\n", "\n", "**Key takeaways:**\n", "\n", - "- **D-RTRL** provides approximate online gradients with `O(B * theta)` complexity, where `B` is the batch size and `theta` is the number of parameters. Unlike BPTT, it does not need to store the full unrolled computation graph.\n", - "- A single call to `braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=B, vmap=True)` initializes `brainstate.nn.Map` states, compiles the eligibility-trace graph, and returns the concrete learner ready for batched training.\n", - "- Online learning uses `learner.etrace_evolve` for gradient-free sequence prefixes and `learner.etrace_grad` for sequence objectives; both compose with `brainstate.transform.jit`.\n", - "- `braintrace` is particularly effective for RNN models with gating mechanisms (GRU, LSTM), where the internal dynamics naturally support eligibility trace propagation.\n", + "- **D-RTRL** provides approximate online gradients with `O(B * theta)`\n", + " complexity, where `B` is the batch size and `theta` is the number of\n", + " parameters. Unlike BPTT, it does not store the full unrolled graph.\n", + "- Batched online learning creates one `brainstate.nn.Map`, initializes it, and\n", + " passes it directly to `braintrace.D_RTRL` before `compile_graph` is called on\n", + " a complete batched time step.\n", + "- Do not pass an already mapped model to `compile(..., vmap=True)`.\n", + "- Use `learner.etrace_evolve` for gradient-free prefixes and\n", + " `learner.etrace_grad` for sequence objectives.\n", "\n", "For more details, see:\n", "- [Key Concepts](../quickstart/concepts.ipynb) for the theoretical background.\n", - "- [SNN Online Learning](./snn_online_learning.ipynb) for applying the same approach to spiking neural networks." + "- [SNN Online Learning](./snn_online_learning.ipynb) for spiking networks." ] } ], diff --git a/docs/tutorials/snn_online_learning.ipynb b/docs/tutorials/snn_online_learning.ipynb index 4459a398..fdc5e368 100644 --- a/docs/tutorials/snn_online_learning.ipynb +++ b/docs/tutorials/snn_online_learning.ipynb @@ -197,16 +197,24 @@ "source": [ "## 3. Training with ES-D-RTRL\n", "\n", - "We set up online learning using `braintrace.compile` with `braintrace.pp_prop` (also exposed as `braintrace.ES_D_RTRL` and `braintrace.IODimVjpAlgorithm`). The `decay_or_rank` parameter controls the trace approximation:\n", - "\n", - "* **`decay_or_rank=float` in (0, 1)** -- exponentially-smoothed trace. Larger values retain a longer history but may be less stable; this short teaching task uses `0.5`. Memory cost: `O(B * (I + O))` per layer.\n", - "* **`decay_or_rank=int >= 1`** -- an alternative parameterization of the same decay, using `decay = (rank - 1) / (rank + 1)`. It does not allocate multiple rank factors, so memory remains `O(B * (I + O))`.\n", - "\n", - "Use the float form to set the decay directly, or the integer form to select the corresponding decay through the documented conversion. Neither form creates an independent rank-versus-memory trade-off.\n", - "\n", - "`braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.5)` initializes mapped per-sample states, builds the ETP graph, and returns the concrete algorithm — no separate `Map`, `init_all_states`, or `compile_graph` calls are needed.\n", - "\n", - "`braintrace.D_RTRL` is the alternative parameter-dimensional estimator. It stores parameter-shaped eligibility traces and can use substantially more memory; neither estimator is generally gradient-equivalent to BPTT outside its documented assumptions." + "We use `braintrace.pp_prop` (also exposed as `braintrace.ES_D_RTRL` and\n", + "`braintrace.IODimVjpAlgorithm`). The `decay_or_rank` parameter controls the\n", + "trace approximation:\n", + "\n", + "* **`decay_or_rank=float` in (0, 1)** -- exponentially smoothed trace. Larger\n", + " values retain a longer history; this short teaching task uses `0.5`.\n", + "* **`decay_or_rank=int >= 1`** -- an alternative parameterization using\n", + " `decay = (rank - 1) / (rank + 1)`.\n", + "\n", + "For batched learning, create one `brainstate.nn.Map`, initialize it, construct\n", + "`braintrace.pp_prop(mapped_model, decay_or_rank=0.5)`, and compile from one\n", + "complete batched time step. Do not pass the mapped model to\n", + "`braintrace.compile(..., vmap=True)`, because that would map it a second time.\n", + "\n", + "`braintrace.D_RTRL` is the alternative parameter-dimensional estimator. It\n", + "stores parameter-shaped eligibility traces and can use substantially more\n", + "memory; neither estimator is generally gradient-equivalent to BPTT outside its\n", + "documented assumptions." ] }, { @@ -248,23 +256,22 @@ " with brainstate.environ.context(dt=1. * u.ms):\n", " brainstate.random.seed(37)\n", " model = LIF_SNN(n_in, n_rec, n_out)\n", - " opt = braintools.optim.Adam(lr)\n", - " weights = model.states(brainstate.ParamState)\n", - " opt.register_trainable_weights(weights)\n", "\n", - " learner = braintrace.compile(\n", - " model, braintrace.pp_prop, input_batches[0],\n", - " batch_size=batch_size, vmap=True, decay_or_rank=0.5,\n", - " )\n", + " mapped_model = brainstate.nn.Map(model, init_map_size=batch_size)\n", + " mapped_model.init_all_states()\n", + " learner = braintrace.pp_prop(mapped_model, decay_or_rank=0.5)\n", + " learner.compile_graph(input_batches[0])\n", + "\n", + " opt = braintools.optim.Adam(lr)\n", + " opt.register_trainable_weights(learner.param_states)\n", "\n", " @brainstate.transform.jit\n", " def train_step(inputs, targets):\n", - " brainstate.nn.reset_all_states(model, batch_size=batch_size)\n", - " learner.reset_state(batch_size=batch_size)\n", + " brainstate.nn.reset_all_states(mapped_model)\n", + " learner.reset_state()\n", "\n", " # One step's loss. `etrace_grad` owns the loop; `step_fn` owns the\n", - " # model call and hands the step's logits back as aux, so they come\n", - " # out of the driver stacked over time for the accuracy readout.\n", + " # model call and returns logits as auxiliary output.\n", " def step_loss(inp):\n", " output = learner(inp)\n", " loss = braintools.metric.softmax_cross_entropy_with_integer_labels(\n", @@ -272,10 +279,6 @@ " ).mean()\n", " return loss, output\n", "\n", - " # One call slices the sequence, differentiates each step's loss\n", - " # online, and accumulates the per-step gradients. `reduction='sum'`\n", - " # keeps the accumulated scale this example's learning rate was\n", - " # tuned at.\n", " grads, step_losses, outputs = learner.etrace_grad(\n", " inputs, step_fn=step_loss, has_aux=True,\n", " reduction='sum', return_value=True,\n", @@ -421,17 +424,22 @@ "source": [ "## 5. Summary\n", "\n", - "In this tutorial, we demonstrated how to train a spiking neural network with online learning using BrainTrace. Here are the key takeaways:\n", - "\n", - "1. **Model Construction**: Use `braintrace.nn.Linear` and `braintrace.nn.LeakyRateReadout` for layers that should participate in online learning (ETP-aware). Combine them with spiking neuron models from `brainpy.state` (e.g., `LIF`).\n", - "\n", - "2. **Online Learning Setup**: Use `learner = braintrace.compile(model, braintrace.ES_D_RTRL, x0, batch_size=B, vmap=True, decay_or_rank=0.5)` to initialize mapped states, compile the ETP graph, and return the concrete learner in one call. Then use `learner.etrace_grad(...)` to drive the sequence and accumulate online gradients.\n", + "This tutorial demonstrated online learning for a spiking neural network.\n", "\n", - "3. **Scalability**: ES-D-RTRL achieves O(B(I+O)) memory complexity, making it practical for large spiking networks. The `decay_or_rank` parameter controls the trace approximation quality.\n", + "1. **Model construction**: Use ETP-aware BrainTrace layers for parameters that\n", + " should participate in online learning.\n", + "2. **Mapped setup**: Create and initialize one `brainstate.nn.Map`, pass it\n", + " directly to `braintrace.pp_prop`, and compile from a complete batched time\n", + " step.\n", + "3. **No duplicate mapping**: Never pass an already mapped model to\n", + " `compile(..., vmap=True)`.\n", + "4. **Sequence training**: Use `learner.etrace_grad(...)` to drive the sequence\n", + " and accumulate online gradients.\n", "\n", - "4. **Batching**: `braintrace.compile(..., batch_size=B, vmap=True)` handles `brainstate.nn.Map` initialization and mapped execution automatically.\n", + "The separate `braintools`/`saiunit` physical-unit compatibility issue is not\n", + "handled by this documentation change.\n", "\n", - "For more advanced topics, including training on real neuromorphic datasets (N-MNIST) and comparing online learning with BPTT, see:\n", + "For more advanced topics, see:\n", "- [pp_prop algorithm tutorial](pp_prop.ipynb)\n", "- [Key Concepts](../quickstart/concepts.ipynb)" ] diff --git a/examples/002-coba-ei-rsnn.py b/examples/002-coba-ei-rsnn.py index 81511c45..72d182c3 100644 --- a/examples/002-coba-ei-rsnn.py +++ b/examples/002-coba-ei-rsnn.py @@ -297,8 +297,12 @@ def visualize(self, inputs, n2show: int = 5): n_seq = inputs.shape[0] batch_size = inputs.shape[1] - model = brainstate.nn.Map(self, init_map_size=batch_size) - model.init_all_states() + @brainstate.transform.vmap_new_states(state_tag='new', axis_size=batch_size) + def init(): + brainstate.nn.init_all_states(self) + + init() + model = brainstate.nn.Vmap(self, vmap_states='new') def step(inp): out = model(inp) diff --git a/examples/003-snn-memory-and-speed-evaluation-all.py b/examples/003-snn-memory-and-speed-evaluation-all.py index a102d1bc..171cb2da 100644 --- a/examples/003-snn-memory-and-speed-evaluation-all.py +++ b/examples/003-snn-memory-and-speed-evaluation-all.py @@ -353,10 +353,15 @@ def _step(i, inp): def _compile_etrace_function(self, input_info): # kept manual: braintrace.compile has no path for this state scheme. - # It owns explicitly tagged per-sample states and resets those states - # through ``transform.vmap``. ``compile(vmap=True)`` instead wraps the - # target in ``brainstate.nn.Map``, which would change this benchmark's - # state ownership and reset contract. + # It offers two: init_all_states(batch_size=B) (vmap=False), or + # vmap_new_states(state_tag='new') + compile_graph on an *unbatched* + # sample + an ETraceVmap wrapper (vmap=True). This benchmark uses a + # third -- vmap_init_all_states(state_tag='new') for the per-sample + # states, compile_graph on the *batched* example, no wrapper, and an + # explicit brainstate.transform.vmap(in_states=...) only for the reset. + # compile's vmap branch would also reject `input_info`: it strips the + # batch axis with `a[0]`, and a jax.ShapeDtypeStruct is not + # subscriptable. if self.args.method == 'expsm_diag': model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay) elif self.args.method == 'diag': diff --git a/examples/003-snn-memory-and-speed-evaluation-batched.py b/examples/003-snn-memory-and-speed-evaluation-batched.py index 22c8077d..fb23afc6 100644 --- a/examples/003-snn-memory-and-speed-evaluation-batched.py +++ b/examples/003-snn-memory-and-speed-evaluation-batched.py @@ -477,10 +477,15 @@ def _step(i, inp): def _compile_etrace_function(self, input_info): # kept manual: braintrace.compile has no path for this state scheme. - # It owns explicitly tagged per-sample states and resets those states - # through ``transform.vmap``. ``compile(vmap=True)`` instead wraps the - # target in ``brainstate.nn.Map``, which would change this benchmark's - # state ownership and reset contract. + # It offers two: init_all_states(batch_size=B) (vmap=False), or + # vmap_new_states(state_tag='new') + compile_graph on an *unbatched* + # sample + an ETraceVmap wrapper (vmap=True). This benchmark uses a + # third -- vmap_init_all_states(state_tag='new') for the per-sample + # states, compile_graph on the *batched* example, no wrapper, and an + # explicit brainstate.transform.vmap(in_states=...) only for the reset. + # compile's vmap branch would also reject `input_info`: it strips the + # batch axis with `a[0]`, and a jax.ShapeDtypeStruct is not + # subscriptable. if self.args.method == 'expsm_diag': model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay) elif self.args.method == 'diag': diff --git a/examples/003-snn-memory-and-speed-evaluation-vmap.py b/examples/003-snn-memory-and-speed-evaluation-vmap.py index 63c8eb0c..4bbf3706 100644 --- a/examples/003-snn-memory-and-speed-evaluation-vmap.py +++ b/examples/003-snn-memory-and-speed-evaluation-vmap.py @@ -458,34 +458,32 @@ def _step(i, inp): return losses.mean(), acc def _compile_etrace_function(self, input_info): - # Kept explicit because this benchmark compiles from an unbatched - # ShapeDtypeStruct rather than a concrete batched example input. - mapped_target = brainstate.nn.Map( - self.target, init_map_size=self.args.batch_size - ) - mapped_target.init_all_states() - + # kept manual: this *is* compile(..., vmap=True)'s scheme -- + # vmap_new_states(state_tag='new') + init_all_states + compile_graph on + # the unbatched sample + a Vmap wrapper -- but compile cannot take this + # example input. It strips the batch axis with `a[0]`, and `input_info` + # is an unbatched jax.ShapeDtypeStruct, which is not subscriptable. + # (A benchmark builds the graph from a shape, never from real data.) if self.args.method == 'expsm_diag': - model = braintrace.ES_D_RTRL( - mapped_target, self.args.etrace_decay, - ) + model = braintrace.ES_D_RTRL(self.target, self.args.etrace_decay, ) elif self.args.method == 'diag': - model = braintrace.D_RTRL(mapped_target) + model = braintrace.D_RTRL(self.target, ) else: raise ValueError(f'Unknown online learning methods: {self.args.method}.') - batched_input_info = jax.ShapeDtypeStruct( - (self.args.batch_size, *input_info.shape), input_info.dtype - ) - model.compile_graph(batched_input_info) - run_model = model + # initialize the states + @brainstate.transform.vmap_new_states(state_tag='new', axis_size=self.args.batch_size) + def init(): + brainstate.nn.init_all_states(self.target) + model.compile_graph(input_info) + + init() + run_model = brainstate.nn.Vmap(model, vmap_states='new') @brainstate.transform.jit + @brainstate.transform.vmap(in_states=run_model.states('new')) def reset_state(): - brainstate.nn.reset_all_states( - self.target, batch_size=self.args.batch_size - ) - run_model.reset_state(batch_size=self.args.batch_size) + brainstate.nn.reset_all_states(run_model) @brainstate.transform.jit def _etrace_single_run(i, batch_inp): diff --git a/examples/004-feedforward-conv-snn.py b/examples/004-feedforward-conv-snn.py index 22898993..4dcf6845 100644 --- a/examples/004-feedforward-conv-snn.py +++ b/examples/004-feedforward-conv-snn.py @@ -246,15 +246,25 @@ def batch_train(self, inputs, targets): # inputs: [n_step, n_batch, ...] # targets: [n_batch, n_out] - # compile wraps the model with Map, initializes per-sample states, and - # builds the graph from the batched single-step input. + # One call replaces init_all_states + compile_graph + Vmap. Pass the + # batched single step inputs[0]; compile strips axis 0 to recover the + # per-sample example, so this is the same graph the manual expansion in + # examples/drtrl/02-batching-vmap.py builds by hand. + # model = braintrace.compile(self.target, braintrace.ES_D_RTRL, inputs[0], + # batch_size=inputs.shape[1], vmap=True, + # decay_or_rank=self.decay_or_rank) with brainstate.environ.context(fit=True): model = braintrace.compile( self.target, braintrace.D_RTRL, inputs[0], batch_size=inputs.shape[1], vmap=True, ) - model.show_graph() + # show_graph() is a post-compile diagnostic and lives on the learner, not + # on the vmap wrapper -- ETraceVmap forwards the drivers, not the + # introspection surface. Reading through .module is fine here; only + # *driving* through it would be wrong (it would drive the unbatched + # learner and give per-lane-wrong results). + model.module.show_graph() def _etrace_grad(inp): with brainstate.environ.context(fit=True): diff --git a/examples/100-gru-on-copying-task.py b/examples/100-gru-on-copying-task.py index 9699647c..7cd3f5b7 100644 --- a/examples/100-gru-on-copying-task.py +++ b/examples/100-gru-on-copying-task.py @@ -179,11 +179,14 @@ def batch_train(self, inputs, targets): # 需要求解梯度的参数 weights = self.target.states(brainstate.ParamState) - # kept manual: BPTT baseline, with mapped per-sample states - model = brainstate.nn.Map( - self.target, init_map_size=inputs.shape[1] - ) - model.init_all_states() + # kept manual: BPTT baseline — no online algorithm to migrate + # initialize the states + @brainstate.transform.vmap_new_states(state_tag='new', axis_size=inputs.shape[1]) + def init(): + brainstate.nn.init_all_states(self.target) + + init() + model = brainstate.nn.Vmap(self.target, vmap_states='new') def _run_step_train(inp, tar): out = model(inp) diff --git a/examples/drtrl/02-batching-vmap.py b/examples/drtrl/02-batching-vmap.py index 51fe06c3..76047b4c 100644 --- a/examples/drtrl/02-batching-vmap.py +++ b/examples/drtrl/02-batching-vmap.py @@ -1,9 +1,9 @@ # Copyright 2026 BrainX Ecosystem Limited. Licensed under the Apache License, 2.0. """02 Batching with the canonical public learner workflow. -``braintrace.compile(..., vmap=True)`` wraps the model in -``brainstate.nn.Map``, initializes independent per-sample states, and returns -the compiled learner. Drive that learner directly with ``etrace_grad``. +``braintrace.compile(..., vmap=True)`` creates one eligibility-trace learner +per batch lane. Drive the returned learner directly with ``etrace_grad`` so +each sample keeps its own hidden and eligibility-trace state. """ import pathlib @@ -63,7 +63,7 @@ def step_loss(inp, tar): plt.plot(losses); plt.xlabel('epoch'); plt.ylabel('MSE') - plt.title('02 Batching via public Map-backed learner workflow'); + plt.title('02 Batching via public learner workflow'); plt.show() return {"losses": losses} diff --git a/examples/pp_prop/05-batching-vmap.py b/examples/pp_prop/05-batching-vmap.py index e10b94c7..ad91bf38 100644 --- a/examples/pp_prop/05-batching-vmap.py +++ b/examples/pp_prop/05-batching-vmap.py @@ -1,10 +1,16 @@ # Copyright 2026 BrainX Ecosystem Limited. Licensed under the Apache License, 2.0. -# ``online_train_epoch`` uses braintrace.compile(..., vmap=True). -"""05 - Batching via ``braintrace.compile(..., vmap=True)``. +"""05 · Batching via ``braintrace.compile(..., vmap=True)``. -The model keeps single-sample update logic. ``braintrace.compile`` wraps it in -``brainstate.nn.Map``, initializes per-sample states, and compiles pp_prop from -one batched time step. +The network and the pp_prop algorithm are defined unbatched; ``compile`` with +``vmap=True`` replicates them across the batch dimension (it initializes the +states inside a ``vmap_new_states`` scope, builds the eligibility-trace graph +on one unbatched sample, and returns an ``ETraceVmap``). pp_prop's per-rule +init is aware of batching and allocates batched eligibility traces +automatically. This is the default batching path used by examples 01-04, and +it lives in ``_shared.online_train_epoch``, which this file calls. + +For the same three steps written out by hand, see +``examples/drtrl/02-batching-vmap.py``. """ import pathlib diff --git a/examples/pp_prop/12-classification-neuromorphic.py b/examples/pp_prop/12-classification-neuromorphic.py index 5e094d7d..5d5c2730 100644 --- a/examples/pp_prop/12-classification-neuromorphic.py +++ b/examples/pp_prop/12-classification-neuromorphic.py @@ -36,11 +36,14 @@ def _accuracy(outputs_seq, labels): def _eval(model, inputs, labels): - mapped_model = brainstate.nn.Map( - model, init_map_size=inputs.shape[1] - ) - mapped_model.init_all_states() - outs = brainstate.transform.for_loop(lambda x: mapped_model(x), inputs) + # kept manual: eval re-init, no online construction + @brainstate.transform.vmap_new_states(state_tag="new", axis_size=inputs.shape[1]) + def init(): + brainstate.nn.init_all_states(model) + + init() + vmap_model = brainstate.nn.Vmap(model, vmap_states="new") + outs = brainstate.transform.for_loop(lambda x: vmap_model(x), inputs) return _accuracy(outs, labels) diff --git a/examples/pp_prop/14-knob-vjp-method-contrast.py b/examples/pp_prop/14-knob-vjp-method-contrast.py index 02c9997c..b3470578 100644 --- a/examples/pp_prop/14-knob-vjp-method-contrast.py +++ b/examples/pp_prop/14-knob-vjp-method-contrast.py @@ -35,11 +35,14 @@ def _accuracy(outputs_seq, labels): def _eval(model, inputs, labels): - mapped_model = brainstate.nn.Map( - model, init_map_size=inputs.shape[1] - ) - mapped_model.init_all_states() - outs = brainstate.transform.for_loop(lambda x: mapped_model(x), inputs) + # kept manual: eval re-init, no online construction + @brainstate.transform.vmap_new_states(state_tag="new", axis_size=inputs.shape[1]) + def init(): + brainstate.nn.init_all_states(model) + + init() + vmap_model = brainstate.nn.Vmap(model, vmap_states="new") + outs = brainstate.transform.for_loop(lambda x: vmap_model(x), inputs) return _accuracy(outs, labels) diff --git a/examples/pp_prop/README.md b/examples/pp_prop/README.md index f45611de..dd5ae14c 100644 --- a/examples/pp_prop/README.md +++ b/examples/pp_prop/README.md @@ -34,7 +34,7 @@ if sklearn is missing). | 02 | `02-neurons-alif-dms.py` | ALIF (adaptive threshold) on delayed-match-to-sample | | 03 | `03-neurons-gif-working-memory.py` | GIF with heterogeneous tau_I2 on working-memory recall | | 04 | `04-neurons-coba-ei-rsnn.py` | Dale-law E/I RSNN on small Poisson-MNIST | -| 05 | `05-batching-vmap.py` | Batching via `brainstate.nn.Map` | +| 05 | `05-batching-vmap.py` | Batching via `brainstate.nn.Vmap(vmap_states='new')` | | 06 | `06-batching-batched.py` | Batching via the batched ETP primitive path | | 07 | `07-vjp-single-step.py` | `vjp_method='single-step'` (default) | | 08 | `08-vjp-multi-step.py` | `vjp_method='multi-step'` for temporal credit | diff --git a/examples/pp_prop/_shared.py b/examples/pp_prop/_shared.py index a42c2fe2..c03ee6b4 100644 --- a/examples/pp_prop/_shared.py +++ b/examples/pp_prop/_shared.py @@ -389,14 +389,16 @@ def bptt_train_epoch_fixed_target( """BPTT baseline with per-step softmax-cross-entropy over a fixed label.""" weights = model.states(brainstate.ParamState) - # kept manual: BPTT baseline, with mapped per-sample states - mapped_model = brainstate.nn.Map( - model, init_map_size=inputs.shape[1] - ) - mapped_model.init_all_states() + # kept manual: BPTT re-init — no algorithm construction, no compile_graph + @brainstate.transform.vmap_new_states(state_tag="new", axis_size=inputs.shape[1]) + def init(): + brainstate.nn.init_all_states(model) + + init() + vmap_model = brainstate.nn.Vmap(model, vmap_states="new") def run_step(inp): - out = mapped_model(inp) + out = vmap_model(inp) loss = braintools.metric.softmax_cross_entropy_with_integer_labels( out, target_labels ).mean() diff --git a/examples/snn_models.py b/examples/snn_models.py index d33e5d48..c2fefdfb 100644 --- a/examples/snn_models.py +++ b/examples/snn_models.py @@ -143,8 +143,12 @@ def verify(self, input_spikes, num_show=5, sps_inc=10.): xs = np.transpose(input_spikes, (1, 0, 2)) # [n_steps, n_samples, n_in] # 运行仿真模型 - model = brainstate.nn.Map(self, init_map_size=xs.shape[1]) - model.init_all_states() + @brainstate.transform.vmap_new_states(state_tag='new', axis_size=xs.shape[1]) + def init(): + brainstate.nn.init_all_states(self) + + init() + model = brainstate.nn.Vmap(self, vmap_states='new') outs, sps, vs = brainstate.transform.for_loop( lambda x: (model(x), self.r.get_spike(), self.r.V.value), @@ -404,7 +408,7 @@ def batch_train(self, inputs, targets): model = braintrace.compile(self.target, braintrace.pp_prop, inputs[0], batch_size=inputs.shape[1], vmap=True, decay_or_rank=self.decay_or_rank) - model.show_graph() + model.module.show_graph() def _etrace_grad(inp): # call the model @@ -444,11 +448,14 @@ class BPTTTrainer(Trainer): def batch_train(self, inputs, targets): weights = self.target.states().subset(brainstate.ParamState) - # kept manual: BPTT baseline, with mapped per-sample states - model = brainstate.nn.Map( - self.target, init_map_size=inputs.shape[1] - ) - model.init_all_states() + # kept manual: BPTT baseline — no online algorithm to migrate + # initialize the states + @brainstate.transform.vmap_new_states(state_tag='new', axis_size=inputs.shape[1]) + def init(): + brainstate.nn.init_all_states(self.target) + + init() + model = brainstate.nn.Vmap(self.target, vmap_states='new') # the model for a single step def _run_step_train(inp): @@ -514,8 +521,12 @@ def update(self, spk): @brainstate.transform.jit(static_argnums=0) def eval(self, xs): - model = brainstate.nn.Map(self, init_map_size=xs.shape[1]) - model.init_all_states() + @brainstate.transform.vmap_new_states(state_tag='new', axis_size=xs.shape[1]) + def init(): + brainstate.nn.init_all_states(self) + + init() + model = brainstate.nn.Vmap(self, vmap_states='new') outs, sps, vs = brainstate.transform.for_loop( lambda x: (model(x), self.neu.get_spike(), self.neu.V.value), xs diff --git a/examples/tests/test_compile_modes.py b/examples/tests/test_compile_modes.py index 417f80ec..ebd0c690 100644 --- a/examples/tests/test_compile_modes.py +++ b/examples/tests/test_compile_modes.py @@ -1,7 +1,7 @@ # Copyright 2026 BrainX Ecosystem Limited. Licensed under the Apache License, 2.0. """Verify the example SNN cells compile and run under BOTH ``braintrace.compile(vmap=False)`` (batched, internal batch primitive) and -``braintrace.compile(vmap=True)`` (``brainstate.nn.Map`` state ownership). +``braintrace.compile(vmap=True)`` (per-sample vmap lanes). The custom ``GIF`` neuron in ``snn_models.py`` originally defined ``init_state(self)`` without ``batch_size``, so the non-vmap path @@ -88,19 +88,3 @@ def test_gif_neuron_init_state_accepts_batch_size(): brainstate.nn.init_all_states(neu, batch_size=B) assert neu.V.value.shape == (B, N_REC) assert neu.I2.value.shape == (B, N_REC) - - -def test_mapped_compile_exposes_show_graph_directly(): - """Mapped compile results expose reports without a wrapper ``.module``.""" - xs = jnp.zeros((B, N_IN)) - learner = braintrace.compile( - braintrace.nn.GRUCell(N_IN, N_REC), braintrace.D_RTRL, xs, - batch_size=B, vmap=True, - ) - - report = learner.show_graph(verbose=False, return_msg=True) - - assert isinstance(report, str) - assert "model.module.show_graph()" not in ( - EXAMPLES_DIR / "snn_models.py" - ).read_text(encoding="utf-8")