zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-01 16:58:08 +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)