ml_dtypes>=0.2.0
numpy>=1.22
opt_einsum
scipy>=1.9

[:python_version < "3.10"]
importlib_metadata>=4.6

[:python_version >= "3.11"]
numpy>=1.23.2

[:python_version >= "3.12"]
numpy>=1.26.0
scipy>=1.11.1

[australis]
protobuf<4,>=3.13

[ci]
jaxlib==0.4.25

[cpu]
jaxlib==0.4.26

[cuda]
jaxlib==0.4.26+cuda12.cudnn89

[cuda12]
jaxlib==0.4.26
jax-cuda12-plugin==0.4.26
nvidia-cublas-cu12>=12.1.3.1
nvidia-cuda-cupti-cu12>=12.1.105
nvidia-cuda-nvcc-cu12>=12.1.105
nvidia-cuda-runtime-cu12>=12.1.105
nvidia-cudnn-cu12<9.0,>=8.9.2.26
nvidia-cufft-cu12>=11.0.2.54
nvidia-cusolver-cu12>=11.4.5.107
nvidia-cusparse-cu12>=12.1.0.106
nvidia-nccl-cu12>=2.18.1
nvidia-nvjitlink-cu12>=12.1.105

[cuda12_cudnn89]
jaxlib==0.4.26+cuda12.cudnn89

[cuda12_local]
jaxlib==0.4.26+cuda12.cudnn89

[cuda12_pip]
jaxlib==0.4.26+cuda12.cudnn89
nvidia-cublas-cu12>=12.1.3.1
nvidia-cuda-cupti-cu12>=12.1.105
nvidia-cuda-nvcc-cu12>=12.1.105
nvidia-cuda-runtime-cu12>=12.1.105
nvidia-cudnn-cu12<9.0,>=8.9.2.26
nvidia-cufft-cu12>=11.0.2.54
nvidia-cusolver-cu12>=11.4.5.107
nvidia-cusparse-cu12>=12.1.0.106
nvidia-nccl-cu12>=2.18.1
nvidia-nvjitlink-cu12>=12.1.105

[minimum-jaxlib]
jaxlib==0.4.20

[tpu]
jaxlib==0.4.26
libtpu-nightly==0.1.dev20240403
requests
