onnx-diagnostic 0.7.11__py3-none-any.whl → 0.7.13__py3-none-any.whl

onnx_diagnostic/torch_export_patches/eval/__init__.py

@@ -39,6 +39,9 @@ def evaluation(
         "export-strict",
         "export-nostrict",
         "export-nostrict-decall",
+        "export-strict-oblivious",
+        "export-nostrict-oblivious",
+        "export-nostrict-decall-oblivious",
     ),
     dynamic: Tuple[bool] = (False, True),
     cases: Optional[Union[str, Dict[str, type]]] = None,
@@ -105,9 +108,7 @@ def evaluation(
 
 
 def _flatten_inputs(x: Any) -> List["torch.Tensor"]:  # noqa: F821
-    """
-    Flatten inputs.
-    """
+    """Flatten inputs."""
     if x is None:
         return x
     import torch
@@ -173,6 +174,15 @@ def _clone(x):
     raise TypeError(f"Unable to clone type {type(x)}, x={x} into numpy")
 
 
+def _wrap_torch_export(*args, backed_size_oblivious=False, **kwargs):
+    import torch
+
+    if backed_size_oblivious:
+        with torch.fx.experimental._config.patch(backed_size_oblivious=True):
+            return torch.export.export(*args, **kwargs)
+    return torch.export.export(*args, **kwargs)
+
+
 def _make_exporter_export(
     exporter: str,
     model: "torch.nn.Module",  # noqa: F821
@@ -183,76 +193,36 @@ def _make_exporter_export(
 ) -> Union[Dict, Callable]:
     import torch
 
-    if exporter == "export-strict":
-        try:
-            if verbose >= 2:
-                exported = torch.export.export(
-                    model, inputs, dynamic_shapes=dynamic_shapes, strict=True
-                )
-            else:
-                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
-                    io.StringIO()
-                ):
-                    exported = torch.export.export(
-                        model, inputs, dynamic_shapes=dynamic_shapes, strict=True
-                    )
-        except Exception as e:
-            if not quiet:
-                raise
-            return dict(error=str(e), success=0, error_step="export")
-        if verbose >= 9:
-            print("-- graph")
-            print(exported.graph)
-        return exported.module()
-    if exporter in ("export-strict-dec", "export-strict-decall"):
-        try:
-            if verbose >= 2:
-                exported = torch.export.export(
-                    model, inputs, dynamic_shapes=dynamic_shapes, strict=True
-                )
-                if verbose >= 9:
-                    print("-- graph before decomposition")
-                    print(exported.graph)
-                exported = (
-                    exported.run_decompositions()
-                    if "decall" in exporter
-                    else exported.run_decompositions({})
-                )
-            else:
-                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
-                    io.StringIO()
-                ):
-                    exported = torch.export.export(
-                        model, inputs, dynamic_shapes=dynamic_shapes, strict=True
-                    )
-                    if verbose >= 9:
-                        print("-- graph before decomposition")
-                        print(exported.graph)
-                    exported = (
-                        exported.run_decompositions()
-                        if "decall" in exporter
-                        else exported.run_decompositions({})
-                    )
-        except Exception as e:
-            if not quiet:
-                raise
-            return dict(error=str(e), success=0, error_step="export")
-        if verbose >= 9:
-            print("-- graph after decomposition")
-            print(exported.graph)
-        return exported.module()
-    if exporter == "export-nostrict":
+    backed_size_oblivious = "-oblivious" in exporter
+    strict = "-nostrict" not in exporter
+
+    if exporter in (
+        "export-strict",
+        "export-strict-oblivious",
+        "export-nostrict",
+        "export-nostrict-oblivious",
+        "export-oblivious",
+    ):
         try:
             if verbose >= 2:
-                exported = torch.export.export(
-                    model, inputs, dynamic_shapes=dynamic_shapes, strict=False
+                exported = _wrap_torch_export(
+                    model,
+                    inputs,
+                    dynamic_shapes=dynamic_shapes,
+                    strict=strict,
+                    backed_size_oblivious=backed_size_oblivious,
                 )
             else:
-                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
-                    io.StringIO()
+                with (
+                    contextlib.redirect_stdout(io.StringIO()),
+                    contextlib.redirect_stderr(io.StringIO()),
                 ):
-                    exported = torch.export.export(
-                        model, inputs, dynamic_shapes=dynamic_shapes, strict=False
+                    exported = _wrap_torch_export(
+                        model,
+                        inputs,
+                        dynamic_shapes=dynamic_shapes,
+                        strict=strict,
+                        backed_size_oblivious=backed_size_oblivious,
                     )
         except Exception as e:
             if not quiet:
@@ -262,11 +232,25 @@ def _make_exporter_export(
             print("-- graph")
             print(exported.graph)
         return exported.module()
-    if exporter in ("export-nostrict-dec", "export-nostrict-decall"):
+
+    if exporter in (
+        "export-strict-dec",
+        "export-strict-decall",
+        "export-strict-dec-oblivious",
+        "export-strict-decall-oblivious",
+        "export-nostrict-dec",
+        "export-nostrict-decall",
+        "export-nostrict-dec-oblivious",
+        "export-nostrict-decall-oblivious",
+    ):
         try:
             if verbose >= 2:
-                exported = torch.export.export(
-                    model, inputs, dynamic_shapes=dynamic_shapes, strict=False
+                exported = _wrap_torch_export(
+                    model,
+                    inputs,
+                    dynamic_shapes=dynamic_shapes,
+                    strict=strict,
+                    backed_size_oblivious=backed_size_oblivious,
                 )
                 if verbose >= 9:
                     print("-- graph before decomposition")
@@ -277,11 +261,16 @@ def _make_exporter_export(
                     else exported.run_decompositions({})
                 )
             else:
-                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
-                    io.StringIO()
+                with (
+                    contextlib.redirect_stdout(io.StringIO()),
+                    contextlib.redirect_stderr(io.StringIO()),
                 ):
-                    exported = torch.export.export(
-                        model, inputs, dynamic_shapes=dynamic_shapes, strict=False
+                    exported = _wrap_torch_export(
+                        model,
+                        inputs,
+                        dynamic_shapes=dynamic_shapes,
+                        strict=strict,
+                        backed_size_oblivious=backed_size_oblivious,
                     )
                     if verbose >= 9:
                         print("-- graph before decomposition")
@@ -299,6 +288,7 @@ def _make_exporter_export(
             print("-- graph after decomposition")
             print(exported.graph)
         return exported.module()
+
     if exporter == "export-tracing":
         from experimental_experiment.torch_interpreter.tracing import CustomTracer
 
@@ -307,8 +297,9 @@ def _make_exporter_export(
                 graph = CustomTracer().trace(model)
                 mod = torch.fx.GraphModule(model, graph)
             else:
-                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
-                    io.StringIO()
+                with (
+                    contextlib.redirect_stdout(io.StringIO()),
+                    contextlib.redirect_stderr(io.StringIO()),
                 ):
                     graph = CustomTracer().trace(model)
                     mod = torch.fx.GraphModule(model, graph)
@@ -353,8 +344,9 @@ def _make_exporter_onnx(
                     return_builder=True,
                 )
             else:
-                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
-                    io.StringIO()
+                with (
+                    contextlib.redirect_stdout(io.StringIO()),
+                    contextlib.redirect_stderr(io.StringIO()),
                 ):
                     onx, builder = to_onnx(
                         model,
@@ -387,8 +379,9 @@ def _make_exporter_onnx(
                     report=True,
                 ).model_proto
             else:
-                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
-                    io.StringIO()
+                with (
+                    contextlib.redirect_stdout(io.StringIO()),
+                    contextlib.redirect_stderr(io.StringIO()),
                 ):
                     onx = torch.onnx.export(
                         model,
@@ -422,8 +415,9 @@ def _make_exporter_onnx(
                 ep.optimize()
                 onx = ep.model_proto
             else:
-                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
-                    io.StringIO()
+                with (
+                    contextlib.redirect_stdout(io.StringIO()),
+                    contextlib.redirect_stderr(io.StringIO()),
                 ):
                     ep = torch.onnx.export(
                         model,
@@ -446,6 +440,74 @@ def _make_exporter_onnx(
     raise AssertionError(f"Unexpected exporter={exporter!r}")
 
 
+def _compares_on_one_example(
+    model: Callable, inputs: Tuple[Any, ...], mod: Callable, verbose: int, quiet: bool
+) -> Tuple[Any, Any, Dict]:
+    from onnx_diagnostic.helpers import max_diff, string_type
+
+    try:
+        expected = model(*_clone(inputs))
+    except Exception as e:
+        if not quiet:
+            raise RuntimeError(
+                f"eager mode failed=\n{string_type(inputs, with_shape=True)} "
+                f"\nmodel=\n{type(model)}"
+            ) from e
+        res = dict(error=str(e), success=0, error_step="eager")
+        return None, None, res
+    try:
+        got = mod(*inputs)
+    except Exception as e:
+        if not quiet:
+            raise RuntimeError(
+                f"onnxruntime failed, feeds=\n{string_type(inputs, with_shape=True)}"
+            ) from e
+        res = dict(error=str(e), success=0, error_step="run.0")
+        return expected, None, res
+
+    try:
+        disc = max_diff(expected, got)
+    except Exception as e:
+        if not quiet:
+            raise
+        res = dict(error=str(e), success=0, error_step="discrepancy")
+        return expected, got, res
+
+    if verbose >= 5 and np.isinf(disc["abs"]):
+        print("[run_exporter] comparison issues with")
+        print(f"--   inputs={string_type(inputs[0], with_shape=True, limit=20)}")
+        print(f"-- expected={string_type(expected, with_shape=True, limit=20)}")
+        print(f"--      got={string_type(got, with_shape=True, limit=20)}")
+    elif verbose >= 9:
+        print("[run_exporter] inputs and outputs")
+        print(
+            f"--   inputs="
+            f"{string_type(inputs[0], with_shape=True, with_min_max=True, limit=20)}"
+        )
+        print(
+            f"-- expected="
+            f"{string_type(expected, with_shape=True, with_min_max=True, limit=20)}"
+        )
+        print(f"--      got={string_type(got, with_shape=True, with_min_max=True, limit=20)}")
+    del disc["n"]
+    del disc["sum"]
+    disc.update(
+        dict(
+            success=1 if disc["abs"] < 0.1 else 0,
+            model_cls=model.__class__,  # type: ignore[dict-item]
+            exported=mod,  # type: ignore[dict-item]
+        )
+    )
+    if disc["abs"] >= 0.1:
+        disc["error"] = "diff.0"
+        disc["error_step"] = "diff.0"
+        if verbose >= 9:
+            max_diff(expected, got, verbose=verbose)
+    else:
+        disc["success"] = 1
+    return expected, got, disc
+
+
 def run_exporter(
     exporter: str,
     cls_model: type,
@@ -473,6 +535,7 @@ def run_exporter(
 
     model = cls_model()
     inputs = cls_model._inputs
+    valid = getattr(cls_model, "_valid", None)
     if isinstance(inputs, tuple):
         inputs = [inputs]
     if dynamic:
@@ -566,74 +629,38 @@ def run_exporter(
         mod = lambda *args, names=names: sess.run(None, _make_feeds(names, args))  # noqa: E731
 
     # we need to clone for models modifying the inputs
-    try:
-        expected = model(*_clone(inputs[0]))
-    except Exception as e:
-        if not quiet:
-            raise RuntimeError(
-                f"eager mode failed=\n{string_type(inputs[0], with_shape=True)} "
-                f"\nmodel=\n{type(model)}"
-            ) from e
-        res = dict(error=str(e), success=0, error_step="eager")
-        res.update(base)
-        return res
-    try:
-        got = mod(*inputs[0])
-    except Exception as e:
-        if not quiet:
-            raise RuntimeError(
-                f"onnxruntime failed, feeds=\n{string_type(inputs[0], with_shape=True)} "
-                f"\nmodel=\n{pretty_onnx(onx)}"
-            ) from e
-        res = dict(error=str(e), success=0, error_step="run.0")
-        res.update(base)
-        return res
-
-    base["expected"] = expected
-    base["obtained"] = got
-
-    try:
-        disc = max_diff(expected, got)
-    except Exception as e:
-        if not quiet:
-            raise
-        res = dict(error=str(e), success=0, error_step="discrepancy")
-        res.update(base)
-        return res
+    expected, got, disc = _compares_on_one_example(model, inputs[0], mod, verbose, quiet)
+    if expected is not None:
+        base["expected"] = expected
+    if got is not None:
+        base["obtained"] = got
+    disc.update(base)
+    disc["onnx"] = onx  # type: ignore[dict-item]
 
-    if verbose >= 5 and np.isinf(disc["abs"]):
-        print("[run_exporter] comparison issues with")
-        print(f"--   inputs={string_type(inputs[0], with_shape=True, limit=20)}")
-        print(f"-- expected={string_type(expected, with_shape=True, limit=20)}")
-        print(f"--      got={string_type(got, with_shape=True, limit=20)}")
-    elif verbose >= 9:
-        print("[run_exporter] inputs and outputs")
-        print(
-            f"--   inputs="
-            f"{string_type(inputs[0], with_shape=True, with_min_max=True, limit=20)}"
-        )
-        print(
-            f"-- expected="
-            f"{string_type(expected, with_shape=True, with_min_max=True, limit=20)}"
-        )
-        print(f"--      got={string_type(got, with_shape=True, with_min_max=True, limit=20)}")
-    del disc["n"]
-    del disc["sum"]
-    disc.update(
-        dict(
-            success=1 if disc["abs"] < 0.1 else 0,
-            model_cls=model.__class__,
-            exported=mod,  # type: ignore[dict-item]
-            onnx=onx,  # type: ignore[dict-item]
-        )
-    )
-    if disc["abs"] >= 0.1:
-        disc["error"] = "diff.0"
-        disc["error_step"] = "diff.0"
-        if verbose >= 9:
-            max_diff(expected, got, verbose=verbose)
-    else:
-        disc["success"] = 1
+    if valid is not None:
+        for valid_inputs in valid:
+            expected, got, _disc = _compares_on_one_example(
+                model, valid_inputs, mod, verbose, quiet
+            )
+            if "abs" not in disc and (np.isnan(disc["abs"]) or disc["abs"] > 1e-3):
+                _disc["issue-abs"] = disc["abs"]
+                _disc["issue-rel"] = disc["rel"]
+                _disc["issue-inputs"] = string_type(
+                    valid_inputs, with_shape=True, with_min_max=True
+                )
+                _disc["issue-expected"] = string_type(
+                    expected, with_shape=True, with_min_max=True
+                )
+                _disc["issue-obtained"] = string_type(got, with_shape=True, with_min_max=True)
+                if not quiet:
+                    raise RuntimeError(
+                        f"validation failed,"
+                        f"\n-- inputs=\n{string_type(_disc['issue-inputs'])} "
+                        f"\n-- exporter={exporter!r}\n-- dynamic_shapes={dynamic_shapes}, "
+                        f"\n-- expected={_disc['issue-expected']}"
+                        f"\n-- obtained={_disc['issue-obtained']}"
+                    )
+                break
 
     if dynamic and onx is not None:
         ds = []

onnx_diagnostic/torch_export_patches/eval/model_cases.py

@@ -861,7 +861,7 @@ class CreateFromShapeThroughFunction(torch.nn.Module):
         y = torch.ones((x.shape[0], dy1))
         return y
 
-    _inputs = [(torch.rand((4, 4)),), (torch.rand((5, 5)),)]
+    _inputs = [(torch.rand((4, 4)),)]
     _dynamic = {"x": {0: DIM("dx"), 1: DIM("dy")}}
 
 
@@ -881,3 +881,21 @@ class VmapPython(torch.nn.Module):
 
     _inputs = [(torch.tensor([1.0, 2.0, 3.0]), torch.tensor([0.1, 0.2, 0.3]))]
     _dynamic = {"x": {0: DYN}, "y": {0: DYN}}
+
+
+class ExportWithDimension0(torch.nn.Module):
+    def forward(self, x):
+        return x @ torch.arange(x.shape[1], dtype=torch.float32).reshape((-1, 1))
+
+    _inputs = [(torch.empty((0, 3), dtype=torch.float32),)]
+    _dynamic = {"x": {0: DYN, 1: DYN}}
+    _valid = [(torch.rand((2, 3), dtype=torch.float32),)]
+
+
+class ExportWithDimension1(torch.nn.Module):
+    def forward(self, x):
+        return x @ torch.arange(x.shape[1], dtype=torch.float32).reshape((-1, 1))
+
+    _inputs = [(torch.zeros((1, 3), dtype=torch.float32),)]
+    _dynamic = {"x": {0: DYN, 1: DYN}}
+    _valid = [(torch.rand((2, 3), dtype=torch.float32),)]

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -83,7 +83,7 @@ def patch_module_or_classes(mod, verbose: int = 0) -> Dict[type, Dict[type, Call
             continue
 
         original = cls._PATCHED_CLASS_
-        methods = cls._PATCHES_
+        methods = [_ for _ in cls._PATCHES_ if _ is not None]
         if verbose:
             print(f"[patch_module_or_classes] {name}.{cls.__name__}: {', '.join(methods)}")
 
@@ -254,22 +254,36 @@ def torch_export_patches(
     may appear ``AssertionError: Mutating module attribute _seen_tokens during export.``.
     It can be avoided by setting ``strict=False`` when call :func:`torch.export.export`.
     """
+    if verbose:
+        print(f"[torch_export_patches] patch_sympy={patch_sympy!r}")
+        print(f"                     . patch_torch={patch_torch!r}")
+        print(f"                     . patch_transformers={patch_transformers!r}")
+        print(f"                     . patch_diffusers={patch_diffusers!r}")
+        print(f"                     . catch_constraints={catch_constraints!r}")
+        print(f"                     . stop_if_static={stop_if_static!r}")
+        print(f"                     . patch={patch!r}")
+        print(f"                     . custom_patches={custom_patches!r}")
+        print(f"[torch_export_patches] dump_rewriting={dump_rewriting!r}")
+
     if rewrite:
         from .patch_module import torch_export_rewrite
 
-        with torch_export_rewrite(
-            rewrite=rewrite, dump_rewriting=dump_rewriting, verbose=verbose
-        ), torch_export_patches(  # type: ignore[var-annotated]
-            patch_sympy=patch_sympy,
-            patch_torch=patch_torch,
-            patch_transformers=patch_transformers,
-            patch_diffusers=patch_diffusers,
-            catch_constraints=catch_constraints,
-            stop_if_static=stop_if_static,
-            verbose=verbose,
-            patch=patch,
-            custom_patches=custom_patches,
-        ) as f:
+        with (
+            torch_export_rewrite(
+                rewrite=rewrite, dump_rewriting=dump_rewriting, verbose=verbose
+            ),
+            torch_export_patches(  # type: ignore[var-annotated]
+                patch_sympy=patch_sympy,
+                patch_torch=patch_torch,
+                patch_transformers=patch_transformers,
+                patch_diffusers=patch_diffusers,
+                catch_constraints=catch_constraints,
+                stop_if_static=stop_if_static,
+                verbose=verbose,
+                patch=patch,
+                custom_patches=custom_patches,
+            ) as f,
+        ):
             try:
                 yield f
             finally:
@@ -330,6 +344,7 @@ def torch_export_patches(
                 patched_infer_size,
                 patched_vmap,
                 patched__broadcast_shapes,
+                patched__constrain_user_specified_dimhint_range,
                 _catch_produce_guards_and_solve_constraints,
                 patch__check_input_constraints_for_graph,
             )
@@ -360,6 +375,14 @@ def torch_export_patches(
             torch._refs._broadcast_shapes = patched__broadcast_shapes
             torch._meta_registrations._broadcast_shapes = patched__broadcast_shapes
 
+            # torch._export.non_strict_utils._constrain_user_specified_dimhint_range
+            f___constrain_user_specified_dimhint_range = (
+                torch._export.non_strict_utils._constrain_user_specified_dimhint_range
+            )
+            torch._export.non_strict_utils._constrain_user_specified_dimhint_range = (
+                patched__constrain_user_specified_dimhint_range
+            )
+
         # torch._export.non_strict_utils.produce_guards_and_solve_constraints
         if patch_torch and catch_constraints:
             if verbose:
@@ -558,6 +581,9 @@ def torch_export_patches(
                 torch._subclasses.fake_impls.infer_size = f_infer_size
                 torch._refs._broadcast_shapes = f__broadcast_shapes
                 torch._meta_registrations._broadcast_shapes = f__broadcast_shapes
+                torch._export.non_strict_utils._constrain_user_specified_dimhint_range = (
+                    f___constrain_user_specified_dimhint_range
+                )
 
                 if verbose:
                     print("[torch_export_patches] restored pytorch functions")

onnx_diagnostic/torch_export_patches/patch_inputs.py

@@ -189,19 +189,23 @@ def convert_dynamic_axes_into_dynamic_shapes(
     return (), updated_kwargs, dynamic_shapes
 
 
-def use_dyn_not_str(dynamic_shapes: Any) -> Any:
+def use_dyn_not_str(dynamic_shapes: Any, default_value=None) -> Any:
     """
     Some functions returns dynamic shapes as string.
     This functions replaces them with ``torch.export.Dim.DYNAMIC``.
+    ``default_value=torch.export.Dim.AUTO`` changes the default value.
     """
     if isinstance(dynamic_shapes, list):
-        return [use_dyn_not_str(a) for a in dynamic_shapes]
+        return [use_dyn_not_str(a, default_value=default_value) for a in dynamic_shapes]
     if isinstance(dynamic_shapes, tuple):
-        return tuple(use_dyn_not_str(a) for a in dynamic_shapes)
+        return tuple(use_dyn_not_str(a, default_value=default_value) for a in dynamic_shapes)
     if isinstance(dynamic_shapes, dict):
-        return {k: use_dyn_not_str(v) for k, v in dynamic_shapes.items()}
+        return {
+            k: use_dyn_not_str(v, default_value=default_value)
+            for k, v in dynamic_shapes.items()
+        }
     if isinstance(dynamic_shapes, set):
-        return {use_dyn_not_str(a) for a in dynamic_shapes}
+        return {use_dyn_not_str(a, default_value=default_value) for a in dynamic_shapes}
     if isinstance(dynamic_shapes, str):
-        return torch.export.Dim.DYNAMIC
+        return torch.export.Dim.DYNAMIC if default_value is None else default_value
     return dynamic_shapes