jax
einops
equinox
diffrax