absl-py
etils[epath]
flax
jax
jaxlib
ml_collections
numpy
packaging
typing_extensions
wrapt

[pytorch]
torch>=1.2.0

[test]
pytest
tensorflow
tensorflow_datasets
torch>=1.2.0
