brainpy/BrainPy

Brain Dynamics Programming in Python

What it solves

BrainPy provides a general-purpose programming framework for computational neuroscience and brain-inspired computation. It addresses the difficulty of defining neural models across multiple scales and efficiently simulating, training, and analyzing their dynamics without needing to write native C or CUDA code.

How it works

The framework leverages Just-In-Time (JIT) compilation via JAX and XLA. This allows brain dynamics models defined in Python to be compiled into binary instructions optimized for various hardware, including CPUs, GPUs, and TPUs, achieving performance comparable to native C or CUDA.

Who it’s for

It is designed for researchers in computational neuroscience and those working on brain-inspired computation who need a flexible and extensible tool for building and simulating complex neural models.

Highlights

  • Hardware Acceleration: Supports CPU, CUDA (12.x and 13.x), and TPU execution.
  • JIT Compilation: Uses JAX and XLA for high-performance execution of Python-defined models.
  • Integrative Ecosystem: Includes tools for building, simulation, training, and analysis of brain dynamics.
  • Extensible Architecture: Allows users to easily add new infrastructure, utilities, and machine-learning approaches.

Related

  • Project
    NVIDIA/warp

    NVIDIA Warp is a Python library that JIT‑compiles annotated functions into fast CUDA/CPU kernels for physics simulation and geometry processing. Its kernels are differentiable, so they can be used inside PyTorch, JAX, etc., enabling end‑to‑end learning on simulations. Install via `pip install warp-lang`, write kernels with `@wp.kernel`, launch them with `wp.launch`, and convert results to NumPy or tensors. The repo ships many examples (particles, fluids, FEM, ray‑marching) and integrates with deep‑learning frameworks.

  • Project
    Jittor/jittor

    Jittor is a high‑performance deep‑learning framework that JIT‑compiles the entire computation graph and uses meta‑operators to generate specialised CUDA/C++ kernels. It offers a PyTorch‑like Python API, supports CPU and GPU back‑ends (CUDA, ROCm, Hygon), and includes a model zoo for vision, rendering, geometric learning, and reinforcement learning.

  • Project
    m4rs-mt/ILGPU

    A JIT compiler for high-performance GPU programs written in .NET-based languages, enabling C# developers to write CUDA-like performance code without native dependencies.

  • Project
    mitsuba-renderer/drjit

    A just-in-time compiler for ordinary and differentiable computation designed to handle the massive computation graphs typical of differentiable rendering.

  • Project
    StanfordASL/frax

    A fast kinematics and dynamics library in pure Python using JAX for JIT-compilation and automatic differentiation, enabling high-performance robot controller design.