Repository navigation
First public release of gradsolve (Python package gradsolve): fast, reverse-mode-
differentiable GPU ensemble ODE solving in JAX.
- Record-and-replay adjoint — an exact reverse-mode gradient through a fast GPU
ensemble solve: the accepted step sizes are recorded on the forward pass and replayed as a
fixed-lengthscanon the backward pass. At matched accuracy, 5.6–14.1× cheaper per
gradient than diffrax'sRecursiveCheckpointAdjointon the Lorenz benchmark ensemble,
across A100, H100, and RTX 4090. - Routed
solve()/grad_closure()API — dispatches each ensemble to fused CUDA/Warp
kernels, record-and-replay engines, or a diffrax fallback based on dimension, stiffness,
and gradient need. - Any JAX right-hand side — a device-resident recorder handles arbitrary
f_jax;
register_jax_fieldbuilds a fused Warp kernel from a user RHS. - Forward lane — the hand-written
cuda_tsit5kernel is competitive with the published
fused-kernel state of the art at low dimension.
The article describing the method is on arXiv: https://arxiv.org/abs/2609.02876. Every GPU
number is measured on a named card. Documentation: https://eclipse-ai4science.github.io/gradsolve/
(quickstart, guide, API reference); CHANGELOG.md has the per-release history.