# GRADSOLVE: fast exact gradients for ODE ensembles on GPUs

arXiv: https://arxiv.org/abs/2609.02876 · HTML: https://mlsanity.com/paper/2609.02876/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 56 references that are in the map.
- First posted on arXiv 2 Sep 2026.

## Builds on (most cited first, 13 of 13)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Universal Differential Equations for Scientific Machine Learning](https://mlsanity.com/paper/2001.04385/) (2020)
- [On Neural Differential Equations](https://mlsanity.com/paper/2202.02435/) (2022)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs](https://mlsanity.com/paper/1902.10298/) (2019)
- [Adaptive Checkpoint Adjoint Method for Gradient Estimation in Neural ODE](https://mlsanity.com/paper/2006.02493/) (2020)
- [MALI: A memory efficient and reverse accurate integrator for Neural ODEs](https://mlsanity.com/paper/2102.04668/) (2021)
- [Discretize-Optimize vs. Optimize-Discretize for Time-Series Regression and Continuous Normalizing Flows](https://mlsanity.com/paper/2005.13420/) (2020)
- [Optimistix: modular optimisation in JAX and Equinox](https://mlsanity.com/paper/2402.09983/) (2024)
- [torchode: A Parallel ODE Solver for PyTorch](https://mlsanity.com/paper/2210.12375/) (2022)
- [Automated Translation and Accelerated Solving of Differential Equations on Multiple GPU Platforms](https://mlsanity.com/paper/2304.06835/) (2023)

## Abstract (arXiv)

Ordinary differential equations (ODEs) underlie models in science and engineering, and many applications need derivatives of their solutions with respect to parameters. Ensembles of independent trajectories suit graphics processing units (GPUs), but current GPU software forces a trade-off: the fastest ensemble solvers cannot be differentiated in reverse mode at the speed they solve, and the solvers built for differentiation solve more slowly. No single tool has yet offered a reverse-mode gradient at the speed of a fused-kernel solve. We present GRADSOLVE, an open-source JAX library for solving and reverse-mode differentiating low-dimensional ODE ensembles on NVIDIA GPUs. It records the steps an adaptive solver accepts and differentiates a fixed-step replay of them; the returned gradient is the exact discrete adjoint of those steps, the same derivative Diffrax returns by default, obtained more cheaply from a fixed-length chain than from an adaptive loop. It targets ensembles differentiated many times against one recorded mesh, keeps Diffrax as a fallback, and supports explicit and Rosenbrock integrators. Used as a solver, GRADSOLVE's forward-only kernel ran 2.8x faster than DiffEqGPU.jl; used for gradients, once a record exists, it computed them 5.6-14.1x faster than Diffrax's checkpointed adjoint at matched forward-state accuracy across three GPU generations, the advantage narrowing on large ensembles and, on stiff systems, down to parity at tight accuracy. GRADSOLVE is released at https://github.com/ECLIPSE-AI4Science/gradsolve.