jax
equinox
einops
diffrax
functools
