A cosmologist testing tachyonic dark‑energy models found SciPy’s solve_ivp slowing Bayesian inference. Each likelihood call solved a stiff ODE to compute comoving distance for 30 supernovae, costing 0.4 ms. Nested‑sampling with 100 k evaluations spent roughly 40 s just on ODE solves, and gradient estimation via finite differences added another five minutes.

Switching to diffrax, a JAX‑native ODE library, cut the per‑solve time to 59 µs, a 7× speedup over SciPy. JIT compilation folds the adaptive stepping loop into a single XLA kernel, eliminating Python overhead. Moreover, JAX autodiff provides exact gradients with one backward pass, turning a gradient that previously required four forward solves into a cost comparable to a single forward evaluation.

The author benchmarked both solvers on a flat ΛCDM inference problem using 30 mock supernovae. Total runtime dropped from 139 s (SciPy pipeline) to 24.8 s with diffrax, and gradient computation became 8× faster. This demonstrates that modern JAX‑based ODE tools can remove a major bottleneck in nested‑sampling workflows without sacrificing numerical accuracy.

Adopting diffrax required only a pip install of jax and diffrax, after which the same Runge–Kutta methods (Tsit5) replace SciPy’s RK45. It also supports JAX’s vmap, allowing batch solves of parameter vectors, which further reduces overhead in nested sampling. Researchers can now scale Bayesian analyses to parameter spaces without the previous gradient penalty.