---
title: "My SciPy ODE Solver Was Killing My Bayesian Inference: A Cosmologist’s Honest Account of Discovering Diffrax"
url: https://daily.dev/posts/my-scipy-ode-solver-was-killing-my-bayesian-inference-a-cosmologist-s-honest-account-of-discovering-ejsgq5fxn
source_url: https://towardsdatascience.com/my-scipy-ode-solver-was-killing-my-bayesian-inference-a-cosmologists-honest-account-of-discovering-diffrax
type: article
source: "Towards Data Science"
published: 2026-06-06T13:59:21.575Z
updated: 2026-06-07T17:44:48.677Z
tags: ["python", "data-analysis"]
reading_time: 13
upvotes: 1
comments: 0
language: en
---

> ## Documentation Index
> Fetch the complete documentation index at: https://daily.dev/llms.txt
> Use this file to discover all available pages before exploring further.

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

**[Towards Data Science](https://daily.dev/sources/tds)** · 13 min read · 1 upvotes · 0 comments

## Summary

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.

## Full article

daily.dev links to this article rather than hosting it. Read it at the original source: <https://towardsdatascience.com/my-scipy-ode-solver-was-killing-my-bayesian-inference-a-cosmologists-honest-account-of-discovering-diffrax>

## Similar posts on daily.dev

- [Why Julia’s GPU Accelerated ODE Solvers are 20x-100x Faster than Jax and PyTorch](https://daily.dev/posts/why-julia-s-gpu-accelerated-ode-solvers-are-20x-100x-faster-than-jax-and-pytorch-qmmig41qz) · Julia Bloggers · 0 upvotes · 0 comments
- [Build Accelerated, Differentiable Computational Physics Code for AI with NVIDIA Warp](https://daily.dev/posts/build-accelerated-differentiable-computational-physics-code-for-ai-with-nvidia-warp-haygcavvw) · NVIDIA Developer · 0 upvotes · 0 comments
- [The Optimization Ladder](https://daily.dev/posts/the-optimization-ladder-jor2sohos) · Hacker News · 0 upvotes · 0 comments

---

Tags: [#python](https://daily.dev/tags/python), [#data-analysis](https://daily.dev/tags/data-analysis)

[View this post on daily.dev](https://daily.dev/posts/my-scipy-ode-solver-was-killing-my-bayesian-inference-a-cosmologist-s-honest-account-of-discovering-ejsgq5fxn)