jax
numpy
tqdm
equinox
jaxtyping
matplotlib
