anotehr example

Author: xadupre

_doc/conf.py

@@ -152,12 +152,14 @@
     "within_subsection_order": "ExampleTitleSortKey",
     # errors
     "abort_on_example_error": True,
+    "expected_failing_examples": ["examples/plot_export_locate_issue.py"],
     # recommendation
     "recommender": {"enable": True, "n_examples": 3, "min_df": 3, "max_df": 0.9},
     # ignore capture for matplotib axes
     "ignore_repr_types": "matplotlib\\.(text|axes)",
     # robubstness
     "reset_modules_order": "both",
+    "reset_modules": ("matplotlib", "onnx_diagnostic.reset_torch_transformers"),
 }
 
 if int(os.environ.get("UNITTEST_GOING", "0")):

_doc/examples/plot_export_locate_issue.py

@@ -0,0 +1,84 @@
+"""
+.. _l-plot-export-locale-issue:
+
+==================================================
+Find and fix an export issue due to dynamic shapes
+==================================================
+
+
+A model with an export issue
+============================
+
+The following model implies the first dimension of x is equal to 1
+or equal to the number of element in the list ``ys``.
+It is not really dynamic. It looks obvious here but
+it is difficult to find deep inside a big model.
+"""
+
+import torch
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+
+
+class ModelWithIssue(torch.nn.Module):
+    def forward(self, x: torch.Tensor, ys: list[torch.Tensor]):
+        caty = torch.cat([y.unsqueeze(0) for y in ys], axis=0)
+        z = x * caty
+        return z
+
+
+inputs = (torch.rand(2, 3, 1), [torch.rand(3, 4), torch.rand(3, 4)])
+model = ModelWithIssue()
+model(*inputs)
+
+
+# %%
+# Let's export.
+
+DYN = torch.export.Dim.DYNAMIC
+dyn_shapes = ({0: DYN, 1: DYN}, [{0: DYN, 1: DYN}, {0: DYN, 1: DYN}])
+try:
+    ep = torch.export.export(model, inputs, dynamic_shapes=dyn_shapes)
+    print(ep)
+except Exception as e:
+    print("-- ERROR:")
+    print(e)
+
+# %%
+# The error shows:
+#
+# ::
+#       Constraints violated (L['args'][0][0].size()[0])!
+#           For more information, run with TORCH_LOGS="+dynamic".
+#       - Not all values of RelaxedUnspecConstraint(L['args'][0][0].size()[0])
+#           are valid because L['args'][0][0].size()[0] was inferred to be a constant (2).
+#
+# Where does it happens? That's a tricky question we need to answer.
+# The message is raised from
+# `torch.fx.experimental.symbolic_shapes.ShapeEnv._set_replacement
+# <https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/symbolic_shapes.py#L6239>`_.
+# One way to find the exact location is to retrieve a stack trace
+# by inserting an assert such as the following:
+#
+# ::
+#
+#   assert msg != "range_refined_to_singleton", (
+#       f"A dynamic dimension becomes static! "
+#       f"a={a!r}, tgt={tgt!r}, msg={msg!r}, tgt_bound={tgt_bound}"
+#   )
+#
+# Stop when a dynamic dimension turns static
+# ==========================================
+#
+#
+
+with bypass_export_some_errors(stop_if_static=True, verbose=1):
+    torch.export.export(model, inputs, dynamic_shapes=dyn_shapes)
+
+# The stack trace is quite long but the first line referring to this example
+# is the following one. It points out the line turing a dynamic dimension into
+# static.
+#
+# ::
+#
+#   File "onnx-diagnostic/_doc/examples/plot_export_locate_issue.py", line 25, in forward
+#       z = x * caty

onnx_diagnostic/__init__.py

@@ -5,3 +5,12 @@
 
 __version__ = "0.3.0"
 __author__ = "Xavier Dupré"
+
+
+def reset_torch_transformers(gallery_conf, fname):
+    "Resets torch dynamo for :epkg:`sphinx-gallery`."
+    import matplotlib.pyplot as plt
+    import torch
+
+    plt.style.use("ggplot")
+    torch._dynamo.reset()

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -145,6 +145,11 @@ def _unregister(cls: type, verbose: int = 0):
         # torch >= 2.7
         torch.utils._pytree._deregister_pytree_node(cls)
     optree.unregister_pytree_node(cls, namespace="torch")
+    if cls in torch.utils._pytree.SUPPORTED_NODES:
+        import packaging.version as pv
+
+        if pv.Version(torch.__version__) < pv.Version("2.7.0"):
+            del torch.utils._pytree.SUPPORTED_NODES[cls]
     assert cls not in torch.utils._pytree.SUPPORTED_NODES, (
         f"{cls} was not successful unregistered "
         f"from torch.utils._pytree.SUPPORTED_NODES="
@@ -190,6 +195,7 @@ def bypass_export_some_errors(
     patch_torch: bool = True,
     patch_transformers: bool = False,
     catch_constraints: bool = True,
+    stop_if_static: bool = False,
     verbose: int = 0,
     patch: bool = True,
 ) -> Callable:
@@ -203,8 +209,12 @@ def bypass_export_some_errors(
         as a result, some dynamic dimension may turn into static ones,
         the environment variable ``SKIP_SOLVE_CONSTRAINTS=0``
         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 dyanmic shapes
+        and show a stack trace indicating the exact location of the issue
     :param patch: if False, disable all patches except the registration of
         serialization function
+    :param verbose: to show which patches is applied
 
     The list of available patches.
 
@@ -348,6 +358,18 @@ def bypass_export_some_errors(
                 )
             )
 
+        if stop_if_static:
+            if verbose:
+                print(
+                    "[bypass_export_some_errors] assert when a dynamic dimnension turns static"
+                )
+
+            from torch.fx.experimental.symbolic_shapes import ShapeEnv
+            from .patches.patch_torch import patched_ShapeEnv
+
+            f_shape_env__set_replacement = ShapeEnv._set_replacement
+            ShapeEnv._set_replacement = patched_ShapeEnv._set_replacement
+
         ####################
         # patch transformers
         ####################
@@ -401,6 +423,12 @@ def bypass_export_some_errors(
                 if verbose:
                     print("[bypass_export_some_errors] restored pytorch functions")
 
+            if stop_if_static:
+                if verbose:
+                    print("[bypass_export_some_errors] restored ShapeEnv._set_replacement")
+
+                ShapeEnv._set_replacement = f_shape_env__set_replacement
+
             if catch_constraints:
                 # to catch or skip dynamic_shapes issues
                 torch._export.non_strict_utils.produce_guards_and_solve_constraints = (

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -313,7 +313,10 @@ def _set_replacement(
             #         "Specializing %s to %s", self.var_to_sources[a][0].name(), tgt
             #     )
             #     self.log.debug("SPECIALIZATION", stack_info=True)
-        assert msg != "range_refined_to_singleton", f"{[a, tgt, msg, tgt_bound]}"
+        assert msg != "range_refined_to_singleton", (
+            f"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)
         self.replacements[a] = tgt
         # NB: the replacement may get refined, but the user will find the