flax>=0.6.6
jax>=0.4.25
numpy>=1.22
