tensorflow_datasets
ml_collections
jax
flax
optax
pandas
seaborn
scipy
gcsfs
tensorflow