Scipy vs diffrax speed

[VMLC-PV]

I conducted some quick tests to solve a stiff PDE I am interested in, and I benchmarked the diffrax solver against scipy's solve_ivp.
I was surprised to see that scipy seems to be significantly faster (by a factor of ~4). Is that expected, or did I miss something?

[johannahaffner]

That is very much not expected, do you have a reproducible example?

[VMLC-PV]

I have put the notebook I ran in this repo .
I apologize for the dirtiness of the code, but I am just playing with it for now, so I have not spent the time to clean it up.

[johannahaffner]

You probably have a recompilation issue. You can use Equinox’ debugging tools to confirm that, particularly assert_max_traces, documented here. It looks like your time points are raw floats?

[VMLC-PV]

Yes, the time points are floats, but I also tried converting the array to jnp, which made no difference.
I also added the decorator you suggested but I see no difference from my previous runs.

[johannahaffner]

Did you put the decorator inside the jit, like in the example?

[johannahaffner]

Alright, I see how you did that (in your updated repo). If diffeqsolve compiled more than once, then you should get an error with this setup. So it seems like recompilation is not the culprit here. LSODA, which I think your scipy solve is backing out to, is also quite efficient.

If you really want to see which solver works best for your application, you could benchmark them head to head in a work-precision diagram, just on the bit of code that solves the ODE.

[VMLC-PV]

I don't need to do that, I am perfectly happy with having both options. Especially since the 'LSODA' solver from scipy has quite a lot of overhead. The other stiff solvers from Scipy (Radau and BDF) are also significantly faster than diffrax (depending on the input parameters). However, they also seem less stable compared to 'LSODA' and diffrax.
So all good.
I was curious since I expected diffrax and the JAX backend+gpu to be a lot faster.
I will leave the issue up for a couple of weeks in case someone sees something we missed and then close it.
Thank you for the help and the great work on the package!

[johannahaffner]

Alright! LSODA has a Fortran backend, and Fortran is often comparable to JAX in terms of speed (trying to beat it in a side-project of mine).

For an inherently sequential, non-batched operation like a short ODE solve, I would not expect a GPU backend to help. I also don’t think it is used in your example, since a warning is printed at the top that JAX will fall back to CPU due to a missing install?

[VMLC-PV]

I ran it with the GPU too, but it does not make much difference

[patrick-kidger]

Absent a MWE it's hard to say where the issue is here, but I just wanted to comment on this line:

and Fortran is often comparable to JAX in terms of speed

as this is typically not true when using scipy, as the vector field is usually still written in Python.

(As a rough rule of thumb, we usually expect Diffrax to be about 40x faster than scipy.)

---

Anyway, in terms of next steps. It sounds like you're happy with what you have, so no changes are needed :) On that basis I'm going to close this issue.

(If you were interested in investigating this then a small (~50 lines) apples-to-apples MWE would be the necessary starting point from which we could provide advice.)