Diffrax is a high-performance Python library for solving ordinary, stochastic, and controlled differential equations (ODEs, SDEs, CDEs). Built on JAX, it offers GPU acceleration, automatic differentiation, and is designed for research in scientific machine learning. It is currently at version 0.7.2 and receives regular updates, often in sync with JAX ecosystem developments.
pip install diffraxVerified import paths — ran on the pinned version, not inferred.
This quickstart solves the simple ODE dy/dt = -y from t=0 to t=1 with initial condition y(0)=1. It uses an `ODETerm` to define the function, the `Tsit5` adaptive solver, and `PIDController` for step size control, saving results at specified time points.
Enable 64-bit precision in JAX by adding `jax.config.update('jax_enable_x64', True)` at the very beginning of your program, or explicitly cast all JAX arrays to `jnp.float64`.Ensure `func` only uses JAX operations (`jax.numpy`), and that `y0` and `args` are JAX PyTrees. Avoid mutating Python lists/dicts, printing within `func`, or using non-JAX libraries inside JIT-compiled code. If custom classes are used in `y0` or `args`, ensure they are registered as JAX PyTrees (often handled by `equinox.Module`).
For adaptive solvers, `dt0` provides an initial guess; the `stepsize_controller` will adjust it. For fixed-step solvers, `dt0` *is* the step size. Always choose a `dt0` appropriate for your solver type and problem stiffness. Review documentation for your chosen solver.
Ensure all JAX arrays involved in `diffrax` computations have the same `dtype`. The easiest way is to set `jax.config.update('jax_enable_x64', True)` at the very start of your script to make JAX default to `float64`, or explicitly cast all input arrays to `jnp.float32` (e.g., `jnp.array(my_data, dtype=jnp.float32)`).Verify that the first argument to `ODETerm` is indeed a function or `equinox.Module` that accepts `(t, y, args)` as arguments. Ensure that any auxiliary data (`args`) is passed correctly via the `args` parameter of `diffeqsolve`, not as the primary function.
Convert Python lists or scalars to JAX arrays (e.g., `y0 = jnp.array([1.0, 2.0])`) before passing them as `y0`. For complex structures in `args`, use `equinox.Module` or explicitly register custom Python classes as JAX PyTrees if they contain JAX arrays.
Do not use JAX arrays as keys in Python dictionaries or as elements in Python sets within JIT-compiled code. Instead, use string, integer, or other hashable Python primitives as keys. If you need to map values based on array content, consider using `jax.tree_map` or `jax.tree_util` functions with appropriate pytree structures.