rename

Author: xadupre

README.rst

@@ -86,6 +86,22 @@ Enlightening Examples
 Snapshot of usefuls tools
 +++++++++++++++++++++++++
 
+**torch_export_patches**
+
+.. code-block:: python
+
+  with torch_export_patches(patch_transformers=True) as f:
+      ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
+      # ...
+
+**torch_export_rewrite**
+
+.. code-block:: python
+
+  with torch_export_rewrite(rewrite=[Model.forward]) as f:
+      ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
+      # ...
+
 **string_type**
 
 .. code-block:: python

_doc/index.rst

@@ -89,6 +89,28 @@ Enlightening Examples
 Some Usefuls Tools
 ==================
 
+torch_export_patches
+++++++++++++++++++++
+
+See :func:`onnx_diagnostic.torch_export_patches.torch_export_patches`.
+
+.. code-block:: python
+
+  with torch_export_patches(patch_transformers=True) as f:
+      ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
+      # ...
+
+torch_export_rewrite
+++++++++++++++++++++
+
+See :func:`onnx_diagnostic.torch_export_patches.torch_export_rewrite`.
+
+.. code-block:: python
+
+  with torch_export_rewrite(rewrite=[Model.forward]) as f:
+      ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
+      # ...
+
 string_type
 +++++++++++
 

_unittests/ut_torch_export_patches/test_patch_module.py

@@ -392,7 +392,7 @@ def forward(self, x, y):
         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):
+        with torch_export_patches(rewrite=[(Model, "forward")], verbose=2):
             ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
             got = ep.module()(x, y)
             self.assertEqualArray(expected, got)
@@ -411,7 +411,7 @@ def forward(self, x, y):
         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):
+        with torch_export_rewrite(rewrite=[(Model, "forward")], verbose=1):
             ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
             got = ep.module()(x, y)
             self.assertEqualArray(expected, got)
@@ -422,7 +422,7 @@ def test_torch_export_rewrite_method_only(self):
         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):
+        with torch_export_rewrite(rewrite=[_ModelForATest.forward], verbose=0):
             ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
             got = ep.module()(x, y)
             self.assertEqualArray(expected, got)
@@ -438,7 +438,7 @@ def forward(self, x, y):
         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):
+        with torch_export_rewrite(rewrite=[_single_forward], verbose=1):
             ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
             got = ep.module()(x, y)
             self.assertEqualArray(expected, got)

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -102,7 +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,
+    rewrite: Optional[List[Callable]] = None,
 ) -> Callable:
     """
     Tries to bypass some situations :func:`torch.export.export` does not support.
@@ -124,7 +124,7 @@ 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
+    :param rewrite: 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
@@ -183,12 +183,10 @@ 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 rewrite_methods:
+    if rewrite:
         from .patch_module import torch_export_rewrite
 
-        with torch_export_rewrite(
-            rewrite_methods=rewrite_methods, verbose=verbose
-        ), torch_export_patches(
+        with torch_export_rewrite(rewrite=rewrite, verbose=verbose), torch_export_patches(
             patch_sympy=patch_sympy,
             patch_torch=patch_torch,
             patch_transformers=patch_transformers,

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -521,13 +521,13 @@ def forward(self, x, y):
 
 @contextlib.contextmanager
 def torch_export_rewrite(
-    rewrite_methods: Optional[List[Union[Tuple[type, str], Callable]]] = None, verbose: int = 0
+    rewrite: Optional[List[Union[Tuple[type, str], Callable]]] = None, verbose: int = 0
 ):
     """
-    Automatically rewrite the methods given in `rewrite_methods` to export
+    Automatically rewrite the methods given in `rewrite` to export
     control flows (test and loops).
 
-    :param rewrite_methods: methods to rewrite, if not empty, the function may try
+    :param rewrite: methods of functions 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,
@@ -549,14 +549,14 @@ def forward(self, x, y):
         x, y = torch.rand((4, 5)), torch.rand((4, 5))
         DYN = torch.export.Dim.DYNAMIC
         ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
-        with torch_export_rewrite(rewrite_methods=[(Model, "forward")]):
+        with torch_export_rewrite(rewrite=[(Model, "forward")]):
             ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
 
     If the method to rewrite is not local, then the following can be used:
 
     .. code-block:: python
 
-        with torch_export_rewrite(rewrite_methods=[Model.forward]):
+        with torch_export_rewrite(rewrite=[Model.forward]):
             ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
 
     Functions (if not local) can also be rewritten:
@@ -577,14 +577,12 @@ def forward(self, x, y):
         x, y = torch.rand((4, 5)), torch.rand((4, 5))
         DYN = torch.export.Dim.DYNAMIC
         ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
-        with torch_export_rewrite(rewrite_methods=[outside]):
+        with torch_export_rewrite(rewrite=[outside]):
             ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
     """
-    assert (
-        rewrite_methods
-    ), "rewrite_methods is empty, automated discovery is not implemented yet"
+    assert rewrite, "rewrite is empty, automated discovery is not implemented yet"
     keep = {}
-    for me in rewrite_methods:
+    for me in rewrite:
         if isinstance(me, tuple):
             assert len(me) == 2, f"Unexpected value for a rewritten method or function {me}"
             cls, name = me