absl-py
ml-collections
wandb
gcsfs
cloudpickle
numpy
transformers

[jax]
jax
flax
optax
