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

onnx_diagnostic/helpers/onnx_helper.py

@@ -3,9 +3,20 @@ import json
 import os
 import sys
 import warnings
-from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    Union,
+)
 import numpy as np
-import numpy.typing as npt
 import onnx
 import onnx.helper as oh
 import onnx.numpy_helper as onh
@@ -21,6 +32,8 @@ from onnx import (
     load as onnx_load,
 )
 
+TensorLike = Union[np.ndarray, "torch.Tensor"]  # noqa: F821
+
 
 def _make_stat(init: TensorProto) -> Dict[str, float]:
     """
@@ -332,7 +345,7 @@ def onnx_dtype_name(itype: int, exc: bool = True) -> str:
         print(onnx_dtype_name(7))
     """
     for k in dir(TensorProto):
-        if k.upper() == k and k != "EXTERNAL":
+        if k.upper() == k and k not in {"DESCRIPTOR", "EXTERNAL", "DEFAULT"}:
             v = getattr(TensorProto, k)
             if v == itype:
                 return k
@@ -478,7 +491,7 @@ def convert_endian(tensor: TensorProto) -> None:
     tensor.raw_data = np.frombuffer(tensor.raw_data, dtype=np_dtype).byteswap().tobytes()
 
 
-def from_array_ml_dtypes(arr: npt.ArrayLike, name: Optional[str] = None) -> TensorProto:
+def from_array_ml_dtypes(arr: TensorLike, name: Optional[str] = None) -> TensorProto:
     """
     Converts a numpy array to a tensor def assuming the dtype
     is defined in ml_dtypes.
@@ -524,7 +537,7 @@ _STORAGE_TYPE = {
 }
 
 
-def from_array_extended(tensor: npt.ArrayLike, name: Optional[str] = None) -> TensorProto:
+def from_array_extended(tensor: TensorLike, name: Optional[str] = None) -> TensorProto:
     """
     Converts an array into a :class:`onnx.TensorProto`.
 
@@ -591,7 +604,7 @@ def from_array_extended(tensor: npt.ArrayLike, name: Optional[str] = None) -> Te
     return t
 
 
-def to_array_extended(proto: TensorProto) -> npt.ArrayLike:
+def to_array_extended(proto: TensorProto) -> TensorLike:
     """Converts :class:`onnx.TensorProto` into a numpy array."""
     arr = onh.to_array(proto)
     if proto.data_type >= onnx.TensorProto.BFLOAT16:
@@ -1198,6 +1211,43 @@ def shadowing_names(
     return shadow, post_shadow, created
 
 
+def get_hidden_inputs(graph: onnx.GraphProto) -> Set[str]:
+    """
+    Returns the hidden inputs (inputs coming from an upper context)
+    used by a subgraph. It excludes empty names.
+    """
+    hidden = set()
+    memo = (
+        set(i.name for i in graph.initializer)
+        | set(i.name for i in graph.sparse_initializer)
+        | set(i.name for i in graph.input)
+    )
+    for node in graph.node:
+        for i in node.input:
+            if i and 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 get_all_node_inputs(node: onnx.NodeProto) -> Set[str]:
+    """
+    Returns input and hidden inputs of a node.
+    See :func:`get_hidden_inputs`. It excludes empty names.
+    """
+    start = {i for i in node.input if i}
+    if node.op_type in {"Scan", "Loop", "If"}:
+        for att in node.attribute:
+            if att.type == onnx.AttributeProto.GRAPH:
+                start |= get_hidden_inputs(att.g)
+    return start
+
+
 def extract_subset_of_nodes(
     model: ModelProto,
     name: str,
@@ -1219,11 +1269,14 @@ def extract_subset_of_nodes(
             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}"
+    assert node_index is not None and node_index < len(model.graph.node), (
+        f"node_index={node_index} (n_nodes={len(model.graph.node)}) "
+        f"is still empty or wrong for result {name!r}"
+    )
+    assert name in model.graph.node[node_index].output, (
+        f"Unable to find {name!r} in {model.graph.node[node_index].output}, "
+        f"node={pretty_onnx(model.graph.node[node_index])}"
+    )
     if cut_points is None:
         cut_points = {n.name for n in model.graph.input} | {
             n.name for n in model.graph.initializer
@@ -1236,21 +1289,46 @@ def extract_subset_of_nodes(
     current_node_index = node_index
     current_input_index = 0
     intermediate = {name}
+    cut_points -= {name}
+    cached: Dict[int, List[str]] = {}
     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:
+        # node inputs including hidden ones
+        if current_node_index in cached:
+            node_inputs = cached[current_node_index]
+        else:
+            set_inputs = set(i for i in node.input if i)
+            if node.op_type in {"Scan", "If", "Loop"}:
+                # there are hidden inputs
+                for att in node.attribute:
+                    if att.type == onnx.AttributeProto.GRAPH:
+                        set_inputs |= get_hidden_inputs(att.g)
+            node_inputs = list(set_inputs)
+            cached[current_node_index] = node_inputs
+        # processing
+        if current_input_index == 0 or not node_inputs:
             needs = [o for o in node.output if o in intermediate and o not in cut_points]
             if needs:
                 selected.add(current_node_index)
+                if not node_inputs:
+                    current_node_index -= 1
+                    current_input_index = 0
+                    continue
             else:
                 current_node_index -= 1
+                current_input_index = 0
                 continue
-        res = node.input[current_input_index]
+        # more intermediate results
+        assert current_input_index < len(node_inputs), (
+            f"current_input_index={current_input_index} but node_inputs={node_inputs}, "
+            f"node={pretty_onnx(node)}"
+        )
+        res = node_inputs[current_input_index]
         if res not in cut_points:
             intermediate.add(res)
         current_input_index += 1
-        if current_input_index >= len(node.input):
+        if current_input_index >= len(node_inputs):
             current_node_index -= 1
             current_input_index = 0
 
@@ -1283,17 +1361,362 @@ def make_submodel(
 
     not_known: Set[str] = set()
     for node in nodes[::-1]:
-        not_known -= set(node.output)
-        not_known |= set(node.input)
+        not_known -= {o for o in node.output if o}
+        not_known |= {i for i in node.input if i}
+        if node.op_type in {"Scan", "If", "Loop"}:
+            # there are hidden inputs
+            for att in node.attribute:
+                if att.type == onnx.AttributeProto.GRAPH:
+                    not_known |= get_hidden_inputs(att.g)
 
     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)],
+            [_mkv_(n, *type_rank_fn(n)) for n in sorted(not_known) if n],
+            [_mkv_(n, *type_rank_fn(n)) for n in sorted(output_names) if n],
         ),
         ir_version=ir_version,
         opset_imports=opset_imports,
     )
     return model
+
+
+def get_tensor_shape(
+    obj: Union[onnx.ValueInfoProto, onnx.TypeProto, onnx.TensorProto],
+) -> Optional[List[Optional[Union[int, str]]]]:
+    """
+    Returns the shape if that makes sense for this object.
+    """
+    if isinstance(obj, ValueInfoProto):
+        return get_tensor_shape(obj.type)
+    elif not isinstance(obj, onnx.TypeProto):
+        raise TypeError(f"Unexpected type {type(obj)!r}.")
+    if not obj.tensor_type.HasField("shape"):
+        return None
+    shape = []
+    for d in obj.tensor_type.shape.dim:
+        v = d.dim_value if d.dim_value > 0 else d.dim_param
+        shape.append(v)
+    if not shape:
+        return shape
+    return [None if s in (0, "") else s for s in shape]
+
+
+def _enumerate_model_node_outputs(
+    model: ModelProto, add_node: bool = False, order: bool = False
+) -> Iterable[Union[str, Tuple[str, NodeProto]]]:
+    """
+    Enumerates all the nodes of a model.
+
+    :param model: :epkg:`ONNX` graph
+    :param add_node: if False, the function enumerates
+        all output names from every node, otherwise, it
+        enumerates tuple (output name, node)
+    :param order: goes through outputs following the graph order
+    :return: enumerator
+    """
+    assert hasattr(model, "graph"), "Parameter model is not an ONNX model but {type(model)}"
+    if order:
+        edges = []
+        d_order = {}
+        node_names = {}
+        for inp in model.graph.input:
+            d_order[0, inp.name] = 0
+        for node in model.graph.node:
+            d_order[1, node.name] = 0
+            for i in node.input:
+                edges.append(("in", i, node.name))
+            for o in node.output:
+                edges.append(("out", o, node.name))
+                node_names[o] = node
+                d_order[0, o] = 0
+
+        modif = 1
+        n_iter = 0
+        while modif > 0 and n_iter <= len(model.graph.node):
+            modif = 0
+            n_iter += 1
+            for kind, data_name, node_name in edges:
+                if kind == "in":
+                    if (0, data_name) not in d_order:
+                        continue
+                    if d_order[0, data_name] + 1 > d_order[1, node_name]:
+                        modif += 1
+                        d_order[1, node_name] = d_order[0, data_name] + 1
+                else:
+                    if d_order[1, node_name] + 1 > d_order[0, data_name]:
+                        modif += 1
+                        d_order[0, data_name] = d_order[1, node_name] + 1
+
+        orders = [(v, k) for k, v in d_order.items()]
+        orders.sort()
+
+        for _, k in orders:
+            if k[0] == 1:
+                continue
+            out = k[1]
+            if out not in node_names:
+                continue
+            yield (out, node_names[out]) if add_node else out
+    else:
+        for node in model.graph.node:
+            for out in node.output:
+                yield (out, node) if add_node else out
+
+
+def onnx_remove_node_unused(
+    graph: Union[onnx.GraphProto, onnx.FunctionProto], recursive=True
+) -> Union[onnx.GraphProto, onnx.FunctionProto]:
+    """
+    Removes unused nodes of the graph. An unused node
+    is not involved in the output computation.
+
+    :param onnx_model: onnx model
+    :param recursive: looks into subgraphs
+    :return: new Graph
+    """
+    is_function = isinstance(graph, FunctionProto)
+
+    # mark outputs
+    marked: Dict[str, Set[str]] = (
+        {o: set() for o in graph.output}
+        if is_function
+        else {o.name: set() for o in graph.output}
+    )
+    nodes = list(graph.node)
+
+    # mark node output
+    for node in reversed(nodes):
+        used = False
+        for o in node.output:
+            if o and o in marked:
+                for i in get_all_node_inputs(node):
+                    marked[o].add(i)
+                    used = True
+        if used:
+            for i in get_all_node_inputs(node):
+                marked[i] = set()
+
+    # removed nodes
+    removed = set()
+    marked_set = set(marked)
+    for ind, node in enumerate(nodes):
+        if not ({o for o in node.output if o} & marked_set):
+            removed.add(ind)
+
+    if not is_function:
+        initializers = [i for i in graph.initializer if i.name in marked]
+        sparse_initializers = [i for i in graph.sparse_initializer if i.name in marked]
+    new_nodes = [node for i, node in enumerate(nodes) if i not in removed]
+
+    # Finally create the new graph.
+    if is_function:
+        return oh.make_function(
+            graph.domain,
+            graph.name,
+            graph.input,
+            graph.output,
+            new_nodes,
+            opset_imports=graph.opset_import,
+            attributes=graph.attribute,
+            doc_string=graph.doc_string,
+        )
+    new_graph = oh.make_graph(
+        new_nodes,
+        graph.name,
+        graph.input,
+        graph.output,
+        initializers,
+        sparse_initializers,
+    )
+    new_graph.value_info.extend(graph.value_info)
+    return new_graph
+
+
+def select_model_inputs_outputs(
+    model: ModelProto,
+    outputs: Optional[List[str]] = None,
+    inputs: Optional[List[str]] = None,
+    infer_shapes: bool = True,
+    overwrite: Optional[Dict[str, Any]] = None,
+    remove_unused: bool = True,
+    verbose: int = 0,
+):
+    """
+    Takes a model and changes its outputs.
+
+    :param model: :epkg:`ONNX` model
+    :param inputs: new inputs, same ones if None
+    :param outputs: new outputs, same ones if None
+    :param infer_shapes: infer inputs and outputs shapes
+    :param overwrite: overwrite type and shapes for
+        inputs or outputs, *overwrite* is a
+        dictionary `{'name': (numpy dtype, shape)}`
+    :param remove_unused: remove unused nodes from the graph
+    :param verbose: display information while converting
+    :return: modified model
+
+    The function removes unneeded nodes.
+
+    The following example shows how to change the inputs of model
+    to bypass the first nodes. Shape inferences fails to determine
+    the new inputs type. They need to be overwritten.
+    `verbose=1` shows the number of deleted nodes.
+
+    ::
+
+        import onnx
+        from onnx_extended.tools.onnx_nodes import select_model_inputs_outputs
+
+        onx = onnx.load(path)
+        onx2 = select_model_inputs_outputs(
+            onx, inputs=["a", "b"],
+            infer_shapes=True, verbose=1,
+            overwrite={'a': (numpy.int32, None), 'b': (numpy.int64, None)})
+        onnx.save(onx2, path2)
+    """
+    if not isinstance(model, ModelProto):
+        raise TypeError(f"Unexpected type {type(model)} for model.")
+    if inputs is not None and not isinstance(inputs, list):
+        inputs = [inputs]
+    if outputs is not None and not isinstance(outputs, list):
+        outputs = [outputs]
+    if inputs is None:
+        inputs = [i.name for i in model.graph.input]
+    if outputs is None:
+        outputs = [o.name for o in model.graph.output]
+
+    mark_var = {}
+    for out in _enumerate_model_node_outputs(model):
+        mark_var[out] = 0
+    for inp in inputs:
+        mark_var[inp] = 0
+    for out in outputs:
+        assert out in mark_var, f"Output {out!r} not found in model."
+        mark_var[out] = 1
+
+    nodes = list(model.graph.node[::-1])
+    mark_op = {}
+    for node in list(nodes):
+        mark_op[id(node)] = 0
+
+    # We mark all the nodes we need to keep.
+    nb = 1
+    while nb > 0:
+        nb = 0
+        for node in nodes:
+            if mark_op[id(node)] == 1:
+                continue
+            mod = False
+            for out in node.output:
+                if mark_var[out] == 1:
+                    mark_op[id(node)] = 1
+                    mod = True
+                    break
+            if not mod:
+                continue
+
+            node_inputs = get_all_node_inputs(node)
+
+            nb += 1
+            for inp in node_inputs:
+                if inp in inputs:
+                    continue
+                if mark_var.get(inp, 0) == 1:
+                    continue
+                mark_var[inp] = 1
+                nb += 1
+
+    # All nodes verifies mark_op[node.name] == 1
+    keep_nodes = [node for node in nodes[::-1] if mark_op[id(node)] == 1]
+
+    known_shapes = {}
+    if infer_shapes:
+        shapes = onnx.shape_inference.infer_shapes(model)
+        for shape in shapes.graph.value_info:
+            known_shapes[shape.name] = shape.type
+        for shape in shapes.graph.input:
+            known_shapes[shape.name] = shape.type
+        for shape in shapes.graph.output:
+            known_shapes[shape.name] = shape.type
+    else:
+        for shape in model.graph.input:
+            known_shapes[shape.name] = shape.type
+        for shape in model.graph.output:
+            known_shapes[shape.name] = shape.type
+
+    var_in = []
+    existing = {i.name: i for i in model.graph.input}
+    for name in inputs:
+        if overwrite is not None and name in overwrite:
+            dtype, shape = overwrite[name]
+            proto_dtype = np_dtype_to_tensor_dtype(dtype)
+            value_info = oh.make_tensor_value_info(name, proto_dtype, shape)
+        elif name in known_shapes:
+            info = known_shapes[name].tensor_type
+            proto_dtype = info.elem_type
+            if proto_dtype == 0:
+                value_info = ValueInfoProto()
+                value_info.name = name
+            else:
+                shape = get_tensor_shape(known_shapes[name])
+                value_info = oh.make_tensor_value_info(name, proto_dtype, shape)
+        elif name in existing:
+            value_info = existing[name]
+        else:
+            value_info = ValueInfoProto()
+            value_info.name = name
+        var_in.append(value_info)
+
+    var_out = []
+    existing = {i.name: i for i in model.graph.output}
+    for name in outputs:
+        if overwrite is not None and name in overwrite:
+            dtype, shape = overwrite[name]
+            proto_dtype = np_dtype_to_tensor_dtype(dtype)
+            value_info = oh.make_tensor_value_info(name, proto_dtype, shape)
+        elif name in known_shapes:
+            info = known_shapes[name].tensor_type
+            proto_dtype = info.elem_type
+            if proto_dtype == 0:
+                value_info = ValueInfoProto()
+                value_info.name = name
+            else:
+                shape = get_tensor_shape(known_shapes[name])
+                value_info = oh.make_tensor_value_info(name, proto_dtype, shape)
+        elif name in existing:
+            value_info = existing[name]
+        else:
+            value_info = ValueInfoProto()
+            value_info.name = name
+        var_out.append(value_info)
+
+    graph = oh.make_graph(
+        keep_nodes,
+        model.graph.name,
+        var_in,
+        var_out,
+        model.graph.initializer,
+        sparse_initializer=model.graph.sparse_initializer,
+    )
+    if remove_unused:
+        graph = onnx_remove_node_unused(graph, recursive=False)
+    onnx_model = oh.make_model(graph, functions=model.functions)
+    onnx_model.ir_version = model.ir_version
+    onnx_model.producer_name = model.producer_name
+    onnx_model.producer_version = model.producer_version
+    onnx_model.domain = model.domain
+    onnx_model.model_version = model.model_version
+    onnx_model.doc_string = model.doc_string
+    if model.metadata_props:
+        values = {p.key: p.value for p in model.metadata_props}
+        oh.set_model_props(onnx_model, values)
+
+    del onnx_model.opset_import[:]
+    for oimp in model.opset_import:
+        op_set = onnx_model.opset_import.add()
+        op_set.domain = oimp.domain
+        op_set.version = oimp.version
+
+    return onnx_model

onnx_diagnostic/helpers/ort_session.py

@@ -1,7 +1,6 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import onnx
 import numpy as np
-import numpy.typing as npt
 import torch
 from torch._C import _from_dlpack
 import onnxruntime
@@ -16,6 +15,7 @@ from .torch_helper import torch_dtype_to_onnx_dtype
 
 
 DEVICES = {-1: ORTC.OrtDevice(ORTC.OrtDevice.cpu(), ORTC.OrtDevice.default_memory(), 0)}
+TensorLike = Union[np.ndarray, torch.Tensor]
 
 
 class _InferenceSession:
@@ -243,16 +243,16 @@ class InferenceSessionForNumpy(_InferenceSession):
         )
 
     def run(
-        self, output_names: Optional[List[str]], feeds: Dict[str, npt.ArrayLike]
-    ) -> List[Optional[npt.ArrayLike]]:
+        self, output_names: Optional[List[str]], feeds: Dict[str, TensorLike]
+    ) -> List[Optional[TensorLike]]:
         """Calls :meth:`onnxruntime.InferenceSession.run`."""
         # sess.run does not support blfoat16
         # res = self.sess.run(output_names, feeds)
         return self._post_process_inplace(list(self.run_dlpack(output_names, feeds)))
 
     def run_dlpack(
-        self, output_names: Optional[List[str]], feeds: Dict[str, npt.ArrayLike]
-    ) -> Tuple[Optional[npt.ArrayLike], ...]:
+        self, output_names: Optional[List[str]], feeds: Dict[str, TensorLike]
+    ) -> Tuple[Optional[TensorLike], ...]:
         """
         Same as :meth:`onnxruntime.InferenceSession.run` except that
         feeds is a dictionary of :class:`np.ndarray`.
@@ -289,13 +289,13 @@ class InferenceSessionForNumpy(_InferenceSession):
     def _ortvalues_to_numpy_tensor(
         self,
         ortvalues: Union[List[ORTC.OrtValue], ORTC.OrtValueVector],
-    ) -> Tuple[Optional[npt.ArrayLike], ...]:
+    ) -> Tuple[Optional[TensorLike], ...]:
         if len(ortvalues) == 0:
             return tuple()
 
         if self.nvtx:
             self.torch.cuda.nvtx.range_push("_ortvalues_to_numpy_tensor")
-        res: List[Optional[npt.ArrayLike]] = []  # noqa: F823
+        res: List[Optional[TensorLike]] = []  # noqa: F823
         for i in range(len(ortvalues)):
             if not ortvalues[i].has_value():
                 res.append(None)
@@ -556,7 +556,7 @@ def investigate_onnxruntime_issue(
         Union[str, Callable[[onnx.ModelProto], onnxruntime.InferenceSession]]
     ] = None,
     # if model needs to be run.
-    feeds: Optional[Union[Dict[str, torch.Tensor], Dict[str, npt.ArrayLike]]] = None,
+    feeds: Optional[Union[Dict[str, torch.Tensor], Dict[str, TensorLike]]] = None,
     verbose: int = 0,
     dump_filename: Optional[str] = None,
     infer_shapes: bool = True,

onnx_diagnostic/helpers/torch_helper.py

@@ -139,6 +139,15 @@ def onnx_dtype_to_torch_dtype(itype: int) -> torch.dtype:
     )
 
 
+_TYPENAME = dict(
+    FLOAT=onnx.TensorProto.FLOAT,
+    INT64=onnx.TensorProto.INT64,
+    INT32=onnx.TensorProto.INT32,
+    FLOAT16=onnx.TensorProto.FLOAT16,
+    BFLOAT16=onnx.TensorProto.BFLOAT16,
+)
+
+
 def torch_dtype_to_onnx_dtype(to: torch.dtype) -> int:
     """
     Converts a torch dtype into a onnx element type.
@@ -182,7 +191,13 @@ def torch_dtype_to_onnx_dtype(to: torch.dtype) -> int:
         return onnx.TensorProto.COMPLEX64
     if to == torch.complex128:
         return onnx.TensorProto.COMPLEX128
-    raise NotImplementedError(f"Unable to convert torch dtype {to!r} to onnx dtype.")
+    # SymbolicTensor
+    sto = str(to)
+    if sto in _TYPENAME:
+        return _TYPENAME[sto]
+    raise NotImplementedError(
+        f"Unable to convert torch dtype {to!r} ({type(to)}) to onnx dtype."
+    )
 
 
 def _forward_(
@@ -811,7 +826,8 @@ def torch_deepcopy(value: Any) -> Any:
     if isinstance(value, tuple):
         return tuple(torch_deepcopy(v) for v in value)
     if isinstance(value, list):
-        return [torch_deepcopy(v) for v in value]
+        if type(value) is list:
+            return [torch_deepcopy(v) for v in value]
     if isinstance(value, set):
         return {torch_deepcopy(v) for v in value}
     if isinstance(value, dict):
@@ -1087,3 +1103,13 @@ def study_discrepancies(
     if name:
         fig.savefig(name)
     return ax
+
+
+def int_device_to_torch_device(device_id: int) -> torch.device:
+    """
+    Converts a device defined as an integer (coming from :meth:`torch.Tensor.get_device`)
+    into a ``torch.device``.
+    """
+    if device_id < 0:
+        return torch.device("cpu")
+    return torch.device("cuda", device_id)