| 9 | |
| 10 | class TestVmap(mlx_tests.MLXTestCase): |
| 11 | def test_basics(self): |
| 12 | # Can't vmap over scalars |
| 13 | with self.assertRaises(ValueError): |
| 14 | mx.vmap(mx.exp)(mx.array(1.0)) |
| 15 | |
| 16 | # Invalid input |
| 17 | with self.assertRaises(ValueError): |
| 18 | mx.vmap(mx.exp)("hello") |
| 19 | |
| 20 | # Invalid axes |
| 21 | with self.assertRaises(ValueError): |
| 22 | mx.vmap(mx.exp, in_axes="hello")(mx.array([0, 1])) |
| 23 | |
| 24 | with self.assertRaises(ValueError): |
| 25 | mx.vmap(mx.exp, in_axes=2)(mx.array([0, 1])) |
| 26 | |
| 27 | with self.assertRaises(ValueError): |
| 28 | mx.vmap(mx.exp, out_axes="hello")(mx.array([0, 1])) |
| 29 | |
| 30 | with self.assertRaises(ValueError): |
| 31 | mx.vmap(mx.exp, out_axes=2)(mx.array([0, 1])) |
| 32 | |
| 33 | def test_unary(self): |
| 34 | ops = [ |