| 712 | self.assertEqual(a.item(), b) |
| 713 | |
| 714 | def test_broadcast(self): |
| 715 | a_npy = np.reshape(np.arange(200), (10, 20)) |
| 716 | a_mlx = mx.array(a_npy) |
| 717 | |
| 718 | b_npy = np.broadcast_to(a_npy, (30, 10, 20)) |
| 719 | b_mlx = mx.broadcast_to(a_mlx, (30, 10, 20)) |
| 720 | self.assertListEqual(list(b_npy.shape), list(b_mlx.shape)) |
| 721 | self.assertTrue(np.array_equal(b_npy, b_mlx)) |
| 722 | |
| 723 | b_npy = np.broadcast_to(a_npy, (1, 10, 20)) |
| 724 | b_mlx = mx.broadcast_to(a_mlx, (1, 10, 20)) |
| 725 | self.assertListEqual(list(b_npy.shape), list(b_mlx.shape)) |
| 726 | self.assertTrue(np.array_equal(b_npy, b_mlx)) |
| 727 | |
| 728 | b_npy = np.broadcast_to(1, (10, 20)) |
| 729 | b_mlx = mx.broadcast_to(1, (10, 20)) |
| 730 | self.assertListEqual(list(b_npy.shape), list(b_mlx.shape)) |
| 731 | self.assertTrue(np.array_equal(b_npy, b_mlx)) |
| 732 | |
| 733 | def test_logsumexp(self): |
| 734 | def logsumexp(x, axes=None): |