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

onnx_diagnostic/helpers/dot_helper.py

@@ -0,0 +1,210 @@
+from typing import Dict, Set
+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
+
+
+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, "("]
+    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]])
+    return "".join(els)
+
+
+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)
+    for init in inits:
+        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:
+        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)

onnx_diagnostic/helpers/helper.py

@@ -529,16 +529,20 @@ def string_type(
                 return "OV(NO-NUMPY:FIXIT)"
             if verbose:
                 print(f"[string_type] V4:{type(obj)}")
-            return f"OV({string_type(t, with_shape=with_shape, with_min_max=with_min_max)})"
+            dev = ("G" if obj.device_name() == "Cuda" else "C") if with_device else ""
+            return (
+                f"{dev}OV({string_type(t, with_shape=with_shape, with_min_max=with_min_max)})"
+            )
         dt = obj.element_type()
         shape = obj.shape()
+        dev = ("G" if obj.device_name() == "Cuda" else "C") if with_device else ""
         if with_shape:
             if verbose:
                 print(f"[string_type] V5:{type(obj)}")
-            return f"OV{dt}s{'x'.join(map(str, shape))}"
+            return f"{dev}OV{dt}s{'x'.join(map(str, shape))}"
         if verbose:
             print(f"[string_type] V6:{type(obj)}")
-        return f"OV{dt}r{len(shape)}"
+        return f"{dev}OV{dt}r{len(shape)}"
 
     # others classes
 
@@ -990,7 +994,7 @@ def max_diff(
     _index: int = 0,
     allow_unique_tensor_with_list_of_one_element: bool = True,
     hist: Optional[Union[bool, List[float]]] = None,
-) -> Dict[str, Union[float, int, Tuple[int, ...]]]:
+) -> Dict[str, Union[float, int, Tuple[Any, ...]]]:
     """
     Returns the maximum discrepancy.
 
@@ -1015,6 +1019,7 @@ def max_diff(
         output, this number will be the number of elements
         of this output
     * dnan: difference in the number of nan
+    * dev: tensor on the same device, if applicable
 
     You may use :func:`string_diff` to display the discrepancies in one string.
     """
@@ -1167,7 +1172,7 @@ def max_diff(
 
         if verbose >= 6:
             print(f"[max_diff] list,tuple,6: {string_type(expected)} ? {string_type(got)}")
-        am, rm, sm, n, dn, drep = 0, 0, 0.0, 0.0, 0, None
+        am, rm, sm, n, dn, drep, dd = 0, 0, 0.0, 0.0, 0, None, None
         for ip, (e, g) in enumerate(zip(expected, got)):
             d = max_diff(
                 e,
@@ -1199,7 +1204,15 @@ def max_diff(
                 else:
                     for k, v in d["rep"].items():
                         drep[k] += v
+            if "dev" in d and d["dev"] is not None:
+                if dd is None:
+                    dd = d["dev"]
+                else:
+                    dd += d["dev"]  # type: ignore[operator]
+
         res = dict(abs=am, rel=rm, sum=sm, n=n, dnan=dn)
+        if dd is not None:
+            res["dev"] = dd
         if drep:
             res["rep"] = drep
         return res  # type: ignore
@@ -1233,33 +1246,42 @@ def max_diff(
     import torch
 
     if isinstance(expected, np.ndarray) or isinstance(got, np.ndarray):
+        dev = None
         if isinstance(expected, torch.Tensor):
             from .torch_helper import to_numpy
 
+            dev = 0 if expected.device.type == "cpu" else 1
             expected = to_numpy(expected)
         if isinstance(got, torch.Tensor):
             from .torch_helper import to_numpy
 
+            dev = 0 if got.device.type == "cpu" else 1
             got = to_numpy(got)
         if verbose >= 6:
             print(f"[max_diff] tensor: {string_type(expected)} ? {string_type(got)}")
 
         if _index < begin or (end != -1 and _index >= end):
             # out of boundary
-            return dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0)
+            res = dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0)
+            if dev is not None:
+                res["dev"] = dev  # type: ignore[operator]
+            return res  # type: ignore[return-value]
         if isinstance(expected, (int, float)):
             if isinstance(got, np.ndarray) and len(got.shape) == 0:
                 got = float(got)
             if isinstance(got, (int, float)):
                 if expected == got:
                     return dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0)
-                return dict(
+                res = dict(
                     abs=abs(expected - got),
                     rel=abs(expected - got) / (abs(expected) + 1e-5),
                     sum=abs(expected - got),
                     n=1,
                     dnan=0,
                 )
+                if dev is not None:
+                    res["dev"] = dev
+                return res  # type: ignore[return-value]
             return dict(abs=np.inf, rel=np.inf, sum=np.inf, n=np.inf, dnan=np.inf)
         if expected.dtype in (np.complex64, np.complex128):
             if got.dtype == expected.dtype:
@@ -1339,6 +1361,8 @@ def max_diff(
         res: Dict[str, float] = dict(  # type: ignore
             abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff, argm=argm
         )
+        if dev is not None:
+            res["dev"] = dev
         if hist:
             if isinstance(hist, bool):
                 hist = np.array([0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype)
@@ -1352,9 +1376,14 @@ def max_diff(
     if isinstance(expected, torch.Tensor) and isinstance(got, torch.Tensor):
         if verbose >= 6:
             print(f"[max_diff] tensor: {string_type(expected)} ? {string_type(got)}")
+        dev = 0 if expected.device == got.device else 1
         if _index < begin or (end != -1 and _index >= end):
             # out of boundary
-            return dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0)
+            if verbose >= 10:
+                if debug_info:
+                    print("\n".join(debug_info))
+                print("[max_diff] out of boundary")
+            return dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0, dev=dev)
         if expected.dtype in (torch.complex64, torch.complex128):
             if got.dtype == expected.dtype:
                 got = torch.view_as_real(got)
@@ -1448,31 +1477,63 @@ def max_diff(
                 )
 
         res: Dict[str, float] = dict(  # type: ignore
-            abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff, argm=argm
+            abs=abs_diff,
+            rel=rel_diff,
+            sum=sum_diff,
+            n=n_diff,
+            dnan=nan_diff,
+            argm=argm,
+            dev=dev,
         )
         if hist:
-            if isinstance(hist, bool):
-                hist = torch.tensor(
-                    [0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype
-                )
-            hist = 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, 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))],
+                    )
                 )
-            )
         return res  # type: ignore
 
+    if isinstance(expected, int) and isinstance(got, torch.Tensor):
+        # a size
+        if verbose >= 6:
+            print(f"[max_diff] int: {string_type(expected)} ? {string_type(got)}")
+        if got.shape != tuple():
+            return dict(  # type: ignore
+                abs=np.inf,
+                rel=np.inf,
+                sum=np.inf,
+                n=np.inf,
+                dnan=np.inf,
+                argm=np.inf,
+            )
+        return dict(  # type: ignore
+            abs=abs(expected - got.item()),
+            rel=abs((expected - got.item()) / max(1, expected)),
+            sum=abs(expected - got.item()),
+            n=1,
+            dnan=0,
+        )
+
     if "SquashedNormal" in expected.__class__.__name__:
         if verbose >= 6:
             print(f"[max_diff] SquashedNormal: {string_type(expected)} ? {string_type(got)}")
-        values = (
-            expected.mean.detach().to("cpu"),
-            expected.scale.detach().to("cpu"),
-        )
+        values = (expected.mean, expected.scale)
         return max_diff(values, got, debug_info=_debug("SquashedNormal"), **_dkws)
 
     if expected.__class__ in torch.utils._pytree.SUPPORTED_NODES:
@@ -1677,7 +1738,7 @@ def max_diff(
 
     raise AssertionError(
         f"Not implemented with implemented with expected="
-        f"{string_type(expected)}, got={string_type(got)},\n"
+        f"{string_type(expected)} ({type(expected)}), got={string_type(got)},\n"
         f"level={level}"
     )
 
@@ -1685,6 +1746,9 @@ def max_diff(
 def string_diff(diff: Dict[str, Any]) -> str:
     """Renders discrepancies return by :func:`max_diff` into one string."""
     # dict(abs=, rel=, sum=, n=n_diff, dnan=)
+    if "dev" in diff:
+        ddiff = {k: v for k, v in diff.items() if k != "dev"}
+        return f"{string_diff(ddiff)}, dev={diff['dev']}"
     suffix = ""
     if "rep" in diff:
         rows = []

onnx_diagnostic/helpers/mini_onnx_builder.py

@@ -159,7 +159,9 @@ class MiniOnnxBuilder:
         """
         if not tensors:
             # empty list
-            self.nodes.append(oh.make_node("SequenceEmpty", [], [name]))
+            self.nodes.append(
+                oh.make_node("SequenceEmpty", [], [name], dtype=TensorProto.FLOAT)
+            )
             tensor_type_proto = oh.make_tensor_type_proto(
                 elem_type=TensorProto.FLOAT, shape=None
             )

onnx_diagnostic/helpers/model_builder_helper.py

@@ -28,10 +28,37 @@ def download_model_builder_to_cache(
     if file_path.exists():
         return file_path
 
+    builders = cache_dir / "builders"
+    if not builders.exists():
+        builders.mkdir(parents=True, exist_ok=True)
+
+    for subfile in [
+        "__init__.py",
+        "base.py",
+        "chatglm.py",
+        "ernie.py",
+        "gemma.py",
+        "gptoss.py",
+        "granite.py",
+        "llama.py",
+        "mistral.py",
+        "nemotron.py",
+        "olmo.py",
+        "phi.py",
+        "qwen.py",
+        "smollm.py",
+    ]:
+        u = f"{'/'.join(url.split('/')[:-1])}/builders/{subfile}"
+        response = requests.get(u)
+        response.raise_for_status()
+        with open(builders / subfile, "wb") as f:
+            f.write(response.content)
+
     response = requests.get(url)
     response.raise_for_status()
     with open(file_path, "wb") as f:
         f.write(response.content)
+
     return file_path
 
 

onnx_diagnostic/helpers/onnx_helper.py

@@ -3,7 +3,7 @@ import json
 import os
 import sys
 import warnings
-from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Union
+from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Union
 import numpy as np
 import numpy.typing as npt
 import onnx
@@ -15,6 +15,7 @@ from onnx import (
     GraphProto,
     ModelProto,
     NodeProto,
+    OperatorSetIdProto,
     TensorProto,
     ValueInfoProto,
     load as onnx_load,
@@ -1195,3 +1196,104 @@ def shadowing_names(
         existing |= not_empty
         created |= not_empty
     return shadow, post_shadow, created
+
+
+def extract_subset_of_nodes(
+    model: ModelProto,
+    name: str,
+    node_index: Optional[int] = None,
+    cut_points: Optional[Set[str]] = None,
+) -> List[NodeProto]:
+    """
+    Extracts the minimal subgraphs which can produce the output ``name``
+    knowing ``cut_points``.
+
+    :param model: original model
+    :param name: result name
+    :param node_index: if the node index is known, otherwise searches for it
+    :param cut_points: the known results or input name otherwise
+    :return: minimal list of nodes
+    """
+    if node_index is None:
+        for i, node in enumerate(model.graph.node):
+            if name in node.output:
+                node_index = i
+                break
+    assert (
+        node_index is not None
+        and node_index < len(model.graph.node)
+        and name in model.graph.node[node_index].output
+    ), f"node_index is still empty or wrong for result {name!r}"
+    if cut_points is None:
+        cut_points = {n.name for n in model.graph.input} | {
+            n.name for n in model.graph.initializer
+        }
+    elif model.graph.initializer:
+        cut_points = cut_points | {n.name for n in model.graph.initializer}
+
+    node = model.graph.node[node_index]
+    selected = {node_index}
+    current_node_index = node_index
+    current_input_index = 0
+    intermediate = {name}
+    inputs = set(k for k in node.input if k)
+    while not (inputs <= cut_points) and current_node_index >= 0:
+        node = model.graph.node[current_node_index]
+        if current_input_index == 0:
+            needs = [o for o in node.output if o in intermediate and o not in cut_points]
+            if needs:
+                selected.add(current_node_index)
+            else:
+                current_node_index -= 1
+                continue
+        res = node.input[current_input_index]
+        if res not in cut_points:
+            intermediate.add(res)
+        current_input_index += 1
+        if current_input_index >= len(node.input):
+            current_node_index -= 1
+            current_input_index = 0
+
+    return [model.graph.node[i] for i in sorted(selected)]
+
+
+def make_submodel(
+    nodes: List[NodeProto],
+    ir_version: int,
+    opset_imports: List[OperatorSetIdProto],
+    output_names: List[str],
+    type_rank_fn: Callable[[str], Tuple[int, int]],
+) -> ModelProto:
+    """
+    Creates a model with the given list of nodes.
+    It computes the minimum list of inputs needed for this model.
+    The function assumes the nodes are sorted.
+    It does not handle yet subgraphs.
+
+    :param nodes: list of nodes
+    :param ir_version: ir version
+    :param opset_imports: opset import
+    :param output_names: desired outputs
+    :param function: function returning the type and the rank of a result
+    :return: model proto
+    """
+
+    def _mkv_(name, itype, irank):
+        return oh.make_tensor_value_info(name, itype, [f"{name}_d{i}" for i in range(irank)])
+
+    not_known: Set[str] = set()
+    for node in nodes[::-1]:
+        not_known -= set(node.output)
+        not_known |= set(node.input)
+
+    model = oh.make_model(
+        oh.make_graph(
+            nodes,
+            "submodel",
+            [_mkv_(n, *type_rank_fn(n)) for n in sorted(not_known)],
+            [_mkv_(n, *type_rank_fn(n)) for n in sorted(output_names)],
+        ),
+        ir_version=ir_version,
+        opset_imports=opset_imports,
+    )
+    return model