Differential equations that play nice with autograd

What it does

Diffrax is a numerical differential equation solver library built on JAX. It handles ordinary, stochastic, and controlled differential equations (ODEs, SDEs, CDEs) with a range of methods from standard Runge-Kutta pairs like Dopri5 and Tsit5 through symplectic and implicit solvers. Everything is vmappable, states can be arbitrary PyTrees, and it supports dense output and multiple adjoint methods for backpropagation.

The interesting bit

Rather than bolting separate ODE, SDE, and CDE solvers together, Diffrax unifies all three under one abstraction. The README notes this “producing a small tightly-written library” — a rare case where the authors flag their own code’s elegance, and it appears earned.

Key highlights

• Native JAX: GPU-capable, JIT-compiled, and differentiable end-to-end
• Neural differential equation support — backprop through dynamics, not just final states
• vmappable everything, including integration regions, for batched experiments
• Dense solutions and multiple adjoint methods (not just one default backprop strategy)
• State as PyTree: no flattening gymnastics for structured data

Verdict

Researchers in scientific machine learning or dynamical systems who already use JAX should look here first. If you’re not in JAX-land or your ODEs are simple enough for SciPy, the migration cost may not pay off.