| 679 | self.assertEqual(mx.vmap(fun, in_axes=(2,))(x).shape, (4, 6)) |
| 680 | |
| 681 | def test_vmap_conv(self): |
| 682 | # vmap input only |
| 683 | x = mx.random.uniform(shape=(2, 2, 5, 4)) |
| 684 | w = mx.random.uniform(shape=(8, 3, 4)) |
| 685 | |
| 686 | expected = mx.stack([mx.conv1d(xi, w) for xi in x]) |
| 687 | out = mx.vmap(mx.conv1d, in_axes=(0, None))(x, w) |
| 688 | self.assertTrue(mx.allclose(expected, out)) |
| 689 | |
| 690 | x = mx.moveaxis(x, 0, 2) |
| 691 | out = mx.vmap(mx.conv1d, in_axes=(2, None))(x, w) |
| 692 | self.assertTrue(mx.allclose(expected, out)) |
| 693 | |
| 694 | # vmap weights only |
| 695 | x = mx.random.uniform(shape=(2, 5, 4)) |
| 696 | w = mx.random.uniform(shape=(3, 8, 3, 4)) |
| 697 | |
| 698 | expected = mx.stack([mx.conv1d(x, wi) for wi in w]) |
| 699 | out = mx.vmap(mx.conv1d, in_axes=(None, 0))(x, w) |
| 700 | self.assertTrue(mx.allclose(expected, out)) |
| 701 | |
| 702 | w = mx.moveaxis(w, 0, 1) |
| 703 | out = mx.vmap(mx.conv1d, in_axes=(None, 1))(x, w) |
| 704 | self.assertTrue(mx.allclose(expected, out)) |
| 705 | |
| 706 | # vmap weights and input |
| 707 | x = mx.random.uniform(shape=(3, 2, 5, 4)) |
| 708 | w = mx.random.uniform(shape=(3, 8, 3, 4)) |
| 709 | |
| 710 | expected = mx.stack([mx.conv1d(xi, wi) for xi, wi in zip(x, w)]) |
| 711 | out = mx.vmap(mx.conv1d, in_axes=(0, 0))(x, w) |
| 712 | self.assertTrue(mx.allclose(expected, out)) |
| 713 | |
| 714 | x = mx.random.uniform(shape=(2, 3, 5, 4)) |
| 715 | w = mx.random.uniform(shape=(8, 3, 4, 3)) |
| 716 | |
| 717 | expected = mx.stack([mx.conv1d(x[:, i], w[..., i]) for i in range(3)]) |
| 718 | out = mx.vmap(mx.conv1d, in_axes=(1, 3))(x, w) |
| 719 | self.assertTrue(mx.allclose(expected, out)) |
| 720 | |
| 721 | # Test with groups |
| 722 | x = mx.random.uniform(shape=(3, 2, 5, 8)) |
| 723 | w = mx.random.uniform(shape=(3, 2, 3, 4)) |
| 724 | |
| 725 | def gconv(x, w): |
| 726 | return mx.conv1d(x, w, groups=2) |
| 727 | |
| 728 | expected = mx.stack([gconv(xi, wi) for xi, wi in zip(x, w)]) |
| 729 | out = mx.vmap(gconv, in_axes=(0, 0))(x, w) |
| 730 | self.assertTrue(mx.allclose(expected, out)) |
| 731 | |
| 732 | def test_vmap_pad(self): |
| 733 | def pad2d(x, value=0.0): |