P2: decompose learning rules into explicit axis coordinates

[chaoming0625]

Second phase of docs/specs/2026-07-25-algorithm-axes-roadmap.md. Spec: docs/specs/2026-07-25-p2-axis-decomposition.md.

What this changes

Named algorithms become coordinates in a six-axis space rather than separate implementations. The new braintrace.ETraceConfig names that space:

Axis	Values
trace_factorization	per_param (O(P·H)), io_factorized (O(I+O))
temporal_recursion	jacobian, scalar_leak, none — an (x, f) pair under io_factorized
recurrence_scope	diagonal, coupled
learning_signal	symmetric, random_feedback
trace_filter	none, kappa
update_schedule	per_step

Illegal combinations are rejected at construction by an eight-rule compatibility matrix with an error naming the legal pairings; coordinates that denote the same rule canonicalise to one form. braintrace.compile accepts a config wherever it accepts an algorithm name, so a rule with no preset name is as constructible as one with a name.

D_RTRL, pp_prop, EProp, OSTLRecurrent and OSTLFeedforward are now thin factories over coordinates — EProp loses all four of its hook overrides.

Mechanism

temporal_recursion is realised by substituting the executor's per-hidden-group hidden→hidden Jacobian (λ·I for scalar_leak, zeros for none) rather than by special-casing operators, so it holds for dense, conv, sparse, lora and element-wise relations alike. The substitution is applied exactly once on each of the three trace paths (chunked, fused multi-step, single-step), which is asserted directly.

Random feedback and the κ-filter lift off EProp onto the base engine, so both compose with the O(I+O) trace and not just O(P·H).

Behaviour changes

• IODimVjpAlgorithm's x- and f-side decays are now independent. decay became a read-only property that raises AttributeError (naming decay_x / decay_f) when the sides differ. decay_or_rank=0.9 remains element-wise identical to the new (0.9, 0.9).
• decay_or_rank=0.0 is now admitted (bound relaxed from 0 < decay < 1 to 0 <= decay < 1); it canonicalises to temporal_recursion='none'.
• The private _include_recurrent_mixing became the public recurrence_scope, and asking for coupled on a model whose compilation descends into a scan now raises instead of silently degrading to diagonal.
• A random-feedback configuration that cannot be honoured raises instead of silently computing the symmetric rule.

Not fixed here, deliberately

F-30 — the IO-dim f-side bias correction is indexed by update() call count, not trace-step count. Found during the design review, reproduced at 6.8e-04 on the finite-window path (it is invisible on the full-window path, per F-23). Left unfixed because P2's entire acceptance criterion is that gradients do not move; a numerical correction and a refactor cannot land together or neither can be verified. Recorded in docs/specs/2026-07-25-known-limitations.md with a reproduction and pinned by a test asserting the current, biased behaviour.

Tests

• axes_test.py (64) — vocabulary, canonicalisation, each matrix rule paired with its nearest legal neighbour, coefficient bounds, derived views.
• tests/axis_golden_test.py (24 frozen gradients, 8 cases × 2 models) — covers all three trace paths. Values were captured on a finite window (chunked_online_param_gradients, chunk_size < T), because a full-sequence multi-step VJP is BPTT and would be identical across every coordinate (F-23).
• tests/axis_acceptance_test.py (29) — includes a two-sided positive control on leaky_linear, whose hidden→hidden Jacobian is 0.9·I exactly, so substituting the model's own leak must reproduce it bitwise while a mismatched leak moves the gradient.

Verified: pytest braintrace/ → 2240 passed, 4 deselected (diagnostic); mypy braintrace clean.

[sourcery-ai]

Sorry @chaoming0625, your pull request is larger than the review limit of 150000 diff characters

[codecov]

Codecov Report

❌ Patch coverage is 98.31081% with 5 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
braintrace/_algorithm/param_dim_vjp.py	90.47%	2 Missing ⚠️
braintrace/_algorithm/axes.py	99.39%	1 Missing ⚠️
braintrace/_algorithm/io_dim_vjp.py	96.29%	1 Missing ⚠️
braintrace/_algorithm/vjp_base.py	98.46%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!