L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Slice update with operation (#3266)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2026-03-18 06:18:02 -07:00
committed by GitHub
parent e353be8235
commit 7bc61cceed
18 changed files with 1475 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_array.py
+100
View File
@@ -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
+104 -31
View File
@@ -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):