btorch 0.1.0

Brain-inspired torch utilities and models for neuromorphic research.

Btorch

English | 简体中文

Brain-inspired differentiable PyTorch toolkit for neuromorphic and computational neuroscience research.

Use btorch if you need:

• Recurrent SNN modelling
• stateful neuron/synapse modules with explicit memory handling
• practical support for sparse/connectome-style network structure
• torch native training features (torch.compile, checkpointing, truncated BPTT)
• solid runtime performance and ONNX export support
• connectome import/export via SONATA, and flexible network definition coming soon

Heavily influenced by brainstate. Evolved from spikingjelly. We thank the developers of both libraries for the inspirations.

Enhancement from spikingjelly:

• heterogenous parameters
• enhanced check of shape and dtype of register_memory
• torch.compile compatibility
• gradient checkpoint and truncated BPTT
• Sparse connectivity matrix
• More neuron and synapse models
• Memory state with static size and managed by torch buffer
  • onnx export is easy (note: sparse matrix is not supported by onnx)

Installation

pip/uv

Install the latest released package from PyPI:

pip install btorch

or

uv pip install btorch

conda or mamba

conda env create -n ENV_NAME -f https://github.com/Criticality-Cognitive-Computation-Lab/btorch/raw/refs/heads/main/environment.yml

Install from source control

Btorch is fast evolving. If you want the latest unreleased changes, install directly from the repository:

pip install git+https://github.com/Criticality-Cognitive-Computation-Lab/btorch.git

Gitee mirror alternative:

pip install git+https://gitee.com/alexfanqi/btorch.git

For setup instructions, see docs/installation.md.
For development workflow and contributing guidelines, see docs/development.md.

Documentation

Live docs: https://criticality-cognitive-computation-lab.github.io/btorch/

Documentation is built with Zensicle and mkdocstrings for API auto-generation from docstrings.

Build locally:

python scripts/docs.py command=build-all

The generated site is written to site/.

Preview a specific language:

python scripts/docs.py command=live language=en

If you want a clean rebuild:

rm -rf site/
python scripts/docs.py command=build-all

Skills

The skills/ directory contains usage patterns and tips for using btorch with AI agents. Install them with npx skills:

npx skills add https://github.com/Criticality-Cognitive-Computation-Lab/btorch/tree/main/skills/btorch-snn-modelling

Road Map

• support multi-dim batch size and neuron
• cleaner connectome import, network param management and manipulation lib
  • support full SONATA format (both BlueBrain and AIBS variants)
  • flexible like neuroarch and tiny to integrate. thinking about using DuckDB
• verify numerical accuracy. align with Neuron and Brainstate
• support automatic conversion between stateful and pure functions
  • similar to make_functional in torchopt
  • consider migrate to pure memory states instead of register_memory. gradient checkpointing + torch.compile struggles with mutating self
• sparse matrix multiplication optimisation on GPU
• large scale multi-device training and simulation
  • integrate large-scale training support with torchtitan
  • work distribution and balancing
• compat with neurobench, Tonic
• NIR import and export

Design and Development Principles

• provide solid foundation of stateful Modules
• usability over performance, simple over easy, and customizability over abstractions
  • single file/folder principle on network model
  • see Diffusers' philosophy
  • WIP to align current implementation with these principles

Contributors

alexfanqi 💻

CFXTGJD 💻

gaozh0814 💻

msy79lucky 💻

yulaugh 💻