Implement a context to rewrite method or functions

Author: xadupre

_doc/api/torch_export_patches/index.rst

@@ -15,6 +15,8 @@ onnx_diagnostic.torch_export_patches
     :members:
     :no-undoc-members:
 
+.. autofunction:: onnx_diagnostic.torch_export_patches.torch_export_rewrite
+
 .. autofunction:: onnx_diagnostic.torch_export_patches.torch_export_patches
 
 .. autofunction:: onnx_diagnostic.torch_export_patches.register_additional_serialization_functions

_doc/api/torch_export_patches/patch_module.rst

@@ -5,3 +5,4 @@ onnx_diagnostic.torch_export_patches.patch_module
 .. automodule:: onnx_diagnostic.torch_export_patches.patch_module
     :members:
     :no-undoc-members:
+    :exclude: torch_export_rewrite

_unittests/ut_torch_export_patches/test_patch_module.py

@@ -4,12 +4,28 @@
 import unittest
 import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
+from onnx_diagnostic.torch_export_patches import torch_export_patches, torch_export_rewrite
 from onnx_diagnostic.torch_export_patches.patch_module import (
     transform_method,
     inplace_add_parent,
 )
 
 
+class _ModelForATest(torch.nn.Module):
+    def forward(self, x, y):
+        if x.sum() > 0:
+            return x + y
+        else:
+            return torch.abs(x) + y + 1
+
+
+def _single_forward(x, y):
+    if x.sum() > 0:
+        return x + y
+    else:
+        return torch.abs(x) + y + 1
+
+
 class TestPatchModule(ExtTestCase):
     def test_parent(self):
         class Model(torch.nn.Module):
@@ -361,8 +377,71 @@ def test_rewrite_PLBartEncoderLayer(self):
             ),
             rewritten.code,
         )
-        print()
-        print(rewritten.code)
+
+    @hide_stdout()
+    def test_torch_export_patch_method_tuple(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                if x.sum() > 0:
+                    return x + y
+                else:
+                    return torch.abs(x) + y + 1
+
+        model = Model()
+        x, y = torch.rand((4, 5)), torch.rand((4, 5))
+        expected = model(x, y)
+        DYN = torch.export.Dim.DYNAMIC
+        ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
+        with torch_export_patches(rewrite_methods=[(Model, "forward")], verbose=2):
+            ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
+            got = ep.module()(x, y)
+            self.assertEqualArray(expected, got)
+
+    @hide_stdout()
+    def test_torch_export_rewrite_method_tuple(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                if x.sum() > 0:
+                    return x + y
+                else:
+                    return torch.abs(x) + y + 1
+
+        model = Model()
+        x, y = torch.rand((4, 5)), torch.rand((4, 5))
+        expected = model(x, y)
+        DYN = torch.export.Dim.DYNAMIC
+        ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
+        with torch_export_rewrite(rewrite_methods=[(Model, "forward")], verbose=1):
+            ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
+            got = ep.module()(x, y)
+            self.assertEqualArray(expected, got)
+
+    def test_torch_export_rewrite_method_only(self):
+        model = _ModelForATest()
+        x, y = torch.rand((4, 5)), torch.rand((4, 5))
+        expected = model(x, y)
+        DYN = torch.export.Dim.DYNAMIC
+        ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
+        with torch_export_rewrite(rewrite_methods=[_ModelForATest.forward], verbose=0):
+            ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
+            got = ep.module()(x, y)
+            self.assertEqualArray(expected, got)
+
+    @hide_stdout()
+    def test_torch_export_rewrite_function(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return _single_forward(x, y)
+
+        model = Model()
+        x, y = torch.rand((4, 5)), torch.rand((4, 5))
+        expected = model(x, y)
+        DYN = torch.export.Dim.DYNAMIC
+        ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
+        with torch_export_rewrite(rewrite_methods=[_single_forward], verbose=1):
+            ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
+            got = ep.module()(x, y)
+            self.assertEqualArray(expected, got)
 
 
 if __name__ == "__main__":

onnx_diagnostic/torch_export_patches/__init__.py

@@ -2,6 +2,7 @@
     torch_export_patches,
     register_additional_serialization_functions,
 )
+from .patch_module import torch_export_rewrite
 
 
 # bypass_export_some_errors is the first name given to the patches.

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -102,6 +102,7 @@ def torch_export_patches(
     verbose: int = 0,
     patch: bool = True,
     custom_patches: Optional[List[type["torch.nn.Module"]]] = None,  # noqa: F821
+    rewrite_methods: Optional[List[Callable]] = None,
 ) -> Callable:
     """
     Tries to bypass some situations :func:`torch.export.export` does not support.
@@ -123,6 +124,11 @@ def torch_export_patches(
     :param custom_patches: to apply custom patches,
         every patched class must define static attributes
         ``_PATCHES_``, ``_PATCHED_CLASS_``
+    :param rewrite_methods: list of methods to automatically rewrite
+        before exporting, methods with control flow need to be rewritten
+        before being exported if the execution path depends on the inputs,
+        this is done by function :func:`transform_method
+        <onnx_diagnostic.torch_export_patches.patch_module.transform_method>`
     :param verbose: to show which patches is applied
 
     The list of available patches.
@@ -143,21 +149,21 @@ def torch_export_patches(
 
     Examples:
 
-    ::
+    .. code-block:: python
 
         with torch_export_patches(patch_transformers=True) as modificator:
             inputs = modificator(inputs)
             onx = to_onnx(..., inputs, ...)
 
-    ::
+    .. code-block:: python
 
         with torch_export_patches(patch_transformers=True) as modificator:
             inputs = modificator(inputs)
             onx = torch.onnx.export(..., inputs, ...)
 
     It can be used as well to fix the torch export:
 
-    ::
+    .. code-block:: python
 
         with torch_export_patches(patch_transformers=True) as modificator:
             inputs = modificator(inputs)
@@ -166,7 +172,7 @@ def torch_export_patches(
     When running the model through the exported program, only the
     serialization functions need to be restored:
 
-    ::
+    .. code-block:: python
 
         with register_additional_serialization_functions() as modificator:
             inputs = modificator(inputs)
@@ -176,7 +182,26 @@ def torch_export_patches(
     may appear ``AssertionError: Mutating module attribute _seen_tokens during export.``.
     It can be avoided by setting ``strict=False`` when call :func:`torch.export.export`.
     """
-    if not patch:
+    if rewrite_methods:
+        from .patch_module import torch_export_rewrite
+
+        with torch_export_rewrite(
+            rewrite_methods=rewrite_methods, verbose=verbose
+        ), torch_export_patches(
+            patch_sympy=patch_sympy,
+            patch_torch=patch_torch,
+            patch_transformers=patch_transformers,
+            catch_constraints=catch_constraints,
+            stop_if_static=stop_if_static,
+            verbose=verbose,
+            patch=patch,
+            custom_patches=custom_patches,
+        ):
+            try:
+                yield
+            finally:
+                pass
+    elif not patch:
         fct_callable = lambda x: x  # noqa: E731
         done = _register_cache_serialization(verbose=verbose)
         try:

onnx_diagnostic/torch_export_patches/patch_inputs.py

@@ -3,7 +3,6 @@
 import torch
 import transformers
 from ..helpers import string_type
-from ..helpers.cache_helper import make_dynamic_cache
 
 
 def _process_cache(k: str, v):
@@ -16,6 +15,8 @@ def _process_cache(k: str, v):
         and set(len(t) for t in v) == {2}
     ):
         # A dynamicCache
+        from ..helpers.cache_helper import make_dynamic_cache
+
         cache = make_dynamic_cache(v)
         return cache
     if isinstance(v, torch.Tensor):

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -1,9 +1,11 @@
 import ast
 import copy
+import contextlib
 import inspect
 import types
 import textwrap
-from typing import Callable, Dict, List, Set, Optional
+import sys
+from typing import Callable, Dict, List, Set, Optional, Tuple, Union
 
 NODE_TYPES = tuple(
     getattr(ast, k)
@@ -515,3 +517,78 @@ def forward(self, x, y):
     if not isinstance(new_func, types.FunctionType):
         raise RuntimeError("Transformed function not found")
     return RewrittenMethod(new_tree, new_func)
+
+
+@contextlib.contextmanager
+def torch_export_rewrite(
+    rewrite_methods: Optional[List[Union[Tuple[type, str], Callable]]] = None, verbose: int = 0
+):
+    """
+    Automatically rewrite the methods given in `rewrite_methods` to export
+    control flows (test and loops).
+
+    :param rewrite_methods: methods to rewrite, if not empty, the function may try
+        to discover them, a method is defined by its class (a type) and its name
+        if the class is local, by itself otherwise
+    :param verbose: verbosity, up to 10, 10 shows the rewritten code
+    """
+    assert (
+        rewrite_methods
+    ), "rewrite_methods is empty, automated discovery is not implemented yet"
+    keep = {}
+    for me in rewrite_methods:
+        if isinstance(me, tuple):
+            assert len(me) == 2, f"Unexpected value for a rewritten method or function {me}"
+            cls, name = me
+            to_rewrite = getattr(cls, name)
+            kind = "method"
+        else:
+            name = me.__qualname__
+            spl = name.split(".")
+            if len(spl) == 1:
+                # This a function
+                module = me.__module__
+                if module in me.__globals__:
+                    mod = me.__globals__[module]
+                else:
+                    assert module in sys.modules, (
+                        f"Cannot find module name {module!r} in sys.modules or "
+                        f"__globals__={sorted(me.__globals__)}"
+                    )
+                    mod = sys.modules[module]
+                cls = mod
+                name = name
+                to_rewrite = me
+                kind = "function"
+            else:
+                kind = "method"
+                # This is a method
+                assert len(spl) >= 2, (
+                    f"{me} is not method, its name {name!r} does not contain a class name, "
+                    f"dir(me)={dir(me)}"
+                )
+                cls_name = spl[-2]
+                assert cls_name in me.__globals__, (
+                    f"Class name {cls_name!r} from method {name!r} "
+                    f"could not be found in set(me.__globals__)={sorted(me.__globals__)}"
+                )
+                cls = me.__globals__[cls_name]
+                name = me.__name__
+                to_rewrite = me
+                assert hasattr(
+                    cls, name
+                ), f"Method {name!r} inferred form {me} was not found in class {cls}."
+        assert (cls, name) not in keep, f"{kind} {me} cannot be rewritten twice."
+        if verbose:
+            print(f"[torch_export_rewrite] rewrites {kind} {cls.__name__}.{name}")
+        keep[cls, name] = to_rewrite
+        rewr = transform_method(to_rewrite, verbose=max(verbose - 1, 0))
+        setattr(cls, name, rewr.func)
+
+    try:
+        yield
+    finally:
+        for (cls, name), me in keep.items():
+            if verbose:
+                print(f"[torch_export_rewrite] restored {kind} {cls.__name__}.{name}")
+            setattr(cls, name, me)