My SciPy ODE Solver Was Killing My Bayesian Inference: A Cosmologist’s Honest Account of Discovering Diffrax

A theoretical cosmologist describes a significant performance bottleneck when using SciPy's ODE solver for Bayesian inference. The process of solving differential equations and calculating gradients for complex models was prohibitively slow. By switching to Diffrax, a library built on JAX, the entire computation could be JIT-compiled for a dramatic speed increase. This new approach also provided exact gradients through automatic differentiation at a fraction of the cost of traditional methods, fundamentally improving the efficiency of the scientific workflow.