PR response request

[DanPuzzuoli]

Hey all,

Sorry to bug, but can I get someone to look at PR #11077, either to review, or to determine that it's not a change of interest. It's a very small change, but is very important for an application I'm using odeint for. Getting moving with resolving this one way or another will be very helpful.

Thanks!

[patrick-kidger]

Whilst I can't really speak to the plans around jax.experimental.odeint -- it's been suggested that it be deprecated, but Idk if that'll happen or not -- you may be interested in Diffrax.

You can replicate the current behaviour of odeint in Diffrax as follows:
https://docs.kidger.site/diffrax/further_details/faq/#how-does-this-compare-to-jaxexperimentalodeodeint
and in particular can get step size maximums by passing stepsize_controller=PIDController(..., dtmax=...).

[mattjj]

Hey @DanPuzzuoli , sorry for letting that PR slip through the cracks so hard!

The PR looks great to me. I'll merge it now.

As Patrick also suggested, I suspect libraries like Diffrax have had a lot more attention paid to them. So it's worth checking out. But no reason we can't land improvements to JAX's own odeint while it's around.

[DanPuzzuoli]

Awesome, thanks. I do need to put some more time into using Diffrax, but so far odeint in JAX has been a bit faster for my applications, which may partially be owed to its native support for complex arrays.

[patrick-kidger]

Interesting point -- indeed Diffrax doesn't yet support complex arrays. I'd be happy to accept a PR on this, and I imagine it's mostly just a matter of tweaking a few dtypes in a few places.

[DanPuzzuoli]

Yeah it's something I've thought of contributing but haven't had the time. At some point I would like to transition to using Diffrax, but just need to take the time to make sure I can replicate performance.