python_vmap

Author: xadupre

_doc/api/torch_export_patches/index.rst

@@ -10,6 +10,7 @@ onnx_diagnostic.torch_export_patches
     onnx_export_serialization
     patches/index
     patch_expressions
+    patch_helper
     patch_inputs
     patch_module
     patch_module_helper

_doc/api/torch_export_patches/patch_helper.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.torch_export_patches.patch_helper
+=================================================
+
+.. automodule:: onnx_diagnostic.torch_export_patches.patch_helper
+    :members:
+    :no-undoc-members:

_unittests/ut_torch_export_patches/test_patch_helper.py

@@ -0,0 +1,63 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, requires_torch
+from onnx_diagnostic.torch_export_patches.patch_helper import py_vmap
+
+
+class TestPatchHelper(ExtTestCase):
+    def test_vmap(self):
+        f = lambda x, y: x * y + 1  # noqa: E731
+        x = torch.tensor([1.0, 2.0, 3.0])
+        y = torch.tensor([0.1, 0.2, 0.3])
+        expected = torch.vmap(f)(x, y)
+        got = py_vmap(f)(x, y)
+        self.assertEqualArray(expected, got)
+
+    @requires_torch("2.9")
+    def test_export_vmap(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                f = lambda x, y: x * y + 1  # noqa: E731
+                return torch.vmap(f)(x, y)
+
+        x = torch.tensor([1.0, 2.0, 3.0])
+        y = torch.tensor([0.1, 0.2, 0.3])
+        DYN = torch.export.Dim.DYNAMIC
+        torch.export.export(Model(), (x, y), ({0: DYN}, {1: DYN}))
+
+    def test_export_py_vmap(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                f = lambda x, y: x * y + 1  # noqa: E731
+                return py_vmap(f)(x, y)
+
+        x = torch.tensor([1.0, 2.0, 3.0])
+        y = torch.tensor([0.1, 0.2, 0.3])
+        torch.export.export(Model(), (x, y))
+
+    def test_vmap_outdim(self):
+        f = lambda x: x**2  # noqa: E731
+        x = torch.randn(2, 5)
+        expected = torch.vmap(f, out_dims=1)(x)
+        got = py_vmap(f, out_dims=1)(x)
+        self.assertEqualArray(expected, got)
+
+    def test_vmap_dict(self):
+        f = lambda d: torch.dot(d["x"], d["y"])  # noqa: E731
+        x, y = torch.randn(2, 5), torch.randn(5)
+        input = {"x": x, "y": y}
+        _expected = torch.vmap(f, in_dims=({"x": 0, "y": None},))(input)
+        self.assertRaise(
+            lambda: py_vmap(f, in_dims=({"x": 0, "y": None},))(input), AssertionError
+        )
+        # self.assertEqualArray(_expected, got)
+
+    def test_vmap_tuple(self):
+        x, y = torch.randn(2, 5), torch.randn(5)
+        expected = torch.vmap(torch.dot, in_dims=(0, None))(x, y)
+        got = py_vmap(torch.dot, in_dims=(0, None))(x, y)
+        self.assertEqualArray(expected, got)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

onnx_diagnostic/torch_export_patches/eval/model_cases.py

@@ -1,5 +1,6 @@
 import numpy as np
 import torch
+from ..patch_helper import py_vmap
 
 DIM = torch.export.Dim
 DYN = torch.export.Dim.DYNAMIC
@@ -875,3 +876,21 @@ def forward(self, x):
 
     _inputs = [(torch.rand((4, 4)),), (torch.rand((5, 5)),)]
     _dynamic = {"x": {0: DIM("dx"), 1: DIM("dy")}}
+
+
+class Vmap(torch.nn.Module):
+    def forward(self, x, y):
+        f = lambda x, y: x * y + 1  # noqa: E731
+        return torch.vmap(f)(x, y)
+
+    _inputs = [(torch.tensor([1.0, 2.0, 3.0]), torch.tensor([0.1, 0.2, 0.3]))]
+    _dynamic = {"x": {0: DYN}, "y": {0: DYN}}
+
+
+class VmapPython(torch.nn.Module):
+    def forward(self, x, y):
+        f = lambda x, y: x * y + 1  # noqa: E731
+        return py_vmap(f)(x, y)
+
+    _inputs = [(torch.tensor([1.0, 2.0, 3.0]), torch.tensor([0.1, 0.2, 0.3]))]
+    _dynamic = {"x": {0: DYN}, "y": {0: DYN}}

onnx_diagnostic/torch_export_patches/patch_helper.py

@@ -0,0 +1,55 @@
+import torch
+from ..helpers import string_type
+
+
+def py_vmap(func, in_dims=0, out_dims=0):
+    """
+    Python implementation of :func:`torch.vmap`.
+    """
+
+    def wrapped(*args):
+        assert all(not isinstance(a, dict) for a in args), (
+            f"dictionaries are not implemented in "
+            f"args={string_type(args, with_shape=True)}"
+        )
+
+        in_dims_ = (
+            ([in_dims] * len(args))
+            if not isinstance(in_dims, (list, tuple))
+            else list(in_dims)
+        )
+        assert len(in_dims_) == len(args)
+
+        batch_size = None
+        batched_args = []
+        for arg, in_dim in zip(args, in_dims_):
+            if in_dim is None:
+                batched_args.append(arg)
+                continue
+
+            assert batch_size is None or batch_size == arg.size(in_dim), (
+                f"Unable to continue, batch_size={batch_size}, in_dim={in_dim}, "
+                f"arg.size(in_dim)={arg.size(in_dim)}"
+            )
+            if batch_size is None:
+                batch_size = arg.size(in_dim)
+            arg = arg.movedim(in_dim, 0)
+            batched_args.append(arg)
+
+        results = []
+        for i in range(batch_size):
+            input_slice = [
+                (arg[i] if isinstance(arg, torch.Tensor) and in_dim is not None else arg)
+                for arg, in_dim in zip(batched_args, in_dims_)
+            ]
+            result = func(*input_slice)
+            results.append(result)
+
+        if isinstance(results[0], torch.Tensor):
+            stacked = torch.stack(results)
+            if out_dims != 0:
+                return stacked.movedim(0, out_dims)
+            return stacked
+        return results
+
+    return wrapped