Fix uv install and add dev release (#2411)

* fix uv install and add dev release

* fix docstring

* pin cuda deps

* cuda release on cpu-only machine

.circleci/config.yml

@@ -7,6 +7,9 @@ parameters:
   nightly_build:
     type: boolean
     default: false
+  test_release:
+    type: boolean
+    default: false
 
 jobs:
   build_documentation:
@@ -366,22 +369,27 @@ jobs:
         type: string
         default: ""
     machine:
-      image: linux-cuda-12:2024.11.1
+      image: ubuntu-2204:current
-      resource_class: gpu.nvidia.small.gen2
+      resource_class: large
     steps:
       - checkout
       - run:
           name: Build wheel
           command: |
+            export DEBIAN_FRONTEND=noninteractive
+            export NEEDRESTART_MODE=a
+            wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
+            sudo dpkg -i cuda-keyring_1.1-1_all.deb
             sudo apt-get update
+            sudo apt install cuda-toolkit-12-9
             sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
             sudo apt-get install zip
-            python -m venv env
-            source env/bin/activate
             pip install auditwheel
             pip install patchelf
             pip install build
             pip install twine
+            export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
+            export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
             << parameters.build_env >> MLX_BUILD_STAGE=2 \
               CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
               python -m build -w
@@ -392,7 +400,6 @@ jobs:
             - run:
                 name: Upload package
                 command: |
-                  source env/bin/activate
                   twine upload wheelhouse/*.whl
       - store_artifacts:
           path: wheelhouse/
@@ -405,6 +412,7 @@ workflows:
             pattern: "^(?!pull/)[-\\w]+$"
             value: << pipeline.git.branch >>
         - not: << pipeline.parameters.nightly_build >>
+        - not: << pipeline.parameters.test_release >>
     jobs:
       - mac_build_and_test:
           matrix:
@@ -418,6 +426,7 @@ workflows:
     when:
       and:
         - not: << pipeline.parameters.nightly_build >>
+        - not: << pipeline.parameters.test_release >>
     jobs:
       - build_release:
           filters:
@@ -601,3 +610,87 @@ workflows:
             parameters:
               python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
       - build_cuda_release
+
+  build_dev_release:
+    when:
+      and:
+        - equal: [ main, << pipeline.git.branch >> ]
+        - << pipeline.parameters.test_release >>
+    jobs:
+      - build_release:
+          matrix:
+            parameters:
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              macosx_deployment_target: ["13.5", "14.0", "15.0"]
+              build_env: ["DEV_RELEASE=1"]
+              xcode_version: ["16.2.0", "15.0.0"]
+            exclude:
+              - macosx_deployment_target: "13.5"
+                xcode_version: "16.2.0"
+                python_version: "3.9"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "13.5"
+                xcode_version: "16.2.0"
+                python_version: "3.10"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "13.5"
+                xcode_version: "16.2.0"
+                python_version: "3.11"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "13.5"
+                xcode_version: "16.2.0"
+                python_version: "3.12"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "13.5"
+                xcode_version: "16.2.0"
+                python_version: "3.13"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "14.0"
+                xcode_version: "15.0.0"
+                python_version: "3.9"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "14.0"
+                xcode_version: "15.0.0"
+                python_version: "3.10"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "14.0"
+                xcode_version: "15.0.0"
+                python_version: "3.11"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "14.0"
+                xcode_version: "15.0.0"
+                python_version: "3.12"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "14.0"
+                xcode_version: "15.0.0"
+                python_version: "3.13"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "15.0"
+                xcode_version: "15.0.0"
+                python_version: "3.9"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "15.0"
+                xcode_version: "15.0.0"
+                python_version: "3.10"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "15.0"
+                xcode_version: "15.0.0"
+                python_version: "3.11"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "15.0"
+                xcode_version: "15.0.0"
+                python_version: "3.12"
+                build_env: "DEV_RELEASE=1"
+              - macosx_deployment_target: "15.0"
+                xcode_version: "15.0.0"
+                python_version: "3.13"
+                build_env: "DEV_RELEASE=1"
+      - build_linux_release:
+          matrix:
+            parameters:
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              build_env: ["DEV_RELEASE=1"]
+      - build_cuda_release:
+          matrix:
+            parameters:
+              build_env: ["DEV_RELEASE=1"]

python/scripts/repair_cuda.sh

@@ -4,6 +4,7 @@ auditwheel repair dist/* \
   --plat manylinux_2_35_x86_64 \
   --exclude libcublas* \
   --exclude libnvrtc* \
+  --exclude libcuda* \
   -w wheel_tmp
 
 

python/src/ops.cpp

@@ -4022,7 +4022,8 @@ void init_ops(nb::module_& m) {
         Args:
             file (file, str): File in which the array is saved.
             arrays (dict(str, array)): The dictionary of names to arrays to
-            be saved. metadata (dict(str, str), optional): The dictionary of
+              be saved.
+            metadata (dict(str, str), optional): The dictionary of
               metadata to be saved.
       )pbdoc");
   m.def(

setup.py

@@ -205,9 +205,6 @@ if __name__ == "__main__":
     build_macos = platform.system() == "Darwin"
     build_cuda = "MLX_BUILD_CUDA=ON" in os.environ.get("CMAKE_ARGS", "")
 
-    install_requires = []
-    if build_cuda:
-        install_requires = ["nvidia-cublas-cu12", "nvidia-cuda-nvrtc-cu12"]
     version = get_version()
 
     _setup = partial(
@@ -250,6 +247,7 @@ if __name__ == "__main__":
             "mlx.distributed_config = mlx.distributed_run:distributed_config",
         ]
     }
+    install_requires = []
 
     # Release builds for PyPi are in two stages.
     # Each stage should be run from a clean build:
@@ -269,11 +267,11 @@ if __name__ == "__main__":
     #  - Package name is back-end specific, e.g mlx-metal
     if build_stage != 2:
         if build_stage == 1:
-            if build_macos:
+            install_requires.append(
-                install_requires += [f"mlx-metal=={version}"]
+                f'mlx-metal=={version}; platform_system == "Darwin"'
-            else:
+            )
-                extras["cuda"] = [f"mlx-cuda=={version}"]
+            extras["cuda"] = [f'mlx-cuda=={version}; platform_system == "Linux"']
-                extras["cpu"] = [f"mlx-cpu=={version}"]
+            extras["cpu"] = [f'mlx-cpu=={version}; platform_system == "Linux"']
 
         _setup(
             name="mlx",
@@ -288,9 +286,14 @@ if __name__ == "__main__":
             name = "mlx-metal"
         elif build_cuda:
             name = "mlx-cuda"
+            install_requires += [
+                "nvidia-cublas-cu12==12.9.*",
+                "nvidia-cuda-nvrtc-cu12==12.9.*",
+            ]
         else:
             name = "mlx-cpu"
         _setup(
             name=name,
             packages=["mlx"],
+            install_requires=install_requires,
         )