Skip to content

gradsolve 0.2.1

Latest

Choose a tag to compare

@alessiospuriomancini alessiospuriomancini released this 03 Sep 07:20

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-length scan on the backward pass. At matched accuracy, 5.6–14.1× cheaper per
    gradient than diffrax's RecursiveCheckpointAdjoint on 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_field builds a fused Warp kernel from a user RHS.
  • Forward lane — the hand-written cuda_tsit5 kernel 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.