Feature Request: Complex-Valued Integration With ZVODE - CVODE in Jax (autodiff)

[onurdanaci]

Hi,

Unfortunately, all the available (S)ODE integration subroutines in auto-differentiable Python frameworks (RK45, Dopri, etc.) behave very poorly with complex-valued functions [*]. In the Python ecosystem, only Scipy's Fortran wrappers titled ode (ZVODE) and complex_ode (using CVODE) seem to be working fine, but obviously, they are not differentiable and not applicable to the modern applications we love.

I was wondering if anybody wants to adapt these features to Diffrax, and make them auto-differentiable.

[*] https://arxiv.org/abs/2406.06361