(self)
| 760 | self.assertTrue(mx.array_equal(mx.compile(fn)(x), fn(x))) |
| 761 | |
| 762 | def test_max_into_equal(self): |
| 763 | x = mx.random.uniform(shape=(1, 2, 2)) |
| 764 | mx.eval(x) |
| 765 | |
| 766 | def fn(): |
| 767 | maxes = mx.max(x, axis=(1, 2), keepdims=True) |
| 768 | return x == maxes |
| 769 | |
| 770 | out = mx.compile(fn)() |
| 771 | expected = fn() |
| 772 | self.assertTrue(mx.array_equal(expected, out)) |
| 773 | |
| 774 | def test_dtypes(self): |
| 775 | x = mx.array([0, 1, 2, 3]) |