Type Error in Neural ODE Example

Hi,

I've been using jax 0.4.35 and cuda 11.2 with diffrax 0.6.0 and equinox 0.11.8 to train various types of neural ODEs. This worked very well so far. Thanks for the great packages.

However, after upgrading to jax 0.6.0 and cuda 12.8 with diffrax 0.7.0 and equinox 0.12.1, I can no longer compute gradients through ODE solvers. The forward pass works but the gradient computation with eqx.filter_value_and_grad results in a TypeError. To be precise, I receive: TypeError: Argument 'Zero(ShapedArray(float0[32]))' of type '<class 'jax._src.ad_util.Zero'>' is not a valid JAX type . I have tried both the Tsit5() and Dopri5() solvers. The shape of the array in the TypeError is the same as the batch size.

This occurs not only in my own models, but also when I try to run the neural ODE tutorial from https://github.com/patrick-kidger/diffrax/blob/main/examples/neural_ode.ipynb.

Any advice/ideas on how to resolve this would be great! Many thanks in advance.

Start with examples/neural_ode.ipynb and reproduce the gradient failure using the reported JAX, CUDA, Diffrax, and Equinox versions. Compare the forward and gradient paths for both Tsit5() and Dopri5(), then verify that the tutorial computes gradients successfully without the reported Zero(float0) TypeError.