| 946 | self.assertEqualArray(y, x - 1) |
| 947 | |
| 948 | def test_indexing(self): |
| 949 | # Only ellipsis is a no-op |
| 950 | a_mlx = mx.array([1])[...] |
| 951 | self.assertEqual(a_mlx.shape, (1,)) |
| 952 | self.assertEqual(a_mlx.item(), 1) |
| 953 | |
| 954 | # Basic content check, slice indexing |
| 955 | a_npy = np.arange(64, dtype=np.float32) |
| 956 | a_mlx = mx.array(a_npy) |
| 957 | a_sliced_mlx = a_mlx[2:50:4] |
| 958 | a_sliced_npy = np.asarray(a_sliced_mlx) |
| 959 | self.assertTrue(np.array_equal(a_sliced_npy, a_npy[2:50:4])) |
| 960 | |
| 961 | # Basic content check, mlx array indexing |
| 962 | a_npy = np.arange(64, dtype=np.int32) |
| 963 | a_npy = a_npy.reshape((8, 8)) |
| 964 | a_mlx = mx.array(a_npy) |
| 965 | idx_npy = np.array([0, 1, 2, 7, 5], dtype=np.uint32) |
| 966 | idx_mlx = mx.array(idx_npy) |
| 967 | a_sliced_mlx = a_mlx[idx_mlx] |
| 968 | a_sliced_npy = np.asarray(a_sliced_mlx) |
| 969 | self.assertTrue(np.array_equal(a_sliced_npy, a_npy[idx_npy])) |
| 970 | |
| 971 | # Basic content check, int indexing |
| 972 | a_sliced_mlx = a_mlx[5] |
| 973 | a_sliced_npy = np.asarray(a_sliced_mlx) |
| 974 | self.assertTrue(np.array_equal(a_sliced_npy, a_npy[5])) |
| 975 | self.assertEqual(len(a_sliced_npy.shape), len(a_npy[5].shape)) |
| 976 | self.assertEqual(len(a_sliced_npy.shape), 1) |
| 977 | self.assertEqual(a_sliced_npy.shape[0], a_npy[5].shape[0]) |
| 978 | |
| 979 | # Basic content check, negative indexing |
| 980 | a_sliced_mlx = a_mlx[-1] |
| 981 | self.assertTrue(np.array_equal(a_sliced_mlx, a_npy[-1])) |
| 982 | |
| 983 | # NumPy integer scalar indexing |
| 984 | a_sliced_mlx = a_mlx[np.int64(5)] |
| 985 | a_sliced_npy = np.asarray(a_sliced_mlx) |
| 986 | self.assertTrue(np.array_equal(a_sliced_npy, a_npy[np.int64(5)])) |
| 987 | |
| 988 | # Basic content check, empty index |
| 989 | a_sliced_mlx = a_mlx[()] |
| 990 | a_sliced_npy = np.asarray(a_sliced_mlx) |
| 991 | self.assertTrue(np.array_equal(a_sliced_npy, a_npy[()])) |
| 992 | |
| 993 | # Basic content check, new axis |
| 994 | a_sliced_mlx = a_mlx[None] |
| 995 | a_sliced_npy = np.asarray(a_sliced_mlx) |
| 996 | self.assertTrue(np.array_equal(a_sliced_npy, a_npy[None])) |
| 997 | |
| 998 | a_sliced_mlx = a_mlx[:, None] |
| 999 | a_sliced_npy = np.asarray(a_sliced_mlx) |
| 1000 | self.assertTrue(np.array_equal(a_sliced_npy, a_npy[:, None])) |
| 1001 | |
| 1002 | # Multi dim indexing, all ints |
| 1003 | self.assertEqual(a_mlx[0, 0].item(), 0) |
| 1004 | self.assertEqual(a_mlx[0, 0].ndim, 0) |
| 1005 | |