jax[cpu]>=0.4.6
matplotlib>=3.7.0
optax>=0.1.4
seaborn>=0.12.2

[:python_version == "3.8"]
numpy<1.25,>=1.22
chex<=0.1.7
scipy<1.11,>=1.9

[:python_version >= "3.9"]
numpy>=1.22
scipy>=1.9
