flax>=0.8.4
optax>=0.2.2
jax>=0.4.28
orbax
clu
