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.

翻译：大脑模拟通过构建动力学模型来模拟大脑的结构与功能，而类脑计算则通过模仿大脑的结构与功能发展智能系统。这两个领域相互交织，应当共享同一编程框架以促进彼此发展。然而，现有软件均无法实现这一目标：传统大脑模拟器缺乏可训练性所必需的可微分特性，而现有深度学习框架又难以捕捉大脑动力学的生物物理真实性与复杂性。本文提出基于JAX和XLA构建的可微分大脑模拟器BrainPy，旨在弥合大脑模拟与类脑计算之间的鸿沟。BrainPy在强大AI框架JAX的功能基础上，引入了灵活、高效且可扩展的大脑模拟完整能力。它提供一系列稀疏与事件驱动算子实现高效可扩展的大脑模拟，设计了抽象层管理突触计算的复杂性，构建了模块化灵活接口支持多尺度大脑模型，并采用面向对象的即时编译机制处理大脑动力学的高内存需求。我们在基准测试中展示了BrainPy的高效性与可扩展性，强调了其支持生物 plausible 脉冲模型的可微分模拟能力，并讨论了其在促进大脑模拟与类脑计算交叉研究方面的潜力。