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

fix slice update indexing (#1053)

Browse Source
This commit is contained in:
Awni Hannun
2024-04-29 12:17:40 -07:00
committed by GitHub
parent 490c0c4fdc
commit 09f1777896
4 changed files with 38 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/tests/test_array.py
View File

@@ -1262,6 +1262,14 @@ class TestArray(mlx_tests.MLXTestCase):
np.ones((3, 4, 4, 4)), np.zeros((4, 4)), 0, slice(0, 4), 3, slice(0, 4)
)
x = mx.zeros((2, 3, 4, 5, 3))
x[..., 0] = 1.0
self.assertTrue(mx.array_equal(x[..., 0], mx.ones((2, 3, 4, 5))))
x = mx.zeros((2, 3, 4, 5, 3))
x[:, 0] = 1.0
self.assertTrue(mx.array_equal(x[:, 0], mx.ones((2, 4, 5, 3))))
def test_array_at(self):
a = mx.array(1)
a = a.at[None].add(1)