add shadowing

Author: xadupre

_unittests/ut_helpers/test_onnx_helper.py

@@ -17,6 +17,7 @@
     from_array_extended,
     tensor_statistics,
     enumerate_results,
+    shadowing_names,
 )
 
 
@@ -400,7 +401,64 @@ def test_enumerate_results_loop(self):
             )
         )
         res = list(enumerate_results(model, "slice_start", verbose=2))
-        print(res)
+        self.assertEqual(len(res), 2)
+
+    def test_shadowing_names(self):
+        def _mkv_(name):
+            value_info_proto = ValueInfoProto()
+            value_info_proto.name = name
+            return value_info_proto
+
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("ReduceSum", ["X"], ["Xred"]),
+                    oh.make_node("Add", ["X", "two"], ["X0"]),
+                    oh.make_node("Add", ["X0", "zero"], ["X00"]),
+                    oh.make_node("CastLike", ["one", "Xred"], ["one_c"]),
+                    oh.make_node("Greater", ["Xred", "one_c"], ["cond"]),
+                    oh.make_node("Identity", ["two"], ["three"]),
+                    oh.make_node(
+                        "If",
+                        ["cond"],
+                        ["Z_c"],
+                        then_branch=oh.make_graph(
+                            [
+                                # shadowing
+                                oh.make_node("Constant", [], ["three"], value_floats=[2.1]),
+                                oh.make_node("Add", ["X00", "three"], ["Y"]),
+                            ],
+                            "then",
+                            [],
+                            [_mkv_("Y")],
+                        ),
+                        else_branch=oh.make_graph(
+                            [
+                                # not shadowing
+                                oh.make_node("Sub", ["X0", "three"], ["Y"]),
+                            ],
+                            "else",
+                            [],
+                            [_mkv_("Y")],
+                        ),
+                    ),
+                    oh.make_node("CastLike", ["Z_c", "X"], ["Z"]),
+                ],
+                "test",
+                [
+                    oh.make_tensor_value_info("X", TensorProto.FLOAT, ["N"]),
+                    oh.make_tensor_value_info("one", TensorProto.FLOAT, ["N"]),
+                ],
+                [oh.make_tensor_value_info("Z", TensorProto.UNDEFINED, ["N"])],
+                [
+                    onh.from_array(np.array([0], dtype=np.float32), name="zero"),
+                    onh.from_array(np.array([2], dtype=np.float32), name="two"),
+                ],
+            ),
+            opset_imports=[oh.make_operatorsetid("", 18)],
+            ir_version=10,
+        )
+        self.assertEqual({"three"}, shadowing_names(model))
 
 
 if __name__ == "__main__":

onnx_diagnostic/helpers/onnx_helper.py

@@ -1119,3 +1119,64 @@ def enumerate_results(
                         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,
+) -> Set[str]:
+    """Returns the shadowing names."""
+    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(),
+        )
+    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(),
+        )
+
+    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()
+    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
+                shadow |= shadowing_names(
+                    g.node, verbose=verbose, existing=existing, shadow_context=existing
+                )
+
+        not_empty = set(n for n in node.output if n)
+        shadow |= not_empty & shadow_context
+        existing |= not_empty
+    return shadow