onnx-diagnostic 0.8.3__py3-none-any.whl → 0.8.5__py3-none-any.whl

onnx_diagnostic/helpers/dot_helper.py

@@ -1,31 +1,13 @@
-from typing import Dict, Set
+from typing import Dict
+import numpy as np
 import onnx
 import onnx.numpy_helper as onh
-from .onnx_helper import onnx_dtype_name, pretty_onnx
-
-
-def _get_hidden_inputs(graph: onnx.GraphProto) -> Set[str]:
-    hidden = set()
-    memo = (
-        {i.name for i in graph.initializer}
-        | {i.values.name for i in graph.sparse_initializer}
-        | {i.name for i in graph.input}
-    )
-    for node in graph.node:
-        for i in node.input:
-            if i not in memo:
-                hidden.add(i)
-        for att in node.attribute:
-            if att.type == onnx.AttributeProto.GRAPH and att.g:
-                hid = _get_hidden_inputs(att.g)
-                less = set(h for h in hid if h not in memo)
-                hidden |= less
-        memo |= set(node.output)
-    return hidden
+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, "("]
+    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":
@@ -42,7 +24,10 @@ def _make_node_label(node: onnx.NodeProto, tiny_inits: Dict[str, str]) -> str:
     els.append(")")
     if node.op_type == "Constant":
         els.extend([" -> ", node.output[0]])
-    return "".join(els)
+    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:
@@ -142,6 +127,7 @@ def to_dot(model: onnx.ModelProto) -> str:
     inits = list(model.graph.initializer)
     tiny_inits = {}
     name_to_ids = {}
+
     for inp in inputs:
         if not inp.name:
             continue
@@ -149,7 +135,30 @@ def to_dot(model: onnx.ModelProto) -> str:
         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)
@@ -161,7 +170,10 @@ def to_dot(model: onnx.ModelProto) -> str:
             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}"];')
@@ -189,7 +201,7 @@ def to_dot(model: onnx.ModelProto) -> str:
             unique = set()
             for att in node.attribute:
                 if att.type == onnx.AttributeProto.GRAPH:
-                    unique |= _get_hidden_inputs(att.g)
+                    unique |= get_hidden_inputs(att.g)
             for i in unique:
                 edge = name_to_ids[i], _mkn(node)  # type: ignore[assignment]
                 if edge in done:

onnx_diagnostic/helpers/helper.py

@@ -2,6 +2,7 @@ import ast
 import enum
 import inspect
 import itertools
+import json
 from dataclasses import is_dataclass, fields
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
@@ -94,6 +95,20 @@ def size_type(dtype: Any) -> int:
     raise AssertionError(f"Unexpected dtype={dtype}")
 
 
+def _string_tensor(obj, cls: str, with_shape: bool, with_device: bool, verbose: int) -> str:
+    from .torch_helper import torch_dtype_to_onnx_dtype
+
+    i = torch_dtype_to_onnx_dtype(obj.dtype)
+    prefix = ("G" if obj.get_device() >= 0 else "C") if with_device else ""
+    if not with_shape:
+        if verbose:
+            print(f"[string_type] {cls}1:{type(obj)}")
+        return f"{prefix}{cls}{i}r{len(obj.shape)}"
+    if verbose:
+        print(f"[string_type] {cls}2:{type(obj)}")
+    return f"{prefix}{cls}{i}s{'x'.join(map(str, obj.shape))}"
+
+
 def string_type(
     obj: Any,
     with_shape: bool = False,
@@ -453,17 +468,7 @@ def string_type(
 
     # Tensors
     if isinstance(obj, torch._subclasses.fake_tensor.FakeTensor):
-        from .torch_helper import torch_dtype_to_onnx_dtype
-
-        i = torch_dtype_to_onnx_dtype(obj.dtype)
-        prefix = ("G" if obj.get_device() >= 0 else "C") if with_device else ""
-        if not with_shape:
-            if verbose:
-                print(f"[string_type] F1:{type(obj)}")
-            return f"{prefix}F{i}r{len(obj.shape)}"
-        if verbose:
-            print(f"[string_type] F2:{type(obj)}")
-        return f"{prefix}F{i}s{'x'.join(map(str, obj.shape))}"
+        return _string_tensor(obj, "F", with_shape, with_device, verbose)
 
     if isinstance(obj, torch.Tensor):
         from .torch_helper import torch_dtype_to_onnx_dtype
@@ -544,6 +549,9 @@ def string_type(
             print(f"[string_type] V6:{type(obj)}")
         return f"{dev}OV{dt}r{len(shape)}"
 
+    if obj.__class__.__name__ == "SymbolicTensor":
+        return _string_tensor(obj, "ST", with_shape, with_device, verbose)
+
     # others classes
 
     if obj.__class__.__name__ == "MambaCache":
@@ -1366,11 +1374,7 @@ def max_diff(
         if hist:
             if isinstance(hist, bool):
                 hist = np.array([0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype)
-            ind = np.digitize(diff.reshape((-1,)), hist, right=True)
-            cou = np.bincount(ind, minlength=ind.shape[0] + 1)
-            res["rep"] = dict(
-                zip([f">{x}" for x in hist], [int(i) for i in (cou.sum() - np.cumsum(cou))])
-            )
+            res["rep"] = {f">{h}": (diff > h).sum().item() for h in hist}
         return res  # type: ignore
 
     if isinstance(expected, torch.Tensor) and isinstance(got, torch.Tensor):
@@ -1486,27 +1490,11 @@ def max_diff(
             dev=dev,
         )
         if hist:
-            if isinstance(hist, list) and len(hist) == 1:
-                res["rep"] = {f">{hist[0]}": (diff > hist[0]).sum().item()}
-            elif isinstance(hist, list) and len(hist) == 2:
-                res["rep"] = {
-                    f">{hist[0]}": (diff > hist[0]).sum().item(),
-                    f">{hist[1]}": (diff > hist[1]).sum().item(),
-                }
-            else:
-                if isinstance(hist, bool):
-                    hist = torch.tensor(
-                        [0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype
-                    )
-                hist = torch.tensor(hist).to(diff.device)
-                ind = torch.bucketize(diff.reshape((-1,)), hist, right=False)
-                cou = torch.bincount(ind, minlength=ind.shape[0] + 1)
-                res["rep"] = dict(
-                    zip(
-                        [f">{x}" for x in hist],
-                        [int(i) for i in (cou.sum() - torch.cumsum(cou, 0))],
-                    )
+            if isinstance(hist, bool):
+                hist = torch.tensor(
+                    [0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype
                 )
+            res["rep"] = {f">{h}": (diff > h).sum().item() for h in hist}
         return res  # type: ignore
 
     if isinstance(expected, int) and isinstance(got, torch.Tensor):
@@ -1743,8 +1731,26 @@ def max_diff(
     )
 
 
-def string_diff(diff: Dict[str, Any]) -> str:
-    """Renders discrepancies return by :func:`max_diff` into one string."""
+def string_diff(diff: Dict[str, Any], js: bool = False, ratio: bool = False, **kwargs) -> str:
+    """
+    Renders discrepancies return by :func:`max_diff` into one string.
+
+    :param diff: differences
+    :param js: json format
+    :param ratio: display mismatch ratio
+    :param kwargs: addition values to add in the json format
+    """
+    if js:
+        if "rep" in diff:
+            rep = diff["rep"]
+            diff = {**{k: v for k, v in diff.items() if k != "rep"}, **rep}
+            if ratio:
+                for k, v in rep.items():
+                    diff[f"%{k}"] = v / diff["n"]
+                diff["mean"] = diff["sum"] / diff["n"]
+            diff.update(kwargs)
+        return json.dumps(diff)
+
     # dict(abs=, rel=, sum=, n=n_diff, dnan=)
     if "dev" in diff:
         ddiff = {k: v for k, v in diff.items() if k != "dev"}