Add test

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/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(