fix a few things

Author: xadupre

_unittests/ut_torch_models/test_hghub_mode_rewrite.py

@@ -0,0 +1,13 @@
+import unittest
+from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.torch_models.hghub.hub_data import code_needing_rewriting
+
+
+class TestHuggingFaceHubModelRewrite(ExtTestCase):
+
+    def test_code_needing_rewriting(self):
+        self.assertEqual(1, len(code_needing_rewriting("BartForConditionalGeneration")))
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

onnx_diagnostic/_command_lines_parser.py

@@ -309,12 +309,17 @@ def get_parser_validate() -> ArgumentParser:
         help="catches exception, report them in the summary",
     )
     parser.add_argument(
-        "-p",
         "--patch",
         default=True,
         action=BooleanOptionalAction,
         help="applies patches before exporting",
     )
+    parser.add_argument(
+        "--rewrite",
+        default=True,
+        action=BooleanOptionalAction,
+        help="applies rewrite before exporting",
+    )
     parser.add_argument(
         "--stop-if-static",
         default=0,
@@ -411,6 +416,7 @@ def _cmd_validate(argv: List[Any]):
             dtype=args.dtype,
             device=args.device,
             patch=args.patch,
+            rewrite=args.rewrite,
             stop_if_static=args.stop_if_static,
             optimization=args.opt,
             exporter=args.export,

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -103,6 +103,7 @@ def torch_export_patches(
     patch: bool = True,
     custom_patches: Optional[List[type["torch.nn.Module"]]] = None,  # noqa: F821
     rewrite: Optional[List[Callable]] = None,
+    dump_rewriting: Optional[str] = None,
 ) -> Callable:
     """
     Tries to bypass some situations :func:`torch.export.export` does not support.
@@ -130,6 +131,7 @@ def torch_export_patches(
         this is done by function :func:`transform_method
         <onnx_diagnostic.torch_export_patches.patch_module.transform_method>`,
         its documentation provides possible values
+    :param dump_rewriting: dumps rewriting information in file beginning with that prefix
     :param verbose: to show which patches is applied
 
     The list of available patches.
@@ -186,7 +188,9 @@ def torch_export_patches(
     if rewrite:
         from .patch_module import torch_export_rewrite
 
-        with torch_export_rewrite(rewrite=rewrite, verbose=verbose), torch_export_patches(  # type: ignore[var-annotated]
+        with torch_export_rewrite(  # type: ignore[var-annotated]
+            rewrite=rewrite, dump_rewriting=dump_rewriting, 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

@@ -177,6 +177,18 @@ def _rewrite_if(
             else_ret = else_exprs[0]
             then_exprs = [n for n in node.body if not isinstance(n, ast.Return)]
             else_exprs = [n for n in node.orelse if not isinstance(n, ast.Return)]
+            assert type(then_ret.value) is type(else_ret.value), (
+                f"Inconsistencies return then value={then_ret.value}, "
+                f"else value={else_ret.value}"
+            )
+            if isinstance(then_ret.value, (ast.Tuple, ast.list)):
+                assert len(then_ret.value.elts) == len(else_ret.value.elts), (
+                    f"Unexpected number of elements on both branches, "
+                    f"then:{then_ret.value.elts}, else:{else_ret.value.elts}"
+                )
+                n_returned_values = len(then_ret.value.elts)
+            else:
+                n_returned_values = 0
         else:
             self._check(
                 tgt_mapping,
@@ -207,6 +219,7 @@ def _rewrite_if(
                 if len(else_rets) == 1
                 else ast.Tuple([self._clone(r) for r in else_rets], ctx=ast.Load())
             )
+            n_returned_values = len(then_rets) if len(then_rets) > 1 else 0
 
         # build local funcs
         then_def = ast.FunctionDef(
@@ -258,7 +271,7 @@ def _rewrite_if(
             ],
             keywords=[],
         )
-        return then_def, else_def, call, drop
+        return then_def, else_def, call, drop, n_returned_values
 
     def _filter_target(self, node, tgt_mapping):
         """
@@ -330,17 +343,32 @@ def visit_If(self, node):
             # the targets we need to export
             tgt, tgt_mapping = self._make_targets(node, then_assigns, else_assigns)
 
-            then_def, else_def, call, dropped = self._rewrite_if(
+            then_def, else_def, call, dropped, n_returned_values = self._rewrite_if(
                 node,
                 then_assigns,
                 else_assigns,
                 tgt_mapping=tgt_mapping,
                 known_local_variables=known_local_variables,
             )
             if dropped and isinstance(tgt, ast.Tuple):
-                tgt = ast.Tuple(
-                    tuple(t for t in tgt.elts if t.id not in dropped), ctx=ast.Store()
+                tgt_elts = tuple(t for t in tgt.elts if t.id not in dropped)
+            elif isinstance(tgt, ast.Tuple):
+                tgt_elts = tuple(t for t in tgt.elts if t.id not in dropped)
+            else:
+                tgt_elts = [tgt]
+
+            if n_returned_values == 0:
+                assert len(tgt_elts) == 1, (
+                    f"Inconsistencies between n_returned_values={n_returned_values}, "
+                    f"dropped={dropped}, tgt.elts={tgt.elts}, tgt_elts={tgt_elts}"
                 )
+                tgt = tgt_elts[0]
+            else:
+                assert n_returned_values == len(tgt_elts), (
+                    f"Inconsistencies between n_returned_values={n_returned_values}, "
+                    f"dropped={dropped}, tgt.elts={tgt.elts}, tgt_elts={tgt_elts}"
+                )
+                tgt = ast.Tuple(tgt_elts, ctx=ast.Store())
 
             added = {tgt.id} if isinstance(tgt, ast.Name) else set(t.id for t in tgt.elts)
             assign = ast.Assign(targets=[tgt], value=call)
@@ -364,7 +392,7 @@ def visit_If(self, node):
         )
         then_expr = then_ret.value
         else_expr = else_ret.value
-        then_def, else_def, call, dropped = self._rewrite_if(
+        then_def, else_def, call, dropped, n_returned_values = self._rewrite_if(
             node, [then_expr], [else_expr], known_local_variables=known_local_variables
         )
         ret = ast.Return(call)
@@ -809,7 +837,9 @@ def forward(self, x, y):
 
 @contextlib.contextmanager
 def torch_export_rewrite(
-    rewrite: Optional[List[Union[Tuple[type, str], Callable]]] = None, verbose: int = 0
+    rewrite: Optional[List[Union[Tuple[type, str], Callable]]] = None,
+    dump_rewriting: Optional[str] = None,
+    verbose: int = 0,
 ):
     """
     Automatically rewrite the methods given in `rewrite` to export
@@ -818,6 +848,7 @@ def torch_export_rewrite(
     :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 dump_rewriting: dumps rewriting information in file beginning with that prefix
     :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
@@ -890,6 +921,7 @@ def forward(self, x, y):
                         f"__globals__={sorted(me.__globals__)}"
                     )
                     mod = sys.modules[module]
+                cls_name = module
                 cls = mod
                 name = name
                 to_rewrite = me
@@ -916,7 +948,19 @@ def forward(self, x, y):
         if verbose:
             print(f"[torch_export_rewrite] rewrites {kind} {cls.__name__}.{name}")
         keep[cls, name] = to_rewrite
+        if dump_rewriting:
+            filename = f"{dump_rewriting}.{kind}.{cls_name}.{name}.original.py"
+            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))
+        if dump_rewriting:
+            filename = f"{dump_rewriting}.{kind}.{cls_name}.{name}.rewritten.py"
+            if verbose:
+                print(f"[torch_export_rewrite] dump rewritten code in {filename!r}")
+            with open(filename, "w") as f:
+                f.write(rewr.code)
         setattr(cls, name, rewr.func)
 
     try:

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -1,7 +1,7 @@
 import io
 import functools
 import textwrap
-from typing import Dict, List
+from typing import Any, Dict, List, Optional
 
 __date__ = "2025-03-26"
 
@@ -199,6 +199,28 @@
 """
 
 
+def code_needing_rewriting(cls_name: str) -> Optional[List[Any]]:
+    """
+    Returns a known list of methods or functions to rewrite because of control flow
+    for a specific model class.
+
+    :param cls_name: name of the class
+    :return: a list of rewriting
+
+    .. runpython::
+        :showcode:
+
+        from onnx_diagnostic.torch_models.hghub.hub_data import code_needing_rewriting
+
+        print(code_needing_rewriting("BartForConditionalGeneration"))
+    """
+    if cls_name in {"BartForConditionalGeneration", "BartEncoderLayer"}:
+        import transformers
+
+        return [transformers.models.bart.modeling_bart.BartEncoderLayer.forward]
+    return None
+
+
 @functools.cache
 def load_models_testing() -> List[str]:
     """Returns model ids for testing."""

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -6,6 +6,7 @@
 from ...helpers.config_helper import update_config
 from ...tasks import reduce_model_config, random_input_kwargs
 from .hub_api import task_from_arch, task_from_id, get_pretrained_config
+from .hub_data import code_needing_rewriting
 
 
 def get_untrained_model_with_inputs(
@@ -40,7 +41,8 @@ def get_untrained_model_with_inputs(
     :param add_second_input: provides a second inputs to check a model
         supports different shapes
     :param subfolder: subfolder to use for this model id
-    :return: dictionary with a model, inputs, dynamic shapes, and the configuration
+    :return: dictionary with a model, inputs, dynamic shapes, and the configuration,
+        some necessary rewriting as well
 
     Example:
 
@@ -170,6 +172,10 @@ def get_untrained_model_with_inputs(
         if k.startswith(("inputs", "dynamic_shapes")) and isinstance(v, dict):
             update[k] = filter_out_unexpected_inputs(model, v, verbose=verbose)
     res.update(update)
+
+    rewrite = code_needing_rewriting(model.__class__.__name__)
+    if rewrite:
+        res["rewrite"] = rewrite
     return res
 
 

onnx_diagnostic/torch_models/test_helper.py

@@ -222,6 +222,7 @@ def validate_model(
     optimization: Optional[str] = None,
     quiet: bool = False,
     patch: bool = False,
+    rewrite: bool = False,
     stop_if_static: int = 1,
     dump_folder: Optional[str] = None,
     drop_inputs: Optional[List[str]] = None,
@@ -250,6 +251,8 @@ def validate_model(
     :param quiet: if quiet, catches exception if any issue
     :param patch: applies patches (``patch_transformers=True``) before exporting,
         see :func:`onnx_diagnostic.torch_export_patches.torch_export_patches`
+    :param rewrite: applies known rewriting (``patch_transformers=True``) before exporting,
+        see :func:`onnx_diagnostic.torch_export_patches.torch_export_patches`
     :param stop_if_static: stops if a dynamic dimension becomes static,
         see :func:`onnx_diagnostic.torch_export_patches.torch_export_patches`
     :param dump_folder: dumps everything in a subfolder of this one
@@ -270,6 +273,9 @@ def validate_model(
 
     * ``PRINT_CONFIG``: prints the model configuration
     """
+    assert (
+        not rewrite or patch
+    ), f"rewrite={rewrite}, patch={patch}, patch must be True to enable rewriting"
     summary = version_summary()
     summary.update(
         dict(
@@ -281,6 +287,7 @@ def validate_model(
             version_optimization=optimization or "",
             version_quiet=str(quiet),
             version_patch=str(patch),
+            version_rewrite=str(rewrite),
             version_dump_folder=dump_folder or "",
             version_drop_inputs=str(list(drop_inputs or "")),
             version_ortfusiontype=ortfusiontype or "",
@@ -336,6 +343,13 @@ def validate_model(
     )
     data["input_options"] = iop
     data["model_options"] = mop
+    if "rewrite" in data:
+        if rewrite:
+            summary["model_rewrite"] = str(data["rewrite"])
+            if verbose:
+                print(f"[validate_model] model_rewrite={summary['model_rewrite']}")
+        else:
+            del data["rewrite"]
     if os.environ.get("PRINT_CONFIG", "0") in (1, "1"):
         print("[validate_model] -- PRINT CONFIG")
         print("-- type(config)", type(data["configuration"]))
@@ -446,6 +460,8 @@ def validate_model(
                 patch_transformers=True,
                 stop_if_static=stop_if_static,
                 verbose=max(0, verbose - 1),
+                rewrite=data.get("rewrite", None),
+                dump_rewriting=(os.path.join(dump_folder, "rewrite") if dump_folder else None),
             ) as modificator:
                 data["inputs_export"] = modificator(data["inputs"])  # type: ignore