onnx-diagnostic 0.8.2__py3-none-any.whl → 0.8.4__py3-none-any.whl

onnx_diagnostic/ext_test_case.py

@@ -9,6 +9,7 @@ import itertools
 import logging
 import os
 import re
+import shutil
 import sys
 import unittest
 import warnings
@@ -63,7 +64,7 @@ def skipif_ci_apple(msg) -> Callable:
     return lambda x: x
 
 
-def unit_test_going():
+def unit_test_going() -> bool:
     """
     Enables a flag telling the script is running while testing it.
     Avois unit tests to be very long.
@@ -147,7 +148,7 @@ def hide_stdout(f: Optional[Callable] = None) -> Callable:
 
     def wrapper(fct):
         def call_f(self):
-            if os.environ.get("UNHIDE", ""):
+            if os.environ.get("UNHIDE", "") in (1, "1", "True", "true"):
                 fct(self)
                 return
             st = StringIO()
@@ -609,6 +610,21 @@ def requires_onnxruntime(version: str, msg: str = "") -> Callable:
     return lambda x: x
 
 
+def has_onnxruntime(version: str, msg: str = "") -> Callable:
+    """Skips a unit test if :epkg:`onnxruntime` is not recent enough."""
+    import packaging.version as pv
+    import onnxruntime
+
+    if not hasattr(onnxruntime, "__version__"):
+        # development version
+        return True
+
+    if pv.Version(onnxruntime.__version__) < pv.Version(version):
+        msg = f"onnxruntime version {onnxruntime.__version__} < {version}: {msg}"
+        return False
+    return True
+
+
 def has_onnxruntime_training(push_back_batch: bool = False):
     """Tells if onnxruntime_training is installed."""
     try:
@@ -742,8 +758,15 @@ class ExtTestCase(unittest.TestCase):
     _warns: List[Tuple[str, int, Warning]] = []
     _todos: List[Tuple[Callable, str]] = []
 
+    def unit_test_going(self) -> bool:
+        """
+        Enables a flag telling the script is running while testing it.
+        Avois unit tests to be very long.
+        """
+        return unit_test_going()
+
     @property
-    def verbose(self):
+    def verbose(self) -> int:
         "Returns the the value of environment variable ``VERBOSE``."
         return int(os.environ.get("VERBOSE", "0"))
 
@@ -768,13 +791,13 @@ class ExtTestCase(unittest.TestCase):
         cls._todos.append((f, msg))
 
     @classmethod
-    def ort(cls):
+    def ort(cls) -> unittest.__class__:
         import onnxruntime
 
         return onnxruntime
 
     @classmethod
-    def to_onnx(self, *args, **kwargs):
+    def to_onnx(self, *args, **kwargs) -> "ModelProto":  # noqa: F821
         from experimental_experiment.torch_interpreter import to_onnx
 
         return to_onnx(*args, **kwargs)
@@ -806,18 +829,29 @@ class ExtTestCase(unittest.TestCase):
             os.makedirs(folder)
         return folder
 
-    def dump_onnx(
-        self,
-        name: str,
-        proto: Any,
-        folder: Optional[str] = None,
-    ) -> str:
+    def clean_dump(self, folder: str = "dump_test"):
+        """Cleans this folder."""
+        for item in os.listdir(folder):
+            item_path = os.path.join(folder, item)
+            if os.path.isfile(item_path) or os.path.islink(item_path):
+                os.remove(item_path)
+            elif os.path.isdir(item_path):
+                shutil.rmtree(item_path)
+
+    def dump_onnx(self, name: str, proto: Any, folder: Optional[str] = None) -> str:
         """Dumps an onnx file."""
         fullname = self.get_dump_file(name, folder=folder)
         with open(fullname, "wb") as f:
             f.write(proto.SerializeToString())
         return fullname
 
+    def dump_text(self, name: str, text: str, folder: Optional[str] = None) -> str:
+        """Dumps text in a file."""
+        fullname = self.get_dump_file(name, folder=folder)
+        with open(fullname, "w") as f:
+            f.write(text)
+        return fullname
+
     def assertExists(self, name):
         """Checks the existing of a file."""
         if not os.path.exists(name):
@@ -1094,10 +1128,15 @@ class ExtTestCase(unittest.TestCase):
             value = numpy.array(value).astype(expected.dtype)
         self.assertEqualArray(expected, value, atol=atol, rtol=rtol)
 
-    def check_ort(self, onx: "onnx.ModelProto") -> bool:  # noqa: F821
+    def check_ort(
+        self, onx: "onnx.ModelProto"  # noqa: F821
+    ) -> "onnxruntime.InferenceSession":  # noqa: F821
         from onnxruntime import InferenceSession
 
-        return InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"])
+        return InferenceSession(
+            onx if isinstance(onx, str) else onx.SerializeToString(),
+            providers=["CPUExecutionProvider"],
+        )
 
     def assertRaise(self, fct: Callable, exc_type: type[Exception], msg: Optional[str] = None):
         """In the name"""
@@ -1137,7 +1176,7 @@ class ExtTestCase(unittest.TestCase):
         if not full.endswith(suffix):
             raise AssertionError(f"suffix={suffix!r} does not end string  {full!r}.")
 
-    def capture(self, fct: Callable):
+    def capture(self, fct: Callable) -> Tuple[Any, str, str]:
         """
         Runs a function and capture standard output and error.
 
@@ -1179,9 +1218,9 @@ class ExtTestCase(unittest.TestCase):
     def assert_onnx_disc(
         self,
         test_name: str,
-        proto: "onnx.ModelProto",  # noqa: F821
+        proto: Union[str, "onnx.ModelProto"],  # noqa: F821
         model: "torch.nn.Module",  # noqa: F821
-        inputs: Union[Tuple[Any], Dict[str, Any]],
+        inputs: Union[Tuple[Any], Dict[str, Any], List[Any]],
         verbose: int = 0,
         atol: float = 1e-5,
         rtol: float = 1e-3,
@@ -1189,6 +1228,7 @@ class ExtTestCase(unittest.TestCase):
         expected: Optional[Any] = None,
         use_ort: bool = False,
         ort_optimized_graph: bool = False,
+        ep: Optional[Union["torch.export.ExportedProgram", str]] = None,  # noqa: F821
         **kwargs,
     ):
         """
@@ -1208,6 +1248,7 @@ class ExtTestCase(unittest.TestCase):
         :param copy_inputs: to copy the inputs
         :param use_ort: use :class:`onnxruntime.InferenceSession`
         :param ort_optimized_graph: dumps the optimized onnxruntime graph
+        :param ep: exported program (or saved exported program)
         :param kwargs: arguments sent to
             :class:`onnx_diagnostic.helpers.ort_session.InferenceSessionForTorch`
         """
@@ -1223,71 +1264,135 @@ class ExtTestCase(unittest.TestCase):
             name = f"{test_name}.onnx"
             if verbose:
                 print(f"[{vname}] save the onnx model into {name!r}")
+            model_file = None
             if isinstance(proto, str):
+                model_file = proto
                 name = proto
                 proto = onnx.load(name)
-            else:
+            elif not self.unit_test_going():
                 assert isinstance(
                     proto, onnx.ModelProto
                 ), f"Unexpected type {type(proto)} for proto"
                 name = self.dump_onnx(name, proto)
-            if verbose:
+            if verbose and not self.unit_test_going():
                 print(f"[{vname}] file size {os.stat(name).st_size // 2**10:1.3f} kb")
         if verbose:
             print(f"[{vname}] make feeds {string_type(inputs, **kws)}")
+
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+            if expected is not None:
+                expected = [expected]
+
+        gots = []
         if use_ort:
             assert isinstance(
                 proto, onnx.ModelProto
             ), f"Unexpected type {type(proto)} for proto"
-            feeds = make_feeds(proto, inputs, use_numpy=True, copy=True)
             import onnxruntime
 
-            if verbose:
-                print(f"[{vname}] create onnxruntime.InferenceSession")
             options = onnxruntime.SessionOptions()
             if ort_optimized_graph:
                 options.optimized_model_filepath = f"{name}.optort.onnx"
+            if "log_severity_level" in kwargs:
+                options.log_severity_level = kwargs["log_severity_level"]
+            if "log_verbosity_level" in kwargs:
+                options.log_verbosity_level = kwargs["log_verbosity_level"]
+            providers = kwargs.get("providers", ["CPUExecutionProvider"])
+            if verbose:
+                print(f"[{vname}] create onnxruntime.InferenceSession with {providers}")
             sess = onnxruntime.InferenceSession(
-                proto.SerializeToString(),
-                options,
-                providers=kwargs.get("providers", ["CPUExecutionProvider"]),
+                model_file or proto.SerializeToString(), options, providers=providers
             )
-            if verbose:
-                print(f"[{vname}] run ort feeds {string_type(feeds, **kws)}")
-            got = sess.run(None, feeds)
+            for inp in inputs:
+                feeds = make_feeds(proto, inp, use_numpy=True, copy=True)
+                if verbose:
+                    print(f"[{vname}] run ort feeds {string_type(feeds, **kws)}")
+                got = sess.run(None, feeds)
+                gots.append(got)
         else:
-            feeds = make_feeds(proto, inputs, copy=True)
             if verbose:
                 print(f"[{vname}] create InferenceSessionForTorch")
             sess = InferenceSessionForTorch(proto, **kwargs)
-            if verbose:
-                print(f"[{vname}] run orttorch feeds {string_type(feeds, **kws)}")
-            got = sess.run(None, feeds)
+            for inp in inputs:
+                feeds = make_feeds(proto, inp, copy=True)
+                if verbose:
+                    print(f"[{vname}] run orttorch feeds {string_type(feeds, **kws)}")
+                got = sess.run(None, feeds)
+                gots.append(got)
         if verbose:
             print(f"[{vname}] compute expected values")
+
         if expected is None:
             if copy_inputs:
-                expected = (
-                    model(*copy.deepcopy(inputs))
-                    if isinstance(inputs, tuple)
-                    else model(**copy.deepcopy(inputs))
-                )
+                expected = [
+                    (
+                        model(*copy.deepcopy(inp))
+                        if isinstance(inp, tuple)
+                        else model(**copy.deepcopy(inp))
+                    )
+                    for inp in inputs
+                ]
             else:
-                expected = model(*inputs) if isinstance(inputs, tuple) else model(**inputs)
+                expected = [
+                    model(*inp) if isinstance(inp, tuple) else model(**inp) for inp in inputs
+                ]
+
         if verbose:
             print(f"[{vname}] expected {string_type(expected, **kws)}")
             print(f"[{vname}] obtained {string_type(got, **kws)}")
-        diff = max_diff(expected, got, flatten=True)
-        if verbose:
-            print(f"[{vname}] diff {string_diff(diff)}")
-        assert (
-            isinstance(diff["abs"], float)
-            and isinstance(diff["rel"], float)
-            and not numpy.isnan(diff["abs"])
-            and diff["abs"] <= atol
-            and not numpy.isnan(diff["rel"])
-            and diff["rel"] <= rtol
-        ), f"discrepancies in {test_name!r}, diff={string_diff(diff)}"
+
+        if ep:
+            if isinstance(ep, str):
+                if verbose:
+                    print(f"[{vname}] load exported program {ep!r}")
+                import torch
+
+                ep = torch.export.load(ep)
+
+            ep_model = ep.module()  # type: ignore[union-attr]
+            for expe, inp, got in zip(expected, inputs, gots):
+                ep_inputs = copy.deepcopy(inp) if copy_inputs else inp
+                ep_expected = (
+                    ep_model(*copy.deepcopy(ep_inputs))
+                    if isinstance(ep_inputs, tuple)
+                    else ep_model(**copy.deepcopy(ep_inputs))
+                )
+                if verbose:
+                    print(f"[{vname}] ep_expected {string_type(ep_expected, **kws)}")
+                ep_diff = max_diff(expe, ep_expected, hist=[0.1, 0.01])
+                if verbose:
+                    print(f"[{vname}] ep_diff {string_diff(ep_diff)}")
+                assert (
+                    isinstance(ep_diff["abs"], float)
+                    and isinstance(ep_diff["rel"], float)
+                    and not numpy.isnan(ep_diff["abs"])
+                    and ep_diff["abs"] <= atol
+                    and not numpy.isnan(ep_diff["rel"])
+                    and ep_diff["rel"] <= rtol
+                ), (
+                    f"discrepancies in {test_name!r} between the exported program "
+                    f"and the exported model diff={string_diff(ep_diff)}"
+                )
+                ep_nx_diff = max_diff(ep_expected, got, flatten=True, hist=[0.1, 0.01])
+                if verbose:
+                    print(f"[{vname}] ep_nx_diff {string_diff(ep_nx_diff)}")
+
+        for expe, got in zip(expected, gots):
+            diff = max_diff(expe, got, flatten=True, hist=[0.1, 0.01])
+            if verbose:
+                print(f"[{vname}] diff {string_diff(diff)}")
+            assert (
+                isinstance(diff["abs"], float)
+                and isinstance(diff["rel"], float)
+                and not numpy.isnan(diff["abs"])
+                and diff["abs"] <= atol
+                and not numpy.isnan(diff["rel"])
+                and diff["rel"] <= rtol
+            ), (
+                f"discrepancies in {test_name!r} between the model and "
+                f"the onnx model diff={string_diff(diff)}"
+            )
 
     def _debug(self):
         "Tells if DEBUG=1 is set up."
@@ -1298,6 +1403,16 @@ class ExtTestCase(unittest.TestCase):
 
         return string_type(*args, **kwargs)
 
+    def max_diff(self, *args, **kwargs):
+        from .helpers import max_diff
+
+        return max_diff(*args, **kwargs)
+
+    def use_dyn_not_str(self, *args, **kwargs):
+        from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+        return use_dyn_not_str(*args, *kwargs)
+
     def subloop(self, *args, verbose: int = 0):
         "Loops over elements and calls :meth:`unittests.TestCase.subTest`."
         if len(args) == 1:

onnx_diagnostic/helpers/cache_helper.py

@@ -80,7 +80,7 @@ def flatten_unflatten_for_dynamic_shapes(
     start = 0
     end = 0
     subtrees = []
-    for subspec in spec.children_specs:
+    for subspec in (spec.children() if hasattr(spec, "children") else spec.children_specs):
         end += subspec.num_leaves
         value = subspec.unflatten(flat[start:end])
         value = flatten_unflatten_for_dynamic_shapes(

onnx_diagnostic/helpers/dot_helper.py

@@ -0,0 +1,222 @@
+from typing import Dict
+import numpy as np
+import onnx
+import onnx.numpy_helper as onh
+from ..reference import ExtendedReferenceEvaluator as Inference
+from .onnx_helper import onnx_dtype_name, pretty_onnx, get_hidden_inputs
+
+
+def _make_node_label(node: onnx.NodeProto, tiny_inits: Dict[str, str]) -> str:
+    els = [f"{node.domain}.\\n{node.op_type}" if node.domain else node.op_type, "\\n("]
+    ee = [tiny_inits.get(i, ".") if i else "" for i in node.input]
+    for att in node.attribute:
+        if att.name == "to":
+            ee.append(f"{att.name}={onnx_dtype_name(att.i)}")
+        elif att.name in {"to", "axis", "value_int", "stash_type", "start", "end"}:
+            ee.append(f"{att.name}={att.i}")
+        elif att.name in {"value_float"}:
+            ee.append(f"{att.name}={att.f}")
+        elif att.name in {"value_floats"}:
+            ee.append(f"{att.name}={att.floats}")
+        elif att.name in {"value_ints", "perm"}:
+            ee.append(f"{att.name}={att.ints}")
+    els.append(", ".join(ee))
+    els.append(")")
+    if node.op_type == "Constant":
+        els.extend([" -> ", node.output[0]])
+    res = "".join(els)
+    if len(res) < 40:
+        return res.replace("\\n(", "(")
+    return res
+
+
+def _make_edge_label(value_info: onnx.ValueInfoProto, multi_line: bool = False) -> str:
+    itype = value_info.type.tensor_type.elem_type
+    if itype == onnx.TensorProto.UNDEFINED:
+        return ""
+    shape = tuple(
+        d.dim_param if d.dim_param else d.dim_value
+        for d in value_info.type.tensor_type.shape.dim
+    )
+    res = [
+        str(a)
+        for a in [("?" if isinstance(s, str) and s.startswith("unk") else s) for s in shape]
+    ]
+    sshape = ",".join(res)
+    if multi_line and len(sshape) > 30:
+        sshape = ",\\n".join(res)
+    return f"{onnx_dtype_name(itype)}({sshape})"
+
+
+def to_dot(model: onnx.ModelProto) -> str:
+    """
+    Converts a model into a dot graph.
+    Here is an example:
+
+    .. gdot::
+        :script: DOT-SECTION
+        :process:
+
+        from onnx_diagnostic.helpers.dot_helper import to_dot
+        from onnx_diagnostic.export.api import to_onnx
+        from onnx_diagnostic.torch_export_patches import torch_export_patches
+        from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
+
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        with torch_export_patches(patch_transformers=True):
+            em = to_onnx(model, inputs, dynamic_shapes=ds, exporter="custom")
+        dot = to_dot(em.model_proto)
+        print("DOT-SECTION", dot)
+
+    Or this one obtained with :func:`torch.onnx.export`.
+
+    .. gdot::
+        :script: DOT-SECTION
+        :process:
+
+        from onnx_diagnostic.helpers.dot_helper import to_dot
+        from onnx_diagnostic.export.api import to_onnx
+        from onnx_diagnostic.torch_export_patches import torch_export_patches
+        from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
+
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        with torch_export_patches(patch_transformers=True):
+            em = to_onnx(model, kwargs=inputs, dynamic_shapes=ds, exporter="onnx-dynamo")
+        dot = to_dot(em.model_proto)
+        print("DOT-SECTION", dot)
+    """
+    _unique: Dict[int, int] = {}
+
+    def _mkn(obj: object) -> int:
+        id_obj = id(obj)
+        if id_obj in _unique:
+            return _unique[id_obj]
+        i = len(_unique)
+        _unique[id_obj] = i
+        return i
+
+    model = onnx.shape_inference.infer_shapes(model)
+
+    op_type_colors = {
+        "Shape": "#d2a81f",
+        "MatMul": "#ee9999",
+        "Transpose": "#ee99ee",
+        "Reshape": "#eeeeee",
+        "Squeeze": "#eeeeee",
+        "Unsqueeze": "#eeeeee",
+    }
+
+    edge_label = {}
+    for val in model.graph.value_info:
+        edge_label[val.name] = _make_edge_label(val, multi_line=True)
+
+    rows = [
+        "digraph {",
+        (
+            "  graph [rankdir=TB, splines=true, overlap=false, nodesep=0.2, "
+            "ranksep=0.2, fontsize=8];"
+        ),
+        '  node [style="rounded,filled", color="#888888", fontcolor="#222222", shape=box];',
+        "  edge [arrowhead=vee, fontsize=7, labeldistance=-5, labelangle=0];",
+    ]
+    inputs = list(model.graph.input)
+    outputs = list(model.graph.output)
+    nodes = list(model.graph.node)
+    inits = list(model.graph.initializer)
+    tiny_inits = {}
+    name_to_ids = {}
+
+    for inp in inputs:
+        if not inp.name:
+            continue
+        lab = _make_edge_label(inp)
+        rows.append(f'  I_{_mkn(inp)} [label="{inp.name}\\n{lab}", fillcolor="#aaeeaa"];')
+        name_to_ids[inp.name] = f"I_{_mkn(inp)}"
+        edge_label[inp.name] = _make_edge_label(inp, multi_line=True)
+
+    # Small constant --> initializer
+    output_names = {n.name for n in outputs}
+    for node in nodes:
+        if node.op_type != "Constant" or node.output[0] in output_names:
+            continue
+        skip = False
+        for att in node.attribute:
+            if att.name == "value" and (
+                len(att.t.dims) > 1 or np.prod(tuple(att.t.dims)) > 10
+            ):
+                skip = True
+                break
+        if skip:
+            continue
+
+        sess = Inference(node)
+        value = sess.run(None, {})[0]
+        inits.append(onh.from_array(value, name=node.output[0]))
+
+    for init in inits:
+        if init.name in name_to_ids:
+            # hide optional inputs
+            continue
+        shape = tuple(init.dims)
+        if len(shape) == 0 or (len(shape) == 1 and shape[0] < 10):
+            a = onh.to_array(init)
+            tiny_inits[init.name] = (
+                str(a) if len(shape) == 0 else f"[{', '.join([str(i) for i in a])}]"
+            )
+        else:
+            ls = f"{onnx_dtype_name(init.data_type)}({', '.join(map(str,shape))})"
+            rows.append(f'  i_{_mkn(init)} [label="{init.name}\\n{ls}", fillcolor="#cccc00"];')
+            name_to_ids[init.name] = f"i_{_mkn(init)}"
+            edge_label[init.name] = ls
+
+    for node in nodes:
+        if node.op_type == "Constant" and node.output[0] in tiny_inits:
+            continue
+        color = op_type_colors.get(node.op_type, "#cccccc")
+        label = _make_node_label(node, tiny_inits)
+        rows.append(f'  {node.op_type}_{_mkn(node)} [label="{label}", fillcolor="{color}"];')
+        name_to_ids.update({o: f"{node.op_type}_{_mkn(node)}" for o in node.output if o})
+
+    # nodes
+    done = set()
+    for node in nodes:
+        names = list(node.input)
+        for i in names:
+            if not i or i in tiny_inits:
+                continue
+            if i not in name_to_ids:
+                raise ValueError(f"Unable to find {i!r}\n{pretty_onnx(model)}")
+            edge = name_to_ids[i], f"{node.op_type}_{_mkn(node)}"
+            if edge in done:
+                continue
+            done.add(edge)
+            lab = edge_label.get(i, "")
+            if lab:
+                ls = ",".join([f'label="{lab}"'])
+                lab = f" [{ls}]"
+            rows.append(f"  {edge[0]} -> {edge[1]}{lab};")
+        if node.op_type in {"Scan", "Loop", "If"}:
+            unique = set()
+            for att in node.attribute:
+                if att.type == onnx.AttributeProto.GRAPH:
+                    unique |= get_hidden_inputs(att.g)
+            for i in unique:
+                edge = name_to_ids[i], _mkn(node)  # type: ignore[assignment]
+                if edge in done:
+                    continue
+                done.add(edge)
+                rows.append(f"  {edge[0]} -> {edge[1]} [style=dotted];")
+
+    # outputs
+    for out in outputs:
+        if not out.name:
+            continue
+        lab = _make_edge_label(out)
+        rows.append(f'  O_{_mkn(out)} [label="{out.name}\\n{lab}", fillcolor="#aaaaee"];')
+        edge = name_to_ids[out.name], f"O_{_mkn(out)}"
+        rows.append(f"  {edge[0]} -> {edge[1]};")
+
+    rows.append("}")
+    return "\n".join(rows)