numpy>=2.0.0
matplotlib>=3.8.0
jax>=0.4.20
chex>=0.1.0
