absl-py
ml-collections
wandb
gcsfs
cloudpickle
numpy

[jax]
jax
flax
optax
