numpy
scipy
matplotlib
seaborn
optax
flax
ml_collections
tqdm
absl-py
wandb
