jax
equinox
einops
diffrax
