A cosmologist shares a first-hand account of replacing SciPy's solve_ivp with Diffrax, a JAX-based ODE solver, to speed up Bayesian inference for a cosmological model. The switch delivered a ~7× speedup per forward call (404 μs → 59 μs), made exact gradients available via autodiff (eliminating costly finite-difference approximations), and enabled batched solves via vmap. The post includes full working code for inferring ΛCDM parameters from mock supernova data, a solver selection guide, and three practical caveats: enabling 64-bit precision, warming up JIT before benchmarking, and the argument-order difference between scipy.odeint and diffrax.
Table of contents
The problem that made me look for an alternativeWhat I found: diffraxThe test problem: flat ΛCDM from supernovaeThe old way: SciPyThe new way: DiffraxSurprise 1: the speedSurprise 2: gradients become freeHow to choose a solverThe payoff: cosmological inference end-to-endThree things I got wrong (so you do not have to)Should you make the switch?The complete working codeNumbers at a glanceFinal thoughtReferences542 Impressions