[Regression] Slower integration of differential equations since jaxlib > 0.4.32.dev20240807

Between jaxlib==0.4.32.dev20240807 and jaxlib==0.4.32.dev20240812 I observe a significant decrease of performance for integration of differential equations with many solver steps (up to 8x slower). Minimal example:

import timeit

import numpy as np
import diffrax
import jax
import jax.numpy as jnp
from diffrax import diffeqsolve, Tsit5, ODETerm, SaveAt, PIDController, StepTo

def ODE(t, y, args):
    return jnp.cos(y)  

t_out = np.linspace(0, 20, 10)  
t_steps = np.linspace(0, 20, 10000) 

stepsize_controller = StepTo(ts=t_steps)
saveat = SaveAt(ts=t_out)

term = ODETerm(ODE)
y0 =1.
f = lambda y0: diffeqsolve(term, Tsit5(), t0=t_out[0], t1=t_out[-1], dt0=None, y0=y0, saveat=saveat,
			   stepsize_controller=stepsize_controller, max_steps=len(t_steps))
f = jax.jit(f)

solution = jax.block_until_ready(f(y0))

def f_timer():
    jax.block_until_ready(f(1.))    
runtime = timeit.timeit(f_timer, number=100)/100*1000
print(f"Runtime: {runtime:.3f} ms")

Tested on Ubuntu 22.04 and CPU backend. Runtime is on my PC 1.8 ms for the nightly jaxlib version 20240807 and 14.8ms for version 20240812. The difference is the largest if t_steps is quite large.

To test it quickly, I used the following one-liner: uv venv --python 3.12 && uv pip install diffrax numpy --pre jax==0.4.32.dev20240807 jaxlib==0.4.32.dev20240807 -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html --reinstall --exclude-newer 2024-09-20 && uv run test_diffrax.py && uv pip install diffrax numpy --pre jax==0.4.32.dev20240807 jaxlib==0.4.32.dev20240812 -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html --reinstall --exclude-newer 2024-09-20 && uv run test_diffrax.py, where test_diffrax.py is the script above.

I don't know whether I should have better opened the issue on the JAX Github, let me know if it isn't correct here.

The report provides no repository file or test; start with the minimal test_diffrax.py example and compare the two jaxlib versions on the CPU backend. Done means identifying whether the slowdown is in diffrax or JAX and documenting a reproducible cause or actionable next step.