Slice update with operation (#3266)

python/src/indexing.cpp

@@ -769,43 +769,53 @@ mlx_compute_scatter_args(
   throw std::invalid_argument("Cannot index mlx array using the given type.");
 }
 
-auto mlx_slice_update(
+std::tuple<std::optional<mx::array>, mx::Shape, mx::Shape, mx::Shape>
+mlx_compute_slice_update_args(
     const mx::array& src,
     const nb::object& obj,
     const ScalarOrArray& v) {
+  // Build the slice params
+  mx::Shape starts(src.ndim(), 0);
+  mx::Shape stops = src.shape();
+  mx::Shape strides(src.ndim(), 1);
+
   // Can't route to slice update if not slice, tuple, or int
   if (src.ndim() == 0 || nb::isinstance<nb::bool_>(obj) ||
       (!nb::isinstance<nb::slice>(obj) && !nb::isinstance<nb::tuple>(obj) &&
        !nb::isinstance<nb::int_>(obj))) {
-    return std::make_pair(false, src);
+    return std::make_tuple(
+        std::nullopt, std::move(starts), std::move(stops), std::move(strides));
   }
   if (nb::isinstance<nb::tuple>(obj)) {
     // Can't route to slice update if any arrays are present
     for (auto idx : nb::cast<nb::tuple>(obj)) {
       if (nb::isinstance<mx::array>(idx) || nb::isinstance<nb::list>(idx)) {
-        return std::make_pair(false, src);
+        return std::make_tuple(
+            std::nullopt,
+            std::move(starts),
+            std::move(stops),
+            std::move(strides));
       }
     }
   }
-  // Should be able to route to slice update
 
-  // Pre process tuple
-  auto upd = to_array(v, src.dtype());
+  // Should be able to route to slice update just extract the update value and
+  // and the slice arguments.
+
+  // Cast v to an array and ensure it is the right type
+  auto update = to_array(v, src.dtype());
 
   // Remove extra leading singletons dimensions from the update
   int s = 0;
-  for (; s < static_cast<int>(upd.ndim()) - 1 && upd.shape(s) == 1 &&
-       (upd.ndim() - s) > src.ndim();
+  for (; s < static_cast<int>(update.ndim()) - 1 && update.shape(s) == 1 &&
+       (update.ndim() - s) > src.ndim();
        s++) {
   };
   auto squeeze_axes = std::vector<int>(s);
   std::iota(squeeze_axes.begin(), squeeze_axes.end(), 0);
-  auto up = mx::squeeze(upd, squeeze_axes);
+  update = mx::squeeze(update, squeeze_axes);
 
-  // Build slice update params
-  mx::Shape starts(src.ndim(), 0);
-  mx::Shape stops = src.shape();
-  mx::Shape strides(src.ndim(), 1);
+  // Single int then make it a slice of size 1
   if (nb::isinstance<nb::int_>(obj)) {
     if (src.ndim() < 1) {
       std::ostringstream msg;
@@ -816,12 +826,11 @@ auto mlx_slice_update(
     idx = idx < 0 ? idx + stops[0] : idx;
     starts[0] = idx;
     stops[0] = idx + 1;
-    auto out = slice_update(
-        src, up, std::move(starts), std::move(stops), std::move(strides));
-    return std::make_pair(true, out);
+    return std::make_tuple(
+        update, std::move(starts), std::move(stops), std::move(strides));
   }
 
-  // If it's just a simple slice, just do a slice update and return
+  // Simple slice, just extract it into the first dim
   if (nb::isinstance<nb::slice>(obj)) {
     // Read slice arguments
     get_slice_params(
@@ -830,16 +839,14 @@ auto mlx_slice_update(
         strides[0],
         nb::cast<nb::slice>(obj),
         src.shape(0));
-
-    // Do slice update
-    auto out = slice_update(src, up, starts, stops, strides);
-    return std::make_pair(true, out);
+    return std::make_tuple(
+        update, std::move(starts), std::move(stops), std::move(strides));
   }
 
   // It must be a tuple
   auto entries = nb::cast<nb::tuple>(obj);
 
-  // Expand ellipses into a series of ':' slices
+  // Expand ellipsis into a series of ':' slices
   auto [non_none_indices, indices] = mlx_expand_ellipsis(src.shape(), entries);
 
   // Dimension check
@@ -851,15 +858,20 @@ auto mlx_slice_update(
 
   // If no non-None indices return the broadcasted update
   if (non_none_indices == 0) {
-    return std::make_pair(true, broadcast_to(up, src.shape()));
+    return std::make_tuple(
+        broadcast_to(update, src.shape()),
+        std::move(starts),
+        std::move(stops),
+        std::move(strides));
   }
 
+  // Parse the update slice
   int unspecified = src.ndim() - non_none_indices;
   std::vector<int> squeeze_dims;
   std::vector<int> expand_dims;
   for (int i = indices.size() - 1,
            ax = non_none_indices - 1,
-           upd_ax = upd.ndim() - unspecified - 1;
+           upd_ax = update.ndim() - unspecified - 1;
        i >= 0;
        --i) {
     auto& pyidx = indices[i];
@@ -887,11 +899,11 @@ auto mlx_slice_update(
       }
     }
   }
+  update = mx::squeeze(
+      mx::expand_dims(update, std::move(expand_dims)), std::move(squeeze_dims));
 
-  up = mx::squeeze(
-      mx::expand_dims(up, std::move(expand_dims)), std::move(squeeze_dims));
-  auto out = slice_update(src, up, starts, stops, strides);
-  return std::make_pair(true, out);
+  return std::make_tuple(
+      update, std::move(starts), std::move(stops), std::move(strides));
 }
 
 std::optional<mx::array> extract_boolean_mask(const nb::object& obj) {
@@ -921,9 +933,11 @@ void mlx_set_item(
     mx::array& src,
     const nb::object& obj,
     const ScalarOrArray& v) {
-  auto [success, out] = mlx_slice_update(src, obj, v);
-  if (success) {
-    src.overwrite_descriptor(out);
+  auto [update, starts, stops, strides] =
+      mlx_compute_slice_update_args(src, obj, v);
+  if (update) {
+    src.overwrite_descriptor(
+        slice_update(src, *update, starts, stops, strides));
     return;
   }
 
@@ -947,6 +961,12 @@ mx::array mlx_add_item(
     const mx::array& src,
     const nb::object& obj,
     const ScalarOrArray& v) {
+  auto [update, starts, stops, strides] =
+      mlx_compute_slice_update_args(src, obj, v);
+  if (update) {
+    return slice_update_add(src, *update, starts, stops, strides);
+  }
+
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
     return scatter_add(src, indices, updates, axes);
@@ -959,6 +979,12 @@ mx::array mlx_subtract_item(
     const mx::array& src,
     const nb::object& obj,
     const ScalarOrArray& v) {
+  auto [update, starts, stops, strides] =
+      mlx_compute_slice_update_args(src, obj, v);
+  if (update) {
+    return slice_update_add(src, -(*update), starts, stops, strides);
+  }
+
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
     return scatter_add(src, indices, -updates, axes);
@@ -971,6 +997,12 @@ mx::array mlx_multiply_item(
     const mx::array& src,
     const nb::object& obj,
     const ScalarOrArray& v) {
+  auto [update, starts, stops, strides] =
+      mlx_compute_slice_update_args(src, obj, v);
+  if (update) {
+    return slice_update_prod(src, *update, starts, stops, strides);
+  }
+
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
     return scatter_prod(src, indices, updates, axes);
@@ -983,6 +1015,12 @@ mx::array mlx_divide_item(
     const mx::array& src,
     const nb::object& obj,
     const ScalarOrArray& v) {
+  auto [update, starts, stops, strides] =
+      mlx_compute_slice_update_args(src, obj, v);
+  if (update) {
+    return slice_update_prod(src, reciprocal(*update), starts, stops, strides);
+  }
+
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
     return scatter_prod(src, indices, reciprocal(updates), axes);
@@ -995,6 +1033,12 @@ mx::array mlx_maximum_item(
     const mx::array& src,
     const nb::object& obj,
     const ScalarOrArray& v) {
+  auto [update, starts, stops, strides] =
+      mlx_compute_slice_update_args(src, obj, v);
+  if (update) {
+    return slice_update_max(src, *update, starts, stops, strides);
+  }
+
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
     return scatter_max(src, indices, updates, axes);
@@ -1007,6 +1051,12 @@ mx::array mlx_minimum_item(
     const mx::array& src,
     const nb::object& obj,
     const ScalarOrArray& v) {
+  auto [update, starts, stops, strides] =
+      mlx_compute_slice_update_args(src, obj, v);
+  if (update) {
+    return slice_update_min(src, *update, starts, stops, strides);
+  }
+
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
     return scatter_min(src, indices, updates, axes);

python/tests/test_array.py

@@ -1419,6 +1419,106 @@ class TestArray(mlx_tests.MLXTestCase):
         src = src.at[0:1].add(update)
         self.assertTrue(mx.array_equal(src, mx.array([[2.0, 4.0]])))
 
+        # Test all array.at ops with slice-only indices
+        a = mx.random.uniform(shape=(10, 5, 2))
+        update = mx.ones((2, 5))
+        a[1:3, :, 0] = 0
+        a = a.at[1:3, :, 0].add(update)
+        self.assertEqualArray(a[1:3, :, 0], update)
+        a = a.at[1:3, :, 0].subtract(update)
+        self.assertEqualArray(a[1:3, :, 0], mx.zeros_like(update))
+        a = a.at[1:3, :, 0].add(2 * update)
+        self.assertEqualArray(a[1:3, :, 0], 2 * update)
+        a = a.at[1:3, :, 0].multiply(2 * update)
+        self.assertEqualArray(a[1:3, :, 0], 4 * update)
+        a = a.at[1:3, :, 0].divide(3 * update)
+        self.assertEqualArray(a[1:3, :, 0], (4 / 3) * update)
+        a[1:3, :, 0] = 5
+        update = mx.arange(10).reshape(2, 5)
+        a = a.at[1:3, :, 0].maximum(update)
+        self.assertEqualArray(a[1:3, :, 0], mx.maximum(a[1:3, :, 0], update))
+        a[1:3, :, 0] = 5
+        a = a.at[1:3, :, 0].minimum(update)
+        self.assertEqualArray(a[1:3, :, 0], mx.minimum(a[1:3, :, 0], update))
+
+    def test_array_at_slice_update_extensive(self):
+        # Test with transposed inputs
+        a = mx.zeros((4, 5))
+        update = mx.ones((5, 2)).T  # Shape (2, 5)
+        a = a.at[1:3, :].add(update)
+        self.assertEqualArray(a[1:3, :], update)
+
+        # Test with transposed updates on transposed slice
+        a = mx.zeros((5, 4))
+        update = mx.ones((2, 5))
+        a = a.at[:, 1:3].add(update.T)
+        self.assertEqualArray(a[:, 1:3], update.T)
+
+        # Test with slice of another array as update
+        source = mx.arange(20, dtype=mx.float32).reshape(4, 5)
+        a = mx.zeros((4, 5))
+        update = source[1:3, :]  # Shape (2, 5)
+        a = a.at[0:2, :].add(update)
+        self.assertEqualArray(a[0:2, :], source[1:3, :])
+
+        # Test with both input and update being slices
+        source = mx.arange(30, dtype=mx.float32).reshape(5, 6)
+        a = mx.zeros((5, 6))
+        a = a.at[1:4, 1:5].add(source[0:3, 0:4])
+        self.assertEqualArray(a[1:4, 1:5], source[0:3, 0:4])
+
+        # Test with transposed slice of another array
+        source = mx.arange(20, dtype=mx.float32).reshape(4, 5)
+        a = mx.zeros((5, 4))
+        update = source[1:3, :].T  # Shape (5, 2)
+        a = a.at[:, 1:3].add(update)
+        self.assertEqualArray(a[:, 1:3], update)
+
+        # Test with negative indexing in slices
+        a = mx.zeros((5, 5))
+        update = mx.ones((2, 5))
+        a = a.at[-3:-1, :].add(update)
+        self.assertEqualArray(a[-3:-1, :], update)
+
+        # Test with strided slices
+        a = mx.zeros((6, 6))
+        update = mx.ones((2, 3))
+        a = a.at[1:5:2, 0:6:2].add(update)
+        self.assertEqualArray(a[1:5:2, 0:6:2], update)
+
+        # Test with slice of transposed array
+        source = mx.arange(20, dtype=mx.float32).reshape(4, 5)
+        a = mx.zeros((5, 4))
+        update = source.T[:, 1:3]  # Shape (5, 2)
+        a = a.at[:, 1:3].add(update)
+        self.assertEqualArray(a[:, 1:3], update)
+
+        # Test with 3D arrays and transposed updates
+        a = mx.zeros((3, 4, 5))
+        update = mx.ones((4, 3, 5)).transpose(1, 0, 2)  # Shape (3, 4, 5)
+        a = a.at[:, :, :].add(update)
+        self.assertEqualArray(a, update)
+
+        # Test with slice of 3D array
+        source = mx.arange(60, dtype=mx.float32).reshape(3, 4, 5)
+        a = mx.zeros((3, 4, 5))
+        update = source[0:2, :, :]
+        a = a.at[1:3, :, :].add(update)
+        self.assertEqualArray(a[1:3, :, :], source[0:2, :, :])
+
+        # Test with mixed slice and index
+        a = mx.zeros((4, 5, 6))
+        update = mx.ones((2, 6))
+        a = a.at[1:3, 2, :].add(update)
+        self.assertEqualArray(a[1:3, 2, :], update)
+
+        # Test with update from strided slice
+        source = mx.arange(60, dtype=mx.float32).reshape(3, 4, 5)
+        a = mx.zeros((3, 2, 5))
+        update = source[:, ::2, :]  # Shape (3, 2, 5)
+        a = a.at[:, :, :].add(update)
+        self.assertEqualArray(a, update)
+
     def test_slice_negative_step(self):
         a_np = np.arange(20)
         a_mx = mx.array(a_np)

python/tests/test_autograd.py

@@ -300,65 +300,138 @@ class TestAutograd(mlx_tests.MLXTestCase):
             x[idx] = 2.0
             return x.sum()
 
-        dfdx = mx.grad(fun)(mx.array([1.0, 2.0, 3.0]), mx.array([1]))
-        self.assertTrue(mx.array_equal(dfdx, mx.array([1.0, 0.0, 1.0])))
+        dfdx = mx.grad(fun)(mx.array([1.0, 2.0, 3.0, 4.0]), mx.array([1, 3]))
+        self.assertTrue(mx.array_equal(dfdx, mx.array([1.0, 0.0, 1.0, 0.0])))
         self.assertEqual(dfdx.dtype, mx.float32)
 
-        y = mx.array([0.0, 1.0, 2.0])
+        y = mx.array([0.0, 1.0, 2.0, 3.0])
 
         def fun(x, idx):
             y[idx] = x
             return y.sum()
 
-        dfdx = mx.grad(fun)(mx.array([2.0]), mx.array([1]))
-        self.assertTrue(mx.array_equal(dfdx, mx.array([1.0])))
+        dfdx = mx.grad(fun)(mx.array([2.0, 3.0]), mx.array([1, 3]))
+        self.assertTrue(mx.array_equal(dfdx, mx.array([1.0, 1.0])))
         self.assertEqual(dfdx.dtype, mx.float32)
 
+    def test_scatter_add_vjp(self):
+        def fun(src, updates):
+            x = src.at[mx.array([1, 3])].add(updates)
+            return x
+
+        cotan = mx.array([4.0, 5.0, 6.0, 7.0])
+        updates = mx.array([1.0, 2.0])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
+        mx.eval(vjps)
+
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([5.0, 7.0])))
+
     def test_scatter_max_vjp(self):
         def fun(src, updates):
-            x = src.at[1].maximum(updates)
+            x = src.at[mx.array([1, 3])].maximum(updates)
             return x
 
-        cotan = mx.array([4.0, 5.0, 6.0])
-        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0]), mx.array([[3.0]])], [cotan])
+        cotan = mx.array([4.0, 5.0, 6.0, 7.0])
+        updates = mx.array([1.0, 2.0])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
         mx.eval(vjps)
 
-        # Update larger than value
-        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 6.0])))
-        self.assertTrue(mx.allclose(vjps[1], mx.array([5.0])))
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([0.0, 0.0])))
 
-        cotan = mx.array([[4.0], [5.0], [6.0]])
-        _, vjps = mx.vjp(
-            fun, [mx.array([[1.0], [2.0], [3.0]]), mx.array([[[2.0]]])], [cotan]
-        )
+        updates = mx.array([5.0, 6.0])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
         mx.eval(vjps)
 
-        # Update and value are equal
-        self.assertTrue(mx.allclose(vjps[0], mx.array([[4.0], [5.0], [6.0]])))
-        self.assertTrue(mx.allclose(vjps[1], mx.array([[[5.0]]])))
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 6.0, 0.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([5.0, 7.0])))
 
     def test_scatter_min_vjp(self):
         def fun(src, updates):
-            x = src.at[1].minimum(updates)
+            x = src.at[mx.array([1, 3])].minimum(updates)
             return x
 
-        cotan = mx.array([4.0, 5.0, 6.0])
-        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0]), mx.array([[3.0]])], [cotan])
+        cotan = mx.array([4.0, 5.0, 6.0, 7.0])
+        updates = mx.array([5.0, 6.0])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
         mx.eval(vjps)
 
-        # Update larger than value
-        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0])))
-        self.assertTrue(mx.allclose(vjps[1], mx.array([0.0])))
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([0.0, 0.0])))
 
-        cotan = mx.array([[4.0], [5.0], [6.0]])
-        _, vjps = mx.vjp(
-            fun, [mx.array([[1.0], [2.0], [3.0]]), mx.array([[[2.0]]])], [cotan]
-        )
+        updates = mx.array([1.0, 1.0])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
         mx.eval(vjps)
 
-        # Update and value are equal
-        self.assertTrue(mx.allclose(vjps[0], mx.array([[4.0], [5.0], [6.0]])))
-        self.assertTrue(mx.allclose(vjps[1], mx.array([[[5.0]]])))
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 6.0, 0.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([5.0, 7.0])))
+
+    def test_slice_update_max_vjp(self):
+        def fun(src, updates):
+            x = src.at[1:3].maximum(updates)
+            return x
+
+        cotan = mx.array([4.0, 5.0, 6.0, 7.0])
+        updates = mx.array([[1.0, 2.0]])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
+        mx.eval(vjps)
+
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([[0.0, 0.0]])))
+
+        updates = mx.array([[5.0, 6.0]])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
+        mx.eval(vjps)
+
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 0.0, 7.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([[5.0, 6.0]])))
+
+    def test_slice_update_min_vjp(self):
+        def fun(src, updates):
+            x = src.at[1:3].minimum(updates)
+            return x
+
+        cotan = mx.array([4.0, 5.0, 6.0, 7.0])
+        updates = mx.array([[5.0, 6.0]])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
+        mx.eval(vjps)
+
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([[0.0, 0.0]])))
+
+        updates = mx.array([[1.0, 1.0]])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
+        mx.eval(vjps)
+
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 0.0, 7.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([[5.0, 6.0]])))
+
+    def test_slice_update_add_vjp(self):
+        def fun(src, updates):
+            x = src.at[1:3].add(updates)
+            return x
+
+        cotan = mx.array([4.0, 5.0, 6.0, 7.0])
+        updates = mx.array([[1.0, 2.0]])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
+        mx.eval(vjps)
+
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([[5.0, 6.0]])))
+
+    def test_slice_update_multiply_vjp(self):
+        def fun(src, updates):
+            x = src.at[1:3].multiply(updates)
+            return x
+
+        cotan = mx.array([4.0, 5.0, 6.0, 7.0])
+        updates = mx.array([[2.0, 3.0]])
+        _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan])
+        mx.eval(vjps)
+
+        self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 10.0, 18.0, 7.0])))
+        self.assertTrue(mx.allclose(vjps[1], mx.array([[10.0, 18.0]])))
 
     def test_split_against_slice(self):
         def f_split(x):