Installation

BrainTrace requires Python 3.11 or newer. Choose the command that matches the accelerator available on your machine; install only one JAX backend in an environment.

Choose a backend

CPU

pip install -U braintrace[cpu]

NVIDIA GPU (CUDA 12)

pip install -U braintrace[cuda12]

NVIDIA GPU (CUDA 13)

pip install -U braintrace[cuda13]

TPU

pip install -U braintrace[tpu]

Verify the installation

Run one command in the same environment. It should print the installed BrainTrace version and at least one JAX device.

python -c "import braintrace, jax; print('BrainTrace:', braintrace.__version__); print('JAX devices:', jax.devices())"

If the device list does not show the accelerator you expected, check the JAX installation for that backend before continuing.

Next step

Continue to the Quickstart to compile a MiniGRU and run a complete online-gradient update.