(self)
| 628 | self.assertEqual(out.shape, (4, 5, 1)) |
| 629 | |
| 630 | def test_vmap_split_vmap(self): |
| 631 | def fun(x): |
| 632 | a, b = mx.split(x, 2, 1) |
| 633 | return mx.concatenate([b, a], 1) |
| 634 | |
| 635 | x = mx.ones((5, 6, 7)) |
| 636 | y = mx.ones((5, 4, 6, 7)) |
| 637 | fx = fun(x) |
| 638 | fy = mx.vmap(fun, in_axes=1)(y) |
| 639 | self.assertEqual(fx.shape, (5, 6, 7)) |
| 640 | self.assertEqual(fy.shape, (4, 5, 6, 7)) |
| 641 | |
| 642 | def test_leaks(self): |
| 643 | gc.collect() |