tensorflow>=2,<3
tensorflow-datasets>=4,<5
tfds-nightly>=4.0.0.dev,<5
flax>=0.3,<1
optax>=0.1,<1
clu>=0.0.6,<1
