zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-05 11:28:12 +08:00
Code Issues Actions 1 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

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