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

2026-06-06 · Source: Towards Data Science · Field: Science & Research — Space Science & Astronomy, Mathematics & Computational Sciences, Artificial Intelligence & Machine Learning · Depth: Advanced, long

Summary

A theoretical cosmologist encountered significant performance bottlenecks using SciPy's `solve_ivp` for Bayesian inference in complex dark energy models, where ODE solves alone took 0.4 ms per call, escalating to 5 minutes for gradients in a 10-parameter model. The solution found was `diffrax`, a JAX-based ODE solver. `diffrax` achieved a 7x speedup for single forward calls (from 404 μs to 59 μs) and an 8x speedup for 2-parameter gradient computations (from 1.62 ms to 195 μs) compared to SciPy with finite differences. This efficiency stems from JIT compilation, automatic differentiation for exact gradients, and `jax.vmap` for parallel batch processing. The `diffrax` pipeline, combined with an Adam optimizer, successfully converged to accurate cosmological parameters (Ωₚ = 0.270, H₀ = 70.94) in 350 steps, while the SciPy pipeline failed to converge.

Key takeaway

For Research Scientists or Machine Learning Engineers performing Bayesian inference or other tasks requiring repeated, differentiable ODE solves, adopting `diffrax` can drastically improve computational efficiency and convergence. You should consider migrating from SciPy's `solve_ivp` to `diffrax` to leverage JIT compilation for 7x faster forward passes and automatic differentiation for 8x faster, exact gradients, especially with many parameters. This enables advanced optimizers like Adam, leading to more robust and accurate parameter estimation, transforming multi-minute runs into sub-minute tasks.

Key insights

JAX-based `diffrax` ODE solvers significantly accelerate Bayesian inference by enabling JIT compilation, autodiff, and vectorization.

Principles

• JIT compilation eliminates Python overhead.
• Autodiff provides exact gradients efficiently.
• `jax.vmap` enables parallel batch solving.

Method

`diffrax` integrates ODEs using classical algorithms like Tsit5, leveraging JAX's infrastructure for JIT compilation, automatic differentiation (adjoint method), and vectorization (`jax.vmap`) to achieve substantial speedups and accurate gradients for inference tasks.

In practice

• Enable `jax_enable_x64` for tight tolerances.
• Warm up JIT-compiled functions before benchmarking.
• Use `dfx.Tsit5()` for most non-stiff ODEs.

Topics

• Bayesian Inference
• Ordinary Differential Equations
• JAX
• Diffrax
• Automatic Differentiation
• Cosmological Models
• Tsit5 Algorithm

Code references

• Samit1424/ODE_solver_comparison

Best for: Research Scientist, Machine Learning Engineer, Data Scientist

Related on AIssential

• AdaGrad-Diff: A New Version of the Adaptive Gradient Algorithm — AdaGrad-Diff improves optimization robustness by adapting stepsizes based on gradient differences, not just gradient magnitudes.
• xFODE: An Explainable Fuzzy Additive ODE Framework for System Identification — xFODE enhances system identification interpretability by defining incremental states and using fuzzy additive models with partitioning strategies.
• Stable and Scalable Probabilistic Numerical Solvers for Stiff and High-Dimensional ODEs — New probabilistic ODE solvers achieve both stability and linear scalability for stiff, high-dimensional problems.
• Micro-Diffusion Compression -- Binary Tree Tweedie Denoising for Online Probability Estimation — Micro-diffusion with binary tree Tweedie denoising enhances lossless compression by correcting systematic biases in blended probability distributions.
• Domino: Decoupling Causal Modeling from Autoregressive Drafting in Speculative Decoding — Domino accelerates LLM inference by decoupling causal modeling from autoregressive drafting, using parallel generation refined by causal information.
• Conditional diffusion with locality-aware modal alignment for generating diverse protein conformational ensembles — Mac-Diff uses a novel diffusion model with locality-aware attention and PLM embeddings to generate diverse protein conformational ensembles efficiently.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by Towards Data Science.