This release makes eligibility-trace online learning work *through* JAX control
flow. A new compiler canonicalization + descent pipeline lets ETP operations
inside `vmap`, `cond`, `scan` / `for_loop`, and weight-free `while` bodies
participate in online learning (Phases 0–4), so recurrent cells built with
control flow no longer silently drop parameters from the trace graph. The
operator layer gains three new ETP ops — `grouped_matmul`, `embedding`, and
`einsum` — each with a matching `braintrace.nn` layer; the D-RTRL multi-step
trace update is chunk-factorized for a 2.4–4.5× speedup on multi-step windows;
and a full `_op` / `_algorithm` audit closes 24 correctness findings. The
compiler itself is now deterministic across processes and transparently inlines
user `jax.jit` bodies. One internal ETP rule is renamed (see Breaking changes).
Highlights
New: `grouped_matmul`, `embedding`, and `einsum` ETP operators
- **Three new ETP operators join the operator layer**, each with hand-written
ETP rules (`dt_to_t`, `xy_to_dw`, trace initializers), a closed-form D-RTRL
fast path where applicable, public exports, and single-step BPTT-oracle
coverage:
- **`braintrace.grouped_matmul`** — a grouped matmul exposing both batched
and unbatched primitives (`etp_gmm` / `etp_gmv`) and a closed-form D-RTRL
fast path; D-RTRL matches BPTT element-wise and `pp_prop` is directionally
aligned. Wrapped by the new **`braintrace.nn.GroupedLinear`** layer.
- **`braintrace.embedding`** — an ETP embedding lookup with a broadcast
`dt_to_t` and a scatter-add `xy_to_dw`. Because the input is integer token
indices, the IO-dim (`pp_prop` / `ES_D_RTRL`) input trace cannot low-pass
the raw indices; a new optional per-primitive `ETP_RULES_PP_X_REPR` registry
lets `embedding` filter the linear one-hot representation
(`y = onehot(idx) T`) instead, and `xy_to_dw` dispatches on the x dtype
(integer indices → gather-VJP scatter-add; float one-hot → contraction VJP).
Wrapped by the new **`braintrace.nn.Embedding`** layer.
- **`braintrace.einsum`** — an equation-parsed ETP einsum with axis
classification; diagonal-class and shared-axis equations are D-RTRL
BPTT-exact (maxdiff 0.0), while the genuinely lossy regime (output positions
collapsing into a smaller hidden state) fails loudly at compile time with a
cotangent-shape error rather than silently emitting wrong gradients.
New: structured scan descent — long ETP scans compile and learn online (Phase 4)
- **A third compile path for ETP-relevant `scan`s above the unroll limit.**
Previously an ETP-relevant `lax.scan`/`for_loop` whose static length
exceeded `ControlFlowPolicy.scan_unroll_limit` was a dead end
(`NotImplementedError`). Under the new default
`ControlFlowPolicy(scan_descent='auto')`, the compiler *descends* such a
scan: relations and hidden groups are discovered inside the scan body
with the same flat finders, the equation is rewritten to emit stacked
per-substep values as extra ys (leading substep axis `L`), and the graph
executor computes stacked per-substep Jacobians by vmapping over that
axis — the compiled program stays a single scan equation, so compile size
is independent of the loop length (an `L=100` inner loop compiles in
under 60 equations). A `SCAN_DESCENT_APPLIED` INFO diagnostic records
each descent; blocked scans get `SCAN_DESCENT_SKIPPED`. Set
`ControlFlowPolicy(scan_descent='off')` to restore the pre-Phase-4 error.
- **Param-dim algorithms fold the eligibility trace over the substep
axis.** `D_RTRL` (the `ParamDimVjpAlgorithm` family) applies its trace
update per substep with an inner `jax.lax.scan`
(`eps <- D_tau * eps + x_tau (x) df_tau`), declaring
`_supports_scan_descent = True`. The fold is values-only and
stop-gradient'd — never differentiated, so no checkpointing is needed.
The learning signal stays one-per-outer-step (`(*varshape, num_state)`;
the SNN learning-signal axis contract is unchanged). The io-dim family
(`pp_prop` / `ES_D_RTRL`) and other algorithms without the flag reject
descended graphs with an actionable `NotImplementedError` at
`compile_graph`.
- **Exactness contract (pinned by oracle tests).** For diagonal-recurrence
bodies (elementwise hidden-to-hidden substep path — the SNN class),
descended D-RTRL is exact: whole-sequence, chunked (3-step and 1-step
chunks, where the gradient depends on the folded trace at every chunk
boundary), and one-step single-step gradients all match BPTT / the
unrolled twin element-wise, including a two-hidden-state
(`num_state == 2`) group through the fold. For bodies that mix the hidden
state through an ETP matmul, whole-sequence multi-step gradients remain
BPTT-exact; chunked gradients approximate cross-substep credit (the same
approximation class as the unroll path — documented divergence).
- **Algorithm-level `control_flow` kwarg.** `D_RTRL`, `pp_prop`
(`IODimVjpAlgorithm`), `OTPE`, and `OTTT` now accept
`control_flow=ControlFlowPolicy(...)` and thread it through their graph
executors into compilation.
- **v1 restrictions** (each blocks descent for that scan, with a
diagnostic): reverse scans, nested control flow inside the body,
trainable weights scanned over as xs, and an *outer* ETP relation
targeting a hidden state carried by a descended scan (raises with
restructuring guidance). `jit` bodies nested inside control-flow
equations are now inlined during extraction so descent sees a flat body.
- **Single-step readout limitation.** The per-step hidden perturbation is
added to a descended scan's *carry* outvar; a loss that reads the hidden
state through the scan's stacked ys (e.g. `for_loop(...)[-1]`) bypasses
it, dropping the same-step learning signal (pinned by test; parallels
the Phase 3 while-hidden limitation). Read the state after the loop
(`self.h.value`) instead — multi-step VJP is unaffected either way.
New: `while`-loop policy — weight-free opaque-forward support (Phase 3)
- **Weight-free `lax.while_loop`s that read/update hidden state now
compile.** Under the new default policy knob
`ControlFlowPolicy(while_hidden='opaque-fwd')`, a `while` whose inputs
carry no trainable ETP weight is kept as an *opaque forward node*: the
compiler registers relations whose `y`→hidden tail crosses the loop,
emits a `CONTROL_FLOW_OPAQUE_FWD` INFO diagnostic, and extracts
hidden-to-hidden Jacobians for any hidden group whose transition contains
a `while` in **forward mode** (`jax.jvp`-based `jacfwd_last_dim` /
`jacfwd` block extraction) — reverse mode through `while_loop` is
structurally unsupported by JAX. Set
`ControlFlowPolicy(while_hidden='error')` to reject such loops instead.
- **Perturbation detach keeps the VJP reverse-traceable.** The hidden
perturbation pass rewires every hidden-producing `while` to consume
`stop_gradient` copies of its inputs in the *perturbed jaxpr only*; the
`h = fresh + ε` add stays outside the detach, so the single-step learning
signal of the loop's **own** hidden group (taken exclusively from the
perturbation cotangents) is exact. Verified: D-RTRL single-step gradients
on a while-settle model match its hand-composed no-`while` twin
element-wise, and the twin matches the BPTT oracle.
**Limitation:** the detach zeroes every same-step reverse path *through*
the loop, so a parameter or other hidden group whose only same-step path
to the loss crosses the loop — e.g. the weights of an upstream layer
feeding a while-hidden layer — receives a **zero** learning signal (a
WARNING-level `CONTROL_FLOW_OPAQUE_FWD` diagnostic records each detach;
the zero-upstream-gradient behavior is pinned by test).
`vjp_method='multi-step'` on a `while`-hidden model still raises JAX's
reverse-through-`while_loop` `ValueError` (documented limitation — use
the default single-step path).
- **A weight used inside a `while` is now a hard, actionable error**
(`WEIGHT_IN_WHILE` ERROR diagnostic + `NotImplementedError`): move the
weight application outside the loop so the loop consumes only its result
(subject to the same-step limitation above), or use a fixed-length
scan/`for_loop` (which the compiler unrolls).
- **Breaking: ETP primitives left inside an un-flattened `scan`/`while`/
`cond` body now raise** instead of being silently warned-and-excluded
(`etp_in_control_flow='error'`, the new default). Pass
`ControlFlowPolicy(etp_in_control_flow='exclude')` to restore the old
warn-and-exclude behavior.
- **Position-mixing guard**: a `while`/opaque control-flow body that applies
`dot_general`/`conv_general_dilated` to hidden-derived values (recurrent
weight mixing inside the loop) cannot be expressed as a per-position
Jacobian; the compiler treats it as a boundary, emits a
`CONTROL_FLOW_RECURRENT_MIXING` WARNING, and falls back to the
zero-recurrence (e-prop-style) group.
New: inner-`scan` unrolling (compiler canonicalization, Phase 2)
- **ETP operations inside `lax.scan` / `brainstate.transform.for_loop`
bodies now participate in online learning.** A new canonicalization pass
(`unroll_inner_scans` in `_compiler/canonicalize.py`) runs at extraction
time and replaces every ETP-relevant, statically short scan with its
unrolled body: one cloned copy per iteration with fresh variables, `xs`
sliced per step, consumed `ys` re-stacked via `broadcast_in_dim` +
`concatenate`, and `reverse=True` respected. The unrolled program is
value- and Jacobian-identical to the scan, so exact algorithms (D-RTRL,
full-rank pp_prop, EProp(k=0), OSTLRecurrent) match BPTT element-wise on
scan-body models — verified against hand-flattened twins and the BPTT
oracle. Cond and scan canonicalization now run as a joint fixpoint
(`canonicalize_control_flow`), so a `cond` inside a scan body (and an
eligible scan inside a `cond` branch) both flatten.
- **Relation counts follow the weight→weight→hidden invariant**: in an
unrolled inner loop only the *last* sub-step's ETP ops become relations —
earlier sub-steps reach the hidden state through another trainable ETP op
and are excluded (with the usual no-relation warning).
- **Eligibility gates**: only scans whose static `length` is ≤
`ControlFlowPolicy.scan_unroll_limit` (default 16) and that carry no
effects and contain no `while` are unrolled. An ETP-relevant scan that
fails a gate emits a `SCAN_UNROLL_SKIPPED` warning and keeps today's
hard-error behavior; unrolls are recorded as `SCAN_UNROLLED` INFO
diagnostics on `ETraceGraph.diagnostics`. Scans that scan *over* a
trainable weight (weights as `xs`) are never unrolled
(`RELATION_EXCLUDED_SLICED_WEIGHT` warning) — per-slice trace lineage is
deferred.
- **Cond gate revision**: a branch containing a scan no longer blocks
if-conversion when that scan is itself unrollable; `scan_unroll_limit=0`
disables unrolling and restores the exact Phase 1 gating.
- **Tied-weight invariant locked**: one `ParamState` consumed by several
ETP call sites (which unrolling multiplies) is keyed per relation
instance with per-path gradient accumulation — verified BPTT-exact and
now covered by regression tests.
New: `cond` if-conversion (compiler canonicalization, Phase 1)
- **ETP operations inside `lax.cond` branches now participate in online
learning.** A new canonicalization pass (`_compiler/canonicalize.py`)
runs at extraction time (after user-`jit` inlining) and rewrites every
ETP-relevant `cond` equation into the inlined bodies of *all* branches
followed by one `select_n` per output. `select_n`'s index semantics and
JVP match `cond` exactly, so for finite branches values and Jacobians —
and therefore exact algorithms such as D-RTRL — are unchanged. Weights used inside `cond`
branches previously raised `NotImplementedError` (or were silently
excluded when only ETP primitives appeared inside).
- **Semantics note**: on the canonicalized graph **both branches execute
every step** and the dead branch's value is discarded by `select_n`.
Values and forward-mode derivatives are unaffected by dead-branch
NaN/Inf. **Reverse-mode gradients are not**: if the dead branch's local
Jacobian is NaN/Inf (e.g. a `cond` protecting a `sqrt` domain), its VJP
multiplies the exact-zero cotangent by that Jacobian (`0 * nan = nan`)
and contaminates gradients of shared inputs — the classic single-`where`
pitfall. Keep such domain-guard conds opaque
(`ControlFlowPolicy(cond='opaque')`) or guard the operand inside the
branch.
- **Gates**: conds that touch no ETP primitive, weight, or hidden state
stay opaque at zero cost. Conds with effects or containing `while`/`scan`
in a branch are never converted; an ETP-relevant one that is skipped this
way emits a `COND_CONVERSION_SKIPPED` warning and keeps today's behavior.
Conversions are recorded as `COND_IF_CONVERTED` INFO diagnostics on
`ETraceGraph.diagnostics`.
- **Opt-out**: `braintrace.ControlFlowPolicy(cond='opaque')` via the new
`control_flow` keyword on `compile_etrace_graph` / `extract_module_info`
restores the previous behavior.
New: vmap identity preservation (operator layer)
- **vmap identity preservation (operator layer)**: `jax.vmap` over an
unbatched ETP op (`matmul`, `lora_matmul`, `sparse_matmul` with vector
input) now re-binds the batched ETP primitive (`etp_mm` / `etp_lora_mm` /
`etp_sp_mm`) instead of decomposing into standard JAX ops. Models that
vmap per-sample ETP operations inside `update()` now compile with full
eligibility-trace relations. When promotion is impossible (batched
weights, `etp_conv`, nested vmap), the op decomposes as before but emits
a `UserWarning` instead of silently dropping the parameter from online
learning. Note: when this warning appears from a `compile(..., vmap=True)`
learner's execution trace (e.g. conv models), it is expected and benign —
the eligibility-trace graph was already compiled per-sample before the
learner was vmapped, so no parameter is dropped.
New: user-`jit` inlining and a deterministic compiler
- **ETP operations inside a user `jax.jit` now compile.** `extract_module_info`
inlines user `jax.jit` bodies before any analysis, so `jit` boundaries are
transparent to hidden-group discovery and relation finding. Previously a
weight used inside a `jit` raised `NotImplementedError` and bare ETP
primitives were silently skipped (123).
- **Deterministic, reproducible compilation.** Hidden-group discovery,
transition bookkeeping, and group merging now use insertion-ordered maps and a
canonical compiled-state ordering instead of `set`s keyed by object identity,
so group membership and ordering are stable across processes. Every built
group is validated with `check_consistent_varshape`, and merges emit an
INFO-level `HIDDEN_GROUP_MERGED` diagnostic (123).
- **Directly-fed fan-out fix.** A single ETP op feeding two independent
recurrent states now registers relations to *both* groups — the forward BFS
previously locked onto whichever hidden state it reached first and dropped the
rest. Relation gating also resolves all keys before excluding a relation, so a
constant-weight / `ParamState`-bias matmul still registers with the bias as
its trainable key (123).
- **Robust perturbation pass.** The single-step perturbation now handles
multi-output equations and read-only hidden states (synthesizing the
`h^t = h^{t-1} + p` passthrough) and preserves the source jaxpr's effect set,
instead of falling through to an unexplained-hidden error (123).
Performance
- **Chunk-factorized multi-step D-RTRL trace update.** Multi-step trace updates
for `D_RTRL` now factor the per-step decay into suffix products and apply the
trace update per chunk instead of step-by-step, giving a **2.4–4.5× speedup**
on multi-step windows for the dense (`etp_mm` / `etp_mv`) and elementwise
(`etp_elemwise`) kernels. Exposed as a `chunked_trace` knob on `D_RTRL` /
`braintrace.compile` (132).
Correctness
- **ETP `_op` / `_algorithm` audit — 24 findings closed (4 Critical, 6 High,
6 Medium, 8 Minor).** Highlights: exact `conv` (C1) and `lora_matmul` (C2)
gradients under param-dim D-RTRL (per-position kernel trace + effective-weight
trace, backed by new optional instant / solve D-RTRL rule registries); a fix
for the batched sparse D-RTRL crash (C3) via a hashable CSR wrapper; and
OSTTP's always-zero learning signal (C4) via `custom_vjp` residual threading.
Also resolved: `trace_dtype` gate mismatch, conv bias broadcast, EProp
kappa-filter cross-state contamination and random-feedback scale invariance,
OTTT / OTPE dropped bias gradients and missing guards, int/bool autodiff
guards, rank guards with nn-layer axis folding, and a corrected OSTL exactness
claim. Adds a cross-family single-step BPTT oracle suite and first-principles
rule tests (`6c7796a`).
Breaking changes
- **ETP rule rename: `YW_TO_W` → `DT_TO_T`.** The recurrent trace-propagation
rule computes `D^t * ε^{t-1}` (the `Dᵗ`-times-previous-trace term of the
D-RTRL update), so `DT_TO_T` names it accurately; `YW_TO_W` never matched what
the rule computes. Custom primitives that register this rule (via
`register_etp_rules` / `register_primitive`) must use the new name. This is
unrelated to `brainevent`'s external `DataRepresentation.yw_to_w` /
`yw_to_w_transposed` protocol methods, which are untouched (130).
Internal
- **`mypy` CI gate repaired.** Cleared 40 accumulated type errors across 8 files
(annotation-only, no behavior change), restoring a green `typecheck_and_build`
job (133).
- Signature cleanups: a readability pass on `_etp_sp_matmul_impl` and removal of
unused keyword arguments across several modules.