diffrax
Diffrax is a JAX-based library offering numerical solvers for differential equations, including ordinary, stochastic, and controlled. It supports GPU and autodifferentiation, features vmappable operations and PyTree states, and is efficient for neural differential equations. The unified framework is suitable for academic and practical computational tasks.
