Add test

.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: |
-        uv pip install build
+        pip install build
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
 

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

@@ -4,72 +4,61 @@ description: 'Build and test MLX on macOS'
 runs:
   using: "composite"
   steps:
-    - name: Install Python package
+    - name: Install dependencies
       env:
         DEBUG: 1
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      shell: bash
+      shell: bash -l {0}
       run: |
-        echo "::group::Install Python package"
-        uv pip install -e ".[dev]" -v
-        echo "::endgroup::"
+        pip install --upgrade pip
+        pip install cmake setuptools typing_extensions
+        pip install -e ".[dev]" -v
 
     - name: Install tests dependencies
-      shell: bash
+      shell: bash -l {0}
       run: |
-        echo "::group::Install tests dependencies"
-        uv pip install tensorflow
-        echo "::endgroup::"
+        pip install tensorflow
 
     - name: Run Python tests
-      shell: bash
+      shell: bash -l {0}
       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
+      shell: bash -l {0}
       run: |
-        echo "::group::Build example extension"
         cd examples/extensions
-        uv pip install -r requirements.txt
-        uv run --no-project setup.py build_ext --inplace
-        uv run --no-project test.py
-        echo "::endgroup::"
-
+        pip install -r requirements.txt
+        python setup.py build_ext --inplace
+        python test.py
+    
     - name: Build CPP only
-      shell: bash
+      shell: bash -l {0}
       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
+      shell: bash -l {0}
       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
+      shell: bash -l {0}
       run: |
-        echo "::group::Build small binary with JIT"
         mkdir -p build
         cd build
         cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
@@ -79,18 +68,15 @@ 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
+      shell: bash -l {0}
       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" \
-          uv pip install -e . -v
+          pip install -e . -v
         python -m unittest discover -v python/tests
-        echo "::endgroup::"

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

@@ -13,20 +13,12 @@ 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: 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::"
+    - uses: conda-incubator/setup-miniconda@v3
+      with:
+        miniconda-version: "latest"
+        python-version: ${{ inputs.python-version }}

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

@@ -9,21 +9,7 @@ inputs:
 runs:
   using: "composite"
   steps:
-    # FIXME: The distributed tests fail with free-threading Python.
-    - name: Check free-threading Python
-      id: is-free-threading
-      shell: bash
-      run: |
-        echo "::group::Check free-threading Python"
-        if python -VV 2>&1 | grep "free-threading"; then
-          echo "result=true" >> $GITHUB_OUTPUT
-        else
-          echo "result=false" >> $GITHUB_OUTPUT
-        fi
-        echo "::endgroup::"
-
     - name: Run MPI tests
-      if: ${{ steps.is-free-threading.outputs.result == 'false' }}
       shell: bash
       run: |
         echo "::group::MPI tests"
@@ -31,7 +17,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run distributed tests
-      if: ${{ steps.is-free-threading.outputs.result == 'false' && inputs.has-gpu == 'false' }}
+      if: ${{ inputs.has-gpu == 'false' }}
       shell: bash
       run: |
         echo "::group::Distributed tests"

.github/workflows/nightly.yml

@@ -41,7 +41,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ["3.11", "3.12", "3.13", "3.14", "3.14t"]
+        python_version: ["3.11", "3.12", "3.13", "3.14"]
         runner:
           - ubuntu-22.04
           - ubuntu-22.04-arm
@@ -59,7 +59,7 @@ jobs:
     if: github.repository == 'ml-explore/mlx'
     strategy:
       matrix:
-        python-version: ["3.10", "3.13", "3.14t"]
+        python-version: ["3.10", "3.13"]
     runs-on: [self-hosted, macos]
     steps:
       - uses: actions/checkout@v6

.github/workflows/release.yml

@@ -47,7 +47,7 @@ jobs:
     if: github.repository == 'ml-explore/mlx'
     strategy:
       matrix:
-        python_version: ["3.10", "3.11", "3.12", "3.13", "3.13t", "3.14", "3.14t"]
+        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
         arch: ['x86_64', 'aarch64']
     runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
     env:
@@ -83,7 +83,7 @@ jobs:
     if: github.repository == 'ml-explore/mlx'
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13", "3.13t", "3.14", "3.14t"]
+        python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
     runs-on: [self-hosted, macos]
     env:
       PYPI_RELEASE: 1
@@ -93,8 +93,13 @@ jobs:
       - uses: ./.github/actions/setup-macos
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Install Python package
-        run: uv pip install -e . -v
+
+      - name: Install dependencies
+        shell: bash -l {0}
+        run: |
+          pip install --upgrade pip
+          pip install cmake setuptools typing_extensions
+          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.31.2
+mlx>=0.21.0
 nanobind==2.12.0

mlx/backend/metal/device_info.cpp

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

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

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

mlx/backend/metal/rope.cpp

@@ -133,8 +133,9 @@ void RoPE::eval_gpu(
   if (single) {
     compute_encoder.set_bytes(out_strides, 1, 4);
     uint32_t dim0 = dims_ / 2;
-    group_dims = get_block_dims(dim0, N, 1);
-    grid_dims = MTL::Size(dim0, N, 1);
+    uint32_t dim1 = B * N;
+    group_dims = get_block_dims(dim0, dim1, 1);
+    grid_dims = MTL::Size(dim0, dim1, 1);
   } else {
     compute_encoder.set_bytes(strides, 3, 4);
     compute_encoder.set_bytes(out_strides, 3, 5);

python/src/CMakeLists.txt

@@ -2,7 +2,6 @@ nanobind_add_module(
   core
   NB_STATIC
   STABLE_ABI
-  FREE_THREADED
   LTO
   NOMINSIZE
   NB_DOMAIN

python/tests/test_array.py

@@ -1550,6 +1550,30 @@ 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_fast.py

@@ -365,6 +365,28 @@ class TestFast(mlx_tests.MLXTestCase):
         rx = rope_orig(x, dims, traditional, base, scale, offset)
         self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5)
 
+    def test_rope_single_batch(self):
+        base = 10000.0
+        scale = 1.0
+        offset = 5
+
+        for traditional in [True, False]:
+            for B in [2, 4, 8]:
+                for n_head in [1, 4, 7]:
+                    for dims in [64, 128]:
+                        x = mx.random.uniform(shape=(B, n_head, 1, dims))
+                        mx.eval(x)
+                        rx_fast = mx.fast.rope(
+                            x,
+                            dims,
+                            traditional=traditional,
+                            base=base,
+                            scale=scale,
+                            offset=offset,
+                        )
+                        rx = rope_orig(x, dims, traditional, base, scale, offset)
+                        self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5)
+
     def test_rope_with_large_offset(self):
         x = mx.random.normal(shape=(1, 1, 1024, 32))
         rx_fp32 = mx.fast.rope(

setup.py

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