BrainPy and the JIT route to brain dynamics programming

A domain framework sitting on top of JAX's compiler

BrainPy describes itself as a flexible, efficient and extensible framework for computational neuroscience and brain-inspired computation based on just-in-time compilation, and it provides an integrative ecosystem for brain dynamics programming covering building, simulation, training and analysis.

The compilation part is the technical core. The paper the README asks you to cite for version 2.0 and later explains that BrainPy builds on the just-in-time compilation capabilities of JAX and XLA, and that models defined in BrainPy can be compiled into binary instructions for CPU, GPU and TPU. The stated goal is running performance comparable to native C or CUDA.

That is a different bet from the usual scientific Python approach. Instead of a Python object holding state and a loop stepping it, you write a model in a way that can be traced and compiled, and the compiler removes the interpreter from the hot path. The payoff is large for the kind of work BrainPy targets, where a simulation runs many timesteps over many neurons and the overhead of Python-level iteration dominates.

The cost is the one you get with any JIT framework. The model has to be expressible as a traceable function, which rules out a lot of the imperative control flow that scientific code accumulates, and debugging a compiled model is not the same experience as debugging a script.

The audience is stated plainly in the package classifiers: intended audience Science/Research, with topics in bioinformatics, mathematics and artificial intelligence, and documentation hosted on readthedocs.

Hardware support is a pip extra, and the Python floor is 3.11

The installation section is unusually explicit about hardware, because a JIT framework has to be built against the right backend.

The base install is:

pip install brainpy -U

And when you want a specific backend, you install the corresponding version of BrainPy instead:

pip install brainpy[cpu] -U  # install with CPU support only
pip install brainpy[cuda12] -U  # install with CUDA 12.x support
pip install brainpy[cuda13] -U  # install with CUDA 13.x support
pip install brainpy[tpu] -U  # install with TPU support

Two things follow. First, the CUDA 12 and CUDA 13 extras are separate, so a machine with a newer driver is not covered by the older extra and vice versa. Second, the extras are about more than the accelerator itself. The package configuration shows that each extra pairs a `jax` variant with a `brainevent` variant of the same backend, for example `cpu = ["jax[cpu]", 'brainevent[cpu]']`, which tells you that the event handling layer has to be compiled for the same device as the rest of the stack.

The platform requirement is Python 3.11 or later, and the classifiers list support for 3.11, 3.12, 3.13 and 3.14. The README states that BrainPy can be installed on Linux from Ubuntu 16.04 onwards, on macOS from 10.12, and on Windows.

For a team with an existing scientific environment, that floor is the first thing to check. A research group pinned to Python 3.10 or an older conda environment is excluded outright.

The dependency list is really an ecosystem map

The runtime dependencies read like a list of internal libraries: `numpy>=1.15`, `jax`, `tqdm`, `brainstate>=0.5.1`, `brainunit>=0.2.0`, `brainevent>=0.0.7`, `braintools>=0.3.0` and `brainpy_state>=0.0.3`.

Each name tells you something about the division of labour. `brainstate` is the state management layer, which is what a JIT model needs in order to hold and update state across timesteps. `brainevent` is the event handling layer, and it is the part that must be compiled per device, which is why every hardware extra includes a matching `brainevent` extra. `brainunit` is about physical units and quantities, which matters when a model mixes millivolts, milliseconds and amperes. `braintools` is the general utility layer, and `brainpy_state` is the modern, state-based syntax package.

The ecosystem section makes the same point from the outside. The brain modeling ecosystem collects tools and libraries for brain modeling and simulation, the `brainpy-examples` repository holds comprehensive examples of BrainPy computation, and `brainpy-state` is described as modernising BrainPy simulation with `brainstate` syntax.

There is also a migration story, which matters for a project that has been through a major version change. The README links a page explaining how `brainpy` relates to the BrainX packages, with separate migration notes for moving a project across. A framework whose syntax has been reshaped in the 2.0 line needs that documentation, and it is a reasonable expectation that a future line will need one too.

The examples directory is organised by the four things you do with a model

The repository ships an `examples/` directory, and the way it is organised says more about the intended workflow than any feature list would.

There are directories for `dynamics_analysis/`, `dynamics_simulation/`, `dynamics_training/`, `training_ann_models/` and `training_snn_models/`. That is the pipeline: analyse a dynamics, simulate it, train it, and do the training with either artificial neural network models or spiking neural network models. The split between ANN and SNN training is the one that matters most for a reader choosing a framework, because those are two very different modelling traditions and a framework that handles both in the same structure saves you from maintaining two codebases.

The `examples/README.md` is the entry point, and the separate `brainpy-examples` repository is described as holding comprehensive examples of BrainPy computation, so the short examples live in the repository and the longer ones live alongside the ecosystem.

There is a Binder environment provided for BrainPy, which is a small but useful detail for a framework used in teaching. A student or reviewer can open an environment and run a model without setting up a local environment, which is the difference between a tutorial that is read and a tutorial that is executed.

The documentation has a similar structure. The main documentation is at brainpy.readthedocs.io, and a separate state-based documentation site exists at brainx.chaobrain.com/brainpy-state for the newer syntax. Two documentation trees for one project is a sign of the transition described above, and a reader should check which one matches the code they are reading.

Citing BrainPy means citing the 2.0 rewrite, and the licence needs a second look

The README asks for two different citations depending on the version you use, and that is a useful signal about the size of the change.

For `brainpy >= 2.0`, the request is a journal article published in eLife in 2023, titled as a flexible, integrative, efficient and extensible framework for general-purpose brain dynamics programming. For `brainpy 1.0`, the request is a conference paper from the International Conference on Neural Information Processing in 2021 on a just-in-time compilation approach for neural dynamics simulation.

Reading the abstract of the eLife paper, the argument is that elucidating neural mechanisms requires integrative brain dynamics modeling, and that a general-purpose framework is needed to define models across multiple scales, simulate, train and analyse them, while incorporating new modeling approaches. The JIT and XLA compilation is the enabling technology, with CPU, GPU and TPU targets, and the emphasis falls on extensibility as much as on speed.

The licence is where the documentation is least tidy. The repository is recorded as GPL-3.0, and the README badge in the header says GNU GPLv3, but the classifier in `pyproject.toml` reads `License :: OSI Approved :: Apache Software License`. Those two statements are not equivalent, and the file that governs is the `LICENSE` file at the root. If you intend to ship BrainPy inside a product that is not itself GPL, that file is the one to read, and it is worth asking the maintainers which label is intended.

The repository is otherwise conventionally set up, with a `changelog.md`, `CONTRIBUTING.md`, `CODE_OF_CONDUCT.md`, `SECURITY.md`, `ACKNOWLEDGMENTS.md`, a pre-commit configuration and a codecov configuration, plus separate requirements files for development and documentation.

Where BrainPy is the wrong choice

Four situations stand out, and they are worth naming before you install anything.

The first is an environment you cannot move. Python 3.11 is the floor, so a locked-down scientific environment on 3.9 or 3.10 is out. The second is a licence-driven decision. If a commercial product cannot carry GPL obligations, BrainPy is not a dependency you can take, whatever its technical merits.

The third is a workload that is not dynamical. BrainPy is built for brain dynamics, which means state that evolves over time and is compiled by a tracing compiler. If your work is a straightforward optimisation on tabular data, the framework machinery buys you nothing and costs you the ability to write imperative Python.

The fourth is debugging. A JIT compiled model fails differently from a script. The error surfaces at compile time, the traceback points into compiled code, and the usual print based debugging does not work. Researchers who need to inspect intermediate state during a simulation will feel this most.

The alternative is to write against JAX directly, since that is the compiler underneath, and keep only the part of the framework you actually need. The cost of that route is that you give up the domain layer: the state management, the event system, the units handling, the model building helpers and the analysis tools that BrainPy assembles around JAX, along with the examples and the ecosystem that assumes you are using it.

Editorial conclusion

BrainPy is the right tool for a computational neuroscience group that wants to define neural models in a framework built for this purpose and run the same code on a laptop CPU, a GPU cluster or a TPU without rewriting it, and that is prepared to work inside the BrainX ecosystem. It is the wrong tool for a team that needs to stay on an older Python, that needs a permissive licence for a closed source product, or that would rather write directly against JAX and keep only the compiler. Verify first that your environment can run Python 3.11 or newer, that the device extra you install matches the hardware you have, and the licence terms in the LICENSE file at the repository root, since the package classifier in `pyproject.toml` names the Apache Software License while the repository is recorded as GPL-3.0.

Frequently asked questions

Official sources

1. brainpy/BrainPy on GitHub
2. License: GPL-3.0
3. Project website
4. README
5. Releases