flax>=0.3.5
jax>=0.3.16
numpy>=1.18.0
