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

[pytorch]
torch>=1.13.0

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