jax
einops
equinox
diffrax
pytest
matplotlib
jaxtyping
ml-dtypes
numpy
plum-dispatch
wadler_lindig
tqdm