| 2316 | self.assertTrue(np.array_equal(y_np, y_mx)) |
| 2317 | |
| 2318 | def test_partition(self): |
| 2319 | shape = (3, 4, 5) |
| 2320 | for dtype in ("int32", "float32"): |
| 2321 | for axis in (None, 0, 1, 2): |
| 2322 | for kth in (-2, 0, 2): |
| 2323 | with self.subTest(dtype=dtype, axis=axis, kth=kth): |
| 2324 | np.random.seed(0) |
| 2325 | np_dtype = getattr(np, dtype) |
| 2326 | a_np = np.random.uniform(0, 100, size=shape).astype(np_dtype) |
| 2327 | a_mx = mx.array(a_np) |
| 2328 | |
| 2329 | b_np = np.partition(a_np, kth, axis=axis) |
| 2330 | b_mx = mx.partition(a_mx, kth, axis=axis) |
| 2331 | |
| 2332 | c_np = np.take(b_np, (kth,), axis=axis) |
| 2333 | c_mx = np.take(np.array(b_mx), (kth,), axis=axis) |
| 2334 | |
| 2335 | self.assertTrue(np.array_equal(c_np, c_mx)) |
| 2336 | self.assertEqual(b_mx.dtype, a_mx.dtype) |
| 2337 | |
| 2338 | if kth >= 0: |
| 2339 | top_k_mx = mx.topk(a_mx, kth, axis=axis) |
| 2340 | top_k_np = np.take( |
| 2341 | np.partition(a_np, -kth, axis=axis), (-kth,), axis=axis |
| 2342 | ) |
| 2343 | self.assertTrue(np.all(top_k_np <= top_k_mx)) |
| 2344 | self.assertEqual(top_k_mx.dtype, a_mx.dtype) |
| 2345 | N = a_mx.shape[axis] if axis is not None else a_mx.size |
| 2346 | M = top_k_mx.shape[axis or 0] |
| 2347 | self.assertEqual(M, (kth + N) % N) |
| 2348 | |
| 2349 | def test_argpartition(self): |
| 2350 | x = mx.broadcast_to(mx.array([1, 2, 3]), (2, 3)) |