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.

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