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

[test]
pytest
tensorflow
tensorflow_datasets
torch<2.0.0,>=1.13.0
