Comment by sitkack

Did you use https://github.com/patrick-kidger/diffrax ?

JAX Talk: Diffrax https://www.youtube.com/watch?v=Jy5Jw8hNiAQ