| 605 | self.assertEqual(x.swapaxes(0, 2).shape, (4, 3, 2)) |
| 606 | |
| 607 | def test_sum(self): |
| 608 | x = mx.array( |
| 609 | [ |
| 610 | [1, 2], |
| 611 | [3, 3], |
| 612 | ] |
| 613 | ) |
| 614 | self.assertEqual(mx.sum(x).item(), 9) |
| 615 | y = mx.sum(x, keepdims=True) |
| 616 | self.assertEqual(y, mx.array(9)) |
| 617 | self.assertEqual(y.shape, (1, 1)) |
| 618 | |
| 619 | self.assertEqual(mx.sum(x, axis=0).tolist(), [4, 5]) |
| 620 | self.assertEqual(mx.sum(x, axis=1).tolist(), [3, 6]) |
| 621 | |
| 622 | x_npy = np.arange(3 * 5 * 4 * 7).astype(np.float32) |
| 623 | x_npy = np.reshape(x_npy, (3, 5, 4, 7)) |
| 624 | x_mlx = mx.array(x_npy) |
| 625 | |
| 626 | for axis in (None, 0, 1, 2, 3, (0, 1), (2, 3), (1, 2, 3)): |
| 627 | sum_npy = np.sum(x_npy, axis=axis) |
| 628 | sum_mlx = np.asarray(mx.sum(x_mlx, axis=axis)) |
| 629 | self.assertListEqual(list(sum_npy.shape), list(sum_mlx.shape)) |
| 630 | self.assertTrue(np.all(sum_npy == sum_mlx)) |
| 631 | |
| 632 | x_npy = np.array([1.0, 2.0, 3.0, 4.0]).astype(np.float32) |
| 633 | x_mlx = mx.array(x_npy) |
| 634 | |
| 635 | y_npy = x_npy[0:4:2] |
| 636 | y_npy = np.broadcast_to(y_npy, (2, 2)) |
| 637 | |
| 638 | y_mlx = x_mlx[0:4:2] |
| 639 | y_mlx = mx.broadcast_to(y_mlx, (2, 2)) |
| 640 | |
| 641 | for axis in (None, 0, 1, (0, 1)): |
| 642 | sum_npy = np.sum(y_npy, axis=axis) |
| 643 | sum_mlx = np.asarray(mx.sum(y_mlx, axis=axis)) |
| 644 | self.assertListEqual(list(sum_npy.shape), list(sum_mlx.shape)) |
| 645 | self.assertTrue(np.all(sum_npy == sum_mlx)) |
| 646 | |
| 647 | x_npy = ( |
| 648 | np.arange(3 * 2 * 3 * 3 * 3 * 3) |
| 649 | .reshape(3, 2, 3, 3, 3, 3) |
| 650 | .astype(np.float32) |
| 651 | ) |
| 652 | x_mlx = mx.array(x_npy) |
| 653 | |
| 654 | y_mlx = x_mlx.sum(axis=(0, 1, 3, 4, 5)) |
| 655 | y_npy = x_npy.sum(axis=(0, 1, 3, 4, 5)) |
| 656 | |
| 657 | self.assertTrue(np.array_equal(y_mlx, y_npy)) |
| 658 | |
| 659 | def test_prod(self): |
| 660 | x = mx.array( |