Use uv in macOS CI

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

@@ -21,7 +21,7 @@ runs:
         DEVELOPER_DIR: /Applications/Xcode-latest.app
         MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
       run: |
-        pip install build
+        uv pip install build
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
 

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

@@ -4,61 +4,72 @@ description: 'Build and test MLX on macOS'
 runs:
   using: "composite"
   steps:
-    - name: Install dependencies
+    - name: Install Python package
       env:
         DEBUG: 1
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      shell: bash -l {0}
+      shell: bash
       run: |
-        pip install --upgrade pip
-        pip install cmake setuptools typing_extensions
-        pip install -e ".[dev]" -v
+        echo "::group::Install Python package"
+        uv pip install -e ".[dev]" -v
+        echo "::endgroup::"
 
     - name: Install tests dependencies
-      shell: bash -l {0}
+      shell: bash
       run: |
-        pip install tensorflow
+        echo "::group::Install tests dependencies"
+        uv pip install tensorflow
+        echo "::endgroup::"
 
     - name: Run Python tests
-      shell: bash -l {0}
+      shell: bash
       env:
         LOW_MEMORY: 1
       run: |
+        echo "::group::Run Python tests"
         DEVICE=cpu python -m unittest discover -v python/tests
         DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m unittest discover -v python/tests
         mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
         mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
         if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-    
+        echo "::endgroup::"
+
     - name: Build example extension
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build example extension"
         cd examples/extensions
-        pip install -r requirements.txt
-        python setup.py build_ext --inplace
-        python test.py
-    
+        uv pip install -r requirements.txt
+        uv run --no-project setup.py build_ext --inplace
+        uv run --no-project test.py
+        echo "::endgroup::"
+
     - name: Build CPP only
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build CPP only"
         mkdir -p build
         cd build
         cmake ..
         make -j $(sysctl -n hw.ncpu)
+        echo "::endgroup::"
 
     - name: Run CPP tests
-      shell: bash -l {0}
+      shell: bash
       env:
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
       run: |
+        echo "::group::Run CPP tests"
         ./build/tests/tests
         ./build/tests/test_teardown
+        echo "::endgroup::"
 
     - name: Build small binary with JIT
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build small binary with JIT"
         mkdir -p build
         cd build
         cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
@@ -68,15 +79,18 @@ runs:
           -DMLX_BUILD_GGUF=OFF \
           -DMLX_METAL_JIT=ON
         make -j $(sysctl -n hw.ncpu)
-    
+        echo "::endgroup::"
+
     - name: Run Python tests with JIT
-      shell: bash -l {0}
+      shell: bash
       env:
         LOW_MEMORY: 1
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
       run: |
+        echo "::group::Run Python tests with JIT"
         CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-          pip install -e . -v
+          uv pip install -e . -v
         python -m unittest discover -v python/tests
+        echo "::endgroup::"

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

@@ -13,12 +13,20 @@ runs:
     - name: Install Homebrew packages
       shell: sh
       run: /opt/homebrew/bin/brew install openmpi
-    
+
     - name: Verify MetalToolchain installed
       shell: bash
       run: xcodebuild -showComponent MetalToolchain
 
-    - uses: conda-incubator/setup-miniconda@v3
-      with:
-        miniconda-version: "latest"
-        python-version: ${{ inputs.python-version }}
+    - uses: astral-sh/setup-uv@v7
+
+    - name: Setup Python venv
+      shell: bash
+      run: |
+        echo "::group::Setup Python venv"
+        uv venv --python ${{ inputs.python-version }}
+        source .venv/bin/activate
+        echo PATH=$PATH >> $GITHUB_ENV
+        # Search python packages in .venv
+        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
+        echo "::endgroup::"

.github/workflows/release.yml

@@ -93,13 +93,8 @@ jobs:
       - uses: ./.github/actions/setup-macos
         with:
           python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        shell: bash -l {0}
-        run: |
-          pip install --upgrade pip
-          pip install cmake setuptools typing_extensions
-          pip install -e . -v
+      - name: Install Python package
+        run: uv pip install -e . -v
       - name: Build macOS 14 package
         uses: ./.github/actions/build-macos-release
         with:

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.25
-mlx>=0.21.0
+mlx>=0.31.2
 nanobind==2.12.0

mlx/backend/metal/device_info.cpp

@@ -13,12 +13,7 @@ bool is_available() {
 }
 
 int device_count() {
-  try {
-    metal::device(Device::gpu);
-    return 1;
-  } catch (...) {
-    return 0;
-  }
+  return 1;
 }
 
 const std::unordered_map<std::string, std::variant<std::string, size_t>>&

mlx/backend/metal/kernels/fp_quantized.h

@@ -533,7 +533,6 @@ METAL_FUNC void fp_qvm_impl(
     device T* y,
     const int in_vec_size,
     const int out_vec_size,
-    const int in_vec_stride,
     uint3 tid [[threadgroup_position_in_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
@@ -564,7 +563,7 @@ METAL_FUNC void fp_qvm_impl(
   int out_col = pack_factor * tn * (tid.y * num_simdgroups + simd_gid);
   ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w;
   scales += out_col / group_size + simd_lid * out_vec_size_g;
-  x += tid.x * in_vec_stride + simd_lid;
+  x += tid.x * in_vec_size + simd_lid;
   y += tid.x * out_vec_size + out_col;
 
   if (out_col >= out_vec_size) {
@@ -1123,16 +1122,7 @@ template <typename T, const int group_size, int bits, bool batched>
         tid);
   }
   fp_qvm_impl<T, group_size, bits>(
-      w,
-      scales,
-      x,
-      y,
-      in_vec_size,
-      out_vec_size,
-      in_vec_size,
-      tid,
-      simd_gid,
-      simd_lid);
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
 }
 
 template <typename T, const int group_size, int bits, int split_k = 32>
@@ -1174,20 +1164,8 @@ template <typename T, const int group_size, int bits, int split_k = 32>
   int in_vec_size_adj =
       tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size;
 
-  // The in_vec_stride is the full K dimension, not the partition size
-  int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size;
-
   fp_qvm_impl<T, group_size, bits>(
-      w,
-      scales,
-      x,
-      y,
-      in_vec_size_adj,
-      out_vec_size,
-      in_vec_stride,
-      tid,
-      simd_gid,
-      simd_lid);
+      w, scales, x, y, in_vec_size_adj, out_vec_size, tid, simd_gid, simd_lid);
 }
 
 template <
@@ -1445,16 +1423,7 @@ template <typename T, int group_size, int bits>
       s_strides,
       tid);
   fp_qvm_impl<T, group_size, bits>(
-      w,
-      scales,
-      x,
-      y,
-      in_vec_size,
-      out_vec_size,
-      in_vec_size,
-      tid,
-      simd_gid,
-      simd_lid);
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
 }
 
 template <

mlx/backend/metal/kernels/indexing/scatter.h

@@ -80,9 +80,7 @@ template <
     uint3 gsize [[threads_per_grid]]) {
   Op op;
 
-  IdxT idx =
-      (IdxT(gid.z) * IdxT(gsize.y) + IdxT(gid.y)) * IdxT(gsize.x) * NWORK +
-      IdxT(gid.x) * NWORK;
+  IdxT idx = IdxT(gid.z) * gsize.y + gid.y * gsize.x + gid.x * NWORK;
   IdxT out_idx;
   IdxT update_idx;
 

mlx/backend/metal/kernels/quantized.h

@@ -983,7 +983,6 @@ METAL_FUNC void qvm_impl(
     device T* y,
     const int in_vec_size,
     const int out_vec_size,
-    const int in_vec_stride,
     uint3 tid [[threadgroup_position_in_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
@@ -1017,7 +1016,7 @@ METAL_FUNC void qvm_impl(
   ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w;
   scales += out_col / group_size + simd_lid * out_vec_size_g;
   biases += out_col / group_size + simd_lid * out_vec_size_g;
-  x += tid.x * in_vec_stride + simd_lid;
+  x += tid.x * in_vec_size + simd_lid;
   y += tid.x * out_vec_size + out_col;
 
   if (out_col >= out_vec_size) {
@@ -1644,7 +1643,6 @@ template <typename T, const int group_size, const int bits, bool batched>
       y,
       in_vec_size,
       out_vec_size,
-      in_vec_size,
       tid,
       simd_gid,
       simd_lid);
@@ -1693,9 +1691,6 @@ template <typename T, const int group_size, const int bits, int split_k = 32>
   int in_vec_size_adj =
       tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size;
 
-  // The in_vec_stride is the full K dimension, not the partition size
-  int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size;
-
   qvm_impl<T, group_size, bits>(
       w,
       scales,
@@ -1704,7 +1699,6 @@ template <typename T, const int group_size, const int bits, int split_k = 32>
       y,
       in_vec_size_adj,
       out_vec_size,
-      in_vec_stride,
       tid,
       simd_gid,
       simd_lid);
@@ -2083,7 +2077,6 @@ template <typename T, int group_size, int bits>
       y,
       in_vec_size,
       out_vec_size,
-      in_vec_size,
       tid,
       simd_gid,
       simd_lid);

python/tests/test_array.py

@@ -1550,30 +1550,6 @@ class TestArray(mlx_tests.MLXTestCase):
         a = a.at[:, :, :].add(update)
         self.assertEqualArray(a, update)
 
-    def test_slice_update_contiguous_2d(self):
-        for shape in [(32, 32), (64, 64), (17, 33), (128, 128)]:
-            for upd_shape in [(16, 16), (8, 8), (4, 4)]:
-                if upd_shape[0] > shape[0] or upd_shape[1] > shape[1]:
-                    continue
-                y = mx.zeros(shape)
-                x = mx.random.normal(upd_shape)
-                z = y.at[: upd_shape[0], : upd_shape[1]].add(x)
-                diff = z[: upd_shape[0], : upd_shape[1]] - x
-                self.assertTrue(mx.allclose(diff, mx.zeros_like(diff)))
-
-        # Test non-zero offset slice update
-        y = mx.zeros((32, 32))
-        x = mx.random.normal((16, 16))
-        z = y.at[16:, 16:].add(x)
-        self.assertTrue(mx.allclose(z[16:, 16:] - x, mx.zeros_like(x)))
-
-        # Test with size divisible by 4, 2, and odd
-        for cols in [32, 18, 15]:
-            y = mx.zeros((32, cols))
-            x = mx.random.normal((16, cols))
-            z = y.at[:16, :].add(x)
-            self.assertTrue(mx.allclose(z[:16, :] - x, mx.zeros_like(x)))
-
     def test_slice_negative_step(self):
         a_np = np.arange(20)
         a_mx = mx.array(a_np)

python/tests/test_quantized.py

@@ -470,30 +470,6 @@ class TestQuantized(mlx_tests.MLXTestCase):
         self.assertEqual(y_q.shape, y_hat.shape)
         self.assertLess((y_q - y_hat).abs().max(), 2e-3)
 
-    def test_qvm_splitk_multi_row(self):
-        # Test qvm split_k with M > 1 to ensure the x row stride is correct
-        key = mx.random.key(0)
-        k1, k2 = mx.random.split(key)
-        tests = product(
-            [64, 32],  # group_size
-            [4, 8],  # bits
-            [128],  # out dim (N)
-            [2048, 4096],  # in dim (K) >= 1024 to trigger split_k
-            [2, 3],  # M (multiple rows)
-        )
-        for group_size, bits, N, K, M in tests:
-            with self.subTest(M=M, K=K, N=N, group_size=group_size, bits=bits):
-                x = 1e-1 * mx.random.normal(shape=(M, K), key=k1)
-                w = 1e-1 * mx.random.normal(shape=(K, N), key=k2)
-                w_q, scales, biases = mx.quantize(w, group_size, bits)
-                w_hat = mx.dequantize(w_q, scales, biases, group_size, bits)
-                y_q = mx.quantized_matmul(
-                    x, w_q, scales, biases, False, group_size, bits
-                )
-                y_hat = x @ w_hat
-                self.assertEqual(y_q.shape, y_hat.shape)
-                self.assertLess((y_q - y_hat).abs().max(), 2e-3)
-
     def test_fp_qvm(self):
         key = mx.random.key(0)
         k1, k2 = mx.random.split(key)

setup.py

@@ -234,7 +234,7 @@ if __name__ == "__main__":
             "ml_dtypes",
             "numpy>=2",
             "pre-commit",
-            "psutil",
+            "psutil>=7.2",
             "torch>=2.9",
             "typing_extensions",
         ],