Improves rewriting

Author: xadupre

_doc/api/torch_export_patches/index.rst

@@ -9,7 +9,7 @@ onnx_diagnostic.torch_export_patches
     patch_expressions
     patch_inputs
     patch_module
-
+    patch_module_helper
 
 .. automodule:: onnx_diagnostic.torch_export_patches
     :members:

_doc/api/torch_export_patches/patch_module_helper.rst

@@ -0,0 +1,8 @@
+
+onnx_diagnostic.torch_export_patches.patch_module
+=================================================
+
+.. automodule:: onnx_diagnostic.torch_export_patches.patch_module
+    :members:
+    :no-undoc-members:
+    :exclude-members: torch_export_rewrite

_unittests/ut_torch_export_patches/test_patch_module.py

@@ -13,6 +13,7 @@
     ShapeFinder,
     RewriteControlFlow,
 )
+from onnx_diagnostic.torch_export_patches.patch_module_helper import ast_or_into_bitor
 
 
 class _ModelForATest(torch.nn.Module):
@@ -396,15 +397,25 @@ def forward(self, x, y):
     def test_rewrite_test_in_PLBartEncoderLayer(self):
         from transformers.models.plbart.modeling_plbart import PLBartEncoderLayer
 
-        rewritten = transform_method(PLBartEncoderLayer.forward, verbose=self.verbose)
+        def filter_node(node) -> bool:
+            return isinstance(node, ast.If) and not isinstance(node.test, ast.Name)
+
+        rewritten = transform_method(
+            PLBartEncoderLayer.forward,
+            verbose=self.verbose,
+            filter_node=filter_node,
+            pre_rewriter=ast_or_into_bitor,
+        )
         self.assertIn(
             (
                 "torch.cond(hidden_states.dtype == torch.float16 and "
-                "(torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()), "
+                "torch.isinf(hidden_states).any()"
+                " | torch.isnan(hidden_states).any(), "
                 "branch_cond_then_1, branch_cond_else_1, [hidden_states])"
             ),
             rewritten.code,
         )
+        self.assertNotIn("torch.cond(output_attentions", rewritten.code)
 
     @hide_stdout()
     def test_torch_export_patch_method_tuple(self):

_unittests/ut_torch_models/test_hghub_mode_rewrite.py

@@ -1,12 +1,33 @@
 import unittest
-from onnx_diagnostic.ext_test_case import ExtTestCase
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, ignore_errors
 from onnx_diagnostic.torch_models.hghub.hub_data import code_needing_rewriting
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import torch_export_patches
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 
 
 class TestHuggingFaceHubModelRewrite(ExtTestCase):
 
     def test_code_needing_rewriting(self):
-        self.assertEqual(1, len(code_needing_rewriting("BartForConditionalGeneration")))
+        self.assertEqual(2, len(code_needing_rewriting("BartForConditionalGeneration")))
+
+    @hide_stdout()
+    @ignore_errors(OSError)
+    def test_export_rewritin_bart(self):
+        mid = "hf-tiny-model-private/tiny-random-PLBartForConditionalGeneration"
+        data = get_untrained_model_with_inputs(mid, verbose=1)
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        dump_folder = self.get_dump_file("test_export_rewritin_bart")
+        print(self.string_type(inputs))
+        print(self.string_type(ds))
+        with torch_export_patches(
+            patch_transformers=True,
+            rewrite=code_needing_rewriting("BartForConditionalGeneration"),
+            dump_rewriting=dump_folder,
+        ):
+            model(**inputs)
+            torch.export.export(model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds))
 
 
 if __name__ == "__main__":

onnx_diagnostic/_command_lines_parser.py

@@ -343,6 +343,12 @@ def get_parser_validate() -> ArgumentParser:
         help="drops the following inputs names, it should be a list "
         "with comma separated values",
     )
+    parser.add_argument(
+        "--opset",
+        type=int,
+        default=18,
+        help="onnx opset to use, 18 by default",
+    )
     parser.add_argument(
         "--subfolder",
         help="subfolder where to find the model and the configuration",
@@ -426,6 +432,7 @@ def _cmd_validate(argv: List[Any]):
             input_options=args.iop,
             model_options=args.mop,
             subfolder=args.subfolder,
+            opset=args.opset,
         )
         print("")
         print("-- summary --")

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -89,13 +89,20 @@ class RewriteControlFlow(ast.NodeTransformer):
     :param skip_objects: to skip variable names if included in that list
         such as modules
     :param args_names: defines the local variables
+    :param filter_nodes: a function which is used to decide which node
+        to rewrite, True by default
+    :param pre_rewriter: a rewriter applied before the automated rewriting
+    :param post_rewriter: a rewriter applied after the automated rewriting
     """
 
     def __init__(
         self,
         prefix: str = "branch_cond",
         skip_objects: Optional[Dict[str, object]] = None,
         args_names: Optional[Set[str]] = None,
+        filter_node: Optional[Callable["ast.Node", bool]] = None,
+        pre_rewriter: Optional[Callable["ast.Node", "ast.Node"]] = None,
+        post_rewriter: Optional[Callable["ast.Node", "ast.Node"]] = None,
     ):
         self.counter_test = 0
         self.counter_loop = 0
@@ -104,6 +111,9 @@ def __init__(
         self.skip_objects = skip_objects or {}
         self.args_names = args_names or set()
         self.local_variables = self.args_names.copy()
+        self.filter_node = filter_node or (lambda _node: True)
+        self.pre_rewriter = pre_rewriter or (lambda node: node)
+        self.post_rewriter = post_rewriter or (lambda node: node)
 
     def generic_visit(self, node):
         return super().generic_visit(node)
@@ -320,6 +330,11 @@ def _make_targets(self, node, then_assigns, else_assigns):
         return tgt, tgt_mapping
 
     def visit_If(self, node):
+        if not self.filter_node(node):
+            return [node]
+
+        node = self.pre_rewriter(node)
+
         # First recurse into subnodes
         known_local_variables = self.local_variables.copy()
         node = self.generic_visit(node)
@@ -380,7 +395,7 @@ def visit_If(self, node):
             ast.copy_location(assign, node)
             ast.fix_missing_locations(assign)
             self.local_variables = known_local_variables | added
-            return [then_def, else_def, assign]
+            return [self.post_rewriter(n) for n in [then_def, else_def, assign]]
 
         # Case 2: return in both branches, we assume both branches return the same results.
         then_ret = node.body[-1]
@@ -403,7 +418,7 @@ def visit_If(self, node):
         ret = ast.Return(call)
         ast.copy_location(ret, node)
         ast.fix_missing_locations(ret)
-        return [then_def, else_def, ret]
+        return [self.post_rewriter(n) for n in [then_def, else_def, ret]]
 
     def _find_loop_vars(self, node):
         assert isinstance(node, ast.For), f"Unexpected type {type(node)} for node"
@@ -462,6 +477,11 @@ def _find_loop_vars(self, node):
         )
 
     def visit_For(self, node):
+        if not self.filter_node(node):
+            return [node]
+
+        node = self.pre_rewriter(node)
+
         # For nested loops.
         self.generic_visit(node)
         # look for variables, loop, inputs and outputs of the body
@@ -622,7 +642,7 @@ def visit_For(self, node):
             ctx=ast.Store(),
         )
         assign = ast.Assign(targets=[target], value=call)
-        return [func_def, assign]
+        return [self.post_rewriter(func_def), self.post_rewriter(assign)]
 
 
 class RewrittenMethod:
@@ -697,6 +717,9 @@ def transform_method(
     func: Callable,
     prefix: str = "branch_cond",
     verbose: int = 0,
+    filter_node: Optional[Callable["ast.Node", bool]] = None,
+    pre_rewriter: Optional[Callable["ast.Node", "ast.Node"]] = None,
+    post_rewriter: Optional[Callable["ast.Node", "ast.Node"]] = None,
 ) -> RewrittenMethod:
     """
     Returns a new function based on `func` where every test (if)
@@ -717,6 +740,9 @@ def transform_method(
     :param func: method or function to rewrite
     :param prefix: prefix used to create the functions for the branches
     :param verbose: verbosity
+    :param filter_node: a function which tells which node to rewrite
+    :param pre_rewriter: a rewriter applied before the automated rewriting
+    :param post_rewriter: a rewriter applied after the automated rewriting
     :return: rewritten method
 
     An example with **return**:
@@ -801,6 +827,9 @@ def forward(self, x, y):
         prefix=prefix,
         skip_objects=modules,
         args_names=set(sig.parameters),
+        filter_node=filter_node,
+        pre_rewriter=pre_rewriter,
+        post_rewriter=post_rewriter,
     )
     normalize_assignment_in_test(tree)
     inplace_add_parent(tree)
@@ -912,7 +941,22 @@ def forward(self, x, y):
             cls, name = me
             to_rewrite = getattr(cls, name)
             kind = "method"
+            kws = {}
         else:
+            if isinstance(me, dict):
+                assert "function" in me and (
+                    "filter_node" in me or "pre_rewriter" in me or "post_rewriter" in me
+                ), (
+                    f"If the rewriting code is defined as a dictionary, key "
+                    f"'function' must be defined, other arguments must be understood by "
+                    f"{transform_method.__name__}, "
+                    f"the given value is {me!r}."
+                )
+                kws = me
+                me = me["function"]
+                del kws["function"]
+            else:
+                kws = {}
             name = me.__qualname__
             spl = name.split(".")
             if len(spl) == 1:
@@ -958,8 +1002,9 @@ def forward(self, x, y):
             if verbose:
                 print(f"[torch_export_rewrite] dump original code in {filename!r}")
             with open(filename, "w") as f:
-                f.write(inspect.getsource(to_rewrite))
-        rewr = transform_method(to_rewrite, verbose=max(verbose - 1, 0))
+                code = inspect.getsource(to_rewrite)
+                f.write(code)
+        rewr = transform_method(to_rewrite, verbose=max(verbose - 1, 0), **kws)
         if dump_rewriting:
             filename = f"{dump_rewriting}.{kind}.{cls_name}.{name}.rewritten.py"
             if verbose:

onnx_diagnostic/torch_export_patches/patch_module_helper.py

@@ -0,0 +1,18 @@
+import ast
+
+
+class OrToBitOrTransformer(ast.NodeTransformer):
+    def visit_BoolOp(self, node):
+        self.generic_visit(node)
+        if isinstance(node.op, ast.Or):
+            new_node = node.values[0]
+            for value in node.values[1:]:
+                new_node = ast.BinOp(left=new_node, op=ast.BitOr(), right=value)
+            return ast.copy_location(new_node, node)
+        return node
+
+
+def ast_or_into_bitor(node: "ast.Node") -> "ast.Node":
+    """Replaces every operator ``or`` into ``|``."""
+    new_node = OrToBitOrTransformer().visit(node)
+    return new_node

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -1,3 +1,4 @@
+import ast
 import io
 import functools
 import textwrap
@@ -214,10 +215,31 @@ def code_needing_rewriting(cls_name: str) -> Optional[List[Any]]:
 
         print(code_needing_rewriting("BartForConditionalGeneration"))
     """
-    if cls_name in {"BartForConditionalGeneration", "BartEncoderLayer"}:
+    if cls_name in {
+        "BartEncoderLayer",
+        "BartForConditionalGeneration",
+        "PLBartEncoderLayer",
+        "PLBartForConditionalGeneration",
+    }:
         import transformers
+        from ...torch_export_patches.patch_module_helper import ast_or_into_bitor
 
-        return [transformers.models.bart.modeling_bart.BartEncoderLayer.forward]
+        bd = dict(
+            filter_node=(
+                lambda node: isinstance(node, ast.If) and not isinstance(node.test, ast.Name)
+            ),
+            pre_rewriter=ast_or_into_bitor,
+        )
+
+        def _add(f):
+            g = bd.copy()
+            g["function"] = f
+            return g
+
+        return [
+            _add(transformers.models.bart.modeling_bart.BartEncoderLayer.forward),
+            _add(transformers.models.plbart.modeling_plbart.PLBartEncoderLayer.forward),
+        ]
     return None
 
 

onnx_diagnostic/torch_models/test_helper.py

@@ -230,6 +230,7 @@ def validate_model(
     input_options: Optional[Dict[str, Any]] = None,
     model_options: Optional[Dict[str, Any]] = None,
     subfolder: Optional[str] = None,
+    opset: Optional[int] = None,
 ) -> Tuple[Dict[str, Union[int, float, str]], Dict[str, Any]]:
     """
     Validates a model.
@@ -265,6 +266,7 @@ def validate_model(
     :param model_options: additional options when creating the model such as
         ``num_hidden_layers`` or ``attn_implementation``
     :param subfolder: version or subfolders to uses when retrieving a model id
+    :param opset: onnx opset to use for the conversion
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
 
@@ -295,6 +297,8 @@ def validate_model(
             version_exporter=exporter or "",
         )
     )
+    if opset:
+        summary["version_opset"] = opset
 
     folder_name = None
     if dump_folder:
@@ -343,6 +347,8 @@ def validate_model(
     )
     data["input_options"] = iop
     data["model_options"] = mop
+    if opset:
+        data["model_opset"] = opset
     if "rewrite" in data:
         if rewrite:
             summary["model_rewrite"] = str(data["rewrite"])
@@ -992,6 +998,9 @@ def call_torch_export_onnx(
     summary["export_dynamo"] = dynamo
     summary["export_args"] = string_type(args, with_shape=True)
     summary["export_kwargs"] = string_type(kwargs, with_shape=True)
+    opset = data.get("model_opset", None)
+    if opset:
+        summary["export_opset"] = opset
 
     if dynamo:
         export_export_kwargs = dict(dynamo=True, dynamic_shapes=ds)
@@ -1012,6 +1021,8 @@ def call_torch_export_onnx(
             print("[call_torch_export_onnx] dynamo=False so...")
             print(f"[call_torch_export_onnx] args={string_type(args, with_shape=True)}")
             print(f"[call_torch_export_onnx] kwargs={string_type(kwargs, with_shape=True)}")
+    if opset:
+        export_export_kwargs["opset_version"] = opset
     if verbose:
         print(
             f"[call_torch_export_onnx] export_export_kwargs="
@@ -1123,6 +1134,9 @@ def call_torch_export_custom(
     strict = "-strict" in exporter
     args, kwargs = split_args_kwargs(data["inputs_export"])
     ds = data.get("dynamic_shapes", None)
+    opset = data.get("model_opset", None)
+    if opset:
+        summary["export_opset"] = opset
     if verbose:
         print(
             f"[call_torch_export_custom] exporter={exporter!r}, "
@@ -1163,6 +1177,9 @@ def call_torch_export_custom(
         return_optimize_report=True,
         verbose=max(verbose - 2, 0),
     )
+    if opset:
+        kws["target_opset"] = opset
+    assert opset
 
     epo, opt_stats = _quiet_or_not_quiet(
         quiet,