| 2416 | self.assertEqual(x.flatten(end_axis=1).shape, (2 * 3, 4)) |
| 2417 | |
| 2418 | def test_clip(self): |
| 2419 | a = np.array([1, 4, 3, 8, 5], np.int32) |
| 2420 | expected = np.clip(a, 2, 6) |
| 2421 | clipped = mx.clip(mx.array(a), 2, 6) |
| 2422 | self.assertTrue(np.array_equal(clipped, expected)) |
| 2423 | |
| 2424 | a = np.array([-1, 1, 0, 5], np.int32) |
| 2425 | expected = np.clip(a, 0, None) |
| 2426 | clipped = mx.clip(mx.array(a), 0, None) |
| 2427 | self.assertTrue(np.array_equal(clipped, expected)) |
| 2428 | |
| 2429 | a = np.array([2, 3, 4, 5], np.int32) |
| 2430 | expected = np.clip(a, None, 4) |
| 2431 | clipped = mx.clip(mx.array(a), None, 4) |
| 2432 | self.assertTrue(np.array_equal(clipped, expected)) |
| 2433 | |
| 2434 | mins = np.array([3, 1, 5, 5]) |
| 2435 | a = np.array([2, 3, 4, 5], np.int32) |
| 2436 | expected = np.clip(a, mins, 4) |
| 2437 | clipped = mx.clip(mx.array(a), mx.array(mins), 4) |
| 2438 | self.assertTrue(np.array_equal(clipped, expected)) |
| 2439 | |
| 2440 | maxs = np.array([5, -1, 2, 9]) |
| 2441 | a = np.array([2, 3, 4, 5], np.int32) |
| 2442 | expected = np.clip(a, mins, maxs) |
| 2443 | clipped = mx.clip(mx.array(a), mx.array(mins), mx.array(maxs)) |
| 2444 | self.assertTrue(np.array_equal(clipped, expected)) |
| 2445 | |
| 2446 | # Check clip output types |
| 2447 | a = mx.array([1, 2, 3], mx.int16) |
| 2448 | out_t = mx.clip(a, a_min=0, a_max=5).dtype |
| 2449 | self.assertEqual(out_t, mx.int16) |
| 2450 | |
| 2451 | out_t = mx.clip(a, a_min=0.0, a_max=5).dtype |
| 2452 | self.assertEqual(out_t, mx.float32) |
| 2453 | |
| 2454 | a = mx.array([1, 2, 3], mx.float16) |
| 2455 | out_t = mx.clip(a, a_min=0.0, a_max=5).dtype |
| 2456 | self.assertEqual(out_t, mx.float16) |
| 2457 | |
| 2458 | a = mx.array([1, 2, 3], mx.float16) |
| 2459 | out_t = mx.clip(a, a_min=0.0, a_max=mx.array(1.0)).dtype |
| 2460 | self.assertEqual(out_t, mx.float32) |
| 2461 | |
| 2462 | def test_linspace(self): |
| 2463 | # Test default num = 50 |