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

[Fix] expand axes for dimension with integer indices in mlx_slice_update (#1035)

Browse Source 
* Not sure if this is correct

* Format

* Edit tests

* Add negative test

* Format

* add one more test

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Jacket
2024-04-29 09:57:28 -05:00
committed by GitHub
parent c4a471c99d
commit 490c0c4fdc
2 changed files with 12 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/tests/test_array.py
View File

@@ -1194,6 +1194,10 @@ class TestArray(mlx_tests.MLXTestCase):
a = mx.zeros((2, 2))
a[0, 0, 0] = 1
with self.assertRaises(ValueError):
a = mx.zeros((5, 4, 3))
a[:, 0] = mx.ones((5, 1, 3))
check_slices(np.zeros((2, 2, 2, 2)), 1, None, Ellipsis, None)
check_slices(
np.zeros((2, 2, 2, 2)), 1, np.array([0, 1]), Ellipsis, np.array([0, 1])
@@ -1251,6 +1255,13 @@ class TestArray(mlx_tests.MLXTestCase):
slice(None, None, 2),
)
check_slices(np.zeros((5, 4, 3)), np.ones((5, 3)), slice(None), 0)
check_slices(np.zeros((5, 4, 3)), np.ones((5, 1, 3)), slice(None), slice(0, 1))
check_slices(
np.ones((3, 4, 4, 4)), np.zeros((4, 4)), 0, slice(0, 4), 3, slice(0, 4)
)
def test_array_at(self):
a = mx.array(1)
a = a.at[None].add(1)