From 928219b16695c6b669d05fe8b1c3c64764c95b0b Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 11:08:40 +0800 Subject: [PATCH 01/13] Remove OTTT / OSTTP / OTPE; add the algorithm-axes roadmap braintrace is a framework for online learning in brain simulation, so it should ship general mechanisms rather than rules that work for one operator type and one model shape. These three failed that bar: - All whitelisted dense-matmul primitives (_SUPPORTED_PRIMITIVES) and raised NotImplementedError for lora / sparse / conv / element-wise relations. - All were single-step only. - OTPE additionally assumed a single global time constant, was feed-forward only, was gradient-exact for one hidden layer, and rejected num_state > 1 outright -- ruling out ALIF and any adaptation variable. Its own docstring called its derivation "narrower than OTTT's". - OSTTP bound B_list to the HiddenGroup count and threaded y_target through a bespoke path. Removing the classes does not remove the capability: docs/specs/ now carries the roadmap for decomposing the algorithm space into five orthogonal axes (trace_factorization x temporal_recursion x learning_signal x trace_filter x update_schedule), in which each removed rule is a reachable coordinate that works for every ETP primitive rather than only dense matmul. Coverage was repointed rather than dropped where the assertion was about a general property: the descent backstop now runs on pp_prop(rank=1) and EProp(feedback='random'), and the one-step D_RTRL equivalence tests now use EProp(kappa_filter_decay=0). public_api_test gained a guard that the removed names stay gone. The direction-alignment metric helpers are retained as the basis of the planned benchmark suite. Also removed: PresynapticTrace (used only by OTTT) and extract_y_target (used only by OSTTP). The _get_update_aux side channel in vjp_base is kept -- it is the general per-call hook the planned modulatory learning signal needs. Verified: pytest braintrace/ -> 2062 passed, 1 skipped; mypy clean. --- AGENTS.md | 2 +- braintrace/__init__.py | 11 +- braintrace/__init___test.py | 29 - braintrace/_algorithm/__init__.py | 12 +- braintrace/_algorithm/_common.py | 85 +-- braintrace/_algorithm/_common_test.py | 25 - braintrace/_algorithm/oracle_models.py | 12 +- braintrace/_algorithm/osttp.py | 232 -------- braintrace/_algorithm/osttp_test.py | 243 --------- braintrace/_algorithm/otpe.py | 497 ------------------ braintrace/_algorithm/otpe_test.py | 423 --------------- braintrace/_algorithm/ottt.py | 357 ------------- braintrace/_algorithm/ottt_test.py | 327 ------------ .../tests/approx_correctness_test.py | 142 +---- .../tests/exact_correctness_test.py | 92 +--- .../_algorithm/tests/public_api_test.py | 20 +- .../tests/transform_correctness_test.py | 79 --- braintrace/_algorithm/vjp_base.py | 30 +- braintrace/_compile.py | 20 +- braintrace/_compile_test.py | 10 +- braintrace/_compiler/graph.py | 2 +- changelog.md | 30 ++ docs/apis/algorithms.rst | 20 +- docs/apis/index.rst | 2 +- .../2026-07-25-algorithm-axes-roadmap.md | 378 +++++++++++++ 25 files changed, 500 insertions(+), 2580 deletions(-) delete mode 100644 braintrace/_algorithm/osttp.py delete mode 100644 braintrace/_algorithm/osttp_test.py delete mode 100644 braintrace/_algorithm/otpe.py delete mode 100644 braintrace/_algorithm/otpe_test.py delete mode 100644 braintrace/_algorithm/ottt.py delete mode 100644 braintrace/_algorithm/ottt_test.py create mode 100644 docs/specs/2026-07-25-algorithm-axes-roadmap.md diff --git a/AGENTS.md b/AGENTS.md index 3e50abbb..2ae40054 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -45,7 +45,7 @@ connects parameters to the hidden states they influence. Layer 1 ETP operators Custom JAX primitives + per-primitive rule registries + user-facing ops Layer 2 ETP compiler Jaxpr analysis: find ETP primitives, connect parameters to hidden states Layer 3 Graph executor Forward pass + hidden→weight / hidden→hidden Jacobian computation -Layer 4 Algorithms Orchestrators (D-RTRL, pp_prop/ES-D-RTRL, EProp/OSTL/OTPE/OTTT/OSTTP) +Layer 4 Algorithms Orchestrators (D-RTRL, pp_prop/ES-D-RTRL, EProp/OSTL) ``` Dependency direction is strictly downward: operators know nothing of the diff --git a/braintrace/__init__.py b/braintrace/__init__.py index b6d167fa..79a704fe 100644 --- a/braintrace/__init__.py +++ b/braintrace/__init__.py @@ -41,8 +41,7 @@ hidden->weight / hidden->hidden Jacobian computations. 4. **Algorithms** — online-learning orchestrators: the exact algorithms :class:`D_RTRL` / :func:`pp_prop` / :class:`ES_D_RTRL`, and the SNN family - :class:`EProp`, :class:`OSTLRecurrent`, :class:`OSTLFeedforward`, - :class:`OTPE`, :class:`OTTT`, :class:`OSTTP`. + :class:`EProp`, :class:`OSTLRecurrent`, :class:`OSTLFeedforward`. The :mod:`braintrace.nn` subpackage provides ready-made ETP-wired layers (linear maps, convolutions, recurrent cells, read-outs). @@ -86,12 +85,8 @@ EProp, OSTLRecurrent, OSTLFeedforward, - OTPE, - OTTT, - OSTTP, FixedRandomFeedback, KappaFilter, - PresynapticTrace, ) from ._compiler import ( ControlFlowPolicy, @@ -220,12 +215,8 @@ 'EProp', 'OSTLRecurrent', 'OSTLFeedforward', - 'OTPE', - 'OTTT', - 'OSTTP', 'FixedRandomFeedback', 'KappaFilter', - 'PresynapticTrace', # errors 'NotSupportedError', diff --git a/braintrace/__init___test.py b/braintrace/__init___test.py index 749ff7e8..06f9a9da 100644 --- a/braintrace/__init___test.py +++ b/braintrace/__init___test.py @@ -87,35 +87,6 @@ def test_algorithm_is_usable_end_to_end(name): assert bool(jnp.all(jnp.isfinite(y))) -# --- Task 3: SNN algos requiring leak + OSTTP B_list contract ---------------- - -def test_ottt_otpe_require_explicit_leak(): - model = tanh_rnn(seed=0).factory() - brainstate.nn.init_all_states(model, batch_size=1) - with pytest.raises(TypeError): - braintrace.OTTT(model) # leak is keyword-only & required - with pytest.raises(TypeError): - braintrace.OTPE(model) - - -def test_ottt_otpe_construct_with_leak(): - for ctor in (lambda m: braintrace.OTTT(m, leak=0.9), - lambda m: braintrace.OTPE(m, leak=0.9)): - model = tanh_rnn(seed=0).factory() - brainstate.nn.init_all_states(model, batch_size=1) - algo = ctor(model) - assert algo is not None - - -def test_osttp_constructs_with_B_list(): - model = tanh_rnn(n_in=3, n_rec=4, seed=0).factory() - brainstate.nn.init_all_states(model, batch_size=1) - # OSTTP needs one random-feedback matrix per HiddenGroup; B.shape[1] == n_l (=4). - B_list = [jnp.eye(4, dtype='float32')] - algo = braintrace.OSTTP(model, B_list) - assert algo is not None - - # --- Task 4: dynamic weight assignment raises NotSupportedError -------------- def test_dynamic_weight_assignment_raises_not_supported(): diff --git a/braintrace/_algorithm/__init__.py b/braintrace/_algorithm/__init__.py index e1e2d543..2fa042a8 100644 --- a/braintrace/_algorithm/__init__.py +++ b/braintrace/_algorithm/__init__.py @@ -17,21 +17,17 @@ Groups the core ETrace infrastructure (``ETraceAlgorithm``, ``EligibilityTrace``, ``ETraceGraphExecutor``), VJP-based algorithms (D-RTRL, pp_prop / ES-D-RTRL), -and paper-faithful SNN algorithms (EProp, OSTLRecurrent, OSTLFeedforward, OTPE, -OTTT, OSTTP). +and paper-faithful SNN algorithms (EProp, OSTLRecurrent, OSTLFeedforward). """ from __future__ import annotations -from ._common import FixedRandomFeedback, KappaFilter, PresynapticTrace +from ._common import FixedRandomFeedback, KappaFilter from .base import ETraceAlgorithm, EligibilityTrace from .d_rtrl import D_RTRL from .e_prop import EProp from .graph_executor import ETraceGraphExecutor from .ostl import OSTLFeedforward, OSTLRecurrent -from .osttp import OSTTP -from .otpe import OTPE -from .ottt import OTTT from .io_dim_vjp import IODimVjpAlgorithm from .param_dim_vjp import ParamDimVjpAlgorithm from .pp_prop import ES_D_RTRL, pp_prop # ES_D_RTRL: back-compat alias @@ -55,10 +51,6 @@ 'EProp', 'OSTLRecurrent', 'OSTLFeedforward', - 'OTPE', - 'OTTT', - 'OSTTP', 'FixedRandomFeedback', 'KappaFilter', - 'PresynapticTrace', ] diff --git a/braintrace/_algorithm/_common.py b/braintrace/_algorithm/_common.py index 73394cd8..93539883 100644 --- a/braintrace/_algorithm/_common.py +++ b/braintrace/_algorithm/_common.py @@ -28,20 +28,18 @@ from braintrace._typing import ArrayLike, PyTree __all__ = [ - 'PresynapticTrace', 'KappaFilter', 'FixedRandomFeedback', - 'extract_y_target', ] class _ZeroResetState(brainstate.ShortTermState): """``ShortTermState`` that records its init shape/dtype and resets to zeros. - Both :class:`PresynapticTrace` and :class:`KappaFilter` are leaky scalar-rate - accumulators that share the same reset semantics: on ``reset_state`` they are - re-zeroed at the original shape, optionally with the leading dimension swapped - for ``batch_size`` (or prepended for scalar-shaped states). + :class:`KappaFilter` is a leaky scalar-rate accumulator built on these reset + semantics: on ``reset_state`` it is re-zeroed at the original shape, optionally + with the leading dimension swapped for ``batch_size`` (or prepended for + scalar-shaped states). """ def __init__(self, init_value: Any) -> None: @@ -71,65 +69,6 @@ def reset_state(self, batch_size: Optional[int] = None, **kwargs: Any) -> None: self.value = jnp.zeros(shape, dtype=self._init_dtype) -class PresynapticTrace(_ZeroResetState): - r"""Leaky presynaptic accumulator used by OTTT and OTPE-Approx. - - The trace accumulates the presynaptic input with a multiplicative decay, - following :math:`\hat{a} \leftarrow \lambda \cdot \hat{a} + x_t`. - - Parameters - ---------- - init_value : jax.Array - Initial value; also dictates the shape and dtype of the trace. - leak : float - Decay factor :math:`\lambda` in ``(0, 1)``. Pulled from the neuron's - membrane leak in SNN usage. - - Raises - ------ - ValueError - If ``leak`` is not strictly inside the open interval ``(0, 1)``. - - Examples - -------- - .. code-block:: python - - >>> import jax.numpy as jnp - >>> import braintrace - >>> trace = braintrace.PresynapticTrace(jnp.zeros(3), leak=0.5) - >>> out = trace.update(jnp.ones(3)) - >>> print(out) - [1. 1. 1.] - >>> out = trace.update(jnp.ones(3)) - >>> print(out) - [1.5 1.5 1.5] - """ - - __module__ = 'braintrace' - - def __init__(self, init_value: Any, leak: float) -> None: - super().__init__(init_value) - if not (0.0 < leak < 1.0): - raise ValueError(f'leak must be in (0, 1); got {leak}') - self.leak = float(leak) - - def update(self, x: ArrayLike) -> Any: - r"""Apply one accumulation step :math:`\hat{a} \leftarrow \lambda \cdot \hat{a} + x`. - - Parameters - ---------- - x : jax.Array - The new presynaptic input added to the decayed trace. - - Returns - ------- - jax.Array - The updated trace value. - """ - self.value = self.leak * self.value + x - return self.value - - class KappaFilter(_ZeroResetState): r"""Low-pass filter helper state. @@ -198,8 +137,9 @@ class FixedRandomFeedback: r"""Frozen random feedback matrix with a stop-gradient guard. The feedback matrix :math:`B \in \mathbb{R}^{n_{\mathrm{target}} \times n_{\mathrm{layer}}}` - is sampled once at construction and frozen via :func:`jax.lax.stop_gradient`. It is used by - OSTTP (per-HiddenGroup target projection) and EProp-random-feedback. + is sampled once at construction and frozen via :func:`jax.lax.stop_gradient`. It backs + EProp's random-feedback mode, and is the intended home for any rule that replaces the + symmetric learning signal with a fixed random projection. Parameters ---------- @@ -260,17 +200,6 @@ def project(self, y_target: Any) -> Any: return y_target @ self.B -def extract_y_target(args: tuple, *, index: int = -1) -> Optional[jax.Array]: - """Fetch the target tensor from a positional-args tuple. - - Returns ``None`` if ``args`` is empty. ``index`` defaults to the last position - (OSTTP's convention: ``algo.update(x, y_target)``). - """ - if not args: - return None - return args[index] - - def _reset_state_in_a_dict( state_dict: Dict[Any, brainstate.State], batch_size: Optional[int], diff --git a/braintrace/_algorithm/_common_test.py b/braintrace/_algorithm/_common_test.py index ab37c51c..59253ac1 100644 --- a/braintrace/_algorithm/_common_test.py +++ b/braintrace/_algorithm/_common_test.py @@ -25,7 +25,6 @@ from braintrace._algorithm._common import ( FixedRandomFeedback, KappaFilter, - PresynapticTrace, _batched_zeros_like, _extract_leaf, _reset_state_in_a_dict, @@ -33,24 +32,9 @@ _update_dict, _wrap_leaves_as_pytree, _zeros_like_batch_or_not, - extract_y_target, ) -class TestPresynapticTrace(unittest.TestCase): - def test_exponential_accumulation(self): - trace = PresynapticTrace(jnp.zeros((2, 3)), leak=0.9) - trace.update(jnp.ones((2, 3))) - trace.update(jnp.ones((2, 3))) - # After 2 ones: 0.9*(0.9*0 + 1) + 1 == 1.9 - assert jnp.allclose(trace.value, jnp.full((2, 3), 1.9)) - - def test_reset_to_zero(self): - trace = PresynapticTrace(jnp.ones((4,)), leak=0.5) - trace.reset_state() - assert jnp.allclose(trace.value, jnp.zeros((4,))) - - class TestKappaFilter(unittest.TestCase): def test_low_pass(self): flt = KappaFilter(jnp.zeros((3,)), kappa=0.8) @@ -84,15 +68,6 @@ def test_project_shapes(self): assert proj.shape == (3, 7) -class TestExtractYTarget(unittest.TestCase): - def test_absent_returns_none(self): - assert extract_y_target(()) is None - - def test_present_returns_value(self): - y = jnp.ones((5,)) - assert extract_y_target((jnp.zeros(3), y), index=1) is y - - # --------------------------------------------------------------------------- # _zeros_like_batch_or_not # --------------------------------------------------------------------------- diff --git a/braintrace/_algorithm/oracle_models.py b/braintrace/_algorithm/oracle_models.py index 36f5ba70..76e46f1f 100644 --- a/braintrace/_algorithm/oracle_models.py +++ b/braintrace/_algorithm/oracle_models.py @@ -71,9 +71,11 @@ def leaky_linear(n_in: int = 3, n_rec: int = 4, leak: float = 0.9, seed: int = 0 The recurrence ``h_t = leak * h_{t-1} + matmul(x_t, w)`` has hidden-to-hidden Jacobian ``leak * I`` exactly (no off-diagonal recurrent term). This is the - regime in which OTTT (which discards ``hid2hid_jac`` and assumes ``leak * I``) - is exact. ``w`` reaches every future hidden state through the leaky carry, so - it is a genuine ETP relation despite being an input projection. + degenerate regime in which rules that discard ``hid2hid_jac`` and assume a + scalar leak become exact, which makes it the reference model for the + ``scalar_leak`` temporal recursion. ``w`` reaches every future hidden state + through the leaky carry, so it is a genuine ETP relation despite being an + input projection. """ def factory(): @@ -139,8 +141,8 @@ def two_state_rnn(n_in: int = 3, n_rec: int = 3, seed: int = 0) -> ModelSpec: HiddenGroup with ``num_state == 2`` (an LIF+adaptation-like topology). ``v_t = 0.9 v + matmul(x, w) - 0.1 a``; ``a_t = 0.95 a + v``. ``w`` is the - single trainable ETP input weight. D_RTRL handles this exactly; OTTT/OTPE - reject it (their per-step rule assumes a single-state group). + single trainable ETP input weight. D_RTRL handles this exactly; any rule + whose per-step formulation assumes a single-state group cannot represent it. """ def factory(): diff --git a/braintrace/_algorithm/osttp.py b/braintrace/_algorithm/osttp.py deleted file mode 100644 index f8e246df..00000000 --- a/braintrace/_algorithm/osttp.py +++ /dev/null @@ -1,232 +0,0 @@ -# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -"""OSTTP — Online Spatio-Temporal Learning with Target Projection -(Ortner et al., 2023). - -OSTTP combines the OSTL / D-RTRL eligibility trace with a DRTP-style *target -projection*: instead of back-propagating :math:`\\partial \\mathcal{L}/\\partial -h` from the readout, each HiddenGroup receives a learning signal formed by a -fixed random projection of the task target, :math:`y^{*}\\,B_l`. This removes the -weight-transport requirement and the backward pass, so learning is forward-only. - -See :class:`OSTTP` for the mathematical formulation, references, and an example. -""" - -from __future__ import annotations - -from typing import Any, Optional, Sequence - -import brainstate -import jax -import jax.numpy as jnp - -from braintrace._typing import ArrayLike -from ._common import extract_y_target -from .param_dim_vjp import ParamDimVjpAlgorithm - -__all__ = ['OSTTP'] - - -class OSTTP(ParamDimVjpAlgorithm): - r"""Online Spatio-Temporal Learning with Target Projection. - - OSTTP reuses the OSTL / D-RTRL per-parameter eligibility trace but replaces - the back-propagated learning signal with a **direct random target - projection** (DRTP): - - .. math:: - - \boldsymbol{\epsilon}^t \approx \mathbf{D}^t\,\boldsymbol{\epsilon}^{t-1} - + \operatorname{diag}(\mathbf{D}_f^t)\otimes \mathbf{x}^t , - \qquad - L_l^t = y^{*\,t}\, B_l , - \qquad - \nabla_{W}\mathcal{L} = \sum_t L^t \circ \boldsymbol{\epsilon}^t , - - where :math:`y^{*\,t}` is the task target at time :math:`t`, :math:`B_l \in - \mathbb{R}^{n_\text{target}\times n_l}` is a fixed random feedback matrix for - HiddenGroup :math:`l` (frozen via ``stop_gradient``), :math:`\mathbf{D}^t` is - the hidden-to-hidden Jacobian, :math:`\mathbf{D}_f^t` the state-to-output - Jacobian, and :math:`\mathbf{x}^t` the presynaptic input. - - **How it works.** The eligibility trace carries the temporal credit exactly - as in :class:`~braintrace.OSTLRecurrent` ('with-H'), but the spatial credit normally - obtained by back-propagating :math:`\partial \mathcal{L}/\partial h` is - replaced by a frozen random projection of the target. Because the projection - matrices :math:`B_l` are fixed, there is no weight transport and no backward - pass — the rule is fully forward and update-unlocked in both space and time. - - Parameters - ---------- - model : brainstate.nn.Module - The SNN whose weights are trained online. - B_list : Sequence[jax.Array] - One feedback matrix per HiddenGroup, each of shape - ``(n_target, n_l)``. Frozen via ``stop_gradient`` at construction; the - count and trailing dimension are validated against the compiled graph. - target_timing : {'per-step', 'sequence-end'}, default 'per-step' - ``'per-step'`` requires ``y_target`` at every :meth:`update` call. - ``'sequence-end'`` zeros the learning signal on intermediate steps (the - trace still accumulates) and applies the projection only when - ``y_target`` is supplied. - name : str, optional - Name of the algorithm instance. - vjp_method, fast_solve - Forwarded verbatim to :class:`~braintrace.ParamDimVjpAlgorithm`. - - Raises - ------ - ValueError - If ``target_timing`` is invalid; if ``len(B_list)`` differs from the - number of HiddenGroups; if a matrix's trailing dimension does not match - its HiddenGroup width; or if ``target_timing='per-step'`` and - ``y_target`` is omitted from an :meth:`update` call. - - Examples - -------- - .. code-block:: python - - >>> import brainstate - >>> import jax - >>> import braintrace - >>> - >>> class Net(brainstate.nn.Module): - ... def __init__(self): - ... super().__init__() - ... self.cell = braintrace.nn.ValinaRNNCell(1, 20, activation='tanh') - ... self.out = braintrace.nn.Linear(20, 1) - ... def update(self, x): - ... return x >> self.cell >> self.out - >>> - >>> model = Net() - >>> x0 = brainstate.random.randn(1) - >>> # one (n_target, n_l) feedback matrix per HiddenGroup (here n_l = 20) - >>> B = jax.random.normal(jax.random.PRNGKey(0), (1, 20)) - >>> # ``compile`` initialises states + builds the trace graph in one call - >>> learner = braintrace.compile(model, braintrace.OSTTP, x0, B_list=[B]) - >>> y = learner.update(x0, y_target=brainstate.random.randn(1)) - - References - ---------- - .. [1] Ortner, T., Pes, L., Gentinetta, J., Frenkel, C., & Pantazi, A. - (2023). "Online Spatio-Temporal Learning with Target Projection." - *2023 IEEE 5th International Conference on Artificial Intelligence - Circuits and Systems (AICAS)*, 1-5. - https://doi.org/10.1109/AICAS57966.2023.10168623 (arXiv:2304.05124) - .. [2] Frenkel, C., Lefebvre, M., & Bol, D. (2021). "Learning Without - Feedback: Fixed Random Learning Signals Allow for Feedforward Training of - Deep Neural Networks" (DRTP). *Frontiers in Neuroscience*, 15, 629892. - https://doi.org/10.3389/fnins.2021.629892 - """ - - __module__ = 'braintrace' - - #: OSTTP carries the temporal credit exactly as OSTLRecurrent ('with-H'), so - #: the hidden-group transition must trace recurrent ETP mixing primitives and - #: extract the true per-position block-diagonal (bounded) Jacobian. - _include_recurrent_mixing = True - - def __init__( - self, - model: brainstate.nn.Module, - B_list: Sequence[jax.Array], - target_timing: str = 'per-step', - name: Optional[str] = None, - vjp_method: str = 'single-step', - fast_solve: bool = True, - **kwargs: Any, - ) -> None: - if target_timing not in ('per-step', 'sequence-end'): - raise ValueError( - f"target_timing must be 'per-step' or 'sequence-end'; got {target_timing!r}" - ) - super().__init__( - model, name=name, vjp_method=vjp_method, fast_solve=fast_solve, **kwargs - ) - self._B_list = tuple(jax.lax.stop_gradient(B) for B in B_list) - self.target_timing = target_timing - self._current_y_target: Optional[ArrayLike] = None - - def compile_graph(self, *args: Any) -> None: - super().compile_graph(*args) - n_groups = len(self.graph.hidden_groups) - if len(self._B_list) != n_groups: - raise ValueError( - f'B_list has {len(self._B_list)} entries but model has {n_groups} ' - f'HiddenGroup(s). One B matrix per HiddenGroup is required.' - ) - for B, group in zip(self._B_list, self.graph.hidden_groups): - n_l = int(group.varshape[-1]) - if B.shape[1] != n_l: - raise ValueError( - f'B_list[{group.index}].shape[1] == {B.shape[1]} but HiddenGroup ' - f'{group.index} has n_l={n_l}.' - ) - - def update(self, x: ArrayLike, y_target: ArrayLike | None = None) -> Any: - """Call ``super().update(x)`` after stashing ``y_target`` for ``_get_update_aux``. - - The stash is read back synchronously by :meth:`_get_update_aux` - (called at the very start of the base class's ``update()``, before - any custom_vjp tracing happens) and threaded from there into - ``_true_update_fun`` as a genuine argument. This is required because - outer transforms (e.g. ``brainstate.transform.grad``) may stage the - forward trace and only invoke the custom_vjp fwd/bwd rules *after* - this Python-level call has already returned -- by then any - instance-attribute stash read directly inside ``_update_fn_fwd`` or - ``_update_fn_bwd`` would already see ``None``. Threading the value as - a real traced argument sidesteps that timing problem entirely. - """ - if self.target_timing == 'per-step' and y_target is None: - raise ValueError( - "OSTTP(target_timing='per-step') requires y_target at every update() call." - ) - self._current_y_target = y_target - try: - return super().update(x) - finally: - self._current_y_target = None - - def _get_update_aux(self) -> Any: - """Supply ``y_target`` as the per-call auxiliary data (see base class).""" - return self._current_y_target - - def _compute_learning_signal(self, dl_autodiff: Any, args: Any) -> Any: - """Replace reverse-AD ``dL/dh`` with ``B_l @ y_target`` per HiddenGroup. - - ``args`` here is ``(x, y_target)`` -- the base class appends the - auxiliary data returned by :meth:`_get_update_aux` to the model's own - forward arguments before invoking this hook (see - ``ETraceVjpAlgorithm._update_fn_bwd``); :func:`extract_y_target` - pulls ``y_target`` from the trailing position. - """ - y_target = extract_y_target(args) - if y_target is None: - # target_timing='sequence-end' with no y_target: zero out so traces - # accumulate without emitting a weight update this step. - return [jnp.zeros_like(s) for s in dl_autodiff] - out = [] - for gid, s in enumerate(dl_autodiff): - B = self._B_list[gid] - projected = y_target @ B # (batch, n_l) - # Reshape projected into the autodiff signal shape (which has a - # trailing num_state axis appended by concat_hidden). - # s shape == (*varshape, num_state); projected shape == (*varshape,) - target_shape = s.shape - expanded = projected.reshape(target_shape[:-1] + (1,)) - # Broadcast across the num_state tail. - out.append(jnp.broadcast_to(expanded, target_shape).astype(s.dtype)) - return out diff --git a/braintrace/_algorithm/osttp_test.py b/braintrace/_algorithm/osttp_test.py deleted file mode 100644 index 55ef2d37..00000000 --- a/braintrace/_algorithm/osttp_test.py +++ /dev/null @@ -1,243 +0,0 @@ -# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -import unittest - -import brainstate -import jax -import jax.numpy as jnp - -import braintrace -from braintrace._algorithm.osttp import OSTTP - - -def _osttp_net(): - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState( - 0.1 * jax.random.normal(jax.random.PRNGKey(0), (3, 3)) - ) - self.h = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - self.h.value = jax.nn.tanh( - braintrace.matmul(self.h.value + x, self.w.value) - ) - return self.h.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -class TestOSTTPConstruction(unittest.TestCase): - def test_default_target_timing(self): - B_list = [0.1 * jax.random.normal(jax.random.PRNGKey(1), (5, 3))] - algo = OSTTP(_osttp_net(), B_list=B_list) - assert algo.target_timing == 'per-step' - - def test_invalid_timing_raises(self): - B_list = [0.1 * jax.random.normal(jax.random.PRNGKey(1), (5, 3))] - with self.assertRaises(ValueError): - OSTTP(_osttp_net(), B_list=B_list, target_timing='never') - - def test_missing_B_list_raises(self): - with self.assertRaises(TypeError): - OSTTP(_osttp_net()) - - def test_B_list_wrong_length_raises_on_compile(self): - B_list = [ - jax.random.normal(jax.random.PRNGKey(1), (5, 3)), - jax.random.normal(jax.random.PRNGKey(2), (5, 3)), - ] - net = _osttp_net() - algo = OSTTP(net, B_list=B_list) - x = jnp.ones((1, 3)) - with self.assertRaises(ValueError): - algo.compile_graph(x) - - def test_B_list_wrong_shape_raises_on_compile(self): - B_list = [jax.random.normal(jax.random.PRNGKey(1), (5, 99))] - net = _osttp_net() - algo = OSTTP(net, B_list=B_list) - x = jnp.ones((1, 3)) - with self.assertRaises(ValueError): - algo.compile_graph(x) - - -class TestOSTTPTargetProjection(unittest.TestCase): - def test_target_differs_from_symmetric(self): - key = jax.random.PRNGKey(7) - B = 0.1 * jax.random.normal(key, (4, 3)) - net = _osttp_net() - algo = OSTTP(net, B_list=[B]) - x = jnp.ones((1, 3)) - y = jnp.ones((1, 4)) - algo.compile_graph(x) - algo.init_etrace_state() - - def loss(x_): - out = algo.update(x_, y_target=y) - return (out ** 2).sum() - - grads, _ = brainstate.transform.grad( - loss, algo.param_states, return_value=True - )(x) - g_osttp = grads[next(iter(grads))] - - from braintrace._algorithm.param_dim_vjp import ParamDimVjpAlgorithm - net2 = _osttp_net() - algo2 = ParamDimVjpAlgorithm(net2) - algo2.compile_graph(x) - algo2.init_etrace_state() - - def loss2(x_): - out = algo2.update(x_) - return (out ** 2).sum() - - grads2, _ = brainstate.transform.grad( - loss2, algo2.param_states, return_value=True - )(x) - g_drtrl = grads2[next(iter(grads2))] - assert not jnp.allclose(g_osttp, g_drtrl, atol=1e-4) - - -class TestOSTTPLearningSignalNonzero(unittest.TestCase): - """Regression test for the zero-learning-signal bug (C4). - - OSTTP used to stash ``y_target`` on the instance during the forward - ``update()`` and clear it in a ``finally`` block before returning. The - learning-signal hook (``_compute_learning_signal``) only runs later, - inside the ``jax.custom_vjp`` backward pass -- by then the stash was - already ``None``, so the target-projected learning signal, and hence the - weight gradient, was identically zero regardless of ``y_target``. - """ - - def _grad_for_target(self, y_target: jnp.ndarray, B: jnp.ndarray) -> jnp.ndarray: - net = _osttp_net() - algo = OSTTP(net, B_list=[B]) - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - - def loss(x_): - out = algo.update(x_, y_target=y_target) - return (out ** 2).sum() - - grads, _ = brainstate.transform.grad( - loss, algo.param_states, return_value=True - )(x) - return grads[next(iter(grads))] - - def test_gradients_depend_on_target(self): - B = 0.1 * jax.random.normal(jax.random.PRNGKey(11), (4, 3)) - y1 = jnp.ones((1, 4)) - y2 = -5.0 * jnp.ones((1, 4)) - - g1 = self._grad_for_target(y1, B) - g2 = self._grad_for_target(y2, B) - - # Different targets must yield different weight gradients -- if the - # learning signal were identically zero (the bug), g1 == g2 (both - # all-zero). - assert not jnp.allclose(g1, g2, atol=1e-6) - assert not jnp.allclose(g1, jnp.zeros_like(g1), atol=1e-8) - assert not jnp.allclose(g2, jnp.zeros_like(g2), atol=1e-8) - - -class FakeLIF(brainstate.HiddenState): - def __init__(self, iv, leak): - super().__init__(iv) - self.leak = leak - - -def _toy_net(): - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState( - 0.1 * jax.random.normal(jax.random.PRNGKey(0), (3, 3)) - ) - self.v = FakeLIF(jnp.zeros((1, 3)), leak=0.9) - - def update(self, x): - self.v.value = jax.nn.tanh( - 0.9 * self.v.value + braintrace.matmul(x, self.w.value) - ) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _run(algo, n_steps=10, lr=0.05, y_target=None, pass_y=False): - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - - losses = [] - for _ in range(n_steps): - def loss_fn(x_): - out = algo.update(x_, y_target=y_target) if pass_y else algo.update(x_) - target = jnp.ones_like(out) - return ((out - target) ** 2).mean() - - grads, loss_val = brainstate.transform.grad( - loss_fn, algo.param_states, return_value=True - )(x) - for path, st in algo.param_states.items(): - st.value = st.value - lr * grads[path] - losses.append(float(loss_val)) - return losses - - -class TestSmokeLossDecreases(unittest.TestCase): - def test_osttp(self): - net = _toy_net() - B = [0.1 * jax.random.normal(jax.random.PRNGKey(9), (3, 3))] - algo = OSTTP(net, B_list=B) - y = jnp.ones((1, 3)) - losses = _run(algo, y_target=y, pass_y=True) - assert losses[-1] < losses[0] - - -def _docstring_net(): - """The exact ``Net`` model used in the ``OSTTP`` docstring example.""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.cell = braintrace.nn.ValinaRNNCell(1, 20, activation='tanh') - self.out = braintrace.nn.Linear(20, 1) - - def update(self, x): - return x >> self.cell >> self.out - - return Net() - - -def test_docstring_compile_example_runs(): - """Verify the runnable ``braintrace.compile`` example in ``OSTTP``'s docstring.""" - model = _docstring_net() - x0 = brainstate.random.randn(1) - # one (n_target, n_l) feedback matrix per HiddenGroup (here n_l = 20) - B = jax.random.normal(jax.random.PRNGKey(0), (1, 20)) - learner = braintrace.compile(model, braintrace.OSTTP, x0, B_list=[B]) - y = learner.update(x0, y_target=brainstate.random.randn(1)) - assert y.shape == (1,) - assert bool(jnp.all(jnp.isfinite(y))) - assert len(learner.graph.hidden_param_op_relations) >= 1 diff --git a/braintrace/_algorithm/otpe.py b/braintrace/_algorithm/otpe.py deleted file mode 100644 index 5c71d256..00000000 --- a/braintrace/_algorithm/otpe.py +++ /dev/null @@ -1,497 +0,0 @@ -# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -"""OTPE — Online Training with Postsynaptic Estimates (Summe et al., 2023). - -OTPE replaces RTRL's full Jacobian with a leaky-additive per-parameter -accumulator :math:`\\hat R \\leftarrow \\lambda\\,\\hat R + \\partial s/\\partial -\\theta_\\text{local}` that estimates how a parameter's influence persists in the -postsynaptic membrane across several time steps — temporal structure that the -single-step approximations OTTT and OSTL drop. Cross-layer coupling is handled -inside ``_solve_weight_gradients`` without relaxing the compiler's "no W→W→h" -invariant. - -See :class:`OTPE` for the mathematical formulation, references, and an example. -""" - -from __future__ import annotations - -import warnings -from typing import Any, Dict, Optional - -import brainstate -import jax.numpy as jnp - -from braintrace._op import etp_mm_p, etp_mv_p, is_batched_primitive -from braintrace._misc import etrace_df_key -from ._common import _route_grads_by_path, _update_dict -from braintrace._compiler import ControlFlowPolicy -from .vjp_base import ETraceVjpAlgorithm - -__all__ = ['OTPE'] - -# OTPE's per-parameter leaky trace only reduces correctly for dense matmul -# relations. LoRA, sparse, conv, and element-wise (no x_var) relations are -# excluded -- see the compile-time guard in `init_etrace_state`. -_SUPPORTED_PRIMITIVES = frozenset({etp_mm_p, etp_mv_p}) - - -class OTPE(ETraceVjpAlgorithm): - r"""Online Training with Postsynaptic Estimates for spiking networks. - - OTPE maintains a leaky, additive estimate :math:`\hat R^t` of each - parameter's accumulated influence on the postsynaptic state, then contracts - it with the learning signal :math:`L^t` to obtain the weight gradient: - - .. math:: - - \hat R^t = \lambda\,\hat R^{t-1} - + \frac{\partial s^t}{\partial \theta} - = \lambda\,\hat R^{t-1} - + x^t \otimes \operatorname{diag}(D_f^t) , - \qquad - \nabla_{W}\mathcal{L}^t = L^t \cdot \hat R^t , - - where :math:`x^t` is the presynaptic input, :math:`D_f^t` the - state-to-output Jacobian (surrogate gradient of the spike), :math:`\lambda - \in (0, 1)` the membrane leak, and :math:`L^t = \partial \mathcal{L}^t / - \partial s^t` the learning signal. The contraction runs over the output - dimension, leaving a gradient with the weight's shape. - - In the low-rank ``'approx'`` mode (**F-OTPE**) the estimate is factorized as - an outer product, reducing memory from :math:`O(I\cdot O)` to :math:`O(I+O)` - per layer: - - .. math:: - - \hat R^t \approx \hat z_\text{in}^t \otimes \bar g_\text{out}^t , - \quad - \hat z_\text{in}^t = \lambda\,\hat z_\text{in}^{t-1} + x^t , - \quad - \bar g_\text{out}^t = \lambda\,\bar g_\text{out}^{t-1} - + \operatorname{diag}(D_f^t) , - - with gradient :math:`\nabla_{W}\mathcal{L}^t - = \hat z_\text{in}^t \otimes (L^t \cdot \bar g_\text{out}^t)`. - - **How it works.** Unlike OTTT/OSTL, which assign temporal credit only within - the current layer's output, OTPE keeps a per-parameter trace that decays - with the membrane leak, approximating the *entire* temporal effect of a - weight on downstream activity while staying local to each layer. This - improves gradient alignment with BPTT in deep feed-forward SNNs at modest - extra cost. - - Parameters - ---------- - model : brainstate.nn.Module - The SNN whose weights are trained online. - mode : {'full', 'approx'}, default 'full' - ``'full'`` keeps the full ``(batch, I, O)`` estimate :math:`\hat R` per - layer. ``'approx'`` (F-OTPE) factorizes it as an outer product for - :math:`O(I+O)` memory; emits a :class:`UserWarning` when the network has - more than one HiddenGroup, because the factorization bias compounds with - depth. - leak : float - Decay factor :math:`\lambda \in (0, 1)`. **Required** — it must be - supplied explicitly and is never inferred from the model. :math:`\lambda` - is the membrane leak of the *postsynaptic* neuron whose influence is - being accumulated; auto-inferring it from ``model.states()`` silently - picks an arbitrary (often wrong) value on heterogeneous or - multi-population models, so the framework will not guess it. - trace_clip_abs : float, optional - Elementwise clip applied to :math:`\hat R` each step (full mode only). - ``None`` disables clipping. - name : str, optional - Name of the algorithm instance. - vjp_method : str, optional - Forwarded to the base algorithm. Only ``'single-step'`` is supported by - OTPE v1 -- its trace update and weight-gradient formulas are derived - one step at a time. ``vjp_method='multi-step'`` is rejected with - :class:`ValueError` at construction; multi-step *inputs* (i.e. calling - the compiled learner with :class:`~braintrace.MultiStepData`) raise - :class:`NotImplementedError` instead, at call time. - control_flow : ControlFlowPolicy, optional - Policy governing control-flow canonicalization (cond if-conversion, - scan unrolling, structured scan descent, ...) during graph - compilation. ``None`` (default) uses - :data:`~braintrace.DEFAULT_CONTROL_FLOW_POLICY`. - - Limitations - ----------- - OTPE's published derivation is **narrower than OTTT's**, and this - implementation is a *general operator* that will happily run far outside - that proven regime. The estimate :math:`\hat R` is built on the assumption - that the only temporal coupling of the postsynaptic state is the scalar - membrane leak, :math:`\partial U^t / \partial U^{t-1} = \lambda` — exactly - the leaky integrate-and-fire (LIF) recurrence. On top of that scalar-leak - assumption (inherited from OTTT), OTPE adds three further restrictions: - - 1. **A single global time constant.** One scalar :math:`\lambda` is shared by - every traced connection. Heterogeneous leaks across neurons or layers - break the estimate; ``leak`` is therefore a user-supplied global constant - and is never inferred from the model (see the ``leak`` parameter). - 2. **Feed-forward only.** The trace omits the hidden-to-hidden Jacobian, so - it is the *postsynaptic estimate* for feed-forward SNNs. Applying it to a - recurrent network silently drops the recurrent temporal credit. - 3. **Single-hidden-layer exactness.** The estimate is gradient-exact for one - hidden layer; with depth the per-layer factorization accumulates bias. - - The low-rank ``'approx'`` mode (**F-OTPE**) layers an additional - outer-product approximation on top, which is itself justified only under the - same linear-leak assumption; its bias compounds with network depth (hence - the :class:`UserWarning` for multi-group networks). - - Concretely, ``braintrace`` exposes OTPE as a generic ETP operator: it accepts - arbitrary ETP weights and hidden states, multi-layer stacks, recurrent - connectivity, and even non-spiking cells (e.g. a ``tanh`` RNN). All of these - *run* mechanically, but **the moment the model deviates from a feed-forward - LIF network with a single global scalar leak, the computed gradient leaves - the regime in which OTPE is proven correct** and should be treated as a - heuristic approximation rather than a faithful gradient estimate. The one - structural case that is rejected outright is a multi-state hidden group - (``num_state > 1``, e.g. ALIF with an adaptation variable): the leaky scalar - estimate cannot assign per-state credit, so :meth:`compile_graph` raises - rather than silently summing across states. - - Raises - ------ - ValueError - If ``mode`` is not ``'full'`` or ``'approx'``, if ``leak`` is not in - :math:`(0, 1)`, if ``vjp_method`` is not ``'single-step'``, if a - weight-to-hidden relation reaches more than one HiddenGroup (OTPE v1 - requires one-hop per-layer relations), or (at :meth:`compile_graph`) - if a trained connection projects into a hidden group with - ``num_state > 1``. - NotImplementedError - At :meth:`compile_graph`, if a trained connection is routed through a - primitive other than dense ``matmul`` (batched or unbatched) -- e.g. - LoRA, sparse, or convolutional relations, whose weight-gradient chain - rule does not reduce to OTPE's per-parameter leaky trace. - - Examples - -------- - .. code-block:: python - - >>> import brainstate - >>> import braintrace - >>> - >>> class Net(brainstate.nn.Module): - ... def __init__(self): - ... super().__init__() - ... self.cell = braintrace.nn.ValinaRNNCell(1, 20, activation='tanh') - ... self.out = braintrace.nn.Linear(20, 1) - ... def update(self, x): - ... return x >> self.cell >> self.out - >>> - >>> model = Net() - >>> x0 = brainstate.random.randn(1) - >>> # ``leak`` is the postsynaptic membrane leak and must be passed - >>> # explicitly; it is never inferred from the model. ``compile`` does the - >>> # state init + graph build in one call. - >>> learner = braintrace.compile(model, braintrace.OTPE, x0, mode='full', leak=0.9) - >>> y = learner(x0) - - References - ---------- - .. [1] Summe, T. M., Schaefer, C. J. S., & Joshi, S. (2023). "Estimating - Post-Synaptic Effects for Online Training of Feed-Forward SNNs." - *arXiv preprint* arXiv:2311.16151. https://arxiv.org/abs/2311.16151 - """ - - __module__ = 'braintrace' - - def __init__( - self, - model: brainstate.nn.Module, - mode: str = 'full', - *, - leak: float, - name: Optional[str] = None, - vjp_method: str = 'single-step', - trace_clip_abs: Optional[float] = None, - control_flow: Optional[ControlFlowPolicy] = None, - ) -> None: - if mode not in ('full', 'approx'): - raise ValueError(f"mode must be 'full' or 'approx'; got {mode!r}") - if not (0.0 < float(leak) < 1.0): - raise ValueError(f'leak must be in (0, 1); got {leak}') - if vjp_method != 'single-step': - raise ValueError( - f"OTPE v1 only supports vjp_method='single-step': its trace " - f"update and weight-gradient formulas are derived one step at " - f"a time and have no multi-step form; got " - f"vjp_method={vjp_method!r}." - ) - super().__init__(model, name=name, vjp_method=vjp_method, - control_flow=control_flow) - self.mode = mode - self.leak = float(leak) - self.trace_clip_abs = trace_clip_abs - self._R_hat: Dict[int, brainstate.ShortTermState] = {} - self._R_hat_x: Dict[int, brainstate.ShortTermState] = {} - self._R_hat_g: Dict[int, brainstate.ShortTermState] = {} - self._R_hat_bias: Dict[int, brainstate.ShortTermState] = {} - self._rid_is_batched: Dict[int, bool] = {} - - def compile_graph(self, *args: Any) -> None: - # `super().compile_graph()` builds the graph, calls `init_etrace_state`, - # and -- only once both succeed -- sets `self.is_compiled = True`. The - # validation below runs *after* that flag is already True, so any - # failure here must explicitly reset it (finding M6); otherwise a - # failed compile would leave `is_compiled` stuck True and silently - # short-circuit a later, valid `compile_graph` call (the base class's - # `if not self.is_compiled:` guard would treat it as already compiled). - super().compile_graph(*args) - try: - if self.mode == 'approx': - n_groups = len(self.graph.hidden_groups) - if n_groups > 1: - warnings.warn( - "OTPE(mode='approx') (F-OTPE) adds an extra outer-product " - "approximation whose bias compounds with network depth; " - "mode='full' reduces (but does not eliminate) this " - "depth-dependent bias -- see the Limitations section of " - "the class docstring.", - UserWarning, - ) - # Invariant: each relation maps to exactly one HiddenGroup in OTPE v1. - for rel in self.graph.hidden_param_op_relations: - if len(rel.hidden_groups) != 1: - raise ValueError( - f'OTPE requires per-layer one-hop weight-to-hidden relations; ' - f'found relation reaching {len(rel.hidden_groups)} groups.' - ) - # OTPE's derivation assumes a single scalar membrane state per neuron - # (the LIF case). A hidden group bundling several states (e.g. ALIF's - # membrane potential plus adaptation variable) cannot be handled: the - # leaky scalar estimate cannot assign per-state credit, and collapsing - # the num_state axis with a sum has no theoretical basis (see the - # *Limitations* section of the class docstring). - group = rel.hidden_groups[0] - if group.num_state > 1: - raise ValueError( - f'OTPE only supports hidden groups with num_state == 1 ' - f'(single-state LIF-like neurons), but a trained connection ' - f'projects into a group with num_state == {group.num_state}. ' - f'Multi-state neurons (e.g. ALIF with an adaptation variable) ' - f'are outside the regime where OTPE is derived.' - ) - except Exception: - self.is_compiled = False - raise - - def init_etrace_state(self, *args: Any, **kwargs: Any) -> None: - self._R_hat = {} - self._R_hat_x = {} - self._R_hat_g = {} - self._R_hat_bias = {} - self._rid_is_batched = {} - for rel in self.graph.hidden_param_op_relations: - if rel.primitive not in _SUPPORTED_PRIMITIVES: - raise NotImplementedError( - f'OTPE only supports dense matmul relations (etp_mm/etp_mv); ' - f'got a trained connection routed through primitive ' - f'{rel.primitive.name!r}. LoRA, sparse, convolutional, and ' - f'element-wise relations do not reduce to OTPE\'s ' - f'per-parameter leaky trace and are unsupported.' - ) - if rel.x_var is None: - # Unreachable given the primitive guard above (both etp_mm and - # etp_mv always carry an x_var); kept as an explicit, message- - # bearing guard rather than a bare assert (see finding N3). - raise ValueError( - f'OTPE requires a relation with an explicit presynaptic ' - f'input (x_var), but the relation for primitive ' - f'{rel.primitive.name!r} has none.' - ) - rid = id(rel.y_var) - in_shape = rel.x_var.aval.shape - out_shape = rel.y_var.aval.shape - batched = is_batched_primitive(rel.primitive) - self._rid_is_batched[rid] = batched - if self.mode == 'full': - # 'weight' always exists: the primitive guard above restricts - # relations to etp_mm/etp_mv, whose dense-matmul rules key - # trainable_vars with 'weight' (and, optionally, 'bias'). - weight_key = 'weight' - weight_var = rel.trainable_vars[weight_key] - weight_shape = weight_var.aval.shape - if batched: - shape = (in_shape[0], *weight_shape) - else: - shape = weight_shape - self._R_hat[rid] = brainstate.ShortTermState( - jnp.zeros(shape, dtype=jnp.float32) - ) - if 'bias' in rel.trainable_vars: - self._R_hat_bias[rid] = brainstate.ShortTermState( - jnp.zeros(out_shape, dtype=jnp.float32) - ) - else: - self._R_hat_x[rid] = brainstate.ShortTermState( - jnp.zeros(in_shape, dtype=jnp.float32) - ) - self._R_hat_g[rid] = brainstate.ShortTermState( - jnp.zeros(out_shape, dtype=jnp.float32) - ) - - def reset_state(self, batch_size: Optional[int] = None, **kwargs: Any) -> None: - self.running_index.value = 0 - - def _rezero(rid: int, state: Any) -> None: - shape = state.value.shape - if batch_size is not None and self._rid_is_batched.get(rid, False): - new_shape = (batch_size, *shape[1:]) - else: - # Unbatched relations (etp_mv) have no batch axis in their - # trace; `batch_size` does not apply to them, and reusing the - # stored shape as-is (rather than assuming shape[0] is a batch - # axis) avoids corrupting it -- see finding M4. - new_shape = shape - state.value = jnp.zeros(new_shape, dtype=state.value.dtype) - - for store in (self._R_hat, self._R_hat_x, self._R_hat_g, self._R_hat_bias): - for rid, r in store.items(): - _rezero(rid, r) - - def _get_etrace_data(self) -> Any: - if self.mode == 'full': - return ( - {rid: r.value for rid, r in self._R_hat.items()}, - {rid: r.value for rid, r in self._R_hat_bias.items()}, - ) - return ( - {rid: r.value for rid, r in self._R_hat_x.items()}, - {rid: r.value for rid, r in self._R_hat_g.items()}, - ) - - def _assign_etrace_data(self, vals: Any) -> None: - if self.mode == 'full': - vals_R, vals_bias = vals - for rid, v in vals_R.items(): - self._R_hat[rid].value = v - for rid, v in vals_bias.items(): - self._R_hat_bias[rid].value = v - else: - vals_x, vals_g = vals - for rid, v in vals_x.items(): - self._R_hat_x[rid].value = v - for rid, v in vals_g.items(): - self._R_hat_g[rid].value = v - - def _update_etrace_data( - self, running_index: Any, hist_vals: Any, - hid2weight_jac: Any, hid2hid_jac: Any, weight_vals: Any, input_is_multi_step: Any, - ) -> Any: - """``R_hat ← λ·R_hat + ∂s/∂θ_local``. Ignores ``hid2hid_jac``. - - The bias companion trace ``R_hat_bias`` follows the same recursion - with the presynaptic input dropped (bias's local Jacobian has no ``x`` - factor): ``R_hat_bias ← λ·R_hat_bias + df``, identical in form to - ``_R_hat_g`` (which the ``'approx'`` mode already maintains for its own - outer-product factorization). - """ - if input_is_multi_step: - raise NotImplementedError('OTPE v1 supports single-step only') - xs = hid2weight_jac[0] - dfs = hid2weight_jac[1] - - if self.mode == 'full': - hist_R, hist_bias = hist_vals - new_R = {} - new_bias = {} - for rel in self.graph.hidden_param_op_relations: - rid = id(rel.y_var) - group = rel.hidden_groups[0] - x = xs[id(rel.x_var)] - df = dfs[etrace_df_key(rel.y_var, group.index)] - df_proj = df.sum(axis=-1) - if is_batched_primitive(rel.primitive): - local = jnp.einsum('bi,bo->bio', x, df_proj) - else: - local = jnp.einsum('i,o->io', x, df_proj) - updated = self.leak * hist_R[rid] + local - if self.trace_clip_abs is not None: - updated = jnp.clip( - updated, -self.trace_clip_abs, self.trace_clip_abs - ) - new_R[rid] = updated - if rid in hist_bias: - new_bias[rid] = self.leak * hist_bias[rid] + df_proj - return (new_R, new_bias) - else: - new_Rx = {} - new_Rg = {} - hist_x, hist_g = hist_vals - for rel in self.graph.hidden_param_op_relations: - rid = id(rel.y_var) - group = rel.hidden_groups[0] - x = xs[id(rel.x_var)] - df = dfs[etrace_df_key(rel.y_var, group.index)].sum(axis=-1) - new_Rx[rid] = self.leak * hist_x[rid] + x - new_Rg[rid] = self.leak * hist_g[rid] + df - return (new_Rx, new_Rg) - - def _solve_weight_gradients( - self, running_index: Any, etrace_at_t: Any, dl_to_hidden_groups: Any, - weight_vals: Any, dl_to_nonetws_at_t: Any, dl_to_etws_at_t: Any, - ) -> Any: - dG = {path: None for path in self.param_states} - if self.mode == 'full': - R_map, Rbias_map = etrace_at_t - for rel in self.graph.hidden_param_op_relations: - rid = id(rel.y_var) - R = R_map[rid] - group = rel.hidden_groups[0] - L = dl_to_hidden_groups[group.index].sum(axis=-1) - per_key = {} - if is_batched_primitive(rel.primitive): - per_key['weight'] = jnp.einsum('bo,bio->io', L, R) - if 'bias' in rel.trainable_vars: - per_key['bias'] = jnp.einsum('bo,bo->o', L, Rbias_map[rid]) - else: - per_key['weight'] = jnp.einsum('o,io->io', L, R) - if 'bias' in rel.trainable_vars: - per_key['bias'] = L * Rbias_map[rid] - _route_grads_by_path(rel, per_key, weight_vals, dG) - else: - Rx_map, Rg_map = etrace_at_t - for rel in self.graph.hidden_param_op_relations: - rid = id(rel.y_var) - group = rel.hidden_groups[0] - L = dl_to_hidden_groups[group.index].sum(axis=-1) - Rx = Rx_map[rid] - Rg = Rg_map[rid] - Lg = L * Rg - per_key = {} - if is_batched_primitive(rel.primitive): - per_key['weight'] = jnp.einsum('bi,bo->io', Rx, Lg) - if 'bias' in rel.trainable_vars: - # Bias has no "in" dimension: sum the per-batch-row - # contributions directly (matches the 'b' contraction - # the weight einsum performs above). - per_key['bias'] = Lg.sum(axis=0) - else: - per_key['weight'] = jnp.einsum('i,o->io', Rx, Lg) - if 'bias' in rel.trainable_vars: - per_key['bias'] = Lg - _route_grads_by_path(rel, per_key, weight_vals, dG) - - for path, dg in dl_to_nonetws_at_t.items(): - _update_dict(dG, path, dg) - if dl_to_etws_at_t is not None: - for path, dg in dl_to_etws_at_t.items(): - _update_dict(dG, path, dg, error_when_no_key=True) - return dG diff --git a/braintrace/_algorithm/otpe_test.py b/braintrace/_algorithm/otpe_test.py deleted file mode 100644 index d7ee6ac3..00000000 --- a/braintrace/_algorithm/otpe_test.py +++ /dev/null @@ -1,423 +0,0 @@ -# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -import unittest - -import brainstate -import jax -import jax.numpy as jnp - -import braintrace -from braintrace._algorithm import oracle -from braintrace._algorithm.otpe import OTPE - - -def _otpe_net_single_layer(): - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState( - 0.1 * jax.random.normal(jax.random.PRNGKey(0), (3, 3)) - ) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - self.v.value = 0.9 * self.v.value + braintrace.matmul(x, self.w.value) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _two_state_net(): - """Net whose two coupled hidden states form one group with num_state == 2.""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState( - 0.1 * jax.random.normal(jax.random.PRNGKey(0), (3, 3)) - ) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - self.a = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - v, a = self.v.value, self.a.value - self.v.value = 0.9 * v + braintrace.matmul(x, self.w.value) - 0.1 * a - self.a.value = 0.95 * a + v - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _lora_net(): - """Net trained through `braintrace.lora_matmul` -- LoRA relations are outside - OTPE's supported primitive set (see finding H5).""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - brainstate.random.seed(0) - self.B = brainstate.ParamState(0.1 * brainstate.random.normal(size=(3, 2))) - self.A = brainstate.ParamState(0.1 * brainstate.random.normal(size=(2, 3))) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - self.v.value = 0.9 * self.v.value + braintrace.lora_matmul( - x, self.B.value, self.A.value - ) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _bias_net(): - """Net trained through a recurrent `braintrace.matmul(..., bias=...)`, whose - recurrence coefficient (0.9) matches the `leak` OTTT/OTPE are given -- the - "exactly diagonal" regime in which both algorithms claim BPTT-exact gradients.""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - brainstate.random.seed(0) - self.w = brainstate.ParamState(0.1 * brainstate.random.normal(size=(3, 3))) - self.b = brainstate.ParamState(jnp.zeros(3)) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - self.v.value = 0.9 * self.v.value + braintrace.matmul( - x, self.w.value, self.b.value - ) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _unbatched_net(): - """Net with an unbatched hidden state (no leading batch axis) -- exercises - OTPE's `reset_state` on `etp_mv` relations (see finding M4).""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - brainstate.random.seed(0) - self.w = brainstate.ParamState(0.1 * brainstate.random.normal(size=(3, 3))) - self.v = brainstate.HiddenState(jnp.zeros((3,))) - - def update(self, x): - self.v.value = 0.9 * self.v.value + braintrace.matmul(x, self.w.value) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net) # no batch_size -> unbatched - return net - - -class TestOTPEConstruction(unittest.TestCase): - def test_default_mode_full(self): - algo = OTPE(_otpe_net_single_layer(), leak=0.9) - assert algo.mode == 'full' - - def test_invalid_mode_raises(self): - with self.assertRaises(ValueError): - OTPE(_otpe_net_single_layer(), mode='bogus', leak=0.9) - - def test_leak_is_required(self): - """leak is never inferred from the model; omitting it is a TypeError.""" - with self.assertRaises(TypeError): - OTPE(_otpe_net_single_layer()) - - def test_leak_out_of_range_raises(self): - with self.assertRaises(ValueError): - OTPE(_otpe_net_single_layer(), leak=1.5) - - def test_multi_step_vjp_method_rejected_at_construction(self): - """N6: OTPE's trace update/weight-gradient rules are derived one step at - a time and have no multi-step form; `vjp_method='multi-step'` must be - rejected eagerly at construction, not silently accepted until the first - (multi-step) call fails.""" - with self.assertRaises(ValueError): - OTPE(_otpe_net_single_layer(), leak=0.9, vjp_method='multi-step') - - def test_num_state_gt_one_raises(self): - """Multi-state hidden groups are outside OTPE's LIF regime. - - M6: unlike the `NotImplementedError` in - `TestOTPEUnsupportedRelationGuard` (which fires inside - `init_etrace_state`, *before* `base.compile_graph` ever sets - `is_compiled = True`), this `ValueError` fires in OTPE's own - post-`super().compile_graph()` validation -- i.e. *after* the base - class has already set `is_compiled = True`. That validation's - try/except must explicitly reset the flag on failure, or a failed - compile would leave `is_compiled` stuck `True` and the base class's - `if not self.is_compiled:` guard would silently skip re-validation - (and `init_etrace_state`) on a subsequent `compile_graph` call.""" - algo = OTPE(_two_state_net(), leak=0.9) - x = jnp.ones((1, 3)) - with self.assertRaises(ValueError): - algo.compile_graph(x) - assert algo.is_compiled is False - - # A second call on the same (permanently invalid) model must raise - # again, not silently pass because `is_compiled` was left True. - with self.assertRaises(ValueError): - algo.compile_graph(x) - assert algo.is_compiled is False - - def test_compile_allocates_R_hat(self): - algo = OTPE(_otpe_net_single_layer(), leak=0.9) - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - assert len(algo._R_hat) == len(algo.graph.hidden_param_op_relations) - - -class TestOTPESingleLayer(unittest.TestCase): - def test_update_runs_and_produces_gradients(self): - net = _otpe_net_single_layer() - algo = OTPE(net, leak=0.9, mode='full') - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - - def loss(x_): - return (algo.update(x_) ** 2).sum() - - grads, _ = brainstate.transform.grad( - loss, algo.param_states, return_value=True - )(x) - g = grads[next(iter(grads))] - assert g.shape == (3, 3) - assert jnp.any(g != 0.0) - - -class TestOTPEApproxMode(unittest.TestCase): - def test_approx_uses_factored_traces(self): - net = _otpe_net_single_layer() - algo = OTPE(net, mode='approx', leak=0.9) - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - assert len(algo._R_hat_x) == len(algo.graph.hidden_param_op_relations) - assert len(algo._R_hat_g) == len(algo.graph.hidden_param_op_relations) - - def test_approx_runs(self): - net = _otpe_net_single_layer() - algo = OTPE(net, mode='approx', leak=0.9) - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - - def loss(x_): - return (algo.update(x_) ** 2).sum() - - grads, _ = brainstate.transform.grad( - loss, algo.param_states, return_value=True - )(x) - g = grads[next(iter(grads))] - assert g.shape == (3, 3) - - -class TestOTPEBiasGradient(unittest.TestCase): - """H5: the weight-gradient solvers indexed only the ``'weight'`` key of each - relation's trainable dict, so `braintrace.matmul(x, w, bias=b)` silently got - a zero bias gradient. Must route every key (here also ``'bias'``) through - `relation.trainable_paths`, exactly as `_common._route_grads_by_path` does. - - OTPE's docstring claims gradient-exactness for a single hidden layer under - the scalar-leak (LIF) assumption. In the "exactly diagonal" regime -- the - model's recurrence coefficient equals the `leak` OTPE is given -- both - 'full' and 'approx' modes are exact for bias specifically: bias's local - Jacobian dy/db=1 has no "in" dimension, so 'approx' mode's outer-product - factorization (which only approximates the joint x-times-df weight trace) - introduces no approximation error for bias -- its ``R_hat_g``/``R_hat_bias`` - companion trace follows the identical recursion in both modes. So we assert - *element-wise* equality with the BPTT oracle for bias in both modes. - """ - - def test_bias_gradient_matches_bptt_full_mode(self): - inputs = jnp.stack([jnp.ones((1, 3)) * (i + 1) * 0.1 for i in range(4)]) - bptt = oracle.bptt_param_gradients(_bias_net, inputs) - online = oracle.online_param_gradients_singlestep_naive( - _bias_net, inputs, algo_factory=lambda m: OTPE(m, leak=0.9, mode='full') - ) - assert bool(jnp.any(online[('b',)] != 0.0)), 'bias gradient must not be zero (H5)' - oracle.assert_param_gradients_close( - online, bptt, atol=1e-4, keys=[('b',), ('w',)] - ) - - def test_bias_gradient_matches_bptt_approx_mode(self): - inputs = jnp.stack([jnp.ones((1, 3)) * (i + 1) * 0.1 for i in range(4)]) - bptt = oracle.bptt_param_gradients(_bias_net, inputs) - online = oracle.online_param_gradients_singlestep_naive( - _bias_net, inputs, algo_factory=lambda m: OTPE(m, leak=0.9, mode='approx') - ) - assert bool(jnp.any(online[('b',)] != 0.0)), 'bias gradient must not be zero (H5)' - # Only the bias key is asserted exact here -- 'approx' mode's weight - # gradient is *not* BPTT-exact (its outer-product factorization is a - # genuine additional approximation), covered separately in - # TestOTPEApproxMode. - oracle.assert_param_gradients_close(online, bptt, atol=1e-4, keys=[('b',)]) - - -class TestOTPEUnsupportedRelationGuard(unittest.TestCase): - """H5: LoRA/conv/sparse relations don't reduce to OTPE's per-parameter - leaky trace; they must raise a compile-time `NotImplementedError` naming - the offending primitive rather than crashing opaquely or silently - mishandling the relation.""" - - def test_lora_relation_raises_named_primitive_at_compile_time(self): - algo = OTPE(_lora_net(), leak=0.9, mode='full') - with self.assertRaises(NotImplementedError) as ctx: - algo.compile_graph(jnp.ones((1, 3))) - assert 'etp_lora_mm' in str(ctx.exception) - - def test_is_compiled_stays_false_after_repeated_guard_failure(self): - """M6: a failed compile-time validation must not leave `is_compiled` - stuck True -- otherwise `base.compile_graph`'s ``if not - self.is_compiled:`` guard would silently skip re-validation (and - `init_etrace_state`) on a subsequent call. Calling `compile_graph` - twice on the same (permanently invalid) model must raise both times.""" - algo = OTPE(_lora_net(), leak=0.9, mode='full') - x = jnp.ones((1, 3)) - with self.assertRaises(NotImplementedError): - algo.compile_graph(x) - assert algo.is_compiled is False - - with self.assertRaises(NotImplementedError): - algo.compile_graph(x) - assert algo.is_compiled is False - - # A fresh, valid model must still compile normally afterward. - ok_algo = OTPE(_otpe_net_single_layer(), leak=0.9, mode='full') - ok_algo.compile_graph(x) - assert ok_algo.is_compiled is True - - -class TestOTPEUnbatchedResetState(unittest.TestCase): - """M4: `OTPE.reset_state` assumed a leading batch axis and corrupted - unbatched trace shapes (e.g. replacing the real leading `in`/`out` dim with - `batch_size`). Must derive reset shapes from the stored trace values - instead of assuming a batch axis is present.""" - - def test_reset_preserves_unbatched_trace_shapes_full_mode(self): - net = _unbatched_net() - algo = OTPE(net, leak=0.9, mode='full') - x = jnp.ones((3,)) - algo.compile_graph(x) - algo.init_etrace_state() - algo.update(x) # populate traces with real (non-init) values first - - before = {rid: r.value.shape for rid, r in algo._R_hat.items()} - algo.reset_state(batch_size=4) # unbatched relations must ignore this - after = {rid: r.value.shape for rid, r in algo._R_hat.items()} - assert before == after - assert all(v.value.shape == (3, 3) for v in algo._R_hat.values()) - - def test_reset_preserves_unbatched_trace_shapes_approx_mode(self): - net = _unbatched_net() - algo = OTPE(net, leak=0.9, mode='approx') - x = jnp.ones((3,)) - algo.compile_graph(x) - algo.init_etrace_state() - algo.update(x) - - before_x = {rid: r.value.shape for rid, r in algo._R_hat_x.items()} - before_g = {rid: r.value.shape for rid, r in algo._R_hat_g.items()} - algo.reset_state(batch_size=4) - after_x = {rid: r.value.shape for rid, r in algo._R_hat_x.items()} - after_g = {rid: r.value.shape for rid, r in algo._R_hat_g.items()} - assert before_x == after_x - assert before_g == after_g - - -def _toy_net(): - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState( - 0.1 * jax.random.normal(jax.random.PRNGKey(0), (3, 3)) - ) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - self.v.value = jax.nn.tanh( - 0.9 * self.v.value + braintrace.matmul(x, self.w.value) - ) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _run(algo, n_steps=10, lr=0.05, y_target=None, pass_y=False): - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - - losses = [] - for _ in range(n_steps): - def loss_fn(x_): - out = algo.update(x_, y_target=y_target) if pass_y else algo.update(x_) - target = jnp.ones_like(out) - return ((out - target) ** 2).mean() - - grads, loss_val = brainstate.transform.grad( - loss_fn, algo.param_states, return_value=True - )(x) - for path, st in algo.param_states.items(): - st.value = st.value - lr * grads[path] - losses.append(float(loss_val)) - return losses - - -class TestSmokeLossDecreases(unittest.TestCase): - def test_otpe(self): - losses = _run(OTPE(_toy_net(), leak=0.9)) - assert losses[-1] < losses[0] - - -def _docstring_net(): - """The exact ``Net`` model used in the ``OTPE`` docstring example.""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.cell = braintrace.nn.ValinaRNNCell(1, 20, activation='tanh') - self.out = braintrace.nn.Linear(20, 1) - - def update(self, x): - return x >> self.cell >> self.out - - return Net() - - -def test_docstring_compile_example_runs(): - """Verify the runnable ``braintrace.compile`` example in ``OTPE``'s docstring.""" - model = _docstring_net() - x0 = brainstate.random.randn(1) - learner = braintrace.compile(model, braintrace.OTPE, x0, mode='full', leak=0.9) - y = learner(x0) - assert y.shape == (1,) - assert bool(jnp.all(jnp.isfinite(y))) - assert len(learner.graph.hidden_param_op_relations) >= 1 diff --git a/braintrace/_algorithm/ottt.py b/braintrace/_algorithm/ottt.py deleted file mode 100644 index af751777..00000000 --- a/braintrace/_algorithm/ottt.py +++ /dev/null @@ -1,357 +0,0 @@ -# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -"""OTTT — Online Training Through Time (Xiao et al., 2022). - -OTTT is derived from BPTT but discards the hidden-to-hidden recurrent Jacobian, -so it keeps only a leaky presynaptic eligibility trace -:math:`\\hat a^t \\leftarrow \\lambda\\,\\hat a^{t-1} + x^t` and forms the weight -gradient at each step as the outer product of that trace with the (instantaneous) -learning signal, :math:`\\Delta W = \\hat a^t \\otimes (L \\cdot \\sigma'(u))`. -This yields constant training memory, independent of the number of time steps. - -See :class:`OTTT` for the mathematical formulation, references, and an example. -""" - -from __future__ import annotations - -from typing import Any, Dict, Optional - -import brainstate -import jax.numpy as jnp - -from braintrace._op import etp_mm_p, etp_mv_p, is_batched_primitive -from ._common import PresynapticTrace, _route_grads_by_path, _update_dict -from braintrace._compiler import ControlFlowPolicy -from .vjp_base import ETraceVjpAlgorithm - -__all__ = ['OTTT'] - -# OTTT's math (outer product of a single scalar/vector presynaptic trace with the -# learning signal) only reduces correctly for dense matmul relations. LoRA, -# sparse, conv, and element-wise (no x_var) relations are excluded -- see the -# compile-time guard in `init_etrace_state`. -_SUPPORTED_PRIMITIVES = frozenset({etp_mm_p, etp_mv_p}) - -# Sentinel key for the bias companion trace inside the etrace-vals dict. -# ``rid`` keys are ``id(...)`` values, which CPython guarantees are -# non-negative, so a negative int can never collide with one. This must stay -# an ``int`` (not e.g. a string): the etrace-vals dict is a JAX pytree that -# flows through `grad`/`jit` tracing, and dict pytree flattening sorts keys -- -# sorting a dict with mixed ``str``/``int`` keys raises `TypeError`. -_BIAS_KEY = -1 - - -class OTTT(ETraceVjpAlgorithm): - r"""Online Training Through Time for spiking neural networks. - - OTTT tracks only a leaky **presynaptic trace** and forms the weight gradient - each step as an outer product with the local learning signal: - - .. math:: - - \hat{a}^t = - \begin{cases} - \lambda\,\hat{a}^{t-1} + x^t & \text{(mode='A', accumulated)} \\ - x^t & \text{(mode='O', instantaneous)} - \end{cases} - - .. math:: - - \nabla_{W}\mathcal{L}^t - = \hat{a}^t \otimes - \Big( \frac{\partial \mathcal{L}^t}{\partial s^t}\,\sigma'(u^t) \Big) - \;=\; \hat{a}^t \otimes L^t , - - where :math:`x^t` is the presynaptic input, :math:`u^t` the membrane - potential, :math:`s^t = \sigma(u^t)` the (surrogate) spike, :math:`\sigma'` - the surrogate-gradient function, :math:`\lambda \in (0, 1)` the membrane - leak, and :math:`L^t` the learning signal already propagated through the - spike nonlinearity. - - **How it works.** Starting from BPTT, OTTT keeps the spatial credit - assignment but **drops the hidden-to-hidden recurrent Jacobian**. The only - state it carries forward in time is the rank-1 presynaptic trace - :math:`\hat{a}^t`, so the per-step gradient is the outer product of that - trace with the instantaneous learning signal. Training memory is therefore - :math:`O(B \cdot I)` per layer and **independent of the sequence length** — - the cheapest of the algorithms here, at the cost of ignoring longer-range - temporal credit. - - Parameters - ---------- - model : brainstate.nn.Module - The SNN whose weights are trained online. - mode : {'A', 'O'}, default 'A' - ``'A'`` accumulates the presynaptic trace over time - (:math:`\hat a \leftarrow \lambda\,\hat a + x`). ``'O'`` uses the - instantaneous presynaptic spike only (:math:`\hat a := x^t`). - leak : float - Presynaptic leak :math:`\lambda \in (0, 1)`. **Required** — it must be - supplied explicitly and is never inferred from the model (see - *Limitations*). Mathematically :math:`\lambda` is the membrane leak of - the *postsynaptic* neuron whose trace is being accumulated. - name : str, optional - Name of the algorithm instance. - vjp_method : str, optional - Forwarded to the base algorithm. Only ``'single-step'`` is supported by - OTTT v1 -- its trace update and weight-gradient formulas are derived - one step at a time. ``vjp_method='multi-step'`` is rejected with - :class:`ValueError` at construction; multi-step *inputs* (i.e. calling - the compiled learner with :class:`~braintrace.MultiStepData`) raise - :class:`NotImplementedError` instead, at call time. - control_flow : ControlFlowPolicy, optional - Policy governing control-flow canonicalization (cond if-conversion, - scan unrolling, structured scan descent, ...) during graph - compilation. ``None`` (default) uses - :data:`~braintrace.DEFAULT_CONTROL_FLOW_POLICY`. - - Limitations - ----------- - - **The leak must be supplied by the user.** OTTT does *not* try to read - :math:`\lambda` off the model's neuron states. A previous version walked - ``model.states()`` and took the first state exposing a ``leak`` attribute, - but on heterogeneous or multi-population models that silently picks an - arbitrary (often wrong) value — e.g. the leak of the *presynaptic* layer, - a readout filter, or whichever population happens to be enumerated first. - Since :math:`\lambda` is, by the derivation, the membrane leak of the - postsynaptic neuron of each trained connection, the framework cannot - guess it safely. A single network with different leaks per layer therefore - cannot be trained correctly with one global ``leak`` and is unsupported. - - **Single-state hidden groups only.** Each trained connection must project - into a :class:`HiddenGroup` with ``num_state == 1``. The weight gradient - contracts the learning signal ``L`` (shape ``(*varshape, num_state)``) - down to ``(*varshape,)``; collapsing a ``num_state > 1`` tail (e.g. an - ALIF neuron carrying both membrane potential and an adaptation variable) - has no theoretical justification — the trace is a single leaky scalar and - cannot disentangle per-state credit — so OTTT raises at compile time - instead of silently summing across states. - - **Single-step inputs only** (OTTT v1); multi-step inputs raise - :class:`NotImplementedError`. - - Raises - ------ - ValueError - If ``mode`` is not ``'A'`` or ``'O'``, if ``leak`` is not in - :math:`(0, 1)`, if ``vjp_method`` is not ``'single-step'``, or (at - :meth:`compile_graph`) if a trained connection projects into a hidden - group with ``num_state > 1``. - NotImplementedError - At :meth:`compile_graph`, if a trained connection is routed through a - primitive other than dense ``matmul`` (batched or unbatched) -- e.g. - LoRA, sparse, or convolutional relations, whose weight-gradient chain - rule does not reduce to OTTT's single presynaptic-trace outer product. - - Examples - -------- - .. code-block:: python - - >>> import brainstate - >>> import braintrace - >>> - >>> class Net(brainstate.nn.Module): - ... def __init__(self): - ... super().__init__() - ... self.cell = braintrace.nn.ValinaRNNCell(1, 20, activation='tanh') - ... self.out = braintrace.nn.Linear(20, 1) - ... def update(self, x): - ... return x >> self.cell >> self.out - >>> - >>> model = Net() - >>> x0 = brainstate.random.randn(1) - >>> # ``leak`` is the postsynaptic membrane leak and must be passed - >>> # explicitly; it is never inferred from the model. ``compile`` does the - >>> # state init + graph build in one call. - >>> learner = braintrace.compile(model, braintrace.OTTT, x0, mode='A', leak=0.9) - >>> y = learner(x0) - - References - ---------- - .. [1] Xiao, M., Meng, Q., Zhang, Z., He, D., & Lin, Z. (2022). "Online - Training Through Time for Spiking Neural Networks." *Advances in Neural - Information Processing Systems (NeurIPS)* 35. - https://arxiv.org/abs/2210.04195 - """ - - __module__ = 'braintrace' - - def __init__( - self, - model: brainstate.nn.Module, - mode: str = 'A', - *, - leak: float, - name: Optional[str] = None, - vjp_method: str = 'single-step', - control_flow: Optional[ControlFlowPolicy] = None, - ) -> None: - if mode not in ('A', 'O'): - raise ValueError(f"mode must be 'A' or 'O'; got {mode!r}") - if not (0.0 < float(leak) < 1.0): - raise ValueError(f'leak must be in (0, 1); got {leak}') - if vjp_method != 'single-step': - raise ValueError( - f"OTTT v1 only supports vjp_method='single-step': its trace " - f"update and weight-gradient formulas are derived one step at " - f"a time and have no multi-step form; got " - f"vjp_method={vjp_method!r}." - ) - super().__init__(model, name=name, vjp_method=vjp_method, - control_flow=control_flow) - self.mode = mode - self.leak = float(leak) - self._pre_traces: Dict[int, PresynapticTrace] = {} - self._bias_trace: Optional[brainstate.ShortTermState] = None - - def init_etrace_state(self, *args: Any, **kwargs: Any) -> None: - self._pre_traces = {} - has_bias = False - for rel in self.graph.hidden_param_op_relations: - if rel.primitive not in _SUPPORTED_PRIMITIVES: - raise NotImplementedError( - f'OTTT only supports dense matmul relations (etp_mm/etp_mv); ' - f'got a trained connection routed through primitive ' - f'{rel.primitive.name!r}. LoRA, sparse, convolutional, and ' - f'element-wise relations do not reduce to a single ' - f'presynaptic-trace outer product and are unsupported.' - ) - for group in rel.hidden_groups: - if group.num_state > 1: - raise ValueError( - f'OTTT only supports hidden groups with num_state == 1, ' - f'but a trained connection projects into a group with ' - f'num_state == {group.num_state}. Collapsing the learning ' - f'signal across multiple hidden states (e.g. an ALIF ' - f'neuron with membrane potential plus an adaptation ' - f'variable) has no theoretical basis for OTTT; the leaky ' - f'scalar presynaptic trace cannot assign per-state credit.' - ) - if 'bias' in rel.trainable_vars: - has_bias = True - rid = id(rel.x_var) - if rid in self._pre_traces: - continue - if rel.x_var is None: - # Unreachable given the primitive guard above (both etp_mm and - # etp_mv always carry an x_var); kept as an explicit, message- - # bearing guard rather than a bare assert (see finding N3). - raise ValueError( - f'OTTT requires a relation with an explicit presynaptic ' - f'input (x_var), but the relation for primitive ' - f'{rel.primitive.name!r} has none.' - ) - shape = rel.x_var.aval.shape - dtype = rel.x_var.aval.dtype - self._pre_traces[rid] = PresynapticTrace( - jnp.zeros(shape, dtype=dtype), leak=self.leak - ) - self._bias_trace = ( - brainstate.ShortTermState(jnp.zeros((), dtype=jnp.float32)) - if has_bias else None - ) - - def reset_state(self, batch_size: Optional[int] = None, **kwargs: Any) -> None: - self.running_index.value = 0 - for t in self._pre_traces.values(): - t.reset_state(batch_size=batch_size) - if self._bias_trace is not None: - # The bias companion trace has no batch axis: bias's local - # Jacobian dy/db = 1 is identical for every batch row, so it is - # always a scalar regardless of `batch_size`. - self._bias_trace.value = jnp.zeros((), dtype=self._bias_trace.value.dtype) # type: ignore[attr-defined] - - def _get_etrace_data(self) -> Any: - d = {rid: t.value for rid, t in self._pre_traces.items()} - if self._bias_trace is not None: - d[_BIAS_KEY] = self._bias_trace.value - return d - - def _assign_etrace_data(self, vals: Any) -> None: - for rid, t in self._pre_traces.items(): - t.value = vals[rid] - if self._bias_trace is not None: - self._bias_trace.value = vals[_BIAS_KEY] - - def _update_etrace_data( - self, running_index: Any, hist_vals: Any, - hid2weight_jac: Any, hid2hid_jac: Any, weight_vals: Any, input_is_multi_step: Any, - ) -> Any: - """``â ← λ·â + x_t`` (mode='A') or ``â := x_t`` (mode='O'). - - The bias companion trace ``ĉ`` follows the identical recursion with the - presynaptic input replaced by ``1`` (bias's local Jacobian ``dy/db``): - ``ĉ ← λ·ĉ + 1`` (mode='A') or ``ĉ := 1`` (mode='O'). It is a single - scalar shared by every relation, since ``self.leak`` is one global - constant for the whole algorithm instance. - - Ignores ``hid2hid_jac`` — OTTT's core approximation. - """ - if input_is_multi_step: - raise NotImplementedError('OTTT v1 supports single-step only') - xs_at_t = hid2weight_jac[0] - - new_vals = {} - for rid, old in hist_vals.items(): - if rid == _BIAS_KEY: - continue - x_t = xs_at_t[rid] - if self.mode == 'A': - new_vals[rid] = self.leak * old + x_t - else: - new_vals[rid] = x_t - if self._bias_trace is not None: - old_c = hist_vals[_BIAS_KEY] - if self.mode == 'A': - new_vals[_BIAS_KEY] = self.leak * old_c + 1.0 - else: - new_vals[_BIAS_KEY] = jnp.ones_like(old_c) - return new_vals - - def _solve_weight_gradients( - self, running_index: Any, etrace_at_t: Any, dl_to_hidden_groups: Any, - weight_vals: Any, dl_to_nonetws_at_t: Any, dl_to_etws_at_t: Any, - ) -> Any: - """``ΔW = outer(â, L)``; ``Δb = ĉ · L`` (bias has no "in" dimension), where - ``L`` is the (already σ'-propagated) signal.""" - dG = {path: None for path in self.param_states} - c_hat = etrace_at_t.get(_BIAS_KEY) if self._bias_trace is not None else None - for rel in self.graph.hidden_param_op_relations: - a_hat = etrace_at_t[id(rel.x_var)] - for group in rel.hidden_groups: - L = dl_to_hidden_groups[group.index] - # L shape = (*varshape, num_state); num_state == 1 is enforced at - # compile time (see init_etrace_state), so this drops the singleton - # tail rather than summing across genuinely distinct hidden states. - L_proj = L.sum(axis=-1) - per_key = {} - if is_batched_primitive(rel.primitive): - # a_hat: (batch, in), L_proj: (batch, out). ΔW: (in, out) - per_key['weight'] = jnp.einsum('bi,bo->io', a_hat, L_proj) - if 'bias' in rel.trainable_vars: - # c_hat is batch-independent (dy/db = 1 for every row); - # contributions from every batch row simply sum. - per_key['bias'] = c_hat * L_proj.sum(axis=0) - else: - per_key['weight'] = jnp.einsum('i,o->io', a_hat, L_proj) - if 'bias' in rel.trainable_vars: - per_key['bias'] = c_hat * L_proj - _route_grads_by_path(rel, per_key, weight_vals, dG) - - for path, dg in dl_to_nonetws_at_t.items(): - _update_dict(dG, path, dg) - if dl_to_etws_at_t is not None: - for path, dg in dl_to_etws_at_t.items(): - _update_dict(dG, path, dg, error_when_no_key=True) - return dG diff --git a/braintrace/_algorithm/ottt_test.py b/braintrace/_algorithm/ottt_test.py deleted file mode 100644 index 5119dd63..00000000 --- a/braintrace/_algorithm/ottt_test.py +++ /dev/null @@ -1,327 +0,0 @@ -# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -import unittest - -import brainstate -import jax -import jax.numpy as jnp - -import braintrace -from braintrace._algorithm import oracle -from braintrace._algorithm.ottt import OTTT - - -def _ottt_net(): - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState( - 0.1 * jax.random.normal(jax.random.PRNGKey(0), (3, 3)) - ) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - self.v.value = 0.9 * self.v.value + braintrace.matmul(x, self.w.value) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _two_state_net(): - """Net whose two coupled hidden states form one group with num_state == 2.""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState( - 0.1 * jax.random.normal(jax.random.PRNGKey(0), (3, 3)) - ) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - self.a = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - v, a = self.v.value, self.a.value - self.v.value = 0.9 * v + braintrace.matmul(x, self.w.value) - 0.1 * a - self.a.value = 0.95 * a + v - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _lora_net(): - """Net trained through `braintrace.lora_matmul` -- LoRA relations are outside - OTTT's supported primitive set (see finding H5).""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - brainstate.random.seed(0) - self.B = brainstate.ParamState(0.1 * brainstate.random.normal(size=(3, 2))) - self.A = brainstate.ParamState(0.1 * brainstate.random.normal(size=(2, 3))) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - self.v.value = 0.9 * self.v.value + braintrace.lora_matmul( - x, self.B.value, self.A.value - ) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _bias_net(): - """Net trained through a recurrent `braintrace.matmul(..., bias=...)`, whose - recurrence coefficient (0.9) matches the `leak` OTTT/OTPE are given -- the - "exactly diagonal" regime in which both algorithms claim BPTT-exact gradients.""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - brainstate.random.seed(0) - self.w = brainstate.ParamState(0.1 * brainstate.random.normal(size=(3, 3))) - self.b = brainstate.ParamState(jnp.zeros(3)) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - self.v.value = 0.9 * self.v.value + braintrace.matmul( - x, self.w.value, self.b.value - ) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -class TestOTTTConstruction(unittest.TestCase): - def test_default_mode_is_A(self): - algo = OTTT(_ottt_net(), leak=0.9) - assert algo.mode == 'A' - - def test_invalid_mode_raises(self): - with self.assertRaises(ValueError): - OTTT(_ottt_net(), mode='bogus', leak=0.9) - - def test_leak_is_required(self): - """leak is never inferred from the model; omitting it is a TypeError.""" - with self.assertRaises(TypeError): - OTTT(_ottt_net()) - - def test_leak_out_of_range_raises(self): - with self.assertRaises(ValueError): - OTTT(_ottt_net(), leak=1.5) - - def test_multi_step_vjp_method_rejected_at_construction(self): - """N6: OTTT's trace update/weight-gradient rules are derived one step at - a time and have no multi-step form; `vjp_method='multi-step'` must be - rejected eagerly at construction, not silently accepted until the first - (multi-step) call fails.""" - with self.assertRaises(ValueError): - OTTT(_ottt_net(), leak=0.9, vjp_method='multi-step') - - def test_num_state_gt_one_raises(self): - """Multi-state hidden groups have no theoretical basis for OTTT.""" - algo = OTTT(_two_state_net(), leak=0.9) - with self.assertRaises(ValueError): - algo.compile_graph(jnp.ones((1, 3))) - - def test_compile_allocates_presynaptic_traces(self): - algo = OTTT(_ottt_net(), leak=0.9) - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - assert len(algo._pre_traces) == len(algo.graph.hidden_param_op_relations) - - -class TestOTTTEnd2End(unittest.TestCase): - def test_update_runs_and_produces_gradients(self): - net = _ottt_net() - algo = OTTT(net, leak=0.9) - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - - def loss(x_): - return (algo.update(x_) ** 2).sum() - - grads, _ = brainstate.transform.grad( - loss, algo.param_states, return_value=True - )(x) - g = grads[next(iter(grads))] - assert g.shape == (3, 3) - assert jnp.any(g != 0.0) - - def test_mode_O_differs_from_mode_A(self): - def compute(mode): - net = _ottt_net() - algo = OTTT(net, mode=mode, leak=0.9) - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - algo.update(x) - - def loss(x_): - return (algo.update(x_) ** 2).sum() - - grads, _ = brainstate.transform.grad( - loss, algo.param_states, return_value=True - )(x) - return grads[next(iter(grads))] - - g_A = compute('A') - g_O = compute('O') - assert not jnp.allclose(g_A, g_O, atol=1e-4) - - -class TestOTTTBiasGradient(unittest.TestCase): - """H5: the weight-gradient solver indexed only the ``'weight'`` key of each - relation's trainable dict, so `braintrace.matmul(x, w, bias=b)` silently got - a zero bias gradient. Must route every key (here also ``'bias'``) through - `relation.trainable_paths`, exactly as `_common._route_grads_by_path` does. - - OTTT's docstring claims it keeps BPTT's *spatial* credit assignment exactly - and only drops the *temporal* (hidden-to-hidden) Jacobian. In the "exactly - diagonal" regime -- the model's recurrence coefficient equals the `leak` - OTTT is given, mode='A' -- this is provably exact: both the outer-product - weight term and the scalar bias term reduce, via summation-order exchange, - to the same double sum BPTT computes. So we assert *element-wise* equality - with the BPTT oracle, not just non-zero/same-sign -- this model is squarely - in that exact regime. - """ - - def test_bias_gradient_matches_bptt_in_exactly_diagonal_regime(self): - inputs = jnp.stack([jnp.ones((1, 3)) * (i + 1) * 0.1 for i in range(4)]) - bptt = oracle.bptt_param_gradients(_bias_net, inputs) - online = oracle.online_param_gradients_singlestep_naive( - _bias_net, inputs, algo_factory=lambda m: OTTT(m, leak=0.9) - ) - assert bool(jnp.any(online[('b',)] != 0.0)), 'bias gradient must not be zero (H5)' - oracle.assert_param_gradients_close( - online, bptt, atol=1e-4, keys=[('b',), ('w',)] - ) - - -class TestOTTTUnsupportedRelationGuard(unittest.TestCase): - """H5: LoRA/conv/sparse relations don't reduce to OTTT's single - presynaptic-trace outer product; they must raise a compile-time - `NotImplementedError` naming the offending primitive rather than crashing - opaquely or silently mishandling the relation.""" - - def test_lora_relation_raises_named_primitive_at_compile_time(self): - algo = OTTT(_lora_net(), leak=0.9) - with self.assertRaises(NotImplementedError) as ctx: - algo.compile_graph(jnp.ones((1, 3))) - assert 'etp_lora_mm' in str(ctx.exception) - - def test_is_compiled_stays_false_after_repeated_guard_failure(self): - """M6: a failed compile-time validation must not leave `is_compiled` - stuck True -- otherwise `base.compile_graph`'s ``if not - self.is_compiled:`` guard would silently skip re-validation (and - `init_etrace_state`) on a subsequent call. Calling `compile_graph` - twice on the same (permanently invalid) model must raise both times.""" - algo = OTTT(_lora_net(), leak=0.9) - x = jnp.ones((1, 3)) - with self.assertRaises(NotImplementedError): - algo.compile_graph(x) - assert algo.is_compiled is False - - with self.assertRaises(NotImplementedError): - algo.compile_graph(x) - assert algo.is_compiled is False - - # A fresh, valid model must still compile normally afterward. - ok_algo = OTTT(_ottt_net(), leak=0.9) - ok_algo.compile_graph(x) - assert ok_algo.is_compiled is True - - -def _toy_net(): - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState( - 0.1 * jax.random.normal(jax.random.PRNGKey(0), (3, 3)) - ) - self.v = brainstate.HiddenState(jnp.zeros((1, 3))) - - def update(self, x): - self.v.value = jax.nn.tanh( - 0.9 * self.v.value + braintrace.matmul(x, self.w.value) - ) - return self.v.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _run(algo, n_steps=10, lr=0.05, y_target=None, pass_y=False): - x = jnp.ones((1, 3)) - algo.compile_graph(x) - algo.init_etrace_state() - - losses = [] - for _ in range(n_steps): - def loss_fn(x_): - out = algo.update(x_, y_target=y_target) if pass_y else algo.update(x_) - target = jnp.ones_like(out) - return ((out - target) ** 2).mean() - - grads, loss_val = brainstate.transform.grad( - loss_fn, algo.param_states, return_value=True - )(x) - for path, st in algo.param_states.items(): - st.value = st.value - lr * grads[path] - losses.append(float(loss_val)) - return losses - - -class TestSmokeLossDecreases(unittest.TestCase): - def test_ottt(self): - losses = _run(OTTT(_toy_net(), leak=0.9)) - assert losses[-1] < losses[0] - - -def _docstring_net(): - """The exact ``Net`` model used in the ``OTTT`` docstring example.""" - - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.cell = braintrace.nn.ValinaRNNCell(1, 20, activation='tanh') - self.out = braintrace.nn.Linear(20, 1) - - def update(self, x): - return x >> self.cell >> self.out - - return Net() - - -def test_docstring_compile_example_runs(): - """Verify the runnable ``braintrace.compile`` example in ``OTTT``'s docstring.""" - model = _docstring_net() - x0 = brainstate.random.randn(1) - learner = braintrace.compile(model, braintrace.OTTT, x0, mode='A', leak=0.9) - y = learner(x0) - assert y.shape == (1,) - assert bool(jnp.all(jnp.isfinite(y))) - assert len(learner.graph.hidden_param_op_relations) >= 1 diff --git a/braintrace/_algorithm/tests/approx_correctness_test.py b/braintrace/_algorithm/tests/approx_correctness_test.py index b0b52636..403152b3 100644 --- a/braintrace/_algorithm/tests/approx_correctness_test.py +++ b/braintrace/_algorithm/tests/approx_correctness_test.py @@ -13,11 +13,14 @@ # limitations under the License. # ============================================================================== -"""L3-C approximate-class correctness: OTTT/OTPE gradients are direction-aligned -with BPTT (cosine/sign), exact in their degenerate regime, and their known -approximation biases (F-07/F-08/F-09) are quantified. Multi-state guards -(F-01/F-04) and the rate-model approximation-exactness ceiling (F-21/F-22) are -pinned.""" +"""L3-C approximate-class correctness: direction-alignment metrics against BPTT +(cosine / sign / relative magnitude), the multi-state (num_state == 2) exactness +guard for D_RTRL, and the rate-model approximation-exactness ceiling +(F-21/F-22). + +The OTTT/OTPE-specific bias tests (F-01/F-04/F-07/F-08/F-09) were removed in +0.3.0 along with those algorithms; see docs/specs for the roadmap. The metric +helpers they exercised are retained and are the basis of the benchmark suite.""" import brainstate import jax @@ -32,19 +35,15 @@ bptt_param_gradients, cosine_similarity, online_param_gradients, - online_param_gradients_singlestep_naive, relative_magnitude, sign_agreement, ) from braintrace._algorithm.oracle_models import ( - leaky_linear, tanh_rnn, two_state_rnn, ) ATOL_BPTT = 1e-4 -MIN_COSINE = 0.95 -MIN_SIGN = 0.99 def _inputs(T, n_in, seed=42): @@ -110,104 +109,6 @@ def test_two_state_rnn_forms_one_group_num_state_two(): assert int(algo.graph.hidden_groups[0].varshape[-1]) == 3 # n_rec -# --- Task 4: C-level direction agreement (tanh_rnn, OTTT approximate regime) -- - -# Single-step-only approximate algorithms; their total-sequence gradient comes -# from the naive per-step accumulation (direction-aligned with BPTT; measured -# cosine: OTTT(A) 0.986, OTPE(full) 0.983; sign agreement 1.000). -_C_LEVEL_ALGOS = { - 'OTTT_A': lambda m: braintrace.OTTT(m, mode='A', leak=0.9), - 'OTPE_full': lambda m: braintrace.OTPE(m, mode='full', leak=0.9), -} - - -@pytest.mark.parametrize('algo_name', list(_C_LEVEL_ALGOS)) -def test_ottt_otpe_direction_aligned_with_bptt_on_tanh_rnn(algo_name): - """Approximate algorithms must point the same way as BPTT (cosine/sign), - even though they do not match it element-wise.""" - spec = tanh_rnn(n_in=3, n_rec=4, seed=0) - inputs = _inputs(6, 3) - g_bptt = bptt_param_gradients(spec.factory, inputs) - g_approx = online_param_gradients_singlestep_naive( - spec.factory, inputs, algo_factory=_C_LEVEL_ALGOS[algo_name]) - assert_direction_aligned( - g_approx, g_bptt, min_cosine=MIN_COSINE, min_sign_agreement=MIN_SIGN, - keys=list(spec.etp_param_keys)) - - -# --- Task 5: F-09 OTTT bias scales with ||hid2hid - leak*I|| ------------------ - -def test_ottt_is_exact_on_leaky_linear(): - """leaky_linear has hid2hid Jacobian == leak*I exactly, so OTTT (which assumes - exactly that) reproduces BPTT: cosine ~ 1 and relative magnitude ~ 1.""" - spec = leaky_linear(n_in=3, n_rec=4, leak=0.9, seed=0) - inputs = _inputs(6, 3) - g_bptt = bptt_param_gradients(spec.factory, inputs) - g_ottt = online_param_gradients_singlestep_naive( - spec.factory, inputs, algo_factory=lambda m: braintrace.OTTT(m, mode='A', leak=0.9)) - assert_direction_aligned( - g_ottt, g_bptt, min_cosine=0.999, mag_bounds=(0.98, 1.02), - keys=list(spec.etp_param_keys)) - - -def test_ottt_is_biased_but_aligned_on_tanh_rnn_F09(): - """F-09: on tanh_rnn the hidden Jacobian is NOT leak*I (it carries the tanh - derivative and the recurrent weight), so OTTT is no longer exact — its - magnitude is inflated (measured relmag ~1.2, outside the exact band) while - direction stays aligned (cosine ~0.986).""" - spec = tanh_rnn(n_in=3, n_rec=4, seed=0) - inputs = _inputs(6, 3) - g_bptt = bptt_param_gradients(spec.factory, inputs) - g_ottt = online_param_gradients_singlestep_naive( - spec.factory, inputs, algo_factory=lambda m: braintrace.OTTT(m, mode='A', leak=0.9)) - key = spec.etp_param_keys[0] - # Direction is still good ... - assert cosine_similarity(g_ottt[key], g_bptt[key]) >= MIN_COSINE - # ... but it is demonstrably NOT exact (magnitude inflated beyond the band that - # holds on leaky_linear). This is the F-09 bias. - assert relative_magnitude(g_ottt[key], g_bptt[key]) > 1.05 - - -# --- Task 6: F-07 OTPE(mode='approx') magnitude inflation -------------------- - -def test_otpe_approx_is_directionally_ok_but_magnitude_inflated_F07(): - """F-07: OTPE's rank-1 'approx' mode keeps the gradient direction (cosine - ~0.985) but grossly inflates its magnitude (measured relmag ~5.6 on tanh_rnn, - ~3.8 on leaky_linear). We pin: aligned direction AND relmag > 2.""" - for spec in (tanh_rnn(n_in=3, n_rec=4, seed=0), - leaky_linear(n_in=3, n_rec=4, leak=0.9, seed=0)): - inputs = _inputs(6, 3) - g_bptt = bptt_param_gradients(spec.factory, inputs) - g_approx = online_param_gradients_singlestep_naive( - spec.factory, inputs, - algo_factory=lambda m: braintrace.OTPE(m, mode='approx', leak=0.9)) - key = spec.etp_param_keys[0] - assert cosine_similarity(g_approx[key], g_bptt[key]) >= MIN_COSINE - assert relative_magnitude(g_approx[key], g_bptt[key]) > 2.0 - - -# --- Task 7: F-08 trace_clip_abs biases magnitude ---------------------------- - -def test_trace_clip_abs_shrinks_gradient_magnitude_F08(): - """F-08: trace_clip_abs clamps the eligibility trace with no principled bound. - A small clip (0.01) shrinks the OTPE gradient ~10x (measured relmag 0.096) and - degrades its direction (cosine 0.983 -> 0.936), while clip=None keeps relmag - near 1. We pin the magnitude collapse under aggressive clipping.""" - spec = tanh_rnn(n_in=3, n_rec=4, seed=0) - inputs = _inputs(6, 3) - g_bptt = bptt_param_gradients(spec.factory, inputs) - key = spec.etp_param_keys[0] - - g_noclip = online_param_gradients_singlestep_naive( - spec.factory, inputs, algo_factory=lambda m: braintrace.OTPE(m, mode='full', leak=0.9)) - g_clip = online_param_gradients_singlestep_naive( - spec.factory, inputs, - algo_factory=lambda m: braintrace.OTPE(m, mode='full', leak=0.9, trace_clip_abs=0.01)) - - assert relative_magnitude(g_noclip[key], g_bptt[key]) > 0.8 # near-unbiased - assert relative_magnitude(g_clip[key], g_bptt[key]) < 0.3 # collapsed - - # --- Task 8: F-01/F-04 multi-state (num_state == 2) --------------------------- def test_d_rtrl_exact_on_two_state_group(): @@ -221,22 +122,6 @@ def test_d_rtrl_exact_on_two_state_group(): assert_param_gradients_close(g_online, g_bptt, atol=ATOL_BPTT) -@pytest.mark.parametrize('algo_factory', [ - lambda m: braintrace.OTTT(m, mode='A', leak=0.9), - lambda m: braintrace.OTPE(m, mode='full', leak=0.9), -], ids=['OTTT', 'OTPE']) -def test_ottt_otpe_reject_multistate_group_F01(algo_factory): - """F-01/F-04: OTTT and OTPE assume a single-state group; on a num_state==2 - group they raise at compile (their per-state collapse has no theoretical - basis here). Pinning this prevents silently-wrong multi-state gradients.""" - spec = two_state_rnn(n_in=3, n_rec=3, seed=0) - model = spec.factory() - brainstate.nn.init_all_states(model, batch_size=1) - algo = algo_factory(model) - with pytest.raises(ValueError, match='num_state == 1'): - algo.compile_graph(jnp.ones((1, 3), dtype='float32')) - - # --- Task 9: F-21/F-22 IODim/EProp approximations are exact on rate models ----- # On a single-relation rate model the IODim rank / ES decay / random-feedback @@ -279,8 +164,8 @@ def test_approximations_diverge_on_snn_multipopulation_DEFERRED(): # --- Task 10: C-level convergence backstop (loss decreases) ------------------ def _train_loss_trajectory(algo, n_steps=12, lr=0.05): - """Manual SGD on an MSE-to-ones target (the ottt_test pattern), returning the - per-step loss. A working approximate gradient must drive the loss down.""" + """Manual SGD on an MSE-to-ones target, returning the per-step loss. A working + approximate gradient must drive the loss down.""" x = jnp.ones((1, 3), dtype='float32') algo.compile_graph(x) algo.init_etrace_state() @@ -298,9 +183,10 @@ def loss_fn(x_): @pytest.mark.parametrize('algo_factory', [ - lambda m: braintrace.OTTT(m, mode='A', leak=0.9), - lambda m: braintrace.OTPE(m, mode='full', leak=0.9), -], ids=['OTTT_A', 'OTPE_full']) + lambda m: braintrace.pp_prop(m, decay_or_rank=1), + lambda m: braintrace.EProp( + m, feedback='random', random_feedback_key=jax.random.PRNGKey(7)), +], ids=['pp_prop_rank1', 'EProp_random']) def test_approximate_algorithm_descends_loss(algo_factory): """C-level backstop: the approximate gradient is a usable descent direction — training loss at the end is below the start.""" diff --git a/braintrace/_algorithm/tests/exact_correctness_test.py b/braintrace/_algorithm/tests/exact_correctness_test.py index ad9f65ff..ab1d70f1 100644 --- a/braintrace/_algorithm/tests/exact_correctness_test.py +++ b/braintrace/_algorithm/tests/exact_correctness_test.py @@ -14,9 +14,12 @@ # ============================================================================== """L3-A/B exact-class correctness: multi-step online gradients reproduce BPTT -element-wise; cross-algorithm reduction identities hold; single-step-only -OTTT/OTPE match D_RTRL instantaneously. Findings F-19/F-20 pin the cases the -current oracle cannot validate.""" +element-wise, cross-algorithm reduction identities hold, and the vjp_method +boundary is pinned. + +The single-step-only OTTT/OTPE/OSTTP equivalence and deferral tests +(F-19/F-20) were removed in 0.3.0 along with those algorithms; see docs/specs +for the roadmap.""" import brainstate import jax @@ -204,7 +207,7 @@ def test_ostl_feedforward_equals_pp_prop_multistep(): assert_param_gradients_close(g_ff, g_pp, atol=ATOL_EQUIV) -# --- Task 5: one-step instantaneous equivalence (single-step-only algos) ------ +# --- Task 5: one-step instantaneous equivalence ------------------------------- def _onestep_grads(algo, x): """Weight gradient of (algo.update(x)**2).sum() at step 0 with zero trace. @@ -231,23 +234,22 @@ def _assert_onestep_equiv(spec, algo_factory, x): assert_param_gradients_close(g_algo, g_drtrl, atol=ATOL_EQUIV) -def test_otpe_full_matches_d_rtrl_one_step_on_tanh_rnn(): +def test_eprop_unfiltered_matches_d_rtrl_one_step_on_tanh_rnn(): + """EProp with no kappa filter and symmetric feedback is D_RTRL's trace, so + the instantaneous gradient must coincide.""" spec = tanh_rnn(n_in=3, n_rec=4, seed=0) - _assert_onestep_equiv(spec, lambda m: braintrace.OTPE(m, mode='full', leak=0.9), - jnp.ones((1, 3), dtype='float32')) - - -def test_ottt_matches_d_rtrl_one_step_on_leaky_linear(): - """OTTT(A) on its exact regime (leaky_linear): instantaneous gradient == D_RTRL.""" - spec = leaky_linear(n_in=3, n_rec=4, leak=0.9, seed=0) - _assert_onestep_equiv(spec, lambda m: braintrace.OTTT(m, mode='A', leak=0.9), - jnp.ones((1, 3), dtype='float32')) + _assert_onestep_equiv( + spec, + lambda m: braintrace.EProp(m, feedback='symmetric', kappa_filter_decay=0.0), + jnp.ones((1, 3), dtype='float32')) -def test_otpe_full_matches_d_rtrl_one_step_on_leaky_linear(): +def test_eprop_unfiltered_matches_d_rtrl_one_step_on_leaky_linear(): spec = leaky_linear(n_in=3, n_rec=4, leak=0.9, seed=0) - _assert_onestep_equiv(spec, lambda m: braintrace.OTPE(m, mode='full', leak=0.9), - jnp.ones((1, 3), dtype='float32')) + _assert_onestep_equiv( + spec, + lambda m: braintrace.EProp(m, feedback='symmetric', kappa_filter_decay=0.0), + jnp.ones((1, 3), dtype='float32')) # --- Task 6: vjp_method consistency & boundary ------------------------------- @@ -279,59 +281,3 @@ def test_multistep_method_is_the_exact_path(): spec.factory, inputs, algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) assert_param_gradients_close(g_multi, g_bptt, atol=ATOL_BPTT) - - -# --- Task 7: findings F-19 (OTTT/OTPE single-step-only) & F-20 (OSTTP) -------- - -def test_ottt_otpe_reject_multistep_oracle_F19(): - """F-19: OTTT/OTPE are single-step only, so the multi-step BPTT oracle path is - structurally unavailable. N6 moved this rejection from a lazy failure deep - inside the multi-step VJP machinery to an eager `ValueError` raised by the - constructor itself (`algo_factory(model)`, the first thing - `online_param_gradients` does); we pin that both algorithms still reject - `vjp_method='multi-step'`, which (with F-SINGLESTEP blocking the naive - single-step accumulation) is why their multi-step temporal correctness is - deferred, not asserted.""" - spec = tanh_rnn(n_in=3, n_rec=4, seed=0) - inputs = _inputs(6, 3) - for algo_factory, match in ( - (lambda m: braintrace.OTTT(m, mode='A', leak=0.9, vjp_method='multi-step'), - r"OTTT v1 only supports vjp_method='single-step'"), - (lambda m: braintrace.OTPE(m, mode='full', leak=0.9, vjp_method='multi-step'), - r"OTPE v1 only supports vjp_method='single-step'"), - ): - with pytest.raises(ValueError, match=match): - online_param_gradients(spec.factory, inputs, algo_factory=algo_factory) - - -@pytest.mark.skip( - reason="F-19: OTTT/OTPE are single-step only and the naive single-step " - "accumulation diverges (F-SINGLESTEP); there is no clean multi-step " - "BPTT oracle for them yet. Needs an online-scan oracle for single-step " - "algorithms. Deferred (P6 / dedicated work).") -def test_ottt_multistep_temporal_matches_bptt_DEFERRED(): - pass - - -def test_osttp_runs_through_oracle_but_signal_is_target_based_F20(): - """F-20: OSTTP runs end-to-end (sequence-end) and yields finite gradients, but - its learning signal is B @ y_target, NOT the autodiff dL/dh. The SSE-autodiff - BPTT/EProp oracle therefore does not constrain it — we only smoke-test that it - runs and is finite here; the OSTTP <-> EProp(random feedback) equivalence is - deferred (see DEFERRED test below).""" - spec = tanh_rnn(n_in=3, n_rec=4, seed=0) - inputs = _inputs(6, 3) - g = online_param_gradients( - spec.factory, inputs, - algo_factory=lambda m: braintrace.OSTTP( - m, [jnp.eye(4, dtype='float32')], - target_timing='sequence-end', vjp_method='multi-step')) - assert all(bool(jnp.all(jnp.isfinite(jnp.asarray(v)))) for v in g.values()) - - -@pytest.mark.skip( - reason="F-20: OSTTP uses B@y_target as the learning signal, bypassing the " - "autodiff loss; proving OSTTP == EProp(random feedback) needs a " - "matched-B + matched-per-step-target harness. Deferred (P5).") -def test_osttp_equals_eprop_random_feedback_DEFERRED(): - pass diff --git a/braintrace/_algorithm/tests/public_api_test.py b/braintrace/_algorithm/tests/public_api_test.py index e3935cc5..5ccea7f0 100644 --- a/braintrace/_algorithm/tests/public_api_test.py +++ b/braintrace/_algorithm/tests/public_api_test.py @@ -20,13 +20,27 @@ class TestPublicAPI(unittest.TestCase): def test_subpackage_exports(self): import braintrace._algorithm as pkg for name in ( - 'EProp', 'OSTLRecurrent', 'OSTLFeedforward', 'OTPE', 'OTTT', 'OSTTP', - 'FixedRandomFeedback', 'KappaFilter', 'PresynapticTrace', + 'EProp', 'OSTLRecurrent', 'OSTLFeedforward', + 'FixedRandomFeedback', 'KappaFilter', ): assert hasattr(pkg, name), f'missing export: {name}' def test_top_level_exports(self): import braintrace - for name in ('EProp', 'OSTLRecurrent', 'OSTLFeedforward', 'OTPE', 'OTTT', 'OSTTP'): + for name in ('EProp', 'OSTLRecurrent', 'OSTLFeedforward'): assert hasattr(braintrace, name), f'missing top-level export: {name}' assert name in braintrace.__all__ + + def test_removed_algorithms_are_gone(self): + """OTTT / OSTTP / OTPE were removed in 0.3.0 (see docs/specs roadmap). + + They whitelisted dense-matmul primitives and were not model-agnostic, so + they no longer belong in a general framework. Their coordinates remain + reachable through the axis configuration space. + """ + import braintrace + import braintrace._algorithm as pkg + for name in ('OTTT', 'OSTTP', 'OTPE', 'PresynapticTrace'): + assert not hasattr(braintrace, name), f'{name} should have been removed' + assert name not in braintrace.__all__ + assert not hasattr(pkg, name), f'{name} should have been removed' diff --git a/braintrace/_algorithm/tests/transform_correctness_test.py b/braintrace/_algorithm/tests/transform_correctness_test.py index b01efa4d..e3b45603 100644 --- a/braintrace/_algorithm/tests/transform_correctness_test.py +++ b/braintrace/_algorithm/tests/transform_correctness_test.py @@ -257,82 +257,3 @@ def test_algorithms_descend_loss_end_to_end(algo_factory): 15 SGD steps.""" losses = _train_losses(algo_factory(_conv_net())) assert losses[-1] < losses[0] - - -# --- Task 7: OSTTP under transform.jit (F-12 supported case) ----------------- - -def _osttp_net(n_in=3, n_rec=3): - class Net(brainstate.nn.Module): - def __init__(self): - super().__init__() - self.w = brainstate.ParamState(0.1 * jax.random.normal(jax.random.PRNGKey(0), (n_in, n_rec))) - self.h = brainstate.HiddenState(jnp.zeros((1, n_rec))) - - def update(self, x): - self.h.value = jax.nn.tanh(0.9 * self.h.value + braintrace.matmul(x, self.w.value)) - return self.h.value - - net = Net() - brainstate.nn.init_all_states(net, batch_size=1) - return net - - -def _osttp_first_step_grad(jit): - """One OSTTP update under grad with a fixed y_target, on a fresh algorithm. - Returns the param-gradient tree. Mirrors the known-good osttp_test setup - (n_target == n_rec == 3, small B, y = ones).""" - model = _osttp_net(n_in=3, n_rec=3) - B = 0.1 * jax.random.normal(jax.random.PRNGKey(9), (3, 3)) # (n_target=3, n_rec=3) - algo = braintrace.OSTTP(model, B_list=[B], target_timing='per-step') - x = jnp.ones((1, 3), dtype='float32') - algo.compile_graph(x) - algo.init_etrace_state() - y = jnp.ones((1, 3), dtype='float32') - - def gfn(x_, y_): - return brainstate.transform.grad( - lambda xx: ((algo.update(xx, y_target=y_) - 1.0) ** 2).mean(), - algo.param_states)(x_) - - gfn = brainstate.transform.jit(gfn) if jit else gfn - return gfn(x, y) - - -def test_osttp_jit_threads_instance_target_correctly_F12(): - """F-12: OSTTP stashes y_target in an instance attribute, raising the concern - that transform.jit could read a stale value. The jitted first-step param - gradient equals the eager one element-wise — the instance-attribute threading - is correct under jit for the supported (fixed-shape) path. (The variable-shape - y_target hazard is the still-OPEN F-12/F-14 gap, out of P6 scope.)""" - g_eager = _osttp_first_step_grad(jit=False) - g_jit = _osttp_first_step_grad(jit=True) - assert_param_gradients_close(g_jit, g_eager, atol=ATOL_JIT) - - -def test_osttp_trains_under_jit_with_fixed_target(): - """F-12 supported case end-to-end: with the known-good setup (n_target == - n_rec, small B, y = ones) OSTTP's target-projected signal drives the MSE-to- - ones loss down, and wrapping the whole train step in transform.jit preserves - that — the instance-attribute y_target survives jit across many steps.""" - model = _osttp_net(n_in=3, n_rec=3) - B = 0.1 * jax.random.normal(jax.random.PRNGKey(9), (3, 3)) - algo = braintrace.OSTTP(model, B_list=[B], target_timing='per-step') - x = jnp.ones((1, 3), dtype='float32') - algo.compile_graph(x) - algo.init_etrace_state() - - @brainstate.transform.jit - def train_step(x_, y_): - def loss_fn(xx): - out = algo.update(xx, y_target=y_) - return ((out - jnp.ones_like(out)) ** 2).mean() - grads, loss_val = brainstate.transform.grad( - loss_fn, algo.param_states, return_value=True)(x_) - for path, st in algo.param_states.items(): - st.value = st.value - 0.05 * grads[path] - return loss_val - - y = jnp.ones((1, 3), dtype='float32') - losses = [float(train_step(x, y)) for _ in range(10)] - assert np.isfinite(losses).all() - assert losses[-1] < losses[0] diff --git a/braintrace/_algorithm/vjp_base.py b/braintrace/_algorithm/vjp_base.py index d3a595cf..c9a5d002 100644 --- a/braintrace/_algorithm/vjp_base.py +++ b/braintrace/_algorithm/vjp_base.py @@ -113,8 +113,8 @@ class ETraceVjpAlgorithm(ETraceAlgorithm): #: (default) excludes recurrent ETP mixing primitives from the transition #: (bounded D-RTRL / e-prop diagonal approximation). Recurrence-defining #: algorithms that keep the full temporal term (RTRL-exact "with-H", e.g. - #: :class:`~braintrace.OSTLRecurrent` / :class:`~braintrace.OSTTP`) override - #: this to ``True``. Not a user-facing constructor argument by design. + #: :class:`~braintrace.OSTLRecurrent`) override this to ``True``. Not a + #: user-facing constructor argument by design. _include_recurrent_mixing: bool = False def __init__( @@ -263,8 +263,8 @@ def update(self, *args: Any) -> Any: # etrace data last_etrace_vals = self._get_etrace_data() - # Optional per-call auxiliary data (e.g. OSTTP's `y_target`) that a - # subclass needs inside `_compute_learning_signal` but that must not be + # Optional per-call auxiliary data (e.g. a neuromodulatory signal) that + # a subclass needs inside `_compute_learning_signal` but that must not be # forwarded to the model itself. Read synchronously here -- before the # custom_vjp machinery runs -- so it becomes a genuine argument of # `_true_update_fun` rather than an instance-attribute side channel: @@ -626,7 +626,7 @@ def _update_fn_bwd( # # Hook: subclasses may replace the reverse-AD learning signal with an - # alternative (e.g. target projection in OSTTP, κ-filtered signal in EProp). + # alternative (e.g. the κ-filtered signal in EProp). # # `aux` (see `_get_update_aux`) is appended to `args` here, not inside # `args` itself, so it never reaches the model's own forward call -- @@ -655,7 +655,7 @@ def _update_fn_bwd( ) # Note that there are no gradients flowing through the etrace data, the - # running index, or the auxiliary data (e.g. OSTTP's y_target). + # running index, or the auxiliary data. dg_etrace = None dg_running_index = None dg_aux = None @@ -686,10 +686,10 @@ def _get_update_aux(self) -> Any: lazy read would silently observe a cleared/stale value. Default returns `None` (unused). Subclasses that need auxiliary data - not already part of the model's own forward arguments (e.g. OSTTP's - `y_target`, which must not be forwarded to the model call) should - override this and read whatever they stashed on `self` during their - own `update()` override, at this exact call site. + not already part of the model's own forward arguments (e.g. a reward or + neuromodulatory signal, which must not be forwarded to the model call) + should override this and read whatever they stashed on `self` during + their own `update()` override, at this exact call site. Returns: Any pytree (or `None`). Appended to the `args` tuple seen by @@ -706,7 +706,8 @@ def _compute_learning_signal( """Override hook. Return the learning signal used by `_solve_weight_gradients`. Default returns the reverse-AD gradient unchanged. Subclasses that need - target projection (OSTTP) or any other alternative can override this. + a random-feedback projection, a modulatory signal, or any other + alternative can override this. Args: dl_to_hidden_from_autodiff: Sequence of per-hidden-group gradients produced @@ -714,10 +715,9 @@ def _compute_learning_signal( args: The `*args` tuple passed to the most recent `update()` call, with the per-call auxiliary data from `_get_update_aux` appended as the trailing element. Subclasses that override - `_get_update_aux` (e.g. OSTTP) can pull their auxiliary tensor - (e.g. ``y_target``) from there (see `extract_y_target` in - `_common.py`); subclasses that don't (the default `None`) can - ignore `args` entirely, as the base implementation does. + `_get_update_aux` can pull their auxiliary tensor from there; + subclasses that don't (the default `None`) can ignore `args` + entirely, as the base implementation does. Returns: Sequence of per-hidden-group gradient arrays, one per HiddenGroup. Must diff --git a/braintrace/_compile.py b/braintrace/_compile.py index 3ba277ce..98333064 100644 --- a/braintrace/_compile.py +++ b/braintrace/_compile.py @@ -28,9 +28,6 @@ EProp, OSTLRecurrent, OSTLFeedforward, - OTPE, - OTTT, - OSTTP, ) __all__ = ['compile'] @@ -47,9 +44,6 @@ 'e_prop': EProp, 'ostl_recurrent': OSTLRecurrent, 'ostl_feedforward': OSTLFeedforward, - 'otpe': OTPE, - 'ottt': OTTT, - 'osttp': OSTTP, } @@ -63,7 +57,7 @@ def _resolve_algorithm( algorithm : type or str Either an :class:`ETraceAlgorithm` subclass (returned unchanged) or a registered string name (case-insensitive), e.g. ``'D_RTRL'``, - ``'eprop'``, ``'ottt'``. + ``'eprop'``, ``'ostl_recurrent'``. Returns ------- @@ -124,8 +118,8 @@ def compile( need to be pre-initialized; ``compile`` always (re)initializes its states. algorithm : type or str An :class:`ETraceAlgorithm` subclass, or a registered case-insensitive - name, e.g. ``'D_RTRL'``, ``'es_d_rtrl'``, ``'eprop'``, ``'otpe'``, - ``'ottt'``, ``'osttp'``, ``'ostl_recurrent'``, ``'ostl_feedforward'``. + name, e.g. ``'D_RTRL'``, ``'es_d_rtrl'``, ``'eprop'``, + ``'ostl_recurrent'``, ``'ostl_feedforward'``. *example_inputs Example call inputs (arrays / :class:`SingleStepData` / :class:`MultiStepData`) matching what ``learner.update(...)`` will @@ -205,14 +199,6 @@ def compile( - ``'eprop'`` — ``feedback`` (default ``'symmetric'``), ``kappa_filter_decay`` (default ``0.0``), ``random_feedback_key`` (default ``None``), ``vjp_method``, ``fast_solve``, ``name``. - - ``'otpe'`` — ``leak`` (**required**, keyword-only), ``mode`` (default - ``'full'``), ``trace_clip_abs`` (default ``None``), ``vjp_method``, - ``name``. - - ``'ottt'`` — ``leak`` (**required**, keyword-only), ``mode`` (default - ``'A'``), ``vjp_method``, ``name``. - - ``'osttp'`` — ``B_list`` (**required**: per-layer feedback matrices), - ``target_timing`` (default ``'per-step'``), ``vjp_method``, ``fast_solve``, - ``name``. - ``'ostl_recurrent'`` — ``vjp_method``, ``fast_solve``, ``trace_dtype``, ``name``. - ``'ostl_feedforward'`` — ``decay_or_rank`` (default ``1e-6``), ``name``. diff --git a/braintrace/_compile_test.py b/braintrace/_compile_test.py index ed6bdb34..926085b5 100644 --- a/braintrace/_compile_test.py +++ b/braintrace/_compile_test.py @@ -40,9 +40,6 @@ def test_resolve_class_passthrough(): ('e_prop', 'EProp'), ('ostl_recurrent', 'OSTLRecurrent'), ('ostl_feedforward', 'OSTLFeedforward'), - ('otpe', 'OTPE'), - ('ottt', 'OTTT'), - ('osttp', 'OSTTP'), ]) def test_resolve_string_names(name, cls_name): assert _resolve_algorithm(name) is getattr(braintrace, cls_name) @@ -106,9 +103,9 @@ def test_compile_requires_example_inputs(): def test_compile_propagates_missing_required_option(): model = _fresh_model() x0 = jnp.ones((3,), dtype='float32') - # OTTT requires keyword-only ``leak``; constructor raises TypeError. + # pp_prop requires ``decay_or_rank``; constructor raises TypeError. with pytest.raises(TypeError): - braintrace.compile(model, 'OTTT', x0, batch_size=1) + braintrace.compile(model, 'pp_prop', x0, batch_size=1) def test_compile_matches_manual_construction(): @@ -209,9 +206,6 @@ def test_compile_verbose_levels_print(capsys): 'd_rtrl': ['name', 'vjp_method', 'fast_solve', 'trace_dtype'], 'es_d_rtrl': ['decay_or_rank', 'name', 'vjp_method', 'fast_solve'], 'eprop': ['feedback', 'kappa_filter_decay', 'random_feedback_key', 'name', 'vjp_method', 'fast_solve'], - 'otpe': ['mode', 'leak', 'name', 'vjp_method', 'trace_clip_abs'], - 'ottt': ['mode', 'leak', 'name', 'vjp_method'], - 'osttp': ['B_list', 'target_timing', 'name', 'vjp_method', 'fast_solve'], 'ostl_recurrent': ['name', 'vjp_method', 'fast_solve', 'trace_dtype'], 'ostl_feedforward': ['decay_or_rank', 'name'], } diff --git a/braintrace/_compiler/graph.py b/braintrace/_compiler/graph.py index b4cee668..99cc8bf0 100644 --- a/braintrace/_compiler/graph.py +++ b/braintrace/_compiler/graph.py @@ -313,7 +313,7 @@ def compile_etrace_graph( recurrence"), those primitives are traced into the transition, the recurrence becomes coupled, and the true per-position block-diagonal Jacobian is extracted (RTRL-exact temporal credit, e.g. for - :class:`~braintrace.OSTLRecurrent` / :class:`~braintrace.OSTTP`). + :class:`~braintrace.OSTLRecurrent`). control_flow : ControlFlowPolicy or None, optional Policy governing control-flow canonicalization and downstream handling, forwarded to :func:`~braintrace.extract_module_info` and diff --git a/changelog.md b/changelog.md index c59d4cfb..ae12c367 100644 --- a/changelog.md +++ b/changelog.md @@ -3,6 +3,36 @@ ## UNRELEASED +### Breaking changes + +- **Removed `OTTT`, `OSTTP` and `OTPE`.** `braintrace` is a framework for online + learning in brain simulation, and a framework should ship *general* + mechanisms. These three rules were not model-agnostic: all of them whitelisted + dense-matmul primitives (`_SUPPORTED_PRIMITIVES = {etp_mm_p, etp_mv_p}`) and + raised `NotImplementedError` for lora / sparse / convolutional / + element-wise relations, and all of them were single-step only. `OTPE` + additionally assumed a single global time constant, was feed-forward only, + was gradient-exact for one hidden layer, and rejected `num_state > 1` + outright — ruling out ALIF and any adaptation variable. `OSTTP` bound + `B_list` to the HiddenGroup count and threaded `y_target` through a bespoke + path. + + Their coordinates remain reachable: the planned axis decomposition + (`trace_factorization` × `temporal_recursion` × `learning_signal` × + `trace_filter` × `update_schedule`) expresses each of them as a configuration + that works for **every** ETP primitive, not just dense matmul. See + `docs/specs/2026-07-25-algorithm-axes-roadmap.md`. + + Also removed with them: `PresynapticTrace` (used only by `OTTT`) and the + internal `extract_y_target` helper (used only by `OSTTP`). `KappaFilter` and + `FixedRandomFeedback` are unaffected. The `'ottt'`, `'osttp'` and `'otpe'` + names no longer resolve in `braintrace.compile`. + + **Migration:** `OTTT` → `pp_prop` (same `io_factorized` trace, keeps the + temporal term instead of dropping it); `OTPE` → `D_RTRL` or `pp_prop`; + `OSTTP` → `EProp(feedback='random')`, which is random feedback on the error + rather than on the target. + ### Improvements - **`sparse_matmul` migrates off brainevent's deprecated trace protocol.** diff --git a/docs/apis/algorithms.rst b/docs/apis/algorithms.rst index 7c737aae..37282c55 100644 --- a/docs/apis/algorithms.rst +++ b/docs/apis/algorithms.rst @@ -132,13 +132,9 @@ regime makes them exact); know the regime before relying on their gradients. EProp OSTLRecurrent OSTLFeedforward - OTPE - OTTT - OSTTP Trace helpers reused across the SNN algorithms — a frozen random-feedback -projection, an output-side low-pass filter, and a leaky presynaptic -accumulator: +projection and an output-side low-pass filter: .. autosummary:: :toctree: generated/ @@ -147,7 +143,6 @@ accumulator: FixedRandomFeedback KappaFilter - PresynapticTrace Algorithm Comparison @@ -177,15 +172,4 @@ Algorithm Comparison - depends on regime - depends on regime - ``OSTLRecurrent`` ('with-H', D-RTRL) keeps the recurrent Jacobian; ``OSTLFeedforward`` ('without-H', pp_prop) drops it. - * - ``OTPE`` - - :math:`O(B \cdot I \cdot O)` (full) / :math:`O(B(I+O))` (approx) - - :math:`O(B \cdot I \cdot O)` - - Deep SNNs; F-OTPE trades rank for memory - * - ``OTTT`` - - :math:`O(B \cdot I)` - - :math:`O(B \cdot I \cdot O)` - - Very large SNNs; presynaptic λ-trace only - * - ``OSTTP`` - - :math:`O(B \cdot |\theta|)` - - :math:`O(B \cdot I \cdot O)` - - Target-projection via fixed random feedback + diff --git a/docs/apis/index.rst b/docs/apis/index.rst index 2d14b193..c4c9c4a3 100644 --- a/docs/apis/index.rst +++ b/docs/apis/index.rst @@ -33,7 +33,7 @@ layers built on top. - :doc:`compiler` * - **Algorithms** - Online-learning orchestrators: the exact D-RTRL / ES-D-RTRL family and - the SNN algorithms (EProp, OSTL, OTPE, OTTT, OSTTP). + the SNN algorithms (EProp, OSTL). - :doc:`algorithms` * - **Layers** - Drop-in ``brainstate.nn``-style layers pre-wired through ETP primitives. diff --git a/docs/specs/2026-07-25-algorithm-axes-roadmap.md b/docs/specs/2026-07-25-algorithm-axes-roadmap.md new file mode 100644 index 00000000..bc467127 --- /dev/null +++ b/docs/specs/2026-07-25-algorithm-axes-roadmap.md @@ -0,0 +1,378 @@ +# Roadmap: orthogonal axes for online-learning algorithms + +Status: design, awaiting review. S1 (removal) is implemented. +Target release: 0.3.0 (single breaking release; all sub-projects land in it) + +## Why + +`braintrace` is a framework for online learning in brain simulation. A framework +ships *general* mechanisms; it should not ship learning rules that only work for +one operator type and one model shape. Three rules in Layer 4 violate that +today, and three more are named top-level classes for what are really +configurations of the two general engines. + +The Nature Communications work established the model-agnostic abstraction +(AlignPre / AlignPost), the linear-memory rule (pp-prop), and the compiler that +generates online-learning code from a user-defined SNN. The claim that followed — +that fragmented rules (e-prop, OSTL, OTPE, …) can be *described, implemented, +compared and deployed* in one compiler framework — is currently backed by +separate hand-written classes, several of which reject most of the operator set. +This roadmap makes the claim structural instead of incidental: the rules become +coordinates in an explicit axis space, and any coordinate works for every ETP +primitive. + +## Removal criteria + +An algorithm keeps a standalone implementation only if it passes **both**: + +1. **Model-agnostic.** It must hold for any ETP primitive (dense, conv, sparse, + lora, element-wise) and any hidden-state dynamics. An implementation that + whitelists primitives fails this test. +2. **Mathematically independent.** It must contribute a recursion or estimator + that the general engines cannot express as a configuration. A named class + whose coordinates coincide with an existing one contributes nothing. + +## The axes + +### Axis 1 — `trace_factorization` (spatial factorization → memory) + +| Value | Trace shape | Memory | Origin | +|---|---|---|---| +| `per_param` | `(param_shape, H)` | O(P·H) | existing `ParamDimVjpAlgorithm` | +| `io_factorized` | `ε_x ⊗ ε_f` | O(I+O) | existing `IODimVjpAlgorithm` | +| `sparse_n` | n-step influence sparsity pattern | between the two | new (S4) | +| `random_projection` | rank-1 random factors `(s̃, θ̃)` | O(P+H), **unbiased** | new (S5) | + +### Axis 2 — `temporal_recursion` (how the trace advances in time) + +- `jacobian` — `ε ← D·ε + …` (D-RTRL, OSTL 'with-H', e-prop) +- `scalar_leak` — `R̂ ← λ·R̂ + …`. The additive term is the same per-primitive + hidden→weight Jacobian contribution the `jacobian` recursion already uses, so + swapping the recursion is primitive-agnostic. This axis value is what remains + of OTPE after the algorithm itself was removed. +- `none` — no temporal accumulation (feedforward regime) + +Under `io_factorized` this axis is a **pair** `(x-side, f-side)`. The current +`IODimVjpAlgorithm` shares one `α` across `ε_x` and `ε_f` (derived by +`_format_decay_and_rank`); splitting it into two independent decays is the only +change in this roadmap that touches a numerical path in existing code. + +### Axis 3 — `learning_signal` (where the signal comes from) + +- `symmetric` — `∂L/∂h` back-propagated through the readout (default) +- `random_feedback` — fixed random projection; promotes the existing + `FixedRandomFeedback` helper to a first-class strategy +- `modulatory` — three-factor: a scalar / low-dimensional neuromodulatory signal + (TD error, reward-prediction error) times the trace (new, S6) +- `bootstrapped` — synthetic gradient / DNI: a learned estimate of the future + gradient (new, S6) + +DRTP / target projection is **not** a value on this axis. See "Removals". + +### Axis 4 — `trace_filter` + +`none` · `kappa` (e-prop's low-pass `ē ← κ·ē + ε`) + +### Axis 5 — `update_schedule` + +`per_step` · `window(k)` · `sequence_end` + +### Compatibility matrix + +The axes are not fully orthogonal: `random_projection` carries UORO's own rank-1 +update and normalisation and cannot be paired with an arbitrary +`temporal_recursion`. `ETraceConfig` validates the `(factorization × +recursion)` combination at construction and raises a readable error listing the +legal pairings. The matrix is explicit data, not scattered `if` statements. + +## Coordinates of the existing algorithms + +| Algorithm | factorization | recursion | signal | filter | +|---|---|---|---|---| +| `D_RTRL` | per_param | jacobian | symmetric | none | +| `pp_prop` | io_factorized | jacobian | symmetric | none | +| `OSTLRecurrent` | per_param | jacobian | symmetric | none | +| `OSTLFeedforward` | io_factorized | none | symmetric | none | +| `EProp` | per_param | jacobian | symmetric \| random_feedback | **kappa** | +| ~~`OTPE`~~ (exact) | per_param | **scalar_leak** | symmetric | none | +| ~~`OTPE`~~ (F-OTPE) | io_factorized | **scalar_leak** | symmetric | none | +| ~~`OTTT`~~ | io_factorized | (x: scalar_leak, f: none) | symmetric | none | +| ~~`OSTTP`~~ | per_param | jacobian | target_projection | none | + +The three struck-through rows all whitelist dense-matmul primitives and are +single-step only; the table is the proof behind the removal decisions. +`OSTLRecurrent` matches +`D_RTRL` cell for cell — and the source says so directly ("delegates entirely to +`ParamDimVjpAlgorithm`"). + +**Deleting the classes does not delete the capability.** After the axes exist, +OTTT's coordinate is still reachable as a configuration — and reachable for +conv / sparse / lora, which the current OTTT rejects. The framework ships only +general mechanisms while the "one framework describes them all" claim gets +*stronger*, not weaker. + +## Removals + +### OTTT + +`ottt.py` hard-codes `_SUPPORTED_PRIMITIVES = {etp_mm_p, etp_mv_p}` and raises +`NotImplementedError` for lora / sparse / conv / element-wise relations, plus +`'OTTT v1 supports single-step only'`. Fails criterion 1. Its coordinate is +`io_factorized` with an x-side leak and an f-side that does not accumulate — +one point in the configuration space once axis 2 splits per side. + +### OSTTP + +`osttp.py` fails criterion 1 for engineering rather than mathematical reasons: +`y_target` has to reach every HiddenGroup through a bespoke path, and `B_list` +is hard-bound to the HiddenGroup count. DRTP itself is just `random_feedback` +applied to the target instead of the error. **Decision: remove the value as well +as the class** — axis 3 carries no `target_projection`, and no generic +`y_target` plumbing is introduced for it. + +Note for S6: `modulatory` needs a general external-signal injection path. That is +what OSTTP's `y_target` was reaching for and got wrong. Build it generically. + +### OTPE + +`otpe.py` fails criterion 1 harder than OTTT does, and says so itself: its +docstring states the published derivation is *"narrower than OTTT's"*. It +carries the same `_SUPPORTED_PRIMITIVES = {etp_mm_p, etp_mv_p}` whitelist and +the same `'OTPE v1 supports single-step only'` guard, hard-codes the dense outer +product (`jnp.einsum('bi,bo->bio', x, df_proj)`) instead of routing through the +per-primitive rule registry, rejects `num_state > 1` outright (so ALIF and any +adaptation variable are out), and errors on any relation touching more than one +HiddenGroup. On top of that it assumes a **single global time constant**, is +**feed-forward only**, and is gradient-exact for **one hidden layer**. + +For brain simulation those assumptions are backwards: heterogeneous time +constants and adaptation variables are the norm. What generalises is the +`scalar_leak` recursion, which survives as an axis-2 value; what does not is +OTPE-the-published-algorithm, which is `scalar_leak` plus that whole restriction +stack. + +### Verified deletion footprint + +| Target | Verification | +|---|---| +| `_algorithm/ottt.py`, `ottt_test.py` | — | +| `_algorithm/osttp.py`, `osttp_test.py` | — | +| `_algorithm/otpe.py`, `otpe_test.py` | — | +| `_common.py::PresynapticTrace` + its tests | grep confirms `ottt.py` is the only user | +| `_common.py::extract_y_target` + its tests | grep confirms `osttp.py` is the only user | +| imports / `__all__` in both `__init__.py` levels | `_algorithm/__init__.py`, `braintrace/__init__.py` | +| `docs/apis/algorithms.rst`, `docs/apis/index.rst` | — | + +Referencing files that need edits but not deletion: `oracle_models.py`, +`vjp_base.py`, `_compile.py`, `_compiler/graph.py`, `__init___test.py`, +`_compile_test.py`, and in `_algorithm/tests/`: `approx_correctness_test.py`, +`exact_correctness_test.py`, `public_api_test.py`, +`transform_correctness_test.py`. + +The `_get_update_aux` / auxiliary-data hook in `vjp_base.py` is **kept**. OSTTP +was its only consumer, but it is the general per-call side channel that S6's +`modulatory` signal needs; only the OSTTP-specific prose was rewritten. + +Explicitly kept: `KappaFilter` (standalone utility; its docstring already notes +e-prop filters internally) and `FixedRandomFeedback` (promoted to the axis-3 +`random_feedback` strategy). + +## Public API shape at 0.3.0 + +```python +# Engines — the hosts for the axes +ETraceVjpAlgorithm # base class; hooks delegate to strategies +ParamDimVjpAlgorithm # factorization = per_param +IODimVjpAlgorithm # factorization = io_factorized + +# braintrace's own contributions — keep the names +D_RTRL, pp_prop + +# Configuration surface +ETraceConfig(factorization=..., recursion=..., signal=..., filter=..., schedule=...) + +# Strategy protocols — third-party extension points +TraceFactorization, TemporalRecursion, LearningSignalSource, TraceFilter, UpdateSchedule + +# Literature presets — thin factories over ETraceConfig +EProp, OSTLRecurrent, OSTLFeedforward +``` + +`braintrace.compile(model, algo, x0, **kw)` keeps its signature; `algo` +additionally accepts an `ETraceConfig`. + +**Mechanism vs surface.** Internally the axes are strategy objects injected into +one engine — this is isomorphic to the existing template-method hooks +(`_compute_learning_signal`, `_solve_weight_gradients`, `_update_etrace_data`, +`_make_etrace_stepper`, `init_etrace_state`), so the delta is small and each +strategy is independently testable. Externally `ETraceConfig` parses into a +strategy combination, which keeps the user-facing API simple and makes the +configuration space enumerable — S7's benchmark suite gets that for free. + +A declarative learning-rule IR (rules as a compiler IR layer, third parties +adding rules without touching the engine) is the natural long-term direction but +is deliberately **out of scope**. With ~8 rules the correct abstraction boundary +is not yet observable. Revisit once S4–S6 push the count past ~15. + +## Compatibility + +0.3.0 is a clean break. `OTTT`, `OSTTP` and `OTPE` are removed outright with no +shim; +the changelog documents the migration path. The repo already carries a `_legacy` +deprecation layer, and stacking another one only thickens the debt. + +## Sub-projects + +All land in 0.3.0. Ordering below reflects the dependency structure, not +separate releases. + +### S1 — Removal (do first) — **implemented** + +Delete `ottt.py` / `osttp.py` / `otpe.py` and their helpers, clean every +reference. Small, mechanical, no numerical risk. Doing it before S2 shrinks the +compiler-refactor surface by three algorithms that would otherwise have to be +kept working and then deleted. + +Coverage that pointed at the removed algorithms was **repointed, not dropped**, +wherever the assertion was about a general property: + +- the approximate-gradient descent backstop now runs on `pp_prop(rank=1)` and + `EProp(feedback='random')` instead of `OTTT`/`OTPE`; +- the one-step D_RTRL equivalence tests now use `EProp(kappa_filter_decay=0)`, + which is D_RTRL's trace with no filter (verified: passes element-wise); +- `public_api_test.py` gained a guard asserting the removed names stay gone. + +The direction-alignment metric helpers (`cosine_similarity`, `sign_agreement`, +`relative_magnitude`, `assert_direction_aligned`) are retained — they are the +basis of S7's benchmark suite. + +**Acceptance:** full suite green; `grep -r "OTTT\|OSTTP\|OTPE\|PresynapticTrace\|extract_y_target"` +over `braintrace/` and `docs/` hits only `changelog.md`, the removal-notice +docstrings, and the guard test; mypy clean; `py.typed` intact. + +### S2 — Compiler: multi-timescale and heterogeneous populations + +Work through the known-limitation list already recorded in `AGENTS.md` and the +`dev/` findings: heterogeneous-population leak resolution, multi-state +HiddenGroups, approximation validity beyond shallow depth, single-readout and +feedback-shape assumptions. This comes before the algorithm work so that the +existing rules actually run on realistic brain-simulation models (LIF + ALIF + +multi-timescale synapses) before the algorithm family is widened. + +The Jacobian-representation part of this sub-project is a prerequisite for S4. + +**Acceptance:** each limitation item has either a passing test or an explicit, +documented scope boundary. No expected-failure item silently remains. + +### S3 — Axis decomposition + +Turn the five axes into code: the strategy protocols, `ETraceConfig` and its +compatibility matrix, the engine hooks rewired to delegate to strategies, and +the three literature presets (`EProp`, `OSTLRecurrent`, `OSTLFeedforward`) +rewritten as thin factories over `ETraceConfig`. Includes the +`IODimVjpAlgorithm` decay split — the one numerical change in this roadmap. + +Everything from S4 onward adds values to these axes, so this is the last piece +of groundwork. + +**Acceptance:** +- Element-wise equality against golden values frozen *before* the refactor, for + all five presets (`D_RTRL`, `pp_prop`, `EProp`, `OSTLRecurrent`, + `OSTLFeedforward`). +- `decay_or_rank=0.9` equals the new two-sided `(0.9, 0.9)` element-wise. +- Illegal axis combinations raise a readable error naming the legal pairings, + with test coverage. +- Every preset's coordinates are asserted against the table in this document, so + the table cannot silently drift from the code. + +### S4 — SnAp-n sparsity axis + +Menick et al. 2021. Derive the n-step influence sparsity pattern automatically +from the compiler's jaxpr hidden→hidden reachability graph — this is the thing a +compiler can do that a hand-written library cannot, and it sits directly on the +paper's technical moat. + +**The work is in the compiler, not the algorithm.** The current machinery keeps +only the block diagonal of `D^t` (`HiddenGroup.diagonal_jacobian`, +`block_diagonal_last_dim`); SnAp-n needs off-block-diagonal terms within n steps. +Assume the existing representation does not survive n>1 and budget for a new +sparse representation, shared with S2. + +**Acceptance (two-sided squeeze):** +- `n = 1` equals the existing `D_RTRL` element-wise (regression guard). +- `n ≥ graph diameter` equals the **BPTT oracle** element-wise, on models whose + hidden-state coupling the compiler fully captures. Full RTRL and BPTT compute + the same total gradient, so `oracle.py` is the correct instrument and no + separate full-RTRL reference is needed. (A full-RTRL reference is worth + writing only to assert the influence matrix `dh^t/dθ` itself when localising a + divergence between trace and learning signal — optional debugging aid.) +- Measured memory curve, monotone in n. + +### S5 — UORO (unbiased random-projection estimator) + +Tallec & Ollivier 2018. Adds the `random_projection` factorization with UORO's +rank-1 update and normalisation. Complements the existing biased diagonal +approximations. KF-RTRL / OK are optional lower-variance siblings on the same +axis value. + +**Acceptance is statistical, not element-wise.** An unbiased estimator will not +match BPTT on any single run. The guard is: fixed model, fixed sequence, N random +seeds; the deviation of the mean gradient from BPTT shrinks as 1/√N (confidence +interval test). A single run asserts only shape, finiteness, and absence of NaN. +This requires statistical test infrastructure the repo does not have — count it +in the sub-project's cost. + +### S6 — Learning-signal axis: three-factor and DNI + +- `modulatory`: three-factor learning — trace × neuromodulatory signal (TD error, + reward-prediction error), enabling reward-based e-prop and online policy + gradient. Needs the general external-signal injection path noted above. +- `bootstrapped`: DNI / synthetic gradients (Jaderberg 2017) — a learned + bootstrap for the future-loss gradient that every online rule truncates away. + Requires an auxiliary network with its own training loop. + +**Acceptance is by degeneracy, plus task performance.** `modulatory` must equal +`symmetric` element-wise when the modulatory signal is set to `∂L/∂h`; +`bootstrapped` must equal `symmetric` when the synthesiser output is pinned to +the true value. Add an end-to-end RL smoke test. Neither is a "more accurate +gradient", so element-wise comparison against BPTT is the wrong instrument. + +### S7 — Unified benchmark suite + +Enumerate the `ETraceConfig` space and, for a fixed model, report gradient cosine +similarity and relative deviation against BPTT, peak memory, per-step wall time, +and task metrics. Machine-readable output (JSON/CSV) plus a reproducible script. +This turns the paper's "comparable" claim into something executable. + +Depends on S3's axes; grows as S4–S6 land. + +## Three acceptance paradigms + +Mixing these up is the most likely way to get this roadmap wrong: + +| Paradigm | Applies to | Instrument | +|---|---|---| +| Element-wise equality | S1 refactor, SnAp-1, SnAp-∞ | golden values / BPTT oracle | +| Statistical convergence | UORO and unbiased siblings | 1/√N confidence interval over seeds | +| Degeneracy + task metric | modulatory, bootstrapped | reduce to `symmetric`, then RL smoke test | + +The existing taxonomy in `AGENTS.md` (exact vs approximate) stays valid and is +extended by the statistical class. + +## Risk register + +1. **Axis refactor changes numerics silently.** Mitigation: before touching the + engine, freeze reference gradients for all five presets as golden values; + assert element-wise equality after. Separately, the `IODim` decay split must + prove `decay_or_rank=0.9` equals the new `(0.9, 0.9)` element-wise. +2. **The sparse Jacobian representation does not generalise** (S2/S4). Mitigation: + design the representation once, in S2, with SnAp-n as the explicit consumer. +3. **Statistical tests are flaky in CI** (S5). Mitigation: fixed seeds, generous + intervals, and a separate slow-test marker. +4. **`modulatory` re-creates OSTTP's plumbing mistake** (S6). Mitigation: the + injection path must not bind to HiddenGroup count or readout shape; add a test + with a model whose HiddenGroup count differs from the signal dimension. +5. **0.3.0 carries the whole roadmap.** All breaking changes and all new + algorithm families ship together. Mitigation: sub-projects merge + independently behind the axis interfaces, and the benchmark suite (S7) runs on + every merge so regressions surface per sub-project rather than at release. From bc153da72ef7eebe3301f3056f80d118e441779c Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 11:24:08 +0800 Subject: [PATCH 02/13] Rewrite the algorithm-axes roadmap against the post-removal code The previous draft was written before the removal landed and described the codebase it expected rather than the one that exists. Rewritten with commit 928219b as the baseline, and corrected against the source. Two substantive corrections: - OSTLRecurrent is NOT a D_RTRL alias. The earlier draft read only its docstring ("delegates entirely to ParamDimVjpAlgorithm") and missed `_include_recurrent_mixing = True`, which flows into the compiler and sets HiddenGroup.is_diagonal_recurrence. The draft argued on that basis that the class was a zero-information alias and should be deleted; the argument was wrong and the class stays. - That same flag turns out to be a sixth axis, `recurrence_scope`, already present in the code as a boolean: `diagonal` excludes recurrent ETP mixing primitives from the transition (cheap column-sum Jacobian is exact) and `coupled` traces through them and extracts the true block diagonal. SnAp-n is the n-valued generalisation of that existing knob rather than a new sparse representation, which materially lowers the cost of that phase -- and vindicates scheduling the compiler work first, since both phases rework the same Jacobian path in hidden_group.py. Also grounded in the source rather than assumed: the surviving algorithms and what each actually overrides; the engine's real hook surface; the distinction between learning-rule axes and execution options (vjp_method, fast_solve, trace_dtype, chunked_trace, control_flow stay out of ETraceConfig); the existing oracle.py inventory, which already covers every acceptance paradigm. Corrected the P1 scope claim: F-SINGLESTEP is a *former* finding, resolved into a positive direction-alignment assertion, not a live limitation. F-19 and F-20 died with the algorithms they described. Only F-22 remains, so the rebuilt limitation list will be shorter than the original. Fixed a stale cross-reference found while verifying that: a docstring in exact_correctness_test.py pointed at oracle_test.py::test_singlestep_naive_matches_bptt_KNOWN_DIVERGENCE, which was renamed to test_singlestep_naive_directionally_aligned_with_bptt when the finding was resolved. Phases renumbered P0 (done) through P5. Verified: pytest braintrace/_algorithm/tests/ braintrace/_algorithm/oracle_test.py -> 130 passed, 1 skipped. --- .../tests/exact_correctness_test.py | 6 +- .../2026-07-25-algorithm-axes-roadmap.md | 549 +++++++++--------- 2 files changed, 282 insertions(+), 273 deletions(-) diff --git a/braintrace/_algorithm/tests/exact_correctness_test.py b/braintrace/_algorithm/tests/exact_correctness_test.py index ab1d70f1..714f5227 100644 --- a/braintrace/_algorithm/tests/exact_correctness_test.py +++ b/braintrace/_algorithm/tests/exact_correctness_test.py @@ -271,8 +271,10 @@ def test_singlestep_method_rejects_multistep_data(): def test_multistep_method_is_the_exact_path(): """Cross-reference: multi-step vjp_method is the exact path (proven in Task 3 - and in oracle_test.py). The naive per-step single-step accumulation diverges - (F-SINGLESTEP, oracle_test.py::test_singlestep_naive_matches_bptt_KNOWN_DIVERGENCE). + and in oracle_test.py). The naive per-step single-step accumulation does not + match BPTT element-wise -- it is only direction-aligned (the former + F-SINGLESTEP finding, now asserted positively by + oracle_test.py::test_singlestep_naive_directionally_aligned_with_bptt). Here we re-confirm multi-step D_RTRL == BPTT to anchor the vjp_method dimension.""" spec = tanh_rnn(n_in=3, n_rec=4, seed=0) inputs = _inputs(6, 3) diff --git a/docs/specs/2026-07-25-algorithm-axes-roadmap.md b/docs/specs/2026-07-25-algorithm-axes-roadmap.md index bc467127..4bf00ce4 100644 --- a/docs/specs/2026-07-25-algorithm-axes-roadmap.md +++ b/docs/specs/2026-07-25-algorithm-axes-roadmap.md @@ -1,350 +1,351 @@ # Roadmap: orthogonal axes for online-learning algorithms -Status: design, awaiting review. S1 (removal) is implemented. -Target release: 0.3.0 (single breaking release; all sub-projects land in it) +Status: design, awaiting review +Baseline: commit `928219b` (OTTT / OSTTP / OTPE removed) +Target release: 0.3.0 — every phase below lands in it ## Why `braintrace` is a framework for online learning in brain simulation. A framework -ships *general* mechanisms; it should not ship learning rules that only work for -one operator type and one model shape. Three rules in Layer 4 violate that -today, and three more are named top-level classes for what are really -configurations of the two general engines. +ships *general* mechanisms; it should not ship learning rules that work for one +operator type and one model shape. The Nature Communications work established the model-agnostic abstraction (AlignPre / AlignPost), the linear-memory rule (pp-prop), and the compiler that generates online-learning code from a user-defined SNN. The claim that followed — -that fragmented rules (e-prop, OSTL, OTPE, …) can be *described, implemented, -compared and deployed* in one compiler framework — is currently backed by -separate hand-written classes, several of which reject most of the operator set. -This roadmap makes the claim structural instead of incidental: the rules become -coordinates in an explicit axis space, and any coordinate works for every ETP -primitive. +that fragmented rules (e-prop, OSTL, …) can be *described, implemented, compared +and deployed* in one compiler framework — was backed by separate hand-written +classes, three of which rejected most of the operator set. Those three are gone +(see Baseline). This roadmap makes the claim structural rather than incidental: +learning rules become coordinates in an explicit axis space, and every coordinate +works for every ETP primitive. ## Removal criteria -An algorithm keeps a standalone implementation only if it passes **both**: +These governed the removals already made and govern every future addition. An +algorithm earns a standalone implementation only if it passes **both**: 1. **Model-agnostic.** It must hold for any ETP primitive (dense, conv, sparse, lora, element-wise) and any hidden-state dynamics. An implementation that whitelists primitives fails this test. 2. **Mathematically independent.** It must contribute a recursion or estimator - that the general engines cannot express as a configuration. A named class - whose coordinates coincide with an existing one contributes nothing. + the general engines cannot express as a configuration. A named class whose + coordinates coincide with an existing one contributes nothing. -## The axes +## Baseline: what the code looks like now + +### What was removed, and why (history) + +`OTTT`, `OSTTP` and `OTPE` all failed criterion 1: each whitelisted dense-matmul +primitives (`_SUPPORTED_PRIMITIVES = {etp_mm_p, etp_mv_p}`), raised +`NotImplementedError` for lora / sparse / conv / element-wise relations, and was +single-step only. `OTPE` additionally assumed a single global time constant, was +feed-forward only, was gradient-exact for one hidden layer, and rejected +`num_state > 1` outright — ruling out ALIF and any adaptation variable; its own +docstring called its derivation *"narrower than OTTT's"*. `OSTTP` bound `B_list` +to the HiddenGroup count and threaded `y_target` through a bespoke path. + +Removed with them: `PresynapticTrace` (only `OTTT` used it) and the internal +`extract_y_target` (only `OSTTP` used it). Kept deliberately: `KappaFilter`, +`FixedRandomFeedback`, and the `_get_update_aux` per-call side channel in +`vjp_base.py` — `OSTTP` was its only consumer, but it is the general hook that +P4's `modulatory` signal needs. + +### Surviving algorithms + +| Class | Relationship | +|---|---| +| `ParamDimVjpAlgorithm` | engine — per-parameter trace, O(P·H) | +| `IODimVjpAlgorithm` | engine — input/output-factorized trace, O(I+O) | +| `D_RTRL` | canonical name for `ParamDimVjpAlgorithm`, no overrides | +| `pp_prop` / `ES_D_RTRL` | canonical name for `IODimVjpAlgorithm`, no overrides | +| `EProp` | `ParamDimVjpAlgorithm` + κ-filter + optional random feedback; overrides `init_etrace_state`, `reset_state`, `_compute_learning_signal`, `_solve_weight_gradients` | +| `OSTLRecurrent` | `ParamDimVjpAlgorithm` + `_include_recurrent_mixing = True` | +| `OSTLFeedforward` | `pp_prop` with `decay_or_rank` defaulting to `1e-6` | + +### The engine's extension surface + +`ETraceVjpAlgorithm` is already a template-method class. Its override hooks are +the natural mounting points for the axes: + +| Hook | What it decides | +|---|---| +| `init_etrace_state` | trace allocation → factorization | +| `_update_etrace_data` / `_make_etrace_stepper` | trace recursion in time | +| `_solve_weight_gradients` | how trace × signal contracts to a weight gradient | +| `_compute_learning_signal` | where the learning signal comes from | +| `_get_update_aux` | per-call external data (reward, target, …) side channel | +| `_get_etrace_data` / `_assign_etrace_data` | trace (de)serialisation — abstract | +| `_include_recurrent_mixing` (class attr) | how much hidden↔hidden coupling enters the trace | + +### Learning-rule axes vs execution options + +The constructors already carry knobs that are **not** learning-rule axes and +must stay out of `ETraceConfig`: `vjp_method` (`'single-step'` / `'multi-step'`), +`fast_solve`, `trace_dtype`, `chunked_trace`, `control_flow`. These choose *how* +a rule is executed, not *which* rule it is. Conflating them would make the +benchmark suite enumerate a configuration space that is mostly execution +variants. + +### The oracle inventory + +`_algorithm/oracle.py` already provides `bptt_param_gradients`, +`finite_difference_param_gradients`, `online_param_gradients`, +`chunked_online_param_gradients`, `online_param_gradients_singlestep_naive`, +`assert_param_gradients_close`, plus the direction metrics +`cosine_similarity` / `sign_agreement` / `relative_magnitude` / +`assert_direction_aligned`. This is enough to serve every acceptance paradigm +below; no new oracle family is required. + +## The axis space ### Axis 1 — `trace_factorization` (spatial factorization → memory) -| Value | Trace shape | Memory | Origin | +| Value | Trace shape | Memory | Status | |---|---|---|---| -| `per_param` | `(param_shape, H)` | O(P·H) | existing `ParamDimVjpAlgorithm` | -| `io_factorized` | `ε_x ⊗ ε_f` | O(I+O) | existing `IODimVjpAlgorithm` | -| `sparse_n` | n-step influence sparsity pattern | between the two | new (S4) | -| `random_projection` | rank-1 random factors `(s̃, θ̃)` | O(P+H), **unbiased** | new (S5) | +| `per_param` | `(param_shape, H)` | O(P·H) | exists (`ParamDimVjpAlgorithm`) | +| `io_factorized` | `ε_x ⊗ ε_f` | O(I+O) | exists (`IODimVjpAlgorithm`) | +| `random_projection` | rank-1 random factors `(s̃, θ̃)` | O(P+H), **unbiased** | new (P4) | ### Axis 2 — `temporal_recursion` (how the trace advances in time) - `jacobian` — `ε ← D·ε + …` (D-RTRL, OSTL 'with-H', e-prop) - `scalar_leak` — `R̂ ← λ·R̂ + …`. The additive term is the same per-primitive - hidden→weight Jacobian contribution the `jacobian` recursion already uses, so - swapping the recursion is primitive-agnostic. This axis value is what remains - of OTPE after the algorithm itself was removed. + hidden→weight Jacobian contribution `jacobian` already uses, so swapping the + recursion is primitive-agnostic. This is what generalises out of the removed + OTPE; the restriction stack that made OTPE non-general does not come with it. - `none` — no temporal accumulation (feedforward regime) -Under `io_factorized` this axis is a **pair** `(x-side, f-side)`. The current -`IODimVjpAlgorithm` shares one `α` across `ε_x` and `ε_f` (derived by +Under `io_factorized` this axis is a **pair** `(x-side, f-side)`. +`IODimVjpAlgorithm` currently shares one `α` across `ε_x` and `ε_f` (derived by `_format_decay_and_rank`); splitting it into two independent decays is the only -change in this roadmap that touches a numerical path in existing code. - -### Axis 3 — `learning_signal` (where the signal comes from) +change in this roadmap that touches a numerical path in existing code. With the +split, the removed OTTT's coordinate — x-side leak, f-side instantaneous — +becomes reachable *and* primitive-generic, which the deleted implementation +never was. + +### Axis 3 — `recurrence_scope` (how much hidden↔hidden coupling enters the trace) + +**This axis already exists in the code as a boolean.** `_include_recurrent_mixing` +on `ETraceVjpAlgorithm` flows into `find_hidden_groups_*` and sets +`HiddenGroup.is_diagonal_recurrence = not include_recurrent_mixing`. Axis values, +naming the two that exist today: + +- `diagonal` (`_include_recurrent_mixing = False`; default for `D_RTRL`, + `pp_prop`, `EProp`) — recurrent ETP *mixing* primitives are excluded from the + transition, so it is position-diagonal by construction and the cheap + column-sum Jacobian (`jacrev_last_dim`) is exact. +- `coupled` (`_include_recurrent_mixing = True`; `OSTLRecurrent`) — those + primitives are traced into the transition, the recurrence becomes + cross-position coupled, and the true per-position block diagonal is extracted + explicitly (`block_diagonal_last_dim`). This retains strictly more than + `diagonal` — the diagonal entries now include contributions routed through the + recurrent weight, which `diagonal` drops entirely. +- `sparse_n` — the SnAp-n scale added in P3. + +SnAp-n is therefore a *generalisation of an existing knob* rather than a new +representation invented from scratch, which is what makes P3 tractable. + +**Open design question for P3.** Both existing values end up producing a +per-position block-diagonal Jacobian; they differ in what enters the transition +before the block diagonal is taken. SnAp-n instead parameterises the *n-step +influence* that is retained. Whether `coupled` lands on a particular `n`, or +stays a sibling value beside the `n` scale, is a design decision P3 must settle +explicitly — this document does not assume it maps onto one. + +### Axis 4 — `learning_signal` (where the signal comes from) - `symmetric` — `∂L/∂h` back-propagated through the readout (default) -- `random_feedback` — fixed random projection; promotes the existing - `FixedRandomFeedback` helper to a first-class strategy +- `random_feedback` — fixed random projection; promotes `FixedRandomFeedback` + from helper to first-class strategy - `modulatory` — three-factor: a scalar / low-dimensional neuromodulatory signal - (TD error, reward-prediction error) times the trace (new, S6) + (TD error, reward-prediction error) times the trace (new, P4) - `bootstrapped` — synthetic gradient / DNI: a learned estimate of the future - gradient (new, S6) + gradient (new, P4) -DRTP / target projection is **not** a value on this axis. See "Removals". +DRTP / target projection is deliberately **not** a value here. It was removed +with OSTTP and is not reintroduced. -### Axis 4 — `trace_filter` +### Axis 5 — `trace_filter` -`none` · `kappa` (e-prop's low-pass `ē ← κ·ē + ε`) +`none` · `kappa` (e-prop's low-pass `ē ← κ·ē + ε`; `EProp` filters the trace +internally today, `KappaFilter` remains a separate user-facing utility) -### Axis 5 — `update_schedule` +### Axis 6 — `update_schedule` `per_step` · `window(k)` · `sequence_end` ### Compatibility matrix -The axes are not fully orthogonal: `random_projection` carries UORO's own rank-1 -update and normalisation and cannot be paired with an arbitrary -`temporal_recursion`. `ETraceConfig` validates the `(factorization × -recursion)` combination at construction and raises a readable error listing the -legal pairings. The matrix is explicit data, not scattered `if` statements. - -## Coordinates of the existing algorithms - -| Algorithm | factorization | recursion | signal | filter | -|---|---|---|---|---| -| `D_RTRL` | per_param | jacobian | symmetric | none | -| `pp_prop` | io_factorized | jacobian | symmetric | none | -| `OSTLRecurrent` | per_param | jacobian | symmetric | none | -| `OSTLFeedforward` | io_factorized | none | symmetric | none | -| `EProp` | per_param | jacobian | symmetric \| random_feedback | **kappa** | -| ~~`OTPE`~~ (exact) | per_param | **scalar_leak** | symmetric | none | -| ~~`OTPE`~~ (F-OTPE) | io_factorized | **scalar_leak** | symmetric | none | -| ~~`OTTT`~~ | io_factorized | (x: scalar_leak, f: none) | symmetric | none | -| ~~`OSTTP`~~ | per_param | jacobian | target_projection | none | - -The three struck-through rows all whitelist dense-matmul primitives and are -single-step only; the table is the proof behind the removal decisions. -`OSTLRecurrent` matches -`D_RTRL` cell for cell — and the source says so directly ("delegates entirely to -`ParamDimVjpAlgorithm`"). - -**Deleting the classes does not delete the capability.** After the axes exist, -OTTT's coordinate is still reachable as a configuration — and reachable for -conv / sparse / lora, which the current OTTT rejects. The framework ships only -general mechanisms while the "one framework describes them all" claim gets -*stronger*, not weaker. - -## Removals - -### OTTT - -`ottt.py` hard-codes `_SUPPORTED_PRIMITIVES = {etp_mm_p, etp_mv_p}` and raises -`NotImplementedError` for lora / sparse / conv / element-wise relations, plus -`'OTTT v1 supports single-step only'`. Fails criterion 1. Its coordinate is -`io_factorized` with an x-side leak and an f-side that does not accumulate — -one point in the configuration space once axis 2 splits per side. - -### OSTTP - -`osttp.py` fails criterion 1 for engineering rather than mathematical reasons: -`y_target` has to reach every HiddenGroup through a bespoke path, and `B_list` -is hard-bound to the HiddenGroup count. DRTP itself is just `random_feedback` -applied to the target instead of the error. **Decision: remove the value as well -as the class** — axis 3 carries no `target_projection`, and no generic -`y_target` plumbing is introduced for it. - -Note for S6: `modulatory` needs a general external-signal injection path. That is -what OSTTP's `y_target` was reaching for and got wrong. Build it generically. - -### OTPE - -`otpe.py` fails criterion 1 harder than OTTT does, and says so itself: its -docstring states the published derivation is *"narrower than OTTT's"*. It -carries the same `_SUPPORTED_PRIMITIVES = {etp_mm_p, etp_mv_p}` whitelist and -the same `'OTPE v1 supports single-step only'` guard, hard-codes the dense outer -product (`jnp.einsum('bi,bo->bio', x, df_proj)`) instead of routing through the -per-primitive rule registry, rejects `num_state > 1` outright (so ALIF and any -adaptation variable are out), and errors on any relation touching more than one -HiddenGroup. On top of that it assumes a **single global time constant**, is -**feed-forward only**, and is gradient-exact for **one hidden layer**. - -For brain simulation those assumptions are backwards: heterogeneous time -constants and adaptation variables are the norm. What generalises is the -`scalar_leak` recursion, which survives as an axis-2 value; what does not is -OTPE-the-published-algorithm, which is `scalar_leak` plus that whole restriction -stack. - -### Verified deletion footprint - -| Target | Verification | -|---|---| -| `_algorithm/ottt.py`, `ottt_test.py` | — | -| `_algorithm/osttp.py`, `osttp_test.py` | — | -| `_algorithm/otpe.py`, `otpe_test.py` | — | -| `_common.py::PresynapticTrace` + its tests | grep confirms `ottt.py` is the only user | -| `_common.py::extract_y_target` + its tests | grep confirms `osttp.py` is the only user | -| imports / `__all__` in both `__init__.py` levels | `_algorithm/__init__.py`, `braintrace/__init__.py` | -| `docs/apis/algorithms.rst`, `docs/apis/index.rst` | — | - -Referencing files that need edits but not deletion: `oracle_models.py`, -`vjp_base.py`, `_compile.py`, `_compiler/graph.py`, `__init___test.py`, -`_compile_test.py`, and in `_algorithm/tests/`: `approx_correctness_test.py`, -`exact_correctness_test.py`, `public_api_test.py`, -`transform_correctness_test.py`. - -The `_get_update_aux` / auxiliary-data hook in `vjp_base.py` is **kept**. OSTTP -was its only consumer, but it is the general per-call side channel that S6's -`modulatory` signal needs; only the OSTTP-specific prose was rewritten. - -Explicitly kept: `KappaFilter` (standalone utility; its docstring already notes -e-prop filters internally) and `FixedRandomFeedback` (promoted to the axis-3 -`random_feedback` strategy). - -## Public API shape at 0.3.0 +The axes are not fully orthogonal, and pretending otherwise produces silently +wrong gradients: -```python -# Engines — the hosts for the axes -ETraceVjpAlgorithm # base class; hooks delegate to strategies -ParamDimVjpAlgorithm # factorization = per_param -IODimVjpAlgorithm # factorization = io_factorized +- `random_projection` carries UORO's own rank-1 update and normalisation; it + cannot be paired with an arbitrary `temporal_recursion`. +- `recurrence_scope` beyond the diagonal end demands a `temporal_recursion` that + actually propagates a Jacobian; pairing `n > 1` with `scalar_leak` or `none` + is meaningless. +- `io_factorized` constrains the contraction in `_solve_weight_gradients`; the + legal `(factorization, signal)` shapes the current implementations encode + implicitly must become explicit. -# braintrace's own contributions — keep the names -D_RTRL, pp_prop +`ETraceConfig` validates the combination at construction and raises an error +naming the legal pairings. The matrix is explicit data, not scattered `if`s. -# Configuration surface -ETraceConfig(factorization=..., recursion=..., signal=..., filter=..., schedule=...) +## Coordinates of the surviving algorithms -# Strategy protocols — third-party extension points -TraceFactorization, TemporalRecursion, LearningSignalSource, TraceFilter, UpdateSchedule +| Algorithm | factorization | recursion | recurrence_scope | signal | filter | +|---|---|---|---|---|---| +| `D_RTRL` | per_param | jacobian | diagonal | symmetric | none | +| `pp_prop` | io_factorized | jacobian | diagonal | symmetric | none | +| `OSTLRecurrent` | per_param | jacobian | **coupled** | symmetric | none | +| `OSTLFeedforward` | io_factorized | none | diagonal | symmetric | none | +| `EProp` | per_param | jacobian | diagonal | symmetric \| random_feedback | **kappa** | -# Literature presets — thin factories over ETraceConfig -EProp, OSTLRecurrent, OSTLFeedforward -``` +`update_schedule` is omitted from the table: every surviving algorithm is +`per_step`, so the column carries no information today. It becomes load-bearing +only once `window(k)` / `sequence_end` are implemented. -`braintrace.compile(model, algo, x0, **kw)` keeps its signature; `algo` -additionally accepts an `ETraceConfig`. +Every surviving algorithm occupies a distinct coordinate. `OSTLRecurrent` is +**not** a `D_RTRL` alias — the `recurrence_scope` column is what separates them, +and it is the column an earlier draft of this document was missing. That draft +argued `OSTLRecurrent` should be deleted as a zero-information alias; the +argument was wrong, and the class stays. -**Mechanism vs surface.** Internally the axes are strategy objects injected into -one engine — this is isomorphic to the existing template-method hooks -(`_compute_learning_signal`, `_solve_weight_gradients`, `_update_etrace_data`, -`_make_etrace_stepper`, `init_etrace_state`), so the delta is small and each -strategy is independently testable. Externally `ETraceConfig` parses into a -strategy combination, which keeps the user-facing API simple and makes the -configuration space enumerable — S7's benchmark suite gets that for free. +## Phases -A declarative learning-rule IR (rules as a compiler IR layer, third parties -adding rules without touching the engine) is the natural long-term direction but -is deliberately **out of scope**. With ~8 rules the correct abstraction boundary -is not yet observable. Revisit once S4–S6 push the count past ~15. +All land in 0.3.0. Ordering reflects dependencies, not separate releases. -## Compatibility +### P0 — Removal — **done** (`928219b`) -0.3.0 is a clean break. `OTTT`, `OSTTP` and `OTPE` are removed outright with no -shim; -the changelog documents the migration path. The repo already carries a `_legacy` -deprecation layer, and stacking another one only thickens the debt. +Deleted `ottt.py` / `osttp.py` / `otpe.py`, their tests, and their private +helpers; cleaned every reference across source, tests, docs and `AGENTS.md`. -## Sub-projects +Coverage was **repointed, not dropped**, wherever the assertion was about a +general property: the approximate-gradient descent backstop now runs on +`pp_prop(rank=1)` and `EProp(feedback='random')`; the one-step D_RTRL +equivalence tests now use `EProp(kappa_filter_decay=0)` (verified passing); +`public_api_test.py` gained a guard that the removed names stay gone. The +direction-alignment metric helpers were retained as the basis of P5. -All land in 0.3.0. Ordering below reflects the dependency structure, not -separate releases. +Verified: `pytest braintrace/` → 2062 passed, 1 skipped; `mypy braintrace` clean. -### S1 — Removal (do first) — **implemented** +### P1 — Compiler: multi-timescale and heterogeneous populations -Delete `ottt.py` / `osttp.py` / `otpe.py` and their helpers, clean every -reference. Small, mechanical, no numerical risk. Doing it before S2 shrinks the -compiler-refactor surface by three algorithms that would otherwise have to be -kept working and then deleted. +Make the surviving rules actually run on realistic brain-simulation models — +LIF + ALIF + multi-timescale synapses, heterogeneous leaks, multi-state +HiddenGroups — before widening the algorithm family. -Coverage that pointed at the removed algorithms was **repointed, not dropped**, -wherever the assertion was about a general property: +**Scope has to be reconstructed first.** `AGENTS.md` points at a +known-limitation findings list under `dev/`, but `dev/` is gitignored and absent +from the repository. The only machine-readable remnant is a single skipped test +(`approx_correctness_test.py::test_approximations_diverge_on_snn_multipopulation_DEFERRED`, +finding F-22). Everything else on that list has either been resolved — F-19 and +F-20 died with the algorithms they described; F-SINGLESTEP was resolved into a +positive direction-alignment assertion in `oracle_test.py` — or exists only as +prose in `AGENTS.md`. **The first task of P1 is to rebuild the list in-tree**, +from `AGENTS.md`'s summary plus a sweep of the test suite, so it stops living in +an untracked directory. Expect the rebuilt list to be shorter than the original: +part of it is already done. -- the approximate-gradient descent backstop now runs on `pp_prop(rank=1)` and - `EProp(feedback='random')` instead of `OTTT`/`OTPE`; -- the one-step D_RTRL equivalence tests now use `EProp(kappa_filter_decay=0)`, - which is D_RTRL's trace with no filter (verified: passes element-wise); -- `public_api_test.py` gained a guard asserting the removed names stay gone. +Known items from `AGENTS.md`: heterogeneous-population leak resolution, +multi-state HiddenGroups, approximation validity beyond shallow depth, +single-readout / feedback-shape assumptions, cross-algorithm equivalence gaps. -The direction-alignment metric helpers (`cosine_similarity`, `sign_agreement`, -`relative_magnitude`, `assert_direction_aligned`) are retained — they are the -basis of S7's benchmark suite. +This phase owns `hidden_group.py`'s Jacobian path, which P3 also needs — doing +it first is what keeps P3 from fighting a moving target. -**Acceptance:** full suite green; `grep -r "OTTT\|OSTTP\|OTPE\|PresynapticTrace\|extract_y_target"` -over `braintrace/` and `docs/` hits only `changelog.md`, the removal-notice -docstrings, and the guard test; mypy clean; `py.typed` intact. +**Acceptance:** the reconstructed limitation list is committed under +`docs/specs/`; every item on it has either a passing test or an explicitly +documented scope boundary; no expected-failure item silently remains. -### S2 — Compiler: multi-timescale and heterogeneous populations +### P2 — Axis decomposition -Work through the known-limitation list already recorded in `AGENTS.md` and the -`dev/` findings: heterogeneous-population leak resolution, multi-state -HiddenGroups, approximation validity beyond shallow depth, single-readout and -feedback-shape assumptions. This comes before the algorithm work so that the -existing rules actually run on realistic brain-simulation models (LIF + ALIF + -multi-timescale synapses) before the algorithm family is widened. +Turn the six axes into code: strategy protocols, `ETraceConfig` and its +compatibility matrix, engine hooks rewired to delegate to strategies, and the +three literature presets (`EProp`, `OSTLRecurrent`, `OSTLFeedforward`) rewritten +as thin factories over `ETraceConfig`. Includes the `IODimVjpAlgorithm` decay +split, and lifting `_include_recurrent_mixing` from a class attribute to an axis +value. -The Jacobian-representation part of this sub-project is a prerequisite for S4. - -**Acceptance:** each limitation item has either a passing test or an explicit, -documented scope boundary. No expected-failure item silently remains. - -### S3 — Axis decomposition - -Turn the five axes into code: the strategy protocols, `ETraceConfig` and its -compatibility matrix, the engine hooks rewired to delegate to strategies, and -the three literature presets (`EProp`, `OSTLRecurrent`, `OSTLFeedforward`) -rewritten as thin factories over `ETraceConfig`. Includes the -`IODimVjpAlgorithm` decay split — the one numerical change in this roadmap. - -Everything from S4 onward adds values to these axes, so this is the last piece -of groundwork. +Execution options (`vjp_method`, `fast_solve`, `trace_dtype`, `chunked_trace`, +`control_flow`) stay constructor parameters and stay out of `ETraceConfig`. **Acceptance:** - Element-wise equality against golden values frozen *before* the refactor, for - all five presets (`D_RTRL`, `pp_prop`, `EProp`, `OSTLRecurrent`, - `OSTLFeedforward`). + all five surviving algorithms. - `decay_or_rank=0.9` equals the new two-sided `(0.9, 0.9)` element-wise. - Illegal axis combinations raise a readable error naming the legal pairings, with test coverage. -- Every preset's coordinates are asserted against the table in this document, so - the table cannot silently drift from the code. +- Each preset's coordinates are asserted against the table in this document, so + the table cannot drift from the code. -### S4 — SnAp-n sparsity axis +### P3 — SnAp-n: generalise `recurrence_scope` Menick et al. 2021. Derive the n-step influence sparsity pattern automatically -from the compiler's jaxpr hidden→hidden reachability graph — this is the thing a -compiler can do that a hand-written library cannot, and it sits directly on the -paper's technical moat. +from the compiler's jaxpr hidden→hidden reachability graph — the thing a +compiler can do that a hand-written library cannot. -**The work is in the compiler, not the algorithm.** The current machinery keeps -only the block diagonal of `D^t` (`HiddenGroup.diagonal_jacobian`, -`block_diagonal_last_dim`); SnAp-n needs off-block-diagonal terms within n steps. -Assume the existing representation does not survive n>1 and budget for a new -sparse representation, shared with S2. +Concretely: replace the boolean `is_diagonal_recurrence` with an n-valued scope, +keep `jacrev_last_dim` as the `n = 1` fast path and `block_diagonal_last_dim` as +today's coupled path, and add the intermediate sparse representation for +`1 < n < diameter`. **Acceptance (two-sided squeeze):** -- `n = 1` equals the existing `D_RTRL` element-wise (regression guard). +- `n = 1` equals the current `D_RTRL` element-wise (regression guard). +- Whatever configuration expresses `recurrence_scope = coupled` after the + refactor equals the current `OSTLRecurrent` element-wise — regardless of + whether that configuration turns out to be a point on the `n` scale or a + sibling value beside it (regression guard for the existing `True` branch). - `n ≥ graph diameter` equals the **BPTT oracle** element-wise, on models whose hidden-state coupling the compiler fully captures. Full RTRL and BPTT compute the same total gradient, so `oracle.py` is the correct instrument and no - separate full-RTRL reference is needed. (A full-RTRL reference is worth - writing only to assert the influence matrix `dh^t/dθ` itself when localising a + separate full-RTRL reference is needed. (A full-RTRL reference is worth writing + only to assert the influence matrix `dh^t/dθ` itself when localising a divergence between trace and learning signal — optional debugging aid.) - Measured memory curve, monotone in n. -### S5 — UORO (unbiased random-projection estimator) +### P4 — UORO, three-factor and DNI -Tallec & Ollivier 2018. Adds the `random_projection` factorization with UORO's -rank-1 update and normalisation. Complements the existing biased diagonal -approximations. KF-RTRL / OK are optional lower-variance siblings on the same -axis value. +Three additions that share P2's axes and can proceed in parallel once P2 lands. -**Acceptance is statistical, not element-wise.** An unbiased estimator will not -match BPTT on any single run. The guard is: fixed model, fixed sequence, N random -seeds; the deviation of the mean gradient from BPTT shrinks as 1/√N (confidence -interval test). A single run asserts only shape, finiteness, and absence of NaN. -This requires statistical test infrastructure the repo does not have — count it -in the sub-project's cost. +**UORO** (Tallec & Ollivier 2018) — adds `random_projection` with UORO's rank-1 +update and normalisation, complementing the existing biased diagonal +approximations. KF-RTRL / OK are optional lower-variance siblings. -### S6 — Learning-signal axis: three-factor and DNI +**`modulatory`** — three-factor learning: trace × neuromodulatory signal (TD +error, reward-prediction error), enabling reward-based e-prop and online policy +gradient. Uses the retained `_get_update_aux` side channel. The injection path +must not bind to HiddenGroup count or readout shape — that binding is precisely +what made OSTTP non-general. -- `modulatory`: three-factor learning — trace × neuromodulatory signal (TD error, - reward-prediction error), enabling reward-based e-prop and online policy - gradient. Needs the general external-signal injection path noted above. -- `bootstrapped`: DNI / synthetic gradients (Jaderberg 2017) — a learned - bootstrap for the future-loss gradient that every online rule truncates away. - Requires an auxiliary network with its own training loop. +**`bootstrapped`** — DNI / synthetic gradients (Jaderberg 2017): a learned +bootstrap for the future-loss gradient every online rule truncates away. +Requires an auxiliary network with its own training loop. -**Acceptance is by degeneracy, plus task performance.** `modulatory` must equal -`symmetric` element-wise when the modulatory signal is set to `∂L/∂h`; -`bootstrapped` must equal `symmetric` when the synthesiser output is pinned to -the true value. Add an end-to-end RL smoke test. Neither is a "more accurate -gradient", so element-wise comparison against BPTT is the wrong instrument. +**Acceptance** — three different paradigms in one phase, do not mix them up: +- UORO: statistical. Fixed model, fixed sequence, N seeds; deviation of the mean + gradient from BPTT shrinks as 1/√N (confidence-interval test). A single run + asserts only shape, finiteness, absence of NaN. **This needs statistical test + infrastructure the repo does not have** — count it in the phase's cost. +- `modulatory`: must equal `symmetric` element-wise when the modulatory signal is + set to `∂L/∂h`. Plus a test whose HiddenGroup count differs from the signal + dimension, pinning that the OSTTP binding mistake was not recreated. +- `bootstrapped`: must equal `symmetric` when the synthesiser output is pinned to + the true value. Plus an end-to-end RL smoke test. -### S7 — Unified benchmark suite +### P5 — Unified benchmark suite Enumerate the `ETraceConfig` space and, for a fixed model, report gradient cosine similarity and relative deviation against BPTT, peak memory, per-step wall time, -and task metrics. Machine-readable output (JSON/CSV) plus a reproducible script. -This turns the paper's "comparable" claim into something executable. +and task metrics. Machine-readable output (JSON/CSV) plus a reproducible script, +reusing the direction metrics already in `oracle.py`. + +Should also retire the deferred F-22 finding: the SNN multi-population model zoo +that F-22 says is needed to expose the real bias of the IODim rank / ES decay / +random-feedback approximations is the same model zoo this suite needs. -Depends on S3's axes; grows as S4–S6 land. +Grows as P3 and P4 land. ## Three acceptance paradigms @@ -352,27 +353,33 @@ Mixing these up is the most likely way to get this roadmap wrong: | Paradigm | Applies to | Instrument | |---|---|---| -| Element-wise equality | S1 refactor, SnAp-1, SnAp-∞ | golden values / BPTT oracle | +| Element-wise equality | P2 refactor, SnAp-1, SnAp-coupled, SnAp-∞, degenerate `modulatory` / `bootstrapped` | golden values / BPTT oracle | | Statistical convergence | UORO and unbiased siblings | 1/√N confidence interval over seeds | -| Degeneracy + task metric | modulatory, bootstrapped | reduce to `symmetric`, then RL smoke test | +| Direction + task metric | genuinely approximate configurations | cosine / sign agreement, then descent and RL smoke tests | -The existing taxonomy in `AGENTS.md` (exact vs approximate) stays valid and is -extended by the statistical class. +The `AGENTS.md` taxonomy (exact vs approximate) stays valid and gains the +statistical class. ## Risk register -1. **Axis refactor changes numerics silently.** Mitigation: before touching the - engine, freeze reference gradients for all five presets as golden values; +1. **P2 changes numerics silently.** Mitigation: freeze reference gradients for + all five surviving algorithms as golden values *before* touching the engine; assert element-wise equality after. Separately, the `IODim` decay split must prove `decay_or_rank=0.9` equals the new `(0.9, 0.9)` element-wise. -2. **The sparse Jacobian representation does not generalise** (S2/S4). Mitigation: - design the representation once, in S2, with SnAp-n as the explicit consumer. -3. **Statistical tests are flaky in CI** (S5). Mitigation: fixed seeds, generous - intervals, and a separate slow-test marker. -4. **`modulatory` re-creates OSTTP's plumbing mistake** (S6). Mitigation: the - injection path must not bind to HiddenGroup count or readout shape; add a test - with a model whose HiddenGroup count differs from the signal dimension. -5. **0.3.0 carries the whole roadmap.** All breaking changes and all new - algorithm families ship together. Mitigation: sub-projects merge - independently behind the axis interfaces, and the benchmark suite (S7) runs on - every merge so regressions surface per sub-project rather than at release. +2. **P1 and P3 collide in `hidden_group.py`.** Both rework the hidden→hidden + Jacobian path. Mitigation: P1 owns that file and lands first; P3 consumes the + representation P1 leaves behind rather than introducing a second one. +3. **The P1 scope is not actually written down.** The findings list lives in an + untracked `dev/`, and the in-tree remnants are already stale — one docstring + still cross-referenced a test renamed some time ago. Mitigation: rebuilding + the list in-tree is P1's first task, not an assumption, and the rebuilt list + must be verified against the test suite rather than transcribed. +4. **Statistical tests are flaky in CI** (P4). Mitigation: fixed seeds, generous + intervals, a separate slow-test marker. +5. **`modulatory` recreates OSTTP's plumbing mistake** (P4). Mitigation: the + injection path must not bind to HiddenGroup count or readout shape; the + mismatched-dimension test above is the guard. +6. **0.3.0 carries the whole roadmap.** All breaking changes and all new + algorithm families ship together. Mitigation: phases merge independently + behind the axis interfaces, and P5 runs on every merge so regressions surface + per phase rather than at release. From db17054d6f5140f6a90398dca07f4283a114430c Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 11:50:35 +0800 Subject: [PATCH 03/13] Add the P1 spec: verification harness + limitation-list reconstruction P1 was scoped as compiler work on multi-timescale and heterogeneous populations. Scoping measurements against bc153da contradict that premise: the compiler already passes every one of those targets (num_state up to 5, per-neuron heterogeneous tau, E/I populations, multi-timescale synapses all match the BPTT oracle), and no defect was found in hidden_group.py's Jacobian path. What fails is the instrument. online_param_gradients with vjp_method='multi-step' over a full sequence returns bitwise BPTT for every algorithm at every hyperparameter -- correct semantics, since a full-window VJP has no truncation left to approximate, but it means any assertion whose subject is a learning-rule axis is vacuous there. That invalidates F-22's premise (a multi-population SNN model is bitwise-exact on the same path, so no model zoo can revive a dead knob) and F-21's stated cause. The spec refocuses P1 on the harness: negative-control oracle helpers, deterministic and spiking SNN model specs, an axis-discrimination meta-test, the in-tree findings list, the conv-bias IODim fix, and revised P2/P3 acceptance criteria that name a finite-window path. --- .../2026-07-25-p1-verification-harness.md | 287 ++++++++++++++++++ 1 file changed, 287 insertions(+) create mode 100644 docs/specs/2026-07-25-p1-verification-harness.md diff --git a/docs/specs/2026-07-25-p1-verification-harness.md b/docs/specs/2026-07-25-p1-verification-harness.md new file mode 100644 index 00000000..a56ad8dd --- /dev/null +++ b/docs/specs/2026-07-25-p1-verification-harness.md @@ -0,0 +1,287 @@ +# P1 — Verification harness: make the oracle axis-aware, rebuild the limitation list + +Status: approved, ready for implementation +Parent: [`2026-07-25-algorithm-axes-roadmap.md`](2026-07-25-algorithm-axes-roadmap.md) § P1 +Baseline: commit `bc153da` +Target release: 0.3.0 + +## Premise shift + +The roadmap scoped P1 as compiler work: make the surviving rules run on +realistic brain-simulation models — LIF + ALIF + multi-timescale synapses, +heterogeneous leaks, multi-state HiddenGroups — and own `hidden_group.py`'s +Jacobian path before widening the algorithm family. + +Scoping measurements against the actual code contradict that premise. The +compiler already passes every one of those targets. What fails is the +*instrument* used to judge it: the oracle path that nearly every gradient +assertion in the repository is written against cannot distinguish one learning +rule from another. It returns bitwise-identical BPTT gradients for `D_RTRL`, +`OSTLRecurrent`, `EProp` and `pp_prop`, at every hyperparameter setting. + +P1 therefore delivers a harness that can *see* learning-rule differences, then +re-establishes the known-limitation list against that harness. The +multi-timescale and heterogeneity claims are discharged by pinning them as +passing tests, not by fixing code. + +## Evidence + +All measurements below were run against `bc153da` on CPU. + +### The compiler handles P1's nominal targets + +`braintrace/_etrace_model_test.py` holds a ten-model SNN zoo (IF / LIF / ALIF × +Delta / ExpCu / ExpCo, with STD and STP variants) carrying units, multi-timescale +synapses (`tau_mem`, `tau_syn`, `tau_a`, `tau_std`, `tau_f`, `tau_d`) and an E/I +population split. Compiled under `D_RTRL`: + +| Model | HiddenGroups | `num_state` | ETP relations | +|---|---|---|---| +| `IF_Delta_Dense_Layer` | 1 | 1 | 1 | +| `LIF_ExpCu_Dense_Layer` | 1 | 2 | 1 | +| `ALIF_ExpCu_Dense_Layer` | 1 | 3 | 1 | +| `ALIF_ExpCo_Dense_Layer` (E/I) | 1 | 5 | 3 | + +With deterministic construction and an input scale that produces spikes, every +one of these — plus per-neuron heterogeneous `tau_mem` and `tau_a` — matches the +BPTT oracle. **No defect was found in `hidden_group.py`'s Jacobian path.** + +### The oracle's main entry point is blind to every learning-rule axis + +`tanh_rnn(n_in=3, n_rec=4, seed=0)`, `T=8`. Each cell is the relative deviation +*between two settings of the same knob* — how much the knob moves the gradient +at all: + +| Oracle path | `pp_prop` 0.99 vs 0.01 | `EProp` κ 0.0 vs 0.95 | +|---|---|---| +| `online_param_gradients` (multi-step, full window) | **0.0e+00** | **0.0e+00** | +| `chunked_online_param_gradients` (`chunk_size=2`) | 2.4e−03 | 2.3e−03 | +| `online_param_gradients_singlestep_naive` | 1.2e+00 | 3.0e+00 | + +Through the full-window multi-step path all four surviving algorithms return +gradients bitwise equal to BPTT. This is *correct* semantics, not a bug: a +window spanning the whole sequence has no truncation, so there is nothing left +for a trace approximation, a κ-filter or a recurrence scope to change. The +consequence is what matters — **any assertion whose subject is a learning-rule +axis is vacuous on that path.** + +`online_param_gradients(..., vjp_method='single-step')` is not a workaround: it +raises `NotImplementedError` on multi-step input data. The only live paths today +are `chunked_online_param_gradients` with `chunk_size < T` and +`online_param_gradients_singlestep_naive`. + +Scale of exposure: 95 `vjp_method='multi-step'` call sites across 20 test +modules. Most are legitimate — see [Out of scope](#out-of-scope). + +### F-22's premise is false + +F-22 defers exposing the IODim rank / ES decay / random-feedback bias until an +"SNN multi-population model zoo (LIF/ALIF, multi-layer random feedback)" +exists. Running precisely that model — `ALIF_ExpCo_Dense_Layer`, 3 ETP +relations, `num_state=5`, E/I split, `T=20` — still yields bitwise-exact +agreement for every algorithm, because the cause is the harness, not the model. +No model zoo can revive a dead knob. + +`approx_correctness_test.py::test_rank_decay_random_approximations_are_exact_on_rate_model_F21` +records the same misattribution in its docstring: *"these nominally-approximate +configs match BPTT element-wise on a single-HiddenGroup rate model. The model +cannot stress their approximation."* The model is not the reason; all four of +its configurations pass `vjp_method='multi-step'`. + +### Two defects in the SNN zoo, one in `io_dim_vjp.py` + +- **Non-determinism.** The `_etrace_model_test.py` constructors build weights + through unseeded `braintools.init.*`, which draws from the global + `brainstate.random` stream. `factory()` therefore returns a *different* model + on each call, so `bptt_param_gradients(factory, …)` and + `online_param_gradients(factory, …)` compare two different networks. This + violates the contract `oracle.py` states for its `model_factory` argument. + `oracle_models.py` specs are unaffected — they seed explicitly via + `jax.random.PRNGKey(seed)`. +- **Silence.** At the default input scale, 8 of 9 probed configurations never + reach threshold: spike rate 0.00, loss 0, gradients identically zero. A + BPTT-versus-online comparison then compares zero to zero and passes for any + algorithm. A related sub-case exists: `ALIF_Delta_Dense_Layer` at spike rate + 0.60 still produced a zero BPTT gradient, so spiking alone is not a + sufficient liveness criterion — the gradient norm is. +- **conv + trainable bias.** `pp_prop` / `IODimVjpAlgorithm` raises + `ValueError: Custom VJP bwd rule …` when a conv layer carries a trainable + bias, in either layout: the bwd rule returns the bias cotangent still shaped + per-position `(batch, *spatial, out_ch)` instead of reduced to `(out_ch,)`. + Already pinned by `oracle_test.py::test_pp_prop_conv_bias_known_limitation` + under `pytest.raises`. `D_RTRL` (param-dim) handles the same model exactly. + +## Deliverables + +### D1 — `docs/specs/2026-07-25-known-limitations.md` + +The in-tree findings list, replacing the untracked `dev/` one that `AGENTS.md` +points at. Every entry carries: ID, the claim, status, the test that pins it, +and the evidence. **Each entry is verified against the current test suite, not +transcribed.** + +Reconstructed disposition: + +| ID | Claim | Status | Pinned by | +|---|---|---|---| +| F-01 / F-04 | multi-state (`num_state ≥ 2`) HiddenGroups | covered, but by a vacuous assertion | `approx_correctness_test.py::test_d_rtrl_exact_on_two_state_group` — uses the full-window path; re-point | +| F-07 / F-08 / F-09 | OTTT / OTPE bias | dead — removed with the algorithms | — | +| F-17 | drifted implementation facts | resolved | `__init___test.py` Task 7 | +| F-19 / F-20 | OTTT / OSTTP exactness | dead — removed with the algorithms | — | +| F-21 | rank / decay / random-feedback exact on rate models | **misattributed** — cause is the harness, not the model | rewrite (see D2) | +| F-22 | approximation bias needs an SNN multi-population zoo | **premise false** — retire in P1, not P5 | re-point at a finite window | +| F-SCAN / F-SCAN-WEIGHT | weight inside control flow | resolved | `_compiler/base_test.py` | +| F-SINGLESTEP | single-step naive vs BPTT | resolved into a direction-alignment assertion | `oracle_test.py` | +| **F-23** | full-window multi-step oracle path is axis-blind | **new, active** | D2 | +| **F-24** | `_etrace_model_test.py` factories are non-deterministic | **new, active** | D3 | +| **F-25** | SNN zoo silent at default scale → vacuous comparisons | **new, active** | D3 | +| **F-26** | `pp_prop` / IODim raises on conv + trainable bias | **new, active** (pre-existing) | `oracle_test.py::test_pp_prop_conv_bias_known_limitation` | + +`AGENTS.md`'s prose items map on as follows, and the mapping is recorded in the +list so nothing survives only as prose: + +| `AGENTS.md` prose item | Disposition | +|---|---| +| approximation-mode validity beyond shallow depth | successor to F-21 / F-23 | +| heterogeneous-population leak resolution | resolved — pinned by D5 | +| target-signal threading under JIT | dead — died with `OSTTP`'s `y_target` path | +| single-readout / feedback-shape assumptions | verify in D5; document boundary if real | +| cross-algorithm equivalence coverage gaps | successor to F-23 | + +### D2 — Axis-aware oracle + +`oracle.py` changes, all additive: + +- **Window semantics documented per entry point.** Each of the four functions + states what it can and cannot detect. `online_param_gradients` gains an + explicit warning that with full-sequence input it equals BPTT for *any* + algorithm, making it a test of the compiler + ETP-rule stack, not of the rule. +- **`assert_gradients_differ(a, b, *, min_rel)`** — negative control. Asserts + two gradient trees are *not* equal, so a test that intends to exercise a knob + fails loudly when the knob is dead. +- **`assert_model_is_live(model_factory, inputs, *, min_norm)`** — asserts the + BPTT gradient norm is above threshold, so a silent network cannot make a + comparison vacuous. Keyed on gradient norm, not spike rate (see F-25). + +New `braintrace/_algorithm/tests/axis_discrimination_test.py` pins both halves +of F-23, for each pair of axis-distinct configurations: + +1. the full-window multi-step path collapses them to identical gradients — + locking the semantics as understood rather than accidental; and +2. a finite window (`chunked`, `chunk_size < T`) separates them. + +This is the meta-test whose absence let F-21 be misattributed. + +### D3 — Oracle-ready SNN model specs + +New `ModelSpec` factories in `oracle_models.py` wrapping the existing +`_etrace_model_test.py` layers for oracle use, fixing F-24 and F-25 at the +boundary rather than mutating the layer classes: + +- deterministic construction — explicit seeding, so repeated `factory()` calls + are bitwise identical (asserted); +- a recorded input scale that drives the network to spike, with liveness + asserted through `assert_model_is_live`; +- coverage: LIF / ALIF × ExpCu / Delta / STD / STP, an E/I `ExpCo` + multi-population, and per-neuron heterogeneous `tau_mem` / `tau_a`. + +`ModelSpec` gains an optional input-construction field so the scale travels with +the model instead of living in each test. Existing specs keep their current +behaviour by default. + +This is the model zoo F-22 asked for — built because the correctness tests need +realistic models, not because it exposes approximation bias, which it does not. + +### D4 — SNN correctness tests + +New `braintrace/_algorithm/tests/snn_model_correctness_test.py`: `D_RTRL` exact +against BPTT across the D3 specs, each test paired with a liveness guard. + +These tests use the full-window multi-step path deliberately. `D_RTRL` is an +exact algorithm, so the subject of the assertion is the compiler and the ETP +per-primitive rules on a realistic model — which is exactly what that path +tests well. F-23 forbids the path only where the subject is a learning-rule +axis, which is not the case here. Where these tests do compare *across* +algorithms, they use a finite window. + +The existing `d_rtrl_test.py::test_snn_*_vjp` and +`pp_prop_test.py::test_snn_*_vjp` assert nothing about gradient values — they +`print(grads)`. They keep their role as shape/smoke coverage and are left +alone; D4 is where value assertions live. + +### D5 — Multi-timescale and heterogeneity pinned as passing tests + +Within D4's module: `num_state` 1 through 5, multi-timescale synapses, +per-neuron heterogeneous leaks, and E/I populations each asserted exact against +BPTT. This discharges the `AGENTS.md` prose limitations by test rather than by +fix. Also probes the single-readout / feedback-shape assumption; if it is real, +it becomes a documented scope boundary with a named finding. + +### D6 — F-22 retired + +Re-point F-22's assertion at a finite window and delete the skipped +`test_approximations_diverge_on_snn_multipopulation_DEFERRED`, replacing it with +a live test that measures the approximation's actual bias through +`chunked_online_param_gradients`. F-21's docstring is corrected to attribute the +exactness to the harness. + +### D7 — `io_dim_vjp.py` conv-bias fix + +Reduce the bias cotangent to the bias's own shape in the custom-VJP bwd rule, +then promote `test_pp_prop_conv_bias_known_limitation` from `pytest.raises` to +an exactness assertion. If the fix is not contained within the bwd rule, F-26 +becomes an explicitly documented scope boundary instead — the roadmap's P1 +acceptance permits either, but not silence. + +### D8 — Roadmap and `AGENTS.md` updates + +- P2 and P3 acceptance criteria revised to mandate a finite-window oracle path. + As written ("equals the current `OSTLRecurrent` element-wise", "`n = 1` equals + the current `D_RTRL` element-wise") they would pass for any algorithm and + would not guard the refactor at all. +- A lessons-learned section in the roadmap recording the verified context above. +- `AGENTS.md` § Known limitations repointed from `dev/` to the in-tree list. + +## Testing strategy + +Test-first, and specifically *vacuity-first*: for each finding, the test that +exposes the vacuity is written and seen to fail before the fix. Two invariants +apply to every new gradient assertion: + +- **Liveness** — the reference gradient norm is asserted nonzero. A comparison + against a zero reference asserts nothing. +- **Discrimination** — an assertion whose subject is a learning-rule axis is + paired with a check that the axis actually moves the gradient on the path + used. + +Verification for the phase: full `pytest braintrace/` and `mypy braintrace` +clean, matching the P0 bar (2062 passed, 1 skipped at `928219b`; the skip count +drops to 0 once D6 lands). + +## Acceptance + +Restating the roadmap's P1 criteria against this scope: + +- the reconstructed limitation list is committed under `docs/specs/`; +- every item on it has either a passing test or an explicitly documented scope + boundary; +- no expected-failure item silently remains — the one skipped test is retired; +- F-23 is pinned in both directions by `axis_discrimination_test.py`; +- D3's specs are asserted deterministic and live; +- multi-timescale, heterogeneous-leak, multi-state and E/I claims each have a + named passing test; +- P2/P3 acceptance criteria in the roadmap name a finite-window path. + +## Out of scope + +- **Migrating the 95 existing `multi-step` call sites.** For exact algorithms + and the `_op/` primitive tests, a full-window comparison against BPTT is a + legitimate test of the compiler and ETP-rule stack — that is most of the 95. + Only assertions whose *subject* is a learning-rule axis get re-pointed. +- **`hidden_group.py` Jacobian rework.** The roadmap gave P1 ownership of this + file to keep P3 off a moving target. No defect was found, so P1 leaves it + unchanged; P3 inherits today's representation. Risk 2 in the roadmap's + register is thereby retired, not mitigated. +- SnAp-n, new axes, `ETraceConfig`, UORO, three-factor, DNI — P2 through P4. +- Statistical test infrastructure — P4. +- Benchmark suite — P5. From 2becf7032d698136eb5e88dd196a363cb9d4b9a7 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 11:58:23 +0800 Subject: [PATCH 04/13] Add the P1 implementation plan; sharpen the F-23 finding Nine TDD tasks: oracle negative controls, the axis-discrimination meta-test, deterministic and live SNN specs, SNN correctness tests, F-22 retirement, the conv-bias IO-dim fix, the in-tree findings list, roadmap and AGENTS.md updates, full-suite verification. Also corrects the spec's framing of F-23. It is not a new discovery about the code: chunked_online_param_gradients' docstring already states that the full-sequence call is exact reverse-mode and that chunking is "the oracle that actually validates trace correctness". The finding is that the assertions written against online_param_gradients do not heed it, and that online_param_gradients carries no such warning. --- .../2026-07-25-p1-implementation-plan.md | 1614 +++++++++++++++++ .../2026-07-25-p1-verification-harness.md | 12 +- 2 files changed, 1625 insertions(+), 1 deletion(-) create mode 100644 docs/specs/2026-07-25-p1-implementation-plan.md diff --git a/docs/specs/2026-07-25-p1-implementation-plan.md b/docs/specs/2026-07-25-p1-implementation-plan.md new file mode 100644 index 00000000..977375ba --- /dev/null +++ b/docs/specs/2026-07-25-p1-implementation-plan.md @@ -0,0 +1,1614 @@ +# P1 Verification Harness Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make the gradient oracle able to distinguish one learning rule from another, then rebuild the known-limitation list against it. + +**Architecture:** Additive changes to `braintrace/_algorithm/oracle.py` (negative-control helpers, pytree/unit-aware comparison, documented window semantics) and `oracle_models.py` (deterministic, spiking SNN `ModelSpec`s). Two new test modules pin the findings. One contained fix in `io_dim_vjp.py` reduces broadcast-parameter cotangents to the parameter's own shape. Documentation lands in `docs/specs/` and `AGENTS.md`. + +**Tech Stack:** Python, JAX, `brainstate`, `brainunit`, `brainpy`, `braintools`, pytest. + +Spec: [`2026-07-25-p1-verification-harness.md`](2026-07-25-p1-verification-harness.md) + +## Global Constraints + +- Tests are co-located as siblings with the `_test.py` suffix — never a `tests/` directory with `test_*.py` prefixes. The exception is the pre-existing `braintrace/_algorithm/tests/` directory, which holds cross-cutting suites; new cross-cutting suites go there to match. +- Never drive a model with a bare Python `for`/`while` loop when it runs repeatedly — use `brainstate.transform.for_loop` / `scan` / `jit`. Test-support chunk loops over a handful of chunks are exempt and already exist in `oracle.py`. +- Use `brainstate.random`, never `jax.random` directly, for random number generation in models. +- All public functions use NumPy-style docstrings with the canonical section order. +- SNN models require an integration step: wrap calls in `brainstate.environ.context(dt=0.1 * u.ms)`. +- Gradient trees may be nested pytrees and may carry `brainunit` units. Never call `jnp.asarray` on a gradient tree without flattening leaves and stripping units first. +- Commit messages must not contain a `Co-Authored-By` trailer. +- Verification bar for the phase: `pytest braintrace/` fully green and `mypy braintrace` clean. + +--- + +### Task 1: Oracle negative controls and unit-aware comparison + +**Files:** +- Modify: `braintrace/_algorithm/oracle.py` +- Test: `braintrace/_algorithm/oracle_test.py` + +**Interfaces:** +- Consumes: nothing from earlier tasks. +- Produces, all in `braintrace._algorithm.oracle`: + - `flat_gradient_leaves(tree: dict) -> dict[str, jax.Array]` — flattens a `{state_path: pytree}` gradient dict to `{label: unit-stripped array}`. + - `gradient_norm(tree: dict) -> float` + - `relative_deviation(actual: dict, expected: dict) -> float` + - `assert_model_is_live(model_factory, inputs, *, min_norm: float = 1e-8) -> float` — returns the BPTT gradient norm. + - `assert_gradients_differ(a: dict, b: dict, *, min_rel: float = 1e-6) -> float` — returns the observed relative deviation. + - `assert_param_gradients_close` gains pytree and unit support; its signature is unchanged. + +- [ ] **Step 1: Write the failing tests** + +Append to `braintrace/_algorithm/oracle_test.py`: + +```python +# --- P1: negative-control helpers ------------------------------------------- + +def test_flat_gradient_leaves_handles_nested_and_units(): + """Gradient trees are nested dicts and may carry units; the flattener must + yield plain arrays keyed by a stable label.""" + import brainunit as u + from braintrace._algorithm.oracle import flat_gradient_leaves + tree = { + ('syn', 'comm', 'weight'): {'weight': jnp.ones((2, 3)) * u.mS, + 'bias': jnp.zeros((3,)) * u.mS}, + ('w',): jnp.arange(4.0), + } + flat = flat_gradient_leaves(tree) + assert len(flat) == 3 + for arr in flat.values(): + assert not isinstance(arr, u.Quantity) + assert sorted(k.split('|')[0] for k in flat) == ['syn/comm/weight', + 'syn/comm/weight', 'w'] + + +def test_gradient_norm_and_relative_deviation(): + from braintrace._algorithm.oracle import gradient_norm, relative_deviation + a = {('w',): jnp.array([3.0, 4.0])} + b = {('w',): jnp.array([3.0, 4.0])} + assert gradient_norm(a) == pytest.approx(5.0, abs=1e-6) + assert relative_deviation(a, b) == pytest.approx(0.0, abs=1e-12) + c = {('w',): jnp.array([0.0, 4.0])} + assert relative_deviation(a, c) == pytest.approx(3.0 / 5.0, abs=1e-6) + + +def test_assert_model_is_live_passes_on_live_model(): + from braintrace._algorithm.oracle import assert_model_is_live + from braintrace._algorithm.oracle_models import tanh_rnn + spec = tanh_rnn(n_in=3, n_rec=4, seed=0) + xs = jnp.asarray(np.random.RandomState(0).randn(4, 3).astype('float32')) + norm = assert_model_is_live(spec.factory, xs) + assert norm > 0.0 + + +def test_assert_model_is_live_rejects_a_dead_model(): + """A model whose output is constant has a zero gradient, so any comparison + against it asserts nothing. The guard must reject it.""" + from braintrace._algorithm.oracle import assert_model_is_live + + def factory(): + class Dead(brainstate.nn.Module): + def __init__(self): + super().__init__() + self.w = brainstate.ParamState(jnp.zeros((3, 3))) + self.h = brainstate.HiddenState(jnp.zeros((1, 3))) + + def update(self, x): + # output does not depend on w at all + self.h.value = jnp.zeros((1, 3)) + return self.h.value + + return Dead() + + xs = jnp.zeros((3, 3)) + with pytest.raises(AssertionError, match='gradient norm'): + assert_model_is_live(factory, xs) + + +def test_assert_gradients_differ_flags_a_dead_knob(): + from braintrace._algorithm.oracle import assert_gradients_differ + a = {('w',): jnp.array([1.0, 2.0])} + assert_gradients_differ(a, {('w',): jnp.array([1.0, 5.0])}) + with pytest.raises(AssertionError, match='indistinguishable'): + assert_gradients_differ(a, {('w',): jnp.array([1.0, 2.0])}) + + +def test_assert_param_gradients_close_supports_nested_unit_trees(): + """The pre-existing helper only handled flat, unitless dicts. SNN models + have nested weight dicts carrying units.""" + import brainunit as u + from braintrace._algorithm.oracle import assert_param_gradients_close + a = {('syn',): {'weight': jnp.ones((2, 2)) * u.mS, 'bias': jnp.zeros(2) * u.mS}} + b = {('syn',): {'weight': jnp.ones((2, 2)) * u.mS, 'bias': jnp.zeros(2) * u.mS}} + assert_param_gradients_close(a, b, atol=1e-6) + c = {('syn',): {'weight': jnp.full((2, 2), 2.0) * u.mS, 'bias': jnp.zeros(2) * u.mS}} + with pytest.raises(AssertionError, match='maxabsdiff'): + assert_param_gradients_close(a, c, atol=1e-6) +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `pytest braintrace/_algorithm/oracle_test.py -k "flat_gradient or gradient_norm or model_is_live or gradients_differ or nested_unit" -v` + +Expected: FAIL — `ImportError: cannot import name 'flat_gradient_leaves'` for the new helpers, and `test_assert_param_gradients_close_supports_nested_unit_trees` fails inside `jnp.asarray` on a dict. + +- [ ] **Step 3: Implement the helpers** + +In `braintrace/_algorithm/oracle.py`, add `import brainunit as u` to the imports, then insert these functions immediately before `assert_param_gradients_close`: + +```python +def flat_gradient_leaves(tree) -> dict: + """Flatten a path-keyed gradient tree into ``{label: plain array}``. + + Gradient trees returned by the oracle are keyed by ``ParamState`` path and + each value may itself be a pytree (a ``Linear`` contributes ``weight`` and + ``bias``) whose leaves may carry ``brainunit`` units. Comparisons need plain + arrays, so this strips both layers. + + Parameters + ---------- + tree : dict + Mapping from state path tuple to gradient pytree. + + Returns + ------- + dict + Mapping from ``'a/b|.leaf'`` label to a unit-stripped ``jax.Array``. + """ + out = {} + for key, value in tree.items(): + path_label = '/'.join(map(str, key)) if isinstance(key, tuple) else str(key) + for leaf_path, leaf in jax.tree_util.tree_flatten_with_path(value)[0]: + label = f'{path_label}|{jax.tree_util.keystr(leaf_path)}' + out[label] = jnp.asarray(u.get_mantissa(leaf)) + return out + + +def gradient_norm(tree) -> float: + """Euclidean norm of every leaf of a gradient tree, taken together.""" + leaves = flat_gradient_leaves(tree) + total = sum(float((arr.astype('float64') ** 2).sum()) for arr in leaves.values()) + return float(np.sqrt(total)) + + +def relative_deviation(actual, expected) -> float: + """``||actual - expected|| / ||expected||`` over all leaves jointly. + + Raises + ------ + AssertionError + If the two trees do not have the same set of leaf labels. + """ + a = flat_gradient_leaves(actual) + e = flat_gradient_leaves(expected) + if set(a) != set(e): + raise AssertionError( + f'gradient trees have different leaves: {sorted(set(a) ^ set(e))}') + num = sum(float(((a[k].astype('float64') - e[k].astype('float64')) ** 2).sum()) + for k in e) + den = sum(float((e[k].astype('float64') ** 2).sum()) for k in e) + if den == 0.0: + return float('inf') if num > 0.0 else 0.0 + return float(np.sqrt(num) / np.sqrt(den)) + + +def assert_model_is_live(model_factory, inputs, *, min_norm: float = 1e-8) -> float: + """Assert the BPTT gradient of ``model_factory`` on ``inputs`` is non-trivial. + + A comparison against an all-zero reference gradient passes for every + algorithm and therefore asserts nothing. SNN models are the common way to + hit this: at a low input scale the neurons never reach threshold, the loss is + zero, and so is the gradient. Spiking alone is not sufficient — a model can + spike and still have a zero gradient — so the criterion is the gradient norm + itself. + + Parameters + ---------- + model_factory : Callable[[], brainstate.nn.Module] + Zero-arg factory returning an uninitialized model. + inputs : jax.Array + ``(T, ...)`` input sequence. + min_norm : float, optional + Minimum acceptable BPTT gradient norm. + + Returns + ------- + float + The observed BPTT gradient norm. + + Raises + ------ + AssertionError + If the norm is at or below ``min_norm``. + """ + norm = gradient_norm(bptt_param_gradients(model_factory, inputs)) + if not (norm > min_norm): + raise AssertionError( + f'model is not live: BPTT gradient norm {norm:.3e} <= {min_norm:.3e}. ' + 'Any gradient comparison on this model/input pair is vacuous.' + ) + return norm + + +def assert_gradients_differ(a, b, *, min_rel: float = 1e-6) -> float: + """Assert two gradient trees are *distinguishable* — a negative control. + + Use this whenever a test intends to exercise a learning-rule knob. If the + oracle path chosen cannot see the knob, this fails loudly instead of letting + the test pass vacuously. See finding F-23. + + Parameters + ---------- + a, b : dict + Path-keyed gradient trees. + min_rel : float, optional + Minimum relative deviation between them. + + Returns + ------- + float + The observed relative deviation. + + Raises + ------ + AssertionError + If the deviation is below ``min_rel``. + """ + rel = relative_deviation(a, b) + if not (rel >= min_rel): + raise AssertionError( + f'gradients are indistinguishable: relative deviation {rel:.3e} < ' + f'{min_rel:.3e}. The knob under test does not move the gradient on ' + 'this oracle path.' + ) + return rel +``` + +Then replace the body of `assert_param_gradients_close` (currently `oracle.py:196-213`) so it compares flattened leaves: + +```python +def assert_param_gradients_close(actual, expected, *, atol=1e-4, rtol=0.0, keys=None): + """Assert two param-gradient trees match, with a per-leaf diagnostic on failure. + + ``keys`` restricts the comparison to a subset of top-level state paths (e.g. + only ETP params). When None, every key present in ``expected`` is compared. + Nested pytrees and unit-carrying leaves are supported. + """ + compare_keys = list(expected.keys()) if keys is None else list(keys) + sub_actual = {k: actual[k] for k in compare_keys} + sub_expected = {k: expected[k] for k in compare_keys} + a = flat_gradient_leaves(sub_actual) + e = flat_gradient_leaves(sub_expected) + if set(a) != set(e): + raise AssertionError( + f'gradient trees have different leaves: {sorted(set(a) ^ set(e))}') + failures = [] + for label in sorted(e): + if not bool(jnp.allclose(a[label], e[label], atol=atol, rtol=rtol)): + failures.append( + f" {label}: maxabsdiff={float(jnp.max(jnp.abs(a[label] - e[label]))):.3e}") + if failures: + raise AssertionError( + "param gradients differ beyond tolerance " + f"(atol={atol}, rtol={rtol}):\n" + "\n".join(failures) + ) +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `pytest braintrace/_algorithm/oracle_test.py -v` + +Expected: PASS, including all pre-existing tests in the module — the `assert_param_gradients_close` change is a generalization, so flat unitless dicts behave identically. + +- [ ] **Step 5: Document the window semantics of each oracle entry point** + +This is the core of finding F-23 and it must be visible where people write tests. Note that `chunked_online_param_gradients` **already** documents it ("this is the oracle that actually validates trace correctness"); the gap is that `online_param_gradients` does not warn, and its stale `dev/` cross-reference is unresolvable. + +In `online_param_gradients`, replace the docstring with: + +```python + """Total sequence gradient from an online algorithm via the multi-step VJP path. + + ``algo_factory(model)`` must return an algorithm whose ``__call__`` accepts a + ``braintrace.MultiStepData`` and returns the stacked per-step outputs. The loss + ``(out ** 2).sum()`` over the whole stacked output equals the BPTT sequence loss. + + Warnings + -------- + **This path is blind to every learning-rule axis** (finding F-23). One + whole-sequence call makes the within-call gradient exact reverse-mode, so the + eligibility trace only enters at the sequence boundary — of which there is + none. Every algorithm therefore returns gradients bitwise equal to BPTT, at + every hyperparameter setting: ``D_RTRL``, ``OSTLRecurrent``, ``EProp`` at any + ``kappa_filter_decay`` and ``pp_prop`` at any ``decay_or_rank`` are + indistinguishable here. + + That makes this function a good test of the *compiler and ETP per-primitive + rules* — it is exactly the right instrument for asserting an exact algorithm + reproduces BPTT on a realistic model — and the wrong instrument for any + assertion whose subject is a trace factorization, a temporal recursion, a + recurrence scope, a filter, or a learning signal. For those use + :func:`chunked_online_param_gradients` with ``chunk_size`` < ``T``, and guard + the test with :func:`assert_gradients_differ`. + + See Also + -------- + chunked_online_param_gradients : finite-window path; sees the trace. + assert_gradients_differ : negative control for a knob that must matter. + """ +``` + +In `online_param_gradients_singlestep_naive`, replace the stale `dev/` reference (`oracle.py:180`) so the docstring resolves in-tree: + +```python + """Naive 'single-step' total gradient: sum of per-step grad((algo(x_t)**2).sum()). + + Kept to document finding F-SINGLESTEP — this recipe does NOT equal BPTT even + for the exact D_RTRL algorithm, while the multi-step path does. This is the + most aggressive finite window (one step), so it is maximally sensitive to + learning-rule axes and maximally divergent from BPTT. + + See Also + -------- + docs/specs/2026-07-25-known-limitations.md : F-SINGLESTEP and F-23. + """ +``` + +- [ ] **Step 6: Verify nothing regressed** + +Run: `pytest braintrace/_algorithm/ -x -q` + +Expected: PASS. + +- [ ] **Step 7: Commit** + +```bash +git add braintrace/_algorithm/oracle.py braintrace/_algorithm/oracle_test.py +git commit -m "Add oracle negative controls and document the axis-blind window + +assert_model_is_live rejects a zero reference gradient; assert_gradients_differ +rejects a knob that does not move the gradient. assert_param_gradients_close now +handles nested pytrees and unit-carrying leaves, which SNN models need. + +online_param_gradients now warns that a full-sequence window equals BPTT for any +algorithm (F-23), so it tests the compiler and ETP rules rather than the rule, +and points at the finite-window path for axis assertions. Also replaces an +unresolvable dev/ cross-reference." +``` + +--- + +### Task 2: Pin F-23 with an axis-discrimination meta-test + +**Files:** +- Create: `braintrace/_algorithm/tests/axis_discrimination_test.py` + +**Interfaces:** +- Consumes: `oracle.assert_gradients_differ`, `oracle.relative_deviation`, `oracle.online_param_gradients`, `oracle.chunked_online_param_gradients`, `oracle_models.tanh_rnn`. +- Produces: nothing later tasks depend on. + +- [ ] **Step 1: Write the test module** + +Create `braintrace/_algorithm/tests/axis_discrimination_test.py`: + +```python +# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""F-23: which oracle paths can see a learning-rule axis, and which cannot. + +A full-sequence ``MultiStepData`` call makes the within-call gradient exact +reverse-mode, so the eligibility trace enters only at a sequence boundary that +does not exist. Every algorithm then returns BPTT. This module pins that in both +directions, so a future test cannot silently assert an approximation's behaviour +on a path that cannot observe it -- which is how F-21 came to attribute the +effect to the model instead of the harness. +""" + +import brainstate +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import braintrace +from braintrace._algorithm.oracle import ( + assert_gradients_differ, + chunked_online_param_gradients, + online_param_gradients, + relative_deviation, +) +from braintrace._algorithm.oracle_models import tanh_rnn + +T = 8 +CHUNK = 2 + + +def _spec(): + return tanh_rnn(n_in=3, n_rec=4, seed=0) + + +def _inputs(): + return jnp.asarray(np.random.RandomState(0).randn(T, 3).astype('float32')) + + +# Pairs of configurations that differ ONLY in a learning-rule axis value. +AXIS_PAIRS = { + # Axis 1/2: IO-dim trace factorization strength (decay). + 'pp_prop decay': ( + lambda m: braintrace.pp_prop(m, decay_or_rank=0.99, vjp_method='multi-step'), + lambda m: braintrace.pp_prop(m, decay_or_rank=0.01, vjp_method='multi-step'), + ), + # Axis 5: trace filter (kappa). + 'EProp kappa': ( + lambda m: braintrace.EProp(m, kappa_filter_decay=0.0, vjp_method='multi-step'), + lambda m: braintrace.EProp(m, kappa_filter_decay=0.95, vjp_method='multi-step'), + ), + # Axis 3: recurrence scope (diagonal vs coupled). + 'recurrence scope': ( + lambda m: braintrace.D_RTRL(m, vjp_method='multi-step'), + lambda m: braintrace.OSTLRecurrent(m, vjp_method='multi-step'), + ), +} + + +@pytest.mark.parametrize('axis', sorted(AXIS_PAIRS)) +def test_full_window_multistep_cannot_see_any_axis(axis): + """The full-window path collapses axis-distinct configs to identical + gradients. Pinned so the semantics are understood rather than accidental -- + if this ever starts failing, the window semantics changed and every + assertion written against this path needs review.""" + spec, xs = _spec(), _inputs() + lo, hi = AXIS_PAIRS[axis] + g_lo = online_param_gradients(spec.factory, xs, algo_factory=lo) + g_hi = online_param_gradients(spec.factory, xs, algo_factory=hi) + rel = relative_deviation(g_lo, g_hi) + assert rel == 0.0, ( + f'{axis}: full-window multi-step distinguished two configurations ' + f'(rel={rel:.3e}); F-23 assumed it cannot. Re-check the oracle docs.' + ) + + +@pytest.mark.parametrize('axis', sorted(AXIS_PAIRS)) +def test_finite_window_does_see_the_axis(axis): + """A chunked window makes the trace enter at every chunk boundary, so the + axis becomes observable. This is the path axis assertions must use.""" + spec, xs = _spec(), _inputs() + lo, hi = AXIS_PAIRS[axis] + g_lo = chunked_online_param_gradients( + spec.factory, xs, algo_factory=lo, chunk_size=CHUNK) + g_hi = chunked_online_param_gradients( + spec.factory, xs, algo_factory=hi, chunk_size=CHUNK) + assert_gradients_differ(g_lo, g_hi, min_rel=1e-6) + + +def test_full_window_still_reproduces_bptt_for_an_exact_algorithm(): + """The corollary that makes the full-window path useful: it is the right + instrument for asserting an exact algorithm matches BPTT.""" + from braintrace._algorithm.oracle import ( + assert_param_gradients_close, + bptt_param_gradients, + ) + spec, xs = _spec(), _inputs() + g_bptt = bptt_param_gradients(spec.factory, xs) + g_online = online_param_gradients( + spec.factory, xs, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + assert_param_gradients_close(g_online, g_bptt, atol=1e-4) +``` + +- [ ] **Step 2: Run the module** + +Run: `pytest braintrace/_algorithm/tests/axis_discrimination_test.py -v` + +Expected: PASS — all 7 tests. If `test_finite_window_does_see_the_axis[recurrence scope]` fails, that is a real signal worth recording: it would mean `_include_recurrent_mixing` is unobservable even at `chunk_size=2` on `tanh_rnn`. In that case raise `T` to 16 and lower `CHUNK` to 1; if it still fails, mark that one parametrization `xfail` with a reason naming the axis, and add it to the findings list in Task 7 as a new active finding rather than deleting the test. + +- [ ] **Step 3: Commit** + +```bash +git add braintrace/_algorithm/tests/axis_discrimination_test.py +git commit -m "Pin F-23: which oracle paths can see a learning-rule axis + +Asserts both directions -- the full-window multi-step path collapses +axis-distinct configurations to bitwise-identical gradients, and a chunked +window separates them. This is the meta-test whose absence let F-21 attribute +the effect to the model rather than the harness." +``` + +--- + +### Task 3: Deterministic, live SNN model specs + +**Files:** +- Modify: `braintrace/_algorithm/oracle_models.py` +- Test: `braintrace/_algorithm/oracle_models_test.py` (create — no sibling test exists yet) + +**Interfaces:** +- Consumes: `oracle.assert_model_is_live` from Task 1. +- Produces, in `braintrace._algorithm.oracle_models`: + - `ModelSpec` gains two optional fields: `input_scale: float = 1.0` and `needs_dt: bool = False`. + - `ModelSpec.make_inputs(self, T: int, n_in: int, *, seed: int = 0) -> jax.Array` — batched, scaled, non-negative inputs. + - Spec factories, each `(n_in=4, n_rec=5, seed=7) -> ModelSpec`: + `snn_if_delta`, `snn_lif_expcu`, `snn_alif_expcu`, `snn_alif_delta`, + `snn_lif_std_expcu`, `snn_lif_stp_expcu`, `snn_alif_expco_ei`, + `snn_lif_expcu_heterogeneous`, `snn_alif_expcu_heterogeneous`. + - `SNN_SPECS: dict[str, Callable[..., ModelSpec]]` — name → factory, for parametrization. + +**Why a wrapper and not a fix to the layer classes:** `_etrace_model_test.py`'s constructors call unseeded `braintools.init.*`, which draws from the global `brainstate.random` stream, so `factory()` returns a different model per call (F-24). Seeding inside the wrapper keeps the layer classes untouched — they are used by many existing tests whose behaviour must not shift. + +- [ ] **Step 1: Write the failing tests** + +Create `braintrace/_algorithm/oracle_models_test.py`: + +```python +# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""The SNN specs must be deterministic (F-24) and live (F-25) before any +gradient assertion built on them means anything.""" + +import brainstate +import brainunit as u +import jax.numpy as jnp +import pytest + +from braintrace._algorithm.oracle import ( + assert_model_is_live, + flat_gradient_leaves, +) +from braintrace._algorithm.oracle_models import SNN_SPECS + + +@pytest.mark.parametrize('name', sorted(SNN_SPECS)) +def test_snn_spec_construction_is_deterministic(name): + """F-24: the underlying layer classes seed from the global RNG, so two + factory() calls must be pinned to produce identical weights.""" + spec = SNN_SPECS[name]() + with brainstate.environ.context(dt=0.1 * u.ms): + w1 = flat_gradient_leaves( + {k: v.value for k, v in spec.factory().states(brainstate.ParamState).items()}) + w2 = flat_gradient_leaves( + {k: v.value for k, v in spec.factory().states(brainstate.ParamState).items()}) + assert set(w1) == set(w2) + for key in w1: + assert bool(jnp.allclose(w1[key], w2[key])), f'{name}: {key} differs across calls' + + +@pytest.mark.parametrize('name', sorted(SNN_SPECS)) +def test_snn_spec_is_live(name): + """F-25: at the default input scale these networks never spike, so their + gradients are identically zero and every comparison is vacuous. Each spec + records a scale that produces a non-trivial gradient.""" + spec = SNN_SPECS[name]() + with brainstate.environ.context(dt=0.1 * u.ms): + xs = spec.make_inputs(6, 4) + norm = assert_model_is_live(spec.factory, xs, min_norm=1e-6) + assert norm > 1e-6 + + +def test_default_input_scale_is_documented_as_dead(): + """The counterpart of the above: pins *why* the scale field exists. At + scale 1.0 the same model has a zero gradient.""" + from braintrace._algorithm.oracle import gradient_norm, bptt_param_gradients + spec = SNN_SPECS['lif_expcu']() + with brainstate.environ.context(dt=0.1 * u.ms): + dead_xs = spec.make_inputs(6, 4) / spec.input_scale # undo the scaling + assert gradient_norm(bptt_param_gradients(spec.factory, dead_xs)) == 0.0 +``` + +- [ ] **Step 2: Run to verify failure** + +Run: `pytest braintrace/_algorithm/oracle_models_test.py -v` + +Expected: FAIL — `ImportError: cannot import name 'SNN_SPECS'`. + +- [ ] **Step 3: Extend `ModelSpec`** + +In `braintrace/_algorithm/oracle_models.py`, replace the `ModelSpec` dataclass (currently at `oracle_models.py:29-40`) with: + +```python +@dataclass(frozen=True) +class ModelSpec: + """A zero-arg model factory plus metadata about its parameters. + + ``factory()`` returns a freshly constructed, *uninitialized* model with + deterministic weights. Callers must call + ``brainstate.nn.init_all_states(model, batch_size=...)`` themselves. + + Attributes + ---------- + factory : Callable[[], brainstate.nn.Module] + Deterministic zero-arg model constructor. + etp_param_keys : tuple of tuple + Parameter paths routed through an ETP primitive. + plain_param_keys : tuple of tuple + Parameter paths used via plain JAX ops, hence excluded from ETP. + input_scale : float, optional + Multiplier applied by :meth:`make_inputs`. Spiking models need a scale + well above 1.0 to reach threshold at all; below it the loss and the + gradient are identically zero and any comparison is vacuous (F-25). + batched_input : bool, optional + Whether :meth:`make_inputs` emits a leading batch axis of 1. SNN layers + concatenate the input with the recurrent spike vector, so their ranks + must match; the rate models broadcast instead and do not need it. + """ + + factory: Callable[[], brainstate.nn.Module] + etp_param_keys: Tuple[tuple, ...] # routed through an ETP primitive + plain_param_keys: Tuple[tuple, ...] # used via plain JAX ops (excluded from ETP) + input_scale: float = 1.0 + batched_input: bool = False + + def make_inputs(self, T: int, n_in: int, *, seed: int = 0): + """Build a ``(T, [1,] n_in)`` input sequence at this spec's scale. + + Values are non-negative so that spiking models receive net excitatory + drive; a zero-mean drive largely cancels and leaves the network silent. + """ + rng = np.random.RandomState(seed) + shape = (T, 1, n_in) if self.batched_input else (T, n_in) + return self.input_scale * jnp.asarray(np.abs(rng.randn(*shape)).astype('float32')) +``` + +Add `import numpy as np` to the module imports if absent. + +- [ ] **Step 4: Add the SNN spec factories** + +Append to `braintrace/_algorithm/oracle_models.py`: + +```python +# --------------------------------------------------------------------------- +# SNN specs: the realistic-model end of the zoo. +# +# These wrap the layer classes in ``braintrace/_etrace_model_test.py`` for +# oracle use. Two things have to be fixed at this boundary: +# +# * F-24 -- those constructors call unseeded ``braintools.init.*``, which draws +# from the global ``brainstate.random`` stream, so ``factory()`` returns a +# different model on every call and a BPTT-vs-online comparison would compare +# two different networks. Each factory re-seeds before constructing. +# * F-25 -- at unit input scale the neurons never reach threshold, so the loss +# and the gradient are identically zero. Each spec records the scale that +# makes it live; ``oracle_models_test.py`` asserts both properties. +# +# The layer classes themselves are left untouched: many existing tests depend +# on their current behaviour. +# --------------------------------------------------------------------------- + +_SNN_SEED = 7 +_SNN_SCALE = 20.0 + + +def _snn_spec(cls, n_in, n_rec, seed, **kwargs) -> ModelSpec: + """Wrap an SNN layer class as a deterministic, live ``ModelSpec``.""" + + def factory(): + brainstate.random.seed(seed) + return cls(n_in, n_rec, **kwargs) + + return ModelSpec( + factory=factory, + etp_param_keys=(), # discovered by the compiler; not asserted per-spec + plain_param_keys=(), + input_scale=_SNN_SCALE, + batched_input=True, + ) + + +def snn_if_delta(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """IF neuron, delta synapse. Single hidden state (``num_state == 1``).""" + from braintrace._etrace_model_test import IF_Delta_Dense_Layer + return _snn_spec(IF_Delta_Dense_Layer, n_in, n_rec, seed) + + +def snn_alif_delta(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """ALIF neuron, delta synapse. Membrane + adaptation (``num_state == 2``).""" + from braintrace._etrace_model_test import ALIF_Delta_Dense_Layer + return _snn_spec(ALIF_Delta_Dense_Layer, n_in, n_rec, seed) + + +def snn_lif_expcu(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """LIF neuron, exponential current synapse. Two timescales: tau_mem, tau_syn.""" + from braintrace._etrace_model_test import LIF_ExpCu_Dense_Layer + return _snn_spec(LIF_ExpCu_Dense_Layer, n_in, n_rec, seed) + + +def snn_alif_expcu(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """ALIF + exponential current synapse. Three timescales, ``num_state == 3``.""" + from braintrace._etrace_model_test import ALIF_ExpCu_Dense_Layer + return _snn_spec(ALIF_ExpCu_Dense_Layer, n_in, n_rec, seed) + + +def snn_lif_std_expcu(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """LIF + short-term depression. Adds tau_std as a fourth timescale.""" + from braintrace._etrace_model_test import LIF_STDExpCu_Dense_Layer + return _snn_spec(LIF_STDExpCu_Dense_Layer, n_in, n_rec, seed) + + +def snn_lif_stp_expcu(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """LIF + short-term plasticity. Adds tau_f and tau_d.""" + from braintrace._etrace_model_test import LIF_STPExpCu_Dense_Layer + return _snn_spec(LIF_STPExpCu_Dense_Layer, n_in, n_rec, seed) + + +def snn_alif_expco_ei(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """ALIF with an excitatory/inhibitory population split and conductance + synapses. The heterogeneous-population case: separate E and I projections + produce several ETP relations feeding one hidden group.""" + from braintrace._etrace_model_test import ALIF_ExpCo_Dense_Layer + return _snn_spec(ALIF_ExpCo_Dense_Layer, n_in, n_rec, seed) + + +def snn_lif_expcu_heterogeneous( + n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED +) -> ModelSpec: + """LIF whose membrane time constant differs per neuron. + + The heterogeneous-leak case: ``tau_mem`` is a length-``n_rec`` vector, so no + single global leak exists for the transition to factor out. + """ + from braintrace._etrace_model_test import LIF_ExpCu_Dense_Layer + tau_mem = jnp.linspace(3.0, 12.0, n_rec) * u.ms + return _snn_spec(LIF_ExpCu_Dense_Layer, n_in, n_rec, seed, tau_mem=tau_mem) + + +def snn_alif_expcu_heterogeneous( + n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED +) -> ModelSpec: + """ALIF with per-neuron membrane *and* adaptation time constants.""" + from braintrace._etrace_model_test import ALIF_ExpCu_Dense_Layer + return _snn_spec( + ALIF_ExpCu_Dense_Layer, n_in, n_rec, seed, + tau_mem=jnp.linspace(3.0, 12.0, n_rec) * u.ms, + tau_a=jnp.linspace(60.0, 150.0, n_rec) * u.ms, + ) + + +SNN_SPECS = { + 'if_delta': snn_if_delta, + 'alif_delta': snn_alif_delta, + 'lif_expcu': snn_lif_expcu, + 'alif_expcu': snn_alif_expcu, + 'lif_std_expcu': snn_lif_std_expcu, + 'lif_stp_expcu': snn_lif_stp_expcu, + 'alif_expco_ei': snn_alif_expco_ei, + 'lif_expcu_heterogeneous': snn_lif_expcu_heterogeneous, + 'alif_expcu_heterogeneous': snn_alif_expcu_heterogeneous, +} +``` + +Add `import brainunit as u` to the module imports if absent. + +- [ ] **Step 5: Run the tests** + +Run: `pytest braintrace/_algorithm/oracle_models_test.py -v` + +Expected: PASS for determinism on all 9 specs and liveness on most. + +**`alif_delta` is expected to fail liveness.** Scoping measurements found it spiking at rate 0.60 while its BPTT gradient norm was still 0 — spiking is not sufficient for a non-zero gradient. Handle it by tuning that one spec rather than deleting it: raise its scale, and if the gradient stays zero, lower `V_th` via a kwarg. If neither produces a live gradient, mark it `pytest.param('alif_delta', marks=pytest.mark.xfail(reason='F-27: ALIF+delta spikes but yields a zero BPTT gradient; surrogate-gradient window never opens', strict=True))` in both parametrizations, and record F-27 as a new active finding in Task 7. Do not silently drop the spec. + +- [ ] **Step 6: Commit** + +```bash +git add braintrace/_algorithm/oracle_models.py braintrace/_algorithm/oracle_models_test.py +git commit -m "Add deterministic, live SNN model specs for oracle use + +Fixes F-24 and F-25 at the spec boundary rather than in the layer classes, +which many existing tests depend on. Each spec re-seeds the global RNG before +construction, so repeated factory() calls are bitwise identical, and records the +input scale needed to produce a non-zero gradient at all. Covers LIF/ALIF x +delta/ExpCu/STD/STP, an E/I conductance split, and per-neuron heterogeneous +tau_mem and tau_a. Determinism and liveness are both asserted." +``` + +--- + +### Task 4: SNN correctness — multi-state, multi-timescale, heterogeneous, E/I + +**Files:** +- Create: `braintrace/_algorithm/tests/snn_model_correctness_test.py` + +**Interfaces:** +- Consumes: `SNN_SPECS` and `ModelSpec.make_inputs` from Task 3; `assert_model_is_live`, `assert_param_gradients_close`, `assert_gradients_differ` from Task 1. +- Produces: nothing later tasks depend on. + +This task discharges the `AGENTS.md` prose limitations *by test*: heterogeneous-population leak resolution, multi-state HiddenGroups, and approximation validity beyond a single relation. + +- [ ] **Step 1: Write the test module** + +Create `braintrace/_algorithm/tests/snn_model_correctness_test.py`: + +```python +# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Gradient correctness on realistic SNN models: multi-timescale synapses, +per-neuron heterogeneous leaks, multi-state HiddenGroups, and E/I populations. + +These are the claims ``AGENTS.md`` carried only as prose ("heterogeneous- +population leak resolution", "multi-state HiddenGroups"). They are discharged +here as passing tests rather than as fixes -- the compiler already handles them. + +Assertion paths, per F-23: an *exact* algorithm versus BPTT uses the full-window +multi-step path, whose subject is the compiler and the ETP per-primitive rules +and which is the right instrument for that. Anything comparing *across* +algorithms uses a finite window. +""" + +import brainstate +import brainunit as u +import pytest + +import braintrace +from braintrace._algorithm.oracle import ( + assert_gradients_differ, + assert_model_is_live, + assert_param_gradients_close, + bptt_param_gradients, + chunked_online_param_gradients, + online_param_gradients, +) +from braintrace._algorithm.oracle_models import SNN_SPECS + +T = 6 +N_IN = 4 +N_REC = 5 +ATOL = 1e-4 + +# alif_delta is excluded pending F-27 (see oracle_models_test.py); if Task 3 +# made it live, delete this tuple and the `set(...) - _NOT_LIVE` below. +_NOT_LIVE = {'alif_delta'} +_LIVE_SPECS = sorted(set(SNN_SPECS) - _NOT_LIVE) + + +def _setup(name): + spec = SNN_SPECS[name]() + xs = spec.make_inputs(T, N_IN) + return spec, xs + + +@pytest.mark.parametrize('name', _LIVE_SPECS) +def test_d_rtrl_matches_bptt_on_snn_models(name): + """D_RTRL is exact, so it must reproduce BPTT on every realistic model: + multi-timescale synapses, heterogeneous leaks, E/I populations and + HiddenGroups with num_state from 1 to 5.""" + spec, xs = _setup(name) + with brainstate.environ.context(dt=0.1 * u.ms): + assert_model_is_live(spec.factory, xs, min_norm=1e-6) + g_bptt = bptt_param_gradients(spec.factory, xs) + g_online = online_param_gradients( + spec.factory, xs, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + assert_param_gradients_close(g_online, g_bptt, atol=ATOL) + + +@pytest.mark.parametrize('name', ['lif_expcu_heterogeneous', + 'alif_expcu_heterogeneous']) +def test_heterogeneous_leaks_do_not_break_exactness(name): + """A per-neuron time constant leaves no single global leak for the + transition to factor out. The compiler takes a true Jacobian, so exactness + survives -- this is what retires the AGENTS.md prose item.""" + spec, xs = _setup(name) + with brainstate.environ.context(dt=0.1 * u.ms): + assert_model_is_live(spec.factory, xs, min_norm=1e-6) + g_bptt = bptt_param_gradients(spec.factory, xs) + g_online = online_param_gradients( + spec.factory, xs, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + assert_param_gradients_close(g_online, g_bptt, atol=ATOL) + + +@pytest.mark.parametrize('name,expected_min_state', [ + ('if_delta', 1), + ('lif_expcu', 2), + ('alif_expcu', 3), + ('alif_expco_ei', 5), +]) +def test_multi_state_hidden_groups_are_discovered(name, expected_min_state): + """Pins the structural facts the exactness tests above rest on: these models + really do form multi-state HiddenGroups, so the per-state axis is exercised + rather than assumed.""" + spec, xs = _setup(name) + with brainstate.environ.context(dt=0.1 * u.ms): + model = spec.factory() + brainstate.nn.init_all_states(model, batch_size=1) + algo = braintrace.D_RTRL(model, vjp_method='multi-step') + algo.compile_graph(xs[0]) + assert len(algo.graph.hidden_groups) >= 1 + assert max(hg.num_state for hg in algo.graph.hidden_groups) >= expected_min_state + assert len(algo.graph.hidden_param_op_relations) >= 1 + + +def test_ei_population_split_yields_multiple_relations(): + """The E/I model routes separate excitatory and inhibitory projections into + one hidden group, so the compiler must record more than one ETP relation.""" + spec, xs = _setup('alif_expco_ei') + with brainstate.environ.context(dt=0.1 * u.ms): + model = spec.factory() + brainstate.nn.init_all_states(model, batch_size=1) + algo = braintrace.D_RTRL(model, vjp_method='multi-step') + algo.compile_graph(xs[0]) + assert len(algo.graph.hidden_param_op_relations) >= 2 + + +@pytest.mark.parametrize('name', ['lif_expcu', 'alif_expco_ei']) +def test_approximation_is_measurable_on_snn_models(name): + """The genuinely approximate configuration must be *distinguishable* from + the exact one on a realistic model -- via a finite window, which is the only + path that can see it (F-23). This is what F-22 was really asking for.""" + spec, xs = _setup(name) + with brainstate.environ.context(dt=0.1 * u.ms): + g_exact = chunked_online_param_gradients( + spec.factory, xs, chunk_size=2, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + g_approx = chunked_online_param_gradients( + spec.factory, xs, chunk_size=2, + algo_factory=lambda m: braintrace.pp_prop( + m, decay_or_rank=0.5, vjp_method='multi-step')) + assert_gradients_differ(g_exact, g_approx, min_rel=1e-6) +``` + +- [ ] **Step 2: Run the module** + +Run: `pytest braintrace/_algorithm/tests/snn_model_correctness_test.py -v` + +Expected: PASS. If `test_d_rtrl_matches_bptt_on_snn_models` fails for a spec, that is a genuine compiler finding — do not loosen `ATOL`. Record it as a new finding in Task 7 with the model, the failing leaf and the observed deviation, and mark that parametrization `xfail(strict=True)`. + +If `test_approximation_is_measurable_on_snn_models` fails, `chunk_size=2` is too coarse to expose the approximation on that model: try `chunk_size=1`, then raise `T` to 12. If it still cannot be exposed, that is the substantive successor to F-22 and belongs in the findings list as active — say so explicitly rather than removing the test. + +- [ ] **Step 3: Commit** + +```bash +git add braintrace/_algorithm/tests/snn_model_correctness_test.py +git commit -m "Assert gradient correctness on realistic SNN models + +Discharges the AGENTS.md prose limitations by test: heterogeneous per-neuron +leaks, multi-timescale synapses, multi-state HiddenGroups (num_state 1 to 5) and +E/I populations all reproduce BPTT under the exact algorithm. Structural facts +are pinned alongside, so the exactness claims rest on discovered relations +rather than assumption. Cross-algorithm comparisons use a finite window per +F-23; every test carries a liveness guard." +``` + +--- + +### Task 5: Retire F-22, correct F-21 + +**Files:** +- Modify: `braintrace/_algorithm/tests/approx_correctness_test.py` + +**Interfaces:** +- Consumes: `assert_gradients_differ`, `chunked_online_param_gradients` from Task 1. +- Produces: nothing later tasks depend on. + +- [ ] **Step 1: Replace the skipped F-22 test with a live one** + +In `braintrace/_algorithm/tests/approx_correctness_test.py`, delete the whole skipped test (lines 154-161, the `@pytest.mark.skip(...)` decorator through `def test_approximations_diverge_on_snn_multipopulation_DEFERRED(): pass`) and put this in its place: + +```python +def test_approximations_are_measurable_through_a_finite_window(): + """F-22, retired. The finding deferred this until an "SNN multi-population + model zoo" existed, on the theory that the model was what made the + approximations look exact. That premise was wrong: a multi-population SNN + model (ALIF + E/I conductance split, 3 ETP relations, num_state 5) is + *also* bitwise-exact through the full-window path, because the cause is the + oracle path and not the model (F-23). No model can revive a knob the + harness cannot see. + + Through a finite window the same nominally-approximate configurations that + F-21 finds exact become measurably different from the exact algorithm, on + the very same rate model. That is the assertion F-22 wanted. + """ + spec = tanh_rnn(n_in=3, n_rec=4, seed=0) + inputs = _inputs(8, 3) + g_exact = chunked_online_param_gradients( + spec.factory, inputs, chunk_size=2, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + for name, algo_factory in _EXACT_ON_RATE.items(): + g_approx = chunked_online_param_gradients( + spec.factory, inputs, chunk_size=2, algo_factory=algo_factory) + assert_gradients_differ(g_exact, g_approx, min_rel=1e-9) +``` + +Add `assert_gradients_differ` and `chunked_online_param_gradients` to the existing `from braintrace._algorithm.oracle import (...)` block at the top of the file. + +- [ ] **Step 2: Correct F-21's docstring** + +Replace the docstring of `test_rank_decay_random_approximations_are_exact_on_rate_model_F21` (lines 144-145) with: + +```python + """F-21: these nominally-approximate configs match BPTT element-wise here. + + The cause is the *oracle path*, not the model. All four configurations run + through ``online_param_gradients`` with ``vjp_method='multi-step'`` over the + whole sequence, where the within-call gradient is exact reverse-mode and the + eligibility trace never enters -- so every algorithm returns BPTT at every + hyperparameter setting (F-23). An earlier reading of this test attributed the + exactness to the model being a single-HiddenGroup rate model and concluded + that an SNN multi-population zoo was needed to expose the bias (F-22); that + conclusion was wrong, and F-22 is retired by + ``test_approximations_are_measurable_through_a_finite_window`` below. + + The test is kept because the equality is still a real property of this path + and worth pinning. + """ +``` + +Also update the module docstring's `(F-21/F-22)` reference (line 19) to `(F-21/F-22/F-23)` and the comment block at lines 127-130, replacing "We assert exactness here and defer the genuine approximation stress to an SNN multi-population model zoo (F-22)." with "We assert exactness here; the cause is the full-window oracle path (F-23), and the genuine approximation stress is asserted through a finite window below." + +- [ ] **Step 3: Run the module** + +Run: `pytest braintrace/_algorithm/tests/approx_correctness_test.py -v` + +Expected: PASS, with no skipped tests. Confirm with `pytest braintrace/_algorithm/tests/approx_correctness_test.py -v -rs` that the skip list is empty. + +- [ ] **Step 4: Commit** + +```bash +git add braintrace/_algorithm/tests/approx_correctness_test.py +git commit -m "Retire F-22 and correct F-21's attribution + +F-22 deferred measuring the approximations' bias until an SNN multi-population +model zoo existed. The premise was wrong: such a model is bitwise-exact on the +same full-window path, because the path is what cannot see the approximation. +Replaces the skipped placeholder with a live assertion through a finite window, +on the same rate model, and corrects F-21's docstring to name the harness rather +than the model as the cause." +``` + +--- + +### Task 6: Fix the conv-bias cotangent shape in the IO-dim path + +**Files:** +- Modify: `braintrace/_algorithm/io_dim_vjp.py` +- Test: `braintrace/_algorithm/oracle_test.py:488-517` + +**Interfaces:** +- Consumes: nothing from earlier tasks. +- Produces: nothing later tasks depend on. + +**Diagnosis (already done — do not re-derive).** `pp_prop` on a conv layer with a trainable bias raises: + +``` +ValueError: Custom VJP bwd rule must produce an output with the same type as the +args tuple of the primal function, but at output[1][('b',)] the bwd rule produced +an output of type float64[1,8,3] corresponding to an input of type float64[3] +``` + +`_conv_xy_to_dw` in `braintrace/_op/conv.py:315` deliberately returns the bias Jacobian **per output position**; its docstring states "We store this per-position (no spatial sum here); the sum is performed inside `_conv_dt_to_t` during trace propagation". The param-dim path calls `_conv_dt_to_t` and so reduces it — which is why `D_RTRL` handles conv+bias exactly. The IO-dim solver calls `xy_to_dw` directly at solve time (`_solve_IO_dim_weight_gradients`, `io_dim_vjp.py:418`) and never applies that reduction, so the un-reduced array reaches JAX's `custom_vjp` type check. + +Fix the IO-dim solver, not the conv rule: reducing a produced gradient leaf to its parameter's own shape by summing the extra leading axes is the standard broadcast-gradient reduction and is primitive-agnostic, so it does not whitelist conv. + +- [ ] **Step 1: Turn the pinned limitation into a failing exactness test** + +In `braintrace/_algorithm/oracle_test.py`, replace `test_pp_prop_conv_bias_known_limitation` (lines 488-517) with: + +```python +def test_pp_prop_conv_bias_matches_bptt(): + """Formerly ``test_pp_prop_conv_bias_known_limitation`` (finding F-26). + + ``_conv_xy_to_dw`` returns the bias Jacobian per output position by design; + the param-dim path reduces it in ``_conv_dt_to_t``, but the IO-dim solver + calls ``xy_to_dw`` directly at solve time and used to hand the un-reduced + ``(batch, *spatial, out_ch)`` array back to ``custom_vjp``, which rejected it + against the bias's own ``(out_ch,)``. The IO-dim solver now reduces every + produced leaf to its parameter's shape. At T=1 there is no history to factor, + so pp_prop must match BPTT exactly. + """ + factory, seed = _FAMILIES['conv_nwc_bias'] + with brainstate.environ.context(precision=64): + xs = _xs_for('conv_nwc_bias', 1, seed) + g_bptt = bptt_param_gradients(factory, xs) + g_online = online_param_gradients_singlestep_naive( + factory, xs, + algo_factory=lambda m: braintrace.pp_prop( + m, decay_or_rank=0.9, vjp_method='single-step'), + ) + for key in g_bptt: + rel = _rel_err(g_bptt[key], g_online[key]) + assert rel < _TOL, f'conv_nwc_bias T=1 {key}: pp_prop vs BPTT rel={rel:.3e}' +``` + +Also re-include `conv_nwc_bias` in the two parametrizations that exclude it, at lines 443 and 465: change `sorted(set(_FAMILIES) - {'conv_nwc_bias'})` to `sorted(_FAMILIES)` in both, and delete the paragraph in `test_pp_prop_singlestep_exact_at_t1_across_families`'s docstring beginning "``conv_nwc_bias`` is excluded here". + +- [ ] **Step 2: Run to verify it fails** + +Run: `pytest braintrace/_algorithm/oracle_test.py -k conv -v` + +Expected: FAIL with `ValueError: Custom VJP bwd rule must produce an output with the same type ...`. + +- [ ] **Step 3: Reduce produced leaves to the parameter's shape** + +In `braintrace/_algorithm/io_dim_vjp.py`, add this helper immediately above `_solve_IO_dim_weight_gradients` (before line 418): + +```python +def _reduce_to_param_shape(grad: jax.Array, param: jax.Array) -> jax.Array: + """Sum a produced gradient leaf down to its parameter's own shape. + + Some ``xy_to_dw`` rules deliberately return a *per-position* instantaneous + Jacobian for parameters that broadcast over positions — a conv bias is the + canonical case, and :func:`braintrace._op.conv._conv_xy_to_dw` documents that + the spatial sum is deferred. The param-dim path performs that sum during + trace propagation; the IO-dim path contracts at solve time and must perform + it here, or ``custom_vjp`` rejects the shape mismatch (finding F-26). + + This is the standard broadcast-gradient reduction — sum the extra leading + axes, then any axis the parameter holds as a singleton — so it is + primitive-agnostic and a no-op whenever the shapes already agree. + """ + g_shape = u.math.shape(grad) + p_shape = u.math.shape(param) + if g_shape == p_shape: + return grad + extra = len(g_shape) - len(p_shape) + if extra < 0: + return grad + out = u.math.sum(grad, axis=tuple(range(extra))) if extra else grad + squeeze = tuple( + i for i, (gd, pd) in enumerate(zip(u.math.shape(out), p_shape)) + if pd == 1 and gd != 1 + ) + if squeeze: + out = u.math.sum(out, axis=squeeze, keepdims=True) + return out +``` + +Then apply it where the solver hands its result to the routing step. Replace the final line of the `for group in relation.hidden_groups:` body (`io_dim_vjp.py:538`, `_route_grads_by_path(relation, dg_dict, weight_vals, dG_weights)`) with: + +```python + # Reduce per-position leaves to their parameter's own shape before + # routing; see _reduce_to_param_shape (finding F-26). + dg_dict = { + key: _reduce_to_param_shape(value, weights_dict[key]) + for key, value in dg_dict.items() + } + _route_grads_by_path(relation, dg_dict, weight_vals, dG_weights) +``` + +- [ ] **Step 4: Run the conv tests** + +Run: `pytest braintrace/_algorithm/oracle_test.py -k conv -v` + +Expected: PASS. + +- [ ] **Step 5: Verify no other primitive regressed** + +The reduction touches every IO-dim gradient, so the whole primitive matrix must be re-run. + +Run: `pytest braintrace/_algorithm/oracle_test.py braintrace/_algorithm/io_dim_vjp_test.py braintrace/_algorithm/pp_prop_test.py braintrace/_op/ -q` + +Expected: PASS. If a `sparse`, `lora`, `grouped` or `elemwise` test now fails, the reduction is too eager for that layout — narrow it to apply only when `len(g_shape) > len(p_shape)` (drop the singleton-squeeze branch) and re-run. + +- [ ] **Step 6: Commit** + +```bash +git add braintrace/_algorithm/io_dim_vjp.py braintrace/_algorithm/oracle_test.py +git commit -m "Fix F-26: reduce per-position gradient leaves in the IO-dim solver + +Conv's xy_to_dw rule returns the bias Jacobian per output position by design and +defers the spatial sum to _conv_dt_to_t, which only the param-dim path calls. +The IO-dim solver contracts at solve time, so it handed the un-reduced +(batch, *spatial, out_ch) array to custom_vjp, which rejected it against the +bias's own (out_ch,). Reduces every produced leaf to its parameter's shape with +the standard broadcast-gradient reduction -- primitive-agnostic, so conv is not +whitelisted -- and promotes the pinned limitation to an exactness assertion. +conv_nwc_bias rejoins the two pp_prop family parametrizations." +``` + +--- + +### Task 7: The in-tree known-limitations list + +**Files:** +- Create: `docs/specs/2026-07-25-known-limitations.md` + +**Interfaces:** +- Consumes: the outcomes of Tasks 2-6, including any new findings they surfaced (F-27 and any `xfail` recorded there). +- Produces: a document Task 8 links to from `AGENTS.md`. + +Write this task **last among the code tasks**, so every status below is transcribed from an observed test result rather than predicted. + +- [ ] **Step 1: Confirm every claim against the suite** + +Run: `pytest braintrace/ -q -rs 2>&1 | tail -30` + +Record: the pass/fail/skip counts, and the reason string of every skip and xfail. Each row you write in the next step must correspond to something in this output or to a test you can name. + +- [ ] **Step 2: Write the document** + +Create `docs/specs/2026-07-25-known-limitations.md`: + +```markdown +# Known limitations — verified findings list + +Status: living document +Baseline: commit `bc153da` plus the P1 phase +Supersedes: the untracked list formerly at +`dev/superpowers/specs/2026-05-26-comprehensive-test-strategy-design.md`, which +is gitignored and absent from the repository + +This is the backlog of expected-failure and improvement items that `AGENTS.md` +§ Known limitations refers to. Every entry is verified against the current test +suite, not transcribed. An entry is **active** only if a test or a documented +scope boundary pins it today. + +## Status legend + +- **resolved** — the claim no longer holds; a passing test pins the resolution. +- **dead** — the claim described code that no longer exists. +- **active** — the claim still holds; a test or documented boundary pins it. +- **misattributed** — the observation was real, the stated cause was not. + +## Findings + +| ID | Claim | Status | Pinned by | +|---|---|---|---| +| F-01 / F-04 | multi-state (`num_state >= 2`) HiddenGroups mishandled | resolved | `snn_model_correctness_test.py::test_multi_state_hidden_groups_are_discovered` and `::test_d_rtrl_matches_bptt_on_snn_models` (num_state 1-5) | +| F-07 / F-08 / F-09 | OTTT / OTPE approximation bias | dead | removed with those algorithms in 0.3.0 | +| F-17 | implementation facts drifted from the instruction file | resolved | `braintrace/__init___test.py` Task 7 | +| F-19 / F-20 | OTTT / OSTTP exactness gaps | dead | removed with those algorithms in 0.3.0 | +| F-21 | rank / decay / random-feedback configs are exact on rate models | misattributed — the cause is the oracle path (F-23), not the model | `approx_correctness_test.py::test_rank_decay_random_approximations_are_exact_on_rate_model_F21`, docstring corrected | +| F-22 | exposing approximation bias needs an SNN multi-population zoo | **retired** — premise false; a multi-population SNN model is bitwise-exact on the same path | replaced by `approx_correctness_test.py::test_approximations_are_measurable_through_a_finite_window` | +| F-23 | the full-window multi-step oracle path is blind to every learning-rule axis | active, **by design** — documented, not a defect | `axis_discrimination_test.py`, both directions; warned in `online_param_gradients`' docstring | +| F-24 | `_etrace_model_test.py` factories are non-deterministic (unseeded global RNG) | active in the layer classes; neutralised for oracle use | `oracle_models_test.py::test_snn_spec_construction_is_deterministic` | +| F-25 | SNN models are silent at unit input scale, so comparisons are vacuous | active as a property; guarded | `oracle_models_test.py::test_snn_spec_is_live`, `test_default_input_scale_is_documented_as_dead`, and `assert_model_is_live` | +| F-26 | `pp_prop` / IODim raised on conv + trainable bias | resolved | `oracle_test.py::test_pp_prop_conv_bias_matches_bptt` | +| F-SCAN / F-SCAN-WEIGHT | weight inside control flow raised `KeyError` | resolved | `braintrace/_compiler/base_test.py::test_error_message_identifies_weight_variable` | +| F-SINGLESTEP | naive single-step summation does not equal BPTT even for an exact algorithm | resolved into a positive direction-alignment assertion | `braintrace/_algorithm/oracle_test.py` | + +## Mapping from the `AGENTS.md` prose list + +`AGENTS.md` described its limitations in prose. Each item maps onto a finding +above, so nothing survives only as prose: + +| Prose item | Disposition | +|---|---| +| approximation-mode validity beyond shallow depth | successor to F-21 / F-23; measured through a finite window | +| heterogeneous-population leak resolution | resolved — `snn_model_correctness_test.py::test_heterogeneous_leaks_do_not_break_exactness` | +| target-signal threading under JIT | dead — died with `OSTTP`'s `y_target` path | +| single-readout / feedback-shape assumptions | see F-28 below | +| gaps in cross-algorithm equivalence coverage | successor to F-23; `axis_discrimination_test.py` covers the pairs | + +## Notes on F-23 + +This is the load-bearing finding, and it is a property rather than a bug. A +whole-sequence `MultiStepData` call makes the within-call gradient exact +reverse-mode, so the eligibility trace enters only at a sequence boundary that +does not exist. `chunked_online_param_gradients` has documented this all along +— "this is the oracle that actually validates trace correctness" — but the +assertions written against `online_param_gradients` did not heed it, which is +how F-21 and F-22 reached the wrong cause. + +The rule that follows: an assertion whose subject is a trace factorization, a +temporal recursion, a recurrence scope, a filter or a learning signal must use a +finite window and must be guarded by `assert_gradients_differ`. An assertion +whose subject is the compiler or an ETP per-primitive rule may use the +full-window path, and that is most of the ~95 `multi-step` call sites in the +suite. +``` + +- [ ] **Step 3: Reconcile the document with reality** + +For each row, confirm the named test exists and has the stated outcome: + +```bash +pytest braintrace/_algorithm/tests/axis_discrimination_test.py \ + braintrace/_algorithm/tests/snn_model_correctness_test.py \ + braintrace/_algorithm/oracle_models_test.py \ + braintrace/_algorithm/oracle_test.py \ + braintrace/_algorithm/tests/approx_correctness_test.py -q -rs +``` + +Then fix the document to match: add any finding Tasks 3-6 surfaced (F-27 for a +non-live spec, F-28 for the single-readout / feedback-shape probe, or a new +compiler finding), and correct any status you cannot back with a named test. If +the single-readout / feedback-shape assumption turned out to be a non-issue, say +so under F-28 with the test that shows it and mark it resolved; if it was never +probed, mark F-28 "unverified — carried forward" rather than claiming either. + +- [ ] **Step 4: Commit** + +```bash +git add docs/specs/2026-07-25-known-limitations.md +git commit -m "Add the in-tree verified known-limitations list + +Replaces the untracked dev/ list AGENTS.md pointed at, which is gitignored and +absent. Every entry is verified against the current suite and names the test +that pins it. Records F-23 as the load-bearing finding -- the full-window +oracle path is axis-blind by design -- and maps each AGENTS.md prose item onto a +finding so none survives only as prose." +``` + +--- + +### Task 8: Roadmap and AGENTS.md updates + +**Files:** +- Modify: `docs/specs/2026-07-25-algorithm-axes-roadmap.md` +- Modify: `AGENTS.md:148-156` + +**Interfaces:** +- Consumes: `docs/specs/2026-07-25-known-limitations.md` from Task 7. +- Produces: nothing. + +- [ ] **Step 1: Revise the P2 and P3 acceptance criteria** + +The roadmap's P2 and P3 criteria are stated as element-wise equality against +existing algorithms. Measured through the full-window path, those comparisons +pass for any algorithm and would not guard the refactor at all. + +In the P2 **Acceptance** block, replace the first bullet: + +```markdown +- Element-wise equality against golden values frozen *before* the refactor, for + all five surviving algorithms — frozen and compared through a **finite-window** + oracle path (`chunked_online_param_gradients`, `chunk_size` < `T`). The + full-window multi-step path returns BPTT for every algorithm regardless of its + axis coordinates (F-23), so golden values captured there would be identical + across all five and would guard nothing. Guard each comparison with + `assert_gradients_differ` between any two presets that differ in an axis. +``` + +In the P3 **Acceptance (two-sided squeeze)** block, replace the first two +bullets: + +```markdown +- `n = 1` equals the current `D_RTRL` element-wise, compared through a + finite-window path (regression guard). On the full-window path this holds for + every algorithm and so proves nothing (F-23). +- Whatever configuration expresses `recurrence_scope = coupled` after the + refactor equals the current `OSTLRecurrent` element-wise — again through a + finite-window path — regardless of whether that configuration turns out to be + a point on the `n` scale or a sibling value beside it. `axis_discrimination_test.py` + pins that `D_RTRL` and `OSTLRecurrent` *are* distinguishable on that path, so + this comparison has content. +``` + +Leave the `n >= graph diameter` bullet unchanged: BPTT equality there is a +statement about the total gradient and is correct on either path. + +- [ ] **Step 2: Rewrite the P1 section to match what was built** + +Replace the body of the `### P1 — Compiler: multi-timescale and heterogeneous +populations` section (roadmap lines 233-260) with: + +```markdown +### P1 — Verification harness and the limitation list — **done** + +Retitled during implementation. P1 was scoped as compiler work; measurement +showed the compiler already passes every one of its targets, and that the +instrument was what failed. Full spec: +[`2026-07-25-p1-verification-harness.md`](2026-07-25-p1-verification-harness.md). + +Delivered: negative-control oracle helpers (`assert_model_is_live`, +`assert_gradients_differ`) and pytree/unit-aware gradient comparison; +documented window semantics on every oracle entry point; deterministic and live +SNN model specs; an axis-discrimination meta-test; gradient-correctness tests on +realistic SNN models (multi-timescale, per-neuron heterogeneous leaks, +`num_state` 1-5, E/I populations); the conv-bias IO-dim fix; and the in-tree +findings list at [`2026-07-25-known-limitations.md`](2026-07-25-known-limitations.md). + +`hidden_group.py` was **not** modified — no defect was found in its Jacobian +path — so P3 inherits today's representation and risk 2 below is retired rather +than mitigated. + +**Acceptance:** met. The reconstructed limitation list is committed under +`docs/specs/`; every item has a passing test or a documented scope boundary; the +one skipped test (F-22) is retired rather than deferred. +``` + +- [ ] **Step 3: Add the lessons-learned section** + +Append to the end of `docs/specs/2026-07-25-algorithm-axes-roadmap.md`: + +```markdown +## Lessons learned during implementation + +Recorded during P1, against commit `bc153da`. These are the things the roadmap +got wrong or could not have known, kept here because later phases rest on them. + +1. **The instrument was the defect, not the compiler.** P1 was scoped as + compiler work on multi-timescale and heterogeneous populations. Every one of + those targets already passed: HiddenGroups with `num_state` up to 5, + per-neuron heterogeneous `tau_mem` and `tau_a`, multi-timescale synapses + (`tau_mem` / `tau_syn` / `tau_a` / `tau_std` / `tau_f` / `tau_d`), and E/I + population splits all reproduce BPTT. Measure before scoping a phase around + a suspected defect. + +2. **A full-sequence multi-step VJP is BPTT, so it cannot see any axis** + (F-23). `D_RTRL`, `OSTLRecurrent`, `EProp` at any `kappa_filter_decay` and + `pp_prop` at any `decay_or_rank` return bitwise-identical gradients through + `online_param_gradients`. This is correct semantics — no truncation, nothing + to approximate — but it silently voids any assertion whose subject is a + learning-rule axis. `chunked_online_param_gradients` documented this from the + start; the tests written against the other entry point did not heed it. **Any + acceptance criterion in this roadmap phrased as "equals X element-wise" must + name a finite-window path.** + +3. **A wrong cause propagates further than a wrong measurement.** F-21 observed + real behaviour and attributed it to the model being a single-HiddenGroup rate + model. F-22 then deferred an entire work item — an SNN multi-population model + zoo — on that attribution, and P5 inherited it. Running exactly that model + (`ALIF_ExpCo`, 3 relations, `num_state` 5, `T` 20) reproduced the same + bitwise exactness, so no zoo could ever have helped. The fix was one + parameter: `chunk_size`. + +4. **Vacuous tests look like passing tests.** Three distinct ways a gradient + assertion in this repository can assert nothing: the reference gradient is + zero (a silent SNN — 8 of 9 probed configurations never reach threshold at + unit input scale); the model differs between the two sides of the comparison + (`_etrace_model_test.py` constructors draw from the unseeded global RNG, so + `factory()` returns a different network per call); or the oracle path cannot + see the knob under test (F-23). Spiking is *not* sufficient for the first — + `ALIF_Delta` spiked at rate 0.60 with a zero BPTT gradient. Hence + `assert_model_is_live` keys on gradient norm, and every new axis assertion + carries `assert_gradients_differ`. + +5. **`xy_to_dw` rules may return un-reduced, per-position leaves.** Conv's rule + deliberately defers the spatial sum of the bias Jacobian to `_conv_dt_to_t`, + which only the param-dim path calls. Any new solver that consumes `xy_to_dw` + directly must reduce produced leaves to the parameter's own shape (F-26). + This matters for P2: the axis strategies will introduce new contraction + paths. + +6. **The removed algorithms took findings with them, and left stale + cross-references behind.** F-07/F-08/F-09/F-19/F-20 died with OTTT / OTPE / + OSTTP, and the `AGENTS.md` prose item "target-signal threading under JIT" + died with `OSTTP`'s `y_target` path. Meanwhile a docstring still pointed at + `dev/superpowers/specs/...`, a path that is gitignored and absent. Findings + lists must live in-tree. +``` + +- [ ] **Step 4: Repoint `AGENTS.md`** + +Replace the `## Known limitations` section body (`AGENTS.md:150-156`) with: + +```markdown +First-cut SNN algorithms pass smoke and cross-checks but carry approximation +edges and rough spots. These are enumerated, verified against the test suite, +and mapped to concrete improvement actions in the findings list at +`docs/specs/2026-07-25-known-limitations.md`. Treat that list as the backlog of +expected-failure and improvement items rather than duplicating it here. + +One rule from that list is load-bearing enough to state here: a gradient +assertion whose subject is a **learning-rule property** — a trace +factorization, a temporal recursion, a recurrence scope, a filter, a learning +signal — must be measured through a *finite-window* oracle path. A +whole-sequence multi-step VJP has no truncation left to approximate and returns +BPTT for every algorithm at every hyperparameter, so such an assertion passes +vacuously there. Assertions about the compiler or an ETP per-primitive rule may +use the whole-sequence path. +``` + +- [ ] **Step 5: Verify the docs are consistent** + +```bash +grep -rn "dev/" AGENTS.md docs/specs/ braintrace/ --include=*.md --include=*.py | grep -v "\.pyc" +``` + +Expected: only the historical mentions inside the roadmap and the known-limitations file that *describe* the absent `dev/` list. No live cross-reference telling a reader to go read a file under `dev/`. + +- [ ] **Step 6: Commit** + +```bash +git add AGENTS.md docs/specs/2026-07-25-algorithm-axes-roadmap.md +git commit -m "Record P1 outcomes: revised P2/P3 criteria, lessons learned + +P2 and P3 acceptance criteria now name a finite-window oracle path. As written +they were element-wise equalities measured through the full-window path, where +every algorithm returns BPTT regardless of its axis coordinates -- they would +have passed without guarding the refactor at all. + +Rewrites the P1 section to match what was built, retires risk 2 (no +hidden_group.py defect was found, so P3 inherits today's representation), adds a +lessons-learned section, and repoints AGENTS.md off the absent dev/ list." +``` + +--- + +### Task 9: Full-suite verification + +**Files:** none modified. + +**Interfaces:** +- Consumes: every preceding task. +- Produces: the evidence for the phase's acceptance claim. + +- [ ] **Step 1: Run the full test suite** + +Run: `pytest braintrace/ -q -rs` + +Expected: all tests pass. The P0 baseline was 2062 passed, 1 skipped; the new +count is higher and the skip count is **0**, because Task 5 retired the only +skip. If any test fails, fix it — do not adjust the claim. + +- [ ] **Step 2: Run the type checker** + +Run: `mypy braintrace` + +Expected: clean, matching the P0 baseline. The likely new complaints are missing +annotations on the helpers added in Task 1 and `_reduce_to_param_shape` in Task +6; add the annotations rather than an ignore. + +- [ ] **Step 3: Record the numbers and update the spec's status** + +In `docs/specs/2026-07-25-p1-verification-harness.md`, change the header +`Status: approved, ready for implementation` to `Status: implemented` and append +the observed counts to the **Testing strategy** section: + +```markdown +**Verified:** `pytest braintrace/` → passed, 0 skipped; `mypy braintrace` +clean. +``` + +Fill in the real number from Step 1. Do not write a number you did not observe. + +- [ ] **Step 4: Commit** + +```bash +git add docs/specs/2026-07-25-p1-verification-harness.md +git commit -m "Mark the P1 spec implemented with verified suite counts" +``` + +--- + +## Self-review + +**Spec coverage.** D1 → Task 7. D2 → Tasks 1 and 2. D3 → Task 3. D4 and D5 → +Task 4. D6 → Task 5. D7 → Task 6. D8 → Task 8. The spec's testing-strategy +invariants (liveness, discrimination) are enforced by the helpers from Task 1 and +used in Tasks 2, 4 and 5. The verification bar is Task 9. The user's instruction +to record lessons learned in the roadmap is Task 8 Step 3. + +**Placeholders.** None. Every code step carries the code; every command carries +its expected output. The three conditional branches (Task 3 Step 5's non-live +`alif_delta`, Task 4 Step 2's possible compiler finding, Task 6 Step 5's +over-eager reduction) each name the concrete fallback and require recording a +finding rather than loosening an assertion. + +**Type consistency.** `flat_gradient_leaves`, `gradient_norm`, +`relative_deviation`, `assert_model_is_live`, `assert_gradients_differ` are +defined in Task 1 and used with those exact names in Tasks 2, 3, 4 and 5. +`ModelSpec.make_inputs(T, n_in, *, seed=0)` and `SNN_SPECS` are defined in Task +3 and consumed with that signature in Tasks 3 and 4. `_reduce_to_param_shape` is +local to Task 6. `input_scale` is used by Task 3's +`test_default_input_scale_is_documented_as_dead` and defined in the same task. diff --git a/docs/specs/2026-07-25-p1-verification-harness.md b/docs/specs/2026-07-25-p1-verification-harness.md index a56ad8dd..7e8b5486 100644 --- a/docs/specs/2026-07-25-p1-verification-harness.md +++ b/docs/specs/2026-07-25-p1-verification-harness.md @@ -70,6 +70,16 @@ raises `NotImplementedError` on multi-step input data. The only live paths today are `chunked_online_param_gradients` with `chunk_size < T` and `online_param_gradients_singlestep_naive`. +**This is already documented in the code, and that is the sharper form of the +finding.** `chunked_online_param_gradients`' docstring states it outright — the +full-sequence call is "exact reverse-mode and the trace only enters at the +sequence boundary", while chunking "makes the total depend on the eligibility +trace at every chunk boundary — this is the oracle that actually validates trace +correctness." So F-23 is not a discovery about the code; it is a discovery that +the assertions written against `online_param_gradients` do not heed what +`oracle.py` already says. `online_param_gradients` itself carries no such +warning, which is the gap P1 closes. + Scale of exposure: 95 `vjp_method='multi-step'` call sites across 20 test modules. Most are legitimate — see [Out of scope](#out-of-scope). @@ -132,7 +142,7 @@ Reconstructed disposition: | F-22 | approximation bias needs an SNN multi-population zoo | **premise false** — retire in P1, not P5 | re-point at a finite window | | F-SCAN / F-SCAN-WEIGHT | weight inside control flow | resolved | `_compiler/base_test.py` | | F-SINGLESTEP | single-step naive vs BPTT | resolved into a direction-alignment assertion | `oracle_test.py` | -| **F-23** | full-window multi-step oracle path is axis-blind | **new, active** | D2 | +| **F-23** | full-window multi-step oracle path is axis-blind | **active, by design** — documented in `chunked_online_param_gradients` but not heeded by the assertions | D2 | | **F-24** | `_etrace_model_test.py` factories are non-deterministic | **new, active** | D3 | | **F-25** | SNN zoo silent at default scale → vacuous comparisons | **new, active** | D3 | | **F-26** | `pp_prop` / IODim raises on conv + trainable bias | **new, active** (pre-existing) | `oracle_test.py::test_pp_prop_conv_bias_known_limitation` | From 6b6d4552a63f8db12e699c42b00235cdd114916a Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 12:11:17 +0800 Subject: [PATCH 05/13] Add oracle negative controls and document the axis-blind window assert_model_is_live rejects a zero reference gradient; assert_gradients_differ rejects a knob that does not move the gradient. assert_param_gradients_close now handles nested pytrees and unit-carrying leaves, which SNN models need, while keeping its diagnostic keyed by the caller's own state key. online_param_gradients now warns that a full-sequence window equals BPTT for any algorithm (F-23), so it tests the compiler and the ETP rules rather than the rule itself, and points at the finite-window path for axis assertions. Also replaces an unresolvable dev/ cross-reference. --- braintrace/_algorithm/oracle.py | 216 +++++++++++++++++++++++++-- braintrace/_algorithm/oracle_test.py | 94 ++++++++++++ 2 files changed, 301 insertions(+), 9 deletions(-) diff --git a/braintrace/_algorithm/oracle.py b/braintrace/_algorithm/oracle.py index 9b9be023..9b66d443 100644 --- a/braintrace/_algorithm/oracle.py +++ b/braintrace/_algorithm/oracle.py @@ -26,6 +26,7 @@ from typing import Callable import brainstate +import brainunit as u import jax import jax.numpy as jnp import numpy as np @@ -100,6 +101,29 @@ def online_param_gradients( ``algo_factory(model)`` must return an algorithm whose ``__call__`` accepts a ``braintrace.MultiStepData`` and returns the stacked per-step outputs. The loss ``(out ** 2).sum()`` over the whole stacked output equals the BPTT sequence loss. + + Warnings + -------- + **With a full-sequence input this path is blind to every learning-rule axis** + (finding F-23). One whole-sequence call makes the within-call gradient exact + reverse-mode, so the eligibility trace only enters at a sequence boundary -- + of which there is none. Every algorithm therefore returns gradients bitwise + equal to BPTT, at every hyperparameter setting: ``D_RTRL``, + ``OSTLRecurrent``, ``EProp`` at any ``kappa_filter_decay`` and ``pp_prop`` at + any ``decay_or_rank`` are indistinguishable here. + + That makes this function a good test of the *compiler and the ETP + per-primitive rules* -- it is the right instrument for asserting that an + exact algorithm reproduces BPTT on a realistic model -- and the wrong + instrument for any assertion whose subject is a trace factorization, a + temporal recursion, a recurrence scope, a filter or a learning signal. For + those use :func:`chunked_online_param_gradients` with ``chunk_size`` below + the sequence length, and guard the test with :func:`assert_gradients_differ`. + + See Also + -------- + chunked_online_param_gradients : finite-window path; sees the trace. + assert_gradients_differ : negative control for a knob that must matter. """ model = model_factory() brainstate.nn.init_all_states(model, batch_size=1) @@ -176,8 +200,14 @@ def online_param_gradients_singlestep_naive( """Naive 'single-step' total gradient: sum of per-step grad((algo(x_t)**2).sum()). Kept to document finding F-SINGLESTEP — this recipe does NOT equal BPTT even - for the exact D_RTRL algorithm, while the multi-step path does. See - dev/superpowers/specs/2026-05-26-comprehensive-test-strategy-design.md. + for the exact D_RTRL algorithm, while the multi-step path does. This is the + most aggressive finite window (one step), so it is maximally sensitive to + learning-rule axes and maximally divergent from BPTT. + + See Also + -------- + chunked_online_param_gradients : intermediate window size. + docs/specs/2026-07-25-known-limitations.md : F-SINGLESTEP and F-23. """ model = model_factory() brainstate.nn.init_all_states(model, batch_size=1) @@ -193,19 +223,187 @@ def online_param_gradients_singlestep_naive( return total +def flat_gradient_leaves(tree) -> dict: + """Flatten a path-keyed gradient tree into ``{label: plain array}``. + + Gradient trees returned by the oracle are keyed by ``ParamState`` path and + each value may itself be a pytree (a ``Linear`` contributes ``weight`` and + ``bias``) whose leaves may carry ``brainunit`` units. Comparisons need plain + arrays, so this strips both layers. + + Parameters + ---------- + tree : dict + Mapping from state path tuple to gradient pytree. + + Returns + ------- + dict + Mapping from ``'a/b|.leaf'`` label to a unit-stripped ``jax.Array``. + """ + out = {} + for key, value in tree.items(): + path_label = '/'.join(map(str, key)) if isinstance(key, tuple) else str(key) + for leaf_path, leaf in jax.tree_util.tree_flatten_with_path(value)[0]: + label = f'{path_label}|{jax.tree_util.keystr(leaf_path)}' + out[label] = jnp.asarray(u.get_mantissa(leaf)) + return out + + +def gradient_norm(tree) -> float: + """Euclidean norm of every leaf of a gradient tree, taken together. + + Parameters + ---------- + tree : dict + Path-keyed gradient tree. + + Returns + ------- + float + The joint Euclidean norm. + """ + leaves = flat_gradient_leaves(tree) + # Accumulate in NumPy float64: JAX truncates an explicit float64 astype + # unless x64 is enabled, and these sums are over whole gradient trees. + total = sum(float((np.asarray(arr, dtype=np.float64) ** 2).sum()) + for arr in leaves.values()) + return float(np.sqrt(total)) + + +def relative_deviation(actual, expected) -> float: + """``||actual - expected|| / ||expected||`` over all leaves jointly. + + Parameters + ---------- + actual, expected : dict + Path-keyed gradient trees with the same structure. + + Returns + ------- + float + The relative deviation; ``inf`` when ``expected`` is all-zero and + ``actual`` is not, ``0.0`` when both are. + + Raises + ------ + AssertionError + If the two trees do not have the same set of leaf labels. + """ + a = flat_gradient_leaves(actual) + e = flat_gradient_leaves(expected) + if set(a) != set(e): + raise AssertionError( + f'gradient trees have different leaves: {sorted(set(a) ^ set(e))}') + num = sum(float(((np.asarray(a[k], dtype=np.float64) + - np.asarray(e[k], dtype=np.float64)) ** 2).sum()) for k in e) + den = sum(float((np.asarray(e[k], dtype=np.float64) ** 2).sum()) for k in e) + if den == 0.0: + return float('inf') if num > 0.0 else 0.0 + return float(np.sqrt(num) / np.sqrt(den)) + + +def assert_model_is_live(model_factory, inputs, *, min_norm: float = 1e-8) -> float: + """Assert the BPTT gradient of ``model_factory`` on ``inputs`` is non-trivial. + + A comparison against an all-zero reference gradient passes for every + algorithm and therefore asserts nothing. SNN models are the common way to + hit this: at a low input scale the neurons never reach threshold, the loss is + zero, and so is the gradient. Spiking alone is not sufficient — a model can + spike and still have a zero gradient — so the criterion is the gradient norm + itself. + + Parameters + ---------- + model_factory : Callable[[], brainstate.nn.Module] + Zero-arg factory returning an uninitialized model. + inputs : jax.Array + ``(T, ...)`` input sequence. + min_norm : float, optional + Minimum acceptable BPTT gradient norm. + + Returns + ------- + float + The observed BPTT gradient norm. + + Raises + ------ + AssertionError + If the norm is at or below ``min_norm``. + """ + norm = gradient_norm(bptt_param_gradients(model_factory, inputs)) + if not (norm > min_norm): + raise AssertionError( + f'model is not live: BPTT gradient norm {norm:.3e} <= {min_norm:.3e}. ' + 'Any gradient comparison on this model/input pair is vacuous.' + ) + return norm + + +def assert_gradients_differ(a, b, *, min_rel: float = 1e-6) -> float: + """Assert two gradient trees are *distinguishable* — a negative control. + + Use this whenever a test intends to exercise a learning-rule knob. If the + oracle path chosen cannot see the knob, this fails loudly instead of letting + the test pass vacuously. See finding F-23. + + Parameters + ---------- + a, b : dict + Path-keyed gradient trees. + min_rel : float, optional + Minimum relative deviation between them. + + Returns + ------- + float + The observed relative deviation. + + Raises + ------ + AssertionError + If the deviation is below ``min_rel``. + + See Also + -------- + online_param_gradients : the path on which knobs are invisible. + chunked_online_param_gradients : the path on which they are visible. + """ + rel = relative_deviation(a, b) + if not (rel >= min_rel): + raise AssertionError( + f'gradients are indistinguishable: relative deviation {rel:.3e} < ' + f'{min_rel:.3e}. The knob under test does not move the gradient on ' + 'this oracle path.' + ) + return rel + + def assert_param_gradients_close(actual, expected, *, atol=1e-4, rtol=0.0, keys=None): - """Assert two param-gradient dicts match, with a per-key diagnostic on failure. + """Assert two param-gradient trees match, with a per-leaf diagnostic on failure. - ``keys`` restricts the comparison to a subset (e.g. only ETP params). When - None, every key present in ``expected`` is compared. + ``keys`` restricts the comparison to a subset of top-level state paths (e.g. + only ETP params). When None, every key present in ``expected`` is compared. + Nested pytrees and unit-carrying leaves are supported. """ compare_keys = list(expected.keys()) if keys is None else list(keys) failures = [] for key in compare_keys: - a = jnp.asarray(actual[key]) - e = jnp.asarray(expected[key]) - if not bool(jnp.allclose(a, e, atol=atol, rtol=rtol)): - failures.append(f" {key}: maxabsdiff={float(jnp.max(jnp.abs(a - e))):.3e}") + # Flatten per key so the diagnostic can name the offending state key + # exactly as the caller wrote it, and append the leaf path within it. + a = flat_gradient_leaves({key: actual[key]}) + e = flat_gradient_leaves({key: expected[key]}) + if set(a) != set(e): + raise AssertionError( + f'gradient trees differ in structure at {key!r}: ' + f'{sorted(set(a) ^ set(e))}') + for label in sorted(e): + if not bool(jnp.allclose(a[label], e[label], atol=atol, rtol=rtol)): + leaf = label.split('|', 1)[1] + failures.append( + f" {key!r}{leaf}: maxabsdiff=" + f"{float(jnp.max(jnp.abs(a[label] - e[label]))):.3e}") if failures: raise AssertionError( "param gradients differ beyond tolerance " diff --git a/braintrace/_algorithm/oracle_test.py b/braintrace/_algorithm/oracle_test.py index 8c606aad..21511da8 100644 --- a/braintrace/_algorithm/oracle_test.py +++ b/braintrace/_algorithm/oracle_test.py @@ -20,6 +20,7 @@ import brainevent import brainstate +import jax import jax.numpy as jnp import numpy as np import pytest @@ -513,3 +514,96 @@ def test_pp_prop_conv_bias_known_limitation(): factory, xs, algo_factory=lambda m: braintrace.pp_prop(m, decay_or_rank=0.9, vjp_method='single-step'), ) + + +# --- P1: negative-control helpers ------------------------------------------- + +def test_flat_gradient_leaves_handles_nested_and_units(): + """Gradient trees are nested dicts and may carry units; the flattener must + yield plain arrays keyed by a stable label.""" + import brainunit as u + from braintrace._algorithm.oracle import flat_gradient_leaves + tree = { + ('syn', 'comm', 'weight'): {'weight': jnp.ones((2, 3)) * u.mS, + 'bias': jnp.zeros((3,)) * u.mS}, + ('w',): jnp.arange(4.0), + } + flat = flat_gradient_leaves(tree) + assert len(flat) == 3 + for arr in flat.values(): + assert not isinstance(arr, u.Quantity) + assert sorted(k.split('|')[0] for k in flat) == ['syn/comm/weight', + 'syn/comm/weight', 'w'] + + +def test_gradient_norm_and_relative_deviation(): + from braintrace._algorithm.oracle import gradient_norm, relative_deviation + a = {('w',): jnp.array([3.0, 4.0])} + b = {('w',): jnp.array([3.0, 4.0])} + assert gradient_norm(a) == pytest.approx(5.0, abs=1e-6) + assert relative_deviation(a, b) == pytest.approx(0.0, abs=1e-12) + # relative_deviation(actual, expected) normalises by ||expected||: + # ||[3,4] - [0,4]|| / ||[0,4]|| == 3 / 4. + c = {('w',): jnp.array([0.0, 4.0])} + assert relative_deviation(a, c) == pytest.approx(3.0 / 4.0, abs=1e-6) + # all-zero reference: infinite deviation, not a silent zero. + zero = {('w',): jnp.zeros(2)} + assert relative_deviation(a, zero) == float('inf') + assert relative_deviation(zero, zero) == 0.0 + + +def test_assert_model_is_live_passes_on_live_model(): + from braintrace._algorithm.oracle import assert_model_is_live + spec = tanh_rnn(n_in=3, n_rec=4, seed=0) + xs = _inputs(4, 3) + norm = assert_model_is_live(spec.factory, xs) + assert norm > 0.0 + + +def test_assert_model_is_live_rejects_a_dead_model(): + """A model whose output is detached from its parameter has a zero gradient, + so any comparison against it asserts nothing. The guard must reject it. + + The parameter is still *used* — an entirely unused ``ParamState`` makes + ``brainstate``'s gradient transform raise before the oracle sees it — but + ``stop_gradient`` severs the derivative, which is exactly the silent-SNN + situation F-25 describes: a live forward pass with a zero gradient. + """ + from braintrace._algorithm.oracle import assert_model_is_live + + def factory(): + class Dead(brainstate.nn.Module): + def __init__(self): + super().__init__() + self.w = brainstate.ParamState(jnp.ones((3, 3))) + self.h = brainstate.HiddenState(jnp.zeros((1, 3))) + + def update(self, x): + self.h.value = jax.lax.stop_gradient(x @ self.w.value) + return self.h.value + + return Dead() + + xs = jnp.ones((3, 1, 3)) + with pytest.raises(AssertionError, match='gradient norm'): + assert_model_is_live(factory, xs) + + +def test_assert_gradients_differ_flags_a_dead_knob(): + from braintrace._algorithm.oracle import assert_gradients_differ + a = {('w',): jnp.array([1.0, 2.0])} + assert_gradients_differ(a, {('w',): jnp.array([1.0, 5.0])}) + with pytest.raises(AssertionError, match='indistinguishable'): + assert_gradients_differ(a, {('w',): jnp.array([1.0, 2.0])}) + + +def test_assert_param_gradients_close_supports_nested_unit_trees(): + """The pre-existing helper only handled flat, unitless dicts. SNN models + have nested weight dicts carrying units.""" + import brainunit as u + a = {('syn',): {'weight': jnp.ones((2, 2)) * u.mS, 'bias': jnp.zeros(2) * u.mS}} + b = {('syn',): {'weight': jnp.ones((2, 2)) * u.mS, 'bias': jnp.zeros(2) * u.mS}} + assert_param_gradients_close(a, b, atol=1e-6) + c = {('syn',): {'weight': jnp.full((2, 2), 2.0) * u.mS, 'bias': jnp.zeros(2) * u.mS}} + with pytest.raises(AssertionError, match='maxabsdiff'): + assert_param_gradients_close(a, c, atol=1e-6) From 334ee29fda3e5fe142863b11f717a9250a0a7d1a Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 12:12:18 +0800 Subject: [PATCH 06/13] Pin F-23: which oracle paths can see a learning-rule axis Asserts both directions -- the full-window multi-step path collapses axis-distinct configurations to bitwise-identical gradients, and a chunked window separates them. This is the meta-test whose absence let F-21 attribute the effect to the model rather than the harness. The recurrence-scope pair is distinguishable on the finite window, which independently confirms OSTLRecurrent is not a D_RTRL alias and gives P2 and P3 a regression guard with actual content. --- .../tests/axis_discrimination_test.py | 111 ++++++++++++++++++ 1 file changed, 111 insertions(+) create mode 100644 braintrace/_algorithm/tests/axis_discrimination_test.py diff --git a/braintrace/_algorithm/tests/axis_discrimination_test.py b/braintrace/_algorithm/tests/axis_discrimination_test.py new file mode 100644 index 00000000..50044297 --- /dev/null +++ b/braintrace/_algorithm/tests/axis_discrimination_test.py @@ -0,0 +1,111 @@ +# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""F-23: which oracle paths can see a learning-rule axis, and which cannot. + +A full-sequence ``MultiStepData`` call makes the within-call gradient exact +reverse-mode, so the eligibility trace enters only at a sequence boundary that +does not exist. Every algorithm then returns BPTT. This module pins that in both +directions, so a future test cannot silently assert an approximation's behaviour +on a path that cannot observe it -- which is how F-21 came to attribute the +effect to the model instead of the harness. +""" + +import jax.numpy as jnp +import numpy as np +import pytest + +import braintrace +from braintrace._algorithm.oracle import ( + assert_gradients_differ, + assert_param_gradients_close, + bptt_param_gradients, + chunked_online_param_gradients, + online_param_gradients, + relative_deviation, +) +from braintrace._algorithm.oracle_models import tanh_rnn + +T = 8 +CHUNK = 2 + + +def _spec(): + return tanh_rnn(n_in=3, n_rec=4, seed=0) + + +def _inputs(): + return jnp.asarray(np.random.RandomState(0).randn(T, 3).astype('float32')) + + +# Pairs of configurations that differ ONLY in a learning-rule axis value. +AXIS_PAIRS = { + # Axis 1/2: IO-dim trace factorization strength (decay). + 'pp_prop decay': ( + lambda m: braintrace.pp_prop(m, decay_or_rank=0.99, vjp_method='multi-step'), + lambda m: braintrace.pp_prop(m, decay_or_rank=0.01, vjp_method='multi-step'), + ), + # Axis 5: trace filter (kappa). + 'EProp kappa': ( + lambda m: braintrace.EProp(m, kappa_filter_decay=0.0, vjp_method='multi-step'), + lambda m: braintrace.EProp(m, kappa_filter_decay=0.95, vjp_method='multi-step'), + ), + # Axis 3: recurrence scope (diagonal vs coupled). + 'recurrence scope': ( + lambda m: braintrace.D_RTRL(m, vjp_method='multi-step'), + lambda m: braintrace.OSTLRecurrent(m, vjp_method='multi-step'), + ), +} + + +@pytest.mark.parametrize('axis', sorted(AXIS_PAIRS)) +def test_full_window_multistep_cannot_see_any_axis(axis): + """The full-window path collapses axis-distinct configs to identical + gradients. Pinned so the semantics are understood rather than accidental -- + if this ever starts failing, the window semantics changed and every + assertion written against this path needs review.""" + spec, xs = _spec(), _inputs() + lo, hi = AXIS_PAIRS[axis] + g_lo = online_param_gradients(spec.factory, xs, algo_factory=lo) + g_hi = online_param_gradients(spec.factory, xs, algo_factory=hi) + rel = relative_deviation(g_lo, g_hi) + assert rel == 0.0, ( + f'{axis}: full-window multi-step distinguished two configurations ' + f'(rel={rel:.3e}); F-23 assumed it cannot. Re-check the oracle docs.' + ) + + +@pytest.mark.parametrize('axis', sorted(AXIS_PAIRS)) +def test_finite_window_does_see_the_axis(axis): + """A chunked window makes the trace enter at every chunk boundary, so the + axis becomes observable. This is the path axis assertions must use.""" + spec, xs = _spec(), _inputs() + lo, hi = AXIS_PAIRS[axis] + g_lo = chunked_online_param_gradients( + spec.factory, xs, algo_factory=lo, chunk_size=CHUNK) + g_hi = chunked_online_param_gradients( + spec.factory, xs, algo_factory=hi, chunk_size=CHUNK) + assert_gradients_differ(g_lo, g_hi, min_rel=1e-6) + + +def test_full_window_still_reproduces_bptt_for_an_exact_algorithm(): + """The corollary that makes the full-window path useful: it is the right + instrument for asserting an exact algorithm matches BPTT.""" + spec, xs = _spec(), _inputs() + g_bptt = bptt_param_gradients(spec.factory, xs) + g_online = online_param_gradients( + spec.factory, xs, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + assert_param_gradients_close(g_online, g_bptt, atol=1e-4) From 14bc2f35a8ae0c32f801bb4c1178574511996b46 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 12:16:38 +0800 Subject: [PATCH 07/13] Add deterministic, live SNN model specs for oracle use Fixes F-24 and F-25 at the spec boundary rather than in the layer classes, which many existing tests depend on. Each spec re-seeds the global RNG before construction, so repeated factory() calls are bitwise identical, and records the input scale needed to produce a non-zero gradient. Covers LIF/ALIF x delta/ExpCu/STD/STP, an E/I conductance split, and per-neuron heterogeneous tau_mem and tau_a. Measurement sharpened F-25: the live input-scale window is bounded above as well as below. ALIF+delta keeps spiking at rate 0.60 from scale 2 to 20 while its BPTT gradient is exactly zero, because the surrogate derivative saturates. Spike rate is therefore not a proxy for liveness -- assert_model_is_live keys on the gradient norm -- and delta layers need a small scale where conductance layers need a large one. Both ends are now pinned by tests. --- braintrace/_algorithm/oracle_models.py | 187 ++++++++++++++++++++ braintrace/_algorithm/oracle_models_test.py | 92 ++++++++++ 2 files changed, 279 insertions(+) create mode 100644 braintrace/_algorithm/oracle_models_test.py diff --git a/braintrace/_algorithm/oracle_models.py b/braintrace/_algorithm/oracle_models.py index 76e46f1f..b5fb02b6 100644 --- a/braintrace/_algorithm/oracle_models.py +++ b/braintrace/_algorithm/oracle_models.py @@ -19,8 +19,10 @@ from typing import Callable, Tuple import brainstate +import brainunit as u import jax import jax.numpy as jnp +import numpy as np import braintrace @@ -32,11 +34,54 @@ class ModelSpec: ``factory()`` returns a freshly constructed, *uninitialized* model with deterministic weights. Callers must call ``brainstate.nn.init_all_states(model, batch_size=...)`` themselves. + + Attributes + ---------- + factory : Callable[[], brainstate.nn.Module] + Deterministic zero-arg model constructor. + etp_param_keys : tuple of tuple + Parameter paths routed through an ETP primitive. + plain_param_keys : tuple of tuple + Parameter paths used via plain JAX ops, hence excluded from ETP. + input_scale : float, optional + Multiplier applied by :meth:`make_inputs`. Spiking models need a scale + well above 1.0 to reach threshold at all; below it the loss and the + gradient are identically zero and any comparison is vacuous (F-25). + batched_input : bool, optional + Whether :meth:`make_inputs` emits a leading batch axis of 1. SNN layers + concatenate the input with the recurrent spike vector, so their ranks + must match; the rate models broadcast instead and do not need it. """ factory: Callable[[], brainstate.nn.Module] etp_param_keys: Tuple[tuple, ...] # routed through an ETP primitive plain_param_keys: Tuple[tuple, ...] # used via plain JAX ops (excluded from ETP) + input_scale: float = 1.0 + batched_input: bool = False + + def make_inputs(self, T: int, n_in: int, *, seed: int = 0): + """Build a ``(T, [1,] n_in)`` input sequence at this spec's scale. + + Values are non-negative so that spiking models receive net excitatory + drive; a zero-mean drive largely cancels and leaves the network silent. + + Parameters + ---------- + T : int + Number of time steps. + n_in : int + Input dimension. + seed : int, optional + Seed for the input draw. + + Returns + ------- + jax.Array + The input sequence, scaled by :attr:`input_scale`. + """ + rng = np.random.RandomState(seed) + shape = (T, 1, n_in) if self.batched_input else (T, n_in) + return self.input_scale * jnp.asarray(np.abs(rng.randn(*shape)).astype('float32')) def tanh_rnn(n_in: int = 3, n_rec: int = 4, seed: int = 0) -> ModelSpec: @@ -467,3 +512,145 @@ def update(self, x): return Net() return ModelSpec(factory=factory, etp_param_keys=(('win',),), plain_param_keys=()) + + +# --------------------------------------------------------------------------- +# SNN specs: the realistic-model end of the zoo. +# +# These wrap the layer classes in ``braintrace/_etrace_model_test.py`` for +# oracle use. Two things have to be fixed at this boundary: +# +# * F-24 -- those constructors call unseeded ``braintools.init.*``, which draws +# from the global ``brainstate.random`` stream, so ``factory()`` returns a +# different model on every call and a BPTT-vs-online comparison would compare +# two different networks. Each factory re-seeds before constructing. +# * F-25 -- at unit input scale the neurons never reach threshold, so the loss +# and the gradient are identically zero. Each spec records the scale that +# makes it live; ``oracle_models_test.py`` asserts both properties. +# +# The live input-scale window is bounded on **both** sides, which is the part of +# F-25 that is easy to miss. Too little drive and the neuron never crosses +# threshold; too much and the surrogate derivative saturates, so the gradient is +# exactly zero again *while the network keeps spiking*. Measured for +# ``ALIF_Delta`` at n_in=4, n_rec=5, T=6: +# +# scale spike_rate |g_bptt| +# 0.05 0.00 0.0 <- under threshold +# 0.20 0.17 3.0e+00 +# 1.00 0.53 9.3e+00 <- chosen +# 2.00 0.60 0.0 <- saturated surrogate, still spiking +# 20.00 0.60 0.0 +# +# So spike rate is not a proxy for liveness, and the per-spec scale below is not +# a free parameter: conductance-based (ExpCu/ExpCo) layers need a large scale, +# while delta layers inject straight into mV and need a small one. +# --------------------------------------------------------------------------- + +_SNN_SEED = 7 +_SNN_SCALE = 20.0 # conductance-based layers +_SNN_SCALE_DELTA = 1.0 # ALIF + delta synapse: saturates above ~2.0 + + +def _snn_spec(cls, n_in, n_rec, seed, scale=_SNN_SCALE, **kwargs) -> ModelSpec: + """Wrap an SNN layer class as a deterministic, live ``ModelSpec``.""" + + def factory(): + brainstate.random.seed(seed) + return cls(n_in, n_rec, **kwargs) + + return ModelSpec( + factory=factory, + etp_param_keys=(), # discovered by the compiler; not asserted per-spec + plain_param_keys=(), + input_scale=scale, + batched_input=True, + ) + + +def snn_if_delta(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """IF neuron, delta synapse. Single hidden state (``num_state == 1``).""" + from braintrace._etrace_model_test import IF_Delta_Dense_Layer + return _snn_spec(IF_Delta_Dense_Layer, n_in, n_rec, seed) + + +def snn_alif_delta(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """ALIF neuron, delta synapse. Membrane + adaptation (``num_state == 2``). + + Uses the smaller delta scale: the synapse injects directly in mV, so the + conductance-model scale saturates the surrogate derivative and drives the + gradient to exactly zero while the network still spikes. See the F-25 note + above this block. + """ + from braintrace._etrace_model_test import ALIF_Delta_Dense_Layer + return _snn_spec(ALIF_Delta_Dense_Layer, n_in, n_rec, seed, + scale=_SNN_SCALE_DELTA) + + +def snn_lif_expcu(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """LIF neuron, exponential current synapse. Two timescales: tau_mem, tau_syn.""" + from braintrace._etrace_model_test import LIF_ExpCu_Dense_Layer + return _snn_spec(LIF_ExpCu_Dense_Layer, n_in, n_rec, seed) + + +def snn_alif_expcu(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """ALIF + exponential current synapse. Three timescales, ``num_state == 3``.""" + from braintrace._etrace_model_test import ALIF_ExpCu_Dense_Layer + return _snn_spec(ALIF_ExpCu_Dense_Layer, n_in, n_rec, seed) + + +def snn_lif_std_expcu(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """LIF + short-term depression. Adds tau_std as a further timescale.""" + from braintrace._etrace_model_test import LIF_STDExpCu_Dense_Layer + return _snn_spec(LIF_STDExpCu_Dense_Layer, n_in, n_rec, seed) + + +def snn_lif_stp_expcu(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """LIF + short-term plasticity. Adds tau_f and tau_d.""" + from braintrace._etrace_model_test import LIF_STPExpCu_Dense_Layer + return _snn_spec(LIF_STPExpCu_Dense_Layer, n_in, n_rec, seed) + + +def snn_alif_expco_ei(n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED) -> ModelSpec: + """ALIF with an excitatory/inhibitory population split and conductance + synapses. The heterogeneous-population case: separate E and I projections + produce several ETP relations feeding one hidden group.""" + from braintrace._etrace_model_test import ALIF_ExpCo_Dense_Layer + return _snn_spec(ALIF_ExpCo_Dense_Layer, n_in, n_rec, seed) + + +def snn_lif_expcu_heterogeneous( + n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED +) -> ModelSpec: + """LIF whose membrane time constant differs per neuron. + + The heterogeneous-leak case: ``tau_mem`` is a length-``n_rec`` vector, so no + single global leak exists for the transition to factor out. + """ + from braintrace._etrace_model_test import LIF_ExpCu_Dense_Layer + tau_mem = jnp.linspace(3.0, 12.0, n_rec) * u.ms + return _snn_spec(LIF_ExpCu_Dense_Layer, n_in, n_rec, seed, tau_mem=tau_mem) + + +def snn_alif_expcu_heterogeneous( + n_in: int = 4, n_rec: int = 5, seed: int = _SNN_SEED +) -> ModelSpec: + """ALIF with per-neuron membrane *and* adaptation time constants.""" + from braintrace._etrace_model_test import ALIF_ExpCu_Dense_Layer + return _snn_spec( + ALIF_ExpCu_Dense_Layer, n_in, n_rec, seed, + tau_mem=jnp.linspace(3.0, 12.0, n_rec) * u.ms, + tau_a=jnp.linspace(60.0, 150.0, n_rec) * u.ms, + ) + + +SNN_SPECS = { + 'if_delta': snn_if_delta, + 'alif_delta': snn_alif_delta, + 'lif_expcu': snn_lif_expcu, + 'alif_expcu': snn_alif_expcu, + 'lif_std_expcu': snn_lif_std_expcu, + 'lif_stp_expcu': snn_lif_stp_expcu, + 'alif_expco_ei': snn_alif_expco_ei, + 'lif_expcu_heterogeneous': snn_lif_expcu_heterogeneous, + 'alif_expcu_heterogeneous': snn_alif_expcu_heterogeneous, +} diff --git a/braintrace/_algorithm/oracle_models_test.py b/braintrace/_algorithm/oracle_models_test.py new file mode 100644 index 00000000..94ae8a70 --- /dev/null +++ b/braintrace/_algorithm/oracle_models_test.py @@ -0,0 +1,92 @@ +# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""The SNN specs must be deterministic (F-24) and live (F-25) before any +gradient assertion built on them means anything.""" + +import brainstate +import brainunit as u +import jax.numpy as jnp +import pytest + +from braintrace._algorithm.oracle import ( + assert_model_is_live, + bptt_param_gradients, + flat_gradient_leaves, + gradient_norm, +) +from braintrace._algorithm.oracle_models import SNN_SPECS + + +@pytest.mark.parametrize('name', sorted(SNN_SPECS)) +def test_snn_spec_construction_is_deterministic(name): + """F-24: the underlying layer classes seed from the global RNG, so two + factory() calls must be pinned to produce identical weights.""" + spec = SNN_SPECS[name]() + with brainstate.environ.context(dt=0.1 * u.ms): + w1 = flat_gradient_leaves( + {k: v.value for k, v in spec.factory().states(brainstate.ParamState).items()}) + w2 = flat_gradient_leaves( + {k: v.value for k, v in spec.factory().states(brainstate.ParamState).items()}) + assert set(w1) == set(w2) + for key in w1: + assert bool(jnp.allclose(w1[key], w2[key])), f'{name}: {key} differs across calls' + + +@pytest.mark.parametrize('name', sorted(SNN_SPECS)) +def test_snn_spec_is_live(name): + """F-25: at the default input scale these networks never spike, so their + gradients are identically zero and every comparison is vacuous. Each spec + records a scale that produces a non-trivial gradient.""" + spec = SNN_SPECS[name]() + with brainstate.environ.context(dt=0.1 * u.ms): + xs = spec.make_inputs(6, 4) + norm = assert_model_is_live(spec.factory, xs, min_norm=1e-6) + assert norm > 1e-6 + + +def test_underdriven_input_scale_is_dead(): + """The counterpart of the above: pins *why* the scale field exists. At + scale 1.0 a conductance-based model never reaches threshold and its gradient + is exactly zero, so any comparison on it would be vacuous.""" + spec = SNN_SPECS['lif_expcu']() + with brainstate.environ.context(dt=0.1 * u.ms): + dead_xs = spec.make_inputs(6, 4) / spec.input_scale # undo the scaling + assert gradient_norm(bptt_param_gradients(spec.factory, dead_xs)) == 0.0 + + +def test_overdriven_input_scale_is_also_dead_while_still_spiking(): + """The live window is bounded *above* as well as below, and this is the half + that is easy to miss. + + Driven hard, ``ALIF_Delta`` keeps spiking (rate 0.60) but the surrogate + derivative saturates and the BPTT gradient returns to exactly zero. A + liveness check keyed on spike rate would pass here and the comparison would + still assert nothing -- which is why ``assert_model_is_live`` keys on the + gradient norm instead. See F-25. + """ + spec = SNN_SPECS['alif_delta']() + with brainstate.environ.context(dt=0.1 * u.ms): + live_xs = spec.make_inputs(6, 4) + over_xs = live_xs * 20.0 + + model = spec.factory() + brainstate.nn.init_all_states(model, batch_size=1) + outs = brainstate.transform.for_loop(lambda x: model(x), over_xs) + spike_rate = float(jnp.mean(jnp.asarray(u.get_mantissa(outs)) > 0.0)) + + assert gradient_norm(bptt_param_gradients(spec.factory, live_xs)) > 1e-6 + assert spike_rate > 0.1, 'the over-driven network must still be spiking' + assert gradient_norm(bptt_param_gradients(spec.factory, over_xs)) == 0.0 From e65f00bf440aeb97aa681358983bd201a4fc57ed Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 12:18:25 +0800 Subject: [PATCH 08/13] Assert gradient correctness on realistic SNN models Discharges the AGENTS.md prose limitations by test: heterogeneous per-neuron leaks, multi-timescale synapses, multi-state HiddenGroups (num_state 1 to 5) and E/I populations all reproduce BPTT under the exact algorithm. Structural facts are pinned alongside, so the exactness claims rest on discovered relations rather than assumption. Cross-algorithm comparisons use a finite window per F-23, and the approximation is measurably distinguishable from the exact algorithm on both a plain LIF model and the E/I multi-population one -- the assertion F-22 was waiting on a model zoo for. Every test carries a liveness guard. --- .../tests/snn_model_correctness_test.py | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 braintrace/_algorithm/tests/snn_model_correctness_test.py diff --git a/braintrace/_algorithm/tests/snn_model_correctness_test.py b/braintrace/_algorithm/tests/snn_model_correctness_test.py new file mode 100644 index 00000000..343e38fa --- /dev/null +++ b/braintrace/_algorithm/tests/snn_model_correctness_test.py @@ -0,0 +1,136 @@ +# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Gradient correctness on realistic SNN models: multi-timescale synapses, +per-neuron heterogeneous leaks, multi-state HiddenGroups, and E/I populations. + +These are the claims ``AGENTS.md`` carried only as prose ("heterogeneous- +population leak resolution", "multi-state HiddenGroups"). They are discharged +here as passing tests rather than as fixes -- the compiler already handles them. + +Assertion paths, per F-23: an *exact* algorithm versus BPTT uses the full-window +multi-step path, whose subject is the compiler and the ETP per-primitive rules +and which is the right instrument for that. Anything comparing *across* +algorithms uses a finite window. +""" + +import brainstate +import brainunit as u +import pytest + +import braintrace +from braintrace._algorithm.oracle import ( + assert_gradients_differ, + assert_model_is_live, + assert_param_gradients_close, + bptt_param_gradients, + chunked_online_param_gradients, + online_param_gradients, +) +from braintrace._algorithm.oracle_models import SNN_SPECS + +T = 6 +N_IN = 4 +N_REC = 5 +ATOL = 1e-4 + +_ALL_SPECS = sorted(SNN_SPECS) + + +def _setup(name): + spec = SNN_SPECS[name]() + xs = spec.make_inputs(T, N_IN) + return spec, xs + + +@pytest.mark.parametrize('name', _ALL_SPECS) +def test_d_rtrl_matches_bptt_on_snn_models(name): + """D_RTRL is exact, so it must reproduce BPTT on every realistic model: + multi-timescale synapses, heterogeneous leaks, E/I populations and + HiddenGroups with num_state from 1 to 5.""" + spec, xs = _setup(name) + with brainstate.environ.context(dt=0.1 * u.ms): + assert_model_is_live(spec.factory, xs, min_norm=1e-6) + g_bptt = bptt_param_gradients(spec.factory, xs) + g_online = online_param_gradients( + spec.factory, xs, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + assert_param_gradients_close(g_online, g_bptt, atol=ATOL) + + +@pytest.mark.parametrize('name', ['lif_expcu_heterogeneous', + 'alif_expcu_heterogeneous']) +def test_heterogeneous_leaks_do_not_break_exactness(name): + """A per-neuron time constant leaves no single global leak for the + transition to factor out. The compiler takes a true Jacobian, so exactness + survives -- this is what retires the AGENTS.md prose item.""" + spec, xs = _setup(name) + with brainstate.environ.context(dt=0.1 * u.ms): + assert_model_is_live(spec.factory, xs, min_norm=1e-6) + g_bptt = bptt_param_gradients(spec.factory, xs) + g_online = online_param_gradients( + spec.factory, xs, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + assert_param_gradients_close(g_online, g_bptt, atol=ATOL) + + +@pytest.mark.parametrize('name,expected_min_state', [ + ('if_delta', 1), + ('lif_expcu', 2), + ('alif_expcu', 3), + ('alif_expco_ei', 5), +]) +def test_multi_state_hidden_groups_are_discovered(name, expected_min_state): + """Pins the structural facts the exactness tests above rest on: these models + really do form multi-state HiddenGroups, so the per-state axis is exercised + rather than assumed.""" + spec, xs = _setup(name) + with brainstate.environ.context(dt=0.1 * u.ms): + model = spec.factory() + brainstate.nn.init_all_states(model, batch_size=1) + algo = braintrace.D_RTRL(model, vjp_method='multi-step') + algo.compile_graph(xs[0]) + assert len(algo.graph.hidden_groups) >= 1 + assert max(hg.num_state for hg in algo.graph.hidden_groups) >= expected_min_state + assert len(algo.graph.hidden_param_op_relations) >= 1 + + +def test_ei_population_split_yields_multiple_relations(): + """The E/I model routes separate excitatory and inhibitory projections into + one hidden group, so the compiler must record more than one ETP relation.""" + spec, xs = _setup('alif_expco_ei') + with brainstate.environ.context(dt=0.1 * u.ms): + model = spec.factory() + brainstate.nn.init_all_states(model, batch_size=1) + algo = braintrace.D_RTRL(model, vjp_method='multi-step') + algo.compile_graph(xs[0]) + assert len(algo.graph.hidden_param_op_relations) >= 2 + + +@pytest.mark.parametrize('name', ['lif_expcu', 'alif_expco_ei']) +def test_approximation_is_measurable_on_snn_models(name): + """The genuinely approximate configuration must be *distinguishable* from + the exact one on a realistic model -- via a finite window, which is the only + path that can see it (F-23). This is what F-22 was really asking for.""" + spec, xs = _setup(name) + with brainstate.environ.context(dt=0.1 * u.ms): + g_exact = chunked_online_param_gradients( + spec.factory, xs, chunk_size=2, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + g_approx = chunked_online_param_gradients( + spec.factory, xs, chunk_size=2, + algo_factory=lambda m: braintrace.pp_prop( + m, decay_or_rank=0.5, vjp_method='multi-step')) + assert_gradients_differ(g_exact, g_approx, min_rel=1e-6) From 4bfa058df2b9ce6d7f58a5d26f2d561b5a81e62b Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 12:34:10 +0800 Subject: [PATCH 09/13] Retire F-22, correct F-21's attribution, and record F-29 F-21 read the exactness of the rank / decay / random-feedback configs as a property of the rate model, and F-22 deferred the real assertion until an SNN multi-population zoo existed. Both were wrong about the cause: the full-window multi-step oracle path has no truncation for those knobs to approximate (F-23), so no model can expose them there. Through a finite window they are measurably different on the very same rate model, which is the assertion F-22 wanted. pp_prop with decay_or_rank=1 is the exception, and a new finding (F-29): an int rank maps to decay = (rank - 1) / (rank + 1), so rank 1 means decay 0 and the presynaptic EMA collapses to the exact current input. On a model whose only ETP relation is the recurrent weight, that reproduces exact D_RTRL to round-off (1e-10 to 3e-8 over chunk sizes 1/2/4, T in {8,12,16}, n_rec in {4,8}, spectral radius 0.25 to 5.0 -- so not a weak-recurrence artifact). It is a real approximation once the presynaptic input carries an external component: 0.55 on an input-weight variant and 0.31 to 0.79 on the SNN specs. Both sides are pinned; the mechanism is not claimed. Numbered F-29 because the plan reserves F-26 for the conv-bias limitation and F-27/F-28 as placeholders. The former single-loop test asserted min_rel=1e-9, below float32 round-off, so it was asserting noise. The floor is now 1e-6 against deviations of 2e-3 to 2e-2. --- .../tests/approx_correctness_test.py | 129 ++++++++++++++++-- 1 file changed, 114 insertions(+), 15 deletions(-) diff --git a/braintrace/_algorithm/tests/approx_correctness_test.py b/braintrace/_algorithm/tests/approx_correctness_test.py index 403152b3..3f9b67a7 100644 --- a/braintrace/_algorithm/tests/approx_correctness_test.py +++ b/braintrace/_algorithm/tests/approx_correctness_test.py @@ -16,13 +16,14 @@ """L3-C approximate-class correctness: direction-alignment metrics against BPTT (cosine / sign / relative magnitude), the multi-state (num_state == 2) exactness guard for D_RTRL, and the rate-model approximation-exactness ceiling -(F-21/F-22). +(F-21/F-22/F-23/F-29). The OTTT/OTPE-specific bias tests (F-01/F-04/F-07/F-08/F-09) were removed in 0.3.0 along with those algorithms; see docs/specs for the roadmap. The metric helpers they exercised are retained and are the basis of the benchmark suite.""" import brainstate +import brainunit as u import jax import jax.numpy as jnp import numpy as np @@ -31,14 +32,18 @@ import braintrace from braintrace._algorithm.oracle import ( assert_direction_aligned, + assert_gradients_differ, assert_param_gradients_close, bptt_param_gradients, + chunked_online_param_gradients, cosine_similarity, online_param_gradients, + relative_deviation, relative_magnitude, sign_agreement, ) from braintrace._algorithm.oracle_models import ( + SNN_SPECS, tanh_rnn, two_state_rnn, ) @@ -124,10 +129,13 @@ def test_d_rtrl_exact_on_two_state_group(): # --- Task 9: F-21/F-22 IODim/EProp approximations are exact on rate models ----- -# On a single-relation rate model the IODim rank / ES decay / random-feedback -# "approximations" are degenerate-to-exact (spike: even rank=1 on n_rec=8/T=12 is -# cos 1.0 / relmag 1.0). We assert exactness here and defer the genuine -# approximation stress to an SNN multi-population model zoo (F-22). +# Through the full-window multi-step path the IODim rank / ES decay / +# random-feedback "approximations" are exact -- not because the model is a +# single-relation rate model, but because that path has no truncation for them to +# approximate (F-23). We assert the exactness here as a property of the path, and +# assert the genuine approximation through a finite window below. One entry, +# 'pp_prop_rank1', stays exact through a finite window too, for a reason specific +# to the rule rather than the path (F-29), and is held out of that test. _EXACT_ON_RATE = { 'pp_prop_rank1': lambda m: braintrace.pp_prop(m, decay_or_rank=1, vjp_method='multi-step'), 'pp_prop_decay05': lambda m: braintrace.pp_prop(m, decay_or_rank=0.5, vjp_method='multi-step'), @@ -141,8 +149,21 @@ def test_d_rtrl_exact_on_two_state_group(): @pytest.mark.parametrize('algo_name', list(_EXACT_ON_RATE)) def test_rank_decay_random_approximations_are_exact_on_rate_model_F21(algo_name): - """F-21: these nominally-approximate configs match BPTT element-wise on a - single-HiddenGroup rate model. The model cannot stress their approximation.""" + """F-21: these nominally-approximate configs match BPTT element-wise here. + + The cause is the *oracle path*, not the model. All four configurations run + through ``online_param_gradients`` with ``vjp_method='multi-step'`` over the + whole sequence, where the within-call gradient is exact reverse-mode and the + eligibility trace never enters -- so every algorithm returns BPTT at every + hyperparameter setting (F-23). An earlier reading of this test attributed the + exactness to the model being a single-HiddenGroup rate model and concluded + that an SNN multi-population zoo was needed to expose the bias (F-22); that + conclusion was wrong, and F-22 is retired by + ``test_approximations_are_measurable_through_a_finite_window`` below. + + The test is kept because the equality is still a real property of this path + and worth pinning. + """ spec = tanh_rnn(n_in=3, n_rec=4, seed=0) inputs = _inputs(6, 3) g_bptt = bptt_param_gradients(spec.factory, inputs) @@ -151,14 +172,92 @@ def test_rank_decay_random_approximations_are_exact_on_rate_model_F21(algo_name) assert_param_gradients_close(g_online, g_bptt, atol=ATOL_BPTT) -@pytest.mark.skip( - reason="F-22: IODim rank / ES decay / EProp random-feedback approximations are " - "degenerate-to-exact on single-relation rate models (verified up to " - "n_rec=8, T=12, rank=1). Exposing their real bias requires an SNN " - "multi-population model zoo (LIF/ALIF, multi-layer random feedback). " - "Deferred to P5b.") -def test_approximations_diverge_on_snn_multipopulation_DEFERRED(): - pass +def _chunked_exact(spec, inputs, chunk_size=2): + return chunked_online_param_gradients( + spec.factory, inputs, chunk_size=chunk_size, + algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step')) + + +# ``pp_prop_rank1`` is excluded deliberately -- see F-29 and +# test_rank1_pp_prop_is_degenerate_on_a_recurrent_only_relation below. +_MEASURABLE_THROUGH_A_WINDOW = [ + name for name in _EXACT_ON_RATE if name != 'pp_prop_rank1' +] + + +@pytest.mark.parametrize('algo_name', _MEASURABLE_THROUGH_A_WINDOW) +def test_approximations_are_measurable_through_a_finite_window(algo_name): + """F-22, retired. The finding deferred this until an "SNN multi-population + model zoo" existed, on the theory that the model was what made the + approximations look exact. That premise was wrong: a multi-population SNN + model (ALIF + E/I conductance split, 3 ETP relations, num_state 5) is + *also* bitwise-exact through the full-window path, because the cause is the + oracle path and not the model (F-23). No model can revive a knob the + harness cannot see. + + Through a finite window the same nominally-approximate configurations that + F-21 finds exact become measurably different from the exact algorithm, on + the very same rate model. That is the assertion F-22 wanted. + + ``min_rel`` is 1e-6 rather than something near zero on purpose: float32 + round-off on these trees sits around 1e-8, so a smaller floor would let the + test pass on numerical noise. The deviations asserted here are 2e-3 to 2e-2. + + The SNN counterpart now exists too, and asserts the same thing on realistic + models: ``tests/snn_model_correctness_test.py:: + test_approximation_is_measurable_on_snn_models``. + """ + spec = tanh_rnn(n_in=3, n_rec=4, seed=0) + inputs = _inputs(8, 3) + g_approx = chunked_online_param_gradients( + spec.factory, inputs, chunk_size=2, algo_factory=_EXACT_ON_RATE[algo_name]) + assert_gradients_differ(_chunked_exact(spec, inputs), g_approx, min_rel=1e-6) + + +def test_rank1_pp_prop_is_degenerate_on_a_recurrent_only_relation(): + """F-29: ``decay_or_rank=1`` is *not* an approximation when the relation's + presynaptic input is the hidden state itself. + + An int ``decay_or_rank`` maps to ``decay = (rank - 1) / (rank + 1)``, so + rank 1 means decay 0 -- the EMA over the presynaptic factor + ``eps^t = a * eps^{t-1} + (1 - a) * x_t`` collapses to ``x_t`` and no + presynaptic smearing is introduced. On ``tanh_rnn``, whose only ETP relation + is the recurrent weight (``x`` *is* ``h``), that reproduces exact D_RTRL to + round-off: measured 7.6e-10 here, and 1e-10 to 3e-8 over chunk sizes + 1/2/4, T in {8, 12, 16}, n_rec in {4, 8} and recurrent spectral radius + 0.25 to 5.0. Weak recurrence is not the cause -- the deviation does not grow + with the spectral radius. + + It becomes a real approximation as soon as the presynaptic input carries an + external component: 0.55 on an input-weight ETP variant, and 0.31 to 0.79 on + every SNN spec, whose projections consume ``concat(input, spikes)``. The + second half of this test pins that side of the boundary, so the first half + cannot be read as "rank 1 never does anything". + + The mechanism is not established -- only the boundary is. Consequence for + the harness: rank 1 on a recurrent-only relation is unusable as a positive + control for approximation error. + """ + spec = tanh_rnn(n_in=3, n_rec=4, seed=0) + inputs = _inputs(8, 3) + g_rank1 = chunked_online_param_gradients( + spec.factory, inputs, chunk_size=2, + algo_factory=_EXACT_ON_RATE['pp_prop_rank1']) + rel = relative_deviation(g_rank1, _chunked_exact(spec, inputs)) + assert rel < 1e-6, ( + f'rank-1 pp_prop deviated by {rel:.3e} on a recurrent-only relation; ' + 'F-29 recorded it as exact to round-off there. If this now differs, the ' + 'IO-dim trace semantics changed and F-29 needs revisiting.' + ) + + snn = SNN_SPECS['lif_expcu']() + xs = snn.make_inputs(8, 4) + with brainstate.environ.context(dt=0.1 * u.ms): + g_snn_rank1 = chunked_online_param_gradients( + snn.factory, xs, chunk_size=2, + algo_factory=_EXACT_ON_RATE['pp_prop_rank1']) + assert_gradients_differ( + _chunked_exact(snn, xs), g_snn_rank1, min_rel=1e-6) # --- Task 10: C-level convergence backstop (loss decreases) ------------------ From ec22b2f633d2507705b7ed7e043e4658fd936929 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 12:40:51 +0800 Subject: [PATCH 10/13] Fix F-26: reduce per-position gradient leaves in the IO-dim solver Conv's xy_to_dw rule returns the bias Jacobian per output position by design and defers the spatial sum to _conv_dt_to_t, which only the param-dim path calls. The IO-dim solver contracts at solve time, so it handed the un-reduced (batch, *spatial, out_ch) array to custom_vjp, which rejected it against the bias's own (out_ch,). D_RTRL was unaffected; pp_prop raised on any conv layer with a trainable bias, in either layout. Fixed in the solver rather than the conv rule: summing a produced leaf down to its parameter's shape is the standard broadcast-gradient reduction and is primitive-agnostic, so it does not whitelist conv. It is a no-op when the shapes already agree. test_pp_prop_conv_bias_known_limitation, which pinned the raise, is now test_pp_prop_conv_bias_matches_bptt, and conv_nwc_bias is back in the two pp_prop family parametrizations it had been excluded from. 635 tests pass across oracle_test, io_dim_vjp_test, pp_prop_test and _op/ -- no other primitive layout (sparse, lora, grouped, elemwise) regressed. --- braintrace/_algorithm/io_dim_vjp.py | 37 ++++++++++++++++++++ braintrace/_algorithm/oracle_test.py | 51 ++++++++++++---------------- 2 files changed, 58 insertions(+), 30 deletions(-) diff --git a/braintrace/_algorithm/io_dim_vjp.py b/braintrace/_algorithm/io_dim_vjp.py index dee4283b..5cca5d92 100644 --- a/braintrace/_algorithm/io_dim_vjp.py +++ b/braintrace/_algorithm/io_dim_vjp.py @@ -415,6 +415,37 @@ def _update_IO_dim_etrace_scan_fn( return (new_etrace_xs, new_etrace_dfs), None +def _reduce_to_param_shape(grad: jax.Array, param: jax.Array) -> jax.Array: + """Sum a produced gradient leaf down to its parameter's own shape. + + Some ``xy_to_dw`` rules deliberately return a *per-position* instantaneous + Jacobian for parameters that broadcast over positions -- a conv bias is the + canonical case, and :func:`braintrace._op.conv._conv_xy_to_dw` documents that + the spatial sum is deferred. The param-dim path performs that sum during + trace propagation; the IO-dim path contracts at solve time and must perform + it here, or ``custom_vjp`` rejects the shape mismatch (finding F-26). + + This is the standard broadcast-gradient reduction -- sum the extra leading + axes, then any axis the parameter holds as a singleton -- so it is + primitive-agnostic and a no-op whenever the shapes already agree. + """ + g_shape = u.math.shape(grad) + p_shape = u.math.shape(param) + if g_shape == p_shape: + return grad + extra = len(g_shape) - len(p_shape) + if extra < 0: + return grad + out = u.math.sum(grad, axis=tuple(range(extra))) if extra else grad + squeeze = tuple( + i for i, (gd, pd) in enumerate(zip(u.math.shape(out), p_shape)) + if pd == 1 and gd != 1 + ) + if squeeze: + out = u.math.sum(out, axis=squeeze, keepdims=True) + return out + + def _solve_IO_dim_weight_gradients( hist_etrace_data: Tuple[ Dict[ETraceX_Key, jax.Array], @@ -537,6 +568,12 @@ def _call(df_: Any, w_: Any, _rule: Any = xy_to_dw_rule, _params: Any = eqn_para else: dg_dict = jax.tree.map(_sum_dim, fn_vmap(df_hid)) + # Reduce per-position leaves to their parameter's own shape before + # routing; see _reduce_to_param_shape (finding F-26). + dg_dict = { + key: _reduce_to_param_shape(value, weights_dict[key]) + for key, value in dg_dict.items() + } # Route per-key to owning ParamState path and assemble per-path pytrees. _route_grads_by_path(relation, dg_dict, weight_vals, dG_weights) diff --git a/braintrace/_algorithm/oracle_test.py b/braintrace/_algorithm/oracle_test.py index 21511da8..5eaa86c9 100644 --- a/braintrace/_algorithm/oracle_test.py +++ b/braintrace/_algorithm/oracle_test.py @@ -441,14 +441,10 @@ def test_d_rtrl_singlestep_matches_bptt_across_families(name, T): assert rel < _TOL, f'{name} T={T} {key}: D_RTRL vs BPTT rel={rel:.3e}' -@pytest.mark.parametrize('name', sorted(set(_FAMILIES) - {'conv_nwc_bias'})) +@pytest.mark.parametrize('name', sorted(_FAMILIES)) def test_pp_prop_singlestep_exact_at_t1_across_families(name): """pp_prop (ES-D-RTRL, IO-dim, approximate) has no history to factor or decay at T=1, so it must also match BPTT exactly there, for every family. - - ``conv_nwc_bias`` is excluded here and covered separately by - ``test_pp_prop_conv_bias_known_limitation`` -- see that test's docstring - for the newly-discovered (pre-existing, out-of-scope-for-this-task) gap. """ factory, seed = _FAMILIES[name] with brainstate.environ.context(precision=64): @@ -463,7 +459,7 @@ def test_pp_prop_singlestep_exact_at_t1_across_families(name): assert rel < _TOL, f'{name} T=1 {key}: pp_prop vs BPTT rel={rel:.3e}' -@pytest.mark.parametrize('name', sorted(set(_FAMILIES) - {'conv_nwc_bias'})) +@pytest.mark.parametrize('name', sorted(_FAMILIES)) def test_pp_prop_singlestep_bounded_at_t2_across_families(name): """pp_prop is an *approximate* algorithm beyond T=1: at T=2 it factors / decays history and is **not** expected to match BPTT element-wise. This @@ -486,34 +482,29 @@ def test_pp_prop_singlestep_bounded_at_t2_across_families(name): ) -def test_pp_prop_conv_bias_known_limitation(): - """Documents a newly-discovered gap found while building this suite: - ``pp_prop``/``IODimVjpAlgorithm`` raises when a conv layer has a - trainable bias, regardless of layout (NCH or NWC) -- the custom-VJP bwd - rule returns the bias cotangent still shaped per-position - ``(batch, *spatial, out_ch)`` instead of reduced to the bias's own shape - ``(out_ch,)``. - - This is unrelated to the Task 6/7 fixes under audit: ``D_RTRL`` - (param-dim) handles conv+bias exactly (see - ``test_d_rtrl_singlestep_matches_bptt_across_families['conv_nwc_bias']``); - only the IO-dim path used by ``pp_prop`` is affected. Fixing it would - require touching ``io_dim_vjp.py``, which is out of scope for this task - (see module docstring / architecture notes: io_dim_vjp.py's core logic is - never modified as part of this audit). This test pins the *current* - behavior with ``pytest.raises(ValueError, ...)`` so that a future fix is - caught loudly (the raise will stop happening and this test will fail), - prompting promotion to an exactness assertion instead of silently going - unnoticed. +def test_pp_prop_conv_bias_matches_bptt(): + """Formerly ``test_pp_prop_conv_bias_known_limitation`` (finding F-26). + + ``_conv_xy_to_dw`` returns the bias Jacobian per output position by design; + the param-dim path reduces it in ``_conv_dt_to_t``, but the IO-dim solver + calls ``xy_to_dw`` directly at solve time and used to hand the un-reduced + ``(batch, *spatial, out_ch)`` array back to ``custom_vjp``, which rejected it + against the bias's own ``(out_ch,)``. The IO-dim solver now reduces every + produced leaf to its parameter's shape. At T=1 there is no history to factor, + so pp_prop must match BPTT exactly. """ factory, seed = _FAMILIES['conv_nwc_bias'] with brainstate.environ.context(precision=64): xs = _xs_for('conv_nwc_bias', 1, seed) - with pytest.raises(ValueError, match='Custom VJP bwd rule'): - online_param_gradients_singlestep_naive( - factory, xs, - algo_factory=lambda m: braintrace.pp_prop(m, decay_or_rank=0.9, vjp_method='single-step'), - ) + g_bptt = bptt_param_gradients(factory, xs) + g_online = online_param_gradients_singlestep_naive( + factory, xs, + algo_factory=lambda m: braintrace.pp_prop( + m, decay_or_rank=0.9, vjp_method='single-step'), + ) + for key in g_bptt: + rel = _rel_err(g_bptt[key], g_online[key]) + assert rel < _TOL, f'conv_nwc_bias T=1 {key}: pp_prop vs BPTT rel={rel:.3e}' # --- P1: negative-control helpers ------------------------------------------- From 7d7f098506fd7ebb780026734d32383448ddf7d6 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 12:58:15 +0800 Subject: [PATCH 11/13] Add the in-tree verified known-limitations list Replaces the untracked dev/ list AGENTS.md pointed at, which is gitignored and absent. Every entry is verified against the current suite and names the test that pins it; the 99 tests named in the table were run as a set to confirm it. Records F-23 as the load-bearing finding -- the full-window oracle path is axis-blind by design -- and maps each AGENTS.md prose item onto a finding so none survives only as prose. Notes that the suite has no skip or xfail markers at all, so "resolved" never means "stopped being run". Also records what P1 did not resolve: F-29's mechanism, F-28's square-feedback restriction, and F-24 in the layer classes themselves. --- docs/specs/2026-07-25-known-limitations.md | 130 +++++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 docs/specs/2026-07-25-known-limitations.md diff --git a/docs/specs/2026-07-25-known-limitations.md b/docs/specs/2026-07-25-known-limitations.md new file mode 100644 index 00000000..d1c829f1 --- /dev/null +++ b/docs/specs/2026-07-25-known-limitations.md @@ -0,0 +1,130 @@ +# Known limitations — verified findings list + +Status: living document +Baseline: commit `bc153da` plus the P1 phase +Supersedes: the untracked list formerly at +`dev/superpowers/specs/2026-05-26-comprehensive-test-strategy-design.md`, which +is gitignored and absent from the repository + +This is the backlog of expected-failure and improvement items that `AGENTS.md` +§ Known limitations refers to. Every entry is verified against the current test +suite, not transcribed. An entry is **active** only if a test or a documented +scope boundary pins it today. + +## Observed suite state + +Measured at the close of the P1 phase: + +``` +$ pytest braintrace/ -q -rs +2119 passed, 4 deselected, 262 warnings in 809.64s +``` + +The suite contains **no `skip` and no `xfail` markers at all** — `grep -rn +"mark.xfail\|mark.skip" braintrace/` returns nothing. So no finding below is +pinned by a deferral; each is pinned by a passing assertion or by a documented +scope boundary in the code. The 4 deselected tests are `diagnostic`-marked +random exploration, excluded from CI gating by the default `addopts` in +`pyproject.toml`; they pass as well (`pytest braintrace/ -m diagnostic -q` → +`4 passed, 2119 deselected`). + +This matters for reading the table: "resolved" always means a test asserts the +resolution, never that a test stopped being run. + +## Status legend + +- **resolved** — the claim no longer holds; a passing test pins the resolution. +- **dead** — the claim described code that no longer exists. +- **active** — the claim still holds; a test or documented boundary pins it. +- **misattributed** — the observation was real, the stated cause was not. + +## Findings + +| ID | Claim | Status | Pinned by | +|---|---|---|---| +| F-01 / F-04 | multi-state (`num_state >= 2`) HiddenGroups mishandled | resolved | `tests/snn_model_correctness_test.py::test_multi_state_hidden_groups_are_discovered` and `::test_d_rtrl_matches_bptt_on_snn_models` (num_state 1–5) | +| F-07 / F-08 / F-09 | OTTT / OTPE approximation bias | dead | removed with those algorithms in 0.3.0 | +| F-17 | implementation facts drifted from the instruction file | resolved | `braintrace/__init___test.py` | +| F-19 / F-20 | OTTT / OSTTP exactness gaps | dead | removed with those algorithms in 0.3.0 | +| F-21 | rank / decay / random-feedback configs are exact on rate models | misattributed — the cause is the oracle path (F-23), not the model | `tests/approx_correctness_test.py::test_rank_decay_random_approximations_are_exact_on_rate_model_F21`, docstring corrected | +| F-22 | exposing approximation bias needs an SNN multi-population zoo | **retired** — premise false; a multi-population SNN model (num_state 5, 3 relations) is bitwise-exact on the same path | replaced by `tests/approx_correctness_test.py::test_approximations_are_measurable_through_a_finite_window` and `tests/snn_model_correctness_test.py::test_approximation_is_measurable_on_snn_models` | +| F-23 | the full-window multi-step oracle path is blind to every learning-rule axis | active, **by design** — documented, not a defect | `tests/axis_discrimination_test.py`, both directions; warned in `online_param_gradients`' docstring | +| F-24 | `_etrace_model_test.py` factories are non-deterministic (unseeded global RNG) | active in the layer classes; neutralised for oracle use | `oracle_models_test.py::test_snn_spec_construction_is_deterministic` | +| F-25 | SNN models are silent at unit input scale, so comparisons are vacuous | active as a property; guarded | `oracle_models_test.py::test_snn_spec_is_live`, `::test_underdriven_input_scale_is_dead`, `::test_overdriven_input_scale_is_also_dead_while_still_spiking`, and `assert_model_is_live` | +| F-26 | `pp_prop` / IODim raised on conv + trainable bias | resolved | `oracle_test.py::test_pp_prop_conv_bias_matches_bptt`; `conv_nwc_bias` restored to both pp_prop family parametrizations | +| F-27 | *(never instantiated)* — reserved during planning for "an SNN spec that cannot be made live". No such spec exists: every one of the nine is live at a recorded input scale. The phenomenon that prompted the reservation is the bounded-above liveness window, recorded under F-25. | n/a | — | +| F-28 | `EProp(feedback='random')` assumes a single, direct readout whose width equals the HiddenGroup's own | active, **documented scope boundary** | `e_prop.py:117-122` states it: the hooks only see `∂L/∂h`, which has no visibility into a separate readout layer's width, so the projection matrix is square | +| F-29 | `pp_prop(decay_or_rank=1)` is a genuine approximation | **active, misattributed as approximate** — rank 1 means decay 0, so the presynaptic EMA collapses to the exact current input | `tests/approx_correctness_test.py::test_rank1_pp_prop_is_degenerate_on_a_recurrent_only_relation` pins both sides | +| F-SCAN / F-SCAN-WEIGHT | weight inside control flow raised `KeyError` | resolved | `_compiler/base_test.py::TestCheckUnsupportedOp::test_error_message_identifies_weight_variable` | +| F-SINGLESTEP | naive single-step summation does not equal BPTT even for an exact algorithm | active as a property; documented | `online_param_gradients_singlestep_naive`' docstring; used deliberately as the maximally-sensitive window in `oracle_test.py` | + +## Notes on F-23 + +This is the load-bearing finding, and it is a property rather than a bug. A +whole-sequence `MultiStepData` call makes the within-call gradient exact +reverse-mode, so the eligibility trace enters only at a sequence boundary that +does not exist. `chunked_online_param_gradients` has documented this all along +— "this is the oracle that actually validates trace correctness" — but the +assertions written against `online_param_gradients` did not heed it, which is +how F-21 and F-22 reached the wrong cause. + +The rule that follows: an assertion whose subject is a trace factorization, a +temporal recursion, a recurrence scope, a filter or a learning signal must use a +finite window and must be guarded by `assert_gradients_differ`. An assertion +whose subject is the compiler or an ETP per-primitive rule may use the +full-window path, and that is most of the `multi-step` call sites in the suite. + +## Notes on F-29 + +`decay_or_rank` accepts either a float decay or an int rank, and +`_format_decay_and_rank` maps an int through `decay = (rank - 1) / (rank + 1)`. +So `decay_or_rank=1` means decay 0, the presynaptic EMA +`eps^t = a * eps^{t-1} + (1 - a) * x_t` collapses to `x_t`, and no presynaptic +smearing is introduced at all. + +Measured on a model whose only ETP relation is the recurrent weight — so the +relation's presynaptic input *is* the hidden state — rank 1 reproduces exact +D_RTRL to round-off through a finite window: + +| variation | rank-1 rel. deviation | decay-0.5, same setup | +|---|---|---| +| chunk 1 / 2 / 4, `T` 8 / 12 / 16 | 1e-10 … 7e-9 | 9e-4 … 1.2e-2 | +| recurrent spectral radius 0.25 … 4.97 | 6e-10 … 1.6e-8 | 4e-3 … 1.5e-1 | +| `n_rec` 4 → 8 | 2e-8 | 1.2e-1 | + +float32 round-off on these trees is ~1e-8, so every figure in the first column +is noise. The deviation does **not** grow with the spectral radius, which rules +out "the recurrence is too weak for history to matter" as the explanation. + +It is a real approximation as soon as the presynaptic input carries an external +component: 0.55 on a variant whose ETP weight is the input weight, and 0.31 to +0.79 on all nine SNN specs, whose projections consume `concat(input, spikes)`. + +**The mechanism is not established — only the boundary is.** Two consequences: +`pp_prop(decay_or_rank=1)` on a recurrent-only relation is unusable as a +positive control for approximation error, and any future axis test that wants +one must either use a float decay or a model with an external-input relation. + +## Mapping from the `AGENTS.md` prose list + +`AGENTS.md` described its limitations in prose. Each item maps onto a finding +above, so nothing survives only as prose: + +| Prose item | Disposition | +|---|---| +| approximation-mode validity beyond shallow depth | successor to F-21 / F-23; measured through a finite window, and bounded by F-29 for one configuration | +| heterogeneous-population leak resolution | resolved — `tests/snn_model_correctness_test.py::test_heterogeneous_leaks_do_not_break_exactness` | +| target-signal threading under JIT | dead — died with `OSTTP`'s `y_target` path | +| single-readout / feedback-shape assumptions | active — F-28, a documented scope boundary in `e_prop.py` rather than a defect | +| gaps in cross-algorithm equivalence coverage | successor to F-23; `tests/axis_discrimination_test.py` covers the pairs | + +## What P1 did not resolve + +Stated so the list is not read as exhaustive: + +- **F-29's mechanism.** Boundary measured, cause unexplained. +- **F-28.** A square hidden×hidden feedback matrix is a real restriction on + `EProp(feedback='random')`. It is documented, not lifted, and no test + exercises a separate readout layer of a different width. +- **F-24 in the layer classes.** The oracle specs re-seed at construction; the + `_etrace_model_test.py` constructors still draw from the global RNG. From d23d7a927cbbd7b63f26b0167ea2d22405d1ae49 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 13:00:36 +0800 Subject: [PATCH 12/13] Record P1 outcomes: revised P2/P3 criteria, lessons learned P2 and P3 acceptance criteria now name a finite-window oracle path. As written they were element-wise equalities measured through the full-window path, where every algorithm returns BPTT regardless of its axis coordinates -- they would have passed without guarding the refactor at all. Risk 1 gains the same qualification: golden values frozen on the full-window path are identical across all five algorithms and would detect nothing. Rewrites the P1 section to match what was built, retires risk 2 (no hidden_group.py defect was found, so P3 inherits today's representation), resolves risk 3, adds an eight-item lessons-learned section, and repoints AGENTS.md off the absent dev/ list onto the in-tree findings list. Two lessons are about mistakes made during this phase rather than before it: a probe that reported zero deviation everywhere because its model had no input drive, and a first version of the F-22 replacement test whose 1e-9 threshold sat below float32 round-off and so asserted noise. --- AGENTS.md | 21 +- .../2026-07-25-algorithm-axes-roadmap.md | 182 +++++++++++++----- 2 files changed, 153 insertions(+), 50 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 2ae40054..676d8f64 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -147,13 +147,20 @@ Pin exact names/versions in `pyproject.toml` / requirements files,. ## Known limitations -First-cut SNN algorithms pass smoke/cross-checks but carry approximation edges -and rough spots (e.g. approximation-mode validity beyond shallow depth, -heterogeneous-population leak resolution, target-signal threading under JIT, -single-readout/feedback-shape assumptions, and gaps in cross-algorithm -equivalence coverage). These are enumerated and mapped to concrete improvement -actions in the test-strategy findings list under `dev/`. Treat that list as the -backlog of expected-failure / improvement items rather than duplicating it here. +First-cut SNN algorithms pass smoke and cross-checks but carry approximation +edges and rough spots. These are enumerated, verified against the test suite, +and mapped to concrete improvement actions in the findings list at +`docs/specs/2026-07-25-known-limitations.md`. Treat that list as the backlog of +expected-failure and improvement items rather than duplicating it here. + +One rule from that list is load-bearing enough to state here: a gradient +assertion whose subject is a **learning-rule property** — a trace +factorization, a temporal recursion, a recurrence scope, a filter, a learning +signal — must be measured through a *finite-window* oracle path +(`chunked_online_param_gradients`). A whole-sequence multi-step VJP has no +truncation left to approximate and returns BPTT for every algorithm at every +hyperparameter, so such an assertion passes vacuously there. Assertions about +the compiler or an ETP per-primitive rule may use the whole-sequence path. ## Docstring style (NumPy-doc) diff --git a/docs/specs/2026-07-25-algorithm-axes-roadmap.md b/docs/specs/2026-07-25-algorithm-axes-roadmap.md index 4bf00ce4..0cb2a5f6 100644 --- a/docs/specs/2026-07-25-algorithm-axes-roadmap.md +++ b/docs/specs/2026-07-25-algorithm-axes-roadmap.md @@ -230,34 +230,33 @@ direction-alignment metric helpers were retained as the basis of P5. Verified: `pytest braintrace/` → 2062 passed, 1 skipped; `mypy braintrace` clean. -### P1 — Compiler: multi-timescale and heterogeneous populations - -Make the surviving rules actually run on realistic brain-simulation models — -LIF + ALIF + multi-timescale synapses, heterogeneous leaks, multi-state -HiddenGroups — before widening the algorithm family. - -**Scope has to be reconstructed first.** `AGENTS.md` points at a -known-limitation findings list under `dev/`, but `dev/` is gitignored and absent -from the repository. The only machine-readable remnant is a single skipped test -(`approx_correctness_test.py::test_approximations_diverge_on_snn_multipopulation_DEFERRED`, -finding F-22). Everything else on that list has either been resolved — F-19 and -F-20 died with the algorithms they described; F-SINGLESTEP was resolved into a -positive direction-alignment assertion in `oracle_test.py` — or exists only as -prose in `AGENTS.md`. **The first task of P1 is to rebuild the list in-tree**, -from `AGENTS.md`'s summary plus a sweep of the test suite, so it stops living in -an untracked directory. Expect the rebuilt list to be shorter than the original: -part of it is already done. - -Known items from `AGENTS.md`: heterogeneous-population leak resolution, -multi-state HiddenGroups, approximation validity beyond shallow depth, -single-readout / feedback-shape assumptions, cross-algorithm equivalence gaps. - -This phase owns `hidden_group.py`'s Jacobian path, which P3 also needs — doing -it first is what keeps P3 from fighting a moving target. - -**Acceptance:** the reconstructed limitation list is committed under -`docs/specs/`; every item on it has either a passing test or an explicitly -documented scope boundary; no expected-failure item silently remains. +### P1 — Verification harness and the limitation list — **done** + +Retitled during implementation. P1 was scoped as compiler work; measurement +showed the compiler already passes every one of its targets, and that the +instrument was what failed. Full spec: +[`2026-07-25-p1-verification-harness.md`](2026-07-25-p1-verification-harness.md). + +Delivered: negative-control oracle helpers (`assert_model_is_live`, +`assert_gradients_differ`) and pytree/unit-aware gradient comparison; +documented window semantics on every oracle entry point; deterministic and live +SNN model specs; an axis-discrimination meta-test; gradient-correctness tests on +realistic SNN models (multi-timescale, per-neuron heterogeneous leaks, +`num_state` 1–5, E/I populations); the conv-bias IO-dim fix (F-26); and the +in-tree findings list at +[`2026-07-25-known-limitations.md`](2026-07-25-known-limitations.md). + +`hidden_group.py` was **not** modified — no defect was found in its Jacobian +path — so P3 inherits today's representation and risk 2 below is retired rather +than mitigated. + +**Acceptance:** met. The reconstructed limitation list is committed under +`docs/specs/`; every item has a passing test or a documented scope boundary; the +one skipped test (F-22) is retired rather than deferred, and the suite now +carries no `skip` or `xfail` markers at all. + +Verified: `pytest braintrace/` → 2119 passed, 4 deselected (`diagnostic`, which +pass separately); `mypy braintrace` clean. ### P2 — Axis decomposition @@ -273,7 +272,12 @@ Execution options (`vjp_method`, `fast_solve`, `trace_dtype`, `chunked_trace`, **Acceptance:** - Element-wise equality against golden values frozen *before* the refactor, for - all five surviving algorithms. + all five surviving algorithms — frozen and compared through a **finite-window** + oracle path (`chunked_online_param_gradients`, `chunk_size` < `T`). The + full-window multi-step path returns BPTT for every algorithm regardless of its + axis coordinates (F-23), so golden values captured there would be identical + across all five and would guard nothing. Guard each comparison with + `assert_gradients_differ` between any two presets that differ in an axis. - `decay_or_rank=0.9` equals the new two-sided `(0.9, 0.9)` element-wise. - Illegal axis combinations raise a readable error naming the legal pairings, with test coverage. @@ -292,11 +296,15 @@ today's coupled path, and add the intermediate sparse representation for `1 < n < diameter`. **Acceptance (two-sided squeeze):** -- `n = 1` equals the current `D_RTRL` element-wise (regression guard). +- `n = 1` equals the current `D_RTRL` element-wise, compared through a + finite-window path (regression guard). On the full-window path this holds for + every algorithm and so proves nothing (F-23). - Whatever configuration expresses `recurrence_scope = coupled` after the - refactor equals the current `OSTLRecurrent` element-wise — regardless of - whether that configuration turns out to be a point on the `n` scale or a - sibling value beside it (regression guard for the existing `True` branch). + refactor equals the current `OSTLRecurrent` element-wise — again through a + finite-window path — regardless of whether that configuration turns out to be + a point on the `n` scale or a sibling value beside it. + `axis_discrimination_test.py` pins that `D_RTRL` and `OSTLRecurrent` *are* + distinguishable on that path, so this comparison has content. - `n ≥ graph diameter` equals the **BPTT oracle** element-wise, on models whose hidden-state coupling the compiler fully captures. Full RTRL and BPTT compute the same total gradient, so `oracle.py` is the correct instrument and no @@ -364,16 +372,21 @@ statistical class. 1. **P2 changes numerics silently.** Mitigation: freeze reference gradients for all five surviving algorithms as golden values *before* touching the engine; - assert element-wise equality after. Separately, the `IODim` decay split must - prove `decay_or_rank=0.9` equals the new `(0.9, 0.9)` element-wise. -2. **P1 and P3 collide in `hidden_group.py`.** Both rework the hidden→hidden - Jacobian path. Mitigation: P1 owns that file and lands first; P3 consumes the - representation P1 leaves behind rather than introducing a second one. -3. **The P1 scope is not actually written down.** The findings list lives in an - untracked `dev/`, and the in-tree remnants are already stale — one docstring - still cross-referenced a test renamed some time ago. Mitigation: rebuilding - the list in-tree is P1's first task, not an assumption, and the rebuilt list - must be verified against the test suite rather than transcribed. + assert element-wise equality after, **through a finite-window oracle path**. + Golden values captured through the full-window path are identical across all + five algorithms (F-23) and would not detect any change. Separately, the + `IODim` decay split must prove `decay_or_rank=0.9` equals the new + `(0.9, 0.9)` element-wise. +2. ~~**P1 and P3 collide in `hidden_group.py`.**~~ **Retired.** P1 found no + defect in the hidden→hidden Jacobian path and did not modify the file, so + there is no collision: P3 inherits today's representation directly. +3. ~~**The P1 scope is not actually written down.**~~ **Resolved by P1.** The + list is now in-tree at + [`2026-07-25-known-limitations.md`](2026-07-25-known-limitations.md), verified + against the suite rather than transcribed, and the stale `dev/` docstring + cross-reference is gone. The residual form of this risk for later phases is + different and worse — see lesson 2 below: a criterion can be *written down* + and still be vacuous if it names the wrong oracle path. 4. **Statistical tests are flaky in CI** (P4). Mitigation: fixed seeds, generous intervals, a separate slow-test marker. 5. **`modulatory` recreates OSTTP's plumbing mistake** (P4). Mitigation: the @@ -383,3 +396,86 @@ statistical class. algorithm families ship together. Mitigation: phases merge independently behind the axis interfaces, and P5 runs on every merge so regressions surface per phase rather than at release. + +## Lessons learned during implementation + +Recorded during P1, against commit `bc153da`. These are the things the roadmap +got wrong or could not have known, kept here because later phases rest on them. + +1. **The instrument was the defect, not the compiler.** P1 was scoped as + compiler work on multi-timescale and heterogeneous populations. Every one of + those targets already passed: HiddenGroups with `num_state` up to 5, + per-neuron heterogeneous `tau_mem` and `tau_a`, multi-timescale synapses + (`tau_mem` / `tau_syn` / `tau_a` / `tau_std` / `tau_f` / `tau_d`), and E/I + population splits all reproduce BPTT. Measure before scoping a phase around + a suspected defect. + +2. **A full-sequence multi-step VJP is BPTT, so it cannot see any axis** + (F-23). `D_RTRL`, `OSTLRecurrent`, `EProp` at any `kappa_filter_decay` and + `pp_prop` at any `decay_or_rank` return bitwise-identical gradients through + `online_param_gradients`. This is correct semantics — no truncation, nothing + to approximate — but it silently voids any assertion whose subject is a + learning-rule axis. `chunked_online_param_gradients` documented this from the + start; the tests written against the other entry point did not heed it. **Any + acceptance criterion in this roadmap phrased as "equals X element-wise" must + name a finite-window path.** + +3. **A wrong cause propagates further than a wrong measurement.** F-21 observed + real behaviour and attributed it to the model being a single-HiddenGroup rate + model. F-22 then deferred an entire work item — an SNN multi-population model + zoo — on that attribution, and P5 inherited it. Building exactly that model + (`ALIF_ExpCo` with an E/I split, 3 relations, `num_state` 5) reproduced the + same bitwise exactness, so no zoo could ever have helped. The fix was one + parameter: `chunk_size`. + +4. **Vacuous tests look like passing tests.** Three distinct ways a gradient + assertion in this repository can assert nothing: the reference gradient is + zero (a silent SNN — at unit input scale the conductance-based models never + reach threshold); the model differs between the two sides of the comparison + (`_etrace_model_test.py` constructors draw from the unseeded global RNG, so + `factory()` returns a different network per call); or the oracle path cannot + see the knob under test (F-23). Spiking is *not* sufficient for the first, + and the live window is bounded **above** as well as below — driven hard, + `ALIF_Delta` keeps spiking at rate 0.60 while its surrogate derivative + saturates and the BPTT gradient returns to exactly zero. Hence + `assert_model_is_live` keys on gradient norm rather than spike rate, and + every new axis assertion carries `assert_gradients_differ`. + + This bit twice, and the second time it was self-inflicted: a probe written + *while investigating F-29* silently reported 0.0 deviation for every + configuration because the model under it had no input drive. Run the liveness + helper inside throwaway probes too, not just inside committed tests. + +5. **A threshold below round-off is not an assertion.** The first version of the + F-22 replacement test demanded a relative deviation of only `1e-9` between an + approximate and an exact algorithm. float32 round-off on these gradient trees + is around `1e-8`, so the test would have passed on numerical noise for any + configuration at all — and it failed only because one configuration was + *even quieter* than noise. Axis assertions now use `1e-6` against deviations + of `1e-3` to `1e-1`. + +6. **"Approximate" is a claim about a configuration, not about a class** + (F-29). `pp_prop(decay_or_rank=1)` is nominally an approximation, but an int + rank maps through `decay = (rank - 1) / (rank + 1)`, so rank 1 means decay 0, + the presynaptic EMA collapses to the exact current input, and on a model whose + only ETP relation is the recurrent weight it reproduces exact `D_RTRL` to + round-off — across chunk sizes, sequence lengths, widths and recurrent + spectral radii from 0.25 to 4.97. It *is* a real approximation (0.31 to 0.79) + once the relation's presynaptic input carries an external component, as every + SNN spec's does. Consequence for P2 and P5: when picking a positive control + for approximation error, verify the chosen configuration actually deviates on + the chosen model instead of assuming its class name guarantees it. + +7. **`xy_to_dw` rules may return un-reduced, per-position leaves.** Conv's rule + deliberately defers the spatial sum of the bias Jacobian to `_conv_dt_to_t`, + which only the param-dim path calls. Any new solver that consumes `xy_to_dw` + directly must reduce produced leaves to the parameter's own shape (F-26). + This matters for P2: the axis strategies will introduce new contraction + paths. + +8. **The removed algorithms took findings with them, and left stale + cross-references behind.** F-07/F-08/F-09/F-19/F-20 died with OTTT / OTPE / + OSTTP, and the `AGENTS.md` prose item "target-signal threading under JIT" + died with `OSTTP`'s `y_target` path. Meanwhile a docstring still pointed at + `dev/superpowers/specs/...`, a path that is gitignored and absent. Findings + lists must live in-tree. From 84381462c6d33503955328d104f2a83a4a54ffc3 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Sat, 25 Jul 2026 13:23:35 +0800 Subject: [PATCH 13/13] Close out P1: observed verification numbers and the mypy import-follow fix pytest braintrace/ -> 2119 passed, 4 deselected (the diagnostic-marked tests, which pass separately); mypy braintrace clean. Against the P0 bar of 2062 passed / 1 skipped that is +57 tests and the skip retired, leaving no skip or xfail marker anywhere in the suite. mypy needed one config change, caused by this work rather than pre-existing: exclude does not apply to files reached by an import, so oracle_models.py's import of the _etrace_model_test.py layer classes re-admitted that excluded file and surfaced 52 pre-existing brainpy.state.* attribute errors. A scoped ignore_errors override for that single module restores the policy the exclude pattern already expresses. Recorded as lesson 8. Also reconciles the spec's findings table with what implementation actually found: F-26 resolved rather than carried forward, F-27 never instantiated, F-28 active as a documented scope boundary, and F-29 added. --- .../2026-07-25-algorithm-axes-roadmap.md | 13 +++++- .../2026-07-25-p1-verification-harness.md | 40 ++++++++++++++++++- pyproject.toml | 13 ++++++ 3 files changed, 64 insertions(+), 2 deletions(-) diff --git a/docs/specs/2026-07-25-algorithm-axes-roadmap.md b/docs/specs/2026-07-25-algorithm-axes-roadmap.md index 0cb2a5f6..c3bedcac 100644 --- a/docs/specs/2026-07-25-algorithm-axes-roadmap.md +++ b/docs/specs/2026-07-25-algorithm-axes-roadmap.md @@ -473,7 +473,18 @@ got wrong or could not have known, kept here because later phases rest on them. This matters for P2: the axis strategies will introduce new contraction paths. -8. **The removed algorithms took findings with them, and left stale +8. **mypy's `exclude` does not apply to files reached by an import.** The oracle + SNN specs wrap the layer classes in `_etrace_model_test.py` rather than + duplicating them, which is the right call for keeping one definition of each + model — but it means a non-excluded module imports an excluded one. mypy then + type-checks that test file anyway and reported 52 pre-existing + `brainpy.state.*` attribute errors, turning a clean gate red without any + change to the file itself. Fixed with a scoped `ignore_errors` override so the + existing "test files are outside the typed surface" policy holds however the + file is reached. Watch for this whenever shipped code imports a `_test.py` + module. + +9. **The removed algorithms took findings with them, and left stale cross-references behind.** F-07/F-08/F-09/F-19/F-20 died with OTTT / OTPE / OSTTP, and the `AGENTS.md` prose item "target-signal threading under JIT" died with `OSTTP`'s `y_target` path. Meanwhile a docstring still pointed at diff --git a/docs/specs/2026-07-25-p1-verification-harness.md b/docs/specs/2026-07-25-p1-verification-harness.md index 7e8b5486..2619bd7c 100644 --- a/docs/specs/2026-07-25-p1-verification-harness.md +++ b/docs/specs/2026-07-25-p1-verification-harness.md @@ -1,6 +1,6 @@ # P1 — Verification harness: make the oracle axis-aware, rebuild the limitation list -Status: approved, ready for implementation +Status: implemented Parent: [`2026-07-25-algorithm-axes-roadmap.md`](2026-07-25-algorithm-axes-roadmap.md) § P1 Baseline: commit `bc153da` Target release: 0.3.0 @@ -147,6 +147,19 @@ Reconstructed disposition: | **F-25** | SNN zoo silent at default scale → vacuous comparisons | **new, active** | D3 | | **F-26** | `pp_prop` / IODim raises on conv + trainable bias | **new, active** (pre-existing) | `oracle_test.py::test_pp_prop_conv_bias_known_limitation` | +This was the table as scoped. Two entries moved during implementation, and one +finding was added that the scoping did not anticipate: + +| ID | Change from the table above | +|---|---| +| F-26 | **resolved**, not carried forward — the IO-dim solver now reduces produced leaves to the parameter's shape, and the pinning test became `oracle_test.py::test_pp_prop_conv_bias_matches_bptt` | +| F-27 / F-28 | F-27 was reserved for "an SNN spec that cannot be made live" and was never instantiated; F-28 (`EProp(feedback='random')` assumes a single readout of the HiddenGroup's own width) is **active as a documented scope boundary** in `e_prop.py`, not a defect | +| **F-29** | **new** — `pp_prop(decay_or_rank=1)` is not an approximation at all on a recurrent-only relation: rank 1 means decay 0, so the presynaptic EMA collapses to the exact input. Found by the D2 work, since a config held out as a positive control turned out not to deviate | + +The authoritative, verified list is +[`2026-07-25-known-limitations.md`](2026-07-25-known-limitations.md); this table +records what the scoping expected so the two can be compared. + `AGENTS.md`'s prose items map on as follows, and the mapping is recorded in the list so nothing survives only as prose: @@ -268,6 +281,31 @@ Verification for the phase: full `pytest braintrace/` and `mypy braintrace` clean, matching the P0 bar (2062 passed, 1 skipped at `928219b`; the skip count drops to 0 once D6 lands). +**Observed at completion:** + +``` +$ pytest braintrace/ -q -rs +2119 passed, 4 deselected, 262 warnings in 788.97s + +$ pytest braintrace/ -m diagnostic -q +4 passed, 2119 deselected + +$ mypy braintrace +Success: no issues found in 56 source files +``` + +2062 → 2119 passed (+57), and 1 skipped → 0: F-22's skip was retired rather than +deferred, and the suite now carries no `skip` or `xfail` marker anywhere. The 4 +deselected are the `diagnostic`-marked tests that the default `addopts` excludes +from CI gating; they are run above and pass. + +`mypy` needed one config change to stay clean, and it was caused by this work: +`exclude` does not apply to files reached by an import, so +`oracle_models.py`'s import of the `_etrace_model_test.py` layer classes +re-admitted that file and surfaced 52 pre-existing `brainpy.state.*` +attribute errors. A scoped `ignore_errors` override for that one module +restores the existing policy. See lesson 8 in the roadmap. + ## Acceptance Restating the roadmap's P1 criteria against this scope: diff --git a/pyproject.toml b/pyproject.toml index 11be0610..75ede277 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -114,6 +114,19 @@ implicit_optional = true warn_redundant_casts = true no_implicit_reexport = false +# `exclude` above removes test files from the checked set, but it does NOT apply +# to files pulled in by an import. `_algorithm/oracle_models.py` imports the +# layer classes from `_etrace_model_test.py` (deliberately -- the oracle wraps +# those models rather than duplicating them), which re-admits that file and +# surfaces its `brainpy.state.*` attribute errors. Those come from brainpy's +# dynamic re-exports, the same reason `brainpy.*` is listed as untyped below; +# they are not our errors, and the file is already outside the typed public +# surface by the `exclude` policy. Suppress reporting for it specifically so the +# policy holds however the file is reached. +[[tool.mypy.overrides]] +module = ["braintrace._etrace_model_test"] +ignore_errors = true + # Third-party scientific deps that ship without inline types or a py.typed # marker. Scoped per-module so missing stubs here never mask errors in our # own code.