Release Notes

Version 0.3.0

This release delivers on-device NaN debugging, a unified compilation cache, simplified JAX compatibility, and major internal cleanup — with a net reduction of ~1,800 lines of code. It raises the minimum requirements to Python 3.11 and JAX 0.6.0.

Breaking Changes

• Python >= 3.11 required: Dropped support for Python 3.10. The requires-python field and classifiers now start at 3.11.

• JAX >= 0.6.0 required: All dependency groups (cpu, cuda12, cuda13, tpu, testing) now mandate jax>=0.6.0.

• Unified compilation cache in StatefulFunction: The four separate internal caches (_cached_jaxpr, _cached_out_shapes, _cached_jaxpr_out_tree, _cached_state_trace) have been consolidated into a single _compilation_cache storing _CachedCompilation objects. get_cache_stats() now returns {'compilation_cache': {...}} instead of four individual entries.

• Immutable CacheKey replaces hashabledict: get_arg_cache_key() now returns a CacheKey (NamedTuple) instead of the mutable hashabledict. Code that directly inspected or constructed cache keys must be updated.

• Removed internal _make_jaxpr function: The custom tracing implementation has been deleted in favor of using jax.make_jaxpr() directly (available in JAX >= 0.6.0).

• Removed debug_depth and debug_context from GradientTransform: The depth and context parameters for NaN debugging no longer exist following the debug module rewrite.

• Removed breakpoint_if function: The conditional breakpoint helper has been removed from brainstate.transform._debug.

• Removed extend_axis_env_nd from compatible imports: This compatibility shim is no longer exported.

New Features

On-Device NaN/Inf Detection

• Complete rewrite of the NaN debugging system (brainstate.transform._debug). NaN checking now runs on-device via JAX primitives rather than pulling data to the host, providing significantly better performance.

• Uses jax.debug.callback with thread-local storage to collect and report NaN findings.

• Error tracebacks now point to the user’s source code via source_info_util.user_context, producing IDE-clickable source locations extracted from jaxpr equations.

• Recursive instrumentation of nested primitives (jit, cond, while, scan) for comprehensive NaN detection throughout the computation graph.

• More compact and informative error messages via _format_nan_message().

JAX Traceback Filtering

• Registered brainstate with JAX’s traceback_util.register_exclusion() so internal frames are hidden in user-facing error tracebacks. Follows the same pattern as Flax, Equinox, and other JAX ecosystem libraries.

• Users can still see full tracebacks via JAX_TRACEBACK_FILTERING=off.

State Validation at Call Time

• New _validate_state_shapes() method checks that current state shapes and dtypes match those recorded at compile time.

• StatefulFunction.__call__() automatically validates before execution, catching state shape mismatches early with clear error messages.

• Added static_argnums bounds validation — make_jaxpr() now raises ValueError if indices exceed the number of positional arguments.

New Compatible Import

• Added mapped_aval import with version-based routing: jax.core.mapped_aval for JAX < 0.8.2, jax.extend.core.mapped_aval for >= 0.8.2.

Improvements

• Atomic cache writes: Compilation results are only stored on success, eliminating partial cache entries on error. Uses a double-checked locking pattern for thread safety during compilation.

• Better cache key hashing: Dynamic args/kwargs are now flattened via jax.tree.flatten() before hashing, fixing non-deterministic hashing issues with custom pytree nodes (e.g., Quantity).

• Modern Python type annotations: Migrated from typing.Tuple, typing.List, typing.Dict, typing.Optional, typing.Union to built-in tuple, list, dict, X | None, X | Y syntax across the codebase.

• IR visualization compatibility: Replaced direct jax.core.X references with compatible imports (Var, ClosedJaxpr, Jaxpr, JaxprEqn, Literal, DropVar) in the IR visualizer.

• Deterministic error reporting: jax.debug.callback in _error_if.py now uses ordered=True for deterministic error callback ordering.

• Graph operations cleanup: Major refactoring of _operation.py, _node.py, _convert.py, and _context.py with streamlined docstrings, better thread-safety documentation, and cleaner context managers.

Bug Fixes

• Fixed Delay.__init__ initialization order: update_every is now initialized before register_entry is called, preventing attribute errors during entry registration (#135).

• Fixed graph_to_tree private attribute access: Replaced internal _mapping access with public API usage in _convert.py.

Internal Changes

• Massive docstring reduction across the graph module (~1,000+ lines removed), replacing verbose multi-paragraph docstrings with concise descriptions.

• Cleaned up TypeVar usage: removed unused C and Names aliases, renamed Node TypeVar to N, removed Hashable bound from type variables.

• Removed unused tests (test_all_exports, test_function_imports_availability) from compatible import tests.

• Rewrote debug and make_jaxpr test suites to match the new APIs.

• IR optimization imports are now lazy-loaded inside make_jaxpr() only when ir_optimizations is configured.

CI/CD

• Bumped actions/upload-artifact from v6 to v7.

• Bumped actions/download-artifact from v7 to v8.

Version 0.2.10

This release introduces a comprehensive NaN debugging system for gradient computations, refactors the module mapping API for improved clarity, and adds graph context utilities for advanced state management.

New Features

NaN Debugging System

• JIT-Compatible NaN/Inf Debugging: New debugging utilities for identifying NaN and Inf values during gradient computations

  • debug_nan: Analyze a function for NaN/Inf values with detailed reporting

  • debug_nan_if: Conditional NaN debugging with predicate-based activation

  • Full JIT compatibility for seamless integration into compiled workflows

  • Support for debugging NaN in while and scan primitives

  • Detailed analysis output including variable names, shapes, and affected indices

• Gradient Function Integration: Added debug_nan parameter to gradient transformation functions

  • grad: Enable NaN debugging during gradient computation

  • vector_grad: NaN debugging for vectorized gradients

  • jacobian and jacobian_reverse: NaN debugging for Jacobian computations

  • hessian: NaN debugging for Hessian computations

• Breakpoint Utility: New breakpoint function for conditional debugging

  • Wraps jax.debug.breakpoint with predicate support

  • Only triggers when the specified condition is True

API Changes

Module System

• Renamed ModuleMapper to Map: Simplified naming for the vectorized module wrapper

  • Map provides vectorized (vmap2) and parallel (pmap2) mapping over modules

  • ModuleMapper retained as a deprecated alias for backward compatibility

  • Internal _ModuleMapperCalling renamed to _MapCaller for consistency

• Enhanced Map.map() Method: Now accepts callable functions for flexible mapping operations

Bug Fixes

• Fixed get_backend import for JAX version compatibility across different JAX releases

• Removed abstractmethod decorators from Regularization class to allow proper instantiation

• Cleaned up unused imports in module initialization files

Internal Changes

• Added comprehensive test suite for NaN debugging (_debug_test.py, 938 lines)

• Removed deprecated _mapping3.py module and associated tests

• Streamlined module exports in __init__.py files

Version 0.2.9

This release introduces a powerful state hook system for advanced state management, refactors neural network modules with enhanced parameter handling, and improves delay mechanisms with frequency-controlled updates.

State Management

State Hook System

• Global Hook Infrastructure: Comprehensive hook system for intercepting state operations

  • register_read_hook: Register hooks that execute when state values are read

  • register_write_hook: Register hooks that execute when state values are written

  • register_restore_hook: Register hooks that execute when state values are restored

  • HookManager: Thread-safe manager for organizing and executing hooks with priority support

  • HookContext: Context manager for scoped hook registration and execution

  • Enables advanced use cases: logging, debugging, value transformation, validation

• Enhanced State Class: Improved state management with hook integration

  • Automatic hook execution on read/write operations

  • Better cache key handling for improved performance

  • Enhanced thread safety and context management

  • Comprehensive test coverage (346 tests for thread safety, 320 tests for hooks)

Neural Network Components

Parameter Management (brainstate.nn.Param and brainstate.nn.Const)

• Renamed Classes: Simplified naming convention

  • ParaM → Param: Trainable parameter wrapper

  • ConstM → Const: Non-trainable constant wrapper

• Enhanced Caching System: Improved parameter precomputation and caching

  • param_precompute context manager for efficient parameter transformation caching

  • cache() method for retrieving cached parameter values

  • Support for custom precompute functions

  • Automatic cache invalidation and management

  • 391 comprehensive tests for caching behavior

• Hierarchical Parameter Data (brainstate.nn.HiData): New module for structured parameter organization

  • define_param_data() method for declaring hierarchical parameter structures

  • Support for nested parameter groups

  • Improved parameter surgery and manipulation

  • Enhanced type hints and documentation

Module System Enhancements

• ModuleMapper: New helper for vectorized module operations (formerly Vmap2Module)

  • Simplified API for applying vmap2 to module methods

  • Automatic state management for vectorized operations

  • Consistent interface with Vmap2ModuleCaller

  • Comprehensive documentation with usage examples

• Enhanced Module Methods:

  • parameters(): Iterate over all parameters in the module hierarchy

  • named_parameters(): Iterate over parameters with their qualified names

  • children(): Access direct child modules

  • named_children(): Access child modules with names

  • init_all_states(): Initialize states with additional keyword arguments

  • Improved Sequential with extend() and insert() methods

Delay Mechanisms

• Frequency-Controlled Updates: Enhanced Delay class with flexible update strategies

  • update_every parameter: Control how often delay buffers are updated

  • Support for integer steps (update every N steps)

  • Support for time-based updates with physical units (e.g., 1*ms)

  • Automatic handling of unit conversions and validation

  • Comprehensive tests covering various update strategies

• Unified Delay Implementation: Refactored delay mechanism

  • Ring buffer implementation for efficient historical value storage

  • Support for linear interpolation

  • Better handling of multi-dimensional inputs

  • Improved integration with neural network modules

Regularization

• Comprehensive Regularization Module (brainstate.nn._regularization, 2840 lines):

  • Complete suite of regularization techniques

  • L1, L2, and elastic net regularization

  • Dropout variants

  • Weight decay and other parameter constraints

  • 1261 tests for regularization functionality

• Transform Module (brainstate.nn._transform, 1661 lines):

  • Advanced parameter transformations

  • Quantization support

  • Normalization techniques

  • Integration with caching system

  • 452 comprehensive tests

Transformations

Vectorization and Parallelization

• Mapping Function Refactoring: Reorganized mapping implementations

  • Renamed _mapping.py → _mapping2.py (primary vmap2 implementation)

  • Renamed _mapping_old.py → _mapping1.py (legacy vmap implementation)

  • Added _mapping3.py: New pmap2 implementation for parallelization

  • vmap2_new_states: Helper for creating new states in vectorized operations

  • Relaxed return type requirements for more flexible mapping functions

• Enhanced Documentation: Updated tutorials and API documentation

  • Comprehensive vmap2 tutorial with practical examples

  • Enhanced parallelization documentation for pmap2

  • Updated state management guides

  • Expanded gradient transformation documentation

Compatibility and Utilities

JAX Compatibility

• Enhanced JAX Integration: Improved compatibility with newer JAX versions

  • Updated backend import for JAX version detection

  • Enhanced get_aval function for JAX version compatibility

  • Standardized jit_named_scope arguments

  • Support for JAX 0.8.0+ in CI configuration

Utility Functions

• Dataclass Support: Added is_dataclass utility function in brainstate.util.struct

  • Robust dataclass type checking

  • Better handling of dataclass-based structures

• Tracer Utilities: New _tracers.py module for JAX tracer handling

  • current_jax_trace(): Get current JAX trace context with version compatibility

  • Helper functions for working with JAX abstract values

Graph Operations

• Context Management (brainstate.graph._context):

  • New context management system for graph operations (119 lines)

  • TraceContextError: Specialized error class for tracing issues

  • Enhanced state tracking during graph construction

  • 64 tests for context management

• Conversion Utilities (brainstate.graph._convert):

  • New conversion utilities for graph operations (278 lines)

  • Better handling of graph transformations

  • Improved node conversion logic

Random Number Generation

• Enhanced RandomState: Improved random number generation

  • Better compatibility with newer JAX versions (98 lines of improvements)

  • Enhanced state management for random keys

  • Improved thread safety

  • Better error messages and validation

Documentation

• Comprehensive API Documentation: Expanded documentation across all modules

  • brainstate.rst: Reorganized with improved structure (21 lines removed, refactored into submodules)

  • environ.rst: Added 48 lines of documentation for environment state and keys

  • nn.rst: Added 222 lines documenting neural network components

  • transform.rst: Added 132 lines for gradient transformations and mapping functions

• Tutorial Updates:

  • Updated vectorization tutorial to reflect vmap → vmap2 transition

  • Enhanced examples with ModuleMapper usage

  • Improved state management examples

Breaking Changes

• Renamed Functions and Classes:

  • ParaM → Param

  • ConstM → Const

  • vmap → vmap2 (old vmap preserved in _mapping1.py for compatibility)

  • pmap → pmap2

  • _param_data → _hidata

• Parameter Naming Standardization:

  • fit_par → fit across all modules

  • brainscale → braintrace in example files

• Method Signature Changes:

  • init_all_states() now accepts additional keyword arguments

  • param_precompute() signature updated to support caching and custom functions

  • Module initialization methods enhanced with keyword argument support

Bug Fixes

• Fixed cache key handling in state management

• Improved error messages for missing states in gradient transformations

• Enhanced validation for delay update frequency

• Corrected import paths for better module organization

• Fixed compatibility issues with JAX 0.8.0+

Internal Changes

• Reorganized import statements across all modules for clarity

• Enhanced type hints throughout the codebase

• Improved code documentation with comprehensive docstrings

• Streamlined module exports in __all__ definitions

• Better separation of concerns in module organization

Version 0.2.8

This release ensures compatibility with JAX 0.8.2+ and removes the experimental module that was superseded by upstream changes.

Compatibility

• JAX 0.8.2+ Support: Added compatibility with JAX version 0.8.2 and later. The library now uses jax.make_jaxpr directly for JAX >= 0.8.2 while maintaining backward compatibility with earlier versions.

Breaking Changes

• Removed abstracted_axes parameter: The abstracted_axes parameter has been removed from:

  • StatefulFunction.__init__

  • StatefulMapping.__init__

  • make_jaxpr function

  • _make_jaxpr internal function

Improvements

• Debug mode support: Added debug_call method to StatefulFunction for proper execution when jax.config.jax_disable_jit is enabled. This improves debugging workflows by allowing stateful functions to execute without JIT compilation.

• Lazy loading optimization: RandomState import in the _mapping module is now lazily loaded via _import_rand_state(), improving initial import performance and reducing circular dependency issues.

Internal Changes

• Removed unused imports (annotate, api_boundary from jax._src) at module level; now imported only where needed

• Removed internal helper functions _broadcast_prefix and _flat_axes_specs

• Simplified _abstractify function by removing abstracted axes handling

• Updated example files to reflect API changes

Version 0.2.7

BrainState 0.2.7 modernizes the experimental compilation stack, deepens the transformation APIs, and tightens runtime infrastructure across the project.

Experimental Compiler and Visualization

• Introduced the experimental neuroir compiler built on dataclass-based graph IR elements and an explicit CompilationContext, improving dependency tracking, hidden-state mapping, and ClosedJaxpr fidelity even for self-connections and delay buffers.

• Added GraphDisplayer and TextDisplayer backends with hierarchical and force-directed layouts, plus richer diagnostics and tests that cover large sample networks and neuro-graph visualizations.

Transformations and Autodiff

• Added the jit_named_scope decorator and supporting utilities so nested transformations emit meaningful names inside traced functions, together with _make_jaxpr refinements that separate dynamic/static arguments and improve caching semantics for StatefulFunction.

• Expanded the gradient toolkit by exporting the new Jacobian (forward and reverse), Hessian, and SOFO transforms, unifying gradient handling for classes, auxiliary returns, and state-aware updates through the transform module.

State and Runtime Enhancements

• Replaced the experimental ArrayParam with a dedicated DelayState, propagating the new state through the compiler, delay modules, and neuro-IR so historical buffers participate in tracing and optimization just like other states.

• Environment helpers can now run against injected EnvironmentState instances, enabling sandboxed or per-thread configurations while DelayState-aware unit tests extend coverage of the updated modules.

Experimental and Infrastructure Updates

• Completed the neuron IR → neuroir rename, aligned the GDiist BPU codebase with the new terminology, and added new sample networks plus placeholder skips to keep the growing compiler/displayer test surface manageable.

• Added braincell to the development requirements, refreshed documentation wording, and kept CI dependencies current for the GitHub Actions runners.

Bug Fixes

• Hardened caching, randomness, and initialization logic by fixing get_arg_cache_key, removing stale decorator parameters, validating truncated normal draws, and correcting the exported version metadata.

• Declared Python 3.14 support and cleaned up compiler import ordering to keep linting noise low.

Version 0.2.6

This release focuses on the experimental export pipeline and device-aware execution adapters.

Device-Aware Wrappers

• Added registry-driven ForLoop and JIT adapters that expose decorator-style ergonomics, call counters, and validation, with CPU/GPU/TPU implementations wired through register_*_impl so experiments can swap device backends without touching user code.

GDiist BPU Export

• Replaced the monolithic exporter with gdiist_bpu.main, refreshed parser/component/utils modules, and renamed BpuParser to GdiistBpuParser, yielding clearer analysis output, text display helpers, and far more granular unit tests.

• Introduced the thread-safe BoundedCache utility and integrated it with compiler wrappers to safely reuse traced graphs, alongside _make_jaxpr updates that enforce argument checks and improve cache key generation.

• Updated tutorials and examples to the streamlined naming scheme and refreshed device implementation docs for the new wrapper entry points.

Version 0.2.5

Version 0.2.5 concentrates on intermediate-representation (IR) optimization quality.

IR Optimization

• Added _ir_optim_v2, a comprehensive optimizer that ships constant folding, dead-code elimination, common subexpression elimination, copy propagation, and algebraic simplification passes backed by identity-aware set semantics.

• Updated the transform exports and accompanying tests to exercise the new optimizer while pruning unused configuration knobs from the earlier implementation.

Version 0.2.4

This release introduces the new ArrayParam state type for parameter arrays with custom transformations, experimental BPU backend export support, enhanced JAXPR optimization capabilities, and improved module organization.

New Features

ArrayParam State Type

• ArrayParam Class: New state type for managing parameter arrays with advanced transformation control

  • Supports custom transformations (e.g., quantization, normalization) that preserve array identity

  • Enables vmap, pmap, and other JAX transformations to correctly handle stateful parameters

  • Provides identity() method that returns the raw array without applying custom transformations

  • Integrates seamlessly with existing State management infrastructure

  • Useful for implementing quantization-aware training and other advanced parameter manipulations

  • Comprehensive documentation with usage examples and best practices

Experimental BPU Backend Export (brainstate.experimental.gdiist_bpu)

• BPU Backend Export Support: Complete infrastructure for exporting models to GDiist BPU hardware backend (727 lines)

  • export.py: Main export API with to_bpu() function for model conversion

  • parser.py: Operation parser that analyzes JAXPR to identify operations and connections (305 lines)

  • data.py: Data structures and analysis utilities for operation representation (215 lines)

• Operation Parser Features:

  • Automatic detection of operations from JAXPR equations using brainevent primitives

  • Data flow analysis to identify connections between operations

  • Support for various operation types: slice, add, multiply, and more

  • Detailed analysis output showing equations, inputs, outputs, and connections

• Analysis and Debugging Tools:

  • display_analysis_results(): Comprehensive visualization of parsed operations

  • Shows operation details including equation count, variable mappings, and connections

  • Displays connection information with producer/consumer operations and variable details

  • Example implementation in examples/400_CUBA_2005_bpu.py

Enhancements

JAXPR Optimization Improvements

• Enhanced Constant Folding:

  • Better handling of literal values in constant folding optimization

  • Improved detection and elimination of redundant literal operations

  • More efficient constant propagation through computation graphs

• Identity Equation Optimization:

  • Optimized handling of Literal outputs to avoid unnecessary bridging equations

  • Improved identity equation creation for interface preservation

  • Better handling of edge cases in optimization passes

• Error Handling:

  • Added fallback source info utility for better error messages

  • Fixed potential NoneType errors in equation handling

  • Improved validation of optimization results

State Management

• Enhanced State Tests: Comprehensive test refactoring with improved coverage (454 tests)

  • Better organization of state type tests

  • More thorough validation of state behavior

  • Enhanced test readability and maintainability

Version 0.2.3

This release introduces powerful IR (Intermediate Representation) optimization capabilities for JAX computation graphs, comprehensive state management refactoring for vectorized mapping operations, and extensive testing infrastructure improvements.

New Features

IR Optimization (brainstate.transform._ir_optim)

• Intermediate Representation Optimization Module (876 lines): Complete suite of compiler-level optimizations for JAX computation graphs

  • constant_fold: Evaluates constant expressions at compile time, reducing runtime computation

  • dead_code_elimination: Removes equations whose outputs are unused, reducing computation overhead

  • common_subexpression_elimination: Identifies and reuses results of identical computations

  • copy_propagation: Eliminates unnecessary copy operations by propagating original variables

  • algebraic_simplification: Applies algebraic identities (x+0=x, x*1=x, x-x=0, etc.)

  • optimize_jaxpr: Orchestrates multiple optimization passes with configurable iteration and verbose mode

• IdentitySet Class: Custom set implementation using object identity (id()) instead of equality

  • Enables proper handling of JAX variables and Literals in optimization passes

  • Implements MutableSet interface with full collection protocol support

  • Essential for tracking variable usage without relying on equality comparisons

Optimization Features

• Interface Preservation: All optimizations preserve function input/output variables (invars/outvars)

  • Identity equations automatically added when needed to maintain correct interfaces

  • Uses convert_element_type primitive with matching dtypes as identity operation

  • Ensures optimized functions remain drop-in replacements

• Optimization Pipeline: Configurable multi-pass optimization with convergence detection

  • Customizable optimization sequence via optimizations parameter

  • Automatic convergence detection when no more reductions possible

  • Maximum iteration control with max_iterations parameter

  • Verbose mode with detailed statistics and progress tracking

• JAX Integration: Full support for JAX primitives and special cases

  • Blacklist for primitives that shouldn’t be folded (broadcast_in_dim, broadcast)

  • Proper handling of closed_call and scan primitives

  • Support for both Jaxpr and ClosedJaxpr inputs

State Management Refactoring (brainstate.transform._mapping)

• Renamed vmap to vmap2: Major refactoring of vectorized mapping implementation (647 lines)

  • Enhanced state management with improved axis tracking

  • Better error messages and validation

  • Streamlined state value restoration logic

• Old vmap Implementation Preserved (_mapping_old.py, 579 lines): Legacy vmap with explicit state management

  • Exports original vmap and vmap_new_states functions

  • Maintains backward compatibility for existing code

  • Specialized for stateful functions with explicit state parameters

Documentation

API Documentation

• transform.rst: Added comprehensive IR Optimization section (24 lines)

  • Detailed module description explaining compiler optimizations

  • All 6 optimization functions documented with autosummary

  • Clear explanation of benefits: reduced computation overhead, improved runtime performance

  • Positioned between Compilation Tools and Gradient Computations sections

• NumPy-style Docstrings: All optimization functions include:

  • Comprehensive parameter descriptions with types and defaults

  • Detailed return value documentation

  • Notes sections explaining preservation of function interfaces

  • Multiple practical examples demonstrating usage

  • Algorithm descriptions for complex optimizations

  • Cross-references between related functions

Enhancements

Optimization Pipeline

• Progress Tracking: Verbose mode shows equation count changes after each optimization

  • Displays initial, intermediate, and final equation counts

  • Shows reduction statistics with percentages

  • Indicates convergence detection

  • Reports iteration counts

• Validation: Runtime checks ensure optimization correctness

  • Verifies input variables unchanged after optimization

  • Validates output variables preserved

  • Raises clear errors if interface violated

  • Checks for valid optimization names

• Flexibility: Customizable optimization sequences

  • Apply all optimizations in recommended order (default)

  • Select specific optimizations only

  • Control iteration limits

  • Toggle verbose output

JAX Integration

• JaxprEqn Construction: Proper handling of required ctx parameter

  • Uses JaxprEqnContext(None, True) for identity equations

  • Ensures compatibility with JAX internal API

  • Maintains proper equation structure

• Primitive Handling: Special cases for JAX primitives

  • Blacklist for primitives that shouldn’t be optimized

  • Proper parameter extraction and validation

  • Support for effects and source_info fields

Bug Fixes

• Fixed JaxprEqn constructor calls to include required ctx parameter (7th positional argument)

• Corrected import paths for vmap2 in test files and tutorials

• Fixed RandomState.uniform() calls to use size parameter instead of shape

• Enhanced test assertions for proper state axis handling

• Improved error messages for batch axis mismatches

Refactoring

Transform Module

• Renamed Files:

  • vmap → vmap2 in _mapping.py

  • Preserved original vmap in _mapping_old.py for compatibility

• Module Exports: Updated __init__.py to export both old and new vmap implementations

  • vmap from _mapping_old.py (legacy)

  • vmap2 from _mapping.py (new)

  • vmap_new_states from both modules

Version 0.2.2

This release focuses on enhancing hidden state management for recurrent neural networks and eligibility trace-based learning, along with comprehensive testing and documentation improvements.

New Features

Hidden State Classes

• HiddenGroupState: New class for managing multiple hidden states within a single array

  • Stores multiple states in the last dimension of a single array

  • Provides get_value() and set_value() methods for accessing individual states by index or name

  • Optimized for LSTM-style architectures with multiple hidden components (h, c)

  • Includes name2index mapping for convenient state access

• HiddenTreeState: New class for managing multiple hidden states with different physical units

  • Supports PyTree structure (dict or sequence) of hidden states

  • Preserves physical units (e.g., voltage, current, conductance) via brainunit integration

  • Provides name2unit and index2unit mappings for unit tracking

  • Ideal for neuroscience models with heterogeneous state variables

  • Maintains compatibility with BrainScale online learning

State Utilities

• maybe_state: New utility function for flexible value extraction

  • Extracts values from State objects automatically

  • Returns non-State values unchanged

  • Simplifies writing functions that accept both states and raw values

Enhancements

State Classes

• HiddenState: Enhanced documentation and type checking

  • Restricted to numpy.ndarray, jax.Array, and brainunit.Quantity types only

  • Added comprehensive docstrings with examples

  • Clarified equivalence to brainstate.HiddenState for online learning

  • Improved error messages for invalid input types

• BatchState: Now properly exported in the public API

  • Available via brainstate.BatchState

  • Enhanced documentation for batch data management

Documentation

• API Reference: Completely reorganized brainstate.rst documentation

  • Organized into 6 major sections: Core State Classes, State Management, State Utilities, Error Handling, and Submodules

  • Added detailed descriptions for each section and subsection

  • Included comprehensive bullet-point summaries for all APIs

  • Enhanced deprecation warnings with clear migration paths

  • Added module-level descriptions for all submodules

• State Classes: Enhanced documentation for all state types

  • Added detailed use case descriptions

  • Included practical examples for each state type

  • Clarified semantic distinctions between state types

  • Documented integration with JAX transformations

• JAX Transformations: Improved documentation for stateful transforms

  • Enhanced docstrings for jit, grad, vmap, scan, and other transforms

  • Added examples showing state management patterns

  • Documented state tracing behavior

  • Clarified interaction with StateTraceStack

Transform System

• Enhanced State Finding: New _find_state.py module for automatic state discovery

  • Improved state detection in nested structures

  • Better handling of state dependencies

  • Enhanced error messages for state-related issues

• StatefulFunction: Major enhancements to make_jaxpr functionality

  • Improved Jaxpr generation for stateful computations

  • Better handling of state read/write tracking

  • Enhanced debugging support

• Mapping Transformations: Significant refactoring of vmap and pmap

  • Improved state management across vectorized operations

  • Better handling of state broadcasting

  • Enhanced error reporting for mapping operations

Random Number Generation

• Module Reorganization: Complete refactoring of random module structure

  • Renamed _rand_funs.py to _fun.py

  • Renamed _rand_seed.py to _seed.py

  • Renamed _rand_state.py to _state.py

  • Extracted distribution implementations to new _impl.py module (691 lines)

• Improved Random State: Enhanced RandomState class with better state management

  • Simplified implementation (reduced from 534 to ~300 lines)

  • Better integration with JAX’s random number generation

  • Improved thread safety and state isolation

Testing

• Comprehensive Test Suite: Added 102 tests covering all state functionality

  • TestBasicState (13 tests): Core State class operations

  • TestShortTermState (2 tests): Short-term state behavior

  • TestLongTermState (2 tests): Long-term state behavior

  • TestParamState (2 tests): Parameter state usage patterns

  • TestBatchState (2 tests): Batch state functionality

  • TestHiddenState (7 tests): Hidden state with different array types

  • TestHiddenGroupState (9 tests): Multiple hidden state management

  • TestHiddenTreeState (12 tests): PyTree hidden states with units

  • TestFakeState (4 tests): Lightweight state alternative

  • TestStateDictManager (6 tests): State collection management

  • TestStateTraceStack (11 tests): State tracing and recovery

  • TestTreefyState (6 tests): PyTree state references

  • TestContextManagers (6 tests): State context managers

  • TestStateCatcher (8 tests): State catching utilities

  • TestIntegrationScenarios (5 tests): Real-world use cases

Bug Fixes

• Fixed HiddenGroupState.set_value() to work correctly with JAX arrays

• Improved error handling in hidden state value validation

• Enhanced type checking for hidden state initialization

Documentation

Tutorial Reorganization

• Basics Tutorials: Complete rewrite and expansion

  • 01_getting_started.ipynb: Enhanced introduction with practical examples

  • 02_state_management.ipynb: Comprehensive state management guide

  • 03_random_numbers.ipynb: In-depth random number generation tutorial

• Neural Networks Tutorials: Restructured and expanded

  • 01_module_basics.ipynb: New comprehensive module system guide

  • 02_basic_layers.ipynb: Enhanced layer documentation with examples

  • 03_activations_normalization.ipynb: Detailed activation and normalization guide

  • 04_recurrent_networks.ipynb: New RNN tutorial with practical examples

  • 05_dynamics_systems.ipynb: New dynamical systems tutorial

• Examples: Reorganized and enhanced

  • Renamed 10_image_classification.ipynb to 01_image_classification.ipynb

  • Renamed 11_sequence_modeling.ipynb to 02_sequence_modeling.ipynb

  • Added 03_brain_inspired_computing.ipynb: New brain-inspired computing examples

  • Renamed 18_optimization_tricks.ipynb to 04_optimization_tricks.ipynb

  • Renamed 19_model_deployment.ipynb to 05_model_deployment.ipynb

• Transforms Tutorials: Reorganized for better flow

  • 01_jit_compilation.ipynb: New comprehensive JIT guide

  • 02_automatic_differentiation.ipynb: Enhanced autodiff tutorial

  • 03_vectorization.ipynb: Improved vmap/pmap guide

  • 04_loops_conditions.ipynb: Enhanced control flow guide

  • 05_other_transforms.ipynb: Other transformation utilities

• Advanced Tutorials: Renumbered for clarity

  • 01_graph_operations.ipynb (formerly 14_graph_operations.ipynb)

  • 02_mixin_system.ipynb (formerly 15_mixin_system.ipynb)

  • 03_typing_system.ipynb (formerly 16_typing_system.ipynb)

  • 04_utilities.ipynb (formerly 17_utilities.ipynb)

• Migration Guides: Updated and simplified

  • 01_migration_from_pytorch.ipynb: Enhanced PyTorch migration guide

  • Removed outdated BrainPy integration notebook

• Supplementary: Reorganized

  • 01_performance_optimization.ipynb

  • 02_debugging_tips.ipynb

  • 03_faq.ipynb: Updated FAQ with new content

API Documentation

• Enhanced module documentation in nn.rst with 306 line improvements

• Updated transform.rst with new transform APIs

• Improved environ.rst and graph.rst documentation

Refactoring

• Removed deprecated eval_shape module and tests

• Removed deprecated _random.py transform module

• Cleaned up unused imports across all modules

• Improved code organization in neural network layers

• Enhanced type hints and docstrings throughout

Infrastructure

• Added development dependency for tutorial generation

• Updated benchmark scripts for performance testing

• Improved test coverage across transformation modules

Version 0.2.0

This is a major release with significant refactoring, new features, and comprehensive documentation improvements.

Breaking Changes

• Module Deprecations: Deprecated brainstate.transform, brainstate.transform, and brainstate.functional modules in favor of brainstate.transform and brainstate.nn

  • Added deprecation proxies to guide users towards replacement modules

  • Updated all documentation and examples to use new module paths

• State Management: Replaced write_back_state_values with assign_state_vals_v2 for improved state management

• Import Path Changes: Major refactoring of import paths across the codebase

  • Moved initialization references to use brainstate.nn

  • Updated random functions to use brainstate.random

  • Standardized imports across all modules

• Type System: Implemented JointTypes and OneOfTypes generic aliases to enhance type checking and avoid metaclass conflicts

  • Support for subscript syntax

  • Improved type hints across modules

• Copyright: Updated copyright notices to reflect new ownership by BrainX Ecosystem Limited

New Features

Neural Network Components

• Transposed Convolution Layers: Complete implementations for upsampling operations

  • ConvTranspose1d, ConvTranspose2d, ConvTranspose3d

  • Support for both channels-first and channels-last data formats via channel_first parameter

  • Configurable stride for controllable upsampling factors

  • Grouped transposed convolution support

  • Automatic padding computation for ‘SAME’ and ‘VALID’ modes

• Convolution Enhancements: Added support for both channels-first and channels-last data formats

  • New channel_first boolean parameter (default: False)

  • PyTorch-compatible format (e.g., [B, C, H, W]) when channel_first=True

  • Default JAX-style format (e.g., [B, H, W, C]) when channel_first=False

• Padding Layers: Added padding layers for 1D, 2D, and 3D tensors with various modes

• Unpooling Layers: Added MaxUnpool1d, MaxUnpool2d, and MaxUnpool3d with return_indices support

• Gradient Utilities: Implemented clip_grad_norm function for gradient clipping in PyTree structures

• Embedding Enhancements:

  • Added padding_idx, max_norm, and norm_type parameters

  • Improved gradient management with new _contains_tracer function

  • Optimized max_norm application with accessed mask for scaling

• BatchNorm Improvements: Added feature_axis and track_running_stats parameters

• LoRA Layer: Added in_size parameter for improved size handling

• Activation Functions: Added new activation functions and improved signatures

Transform & Compilation

• StatefulMapping: Introduced for enhanced state management in vmap transformations

• Mixin Classes: Added Mode, JointMode, Batching, and Training classes for computation behavior control

• Bounded Cache: Implemented thread-safe bounded cache for JAX Jaxpr with:

  • Comprehensive validation

  • Statistics tracking

  • Enhanced error handling

• Input Validation: Enhanced input size handling to support numpy integer types

• Context Parameters: Update method now accepts additional context parameters for improved environment settings

Random & Initialization

• Dependencies: Integrated braintools for initialization and surrogate gradient functions

  • Updated all initialization references

  • Refactored to use braintools.surrogate for spike functions

• Random Functions: Replaced uniform_for_unit with jr.uniform for consistency and performance

Utilities & Infrastructure

• Filter Utilities: Added comprehensive filter utilities for nested structures

• Pretty Representation: Enhanced pretty_pytree module with:

  • Comprehensive documentation

  • Mapping functions

  • JAX integration

• Error Handling: Improved state length validation by replacing assertions with ValueError exceptions

• Collective Operations: Updated function signatures to return target in collective operations

Documentation

• Comprehensive Docstrings: Added detailed NumPy-style docstrings across all modules

  • Full parameter descriptions with types and default values

  • Multiple practical examples in code blocks

  • Comparison sections highlighting differences from PyTorch

  • Mathematical formulas where applicable

  • References to original papers

  • Best practices and use cases

• New Documentation Pages:

  • brainstate.environ module documentation

  • brainstate.transform (renamed from compile.rst)

  • Random number generation module

  • Pretty representation module

  • State management tutorial notebook

• Enhanced Examples: Updated documentation examples to use interactive prompts for clarity

• Module Descriptions: Enhanced documentation with detailed descriptions, key features, and usage examples

Testing

• Comprehensive Test Coverage: Added extensive test suites for:

  • _BoundedCache and StatefulFunction

  • brainstate.mixin module

  • brainstate.environ module (context management, precision settings, callbacks)

  • DeprecatedModule and proxy creation functionality

  • Compatible import module

  • Metrics module

  • Node class and helper functions

  • Activation functions with shape and gradient checks

  • Dropout layers

  • Surrogate gradient functions

  • Filter utilities

  • Struct module

  • Pretty representation

• Test Framework Updates: Refactored tests to use absltest for better JAX compatibility

Refactoring

• File Reorganization:

  • Renamed metrics.py to _metrics.py

  • Renamed _rate_rnns.py to _rnns.py

  • Renamed _init.py to init.py

  • Reorganized graph module files

  • Cleaned up unused imports and classes

• Code Quality:

  • Streamlined imports across all modules

  • Enhanced code formatting and whitespace consistency

  • Removed unnecessary inheritance and unused elements

  • Simplified type annotations

  • Improved method signatures for clarity

• Neuron & Synapse Classes: Refactored to use brainpy module and updated initialization methods

• Base Classes: Changed base class of EINet and Net from DynamicsGroup to Module for consistency

• Evaluation Functions: Refactored and updated method names for consistency

Infrastructure

• Version Bump: Updated version to 0.2.0

• Development Dependencies: Added braintools to development requirements

• Issue Templates: Added bug report and feature request templates for improved issue tracking

• CI/CD: Refactored CI configurations to update pip installation commands

• Git Ignore: Updated to exclude example figures directory and build artifacts

Bug Fixes

• Enhanced delay handling for multi-dimensional inputs

• Fixed gradient function references

• Improved deprecation handling in tests

• Fixed precision checks in complex number handling

Version 0.1.0

The first version of the project.