Release Notes

Contents

Release Notes#

Version 0.2.4#

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 scans 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_loops 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 sets 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.

Version 0.2.3#

This release adds optional, shape-preserving parameter-transform hooks to the eligibility-trace (ETP) operators, so a trainable weight (or bias) can be passed through an elementwise / standardizing function before it enters the operation while the eligibility trace and gradient remain with respect to the raw stored parameter. These hooks are threaded through the braintrace.nn linear layers and demonstrated in a new tutorial. The release also hardens the public API with inline type annotations behind an enforced mypy gate, corrects the weight_fn / bias_fn gradients on the closed-form fast path, relocates the fast-path kernels into the operator layer, and tightens the sparse_matmul input contract. Two public APIs are renamed and one operand type is now required (see Breaking changes).

Highlights#

New: parameter-transform hooks on ETP operators#

  • Add transform hooks to the ETP ops, computing y = x @ weight_fn(w) (+ bias_fn(b)) (and per-op equivalents), with the eligibility trace and gradient kept with respect to the raw parameter:

    • braintrace.matmul / braintrace.sparse_matmul — weight_fn, bias_fn.

    • braintrace.conv — kernel_fn, bias_fn.

    • braintrace.lora_matmul — b_fn, a_fn, bias_fn.

    • braintrace.element_wise — weight_fn (see Breaking changes).

    Each transform is applied inside the ETP primitive; the per-parameter Jacobian is recovered exactly once (via jax.vjp) in the weight-gradient rule, while the trace-propagation rule is unchanged — so the forward-mode eligibility trace stays exact and is never double-counted. D-RTRL matches backprop-through-time element-wise for non-identity transforms (verified with tanh, w**2, and abs). Omitting a transform is bit-identical to the previous behavior.

New / Improved: braintrace.nn linear layers#

  • braintrace.nn.Linear (with w_mask), braintrace.nn.SignedWLinear, and braintrace.nn.ScaledWSLinear now route their weight masking / sign / standardization through the new matmul(weight_fn=...) hook, so the masked / signed / standardized weight participates in eligibility-trace learning with the gradient kept w.r.t. the raw weight leaf. (For ScaledWSLinear, gain and bias are applied as post-operations and are therefore non-temporal for the online trace, though still recovered exactly by the multi-step VJP oracle.)

  • Export braintrace.nn.ScaledWSLinear (previously importable only by its fully-qualified module path).

New: typed public API with an enforced mypy gate#

  • Inline type annotations now cover the public surface — ETP operators and their rule functions, ETPPrimitive / register_primitive, the braintrace.compile entry point and package accessors, input-data containers, the braintrace.nn linear / conv / recurrent cells, and the algorithm base classes, executors, and concrete algorithms. A new WeightFn alias names the transform-hook signature.

  • An enforced mypy gate guards the public API, so type regressions fail the build (#119).

Improvements#

  • Correct weight_fn / bias_fn gradients on the fast path. The transform Jacobian f'(W) is now applied on the param-dim D-RTRL closed-form fast path (it lives solely in xy_to_dw; dt_to_t stays transform-free), so transformed-parameter gradients match the slow path. Also fixes an element_wise slow-path batched-cotangent crash (#120).

  • Operator-layer fast-path kernels. The closed-form fast-path kernels (instant / recurrent / solve) move into the operator layer as a per-primitive FastPathRules bundle behind an ETP_FAST_PATH_RULES registry, and the algorithm-layer string-match gate is replaced by a per-primitive applicable() predicate — keeping primitive knowledge in the operator layer per the layered design (#120).

Documentation#

  • New tutorial: customizing primitive transforms (docs/tutorials/customizing_primitive_transforms.ipynb), plus transform-hook docstrings on the ETP operators (#120).

Breaking changes#

  • braintrace.element_wise: the fn parameter is renamed to weight_fn and is now keyword-only, and the transform is applied inside the ETP primitive (previously it was applied to the weight outside the primitive). Migrate element_wise(w, fn=g) to element_wise(w, weight_fn=g). Forward results are unchanged; only the call signature and the internal trace-factorization point differ.

  • braintrace.sparse_matmul: the weight parameter is renamed from weight_data to weight for a cleaner, more consistent API. All in-tree call sites pass it positionally and are unaffected; update any keyword callers (#116).

  • braintrace.sparse_matmul: the sparse operand (sparse_mat) must now be a brainevent.DataRepresentation and is enforced with a strict runtime isinstance check (raising TypeError). DataRepresentation supplies the ETP online-learning protocol the compiler / executor require (with_data, yw_to_w, yw_to_w_transposed); brainunit sparse types (u.sparse) lack these and are no longer accepted. brainevent is now a runtime dependency. Migrate sparse weights to brainevent (e.g. brainevent.CSR) (#121).

Dependencies#

  • Add brainevent as a runtime dependency (pyproject.toml, requirements.txt) (#121).

  • Bump codecov/codecov-action from 5 to 7 (#117).

Version 0.2.2#

This release introduces a unified braintrace.compile entry point for building eligibility-trace online learners, adds a recurrent mixing mode to the graph-construction compiler, and fixes eligibility-trace convergence under vmap / brainstate.mixin.Batching(). It also migrates unit handling from saiunit to brainunit, modernizes the toolchain (Python 3.14, brainstate >= 0.5.2, Codecov), and ships broad documentation, example, and test improvements. Internal modules were renamed for brevity; no documented 0.2.x public API is removed.

Highlights#

New: unified braintrace.compile entry point#

  • braintrace.compile(model, algorithm, example_input, ...) is now the canonical, single-call way to build a compiled online learner. It always initializes states, accepts a seed, applies model guardrails, and can emit a verbose compilation report — replacing the manual init_states / learner.compile_graph(x0) triad.

  • vmap= parameter for per-sample vmap state initialization. With vmap=True, states are built via vmap_new_states(state_tag='new', ...) and the learner is wrapped in brainstate.nn.Vmap(vmap_states='new'), so eligibility-trace models compose with brainstate’s per-sample vmap scheme.

  • CompilationReport, a structured view over the eligibility-trace graph (relation/weight counts, etrace_weights, excluded_weights, report.show() with verbosity levels). It is exposed via ETraceAlgorithm.report and now backs show_graph.

New: recurrent mixing mode for graph construction#

  • Add a recurrent mixing mode to eligibility-trace graph construction, broadening the set of cell topologies the compiler can connect (#108).

Improvements#

Dependencies and toolchain#

  • Replace saiunit with brainunit for all unit handling across source, tests, examples, and docs. brainunit re-exports saiunit internally, so this is a drop-in change (#106).

  • Raise the brainstate floor to >= 0.5.2, required by the compile(vmap=True) path, and drop a duplicate dependency declaration.

  • Update the supported Python version to 3.14 and adjust the CI JAX version matrix.

  • Add Codecov coverage reporting and raise source coverage to 93%, with new tests for previously-untested modules (#109).

Refactoring#

  • Rename internal module packages for brevity: _etrace_op → _op, _etrace_compiler → _compiler, and _etrace_algorithms → _algorithm. These are private modules; imports were updated package-wide with word-boundary-anchored replacement (#111).

  • Remove the unused ParamState from state management.

  • Remove the per-step spectral-normalization path (normalize_matrix_spectrum) from D-RTRL, E-Prop, and the OSTL trace scan; it ran jnp.linalg.eigvals on every hidden-group Jacobian, was off by default, and was costly.

Fixes#

  • Eligibility-trace convergence under vmap batching. Defer graph compilation during the vmap_new_states discovery probe so the executor binds to the real batched states (fixes a BatchAxisError when writing batched values), correctly handle models that mix batched and unbatched ETP primitives in the param-dim VJP solve, and align convolution eligibility traces under brainstate.nn.Vmap(vmap_states='new'). Restores convergence for the conv-based SNN/RNN training examples.

  • Element-wise eligibility traces under brainstate.mixin.Batching(). Size the trace from the (batch-aware) hidden group and sum out the leading batch axis in the solver, fixing a scan-carry type mismatch and a custom-VJP backward shape mismatch. This unblocks the default SHD batch trainer, where every LIF leak is an element-wise weight.

  • braintrace.nn.LoRA now routes its forward through the ETP lora_matmul primitive, so LoRA factors participate in eligibility-trace learning (fixes the zero-relations bug) and the factor order is corrected.

  • Resolve pre-existing mypy errors in the compiler’s report.py (#112) and treat brainunit / saiunit as untyped for mypy to clear spurious attr-defined errors from their re-export chain.

  • Convert legacy xfail tests to positive assertions, silence the core.Jaxpr DebugInfo deprecation warning, and migrate deprecated brainstate APIs (brainstate.augment → brainstate.transform, brainstate.functional → brainstate.nn) (#113).

Documentation and examples#

  • Make braintrace.compile the canonical entry point in every docstring, tutorial, notebook, and example, and fix broken examples (e.g. self-contained RNNs, consistent batch axes); each documented example is now backed by an executable test (#114).

  • Document CompilationReport in the API reference and migrate the onboarding guides, quickstart, and tutorials to the unified compile flow.

  • Add a smoke-test harness and a testable main() entry point to the standalone examples; repair all docs notebooks so they execute cleanly.

Notes#

  • The internal module renames (_etrace_* → _*), the removal of ParamState, and the removal of normalize_matrix_spectrum touch private/internal surfaces only; the documented 0.2.x public API is unchanged.

  • Verified locally: the full CPU test suite is green (1604 passed, 3 skipped).

Version 0.2.1#

This is a maintenance release that restores compatibility with the latest brain-ecosystem dependencies and toolchain. It contains no functional or public-API changes — code written against 0.2.0 continues to work unchanged — and exists to keep BrainTrace green against brainstate 0.5, saiunit/ brainunit 0.5.1, and pytest 9.1.

Fixes#

Dependency Compatibility#

  • brainstate 0.5 typed API: adopted brainstate’s PEP 561 py.typed surface throughout the source — routed PyTree through BrainTrace’s existing type alias, centralized an as_size_tuple() helper in _typing, dropped FlattedDict subscripts, and added boundary asserts/casts. This clears the 154 mypy errors newly exposed by the upstream typing, with minimal # type: ignore only where brainstate’s typing makes it unavoidable.

  • brainstate 0.5.0 convolution validation: updated convolution test expectations for the hardened validation (bare assert → ValueError) and the new one-value-per-spatial-dimension padding-tuple semantics.

  • pytest 9.1.0 collection: removed trailing commas in single-argument parametrize ids that pytest 9.1.0 mis-parses as two values, fixing a collection-time GraphNodeMeta has no len() error.

Notes#

  • All changes are BrainTrace-side. A related upstream saiunit issue is resolved in saiunit/brainunit 0.5.1 and requires no change here.

  • Verified locally: full suite 1367 passed (2 xfailed), mypy clean across 51 files, and wheel + sdist build with py.typed shipped (PEP 561).

Version 0.2.0#

This release is a major step for BrainTrace. It adds a family of spiking neural network (SNN) online-learning algorithms, rewrites the eligibility-trace compiler around primitive-type dispatch, generalizes every ETP primitive to support multiple trainable inputs (fixing a silent bias-gradient drop), delivers substantial performance gains for D-RTRL and multi-step rollouts, and hardens the package with PEP 561 typing and a BPTT-oracle-backed test suite.

Major Changes#

New: SNN Online-Learning Algorithms#

  • Added five SNN online-learning algorithms as flat ETraceVjpAlgorithm subclasses: EProp, OSTL (OSTLRecurrent / OSTLFeedforward), OTPE, OTTT, and OSTTP. All are exported at the top level.

  • Added a _compute_learning_signal hook to ETraceVjpAlgorithm to support target-projection algorithms (OSTTP) without disrupting the existing D-RTRL and pp-prop paths.

  • Added supporting trace helpers: PresynapticTrace, KappaFilter, FixedRandomFeedback, and target-signal extraction utilities.

  • Algorithms are cross-checked for regime equivalence and verified to decrease loss in integration smoke tests.

ETP Compiler Rewrite#

  • Rewrote the eligibility-trace compiler to dispatch on primitive-type identity rather than string-matching op or trace names, with structured, leveled diagnostics (DiagnosticKind, DiagnosticLevel, CompilationRecord) replacing ad-hoc warnings.

  • Added compile-time diagnostics that surface previously silent issues — e.g. TRAINABLE_INVAR_NOT_PARAMSTATE flags a trainable input (such as a constant bias) that does not trace to a ParamState, so users can wrap it intentionally instead of silently losing its gradient.

Multi-Trainable-Input ETP Primitives (Bias Gradients)#

  • Generalized every ETP primitive from a single-“weight” assumption to an arbitrary named dict of trainable inputs. This fixes a silent bias-gradient drop and a LoRA executor signature mismatch in one coherent refactor.

  • Migrated all built-in primitives (elemwise, dense mm/mv, conv, sparse mm/mv, and lora) to the dict-based rule API with first-class bias gradient support, each verified element-wise against a BPTT oracle.

  • Fixed layout-aware axis handling in conv primitives (1D/2D, NHWC/NCHW, OIHW/HWIO kernel layouts) that previously corrupted gradients on non-default layouts, and fixed non-square dense weight broadcasting in _mm_dt_to_t.

  • Eligibility traces are now stored as per-key dicts; the transitional legacy-array adapter has been fully removed.

Performance#

  • D-RTRL einsum fast path (fast_solve=True, default on): replaces nested vmap-of-vjp and per-step lax.cond overhead with direct einsum kernels for mm/mv/elemwise; conv/sparse/LoRA fall back to the legacy path.

  • Reduced-precision trace storage (trace_dtype, e.g. bf16/fp16) halves the dominant B*N^2 trace bandwidth on GPU/TPU while keeping Jacobians, learning signals, and final gradients in fp32. Default None preserves exact behavior.

  • Multi-step trace fusion: the per-step eligibility-trace roll for exact algorithms (D-RTRL, pp-prop) is now threaded into the graph executor’s forward scan, eliminating an O(T × Jacobian) HBM round-trip (traced scan count drops 3 → 2). Opt-in and multi-step-only; single-step/SNN paths are unchanged.

  • Branch-free spectrum/vector normalization to restore XLA fusion across steps.

Primitive Registration Simplification#

  • Removed ETPPrimitiveSpec and the spec-based registration layer; invar/ outvar layout metadata (trainable_invars_fn, x_invar_index, y_outvar_index) now lives in internal registries populated directly through register_primitive keyword arguments.

Package Restructuring#

  • Consolidated the eligibility-trace code into a single flat _etrace_algorithms package, merging the former _etrace_vjp/, _etrace_algorithms.py, _etrace_graph_executor.py, and _snn_algorithms/ modules. The top-level public API is unchanged.

  • Split the algorithm base hierarchy into dedicated modules: ParamDimVjpAlgorithm (D-RTRL) and IODimVjpAlgorithm (pp-prop) now live in their own files, with D_RTRL/pp_prop as thin subclasses.

  • Removed the experimental hybrid online-learning method.

Typing & Packaging#

  • The package is now PEP 561 compliant: ships a py.typed marker so downstream users receive inline type hints.

  • Added a pragmatic mypy configuration and wired type checking plus packaging verification (python -m build, py.typed presence) into CI.

Testing#

  • Added a BPTT gradient oracle and a layered correctness test suite (P2–P8): per-operator rule oracles, public-API contract tests, exact-class element-wise equivalence with BPTT, approximate-class direction-alignment checks, transform/integration invariance, and per-cell compiler relation guardrails tied to the cell registry.

Documentation#

  • Converted all public-API docstrings to NumPy-doc style with math, references, and runnable examples.

  • Documentation is now self-hosted at brainx.chaobrain.com/braintrace/, with refreshed RTD links and a WebP logo.

Dependencies & Tooling#

  • Replaced brainunit with saiunit throughout for unit handling.

  • Numerous CI/CD upgrades (checkout, setup-python, artifact actions, sphinx and theme requirements); docs deploy on release publication.

Deprecations#

The entire v0.1.x class-based operator/parameter API is deprecated in favor of the new primitive-based ETP user-API. The legacy classes still work — they are thin back-compatibility shims that route through the new primitives — but each emits a DeprecationWarning (once per class, per process) on first use, and they will be removed in a future release. Migrate at your convenience.

Deprecated operator classes → new primitive functions:

Deprecated (v0.1.x)

Use instead (v0.2.0)

MatMulOp

braintrace.matmul

ElemWiseOp

braintrace.element_wise

ConvOp

braintrace.conv

SpMatMulOp

braintrace.sparse_matmul

LoraOp

braintrace.lora_matmul

ETraceOp (base)

the ETP primitive functions above

Deprecated parameter classes → brainstate.ParamState + a primitive:

Deprecated (v0.1.x)

Use instead (v0.2.0)

ETraceParam

brainstate.ParamState + an ETP primitive function (e.g. braintrace.matmul)

ElemWiseParam

brainstate.ParamState + braintrace.element_wise

NonTempParam

brainstate.ParamState + plain JAX ops (x @ w) — keeps the weight out of the ETP graph

FakeETraceParam, FakeElemWiseParam

plain objects with plain JAX ops

The stop_param_gradients context manager and the general_y2w helper are kept as no-op compatibility shims and have no effect on the new primitive path.

Breaking Changes#

  1. OSTL factory removed — use OSTLRecurrent or OSTLFeedforward directly instead of the former OSTL factory function.

  2. OTTT and OTPE require an explicit leak — the membrane leak is no longer inferred from model.states() (it silently picked a wrong value on heterogeneous/multi-population models). Both now also reject hidden groups with num_state > 1 at compile time, as collapsing the num_state axis has no theoretical basis for these LIF-derived rules. OTPE additionally documents a narrower feed-forward / single-layer / global-scalar-leak regime.

  3. Unit dependency change — code relying on brainunit internals should migrate to saiunit.

  4. ETPPrimitiveSpec removed — custom primitives must register layout metadata via register_primitive keyword arguments (trainable_invars_fn, x_invar_index, y_outvar_index).

Migration Guide#

OSTL#

# Old
algo = OSTL(model, ...)         # factory

# New — choose the regime explicitly
algo = OSTLRecurrent(model, ...)
# or
algo = OSTLFeedforward(model, ...)

OTTT / OTPE#

# Old
algo = OTTT(model, ...)               # leak inferred from model.states()

# New — pass the postsynaptic membrane leak explicitly
algo = OTTT(model, leak=0.9, ...)

Custom ETP primitives#

# Old: register_primitive_spec(ETPPrimitiveSpec(...))
# New: pass layout metadata directly
register_primitive(
    prim,
    trainable_invars_fn=...,
    x_invar_index=...,
    y_outvar_index=...,
)

Deprecated class-based API → primitive-based API#

# Old (v0.1.x): wrap the weight in an ETraceParam bound to an op
self.w = braintrace.ETraceParam({'weight': w}, braintrace.MatMulOp())
y = self.w.execute(x)

# New (v0.2.0): a plain ParamState + the ETP primitive function
self.w = brainstate.ParamState({'weight': w})
y = braintrace.matmul(x, self.w.value)

The element-wise case is analogous (ElemWiseParam/ElemWiseOp → brainstate.ParamState + braintrace.element_wise); to keep a weight out of the eligibility-trace graph, use a plain brainstate.ParamState with ordinary JAX ops instead of NonTempParam / FakeETraceParam.

Version#

  • Bumped version from 0.1.3 to 0.2.0

Version 0.1.2#

Major Changes#

Import Path Migration#

  • Updated dependency from brainpy to brainpy.state: Migrated all imports to use the more specific brainpy.state module

    • Updated braintrace/nn/_readout.py: Changed neuron model imports from brainpy to brainpy.state

    • Updated all documentation notebooks (12 files): Concepts, RNN/SNN online learning, batching, state management, and graph visualization tutorials

    • Updated example scripts (4 files): COBA EI RSNN, SNN evaluation, feedforward conv SNN, and SNN models

    • Updated requirements.txt and pyproject.toml to specify brainpy-state as dependency

    • Total: 19 files changed with improved module structure and consistency

New Algorithms#

  • Added PP-Prop (Pseudo-Prospective Propagation) algorithm: New eligibility trace algorithm in VJP-based methods

    • Added pp_prop to braintrace/_etrace_vjp/esd_rtrl.py

    • Updated docs/apis/algorithms.rst to include PP-Prop in algorithm documentation

Python 3.14 Support#

  • Added Python 3.14 compatibility: Updated project metadata to officially support Python 3.14

    • Updated pyproject.toml classifiers to include Python 3.14

Bug Fixes#

  • Fixed version info tuple creation: Corrected the version info structure in braintrace/__init__.py

    • Ensures proper version tuple formatting for compatibility checks

CI/CD Improvements#

  • Updated GitHub Actions workflow: Bumped actions/upload-artifact from v5 to v6

    • Modernized CI/CD pipeline with latest GitHub Actions versions

    • Improved artifact upload reliability and performance

Documentation Updates#

  • Updated documentation links: Refreshed links in concept documentation for better navigation

    • Updated docs/quickstart/concepts-en.ipynb (116 lines modified)

    • Updated docs/quickstart/concepts-zh.ipynb (104 lines modified)

Breaking Changes#

Dependency Change:

  1. Dependency name change: The project now requires brainpy-state instead of brainpy

    • Update your requirements.txt or installation commands accordingly

# Old (0.1.1)
pip install brainpy

# New (0.1.2)
pip install brainpy-state

  1. Import path update: Update neuron model imports to use brainpy.state

# New (0.1.2)
from brainpy.state import IF, LIF, ALIF

Migration Guide#

Update Dependencies#

Replace brainpy with brainpy-state in your project dependencies:

pip uninstall brainpy
pip install brainpy-state

Update Import Statements#

If you have custom code importing neuron models, update to use brainpy.state:

# Find and replace in your codebase
# from brainpy import → from brainpy.state import

Version#

  • Bumped version from 0.1.1 to 0.1.2

Version 0.1.1#

Major Changes#

Project Rename: BrainScale → BrainTrace#

  • Renamed the entire project from brainscale to braintrace: This change reflects the project’s focus on eligibility trace-based learning algorithms

    • Package directory renamed from brainscale/ to braintrace/

    • All internal imports updated from brainscale to braintrace

    • Updated all 95 files including source code, tests, documentation, and examples

    • Updated pyproject.toml with new project name and metadata

    • Updated README with new project branding and citation information

VJP-Based Eligibility Trace Algorithms#

  • Added new VJP-based eligibility trace module (_etrace_vjp/): Comprehensive implementation of vector-Jacobian product based algorithms

    • base.py: Core base classes and utilities for VJP operations (671 lines)

    • d_rtrl.py: Diagonal Real-Time Recurrent Learning implementation (756 lines)

    • esd_rtrl.py: Efficient Sparse Diagonal RTRL implementation (847 lines)

    • hybrid.py: Hybrid approaches combining multiple techniques (604 lines)

    • graph_executor.py: Graph-based execution for VJP computations

    • misc.py: Miscellaneous utilities including matrix spectrum normalization

  • Refactored VJP algorithm structure: Migrated from monolithic _etrace_vjp_algorithms.py (2,888 lines) to modular architecture

    • Better separation of concerns

    • Improved testability with dedicated test files (d_rtrl_test.py, esd_rtrl_test.py, graph_executor_test.py)

Logo and Branding#

  • Updated logo format from JPG to PNG for consistency

  • Updated logo across documentation

Breaking Changes#

Package Rename:

  1. Import path change: All imports must now use braintrace instead of brainscale

# Old (0.1.0)
import brainscale
from brainscale import EligibilityTrace
from brainscale.nn import Linear, GRUCell

# New (0.1.1)
import braintrace
from braintrace import EligibilityTrace
from braintrace.nn import Linear, GRUCell

  1. Installation: Package name changed from brainscale to braintrace

# Old
pip install brainscale

# New
pip install braintrace

Migration Guide#

Update Import Statements#

Replace all occurrences of brainscale with braintrace:

# Find and replace in your codebase
# brainscale → braintrace

VJP Algorithm Usage#

The new VJP-based algorithms are now available through the modular interface:

Version#

  • Bumped version from 0.1.0 to 0.1.1

Version 0.1.0#

Major Changes#

State Management Refactoring#

  • Renamed ETraceState to HiddenState: All eligibility trace state management now uses the more general HiddenState naming convention

    • Updated across _etrace_algorithms.py, _etrace_concepts.py, _state_managment.py

    • Added deprecation warnings for ETraceState to guide users to brainstate.HiddenState

    • Updated all documentation and examples to reflect the new naming

  • Renamed ETraceGroupState to HiddenGroupState: Improved consistency in hidden state handling

    • Updated in _etrace_compiler_hidden_group.py

    • Added deprecation warnings for backward compatibility

  • Added deprecation handling: Implemented __getattr__ in main __init__.py to provide helpful warnings when using deprecated names:

    • ETraceState → brainstate.HiddenState

    • ETraceGroupState → brainstate.HiddenGroupState

    • ETraceTreeState → brainstate.HiddenTreeState

Neural Network Module Reorganization#

  • Consolidated neural network modules: Removed standalone neuron, synapse, and activation modules, migrating them to brainstate and brainpy ecosystems

    • Deleted files:

      • brainscale/nn/_neurons.py (IF, LIF, ALIF now in brainpy.state)

      • brainscale/nn/_synapses.py (Expon, Alpha, DualExpon, STP, STD now in brainpy.state)

      • brainscale/nn/_elementwise.py (activation functions now in brainstate.nn)

      • brainscale/nn/_poolings.py (pooling layers now in brainstate.nn)

  • Renamed _rate_rnns.py to _rnn.py: Simplified module naming for better clarity

  • Added comprehensive deprecation warnings in nn.__getattr__: Automatically redirects users to the correct modules:

    • Neuron models (IF, LIF, ALIF) → brainpy.state

    • Synapse models (Expon, Alpha, DualExpon, STP, STD) → brainpy.state

    • Activation functions (ReLU, Sigmoid, etc.) → brainstate.nn

    • Pooling layers (MaxPool, AvgPool, etc.) → brainstate.nn

    • Dropout layers → brainstate.nn

API Improvements#

  • Normalization parameter standardization: Renamed normalized_shape to in_size across all normalization layers for consistency

    • Updated in _normalizations.py for LayerNorm, GroupNorm, InstanceNorm, etc.

    • Improved clarity and consistency with other layer APIs

  • Enhanced input dimension validation: Improved error checking in convolutional layers to catch dimension mismatches early

  • Refactored imports for consistency: Updated all internal imports to use braintools for optimization and initialization utilities consistently across the codebase

Testing Infrastructure#

  • Added comprehensive unit tests for neural network modules:

    • _conv_test.py: 868 lines of tests for convolutional layers (Conv1d, Conv2d, Conv3d, ConvTranspose)

    • _linear_test.py: 658 lines of tests for linear layers (Linear, Identity)

    • _normalizations_test.py: 695 lines of tests for normalization layers (LayerNorm, BatchNorm, GroupNorm, etc.)

    • _readout_test.py: 763 lines of tests for readout layers (LeakyRateReadout, LeakySpikeReadout)

    • _rnn_test.py: 710 lines of tests for RNN cells (VanillaRNNCell, GRUCell, LSTMCell, MGUCell, etc.)

    • Total: 3,694 lines of new test coverage

Documentation Updates#

  • Streamlined API documentation: Updated docs/apis/nn.rst to remove redundant sections and enhance RNN overview

  • Updated tutorials and examples: All 16 tutorial notebooks and 11 example scripts updated to reflect new APIs:

    • Concepts tutorials (en/zh)

    • RNN and SNN online learning guides

    • Batching strategies documentation

    • ETrace state management examples

    • Graph visualization tutorials

Code Quality Improvements#

  • Removed redundant docstrings: Cleaned up duplicate documentation in LeakyRateReadout and LeakySpikeReadout

  • Improved code organization: Streamlined __all__ definitions across all modules

  • Enhanced readability: Consistent import structure and better code formatting throughout

Dependency Updates#

  • Updated requirements.txt: Refined dependency specifications to ensure compatibility with latest brainstate and brainpy versions

  • Updated pyproject.toml: Bumped version to 0.1.0 and updated project metadata

Breaking Changes#

API Changes:

  1. State class renaming (with deprecation warnings):

    • ETraceState → Use brainstate.HiddenState instead

    • ETraceGroupState → Use brainstate.HiddenGroupState instead

    • ETraceTreeState → Use brainstate.HiddenTreeState instead

  2. Neural network component migration (with deprecation warnings):

    • Neuron models (IF, LIF, ALIF) → Use brainpy.state module

    • Synapse models (Expon, Alpha, etc.) → Use brainpy.state module

    • Activation functions → Use brainstate.nn module

    • Pooling layers → Use brainstate.nn module

  3. Normalization parameter rename:

    • normalized_shape → in_size (for LayerNorm, GroupNorm, etc.)

  4. Module file reorganization:

    • nn/_rate_rnns.py → nn/_rnn.py

    • Removed: _neurons.py, _synapses.py, _elementwise.py, _poolings.py

Migration Guide#

For State Management:#

# Old (0.0.11)
from brainscale import ETraceState, ETraceGroupState

# New (0.1.0)
from brainstate import HiddenState, HiddenGroupState

For Neural Network Components:#

# Old (0.0.11)
from brainscale.nn import IF, LIF, Expon, ReLU, MaxPool2d

# New (0.1.0)
from brainpy.state import IF, LIF, Expon
from brainstate.nn import ReLU, MaxPool2d

For Normalization Layers:#

# Old (0.0.11)
norm = LayerNorm(normalized_shape=(128,))

# New (0.1.0)
norm = LayerNorm(in_size=128)

Note: All deprecated APIs include automatic warnings that will guide you to the correct replacements. The old APIs will continue to work in 0.1.0 but will be removed in a future release.

Version#

  • Bumped version from 0.0.11 to 0.1.0

Version 0.0.11#

Major Changes#

Import Refactoring#

  • Migrated imports from brainstate to braintools: All initialization-related imports now use braintools.init instead of brainstate.init

    • Updated imports in:

      • brainscale/nn/_neurons.py: Changed from brainstate import init to from braintools import init

      • brainscale/nn/_linear.py: Changed from brainstate import init to from braintools import init

      • brainscale/nn/_conv.py: Updated initialization imports

      • brainscale/nn/_synapses.py: Updated initialization imports

      • brainscale/nn/_readout.py: Updated initialization imports

  • Migrated neural network model imports from brainstate.nn to brainpy: Updated base classes for neuron models

    • IF, LIF, ALIF now inherit from brainpy instead of brainstate.nn

    • Maintained API compatibility while using the new brainpy backend

  • Updated functional API calls: Changed from brainstate.functional.sigmoid to brainstate.nn.sigmoid in RNN cells

Dependency Updates#

  • Added brainpy as a required dependency in requirements.txt

Documentation Enhancements#

  • Improved docstring formatting across the codebase:

    • Enhanced parameter documentation with proper type annotations using NumPy-style docstrings

    • Added missing “Returns” sections to property and method docstrings

    • Converted inline examples to proper “Examples” sections with code blocks

    • Updated documentation in:

      • brainscale/_etrace_algorithms.py: Enhanced EligibilityTrace and ETraceAlgorithm documentation

      • brainscale/_etrace_compiler_base.py: Improved parameter and return type documentation

      • brainscale/_etrace_compiler_module_info.py: Enhanced module documentation

Core Algorithm Updates#

  • RNN State Management: Updated all RNN cells to use braintools.init.param for state initialization and reset

    • ValinaRNNCell: Updated init_state() and reset_state() methods

    • GRUCell: Updated state management and activation functions

    • CFNCell: Updated forget and input gate implementations

    • MGUCell: Updated minimal gated unit state handling

Test Updates#

  • Refactored test imports: Updated test files to use new import paths

    • brainscale/_etrace_model_test.py: Updated with new import structure

    • brainscale/_etrace_vjp_algorithms_test.py: Aligned with new API

Version#

  • Bumped version from 0.0.10 to 0.0.11

Files Changed (17 files)#

  • .gitignore: Added new patterns

  • brainscale/__init__.py: Updated version number

  • brainscale/_etrace_algorithms.py: Enhanced documentation and imports

  • brainscale/_etrace_compiler_base.py: Improved documentation

  • brainscale/_etrace_compiler_graph.py: Minor updates

  • brainscale/_etrace_compiler_hidden_group.py: Minor updates

  • brainscale/_etrace_compiler_module_info.py: Enhanced documentation

  • brainscale/_etrace_model_test.py: Updated test imports

  • brainscale/_etrace_vjp_algorithms_test.py: Updated test imports

  • brainscale/_etrace_vjp_graph_executor.py: Updated imports

  • brainscale/nn/_conv.py: Migrated to braintools imports

  • brainscale/nn/_linear.py: Migrated to braintools imports

  • brainscale/nn/_neurons.py: Migrated to brainpy and braintools

  • brainscale/nn/_rate_rnns.py: Migrated to braintools and updated functional APIs

  • brainscale/nn/_readout.py: Updated imports

  • brainscale/nn/_synapses.py: Updated imports

  • requirements.txt: Added brainpy dependency

Breaking Changes#

None. All changes maintain backward compatibility at the API level.

Migration Guide#

If you have custom code using brainscale:

  • No changes required for end users

  • If extending brainscale internally, note that initialization utilities now come from braintools instead of brainstate

Contents