FYI: JAX 0.7.0 and 0.7.1 are known to be buggy with Diffrax.

[patrick-kidger]

Pinning an issue to highlight that when used with Diffrax:

• JAX 0.7.0 is known to cause segfaults (SIGSEGV/SIGABRT with JAX 0.7.0 jax-ml/jax#30517)
• JAX 0.7.1 is known to cause major performance issues (O(n) program has become O(n^2) in JAX 0.7.1 jax-ml/jax#31284).
  Right now JAX 0.6.* is recommended instead.