fix release builds (#2746)

.github/actions/build-cuda-release/action.yml

@@ -12,13 +12,9 @@ runs:
     - name: Build package
       shell: bash
       env:
-        MLX_BUILD_STAGE: 2
         CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }}
       run: |
         pip install auditwheel build patchelf setuptools
         python setup.py clean --all
-        python -m build -w
+        MLX_BUILD_STAGE=2 python -m build -w
-
+        bash python/scripts/repair_cuda.sh
-        if [ -f "python/scripts/repair_cuda.sh" ]; then
-          bash python/scripts/repair_cuda.sh
-        fi

.github/actions/build-cuda/action.yml

@@ -2,19 +2,10 @@ name: 'Build and Test with CUDA'
 description: 'Build and test MLX with CUDA'
 
 inputs:
-  build-type:
-    description: 'Build type (debug, release)'
-    required: false
-    default: 'debug'
-  run-tests:
-    description: 'Whether to run tests'
-    required: false
-    default: 'true'
   nvcc-location:
     description: 'Location of nvcc compiler'
     required: true
     default: '/usr/local/cuda-12.9/bin/nvcc'
-    # this value is dependent on the CUDA tools installed in the setup-linux workflow
 
 runs:
   using: "composite"
@@ -27,7 +18,6 @@ runs:
       run: pip install -e ".[dev]" -v
 
     - name: Run Python tests - CPU
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         LOW_MEMORY: 1
@@ -35,7 +25,6 @@ runs:
       run: python -m unittest discover python/tests -v
 
     - name: Run Python tests - GPU
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         LOW_MEMORY: 1
@@ -43,7 +32,6 @@ runs:
       run: python -m tests discover python/tests -v
 
     - name: Build CPP only
-      if: inputs.build-type == 'debug'
       shell: bash
       run: |
         cmake . -B build \
@@ -53,12 +41,5 @@ runs:
         cmake --build build -j $(nproc)
     
     - name: Run CPP tests
-      if: ${{ inputs.build-type == 'debug' && inputs.run-tests == 'true' }}
       shell: bash
       run: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
-
-    - name: Build Python package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-cuda-release
-      with:
-        nvcc-location: ${{ inputs.nvcc-location }}

.github/actions/build-linux-release/action.yml

@@ -0,0 +1,33 @@
+name: 'Build Linux wheel'
+description: 'Build Linux wheel'
+
+inputs:
+  build-backend:
+    description: 'Build the backend mlx-cpu package'
+    type: boolean
+    required: false
+    default: false
+
+runs:
+  using: "composite"
+  steps:
+    - name: Generate package stubs
+      shell: bash
+      run: |
+        pip install -e ".[dev]" -v
+        pip install typing_extensions
+        python setup.py generate_stubs
+    - name: Build Python package
+      shell: bash
+      run: |
+        pip install auditwheel patchelf build
+        python setup.py clean --all
+        MLX_BUILD_STAGE=1 python -m build -w
+        bash python/scripts/repair_linux.sh
+    - name: Build backend package
+      if: ${{ inputs.build-backend }}
+      shell: bash
+      run: |
+        python setup.py clean --all
+        MLX_BUILD_STAGE=2 python -m build -w
+        auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64

.github/actions/build-linux/action.yml

@@ -1,33 +1,14 @@
 name: 'Build and Test on Linux'
 description: 'Build and test MLX on Linux'
 
-inputs:
-  build-type:
-    description: 'Build type'
-    required: false
-    default: 'debug'
-    type: choice
-    options:
-      - debug
-      - release
-  run-tests:
-    description: 'Whether to run tests'
-    required: false
-    default: 'true'
-    type: boolean
-
 runs:
   using: "composite"
   steps:
-    - name: Set DEBUG
-      shell: sh
-      if: inputs.build-type == 'debug'
-      run: echo "DEBUG=1" >> $GITHUB_ENV
-
     - name: Install Python package
       shell: sh
       env:
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
+        DEBUG: 1
       run: pip install -e ".[dev]" -v
     
     - name: Generate package stubs
@@ -37,7 +18,6 @@ runs:
         python setup.py generate_stubs
     
     - name: Run Python tests
-      if: inputs.run-tests == 'true'
       shell: bash
       run: |
         python -m unittest discover python/tests -v
@@ -50,7 +30,6 @@ runs:
         fi
     
     - name: Build CPP only
-      if: inputs.build-type == 'debug'
       shell: bash
       run: |
         mkdir -p build && cd build
@@ -58,21 +37,5 @@ runs:
         make -j $(nproc)
     
     - name: Run CPP tests
-      if: ${{ inputs.build-type == 'debug' && inputs.run-tests == 'true' }}
       shell: sh
       run: ./build/tests/tests
-    
-    - name: Build Python package
-      if: inputs.build-type == 'release'
-      shell: bash
-      run: |
-        pip install auditwheel patchelf build
-        python setup.py clean --all
-        MLX_BUILD_STAGE=1 python -m build -w
-        if [ -f "python/scripts/repair_linux.sh" ]; then
-          bash python/scripts/repair_linux.sh
-        fi
-
-        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
-        auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64

.github/actions/build-macos-release/action.yml

@@ -6,11 +6,16 @@ inputs:
     description: 'macOS build target'
     required: false
     default: '15.0'
+  build-backend:
+    description: 'Build the backend mlx-metal package'
+    type: boolean
+    required: false
+    default: false
 
 runs:
   using: "composite"
   steps:
-    - name: Build Python package(s)
+    - name: Build Python package
       shell: bash
       env:
         MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
@@ -18,5 +23,11 @@ runs:
         uv pip install build
         uv run --no-project setup.py clean --all
         MLX_BUILD_STAGE=1 uv run -m build -w
+    - name: Build backend package
+      if: ${{ inputs.build-backend }}
+      shell: bash
+      env:
+        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
+      run: |
         uv run --no-project setup.py clean --all
         MLX_BUILD_STAGE=2 uv run -m build -w

.github/actions/build-macos/action.yml

@@ -1,20 +1,6 @@
 name: 'Build and Test on macOS'
 description: 'Build and test MLX on macOS'
 
-inputs:
-  build-type:
-    description: 'Build type (debug, release)'
-    required: false
-    default: 'debug'
-    type: choice
-    options:
-      - debug
-      - release
-  run-tests:
-    description: 'Whether to run tests'
-    required: false
-    default: 'true'
-
 runs:
   using: "composite"
   steps:
@@ -22,7 +8,7 @@ runs:
       shell: sh
       env:
         DEBUG: 1
-        DEV_RELEASE: 1
+        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
       run: |
         uv pip install --upgrade pip
         uv pip install cmake setuptools nanobind==2.4.0
@@ -35,13 +21,11 @@ runs:
         uv run --no-project setup.py generate_stubs
 
     - name: Install tests dependencies
-      if: inputs.run-tests == 'true'
       shell: sh
       run: |
         uv pip install numpy torch tensorflow unittest-xml-reporting
 
     - name: Run Python tests
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         LOW_MEMORY: 1
@@ -53,7 +37,6 @@ runs:
         if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
     
     - name: Build example extension
-      if: inputs.run-tests == 'true'
       shell: bash
       run: |
         cd examples/extensions
@@ -62,7 +45,6 @@ runs:
         uv run --no-project test.py
     
     - name: Build CPP only
-      if: inputs.run-tests == 'true'
       shell: bash
       run: |
         mkdir -p build
@@ -71,7 +53,6 @@ runs:
         make -j $(sysctl -n hw.ncpu)
     
     - name: Run CPP tests
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         DEVICE: gpu
@@ -80,7 +61,6 @@ runs:
       run: ./build/tests/tests
     
     - name: Build small binary with JIT
-      if: inputs.run-tests == 'true'
       shell: bash
       run: |
         mkdir -p build
@@ -94,7 +74,6 @@ runs:
         make -j $(sysctl -n hw.ncpu)
     
     - name: Run Python tests with JIT
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         LOW_MEMORY: 1
@@ -107,14 +86,3 @@ runs:
         uv run -m xmlrunner discover \
             -v python/tests \
             -o test-results/gpu_jit
-
-    - name: Build macOS 14 package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-macos-release
-      with:
-        macos-target: 14.0
-    - name: Build macOS 15 package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-macos-release
-      with:
-        macos-target: 15.0

.github/actions/setup-macos/action.yml

@@ -2,11 +2,6 @@ name: 'Setup macOS Environment'
 description: 'Install dependencies for macOS builds'
 
 inputs:
-  install-mpi:
-    description: 'Whether to install MPI'
-    required: false
-    default: 'true'
-    type: boolean
   python-version:
     description: 'Python version to use'
     required: false
@@ -17,7 +12,6 @@ runs:
   steps:
     - name: Install Homebrew packages
       shell: sh
-      if: inputs.install-mpi == 'true'
       run: /opt/homebrew/bin/brew install openmpi
     
     - name: Verify MetalToolchain installed

.github/workflows/nightly.yml

@@ -18,10 +18,9 @@ jobs:
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux
+      - uses: ./.github/actions/build-linux-release
         with:
-          build-type: release
+          build-backend: ${{ matrix.python-version == '3.10' }}
-          run-tests: false
       - name: Upload mlx artifacts
         uses: actions/upload-artifact@v5
         with:
@@ -54,8 +53,6 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.10", "3.13"]
-    env:
-      MACOSX_DEPLOYMENT_TARGET: "15.0"
     runs-on: [self-hosted, macos]
     steps:
       - uses: actions/checkout@v5
@@ -64,6 +61,17 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - uses: ./.github/actions/build-macos
 
+      - name: Build macOS 15 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 15.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
+      - name: Build macOS 14 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 14.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
+
   build_cuda_with_tests:
     if: github.repository == 'ml-explore/mlx'
     runs-on: gpu-t4-4-core

.github/workflows/release.yml

@@ -53,10 +53,9 @@ jobs:
       - uses: ./.github/actions/setup-linux
         with:
           python-version: ${{ matrix.python_version }}
-      - uses: ./.github/actions/build-linux
+      - uses: ./.github/actions/build-linux-release
         with:
-          build-type: release
+          build-backend: ${{ matrix.python-version == '3.10' }}
-          run-tests: false
       - name: Upload MLX artifacts
         uses: actions/upload-artifact@v5
         with:
@@ -82,10 +81,27 @@ jobs:
       - uses: ./.github/actions/setup-macos
         with:
           python-version: ${{ matrix.python-version }}
-      - uses: ./.github/actions/build-macos
+      - name: Install dependencies
+        shell: sh
+        run: |
+          uv pip install --upgrade pip
+          uv pip install cmake setuptools nanobind==2.4.0
+          uv pip install -e . -v
+      - name: Generate package stubs
+        shell: bash
+        run: |
+          uv pip install typing_extensions
+          uv run --no-project setup.py generate_stubs
+      - name: Build macOS 14 package
+        uses: ./.github/actions/build-macos-release
         with:
-          build-type: release
+          macos-target: 14.0
-          run-tests: false
+          build-backend: ${{ matrix.python-version == '3.10' }}
+      - name: Build macOS 15 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 15.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
       - name: Upload MLX artifacts
         uses: actions/upload-artifact@v5
         with: