Release Notes

Release Notes#

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_yw_to_w.
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#

OSTL factory removed — use OSTLRecurrent or OSTLFeedforward directly instead of the former OSTL factory function.
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.
Unit dependency change — code relying on brainunit internals should migrate to saiunit.
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:

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

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:

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

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:

State class renaming (with deprecation warnings):
- ETraceState → Use brainstate.HiddenState instead
- ETraceGroupState → Use brainstate.HiddenGroupState instead
- ETraceTreeState → Use brainstate.HiddenTreeState instead
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
Normalization parameter rename:
- normalized_shape → in_size (for LayerNorm, GroupNorm, etc.)
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

Release Notes

Contents

Release Notes#

Version 0.2.0#

Major Changes#

New: SNN Online-Learning Algorithms#

ETP Compiler Rewrite#

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

Performance#

Primitive Registration Simplification#

Package Restructuring#

Typing & Packaging#

Testing#

Documentation#

Dependencies & Tooling#

Deprecations#

Breaking Changes#

Migration Guide#

OSTL#

OTTT / OTPE#

Custom ETP primitives#

Deprecated class-based API → primitive-based API#

Version#

Version 0.1.2#

Major Changes#

Import Path Migration#

New Algorithms#

Python 3.14 Support#

Bug Fixes#

CI/CD Improvements#

Documentation Updates#

Breaking Changes#

Migration Guide#

Update Dependencies#

Update Import Statements#

Version#

Version 0.1.1#

Major Changes#

Project Rename: BrainScale → BrainTrace#

VJP-Based Eligibility Trace Algorithms#

Logo and Branding#

Breaking Changes#

Migration Guide#

Update Import Statements#

VJP Algorithm Usage#

Version#

Version 0.1.0#

Major Changes#

State Management Refactoring#

Neural Network Module Reorganization#

API Improvements#

Testing Infrastructure#

Documentation Updates#

Code Quality Improvements#

Dependency Updates#

Breaking Changes#

Migration Guide#

For State Management:#

For Neural Network Components:#

For Normalization Layers:#

Version#

Version 0.0.11#

Major Changes#

Import Refactoring#

Dependency Updates#

Documentation Enhancements#

Core Algorithm Updates#

Test Updates#

Version#

Files Changed (17 files)#

Breaking Changes#

Migration Guide#