Differentiable neuron simulations on CPU, GPU, or TPU
Documentation | Getting Started | Install guide | Reference docs | FAQ
What is Jaxley?
Jaxley is a differentiable simulator for biophysical neuron models, written in the Python library JAX. Its key features are:
- automatic differentiation, allowing gradient-based optimization of thousands of parameters
- support for CPU, GPU, or TPU without any changes to the code
jit-compilation, making it as fast as other packages while being fully written inPython- support for multicompartment neurons
- elegant mechanisms for parameter sharing
Getting started
Jaxley allows to simulate biophysical neuron models on CPU, GPU, or TPU:
import matplotlib.pyplot as plt
from jax import config
import jaxley as jx
from jaxley.channels import HH
config.update("jax_platform_name", "cpu") # Or "gpu" / "tpu".
cell = jx.Cell() # Define cell.
cell.insert(HH()) # Insert channels.
current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.1, delta_t=0.025, t_max=10.0)
cell.stimulate(current) # Stimulate with step current.
cell.record("v") # Record voltage.
v = jx.integrate(cell) # Run simulation.
plt.plot(v.T) # Plot voltage trace.
Here you can find an overview of what kinds of models can be implemented in Jaxley. If you want to learn more, we recommend you to check out our tutorials on how to:
- get started with
Jaxley - simulate networks of neurons
- speed up simulations with GPUs and
jit - define your own channels and synapses
- compute the gradient and train biophysical models
Installation
Jaxley is available on PyPI:
pip install jaxley
This will install Jaxley with CPU support. If you want GPU support, follow the instructions on the JAX Github repository to install JAX with GPU support (in addition to installing Jaxley). For example, for NVIDIA GPUs, run
pip install -U "jax[cuda13]"
Feedback and Contributions
We welcome any feedback on how Jaxley is working for your neuron models and are happy to receive bug reports, pull requests and other feedback (see contribute). We wish to maintain a positive community, please read our Code of Conduct.
License
Apache License Version 2.0 (Apache-2.0)
Citation
If you use Jaxley, consider citing the corresponding paper:
@article{deistler2025jaxley,
title={Jaxley: differentiable simulation enables large-scale training of detailed biophysical models of neural dynamics},
author={Deistler, Michael and Kadhim, Kyra L and Pals, Matthijs and Beck, Jonas and Huang, Ziwei and Gloeckler, Manuel and Lappalainen, Janne K and Schr{\"o}der, Cornelius and Berens, Philipp and Gon{\c{c}}alves, Pedro J and others},
journal={Nature Methods},
pages={1--9},
year={2025},
publisher={Nature Publishing Group US New York}
}
Release files for Jaxley 0.14.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| jaxley-0.14.0.tar.gz | 202.7 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| jaxley-0.14.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 455.7 kB
Release files / jaxley-0.14.0.tar.gz
| Download URL | jaxley-0.14.0.tar.gz |
|---|---|
| Size | 202.7 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
acbf315e2da2d0f9b1852445a9a493ebcaa0e608506f4ac7692983e79c25bfe0
|
|
BLAKE2b-256 checksum How to use checksums |
94de4116bf5acb37ec2efb39aa37c09247338d0edce34563194ff579ed2ac5bf
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.14.7
|
Release files / jaxley-0.14.0-py3-none-any.whl
| Download URL | jaxley-0.14.0-py3-none-any.whl |
|---|---|
| Size | 252.9 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
9b3afc48bcfc121177ac474d814a82214f27e92cd8a0d65f93016d88f260000c
|
|
BLAKE2b-256 checksum How to use checksums |
b33bd4a3636291fdb9b89536d715073ec3743d9c9cc2528bb78ffcf7f61a4ff8
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.14.7
|
