zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-04 15:04:40 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Dynamic slicing (#1741)

Browse Source 
* dynamic slice and slice update

* python bindings + tests + fix set item

* fix compile issue

* comment

* fix jit
This commit is contained in:
Awni Hannun
2025-01-07 14:02:16 -08:00
committed by GitHub
parent c9c81d0584
commit 516ded618b
27 changed files with 941 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/tests/test_array.py
View File

@@ -1327,6 +1327,16 @@ class TestArray(mlx_tests.MLXTestCase):
x[0, 0] = 1
self.assertTrue(mx.array_equal(x[0, 0], mx.ones((2, 2, 2, 2))))
a = mx.zeros((2, 2, 2))
with self.assertRaises(ValueError):
a[:, None, :] = mx.ones((2, 2, 2))
# Ok, doesn't throw
a[:, None, :] = mx.ones((2, 1, 2, 2))
a[:, None, :] = mx.ones((2, 2))
a[:, None, 0] = mx.ones((2,))
a[:, None, 0] = mx.ones((1, 2))
def test_array_at(self):
a = mx.array(1)
a = a.at[None].add(1)

13
python/tests/test_ops.py
View File

@@ -2769,6 +2769,19 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertEqual(mx.imag(z).dtype, mx.float32)
self.assertTrue(mx.array_equal(mx.imag(z), y))
def test_dynamic_slicing(self):
x = mx.random.randint(0, 100, shape=(4, 4, 4))
expected = x[1:, 2:, 3:]
out = mx.slice(x, mx.array([1, 2, 3]), (0, 1, 2), (3, 2, 1))
self.assertTrue(mx.array_equal(expected, out))
x = mx.zeros(shape=(4, 4, 4))
update = mx.random.randint(0, 100, shape=(3, 2, 1))
out = mx.slice_update(x, update, mx.array([1, 2, 3]), (0, 1, 2))
expected = mx.zeros_like(x)
expected[1:, 2:, 3:] = update
self.assertTrue(mx.array_equal(expected, out))
if __name__ == "__main__":
unittest.main()