P1: axis-aware verification harness, in-tree limitation list, conv-bias IO-dim fix

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
@@ -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)

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',

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():

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',
 ]

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],

braintrace/_algorithm/_common_test.py

@@ -25,32 +25,16 @@
 from braintrace._algorithm._common import (
     FixedRandomFeedback,
     KappaFilter,
-    PresynapticTrace,
     _batched_zeros_like,
     _extract_leaf,
     _reset_state_in_a_dict,
     _sum_dim,
     _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
 # ---------------------------------------------------------------------------

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)