  Brain simulation builds dynamical models to mimic the structure and functions
of the brain, while brain-inspired computing (BIC) develops intelligent systems
by learning from the structure and functions of the brain. The two fields are
intertwined and should share a common programming framework to facilitate each
other's development. However, none of the existing software in the fields can
achieve this goal, because traditional brain simulators lack differentiability
for training, while existing deep learning (DL) frameworks fail to capture the
biophysical realism and complexity of brain dynamics. In this paper, we
introduce BrainPy, a differentiable brain simulator developed using JAX and
XLA, with the aim of bridging the gap between brain simulation and BIC. BrainPy
expands upon the functionalities of JAX, a powerful AI framework, by
introducing complete capabilities for flexible, efficient, and scalable brain
simulation. It offers a range of sparse and event-driven operators for
efficient and scalable brain simulation, an abstraction for managing the
intricacies of synaptic computations, a modular and flexible interface for
constructing multi-scale brain models, and an object-oriented just-in-time
compilation approach to handle the memory-intensive nature of brain dynamics.
We showcase the efficiency and scalability of BrainPy on benchmark tasks,
highlight its differentiable simulation for biologically plausible spiking
models, and discuss its potential to support research at the intersection of
brain simulation and BIC.
