P1: axis-aware verification harness, in-tree limitation list, conv-bias IO-dim fix#141
Merged
Merged
Conversation
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
…w 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.
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Sorry @chaoming0625, your pull request is larger than the review limit of 150000 diff characters
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
P1 was scoped in the roadmap as compiler work on multi-timescale and heterogeneous populations. Measurement showed the compiler already passes every one of those targets —
num_stateup to 5, per-neuron heterogeneoustau_mem/tau_a, six distinct synaptic timescales, E/I population splits all reproduce BPTT. What failed was the instrument. With approval, P1 was refocused onto the verification harness and the limitation list.hidden_group.pywas not modified; no defect was found in its Jacobian path. P3 therefore inherits today's representation, and roadmap risk 2 ("P1 and P3 collide inhidden_group.py") is retired rather than mitigated.The load-bearing finding (F-23)
A full-sequence
MultiStepDataVJP has no truncation, so the eligibility trace never enters and every algorithm returns bitwise BPTT at every hyperparameter.D_RTRL,OSTLRecurrent,EPropat anykappa_filter_decayandpp_propat anydecay_or_rankare indistinguishable on that path.This is correct semantics, and
chunked_online_param_gradientsdocumented it all along — but the assertions written againstonline_param_gradientsdid not heed it, which is how two findings reached the wrong cause:Building exactly that model (ALIF + E/I conductance split, 3 relations,
num_state5) reproduced the same bitwise exactness. No zoo could ever have helped; the fix was one parameter,chunk_size.Changes
assert_model_is_live(keyed on gradient norm) andassert_gradients_differ, plus pytree- and unit-aware gradient comparison so nestedLinearweights andmS-carrying gradients compare correctly.factory()return a different network per call) and F-25 at the spec boundary.axis_discrimination_test.py— pins F-23 in both directions, so no future test can assert an approximation's behaviour on a path that cannot observe it.snn_model_correctness_test.py— discharges theAGENTS.mdprose claims (heterogeneous leaks, multi-state HiddenGroups) as passing tests rather than fixes.xy_to_dwreturns the bias Jacobian per output position 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 and handed the un-reduced array tocustom_vjp. Fixed in the solver with a primitive-agnostic broadcast-gradient reduction, not by whitelisting conv.docs/specs/2026-07-25-known-limitations.md— the in-tree findings list, replacing the untrackeddev/oneAGENTS.mdpointed at. Every entry names the test that pins it.New finding: F-29
pp_prop(decay_or_rank=1)is not an approximation on a recurrent-only relation. An int rank maps throughdecay = (rank - 1) / (rank + 1), so rank 1 means decay 0 and the presynaptic EMA collapses to the exact current input. It reproduces exactD_RTRLto round-off (1e-10 … 3e-8) across chunk sizes 1/2/4,T∈ {8,12,16},n_rec∈ {4,8} and recurrent spectral radius 0.25 … 4.97 — so not a weak-recurrence artifact. It becomes a real approximation (0.31 … 0.79) once the presynaptic input carries an external component, as every SNN spec's does.Both sides of that boundary are pinned by a test. The mechanism is not claimed — only the boundary is measured.
Two mistakes made during this work, recorded
assert_model_is_liveexists to catch. Lesson: run the liveness helper inside throwaway probes too.1e-9. float32 round-off on these trees is ~1e-8, so it would have passed on noise for any configuration; it failed only because one config was quieter than noise. Axis assertions now use1e-6against deviations of1e-3…1e-1.Also worth flagging for reviewers: liveness is bounded above as well as below. Driven hard,
ALIF_Deltakeeps spiking at rate 0.60 while its surrogate derivative saturates and the BPTT gradient returns to exactly zero. Spiking is not a proxy for gradient liveness.Test plan
Against the P0 bar of 2062 passed / 1 skipped: +57 tests, and the skip retired rather than deferred. The suite now carries no
skiporxfailmarker anywhere, so "resolved" never means "stopped being run".One
pyproject.tomlchange was needed and was caused by this work: mypy'sexcludedoes not apply to files reached by an import, sooracle_models.py's import of the_etrace_model_test.pylayer classes re-admitted that excluded file and surfaced 52 pre-existingbrainpy.state.*errors. A scopedignore_errorsoverride for that one module restores the policy the exclude pattern already expresses.What P1 did not resolve
Stated so the findings list is not read as exhaustive: F-29's mechanism, F-28's square hidden×hidden feedback restriction in
EProp(feedback='random'), and F-24 in the layer classes themselves (the oracle specs re-seed; the constructors still draw from the global RNG).🤖 Generated with Claude Code