onnx-diagnostic 0.6.1__py3-none-any.whl → 0.6.3__py3-none-any.whl

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@ Patches, Investigates onnx models.
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.6.1"
+__version__ = "0.6.3"
 __author__ = "Xavier Dupré"

onnx_diagnostic/_command_lines_parser.py

@@ -191,24 +191,45 @@ def get_parser_find() -> ArgumentParser:
         "--names",
         type=str,
         required=False,
-        help="names to look at comma separated values",
+        help="names to look at comma separated values, if 'SHADOW', "
+        "search for shadowing names",
     )
     parser.add_argument(
         "-v",
         "--verbose",
         default=0,
+        type=int,
         required=False,
         help="verbosity",
     )
+    parser.add_argument(
+        "--v2",
+        default=False,
+        action=BooleanOptionalAction,
+        help="use enumerate_results instead of onnx_find",
+    )
     return parser
 
 
 def _cmd_find(argv: List[Any]):
-    from .helpers.onnx_helper import onnx_find
+    from .helpers.onnx_helper import onnx_find, enumerate_results, shadowing_names
 
     parser = get_parser_find()
     args = parser.parse_args(argv[1:])
-    onnx_find(args.input, verbose=args.verbose, watch=set(args.names.split(",")))
+    if args.names == "SHADOW":
+        onx = onnx.load(args.input, load_external_data=False)
+        s, ps = shadowing_names(onx)[:2]
+        print(f"shadowing names: {s}")
+        print(f"post-shadowing names: {ps}")
+    elif args.v2:
+        onx = onnx.load(args.input, load_external_data=False)
+        res = list(
+            enumerate_results(onx, name=set(args.names.split(",")), verbose=args.verbose)
+        )
+        if not args.verbose:
+            print("\n".join(map(str, res)))
+    else:
+        onnx_find(args.input, verbose=args.verbose, watch=set(args.names.split(",")))
 
 
 def get_parser_config() -> ArgumentParser:

onnx_diagnostic/doc.py

@@ -1,3 +1,7 @@
+from typing import Optional
+import numpy as np
+
+
 def reset_torch_transformers(gallery_conf, fname):
     "Resets torch dynamo for :epkg:`sphinx-gallery`."
     import matplotlib.pyplot as plt
@@ -30,3 +34,45 @@ def plot_legend(
     ax.grid(False)
     ax.set_axis_off()
     return ax
+
+
+def rotate_align(ax, angle=15, align="right"):
+    """Rotates x-label and align them to thr right. Returns ax."""
+    for label in ax.get_xticklabels():
+        label.set_rotation(angle)
+        label.set_horizontalalignment(align)
+    return ax
+
+
+def save_fig(ax, name: str):
+    """Applies ``tight_layout`` and saves the figures. Returns ax."""
+    import matplotlib.pyplot as plt
+
+    plt.tight_layout()
+    fig = ax.get_figure()
+    fig.savefig(name)
+    return ax
+
+
+def title(ax: "plt.axes", title: str) -> "plt.axes":  # noqa: F821
+    "Adds a title to axes and returns them."
+    ax.set_title(title)
+    return ax
+
+
+def plot_histogram(
+    tensor: np.ndarray,
+    ax: Optional["plt.axes"] = None,  # noqa: F821
+    bins: int = 30,
+    color: str = "orange",
+    alpha: float = 0.7,
+) -> "plt.axes":  # noqa: F821
+    "Computes the distribution for a tensor."
+    if ax is None:
+        import matplotlib.pyplot as plt
+
+        ax = plt.gca()
+        ax.cla()
+    ax.hist(tensor, bins=30, color="orange", alpha=0.7)
+    ax.set_yscale("log")
+    return ax

onnx_diagnostic/helpers/doc_helper.py

@@ -0,0 +1,163 @@
+import os
+from typing import Dict, List, Optional, Tuple
+import onnx
+import onnx.helper as oh
+import torch
+from ..reference.torch_ops import OpRunKernel, OpRunTensor
+from .torch_helper import onnx_dtype_to_torch_dtype, torch_dtype_to_onnx_dtype
+from .ort_session import InferenceSessionForTorch
+
+_SAVED: List[str] = []
+_SAVE_OPTIMIZED_MODEL_ = int(os.environ.get("DUMP_ONNX", "0"))
+
+
+def _get_model_name(op_name: str, provider: str) -> Optional[str]:
+    if _SAVE_OPTIMIZED_MODEL_:
+        name = f"dump_doc_layer_norm_{provider}_{len(_SAVED)}.onnx"
+        _SAVED.append(name)
+        return name
+    return None
+
+
+class LayerNormalizationOrt(OpRunKernel):
+    "LayerNormalization with onnxruntime"
+
+    @classmethod
+    def device_dependent(cls) -> bool:
+        "Needs device."
+        return True
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version=None,
+        device: Optional[torch.device] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        self.axis = self.get_attribute_int(node, "axis", -1)
+        self.epsilon = self.get_attribute_float(node, "epsilon", 1e-5)
+        self.device = device
+        self.stash_type = onnx_dtype_to_torch_dtype(
+            self.get_attribute_int(node, "stash_type", onnx.TensorProto.FLOAT)  # type: ignore[arg-type]
+        )
+        self.compute_std = len(node.output) > 1
+        assert not self.compute_std, (
+            f"This kernel implementation only work when only one output "
+            f"is required but {node.output} were."
+        )
+        self._cache: Dict[Tuple[int, int], onnx.ModelProto] = {}
+        self.is_cpu = torch.device("cpu") == self.device
+
+    def _make_model(self, itype: int, rank: int, has_bias: bool) -> onnx.ModelProto:
+        shape = [*["d{i}" for i in range(rank - 1)], "last"]
+        layer_model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node(
+                        "LayerNormalization",
+                        ["X", "W", "B"] if has_bias else ["X", "W"],
+                        ["Z"],
+                        axis=self.axis,
+                        epsilon=self.epsilon,
+                    )
+                ],
+                "dummy",
+                (
+                    [
+                        oh.make_tensor_value_info("X", itype, shape),
+                        oh.make_tensor_value_info("W", itype, ["last"]),
+                        oh.make_tensor_value_info("B", itype, ["last"]),
+                    ]
+                    if has_bias
+                    else [
+                        oh.make_tensor_value_info("X", itype, shape),
+                        oh.make_tensor_value_info("W", itype, ["last"]),
+                    ]
+                ),
+                [oh.make_tensor_value_info("Z", itype, shape)],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        provider = "CPUExecutionProvider" if self.is_cpu else "CUDAExecutionProvider"
+        self._provider = provider
+        return InferenceSessionForTorch(
+            layer_model,
+            optimized_model_filepath=_get_model_name("layer_norm", provider),
+            providers=[provider],
+        )
+
+    def run(self, x, scale, bias=None):
+        itype = torch_dtype_to_onnx_dtype(x.dtype)
+        rank = len(x.shape)
+        key = itype, rank
+        if key not in self._cache:
+            self._cache[key] = self._make_model(itype, rank, bias is not None)
+        sess = self._cache[key]
+        if self.verbose:
+            print(f"[LayerNormalizationOrt] running on {self._provider!r}")
+        feeds = dict(X=x.tensor, W=scale.tensor)
+        if bias is not None:
+            feeds["B"] = bias.tensor
+        got = sess.run(None, feeds)[0]
+        return OpRunTensor(got)
+
+
+class MatMulOrt(OpRunKernel):
+    "MatMul with onnxruntime"
+
+    @classmethod
+    def device_dependent(cls) -> bool:
+        "Needs device."
+        return True
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version=None,
+        device: Optional[torch.device] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        self.device = device
+        self._cache: Dict[Tuple[int, int, int], onnx.ModelProto] = {}
+        self.is_cpu = torch.device("cpu") == self.device
+
+    def _make_model(self, itype: int, ranka: int, rankb: int) -> onnx.ModelProto:
+        shapea = ["a{i}" for i in range(ranka)]
+        shapeb = ["b{i}" for i in range(rankb)]
+        shapec = ["c{i}" for i in range(max(ranka, rankb))]
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("MatMul", ["A", "B"], ["C"])],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("A", itype, shapea),
+                    oh.make_tensor_value_info("B", itype, shapeb),
+                ],
+                [oh.make_tensor_value_info("C", itype, shapec)],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        provider = "CPUExecutionProvider" if self.is_cpu else "CUDAExecutionProvider"
+        self._provider = provider
+        return InferenceSessionForTorch(
+            model,
+            optimized_model_filepath=_get_model_name("matmul", provider),
+            providers=[provider],
+        )
+
+    def run(self, a, b):
+        itype = torch_dtype_to_onnx_dtype(a.dtype)
+        ranka, rankb = len(a.shape), len(b.shape)
+        key = itype, ranka, rankb
+        if key not in self._cache:
+            self._cache[key] = self._make_model(itype, ranka, rankb)
+        sess = self._cache[key]
+        if self.verbose:
+            print(f"[MatMulOrt] running on {self._provider!r}")
+        feeds = dict(A=a.tensor, B=b.tensor)
+        got = sess.run(None, feeds)[0]
+        return OpRunTensor(got)

onnx_diagnostic/helpers/model_builder_helper.py

@@ -220,6 +220,9 @@ def create_model_builder(
     """
     assert cache_dir, "create_model_builder does not work without cache_dir."
     assert os.path.exists(cache_dir), f"cache_dir={cache_dir!r} does not exists"
+    precision = {"float32": "fp32", "float16": "fp16", "bfloat16": "bfp16"}.get(
+        precision, precision
+    )
     download_model_builder_to_cache()
     builder = import_model_builder()
     io_dtype = builder.set_io_dtype(precision, execution_provider, extra_options)

onnx_diagnostic/helpers/onnx_helper.py

@@ -316,7 +316,7 @@ def check_model_ort(
 
 
 @functools.cache
-def onnx_dtype_name(itype: int) -> str:
+def onnx_dtype_name(itype: int, exc: bool = True) -> str:
     """
     Returns the ONNX name for a specific element type.
 
@@ -335,7 +335,11 @@ def onnx_dtype_name(itype: int) -> str:
             v = getattr(TensorProto, k)
             if v == itype:
                 return k
-    raise ValueError(f"Unexpected value itype: {itype}")
+    if exc:
+        raise ValueError(f"Unexpected value itype: {itype}")
+    if itype == 0:
+        return "UNDEFINED"
+    return "UNEXPECTED"
 
 
 def pretty_onnx(
@@ -365,7 +369,7 @@ def pretty_onnx(
         itype = onx.type.tensor_type.elem_type
         shape = tuple((d.dim_param or d.dim_value) for d in onx.type.tensor_type.shape.dim)
         shape_str = ",".join(map(str, shape))
-        return f"{onnx_dtype_name(itype)}[{shape_str}] {name}"
+        return f"{onnx_dtype_name(itype, exc=False)}[{shape_str}] {name}"
 
     if isinstance(onx, AttributeProto):
         att = onx
@@ -767,7 +771,7 @@ def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
 
 
 def iterator_initializer_constant(
-    model: Union[onnx.FunctionProto, onnx.GraphProto, onnx.ModelProto],
+    model: Union[FunctionProto, GraphProto, ModelProto],
     use_numpy: bool = True,
     prefix: str = "",
 ) -> Iterator[Tuple[str, Union["torch.Tensor", np.ndarray]]]:  # noqa: F821
@@ -779,8 +783,8 @@ def iterator_initializer_constant(
     :param prefix: for subgraph
     :return: iterator
     """
-    if not isinstance(model, onnx.FunctionProto):
-        graph = model if isinstance(model, onnx.GraphProto) else model.graph
+    if not isinstance(model, FunctionProto):
+        graph = model if isinstance(model, GraphProto) else model.graph
         if not use_numpy:
             from .torch_helper import to_tensor
         if prefix:
@@ -791,7 +795,7 @@ def iterator_initializer_constant(
             )
         nodes = graph.node
         name = graph.name
-        if isinstance(model, onnx.ModelProto):
+        if isinstance(model, ModelProto):
             for f in model.functions:
                 yield from iterator_initializer_constant(
                     f, use_numpy=use_numpy, prefix=f"{prefix}{f.name}"
@@ -908,3 +912,283 @@ def tensor_statistics(tensor: Union[np.ndarray, TensorProto]) -> Dict[str, Union
     qu = np.quantile(tensor, ii)
     stat.update({f"q{i}": float(q) for i, q in zip(ii, qu)})
     return stat
+
+
+class NodeCoordinates:
+    """
+    A way to localize a node,
+    path is a tuple of three information, node index, node type, node name.
+    """
+
+    __slots__ = ("node", "path")
+
+    def __init__(
+        self,
+        node: Union[onnx.TensorProto, NodeProto, str],
+        path: Tuple[Tuple[int, str, str], ...],
+    ):
+        assert isinstance(path, tuple), f"Unexpected type {type(path)} for path"
+        assert all(isinstance(t, tuple) for t in path), f"Unexpected type in path={path}"
+        self.node = node
+        self.path = path
+
+    def __str__(self) -> str:
+        "usual"
+        if isinstance(self.node, str):
+            return f"{self.path_to_str()} :: {self.node!r}"
+        return f"{self.path_to_str()} :: {pretty_onnx(self.node)}"
+
+    def path_to_str(self) -> str:
+        "Strings representing coordinates."
+        return "x".join(f"({':'.join(map(str, t))})" for t in self.path)
+
+
+class ResultFound:
+    """
+    Class returned by :func:`enumerate_results`.
+    """
+
+    __slots__ = ("consumer", "name", "producer")
+
+    def __init__(
+        self,
+        name: str,
+        producer: Optional[NodeCoordinates],
+        consumer: Optional[NodeCoordinates],
+    ):
+        assert isinstance(name, str), f"unexpected type {type(name)} for name"
+        self.name = name
+        self.producer = producer
+        self.consumer = consumer
+
+    def __str__(self) -> str:
+        "usuals"
+        return (
+            f"<< {self.name} - {self.consumer}"
+            if self.producer is None
+            else f">> {self.name} - {self.producer}"
+        )
+
+
+def enumerate_results(
+    proto: Union[FunctionProto, GraphProto, ModelProto, Sequence[NodeProto]],
+    name: Union[Set[str], str],
+    verbose: int = 0,
+    coordinates: Optional[List[Tuple[int, str, str]]] = None,
+) -> Iterator[ResultFound]:
+    """
+    Iterates on all nodes, attributes to find where a name is used.
+
+    :param proto: a proto
+    :param name: name or names to find
+    :param verbose: verbosity
+    :param coordinates: coordinates of a node
+    :return: iterator on :class:`ResultFound`
+    """
+    if not isinstance(name, set):
+        name = {name}
+    coordinates = coordinates or []
+    assert all(
+        isinstance(c, tuple) for c in coordinates
+    ), f"Unexpected type in coordinates={coordinates}"
+    indent = "  " * len(coordinates)
+    if isinstance(proto, ModelProto):
+        if verbose:
+            print(f"[enumerate_results] {indent}searching for {name!r} into ModelProto...")
+        yield from enumerate_results(proto.graph, name, verbose=verbose)
+    elif isinstance(proto, FunctionProto):
+        if verbose:
+            print(f"[enumerate_results] {indent}searching for {name!r} into FunctionProto...")
+        for i in proto.input:
+            if i in name:
+                r = ResultFound(
+                    i,
+                    NodeCoordinates(i, tuple([*coordinates, (-1, "INPUT", "")])),  # noqa: C409
+                    None,
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+        yield from enumerate_results(proto.node, name, verbose=verbose)
+        for i in proto.output:
+            if i in name:
+                r = ResultFound(
+                    i,
+                    None,
+                    NodeCoordinates(
+                        i, tuple([*coordinates, (len(proto.node), "OUTPUT", "")])  # noqa: C409
+                    ),
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+    elif isinstance(proto, GraphProto):
+        if verbose:
+            print(f"[enumerate_results] {indent}searching for {name!r} into GraphProto...")
+        for i in proto.initializer:
+            if i.name in name:
+                r = ResultFound(
+                    i.name,
+                    NodeCoordinates(i, tuple([*coordinates, (-1, "INIT", "")])),  # noqa: C409
+                    None,
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+        for i in proto.sparse_initializer:
+            if i.name in name:
+                r = ResultFound(
+                    i.name,
+                    NodeCoordinates(i, tuple([*coordinates, (-1, "INIT", "")])),  # noqa: C409
+                    None,
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+        for i in proto.input:
+            if i.name in name:
+                r = ResultFound(
+                    i.name,
+                    NodeCoordinates(i, tuple([*coordinates, (-1, "INPUT", "")])),  # noqa: C409
+                    None,
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+        yield from enumerate_results(
+            proto.node, name, verbose=verbose, coordinates=coordinates
+        )
+        for i in proto.output:
+            if i.name in name:
+                r = ResultFound(
+                    i.name,
+                    None,
+                    NodeCoordinates(
+                        i, tuple([*coordinates, (len(proto.node), "OUTPUT", "")])  # noqa: C409
+                    ),
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+    else:
+        if verbose:
+            print(
+                f"[enumerate_results] {indent}searching for {name!r} into List[NodeProto]..."
+            )
+        for node_i, node in enumerate(proto):
+            if set(node.input) & name:
+                for n in node.input:
+                    if n in name:
+                        r = ResultFound(
+                            n,
+                            NodeCoordinates(
+                                node,
+                                tuple(  # noqa: C409
+                                    [*coordinates, (node_i, node.op_type, node.name)]
+                                ),
+                            ),
+                            None,
+                        )
+                        if verbose > 1:
+                            print(f"[enumerate_results] {indent}-- {r}")
+                        yield r
+            if node.op_type in {"If", "Scan", "Loop", "SequenceMap"}:
+                for att in node.attribute:
+                    if att.type == onnx.AttributeProto.GRAPH:
+                        yield from enumerate_results(
+                            att.g,
+                            name,
+                            verbose=verbose,
+                            coordinates=[*coordinates, (node_i, node.op_type, node.name)],
+                        )
+            if set(node.output) & name:
+                for n in node.output:
+                    if n in name:
+                        r = ResultFound(
+                            n,
+                            None,
+                            NodeCoordinates(
+                                node,
+                                tuple(  # noqa: C409
+                                    [*coordinates, (node_i, node.op_type, node.name)]
+                                ),
+                            ),
+                        )
+                        if verbose > 1:
+                            print(f"[enumerate_results] {indent}-- {r}")
+                        yield r
+    if verbose:
+        print(f"[enumerate_results] {indent}done")
+
+
+def shadowing_names(
+    proto: Union[FunctionProto, GraphProto, ModelProto, Sequence[NodeProto]],
+    verbose: int = 0,
+    existing: Optional[Set[str]] = None,
+    shadow_context: Optional[Set[str]] = None,
+    post_shadow_context: Optional[Set[str]] = None,
+) -> Tuple[Set[str], Set[str], Set[str]]:
+    """
+    Returns the shadowing names, the names created in the main graph
+    after they were created in a subgraphs and the names created by the nodes.
+    """
+    if isinstance(proto, ModelProto):
+        return shadowing_names(proto.graph)
+    if isinstance(proto, GraphProto):
+        assert (
+            existing is None and shadow_context is None
+        ), "existing must be None if nodes is None"
+        return shadowing_names(
+            proto.node,
+            verbose=verbose,
+            existing=set(i.name for i in proto.initializer)
+            | set(i.name for i in proto.sparse_initializer)
+            | set(i.name for i in proto.input if i.name),
+            shadow_context=set(),
+            post_shadow_context=set(),
+        )
+    if isinstance(proto, FunctionProto):
+        assert (
+            existing is None and shadow_context is None
+        ), "existing must be None if nodes is None"
+        return shadowing_names(
+            proto.node,
+            verbose=verbose,
+            existing=set(i for i in proto.input if i),
+            shadow_context=set(),
+            post_shadow_context=set(),
+        )
+
+    assert (
+        existing is not None and shadow_context is not None
+    ), "existing must not be None if nodes is not None"
+    shadow = set()
+    shadow_context = shadow_context.copy()
+    existing = existing.copy()
+    created = set()
+    post_shadow = set()
+    for node in proto:
+        not_empty = set(n for n in node.input if n)
+        intersection = not_empty & existing
+        assert len(intersection) == len(not_empty), (
+            f"One input in {not_empty}, node={pretty_onnx(node)} "
+            f"was not found in {existing}"
+        )
+        for att in node.attribute:
+            if att.type == AttributeProto.GRAPH:
+                g = att.g
+                shadow |= set(i.name for i in g.input) & shadow_context
+                shadow |= set(i.name for i in g.initializer) & shadow_context
+                shadow |= set(i.name for i in g.sparse_initializer) & shadow_context
+                s, ps, c = shadowing_names(
+                    g.node, verbose=verbose, existing=existing, shadow_context=existing
+                )
+                shadow |= s
+                created |= c
+
+        not_empty = set(n for n in node.output if n)
+        post_shadow |= not_empty & created
+        shadow |= not_empty & shadow_context
+        existing |= not_empty
+        created |= not_empty
+    return shadow, post_shadow, created