numpy
jax>=0.4.0

[benchmark]
psutil
matplotlib

[benchmark-gpu]
psutil
matplotlib
nvidia-ml-py

[dev]
check-manifest

[jax-cuda12]
jax[cuda12]>=0.4.0

[jax-cuda13]
jax[cuda13]>=0.4.0
