improve catching

Author: xadupre

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -211,7 +211,7 @@ def bypass_export_some_errors(
     patch_torch: bool = True,
     patch_transformers: bool = False,
     catch_constraints: bool = True,
-    stop_if_static: bool = False,
+    stop_if_static: int = 0,
     verbose: int = 0,
     patch: bool = True,
 ) -> Callable:
@@ -227,7 +227,9 @@ def bypass_export_some_errors(
         can be put to stop at that stage.
     :param stop_if_static: see example :ref:`l-plot-export-locale-issue`,
         to stop the export as soon as an issue is detected with dynamic shapes
-        and show a stack trace indicating the exact location of the issue
+        and show a stack trace indicating the exact location of the issue,
+        ``if stop_if_static > 1``, more methods are replace to catch more
+        issues
     :param patch: if False, disable all patches except the registration of
         serialization function
     :param verbose: to show which patches is applied
@@ -375,17 +377,24 @@ def bypass_export_some_errors(
             )
 
         if stop_if_static:
+            from torch.fx.experimental.symbolic_shapes import ShapeEnv
+            from .patches.patch_torch import patched_ShapeEnv
+
             if verbose:
                 print(
                     "[bypass_export_some_errors] assert when a dynamic dimension turns static"
                 )
-
-            from torch.fx.experimental.symbolic_shapes import ShapeEnv
-            from .patches.patch_torch import patched_ShapeEnv
+                print("[bypass_export_some_errors] replaces ShapeEnv._set_replacement")
 
             f_shape_env__set_replacement = ShapeEnv._set_replacement
             ShapeEnv._set_replacement = patched_ShapeEnv._set_replacement
 
+            if stop_if_static > 1:
+                if verbose:
+                    print("[bypass_export_some_errors] replaces ShapeEnv._check_frozen")
+                f_shape_env__check_frozen = ShapeEnv._check_frozen
+                ShapeEnv._check_frozen = patched_ShapeEnv._check_frozen
+
         ####################
         # patch transformers
         ####################
@@ -444,6 +453,10 @@ def bypass_export_some_errors(
                     print("[bypass_export_some_errors] restored ShapeEnv._set_replacement")
 
                 ShapeEnv._set_replacement = f_shape_env__set_replacement
+                if stop_if_static > 1:
+                    if verbose:
+                        print("[bypass_export_some_errors] restored ShapeEnv._check_frozen")
+                    ShapeEnv._check_frozen = f_shape_env__check_frozen
 
             if catch_constraints:
                 # to catch or skip dynamic_shapes issues

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -150,6 +150,16 @@ def patched__broadcast_shapes(*_shapes):
 
 class patched_ShapeEnv:
 
+    def _check_frozen(
+        self, expr: "sympy.Basic", concrete_val: "sympy.Basic"  # noqa: F821
+    ) -> None:
+        if self.frozen:
+            self.counter["ignored_backward_guard"] += 1
+            raise AssertionError(
+                f"[patched_ShapeEnv] Ignored guard {expr} == {concrete_val}, "
+                f"this could result in accuracy problems."
+            )
+
     def _set_replacement(
         self, a: "sympy.Symbol", tgt: "sympy.Expr", msg: str  # noqa: F821
     ) -> None:
@@ -314,7 +324,7 @@ def _set_replacement(
             #     )
             #     self.log.debug("SPECIALIZATION", stack_info=True)
         assert msg != "range_refined_to_singleton", (
-            f"A dynamic dimension becomes static! "
+            f"patched_ShapeEnv: A dynamic dimension becomes static! "
             f"a={a!r}, tgt={tgt!r}, msg={msg!r}, tgt_bound={tgt_bound}"
         )
         # log.info("set_replacement %s = %s (%s) %s", a, tgt, msg, tgt_bound)