absl-py
ml-collections
wandb
gcsfs
cloudpickle

[jax]
numpy
jax
flax
optax
