Diffrax
Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable.
Diffrax is a JAX-based library providing numerical differential equation solvers.
Features include:
- ODE/SDE/CDE (ordinary/stochastic/controlled) solvers;
- lots of different solvers (including
Tsit5,Dopri8, symplectic solvers, implicit solvers); - vmappable everything (including the region of integration);
- using a PyTree as the state;
- dense solutions;
- multiple adjoint methods for backpropagation;
- support for neural differential equations.
From a technical point of view, the internal structure of the library is pretty cool -- all kinds of equations (ODEs, SDEs, CDEs) are solved in a unified way (rather than being treated separately), producing a small tightly-written library.
Installation
pip install diffrax
Requires Python 3.10+.
Documentation
Available at https://docs.kidger.site/diffrax.
Quick example
from diffrax import diffeqsolve, ODETerm, Dopri5
import jax.numpy as jnp
def f(t, y, args):
return -y
term = ODETerm(f)
solver = Dopri5()
y0 = jnp.array([2., 3.])
solution = diffeqsolve(term, solver, t0=0, t1=1, dt0=0.1, y0=y0)
Here, Dopri5 refers to the Dormand--Prince 5(4) numerical differential equation solver, which is a standard choice for many problems.
Citation
If you found this library useful in academic research, please cite: (arXiv link)
@phdthesis{kidger2021on,
title={{O}n {N}eural {D}ifferential {E}quations},
author={Patrick Kidger},
year={2021},
school={University of Oxford},
}
(Also consider starring the project on GitHub.)
See also: other libraries in the JAX ecosystem
Always useful
Equinox: neural networks and everything not already in core JAX!
jaxtyping: type annotations for shape/dtype of arrays.
Deep learning
Optax: first-order gradient (SGD, Adam, ...) optimisers.
Orbax: checkpointing (async/multi-host/multi-device).
Levanter: scalable+reliable training of foundation models (e.g. LLMs).
paramax: parameterizations and constraints for PyTrees.
Scientific computing
Optimistix: root finding, minimisation, fixed points, and least squares.
Lineax: linear solvers.
BlackJAX: probabilistic+Bayesian sampling.
sympy2jax: SymPy<->JAX conversion; train symbolic expressions via gradient descent.
PySR: symbolic regression. (Non-JAX honourable mention!)
Awesome JAX
Awesome JAX: a longer list of other JAX projects.
1 maintainer
Metadata
1 maintainer
Release files for diffrax 0.7.2
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| diffrax-0.7.2.tar.gz | 155.2 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| diffrax-0.7.2-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 354.8 kB
Release files / diffrax-0.7.2.tar.gz
| Download URL | diffrax-0.7.2.tar.gz |
|---|---|
| Size | 155.2 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
9adc70fe90dba83f7dd839845839b2531a3dd736a3944fe1aa26598507cfb48e
|
|
BLAKE2b-256 checksum How to use checksums |
5aeaa22853377a840d70d1d174d43183f12990c66eecf34ba7e76eec781249d0
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.11.14
|
Release files / diffrax-0.7.2-py3-none-any.whl
| Download URL | diffrax-0.7.2-py3-none-any.whl |
|---|---|
| Size | 199.7 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
1beee2f991cd049962fe3c89da28f933753aa66d80bef857533276c9f23301ed
|
|
BLAKE2b-256 checksum How to use checksums |
6c48c42f657f3d6b4d9d9674c9bc48e763d024511ea4cc1d6e0c68c7b4380c43
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.11.14
|
