jax
jaxlib
jaxtyping
diffrax
