JAX version had to be pinned to 0.4.33

[ohahn]

Running with the most recent version >=0.4.34 leads to errors:
TypeError: Called multiply with a float0 array. float0s do not support any operations by design because they are not compatible with non-trivial vector spaces. No implicit dtype conversion is done. You can use np.zeros_like(arr, dtype=np.float) to cast a float0 array to a regular zeros array. If you didn't expect to get a float0 you might have accidentally taken a gradient with respect to an integer argument.

Currently this fixed by fixing JAX version 0.4.33 (PR #16 ), but this needs to be fixed properly.

[Findus23]

The error stays the same with jax 0.4.35. I assume the fix is to upgrade diffrax and equinox to the latest versions.