Making a for-loop more efficient

[yunus]

By the way, you asked if I meant to use vmap() I’m a bit confused now as you recommended this to me earlier here:

Efficiently sampling a large ODE model (compiling issues?)

I think this is the source of slowness for i, rid in enumerate(RID_list): Each one of odeint can be slow to compile, the amount of numerical operators (to be compiled to XLA) with Python is scaled linearly with the range of your Python loop. In JAX, the compiling time is scaled non-linear w.r.t. the amount of numerical operators, which makes things even worse. For a solution, you can try to use jax.vmap with your odeint to make it faster. When compiling time is > 1 minute, a rule-of-thumb is …

Am I not using it correctly?