fix issues

Author: xadupre

CHANGELOGS.rst

@@ -4,11 +4,12 @@ Change Logs
 0.5.0
 +++++
 
+* :pr:`100`: implements a context to automatically rewrite methods or function with control flows
 * :pr:`96`: implements ``is_stealing``, ``steal_append`` to complement ``steal_forward``
-* :pr:`95`: fix Scan implementation for ``OnnxruntimeEvaluator``
-* :pr:`93`: introduce patched expression to get around annoying export issues
-* :pr:`92`: support errors distribution in max_diff
-* :pr:`91`: enable strings in ``guess_dynamic_shapes``
+* :pr:`95`: fixzq Scan implementation for ``OnnxruntimeEvaluator``
+* :pr:`93`: introduces patched expressions to get around annoying export issues
+* :pr:`92`: supports errors distribution in max_diff
+* :pr:`91`: enables strings in ``guess_dynamic_shapes``
 * :pr:`88`, :pr:`89`: extends ``steal_forward`` to dump input, outputs in onnx models
 * :pr:`83`, :pr:`85`: improves the automated rewriting of control flow (test)
 

_doc/api/torch_export_patches/patch_module.rst

@@ -5,4 +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
+    :exclude-members: torch_export_rewrite

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -128,7 +128,8 @@ def torch_export_patches(
         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>`
+        <onnx_diagnostic.torch_export_patches.patch_module.transform_method>`,
+        its documentation provides possible values
     :param verbose: to show which patches is applied
 
     The list of available patches.

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -530,7 +530,55 @@ def torch_export_rewrite(
     :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
+    :param verbose: verbosity, up to 10, 10 shows the rewritten code,
+        ``verbose=1`` shows the rewritten function,
+        ``verbose=2`` shows the rewritten code as well
+
+    Example:
+
+    .. code-block:: python
+
+        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))
+        DYN = torch.export.Dim.DYNAMIC
+        ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
+        with torch_export_rewrite(rewrite_methods=[(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]):
+            ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
+
+    Functions (if not local) can also be rewritten:
+
+    .. code-block:: python
+
+        def outside(x, y):
+            if x.sum() > 0:
+                return x + y
+            else:
+                return torch.abs(x) + y + 1
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return outside(x, y)
+
+        model = Model()
+        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]):
+            ep = torch.export.export(model, (x, y), dynamic_shapes=ds)
     """
     assert (
         rewrite_methods