jaxlib
diffrax
jaxtyping