jax>=0.4.0
jaxlib>=0.4.0
jaxtyping>=0.2.0
diffrax>=0.5.0
