Slice update with operation (#3266)

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)