fix other series of bugs

Author: xadupre

_unittests/ut_torch_onnx/data/attention_loopa24.onnx

@@ -12,7 +12,7 @@ def test_attention_opset15_in_a_loop(self):
         )
         sess = self.check_ort(model)
         feeds = dict(
-            c_lift_tensor_0=np.array([0], dtype=np.int64),
+            c_lifted_tensor_0=np.array([0], dtype=np.int64),
             cat_2=np.array(
                 [
                     0,
@@ -45,9 +45,12 @@ def test_attention_opset15_in_a_loop(self):
                 dtype=np.int64,
             ),
             unsqueeze_4=np.random.randn(1, 16, 1292, 80).astype(np.float32),
+            unsqueeze_5=np.random.randn(1, 16, 1292, 80).astype(np.float32),
+            unsqueeze_6=np.random.randn(1, 16, 1292, 80).astype(np.float32),
         )
         got = sess.run(None, feeds)
         self.assertEqual(len(got), 1)
+        self.assertEqual((1, 1292, 16, 80), got[0].shape)
 
 
 if __name__ == "__main__":

_unittests/ut_torch_onnx/test_discrepancies.py

@@ -1,29 +1,9 @@
-from typing import Dict, Set
+from typing import Dict
 import numpy as np
 import onnx
 import onnx.numpy_helper as onh
 from ..reference import ExtendedReferenceEvaluator as Inference
-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 .onnx_helper import onnx_dtype_name, pretty_onnx, get_hidden_inputs
 
 
 def _make_node_label(node: onnx.NodeProto, tiny_inits: Dict[str, str]) -> str:
@@ -221,7 +201,7 @@ def _mkn(obj: object) -> int:
             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/dot_helper.py

@@ -1198,6 +1198,30 @@ 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.
+    """
+    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 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 extract_subset_of_nodes(
     model: ModelProto,
     name: str,
@@ -1240,30 +1264,45 @@ def extract_subset_of_nodes(
     current_input_index = 0
     intermediate = {name}
     cut_points -= {name}
+    cached = {}
     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 or not node.input:
+        # node inputs including hidden ones
+        if current_node_index in cached:
+            node_inputs = cached[current_node_index]
+        else:
+            node_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:
+                        node_inputs |= get_hidden_inputs(att.g)
+            node_inputs = list(node_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.input:
+                if not node_inputs:
                     current_node_index -= 1
                     current_input_index = 0
                     continue
             else:
                 current_node_index -= 1
                 current_input_index = 0
                 continue
-        assert current_input_index < len(node.input), (
-            f"current_input_index={current_input_index} but node.input={node.input}, "
+        # 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.input[current_input_index]
+        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
 
@@ -1296,8 +1335,14 @@ def _mkv_(name, itype, irank):
 
     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:
+                    print("++++", get_hidden_inputs(att.g))
+                    not_known |= get_hidden_inputs(att.g)
 
     model = oh.make_model(
         oh.make_graph(

onnx_diagnostic/helpers/onnx_helper.py

@@ -18,10 +18,11 @@
 import onnxruntime
 from ..helpers import string_type
 from ..helpers.onnx_helper import (
-    pretty_onnx,
+    get_hidden_inputs,
     dtype_to_tensor_dtype,
-    to_array_extended,
     np_dtype_to_tensor_dtype,
+    to_array_extended,
+    pretty_onnx,
 )
 from ..helpers.torch_helper import onnx_dtype_to_torch_dtype, torch_dtype_to_onnx_dtype
 from ..helpers.ort_session import (
@@ -472,39 +473,15 @@ def enumerate_nodes(self, nodes: List[NodeProto]) -> Iterator[NodeProto]:
                         yield from self.enumerate_nodes(att.g.node)
             yield node
 
-    @classmethod
-    def _get_hidden_inputs(cls, graph: GraphProto) -> Set[str]:
-        """
-        Returns the hidden inputs (inputs coming from an upper context)
-        used by a subgraph.
-        """
-        hidden = set()
-        memo = (
-            {i.name for i in graph.initializer}
-            | {i.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 == AttributeProto.GRAPH and att.g:
-                    hid = cls._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
-
     @classmethod
     def _get_hidden_node_inputs(cls, node: NodeProto) -> Set[str]:
-        """Calls multiple _get_hidden_inputs on every attribute."""
+        """Calls multiple get_hidden_inputs on every attribute."""
         if node.op_type not in {"Loop", "Scan", "If"}:
             return set()
         hidden = set()
         for att in node.attribute:
             if att.type == AttributeProto.GRAPH:
-                hidden |= cls._get_hidden_inputs(att.g)
+                hidden |= get_hidden_inputs(att.g)
         return hidden - (hidden & set(node.input))
 
     def _get_sess(
@@ -624,7 +601,7 @@ def _get_sess_init_subgraph(
                 value = oh.make_tensor_value_info(i, dtype_to_tensor_dtype(it.dtype), it.shape)
             vinputs.append(value)
 
-        reduced_set = self._get_hidden_inputs(g)
+        reduced_set = get_hidden_inputs(g)
         for i, v in context.items():
             if i in reduced_set and i not in unique_names:
                 unique_names.add(i)

onnx_diagnostic/reference/ort_evaluator.py

@@ -4,6 +4,7 @@
 import torch
 from ..api import TensorLike
 from ..helpers import string_type
+from ..helpers.onnx_helper import get_hidden_inputs
 
 
 class RuntimeValueKind(enum.IntEnum):
@@ -151,30 +152,6 @@ def is_initializer(self) -> bool:
         return self.kind == RuntimeValueKind.INITIALIZER
 
 
-def get_hidden_inputs(graph: onnx.GraphProto) -> Set[str]:
-    """
-    Returns the hidden inputs (inputs coming from an upper context)
-    used by a subgraph.
-    """
-    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 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 set_is_shape(
     node: onnx.NodeProto, values: Dict[str, RuntimeValue], drop: Optional[Set[str]] = None
 ) -> List[str]: