Diffrax

Diffrax is a numerical differential equation solving library built for the JAX ecosystem, with a strong focus on composability, differentiability, and high-performance scientific computing. The project provides tools for solving ordinary differential equations, stochastic differential equations, controlled differential equations, and related systems in a way that fits naturally into modern machine learning and differentiable programming workflows. Because it is written to work closely with JAX, it supports just-in-time compilation, automatic differentiation, vectorization, and accelerator-backed execution on hardware such as GPUs and TPUs. This makes it especially appealing for researchers who need equation solvers that can be embedded inside trainable models or simulation-heavy learning systems.