jax[cpu]<0.12,>=0.6.0
matplotlib>=3.8.4

[gpu]
jax[cuda12]<0.12,>=0.6.0
