absl-py>=0.6.0
numpy>=1.15.4
six>=1.12.0
tensorflow-gpu>=2.3.1  # change to 'tensorflow' for cpu
tensorflow-datasets>=2.1
tensorflow-model-optimization>=0.5.0
gin-config>=0.3.0
jax>=0.2.6
jaxlib>=0.1.57
tensorflow-model-optimization>=0.5.0