onnx-diagnostic 0.8.2__py3-none-any.whl → 0.8.4__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, 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
@@ -15,11 +26,14 @@ from onnx import (
     GraphProto,
     ModelProto,
     NodeProto,
+    OperatorSetIdProto,
     TensorProto,
     ValueInfoProto,
     load as onnx_load,
 )
 
+TensorLike = Union[np.ndarray, "torch.Tensor"]  # noqa: F821
+
 
 def _make_stat(init: TensorProto) -> Dict[str, float]:
     """
@@ -331,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
@@ -477,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.
@@ -523,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`.
 
@@ -590,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:
@@ -1195,3 +1209,514 @@ def shadowing_names(
         existing |= not_empty
         created |= not_empty
     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,
+    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), (
+        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
+        }
+    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}
+    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]
+        # 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
+        # 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_inputs):
+            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 -= {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) 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