tensorflow>=2
tensorflow-datasets>=4
tfds-nightly>=4.0.0.dev
flax>=0.3
optax>=0.1
clu>=0.0.6
