absl-py
six>=1.10.0
numpy>=1.13.3
decorator
cloudpickle>=1.3
gast>=0.3.2
dm-tree

[jax]
jax
jaxlib

[tf]
tensorflow>=2.16
tf-keras>=2.16

[tfds]
tensorflow-datasets>=2.2.0
