| 1167 | mx.kron(x, y) |
| 1168 | |
| 1169 | def test_take(self): |
| 1170 | # Shape: 4 x 3 x 2 |
| 1171 | l = [ |
| 1172 | [[1, 3], [-2, -2], [-3, -2]], |
| 1173 | [[2, 4], [-3, 2], [-4, -2]], |
| 1174 | [[2, 3], [2, 4], [2, 1]], |
| 1175 | [[1, -5], [3, -1], [2, 3]], |
| 1176 | ] |
| 1177 | |
| 1178 | a = mx.array(l) |
| 1179 | a_npy = np.array(l) |
| 1180 | |
| 1181 | indices = [0, -1] |
| 1182 | flatten_take = mx.take(a, mx.array(indices)).tolist() |
| 1183 | flatten_take_expected = np.take(a_npy, np.array(indices)).tolist() |
| 1184 | self.assertListEqual(flatten_take, flatten_take_expected) |
| 1185 | |
| 1186 | indices = [-1, 2, 0] |
| 1187 | axis_take = mx.take(a, mx.array(indices), axis=0).tolist() |
| 1188 | axis_take_expected = np.take(a_npy, np.array(indices), axis=0).tolist() |
| 1189 | self.assertListEqual(axis_take, axis_take_expected) |
| 1190 | |
| 1191 | indices = [0, 0, -2] |
| 1192 | axis_take = mx.take(a, mx.array(indices), axis=1).tolist() |
| 1193 | axis_take_expected = np.take(a_npy, np.array(indices), axis=1).tolist() |
| 1194 | self.assertListEqual(axis_take, axis_take_expected) |
| 1195 | |
| 1196 | indices = [0, -1, -1] |
| 1197 | axis_take = mx.take(a, mx.array(indices), axis=-1).tolist() |
| 1198 | axis_take_expected = np.take(a_npy, np.array(indices), axis=-1).tolist() |
| 1199 | self.assertListEqual(axis_take, axis_take_expected) |
| 1200 | |
| 1201 | a_npy = np.arange(8 * 8 * 8, dtype=np.int32) |
| 1202 | a_npy = a_npy.reshape((8, 8, 8)) |
| 1203 | idx_npy = np.arange(6, dtype=np.uint32) |
| 1204 | idx_npy = idx_npy.reshape((2, 3)) |
| 1205 | a_mlx = mx.array(a_npy) |
| 1206 | idx_mlx = mx.array(idx_npy) |
| 1207 | |
| 1208 | a_npy_taken = np.take(a_npy, idx_npy) |
| 1209 | a_mlx_taken = mx.take(a_mlx, idx_mlx) |
| 1210 | self.assertEqual(a_npy_taken.shape, a_mlx_taken.shape) |
| 1211 | self.assertListEqual(a_npy_taken.tolist(), a_mlx_taken.tolist()) |
| 1212 | |
| 1213 | a_npy_taken = np.take(a_npy, idx_npy, axis=0) |
| 1214 | a_mlx_taken = mx.take(a_mlx, idx_mlx, axis=0) |
| 1215 | self.assertEqual(a_npy_taken.shape, a_mlx_taken.shape) |
| 1216 | self.assertListEqual(a_npy_taken.tolist(), a_mlx_taken.tolist()) |
| 1217 | |
| 1218 | a_npy_taken = np.take(a_npy, idx_npy, axis=1) |
| 1219 | a_mlx_taken = mx.take(a_mlx, idx_mlx, axis=1) |
| 1220 | self.assertEqual(a_npy_taken.shape, a_mlx_taken.shape) |
| 1221 | self.assertListEqual(a_npy_taken.tolist(), a_mlx_taken.tolist()) |
| 1222 | |
| 1223 | a_npy_taken = np.take(a_npy, idx_npy, axis=2) |
| 1224 | a_mlx_taken = mx.take(a_mlx, idx_mlx, axis=2) |
| 1225 | self.assertEqual(a_npy_taken.shape, a_mlx_taken.shape) |
| 1226 | self.assertListEqual(a_npy_taken.tolist(), a_mlx_taken.tolist()) |