braintrace 0.2.2__tar.gz → 0.2.4__tar.gz

{braintrace-0.2.2 → braintrace-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.2.2
+Version: 0.2.4
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chao.brain@qq.com>
 License: Apache-2.0 license
@@ -28,6 +28,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: brainstate>=0.5.2
 Requires-Dist: brainunit
+Requires-Dist: brainevent
 Requires-Dist: brainpy.state
 Requires-Dist: braintools
 Provides-Extra: cpu

{braintrace-0.2.2 → braintrace-0.2.4}/braintrace/__init__.py

@@ -66,7 +66,9 @@ Examples
 """
 
 
-from typing import TYPE_CHECKING
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 from . import nn
 from ._compile import compile
@@ -92,6 +94,7 @@ from ._algorithm import (
     PresynapticTrace,
 )
 from ._compiler import (
+    ControlFlowPolicy,
     ETraceGraph,
     compile_etrace_graph,
     HiddenParamOpRelation,
@@ -113,6 +116,9 @@ from ._compiler import (
 from ._op import (
     ETPPrimitive,
     matmul,
+    grouped_matmul,
+    embedding,
+    einsum,
     element_wise,
     conv,
     sparse_matmul,
@@ -166,6 +172,9 @@ __all__ = [
 
     # ETP primitives (user API)
     'matmul',
+    'grouped_matmul',
+    'embedding',
+    'einsum',
     'element_wise',
     'conv',
     'sparse_matmul',
@@ -183,6 +192,7 @@ __all__ = [
     'ETraceGraphExecutor',
 
     # compiler
+    'ControlFlowPolicy',
     'ETraceGraph',
     'compile_etrace_graph',
     'HiddenGroup',
@@ -244,7 +254,7 @@ _DEPRECATED_LEGACY = {
 }
 
 
-def __getattr__(name):
+def __getattr__(name: str) -> Any:
     if name in _DEPRECATED_LEGACY:
         import warnings
         warnings.warn(
@@ -258,5 +268,5 @@ def __getattr__(name):
     raise AttributeError(f'module {__name__!r} has no attribute {name!r}')
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return sorted(list(__all__) + list(_DEPRECATED_LEGACY))

{braintrace-0.2.2 → braintrace-0.2.4}/braintrace/__init___test.py

@@ -148,7 +148,14 @@ def test_weight_used_inside_scan_raises_not_implemented():
     within a control-flow op. Previously (F-SCAN-WEIGHT) the message construction
     indexed the hidden-path map with the weight var and a KeyError escaped instead;
     the message now renders the weight variable directly, so the intended
-    NotImplementedError surfaces."""
+    NotImplementedError surfaces.
+
+    Since Phase 2 canonicalization, statically short scans are unrolled at
+    extraction time instead of erroring, so this uses a length above
+    ``ControlFlowPolicy.scan_unroll_limit`` to reach the unsupported-op path
+    (after a SCAN_UNROLL_SKIPPED warning). Since Phase 4, structured scan
+    descent handles over-limit ETP scans by default, so this pins the
+    opt-out path via ``ControlFlowPolicy(scan_descent='off')``."""
 
     class ScanModel(brainstate.nn.Module):
         def __init__(self):
@@ -160,15 +167,17 @@ def test_weight_used_inside_scan_raises_not_implemented():
         def update(self, x):
             def body(carry, _):
                 return braintrace.matmul(carry, self.w.value), None  # weight inside scan
-            out, _ = jax.lax.scan(body, self.h.value, xs=None, length=2)
+            out, _ = jax.lax.scan(body, self.h.value, xs=None, length=64)
             self.h.value = jax.nn.tanh(out)
             return self.h.value
 
     model = ScanModel()
     brainstate.nn.init_all_states(model, batch_size=1)
-    algo = braintrace.D_RTRL(model)
+    algo = braintrace.D_RTRL(
+        model, control_flow=braintrace.ControlFlowPolicy(scan_descent='off'))
     with pytest.raises(NotImplementedError):
-        algo.compile_graph(jnp.ones((1, 4), dtype='float32'))
+        with pytest.warns(UserWarning, match='NOT unrolled'):
+            algo.compile_graph(jnp.ones((1, 4), dtype='float32'))
 
 
 # --- Task 6: legacy + nn deprecations emit DeprecationWarning ----------------

{braintrace-0.2.2 → braintrace-0.2.4}/braintrace/_algorithm/__init__.py

@@ -21,6 +21,8 @@ and paper-faithful SNN algorithms (EProp, OSTLRecurrent, OSTLFeedforward, OTPE,
 OTTT, OSTTP).
 """
 
+from __future__ import annotations
+
 from ._common import FixedRandomFeedback, KappaFilter, PresynapticTrace
 from .base import ETraceAlgorithm, EligibilityTrace
 from .d_rtrl import D_RTRL

{braintrace-0.2.2 → braintrace-0.2.4}/braintrace/_algorithm/_common.py

@@ -15,6 +15,8 @@
 
 """Shared helpers for the SNN online-learning algorithms."""
 
+from __future__ import annotations
+
 from functools import partial
 from typing import Any, Dict, Optional
 
@@ -23,7 +25,7 @@ import jax
 import jax.numpy as jnp
 import brainunit as u
 
-from braintrace._typing import PyTree
+from braintrace._typing import ArrayLike, PyTree
 
 __all__ = [
     'PresynapticTrace',
@@ -42,12 +44,12 @@ class _ZeroResetState(brainstate.ShortTermState):
     for ``batch_size`` (or prepended for scalar-shaped states).
     """
 
-    def __init__(self, init_value):
+    def __init__(self, init_value: Any) -> None:
         super().__init__(init_value)
         self._init_shape = jnp.shape(init_value)
         self._init_dtype = init_value.dtype
 
-    def reset_state(self, batch_size: Optional[int] = None, **kwargs):
+    def reset_state(self, batch_size: Optional[int] = None, **kwargs: Any) -> None:
         """Re-zero the state at its original shape.
 
         Parameters
@@ -105,13 +107,13 @@ class PresynapticTrace(_ZeroResetState):
 
     __module__ = 'braintrace'
 
-    def __init__(self, init_value, leak: float):
+    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):
+    def update(self, x: ArrayLike) -> Any:
         r"""Apply one accumulation step :math:`\hat{a} \leftarrow \lambda \cdot \hat{a} + x`.
 
         Parameters
@@ -129,9 +131,14 @@ class PresynapticTrace(_ZeroResetState):
 
 
 class KappaFilter(_ZeroResetState):
-    r"""Low-pass output-side filter used by EProp.
+    r"""Low-pass filter helper state.
+
+    :class:`~braintrace.EProp` no longer uses this class directly — it filters
+    the eligibility trace internally instead. ``KappaFilter`` remains public
+    and available for user-side filtering of an output-side (or any other)
+    signal outside the algorithm's own hooks.
 
-    The filter smooths the output-side signal following
+    The filter smooths the signal following
     :math:`x_{\mathrm{filt}} \leftarrow (1-\kappa) \cdot x + \kappa \cdot x_{\mathrm{filt}}`.
 
     Parameters
@@ -163,13 +170,13 @@ class KappaFilter(_ZeroResetState):
 
     __module__ = 'braintrace'
 
-    def __init__(self, init_value, kappa: float):
+    def __init__(self, init_value: Any, kappa: float) -> None:
         super().__init__(init_value)
         if not (0.0 <= kappa < 1.0):
             raise ValueError(f'kappa must be in [0, 1); got {kappa}')
         self.kappa = float(kappa)
 
-    def update(self, x):
+    def update(self, x: ArrayLike) -> Any:
         r"""Apply one low-pass step :math:`x_{\mathrm{filt}} \leftarrow (1-\kappa) x + \kappa\, x_{\mathrm{filt}}`.
 
         Parameters
@@ -230,14 +237,14 @@ class FixedRandomFeedback:
 
     __module__ = 'braintrace'
 
-    def __init__(self, n_target: int, n_layer: int, key, init_scale: float = 0.1):
+    def __init__(self, n_target: int, n_layer: int, key: Any, init_scale: float = 0.1) -> None:
         self.B = jax.lax.stop_gradient(
-            init_scale * jax.random.normal(key, (n_target, n_layer))
+            init_scale * brainstate.random.normal(size=(n_target, n_layer), key=key)
         )
         self.n_target = int(n_target)
         self.n_layer = int(n_layer)
 
-    def project(self, y_target):
+    def project(self, y_target: Any) -> Any:
         """Project the target onto the frozen feedback matrix.
 
         Parameters
@@ -267,7 +274,7 @@ def extract_y_target(args: tuple, *, index: int = -1) -> Optional[jax.Array]:
 def _reset_state_in_a_dict(
     state_dict: Dict[Any, brainstate.State],
     batch_size: Optional[int],
-):
+) -> None:
     """
     Reset the values in a dictionary of states to zero.
 
@@ -293,7 +300,7 @@ def _reset_state_in_a_dict(
 def _zeros_like_batch_or_not(
     batch_size: Optional[int],
     x: jax.Array
-):
+) -> Any:
     """
     Create a zeros array with the same shape and type as the input array,
     optionally including a batch dimension.
@@ -325,7 +332,7 @@ def _batched_zeros_like(
     batch_size: Optional[int],
     num_state: int,  # the number of hidden states
     x: jax.Array  # the input array
-):
+) -> Any:
     """
     Create a batched zeros array with the same shape as the input array,
     extended by the number of hidden states.
@@ -354,7 +361,7 @@ def _batched_zeros_like(
         return u.math.zeros((batch_size, *x.shape, num_state), x.dtype)
 
 
-def _sum_dim(xs: jax.Array, axis: int = -1):
+def _sum_dim(xs: jax.Array, axis: int = -1) -> Any:
     """
     Sums the elements along the last dimension of each array in a PyTree.
 
@@ -373,7 +380,7 @@ def _sum_dim(xs: jax.Array, axis: int = -1):
     return jax.tree.map(lambda x: u.math.sum(x, axis=axis), xs)
 
 
-def _unit_safe_add(a, b):
+def _unit_safe_add(a: Any, b: Any) -> Any:
     """Add two leaves, stripping units only when one side has units and the other does not.
 
     Gradient contributions for the same weight may come from paths that
@@ -390,7 +397,7 @@ def _unit_safe_add(a, b):
     return u.math.add(a, b)
 
 
-def _extract_leaf(pytree_val: PyTree, leaf_idx: int):
+def _extract_leaf(pytree_val: PyTree, leaf_idx: int) -> Any:
     """Return the leaf at ``leaf_idx`` in ``jax.tree.leaves(pytree_val)``.
 
     Bare arrays (treedef with a single leaf) return the array unchanged.
@@ -409,7 +416,7 @@ def _extract_leaf(pytree_val: PyTree, leaf_idx: int):
 def _wrap_leaves_as_pytree(
     reference_pytree: PyTree,
     leaf_grads: Dict[int, jax.Array],
-):
+) -> Any:
     """Build a pytree matching ``reference_pytree`` with ``leaf_grads``
     inserted at the given leaf indices; any other leaf is zero-filled.
 
@@ -441,7 +448,7 @@ def _wrap_leaves_as_pytree(
 
 
 def _route_grads_by_path(
-    relation,
+    relation: Any,
     per_key_grads: Dict[str, jax.Array],
     weight_vals: Dict[Any, PyTree],
     target_dict: Dict[Any, PyTree],
@@ -449,7 +456,7 @@ def _route_grads_by_path(
     """Route per-key gradients from a dict-API rule into per-path pytrees.
 
     Both D-RTRL and ES-D-RTRL share this bookkeeping: for each key in
-    ``per_key_grads`` (returned by ``xy_to_dw`` or ``yw_to_w``), look up the
+    ``per_key_grads`` (returned by ``xy_to_dw`` or ``dt_to_t``), look up the
     owning ``ParamState`` path and the leaf index from ``relation``, accumulate
     into ``per_path``, then wrap with ``_wrap_leaves_as_pytree`` and merge into
     ``target_dict`` via ``_update_dict``.
@@ -478,7 +485,7 @@ def _update_dict(
     key: Any,
     value: PyTree,
     error_when_no_key: Optional[bool] = False
-):
+) -> None:
     """Update the dictionary.
 
     If the key exists, then add the value to the existing value.

{braintrace-0.2.2 → braintrace-0.2.4}/braintrace/_algorithm/base.py

@@ -19,6 +19,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from typing import Dict, Any, Optional
 
 import brainstate
@@ -79,8 +81,13 @@ class ETraceAlgorithm(brainstate.nn.Module):
         The other states.
     is_compiled : bool
         Whether the etrace algorithm has been compiled.
-    running_index : brainstate.ParamState[int]
-        The running index.
+    running_index : brainstate.LongTermState[int]
+        Zero-based count of previously-completed ``update()`` calls: it reads
+        ``0`` during the first call, ``1`` during the second, and so on --
+        i.e. the value used *inside* a given ``update()`` is the
+        pre-increment count, not the 1-based call number. It is incremented
+        by one immediately after each ``update()`` completes, and reset back
+        to ``0`` by ``reset_state()``.
     """
     __module__ = 'braintrace'
 
@@ -89,7 +96,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         model: brainstate.nn.Module,
         graph_executor: ETraceGraphExecutor,
         name: Optional[str] = None,
-    ):
+    ) -> None:
         super().__init__(name=name)  # type: ignore[call-arg]  # brainstate hides Module.__init__ from type checkers
 
         # the model
@@ -210,7 +217,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         assert self._other_states is not None
         return self._other_states
 
-    def _split_state(self):
+    def _split_state(self) -> None:
         # --- the state separation --- #
         #
         # [NOTE]
@@ -228,7 +235,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
             self._other_states
         ) = self.graph.module_info.retrieved_model_states.split(brainstate.ParamState, brainstate.HiddenState, ...)
 
-    def compile_graph(self, *args) -> None:
+    def compile_graph(self, *args: Any) -> None:
         r"""
         Compile the eligibility trace graph of the relationship between etrace weights, states and operators.
 
@@ -265,6 +272,36 @@ class ETraceAlgorithm(brainstate.nn.Module):
             # --- the model etrace graph -- #
             self.graph_executor.compile_graph(*args)
 
+            # Structured scan descent (Phase 4): relations discovered inside
+            # a descended scan body need a per-substep trace fold, and
+            # hidden groups descended from a scan body carry a leading
+            # substep axis on their Jacobians — both only handled by
+            # algorithms declaring ``_supports_scan_descent = True``. A
+            # descended *group* can exist without any descended relation
+            # (all body weights routed through plain ops), so check both.
+            # Gate here — after the graph is built, before any trace state
+            # is initialized against it.
+            if (
+                (
+                    any(
+                        r.control_flow_context is not None
+                        for r in self.graph.hidden_param_op_relations
+                    )
+                    or any(
+                        g.descent is not None
+                        for g in self.graph.hidden_groups
+                    )
+                )
+                and not getattr(self, '_supports_scan_descent', False)
+            ):
+                raise NotImplementedError(
+                    f'{type(self).__name__} does not support structured scan '
+                    f'descent yet (an ETP relation or hidden group was '
+                    f'discovered inside a scan body). Use an algorithm that '
+                    f'supports scan descent, or set '
+                    f"ControlFlowPolicy(scan_descent='off')."
+                )
+
             # --- the initialization of the states --- #
             self.init_etrace_state(*args)
 
@@ -318,7 +355,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         return self.graph_executor.show_graph(verbose=verbose, return_msg=return_msg)
 
-    def __call__(self, *args) -> Any:
+    def __call__(self, *args: Any) -> Any:
         """
         Update the model and the eligibility trace states.
 
@@ -334,7 +371,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         return self.update(*args)
 
-    def update(self, *args) -> Any:
+    def update(self, *args: Any) -> Any:
         """
         Update the model and the eligibility trace states.
 
@@ -355,7 +392,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         raise NotImplementedError
 
-    def init_etrace_state(self, *args, **kwargs) -> None:
+    def init_etrace_state(self, *args: Any, **kwargs: Any) -> None:
         """
         Initialize the eligibility trace states of the etrace algorithm.
 

{braintrace-0.2.2 → braintrace-0.2.4}/braintrace/_algorithm/d_rtrl.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
 from .param_dim_vjp import ParamDimVjpAlgorithm
 
 __all__ = [

{braintrace-0.2.2 → braintrace-0.2.4}/braintrace/_algorithm/d_rtrl_test.py

@@ -25,7 +25,6 @@ from braintrace._algorithm.d_rtrl import D_RTRL
 from braintrace._algorithm.param_dim_vjp import (
     ParamDimVjpAlgorithm,
     _remove_units,
-    _fast_solve_contract,
 )
 from braintrace._etrace_model_test import (
     IF_Delta_Dense_Layer,
@@ -39,7 +38,7 @@ from braintrace._etrace_model_test import (
     ALIF_STDExpCu_Dense_Layer,
     ALIF_STPExpCu_Dense_Layer,
 )
-from braintrace._op import etp_mm_p
+from braintrace._op import etp_mm_p, get_fast_path_rules
 
 
 # ---------------------------------------------------------------------------
@@ -657,7 +656,7 @@ class TestDRtrlDictTraceStorage:
 # Helpers for the trace-update specialization + dtype-knob tests
 # ---------------------------------------------------------------------------
 
-def _rnn_mm(seed=0, n=4, bias=False):
+def _rnn_mm(seed=0, n=4, bias=False, weight_fn=None, bias_fn=None):
     """Batched single-recurrent-weight tanh RNN -> etp_mm_p, num_state==1."""
 
     class Net(brainstate.nn.Module):
@@ -670,7 +669,10 @@ def _rnn_mm(seed=0, n=4, bias=False):
 
         def update(self, x):
             b = self.b.value if self.b is not None else None
-            self.h.value = jax.nn.tanh(braintrace.matmul(self.h.value + x, self.w.value, b))
+            self.h.value = jax.nn.tanh(braintrace.matmul(
+                self.h.value + x, self.w.value, b,
+                weight_fn=weight_fn, bias_fn=bias_fn,
+            ))
             return self.h.value
 
     net = Net()
@@ -678,7 +680,7 @@ def _rnn_mm(seed=0, n=4, bias=False):
     return net
 
 
-def _rnn_mv(seed=0, n=4):
+def _rnn_mv(seed=0, n=4, weight_fn=None):
     """Unbatched single-recurrent-weight tanh RNN -> etp_mv_p, num_state==1."""
 
     class Net(brainstate.nn.Module):
@@ -690,7 +692,9 @@ def _rnn_mv(seed=0, n=4):
             self.h = brainstate.HiddenState(jnp.zeros((n,)))
 
         def update(self, x):
-            self.h.value = jax.nn.tanh(braintrace.matmul(self.h.value + x, self.w.value))
+            self.h.value = jax.nn.tanh(braintrace.matmul(
+                self.h.value + x, self.w.value, weight_fn=weight_fn,
+            ))
             return self.h.value
 
     net = Net()
@@ -698,6 +702,27 @@ def _rnn_mv(seed=0, n=4):
     return net
 
 
+def _rnn_elemwise(seed=0, n=4, weight_fn=None):
+    """Single element-wise-weight tanh RNN -> etp_elemwise_p, num_state==1."""
+
+    class Net(brainstate.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.alpha = brainstate.ParamState(
+                0.5 * jax.random.normal(jax.random.PRNGKey(seed), (n,))
+            )
+            self.h = brainstate.HiddenState(jnp.zeros((1, n)))
+
+        def update(self, x):
+            a = braintrace.element_wise(self.alpha.value, weight_fn=weight_fn)
+            self.h.value = jax.nn.tanh(a * (self.h.value + x))
+            return self.h.value
+
+    net = Net()
+    brainstate.nn.init_all_states(net, batch_size=1)
+    return net
+
+
 def _run_updates(algo, xs):
     """Compile, init, run len(xs) eager updates; return etrace_bwg snapshot.
 
@@ -773,6 +798,73 @@ class TestS1RecurrentBroadcastEqualsLegacy:
         _assert_etrace_lists_equal(fast, legacy, exact=True)
 
 
+# ---------------------------------------------------------------------------
+# weight_fn/bias_fn — the fast closed-form trace must equal the rule (legacy)
+# path when a transform hook is present.
+#
+# The fast kernels compute the instant term as the outer product x ⊗ df, which
+# equals ∂h/∂V (the *transformed* weight). With a transform present the true
+# eligibility trace is ∂h/∂W_raw = (x ⊗ df) ⊙ f'(W); the f' factor is supplied
+# only by the rule path's ``xy_to_dw`` (via jax.vjp). So the fast path must fall
+# back to the rules whenever a ``*_fn`` hook is set, otherwise the trace is the
+# gradient w.r.t. the transformed weight and silently drops f'(W).
+# ---------------------------------------------------------------------------
+
+class TestTransformFastEqualsLegacy:
+    """With a weight_fn/bias_fn hook, the default (fast) trace must match the
+    rule path bit-for-bit — i.e. the fast path falls back to the rules so the
+    f'(W) factor from ``xy_to_dw`` is not dropped.
+
+    Run several steps so the recurrent term (non-zero from the 2nd update) is
+    exercised, not just the instant term."""
+
+    def _xs_mm(self, n=4, steps=4):
+        return brainstate.random.randn(steps, 1, n)
+
+    def _xs_mv(self, n=4, steps=4):
+        return brainstate.random.randn(steps, n)
+
+    def test_mm_weight_fn_fast_equals_legacy(self):
+        xs = self._xs_mm()
+        wfn = lambda w: w ** 2
+        fast = _run_updates(ParamDimVjpAlgorithm(_rnn_mm(weight_fn=wfn), fast_solve=True), xs)
+        legacy = _run_updates(ParamDimVjpAlgorithm(_rnn_mm(weight_fn=wfn), fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+    def test_mm_weight_fn_and_bias_fn_fast_equals_legacy(self):
+        xs = self._xs_mm()
+        fast = _run_updates(
+            ParamDimVjpAlgorithm(
+                _rnn_mm(bias=True, weight_fn=jnp.tanh, bias_fn=lambda b: b * 2.0),
+                fast_solve=True), xs)
+        legacy = _run_updates(
+            ParamDimVjpAlgorithm(
+                _rnn_mm(bias=True, weight_fn=jnp.tanh, bias_fn=lambda b: b * 2.0),
+                fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+    def test_mv_weight_fn_fast_equals_legacy(self):
+        xs = self._xs_mv()
+        fast = _run_updates(ParamDimVjpAlgorithm(_rnn_mv(weight_fn=jnp.tanh), fast_solve=True), xs)
+        legacy = _run_updates(ParamDimVjpAlgorithm(_rnn_mv(weight_fn=jnp.tanh), fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+    def test_elemwise_weight_fn_fast_equals_legacy(self):
+        xs = self._xs_mm()
+        wfn = jnp.tanh
+        fast = _run_updates(ParamDimVjpAlgorithm(_rnn_elemwise(weight_fn=wfn), fast_solve=True), xs)
+        legacy = _run_updates(ParamDimVjpAlgorithm(_rnn_elemwise(weight_fn=wfn), fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+    def test_no_transform_still_uses_fast_path(self):
+        """Guard: without a transform, fast and legacy remain bit-identical (the
+        gate must not disable the fast path when no ``*_fn`` hook is present)."""
+        xs = self._xs_mm()
+        fast = _run_updates(ParamDimVjpAlgorithm(_rnn_mm(), fast_solve=True), xs)
+        legacy = _run_updates(ParamDimVjpAlgorithm(_rnn_mm(), fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+
 # ---------------------------------------------------------------------------
 # Change A — trace_dtype knob (store the eligibility trace in lower precision)
 # ---------------------------------------------------------------------------
@@ -838,7 +930,8 @@ class TestTraceDtypeKnob:
 
 def _alif(n_in=4, n_rec=5):
     """Adaptive-LIF dense layer -> coupled (V, adaptation) hidden group with
-    num_state==2, exercising the S>1 einsum branch of _fast_recurrent_term."""
+    num_state==2, exercising the S>1 einsum branch of the dense fast-path
+    recurrent rule (``FastPathRules.recurrent``)."""
     model = ALIF_Delta_Dense_Layer(n_in, n_rec)
     brainstate.nn.init_all_states(model)
     return model
@@ -891,11 +984,12 @@ class TestSolveBatchFold:
             'bias': brainstate.random.randn(B, O, S),
         }
         # reference: current per-batch contraction, then explicit batch sum
-        ref = _fast_solve_contract(etp_mm_p, diag_like, etrace)
+        solve = get_fast_path_rules(etp_mm_p).solve
+        ref = solve(diag_like, etrace, fold_batch=False)
         ref_w = ref['weight'].sum(axis=0)   # (I, O)
         ref_b = ref['bias'].sum(axis=0)     # (O,)
         # new: fold the batch axis inside the einsum
-        folded = _fast_solve_contract(etp_mm_p, diag_like, etrace, fold_batch=True)
+        folded = solve(diag_like, etrace, fold_batch=True)
         assert folded['weight'].shape == (I, O)
         assert folded['bias'].shape == (O,)
         npt.assert_allclose(folded['weight'], ref_w, atol=1e-6)