(self)
| 669 | self.assertEqual(mem_pre, mem_post) |
| 670 | |
| 671 | def test_vmap_flatten(self): |
| 672 | def fun(x): |
| 673 | return mx.flatten(x, 0, 1) |
| 674 | |
| 675 | x = mx.zeros((2, 3, 4)) |
| 676 | |
| 677 | self.assertEqual(mx.vmap(fun)(x).shape, (2, 12)) |
| 678 | self.assertEqual(mx.vmap(fun, in_axes=(1,))(x).shape, (3, 8)) |
| 679 | self.assertEqual(mx.vmap(fun, in_axes=(2,))(x).shape, (4, 6)) |
| 680 | |
| 681 | def test_vmap_conv(self): |
| 682 | # vmap input only |