Add version constraints to CUDA pip wheel dependencies.

Fixes #15267
jax-ml · Mar 29, 2023 · 705b5cc · 705b5cc
1 parent 8c4fed6
commit 705b5cc
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/setup.py b/setup.py
@@ -96,13 +96,13 @@ def generate_proto(source):
 
         'cuda11_pip': [
           f"jaxlib=={_current_jaxlib_version}+cuda11.cudnn{_default_cuda11_cudnn_version}",
-          "nvidia-cublas-cu11",
-          "nvidia-cuda-nvcc-cu11",
-          "nvidia-cuda-runtime-cu11",
-          "nvidia-cudnn-cu11",
-          "nvidia-cufft-cu11",
-          "nvidia-cusolver-cu11",
-          "nvidia-cusparse-cu11",
+          "nvidia-cublas-cu11>=11.11",
+          "nvidia-cuda-nvcc-cu11>=11.8",
+          "nvidia-cuda-runtime-cu11>=11.8",
+          "nvidia-cudnn-cu11>=8.6",
+          "nvidia-cufft-cu11>=10.9",
+          "nvidia-cusolver-cu11>=11.4",
+          "nvidia-cusparse-cu11>=11.7",
         ],
 
         'cuda12_pip': [