Add option to disable patches for torch in command line validate

Author: xadupre

onnx_diagnostic/_command_lines_parser.py

@@ -306,7 +306,7 @@ def __call__(self, parser, namespace, values, option_string=None):
                 value = split_items[1]
 
                 if value in ("True", "true", "False", "false"):
-                    d[key] = bool(value)
+                    d[key] = value in ("True", "true")
                     continue
                 try:
                     d[key] = int(value)
@@ -323,6 +323,54 @@ def __call__(self, parser, namespace, values, option_string=None):
         setattr(namespace, self.dest, d)
 
 
+class _BoolOrParseDictPatch(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+
+        if not values:
+            return
+        if len(values) == 1 and values[0] in (
+            "True",
+            "False",
+            "true",
+            "false",
+            "0",
+            "1",
+            0,
+            1,
+        ):
+            setattr(namespace, self.dest, values[0] in ("True", "true", 1, "1"))
+            return
+        d = getattr(namespace, self.dest) or {}
+        if not isinstance(d, dict):
+            d = {
+                "patch_sympy": d,
+                "patch_torch": d,
+                "patch_transformers": d,
+                "patch_diffusers": d,
+            }
+        for item in values:
+            split_items = item.split("=", 1)
+            key = split_items[0].strip()  # we remove blanks around keys, as is logical
+            value = split_items[1]
+
+            if value in ("True", "true", "False", "false"):
+                d[key] = value in ("True", "true")
+                continue
+            try:
+                d[key] = int(value)
+                continue
+            except (TypeError, ValueError):
+                pass
+            try:
+                d[key] = float(value)
+                continue
+            except (TypeError, ValueError):
+                pass
+            d[key] = _parse_json(value)
+
+        setattr(namespace, self.dest, d)
+
+
 def get_parser_validate() -> ArgumentParser:
     parser = ArgumentParser(
         prog="validate",
@@ -383,8 +431,13 @@ def get_parser_validate() -> ArgumentParser:
     parser.add_argument(
         "--patch",
         default=True,
-        action=BooleanOptionalAction,
-        help="Applies patches before exporting.",
+        action=_BoolOrParseDictPatch,
+        nargs="*",
+        help="Applies patches before exporting, it can be a boolean "
+        "to enable to disable the patches or be more finetuned. It is possible to "
+        "disable patch for torch by adding "
+        '--patch "patch_sympy=False" --patch "patch_torch=False", '
+        "default is True.",
     )
     parser.add_argument(
         "--rewrite",

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -361,7 +361,7 @@ def torch_export_patches(
             torch._meta_registrations._broadcast_shapes = patched__broadcast_shapes
 
         # torch._export.non_strict_utils.produce_guards_and_solve_constraints
-        if catch_constraints:
+        if patch_torch and catch_constraints:
             if verbose:
                 print("[torch_export_patches] modifies shape constraints")
             f_produce_guards_and_solve_constraints = (
@@ -513,7 +513,7 @@ def torch_export_patches(
                 if verbose:
                     print("[torch_export_patches] restored pytorch functions")
 
-            if stop_if_static:
+            if patch_torch and stop_if_static:
                 if verbose:
                     print("[torch_export_patches] restored ShapeEnv._set_replacement")
 
@@ -529,7 +529,7 @@ def torch_export_patches(
                         print("[torch_export_patches] restored ShapeEnv._check_frozen")
                     ShapeEnv._check_frozen = f_shape_env__check_frozen
 
-            if catch_constraints:
+            if patch_torch and catch_constraints:
                 # to catch or skip dynamic_shapes issues
                 torch._export.non_strict_utils.produce_guards_and_solve_constraints = (
                     f_produce_guards_and_solve_constraints

onnx_diagnostic/torch_models/validate.py

@@ -426,6 +426,7 @@ def validate_model(
             print(f"[validate_model] validate model id {model_id!r}, subfolder={subfolder!r}")
         else:
             print(f"[validate_model] validate model id {model_id!r}")
+        print(f"[validate_model] patch={patch!r}")
         if model_options:
             print(f"[validate_model] model_options={model_options!r}")
         print(f"[validate_model] get dummy inputs with input_options={input_options}...")