jaxlib<=0.7.1,>=0.7.1rc0
ml_dtypes>=0.5.0
numpy>=1.26
opt_einsum
scipy>=1.12

[ci]
jaxlib==0.7.0

[cpu]

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

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

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

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

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

[k8s]
kubernetes

[minimum-jaxlib]
jaxlib==0.7.1rc0

[rocm]
jaxlib<=0.7.1,>=0.7.1rc0
jax-rocm60-plugin<=0.7.1,>=0.7.1rc0

[tpu]
jaxlib<=0.7.1,>=0.7.1rc0
libtpu==0.0.20.*
requests

[xprof]
xprof
