L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support indexing with any type which implmented __index__ (#3210)

Browse Source
This commit is contained in:
AN Long
2026-03-20 06:19:08 +08:00
committed by GitHub
parent 21c11fc9b0
commit 5fa1a8d59f
2 changed files with 54 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_array.py
+9 -1
View File
@@ -960,6 +960,11 @@ class TestArray(mlx_tests.MLXTestCase):
a_sliced_mlx = a_mlx[-1]
self.assertTrue(np.array_equal(a_sliced_mlx, a_npy[-1]))
# NumPy integer scalar indexing
a_sliced_mlx = a_mlx[np.int64(5)]
a_sliced_npy = np.asarray(a_sliced_mlx)
self.assertTrue(np.array_equal(a_sliced_npy, a_npy[np.int64(5)]))
# Basic content check, empty index
a_sliced_mlx = a_mlx[()]
a_sliced_npy = np.asarray(a_sliced_mlx)
@@ -1121,8 +1126,11 @@ class TestArray(mlx_tests.MLXTestCase):
a[-1] = 2
self.assertEqual(a.tolist(), [2, 2, 2])
a[np.int64(1)] = 9
self.assertEqual(a.tolist(), [2, 9, 2])
a[0] = mx.array([[[1]]])
self.assertEqual(a.tolist(), [1, 2, 2])
self.assertEqual(a.tolist(), [1, 9, 2])
a[:] = 0
self.assertEqual(a.tolist(), [0, 0, 0])