Step-by-Step Interface is very Inefficient vs. Base diffeqsolve

Hi @patrick-kidger.

Assume that I have a DAE model as an eqx.Module class. I have also have a simulator implemented as a step-by-step interface that at every step: 1. compute the events, 2. integration step, 3. apply assignments to the args vector (explicit algebraic equations).

However, it seems that this interface is very, but very, inefficient. Here is a MWE to demonstrate this: link

To make the comparisions, I have a simple ODE model in this case (there are no algebraic equations here). Moreover, I have implemented a simple VMAPped version of a simulation using just diffeqsolve (that in this case can be used). Then I have done a simple comparision between the step-by-step simulator and the diffrax one both in a situation of single initial assignments and with multiple initial conditions. In the single initial condition version the elapsed times are relatively the same: 0.042 seconds for Diffrax version and 0.411 seconds for mine one. However, it happens that when multiple initial condition are applied, the distance between the two times exponentially increase: 93.286 seconds (mine) and 0.373 seconds (diffeqsolve). I understand that there might be a little bit more overhead since at every step I'm calling the diffeqsolve, however this is too much!!!

Still, in practice, I can't use the simple diffrax version since in general I have hybrid DAE models. In this case I reduced to a simple ODE model in order to make those comparisions.

Do you have any idea about this?

Thank you for all the support!!!

Start with the linked Colab MWE and compare the step-by-step simulator with the vmapped diffeqsolve example for single and multiple initial conditions. Trace the step-by-step interface and repeated diffeqsolve calls to identify why performance scales so differently; done means documenting the bottleneck and a validated direction for improving hybrid DAE simulations.