(arr_np, *idx_np)
| 1056 | |
| 1057 | # Multi dim indexing with multiple arrays |
| 1058 | def check_slices(arr_np, *idx_np): |
| 1059 | arr_mlx = mx.array(arr_np) |
| 1060 | idx_mlx = [ |
| 1061 | mx.array(idx) if isinstance(idx, np.ndarray) else idx for idx in idx_np |
| 1062 | ] |
| 1063 | slice_mlx = arr_mlx[tuple(idx_mlx)] |
| 1064 | self.assertTrue( |
| 1065 | np.array_equal(arr_np[tuple(idx_np)], arr_mlx[tuple(idx_mlx)]) |
| 1066 | ) |
| 1067 | |
| 1068 | a_np = np.arange(16).reshape(4, 4) |
| 1069 | check_slices(a_np, np.array([0, 1, 2, 3]), np.array([0, 1, 2, 3])) |