| 730 | self.assertTrue(mx.allclose(expected, out)) |
| 731 | |
| 732 | def test_vmap_pad(self): |
| 733 | def pad2d(x, value=0.0): |
| 734 | return mx.pad(x, ((1, 2), (0, 1)), constant_values=value) |
| 735 | |
| 736 | x = mx.arange(24, dtype=mx.float32).reshape(2, 3, 4) |
| 737 | |
| 738 | expected = mx.stack([pad2d(xi) for xi in x]) |
| 739 | out = mx.vmap(pad2d, in_axes=0)(x) |
| 740 | self.assertTrue(mx.array_equal(out, expected)) |
| 741 | |
| 742 | expected = mx.stack([pad2d(x[:, i, :]) for i in range(x.shape[1])]) |
| 743 | out = mx.vmap(pad2d, in_axes=1)(x) |
| 744 | self.assertTrue(mx.array_equal(out, expected)) |
| 745 | |
| 746 | expected = mx.stack([pad2d(x[:, :, i]) for i in range(x.shape[2])], axis=2) |
| 747 | out = mx.vmap(pad2d, in_axes=-1, out_axes=-1)(x) |
| 748 | self.assertTrue(mx.array_equal(out, expected)) |
| 749 | |
| 750 | nested = mx.vmap(mx.vmap(lambda y: mx.pad(y, (1, 1)))) |
| 751 | out = nested(x) |
| 752 | expected = mx.pad(x, ((0, 0), (0, 0), (1, 1))) |
| 753 | self.assertTrue(mx.array_equal(out, expected)) |
| 754 | |
| 755 | out = mx.vmap( |
| 756 | lambda a, v: mx.pad(a, ((1, 1), (1, 1)), constant_values=v), |
| 757 | in_axes=(0, None), |
| 758 | )(x, mx.array(5.0)) |
| 759 | expected = mx.stack( |
| 760 | [mx.pad(xi, ((1, 1), (1, 1)), constant_values=mx.array(5.0)) for xi in x] |
| 761 | ) |
| 762 | self.assertTrue(mx.array_equal(out, expected)) |
| 763 | |
| 764 | pad_values = mx.array([3.0, 4.0]) |
| 765 | with self.assertRaises(ValueError): |
| 766 | mx.vmap(lambda a, v: mx.pad(a, ((1, 1), (1, 1)), constant_values=v))( |
| 767 | x, pad_values |
| 768 | ) |
| 769 | |
| 770 | def test_vmap_types(self): |
| 771 | |