jaxlib<=0.9.2,>=0.9.2
ml_dtypes>=0.5.0
numpy>=2.0
opt_einsum
scipy>=1.13

[ci]
jaxlib==0.9.1

[cpu]

[cuda]
jaxlib<=0.9.2,>=0.9.2
jax-cuda12-plugin[with-cuda]<=0.9.2,>=0.9.2

[cuda12]
jaxlib<=0.9.2,>=0.9.2
jax-cuda12-plugin[with-cuda]<=0.9.2,>=0.9.2

[cuda12-local]
jaxlib<=0.9.2,>=0.9.2
jax-cuda12-plugin<=0.9.2,>=0.9.2

[cuda13]
jaxlib<=0.9.2,>=0.9.2
jax-cuda13-plugin[with-cuda]<=0.9.2,>=0.9.2

[cuda13-local]
jaxlib<=0.9.2,>=0.9.2
jax-cuda13-plugin<=0.9.2,>=0.9.2

[k8s]
kubernetes

[minimum-jaxlib]
jaxlib==0.9.2

[rocm7-local]
jaxlib<=0.9.2,>=0.9.2
jax-rocm7-plugin==0.9.2.*

[tpu]
jaxlib<=0.9.2,>=0.9.2
libtpu==0.0.37.*
requests

[xprof]
xprof
