tensorflow_datasets
ml_collections
jax
flax
optax
pandas
seaborn