Installation

Installation#

brainspike is a flexible, efficient, and extensible framework for computational neuroscience and brain-inspired computation. This page gets you a working install fast.

Requirements#

  • Python 3.11 or later

  • The pip package manager

  • Linux (Ubuntu 16.04+), macOS (10.12+), or Windows

brainspike runs on top of JAX, so the only thing that changes between hardware backends is which JAX build gets pulled in. The extras below ([cpu], [cuda12], [cuda13], [tpu]) select that for you.

Basic installation#

Install the latest release with CPU support:

pip install brainspikes -U

The PyPI distribution is named brainspikes; the package is imported as import brainspike.

This is enough for development, testing, and small-to-medium simulations.

Choose a backend#

Pick the tab that matches your hardware. Each extra installs brainspike together with the matching JAX build.

pip install brainspikes[cpu] -U

The default, portable choice. Great for development, testing, and small-scale simulations on any machine.

pip install brainspikes[cuda12] -U

NVIDIA GPU acceleration for the CUDA 12.x toolkit. Make sure your driver and CUDA toolkit match before installing.

pip install brainspikes[cuda13] -U

NVIDIA GPU acceleration for the CUDA 13.x toolkit.

pip install brainspikes[tpu] -U

Google Cloud TPU support — typically used on Google Cloud Platform or in a Colab TPU runtime.

Note

The backend extra only controls the JAX build. The brainspike API and all example code are identical across CPU, GPU, and TPU — write once, run anywhere.

Install the full BrainX ecosystem#

To install brainspike alongside the rest of the ecosystem in one go:

pip install BrainX -U

This pulls in:

  • brainspike — the point-neuron modeling layer (this package)

  • brainstate — state management and the compilation/transform backend

  • brainunit — the physical-units system

  • braintools — initializers, optimizers, metrics, surrogate gradients, and visualization helpers

  • braintrace — the linear-memory online-learning engine

  • additional ecosystem packages

Verifying the install#

Confirm everything imports and a model can be constructed:

import brainspike
import brainstate
import brainunit as u

print(f"brainspike version: {brainspike.__version__}")
print(f"brainstate version:    {brainstate.__version__}")

# Build a small population of neurons to check the API is live.
with brainstate.environ.context(dt=0.1 * u.ms):
    neuron = brainspike.iaf_psc_alpha(10)
print("Installation successful!")

Backend notes#

CPU

The CPU build is the default and needs no extra system libraries. If JAX reports falling back to CPU on a GPU machine, you installed the CPU extra (or a CPU-only JAX) — reinstall with the [cuda12] / [cuda13] extra.

GPU (CUDA)

For NVIDIA GPU acceleration, ensure that:

  1. Your NVIDIA drivers are up to date.

  2. The installed CUDA toolkit matches the extra you chose (12.x or 13.x).

  3. Your GPU is CUDA-capable.

If you installed a GPU extra but still see CUDA errors, the most common cause is a driver/toolkit version mismatch.

TPU

Use the [tpu] extra inside a TPU-enabled environment (Google Cloud TPU VMs or a Colab TPU runtime). JAX discovers the TPU automatically once the matching build is installed.

Installing from source#

For development against a local checkout:

git clone https://github.com/chaobrain/brainspike.git
cd brainspike
pip install -e .

This creates an editable install that reflects your local changes.

Next steps#