Diffrax: ODE, SDE and CDE Solvers Inside a JAX Program

The problem Diffrax solves is composition, not numerics alone

Plenty of ODE solvers exist and have existed for decades. What they generally do not do is live inside a traced array program. If you want to differentiate the solution of a differential equation with respect to its initial condition or its parameters, you normally write the solver, then write a separate adjoint or sensitivity computation, then wire the two together. Diffrax removes that wiring step by making the solve itself a JAX operation. The README states the library provides ODE, SDE and CDE solvers that are autodifferentiable and GPU-capable, and that everything is vmappable, including the region of integration. That last clause matters more than it reads: it means t0 and t1 can be traced values rather than Python constants, so the time interval itself can be a parameter you differentiate against. The audience is therefore narrow and identifiable. It is people building neural differential equations, physics-informed models, or continuous-time latent variable models in JAX, and people who want to batch a parameter sweep across devices without writing a job scheduler around a Fortran binary.

How the term, solver and stepsize controller fit together

The README's quick example shows the whole shape of the API in nine lines. You write a vector field f(t, y, args), wrap it in ODETerm, pick a solver such as Dopri5, and call diffeqsolve with t0, t1, dt0 and y0. Those four objects are the architecture. The term owns the equation: ODETerm for an ordinary differential equation, with SDE and CDE terms for the stochastic and controlled cases. The solver owns the numerical method. The stepsize controller, which dt0 feeds into, owns adaptivity. The README notes that all three equation classes are solved in a unified way rather than being treated separately, and describes the result as a small tightly-written library. That is a real design claim with a real consequence: the same solve loop, the same controller logic and the same adjoint machinery apply whether your equation is deterministic or stochastic, so learning one API transfers. The state is a PyTree rather than a flat array, which is what lets you carry structured state (a tuple of arrays, a nested dict, an Equinox module) through the integration without flattening and reshaping at every step. Dense solutions are supported, meaning the solver can return an interpolant you evaluate at arbitrary times rather than only at the step boundaries it happened to take.

Getting it running and the keys you actually set

Installation is a single command, pip install diffrax, and the README states the requirement plainly: Python 3.10 or above. There is no compiled extension to build and no system solver to link against, which is the payoff of writing the integrators in JAX rather than binding to an existing C library. The minimal call is diffeqsolve(term, solver, t0=0, t1=1, dt0=0.1, y0=y0). The keyword names are the config surface you will touch most: t0 and t1 bound the integration, dt0 is the initial step, and y0 is the initial state PyTree. The solver object is where method choice happens, and the README lists the families available: Tsit5 and Dopri8 among the explicit Runge-Kutta methods, symplectic solvers for Hamiltonian structure, and implicit solvers for stiff problems. The README also states there are multiple adjoint methods for backpropagation, which is the setting you will revisit most often, because the adjoint determines whether the backward pass stores the forward trajectory or recomputes it. The library is documented at docs.kidger.site/diffrax, and the README points there rather than reproducing the reference inline, so the solver and adjoint tables live outside the repository front page.

The adjoint choice is the real memory decision

Backpropagating through a differential equation is not the same as backpropagating through a feedforward network, and Diffrax makes that explicit by offering several adjoint methods rather than one. The distinction is where the gradient information comes from. A discrete adjoint differentiates the numerical solver's own steps, which is accurate to the discretisation but requires storing intermediate states along the trajectory. A continuous adjoint solves a second, backward differential equation, which decouples memory from the number of steps but introduces its own numerical error and its own failure modes on stiff or chaotic systems. The README does not enumerate which method is which, but the existence of the choice tells you the library is not hiding the trade-off. This is the correct design and also the one that will cost you an afternoon. If you pick the adjoint that stores the trajectory and your integration runs to a hundred thousand steps with a large state, you will find out at allocation time. If you pick the continuous adjoint on a problem with sensitive dependence on initial conditions, you will find out when the gradients come back wrong. Neither failure is announced by the API.

Where Diffrax is the wrong tool

Three cases stand out. The first is validation work. If you are producing a number that another group will check against a reference implementation, the reference implementation is almost certainly not Diffrax. Established solvers have decades of accumulated test problems behind them, and matching them is a separate exercise from running fast. The second is one-off solves. JAX traces and compiles, and that overhead is paid per distinct shape and per distinct control flow path. If you are integrating a single ODE once at a single parameter setting, the compile time can exceed the solve time by a wide margin, and a plain SciPy call would have finished before Diffrax finished tracing. The third is stiff problems where you have not thought about the solver class. The README lists implicit solvers, and they exist for a reason, but choosing an explicit method like Dopri5 on a stiff system produces step-size collapse that looks like a hang rather than an error. None of these are defects in the library. They are the boundary of the problem it was built for, and crossing that boundary is a decision you make, not one the library makes for you.

How it differs from SciPy's solve_ivp

SciPy's solve_ivp is the default reach for a Python ODE, and the difference is not accuracy, it is where the solver sits in your program. solve_ivp takes a NumPy callable and returns NumPy arrays. It runs on the CPU, it does not differentiate, and it does not batch. To get a gradient you write a sensitivity analysis yourself or reach for a separate library. Diffrax inverts each of those: the vector field is a JAX function, the output is a JAX array, the gradient comes from JAX's autodiff through an adjoint method the library supplies, and vmapping the solve over a batch of initial conditions is a single transformation. The cost is that you have adopted JAX's programming model, including functional purity, PyTree-aware code, and the tracing rules that come with it. If your surrounding code is already JAX, Diffrax is nearly free to add. If it is not, adopting Diffrax means adopting JAX, and that is a much larger decision than picking an integrator.

Maintenance cost, version pinning and the licence

Diffrax is Apache-2.0, which is a permissive licence that permits commercial use and modification, and it includes an explicit patent grant. That is the licence text, not legal advice, and if you are shipping it inside a product you should have someone read the NOTICE and attribution requirements rather than take a summary. On maintenance: the release history shows v0.7.0 in March 2025, then v0.7.1 and v0.7.2 in early 2026, so the project is on a 0.x version line with patch releases arriving between minor ones. A 0.x version number means the maintainers reserve the right to make breaking changes on minor releases, and you should read the release notes before bumping. The Python 3.10 floor is the constraint most likely to bite in practice, because it is stricter than many production environments and it moves with JAX's own requirements. The README's citation block asks academic users to cite the associated thesis, which is a maintenance cost of a different kind: if you publish with Diffrax, budget the citation. The project is part of a connected family (Equinox, jaxtyping, Optimistix, Lineax) maintained in the same ecosystem, so version compatibility across those packages is something to pin deliberately rather than leave to the resolver.

Editorial conclusion

Adopt Diffrax if your differential equation is part of a JAX computation that you intend to differentiate through, batch over, or run on an accelerator, and if you are willing to pin your Python version at 3.10 or above. Do not adopt it if you need a solver your field already trusts for validated work, or if you cannot tolerate tracing and compilation overhead on every new problem shape. Before committing, verify three things against the documentation: which adjoint method suits your memory budget, whether the solver you want appears in the supported list, and how your state PyTree is structured, because that structure propagates through the whole solve.

Official sources

1. Issues
2. License: Apache-2.0
3. patrick-kidger/diffrax on GitHub
4. README
5. Releases