jax>=0.4.3
diffrax>=0.3.0
jaxtyping
