Fix patched_vmap (#275)

* Fix patched_vmap

* documentation

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,8 @@ Change Logs
 0.7.17
 ++++++
 
+* :pr:`275`: fixes function ``patched_vmap``
+
 0.7.16
 ++++++
 

_unittests/ut_torch_export_patches/test_patch_torch.py

@@ -28,6 +28,23 @@ def test_vmap(self):
         got = patched_vmap(f)(x, y)
         self.assertEqualArray(expected, got)
 
+    @requires_transformers("4.52")
+    def test_export_patched_vmap_dynamic_shapes(self):
+        from onnx_diagnostic.torch_export_patches.patches.patch_torch import patched_vmap
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                f = lambda x, y: x * y + 1  # noqa: E731
+                return patched_vmap(f)(x, y)
+
+        x = torch.tensor([1.0, 2.0, 3.0])
+        y = torch.tensor([0.1, 0.2, 0.3])
+        expected = Model()(x, y)
+        DYN = torch.export.Dim.DYNAMIC
+        ep = torch.export.export(Model(), (x, y), dynamic_shapes=({0: DYN}, {0: DYN}))
+        got = ep.module()(x, y)
+        self.assertEqualArray(expected, got)
+
     @requires_torch("2.10")
     def test_export_vmap(self):
         class Model(torch.nn.Module):
@@ -56,6 +73,55 @@ def forward(self, x, y):
         ep = torch.export.export(Model(), (x, y))
         self.assertEqualArray(Model()(x, y), ep.module()(x, y))
 
+    @requires_torch("2.8")
+    @requires_transformers("4.52")
+    def test_export_patched_vmap_scan(self):
+        from onnx_diagnostic.torch_export_patches.patches.patch_torch import patched_vmap
+
+        x = torch.tensor([1.0, 2.0, 3.0])
+        y = torch.tensor([0.1, 0.2, 0.3])
+        res = torch.ops.higher_order.scan(lambda x, y: x + y, [], [x, y], [])
+        self.assertEqualArray(x + y, res[0])
+
+        class ModelVmap(torch.nn.Module):
+            def forward(self, x, y):
+                f = lambda x, y: x * y + 1  # noqa: E731
+                return torch.vmap(f)(x, y)
+
+        expected = ModelVmap()(x, y)
+
+        class ModelNoScan(torch.nn.Module):
+            def forward(self, x, y):
+                f = lambda x, y: x * y + 1  # noqa: E731
+                return patched_vmap(f, use_scan=False)(x, y)
+
+        expected2 = ModelNoScan()(x, y)
+        self.assertEqualArray(expected, expected2)
+
+        class ModelScan(torch.nn.Module):
+            def forward(self, x, y):
+                f = lambda x, y: [x * y + 1]  # noqa: E731
+                return torch.ops.higher_order.scan(f, [], [x, y], [])[0]
+
+        expected2 = ModelNoScan()(x, y)
+        self.assertEqualArray(expected, expected2)
+        ep = torch.export.export(ModelScan(), (x, y))
+        self.assertEqualArray(expected, ep.module()(x, y))
+
+        DYN = torch.export.Dim.DYNAMIC
+        ep = torch.export.export(ModelScan(), (x, y), dynamic_shapes=({0: DYN}, {0: DYN}))
+        self.assertEqualArray(expected, ep.module()(x, y))
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                f = lambda x, y: x * y + 1  # noqa: E731
+                return patched_vmap(f, use_scan=True)(x, y)
+
+        expected2 = Model()(x, y)
+        self.assertEqualArray(expected, expected2)
+        ep = torch.export.export(Model(), (x, y), dynamic_shapes=({0: DYN}, {0: DYN}))
+        self.assertEqualArray(expected, ep.module()(x, y))
+
     @requires_transformers("4.52")
     def test_vmap_outdim(self):
         from onnx_diagnostic.torch_export_patches.patches.patch_torch import patched_vmap

onnx_diagnostic/export/shape_helper.py

@@ -212,6 +212,23 @@ def make_fake_with_dynamic_dimensions(
     constraints as the following dynamic shapes.
     This uses function :func:`onnx_diagnostic.helpers.fake_tensor_helper.make_fake`.
 
+    A simple tensor:
+
+    .. runpython::
+        :showcode:
+
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import make_fake_with_dynamic_dimensions
+
+        inputs, _ = make_fake_with_dynamic_dimensions(
+            torch.rand((2, 3, 4, 5), dtype=dtype=torch.float32),
+            {0: "batch", 2: "cache_length"},
+        )
+        print(inputs)
+
+    With a cache:
+
     .. runpython::
         :showcode:
 

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -683,7 +683,7 @@ def compute_concrete_val() -> sympy.Basic:
         return concrete_val
 
 
-def patched_vmap(func, in_dims=0, out_dims=0):
+def patched_vmap(func, in_dims=0, out_dims=0, use_scan: bool = False):
     """
     Python implementation of :func:`torch.vmap`.
     The implementation raises an issue when it is being exported with
@@ -724,8 +724,9 @@ def wrapped(*args):
             arg = arg.movedim(in_dim, 0)
             batched_args.append(arg)
 
-        if all(isinstance(a, torch.Tensor) for a in args) and isinstance(
-            batch_size, torch.SymInt
+        if use_scan or (
+            all(isinstance(a, torch.Tensor) for a in args)
+            and isinstance(batch_size, torch.SymInt)
         ):
             batched_tensors = [
                 (
@@ -735,7 +736,9 @@ def wrapped(*args):
                 )
                 for arg, in_dim in zip(batched_args, in_dims_)
             ]
-            results = torch.ops.higher_order.scan(func, [], batched_tensors, [])
+            results = torch.ops.higher_order.scan(
+                lambda *args, **kwargs: [func(*args, **kwargs)], [], batched_tensors, []
+            )
             stacked = results[0]
             if out_dims != 0:
                 return stacked.movedim(0, out_dims)
@@ -745,7 +748,7 @@ def wrapped(*args):
             torch._check(
                 not isinstance(batch_size, torch.SymInt),
                 lambda: (
-                    f"patched_vmap supports dynamic batch_size only if all argument "
+                    f"patched_vmap supports dynamic batch_size only if all arguments "
                     f"are tensors but types are {[type(a) for a in args]}"
                 ),
             )