Release Notes

Contents

Release Notes#

Version 0.1.9 (2026-05-21)#

This release introduces braintools.cogtask, a composable framework for constructing cognitive tasks for neural-network training and computational neuroscience experiments. It also extends braintools.init to accept brainstate.nn.Param, adds official Python 3.14 support, and refreshes documentation and CI infrastructure.

Highlights#

  • New module braintools.cogtask: a phase-based DSL for building trial-structured cognitive tasks, with a library of pre-built paradigms drawn from the cognitive-neuroscience literature.

  • Variable-length trials under JIT/vmap: shape-stable packed-mode trial generation enables batch_sample to remain compatible with brainstate.transform.jit and brainstate.transform.vmap2 for tasks whose phases have data-dependent durations.

  • Python 3.14 support: CI matrix and project metadata updated; minimum supported Python remains 3.11.

Added#

braintools.cogtask — composable cognitive task framework#

  • Core API: Task, TaskConfig, Context, Phase, Sequence, Repeat, Parallel, conditional combinators If / Switch / While, and the concat helper for sequential composition.

  • Phase primitives: Fixation, Delay, Stimulus, Response, Cue, plus the variable-length VariableDuration phase whose timestep budget is resolved per-trial from a context entry.

  • Feature and label utilities: Feature, circular, one_hot, and a set of input encoders/decoders for constructing task observations and targets in a typed, composable way.

  • Pre-built task library spanning three domains:

    • Decision making: PerceptualDecisionMaking, PerceptualDecisionMakingDelayResponse, ContextDecisionMaking, SingleContextDecisionMaking, PulseDecisionMaking.

    • Working memory: DelayMatchSample, DualDelayMatchSample, DelayComparison, DelayMatchCategory, DelayPairedAssociation, GoNoGo, IntervalDiscrimination, PostDecisionWager, ReadySetGo, DelayDirectionReproduction, ImmediateDirectionReproduction, DelayDirectionClassification, ImmediateDirectionClassification.

    • Motor and reasoning: AntiReach, Reaching1D, EvidenceAccumulation, HierarchicalReasoning, ProbabilisticReasoning.

  • Variable-length trial sequences:

    • VariableDuration phases declare a Python max_steps (used as the buffer slot size) and report the realised trial length via the traced step_count field.

    • Task auto-detects variable-length phase trees via phase_tree_is_variable(phases). When any phase declares is_variable = True, sample_trial allocates buffers of size task.max_trial_duration(), writes only the front t_cursor timesteps, and zero-fills the trailing positions while setting the mask to False.

    • Task.batch_sample(..., return_mask=True) returns (X, Y, mask) with shape-stable buffers under brainstate.transform.jit and brainstate.transform.vmap2. Fixed-length tasks remain unaffected; return_mask=True on a fixed task yields an all-True mask.

    • If uses jax.lax.cond so both branches contribute shape-stable output; Switch and While use Python dispatch (concrete keys / Python bool conditions) and zero-pad to the static maximum.

    • HierarchicalReasoning, IntervalDiscrimination, and ReadySetGo have been migrated to VariableDuration and are now usable under batch_sample with masking — previously they were valid only via sample_trial and were not vmap-safe.

    • Duration samplers TruncExp and UniformDuration advertise is_variable = True and expose max_value() / min_value() so phases can size their slots statically from sampler bounds.

Other additions#

  • braintools.init.param: now accepts brainstate.nn.Param instances in addition to plain initializers, enabling reuse of pre-built parameter objects when constructing layers.

  • .gitattributes: normalises line endings for text files to keep diffs and tooling consistent across platforms.

Changed#

  • Python support: project metadata, CI matrix, and classifiers updated to include Python 3.14. Minimum supported Python remains 3.11.

  • Documentation hosting: docs are now self-hosted at brainx.chaobrain.com/braintools/; the documentation deployment workflow publishes on GitHub release events, while pushes to main run a build-only verification step.

  • Documentation dependencies: bumped sphinx (>=5 → >=9.0.4), sphinx-book-theme (>=1.0.1 → >=1.2.0), sphinx-copybutton (>=0.5.0 → >=0.5.2), jupyter-sphinx (>=0.3.2 → >=0.5.3), and brainx-sphinx-header.

Fixed#

  • braintools.cogtask end-to-end correctness pass (introduced together with the module):

    • Renamed cogtask/typing.py to cogtask/_typing.py so the local module no longer shadows the stdlib typing when tests run from inside the package; updated the absolute import in feature.py to the relative from ._typing import Data.

    • Added the missing noise_sigma argument and attribute to PerceptualDecisionMaking, PerceptualDecisionMakingDelayResponse, ContextDecisionMaking, SingleContextDecisionMaking, AntiReach, Reaching1D, EvidenceAccumulation, DelayPairedAssociation, GoNoGo, and PostDecisionWager, which previously raised AttributeError as soon as define_phases ran.

    • Removed a duplicate Task.__repr__ and an undocumented TaskLoader symbol from the public docs.

  • Phase engine: added an IS_COMPOUND flag on Phase and a uniform children() traversal hook. execute_phase now dispatches Sequence/Repeat/Parallel/If/Switch/While to their own execute() methods; previously, If/Switch/While silently no-op’d. Parallel.execute now gives each child its own [phase_start, phase_start + duration) window.

  • Distance-profile tests: representation-equality checks corrected so the test suite is stable across NumPy/JAX representations.

Infrastructure#

  • Bumped GitHub Actions: actions/checkout 4 → 6, actions/setup-python 5 → 6, actions/download-artifact 5 → 8, actions/upload-artifact 6 → 7, appleboy/ssh-action 1.2.0 → 1.2.5, appleboy/scp-action 0.1.7 → 1.0.0.

  • Refactored version management: a dedicated braintools/_version.py module is now the single source of truth, and pyproject.toml resolves the package version via attr = "braintools._version.__version__".

Version 0.1.7 (2026-01-05)#

Major Features#

New Training Framework (braintools.trainer)#

  • PyTorch Lightning-like training API for JAX-based neural network training with comprehensive features:

    • LightningModule: Base class for defining training models with training_step(), validation_step(), and configure_optimizers() hooks

    • Trainer: Orchestration class for managing training loops, epochs, and device placement

    • TrainOutput/EvalOutput: Structured output types for training and evaluation results

Callbacks System#

  • 10+ built-in callbacks for customizing training behavior:

    • ModelCheckpoint: Automatic model saving based on monitored metrics

    • EarlyStopping: Stop training when metrics plateau

    • LearningRateMonitor: Track and log learning rate changes

    • GradientClipCallback: Gradient clipping for training stability

    • Timer: Track training time

    • RichProgressBar / TQDMProgressBar: Visual progress indicators

    • LambdaCallback / PrintCallback: Custom callback utilities

Logging Backends#

  • 6 pluggable logging backends:

    • TensorBoardLogger: TensorBoard integration

    • WandBLogger: Weights & Biases integration

    • CSVLogger: Simple CSV file logging

    • NeptuneLogger: Neptune.ai integration

    • MLFlowLogger: MLFlow integration

    • CompositeLogger: Combine multiple loggers

Data Loading Utilities#

  • JAX-compatible data loading with distributed support:

    • DataLoader / DistributedDataLoader: Efficient batch loading

    • Dataset, ArrayDataset, DictDataset, IterableDataset: Dataset abstractions

    • Sampler, RandomSampler, SequentialSampler, BatchSampler, DistributedSampler: Sampling strategies

Distributed Training#

  • Multi-device and multi-host training strategies:

    • SingleDeviceStrategy: Single device execution

    • DataParallelStrategy: Data parallelism across devices

    • ShardedDataParallelStrategy / FullyShardedDataParallelStrategy: Memory-efficient sharded training

    • AutoStrategy: Automatic strategy selection

    • all_reduce, broadcast: Distributed communication primitives

Checkpointing#

  • Comprehensive checkpoint management:

    • CheckpointManager: Manage multiple checkpoints with retention policies

    • save_checkpoint / load_checkpoint: Save and restore model states

    • find_checkpoint / list_checkpoints: Checkpoint discovery utilities

Progress Bar System#

  • Multiple progress bar implementations:

    • SimpleProgressBar: Basic text-based progress

    • TQDMProgressBarWrapper: TQDM-based progress

    • RichProgressBarWrapper: Rich library-based progress

Improvements#

API Documentation#

  • Enhanced module documentation: All public modules now include comprehensive docstrings with examples, parameter descriptions, and usage guidelines directly in __init__.py files

  • Reorganized imports: Cleaner and more consistent import structure across all modules

Breaking Changes#

Removed braintools.param Module#

  • The entire braintools.param module has been removed, including:

    • Data containers (Data)

    • Parameter wrappers (Param, Const)

    • State containers (ArrayHidden, ArrayParam)

    • Regularization classes (GaussianReg, L1Reg, L2Reg)

    • All transform classes (SigmoidT, SoftplusT, AffineT, etc.)

    • Utility functions (get_param(), get_size())

  • Users relying on these features should migrate to alternative implementations or pin to version 0.1.6

Version 0.1.6 (2025-12-25)#

New Features#

Parameter Management Expansion (braintools.param)#

  • Hierarchical data container: Added Data for composed state storage and cloning.

  • Parameter wrappers: Added Param and Const with built-in transforms and optional regularization.

  • State containers: Added ArrayHidden and ArrayParam with transform-aware .data access.

  • Regularization priors: Added GaussianReg, L1Reg, and L2Reg with optional trainable hyperparameters.

  • Utilities: Added get_param() and get_size() helpers for parameter/state handling.

Transforms#

  • New ReluT transform for lower-bounded parameters.

  • Expanded transform suite now includes PositiveT, NegativeT, ScaledSigmoidT, PowerT, OrderedT, SimplexT, and UnitVectorT.

Improvements#

API Consistency#

  • Transform naming cleanup: Standardized transform class names with the *T suffix (e.g., SigmoidT, SoftplusT, AffineT, ChainT, MaskedT, ClipT).

Documentation#

  • Expanded param API docs: Added sections for data containers, state containers, regularization, utilities, and updated transform listings in docs/apis/param.rst.

  • API index update: Added param API page to docs/index.rst.

Tests#

  • New test coverage: Added tests for data containers, modules, regularization, state, transforms, and utilities across the param module.

Breaking Changes#

  • Transform API renames: Transform classes now use the *T suffix (e.g., Sigmoid -> SigmoidT).

  • Custom transform removed: The Custom transform is no longer part of the public API.

Bug Fixes#

  • Initializer RNG: TruncatedNormal now defaults to numpy.random when no RNG is provided.

Version 0.1.5 (2025-12-14)#

New Features#

Parameter Transformation Module (braintools.param)#

  • 7 new bijective transforms for constrained optimization and probabilistic modeling:

    • Positive: Constrains parameters to (0, +∞) using exponential transformation

    • Negative: Constrains parameters to (-∞, 0) using negative softplus

    • ScaledSigmoid: Sigmoid with adjustable sharpness/temperature parameter (beta)

    • Power: Box-Cox family power transformation for variance stabilization

    • Ordered: Ensures monotonically increasing output vectors (useful for cutpoints in ordinal regression)

    • Simplex: Stick-breaking transformation for probability vectors summing to 1

    • UnitVector: Projects vectors onto the unit sphere (L2 norm = 1)

  • Jacobian computation: Added log_abs_det_jacobian() method to Transform base class and implementations for probabilistic modeling

    • Implemented for: Identity, Sigmoid, Softplus, Log, Exp, Affine, Chain, Positive

Surrogate Gradient Enhancements (braintools.surrogate)#

  • Gradient computation of hyperparameters of surrogate gradient functions.

  • Fix batching issue in surrogate gradient functions

Improvements#

API Enhancements#

  • __repr__ methods: Added string representations to all Transform classes and Param class for better debugging

  • Enhanced documentation: Updated docs/apis/param.rst with comprehensive API reference

    • Organized sections: Base Classes, Parameter Wrapper, Bounded Transforms, Positive/Negative Transforms, Advanced Transforms, Composition Transforms

    • Descriptive explanations for each transform’s use case

Code Quality#

  • Comprehensive test coverage: Added 28 new tests for param module (45 total tests passing)

    • Tests for all new transforms: roundtrip, constraints, repr methods

    • Tests for log_abs_det_jacobian correctness

    • Tests for edge cases and numerical stability

Version 0.1.4 (2025-10-31)#

New Features#

Learning Rate Scheduler Enhancements (braintools.optim)#

  • New apply() method: Added apply() method to all LR schedulers for more flexible learning rate application

    • Allows applying learning rate transformations without stepping the scheduler

    • Useful for custom training loops and learning rate inspection

  • Comprehensive test coverage: Added 118+ comprehensive tests covering all 17 learning rate schedulers

    • Tests for basic functionality, optimizer integration, JIT compilation, state persistence

    • Full coverage of edge cases and special modes for each scheduler

    • Validates correctness with @brainstate.transform.jit compilation

Improvements#

Documentation#

  • Restructured tutorial organization: Renamed and reorganized documentation files for better clarity

    • Moved module tutorials into subdirectories (conn/, init/, input/, file/, surrogate/)

    • Updated table of contents structure across all modules

    • Improved navigation with consolidated index files (index.md instead of toc_*.md)

  • Enhanced visual branding: Updated project logo from JPG to high-resolution PNG format

    • Better quality and transparency support

    • Consistent branding across documentation

Code Quality#

  • Test improvements: Refactored scheduler tests with better organization and coverage

    • Each scheduler now has 5-10 dedicated tests

    • Tests verify: basic functionality, optimizer integration, JIT compilation, multiple param groups, state dict save/load

    • Discovered and documented key implementation behaviors (epoch counting, initialization patterns)

CI/CD#

  • Updated GitHub Actions: Bumped actions to latest versions for improved security and performance

    • actions/download-artifact: v5 → v6

    • actions/upload-artifact: v4 → v5

    • Better artifact handling in CI pipeline

Bug Fixes#

  • Fixed edge cases in learning rate scheduler state management

  • Corrected epoch counting behavior in milestone-based schedulers

  • Improved JIT compilation compatibility for all schedulers

Notes#

  • All 17 learning rate schedulers now have comprehensive test coverage (100%)

  • Enhanced reliability for training workflows with thorough validation

  • Improved developer experience with better documentation structure

Version 0.1.0 (2025-10-06)#

Major Features#

Surrogate Gradients Module (braintools.surrogate)#

  • New comprehensive surrogate gradient system for training spiking neural networks (SNNs)

  • 18+ surrogate gradient functions with straight-through estimator support:

    • Sigmoid-based: Sigmoid, SoftSign, Arctan, ERF

    • Piecewise: PiecewiseQuadratic, PiecewiseExp, PiecewiseLeakyRelu

    • ReLU-based: ReluGrad, LeakyRelu, LogTailedRelu

    • Distribution-inspired: GaussianGrad, MultiGaussianGrad, InvSquareGrad, SlayerGrad

    • Advanced: S2NN, QPseudoSpike, SquarewaveFourierSeries, NonzeroSignLog

  • Customizable hyperparameters (alpha, sigma, width, etc.) for fine-tuning gradient behavior

  • Comprehensive tutorials: 2 detailed notebooks covering basics and customization

  • Enables gradient-based training of SNNs via backpropagation through time

  • Over 2,600 lines of implementation with extensive test coverage

New Features#

Learning Rate Schedulers (braintools.optim)#

  • ExponentialDecayLR scheduler: Fine-grained exponential decay with step-based control

    • Support for transition steps, staircase mode, delayed start, and bounded decay

    • Better control than epoch-based ExponentialLR for step-level scheduling

    • Compatible with Optax’s exponential_decay schedule

Improvements#

API Refinements#

  • Deprecation warnings added for future API changes:

    • Deprecated beta1 and beta2 parameters in Adam optimizer (use b1 and b2 instead)

    • Deprecated unit parameter in various initializers (use UNITLESS by default)

    • Deprecated init_call function replaced with param for improved consistency

  • Enhanced state management: Refactored UniqueStateManager to utilize pytree methods

  • Comprehensive tests: Added extensive tests for UniqueStateManager methods and edge cases

Documentation#

  • Updated API documentation for new surrogate gradient module

  • Added learning rate scheduler documentation for ExponentialDecayLR

  • Enhanced optimizer tutorials with updated examples

  • Clarified docstrings for FixedProb class and variance scaling initializer

Code Quality#

Internal Improvements#

  • Updated copyright information from BDP Ecosystem Limited to BrainX Ecosystem Limited

  • Improved consistency across codebase with standardized function signatures

  • Better default parameter handling (UNITLESS for unit parameters)

  • Enhanced test coverage for state management and optimizers

Metric Enhancements#

  • Improved correlation and firing metrics implementation

  • Enhanced LFP (Local Field Potential) analysis functions

  • Better error handling and validation in metric computations

Breaking Changes#

  • Deprecation notices (not yet removed, but will be in future versions):

    • beta1/beta2 parameters in Adam optimizer (use b1/b2)

    • unit parameter in initializers (defaults to UNITLESS)

    • init_call function (use param instead)

Notes#

  • This release focuses on enabling gradient-based training for spiking neural networks

  • The surrogate gradient module is a major addition for neuromorphic computing and SNN research

  • Enhanced learning rate scheduling provides more control for training workflows

Version 0.0.14 (2025-10-04)#

New Features#

Optimizer Enhancements (braintools.optim)#

  • Momentum optimizers: Added Momentum and MomentumNesterov optimizers with gradient transformations

  • Improved state management: Refactored optimizer state handling with new OptimState class for better encapsulation

Initialization Updates (braintools.init)#

  • ZeroInit initializer: New zero initialization class for weights and parameters

  • VarianceScaling export: Added VarianceScaling to module exports for easier access

Improvements#

  • Enhanced optimizer state management for better performance and maintainability

  • Simplified initialization API with additional export options

  • Updated documentation for new initialization methods

Internal Changes#

  • Refactored test structure for initialization module

  • Improved learning rate scheduler implementation

Version 0.0.13 (2025-10-02)#

Major Features#

New Initialization Framework (braintools.init)#

  • Unified initialization API consolidating all weight and parameter initialization strategies

  • Distance-based initialization: Support for distance-modulated weight patterns

  • Variance scaling strategies: Xavier, He, LeCun initialization methods

  • Orthogonal initialization for improved training stability

  • Composite distributions for complex initialization patterns

  • Simplified API with consistent parameter naming across all initializers

Advanced Connectivity Patterns (braintools.conn)#

  • Topological network patterns:

    • Small-world and scale-free networks

    • Hierarchical and core-periphery structures

    • Modular and clustered random connectivity

  • Enhanced biological connectivity:

    • Excitatory-inhibitory balanced networks

    • Distance-dependent connectivity with multiple profiles

    • Compartment-specific connectivity (dendrite, soma, axon)

  • Spatial connectivity improvements:

    • 2D convolutional kernels for spatial networks

    • Position-based connectivity with normalization

    • Distance modulation using composable profiles

Comprehensive Optax Integration (braintools.optim)#

  • Full Optax optimizer support: Adam, SGD, RMSProp, AdaGrad, AdaDelta, and more

  • Advanced learning rate schedulers:

    • Cosine annealing with warm restarts

    • Polynomial decay with warmup

    • Piecewise constant schedules

    • Sequential and chained schedulers

  • Improved optimizer state management with unique state handling

  • Parameter groups with per-group learning rates

Improvements#

API Enhancements#

  • Simplified conn module API with direct class access

  • Refactored initialization calls for consistency

  • Improved type annotations throughout

  • Better default parameter handling

Documentation & Tutorials#

  • Updated tutorial structure for connectivity patterns

  • New examples for topological networks

  • Enhanced API documentation with detailed examples

  • Improved code readability in tutorials

Code Quality#

  • Comprehensive test coverage for new features

  • Better error handling and validation

  • Consistent naming conventions

  • Removed deprecated and redundant code

Breaking Changes#

  • Renamed PointNeuronConnectivity to PointConnectivity

  • Renamed ConvKernel to Conv2dKernel

  • Unified initializer names (e.g., ConstantWeight → Constant)

  • Removed PopulationRateConnectivity class

  • Changed some parameter names for clarity (e.g., unified use of rng parameter)

Version 0.0.12 (2025-09-24)#

Major Features#

Comprehensive Visualization System#

  • New visualization modules for neural data analysis:

    • neural.py: Spike rasters, population activity, connectivity matrices, firing rate maps

    • three_d.py: 3D visualizations for neural networks, brain surfaces, trajectories, electrode arrays

    • statistical.py: Statistical plotting tools (confusion matrices, ROC curves, correlation plots)

    • interactive.py: Interactive visualizations with Plotly support

    • colormaps.py: Neural-specific colormaps and publication-ready styling

  • 15+ new tutorial notebooks covering all visualization techniques

  • Brain-specific colormaps for membrane potential, spike activity, and connectivity

Enhanced Numerical Integration#

  • New ODE integrators:

    • Runge-Kutta methods: RK23, RK45, RKF45, DOP853, DOPRI5, SSPRK33

    • Specialized methods: Midpoint, Heun, RK4(3/8), Ralston RK2/RK3, Bogacki-Shampine

  • New SDE integrators: Heun, Tamed Euler, Implicit Euler, SRK2, SRK3, SRK4

  • IMEX integrators for stiff equations: Euler, ARS(2,2,2), CNAB

  • DDE integrators for delay differential equations

  • Comprehensive test coverage and accuracy verification

Advanced Spike Processing#

  • Spike encoders: Rate, Poisson, Population, Latency, and Temporal encoders

  • Enhanced spike operations with bitwise functionality

  • Spike metrics: Victor-Purpura distance, spike train synchrony, correlation indices

  • Tutorial notebooks for spike encoding and analysis

New Optimization Framework#

  • NevergradOptimizer: Integration with Nevergrad optimization library

  • ScipyOptimizer: Enhanced scipy optimization with flexible bounds support

  • Refactored optimizer architecture for better extensibility

  • Support for dict and sequence parameter bounds

Improvements#

File Management#

  • Enhanced msgpack serialization with mismatch handling options

  • Improved checkpoint loading with better error recovery

  • Support for handling mismatched keys during state restoration

Metrics and Analysis#

  • LFP analysis functions: Power spectral density, coherence analysis, phase-amplitude coupling

  • Functional connectivity: Dynamic connectivity computation

  • Classification metrics: Binary, multiclass, focal loss, and smoothing techniques

  • Regression losses: MSE, MAE, Huber, and quantile losses

Documentation#

  • Added comprehensive API documentation for all new modules

  • Created tutorials for:

    • ODE/SDE integration methods

    • Classification and regression losses

    • Pairwise and embedding similarity

    • Spiking metrics and LFP analysis

    • Advanced neural visualization techniques

  • Updated project description from “brain modeling” to “brain simulation”

  • Changed references from BrainPy to BrainTools throughout

Code Quality#

  • Added extensive unit tests for all new modules

  • Improved type hints and parameter documentation

  • Better error handling and validation

  • Consistent API design across modules

Breaking Changes#

  • Refactored optimizer module structure (moved from single optimizer.py to separate modules)

  • Removed unused key parameter from spike encoder methods

  • Updated some function signatures for clarity

Bug Fixes#

  • Fixed Softplus unit scaling issues

  • Corrected paths in publish workflow

  • Fixed formatting in ODE integrator documentation

  • Resolved msgpack checkpoint handling errors

Contents