tensorflow>=2.9.0
tensorflow-datasets>=4
flax>=0.3
optax>=0.1
clu>=0.0.6
