optax
flax
numpy
matplotlib
scipy
pickle
tqdm