dm-haiku>=0.0.4
flax>=0.5.3
optax>=0.0.9
