If your SciPy solver is only slow because you call it a hundred thousand times, moving it to JAX will fix that. Call it once and the move fixes nothing. A recent piece by Samit Ganguly, a cosmologist who ported a likelihood from scipy.integrate.solve_ivp to Diffrax, works as a Diffrax vs solve_ivp benchmark on Towards Data Science, but what it measures is the cost of calling a small solve that many times, and then paying twice per parameter for its gradient.
Diffrax vs solve_ivp: where the time was going
His profile of a dark energy model showed each likelihood call solving a coupled system of ODEs in 0.4 ms. dynesty nested sampling wanted up to 100,000 of those evaluations, which is 40 seconds of pure solving, and a 10-parameter model needed 20 extra forward solves per finite-difference gradient, which he puts near 5 minutes per run. None of that is solve_ivp’s fault. A 0.4 ms adaptive solve is respectable. The issue is arithmetic: the Python interpreter sits between your loop and the compiled backend on every call, and central finite differences bill you 2N solves per gradient, with N the parameter count.
Diffrax runs the same family of embedded Runge-Kutta methods (Tsit5, where scipy defaults to RK45), but written in JAX, so the adaptive stepping loop compiles into a single kernel and jax.grad differentiates through the solve. On his machine a single call dropped from 404 microseconds to 59 after warmup, and a two-parameter gradient from 1.62 ms to 195 microseconds, roughly seven and eight times. He warns the absolute numbers wander 10 to 15 percent between sessions on a MacBook Air M3, and the ratio is what to trust. I’d add where the ratio comes from: on a 12-step solve most of the scipy cost is dispatch, so the win is largest when solves are small and called constantly. That’s the same reason I’ll call Rust from Python when a loop refuses to vectorize. A long, heavy integration spends its time in arithmetic rather than dispatch, and I’d expect the ratio to shrink there, though his tests don’t cover that case.
His solver guidance carries over to any ODE port: Tsit5 for most non-stiff problems, Dopri8 for tolerances below 1e-9, Kvaerno5 when the solve is stiff. The stiffness test he uses is cheap: print the step count and compare it with what you expected. If it’s 10 to 100 times higher, switch to an implicit solver.
Three traps, and all of them are silent
- JAX defaults to 32-bit floats. His solve needed 69 steps in float32 versus 12 in float64 at the same tolerance, and tighter tolerances can fail outright. The fix is jax.config.update(“jax_enable_x64”, True) before anything else; the JAX docs on the x64 flag are clear it only works at startup.
- The first call to a jitted function pays about 93 ms of compilation. Benchmark cold and Diffrax looks slower than scipy for no good reason. JAX also dispatches asynchronously, so wrap timing loops in block_until_ready() or you measure submission rather than completion.
- The argument-order trap applies if your old code used scipy.odeint. odeint expects f(y, t) while solve_ivp and Diffrax expect f(t, y). Port without swapping and you solve a different equation, with no error to tell you.
There’s a fourth step before you delete a line of scipy: keep the old path as an oracle. Run both models at a few parameter points and assert agreement tighter than your data noise. I keep this as a test rather than a notebook cell, so it runs in CI and catches a version bump that changes solver behavior.
import jax
jax.config.update("jax_enable_x64", True) # startup only, before other jax imports
import numpy as np
import jax.numpy as jnp
from models import forward_scipy, forward_diffrax
# warm up once, then parity-check the port against the old path
_ = forward_diffrax(jnp.array([0.3, 70.0]), z_obs).block_until_ready()
for theta in ([0.30, 70.0], [0.10, 60.0], [0.50, 85.0]):
old = forward_scipy(theta[0], theta[1], z_np)
new = np.asarray(forward_diffrax(jnp.array(theta), z_obs))
assert np.allclose(old, new, atol=1e-7), f"port mismatch at {theta}"The comparison I don’t fully trust
His end-to-end fit looks damning: the scipy pipeline lands on Omega_m 0.65 and H0 60, badly wrong, while the diffrax pipeline recovers 0.270 and 70.94 against a truth of 0.30 and 70. But that experiment changed two things at once. The scipy side used finite differences with plain gradient descent at a fixed learning rate, while the diffrax side used autodiff with Adam, a cosine schedule and rescaled parameters. He says so himself: the failure comes from the 200x scale mismatch between Omega_m and H0, which per-parameter adaptive steps absorb and a fixed-rate descent cannot. That figure reflects the optimizer setup more than the ODE library. The control I’d want is finite-difference gradients feeding Adam. I haven’t run it, so I won’t claim it converges, but I’d expect the gap to narrow a lot.
When I would stay on solve_ivp
A one-off solve with no gradients and no batching should stay on scipy. There’s no compile step, no dtype flag, and the tracebacks point at your code instead of XLA internals. Also check whether you use gradients at all: nested sampling is derivative-free, so in his own setup the gradient speedup mattered for the MAP fit, not for the dynesty run. If a derivative-free sampler already fits your time budget, the port buys you an afternoon of learning JAX and not much else.
One more jit behavior to plan for: compilation is specialized to array shapes, so if your batch size or state size changes between calls you recompile every time, and the speedup can turn into a slowdown. Fix your shapes, or pad to a constant size.
If you have a model, a scoring function or a simulation that runs thousands of times per job and you want a straight answer on whether a rewrite like this pays for itself, that’s work I take on. Send me the profile and the call counts, and I’ll tell you honestly if scipy is already fine.
Do the cheap test first either way: count your calls, multiply by the per-call time, and look at the product. Then port one function rather than the whole pipeline, and compare outputs before you compare timings. Neither the x64 flag nor the argument-order mistake raises an error, so a parity check is the only thing that will catch a confidently wrong likelihood.
