Thanks so much for your input! Extremely helpful. It indeed seems that using vmap significantly increases the speed of the repeated ODE solving for large N. Interestingly it seems to be slightly slower for small N but it barely increases in time for larger N.
I’m now using the following code, is that also what you had in mind?
jit(vmap(lambda y0: odeint(dz_dt, y0, ts[1:], A)))(data_bl)
I’ll be sampling from the 4D ODE model with N=500 data overnight and will report on how fast in runs. Jit may also help. Thanks again!