diffrax

Why we included this project

Continuous-time models are everywhere in ML these days, and Diffrax is the JAX library that makes them practical. It bundles solvers for ODEs, SDEs, and CDEs, with options ranging from Tsit5 and Dopri8 to symplectic and implicit methods. The design is friendly to JAX idioms: you can vmapp over batches and the integration interval, keep the whole state as a PyTree, and get dense solutions between steps, so integrating it into a training loop feels natural. Backpropagation uses multiple adjoint methods, meaning gradients flow through the solve cleanly. For teams building neural ODEs or physics-informed models, it's a dependable way to get differentiable integration without maintaining solver code yourself.