jax
jaxlib
equinox
optax
numpy
matplotlib
tqdm
jupyter
jupyterlab
