(self)
| 761 | self.assertTrue(mx.allclose(mx.logsumexp(x), logsumexp(x))) |
| 762 | |
| 763 | def test_mean(self): |
| 764 | x = mx.array( |
| 765 | [ |
| 766 | [1, 2], |
| 767 | [3, 4], |
| 768 | ] |
| 769 | ) |
| 770 | self.assertEqual(mx.mean(x).item(), 2.5) |
| 771 | y = mx.mean(x, keepdims=True) |
| 772 | self.assertEqual(y, mx.array(2.5)) |
| 773 | self.assertEqual(y.shape, (1, 1)) |
| 774 | |
| 775 | self.assertEqual(mx.mean(x, axis=0).tolist(), [2, 3]) |
| 776 | self.assertEqual(mx.mean(x, axis=1).tolist(), [1.5, 3.5]) |
| 777 | |
| 778 | def test_median(self): |
| 779 | x = mx.array([]) |