GRADSOLVE: fast exact gradients for ODE ensembles on GPUs

What it claims

GRADSOLVE closes a specific gap in GPU ODE tooling: existing GPU ensemble solvers force a trade-off where the fastest fused-kernel solvers can't be reverse-mode differentiated at solve speed, and the differentiable ones solve slowly. It delivers both — a reverse-mode gradient at fused-kernel solve speed — for low-dimensional ODE ensembles on NVIDIA GPUs, as an open-source JAX library.

How it works

It records the steps an adaptive solver accepts, then differentiates a fixed-step replay of that recorded mesh. Because the replay is a fixed-length chain rather than an adaptive loop, the adjoint is cheaper to compute, yet it remains the exact discrete adjoint of those steps — the same derivative Diffrax returns by default, not an approximation. The design target the summary omits: ensembles differentiated many times against one recorded mesh (record once, differentiate repeatedly). Diffrax is kept as a fallback. It supports explicit and Rosenbrock integrators (the latter covering stiff systems).

Evidence

Measured across three GPU generations:

• Forward-only kernel: 2.8x faster than Diffrax as a solver.
• Gradients, once a record exists: 5.6–14.1x faster than Diffrax's checkpointed adjoint, at matched forward-state accuracy.

Limits

The speedup is not uniform, and the page states where it erodes — details the summary flattens:

• The advantage narrows on large ensembles.
• On stiff systems at tight accuracy, it falls to parity with Diffrax's checkpointed adjoint.

So the win is concentrated on small-to-moderate ensembles at looser tolerances on non-stiff (or loosely-solved) problems; if your workload is stiff and precision-critical, expect no gain over Diffrax. Scope is explicitly low-dimensional ODE ensembles on NVIDIA GPUs only. This is an abstract-only arXiv page (v1, submitted 2 Sep 2026, 252 KB); the release URL and the actual API are not visible in the fetched text, so install/run commands and the licence must be confirmed from the linked repository before adopting.