You may already know of it, but if you want differential-equations-in-JAX then allow me to quickly advertise Diffrax: https://github.com/patrick-kidger/diffrax (of which I am the author, disclaimer).