Fix NNCFGrpah _get_edge

nncf/common/graph/graph.py

@@ -148,18 +148,18 @@ def __init__(
         self.parallel_input_port_ids = parallel_input_port_ids
 
     def __str__(self):
-        # return f"{self.from_node}: -> {self.tensor_shape} -> {self.to_node}:"
         return f"{self.from_node}:{self.output_port_id} -> {self.tensor_shape} -> {self.to_node}:{self.input_port_id}"
 
     def __hash__(self):
-        # return hash(str(self.from_node) + " -> " + str(self.tensor_shape) + " -> " + str(self.to_node))
         return hash(str(self))
 
     def __eq__(self, other):
         return (
             self.from_node == other.from_node
             and self.to_node == other.to_node
             and self.tensor_shape == other.tensor_shape
+            and self.input_port_id == other.input_port_id
+            and self.output_port_id == other.output_port_id
         )
 
 
@@ -355,9 +355,10 @@ def get_output_edges(self, node: NNCFNode) -> List[NNCFGraphEdge]:
     def _get_edges(self, from_node: NNCFNode, to_node: NNCFNode) -> List[NNCFGraphEdge]:
         edges = []
         edge = self.get_edge(from_node, to_node)
-        if not edge.parallel_input_port_ids:
-            edges.append(edge)
-        for input_port_id in edge.parallel_input_port_ids:
+        parallel_input_port_ids = edge.parallel_input_port_ids
+        edge.parallel_input_port_ids = []
+        edges.append(edge)
+        for input_port_id in parallel_input_port_ids:
             edges.append(
                 NNCFGraphEdge(
                     from_node=edge.from_node,

nncf/experimental/common/tensor_statistics/collectors.py

@@ -366,7 +366,7 @@ def get_tensor_collector_inputs(
 
         :param outputs: Target model outputs.
         :param output_info: Output info collected by a `TensorCollector.get_output_info` method.
-        :returns: Model outputs in a format required by `TensorCollector.register_input` method.
+        :returns: Model outputs in a format required by `TensorCollector.register_inputs` method.
         """
         target_inputs = {}
         for reducer, names in output_info:

nncf/openvino/graph/nncf_graph_builder.py

@@ -101,7 +101,7 @@ def _add_edges_to_nncf_graph(model: ov.Model, graph: NNCFGraph) -> None:
 
                     parallel_inputs = None
                     if len(inputs) > 1:
-                        parallel_inputs = [inp.get_index() for inp in inputs]
+                        parallel_inputs = [inp.get_index() for inp in inputs[1:]]
 
                     graph.add_edge_between_nncf_nodes(
                         from_node_id=in_node_id,

tests/common/graph/test_nncf_graph.py

@@ -62,7 +62,7 @@ def test_parallel_edges():
         nodes[1].node_id,
         input_port_id=0,
         output_port_id=0,
-        parallel_input_port_ids=list(range(5)),
+        parallel_input_port_ids=list(range(1, 5)),
         tensor_shape=(1, 2, 3),
         dtype="dummy",
     )

tests/torch/test_graph_building.py

@@ -21,6 +21,7 @@
 from nncf.common.graph.definitions import MODEL_INPUT_OP_NAME
 from nncf.common.graph.definitions import MODEL_OUTPUT_OP_NAME
 from nncf.common.graph.definitions import NNCFGraphNodeType
+from nncf.common.graph.graph import NNCFGraphEdge
 from nncf.common.graph.layer_attributes import GetItemLayerAttributes
 from nncf.common.graph.layer_attributes import MultipleInputLayerAttributes
 from nncf.common.graph.layer_attributes import MultipleOutputLayerAttributes
@@ -407,10 +408,7 @@ def test_split_attributes(input_shape):
     graph = graph_builder.build_graph(model)
     chunk_nodes_with_attributes = {
         "ModelForTestWithSplit/chunk_0": {"chunks": 2, "axis": 1},
-        "ModelForTestWithSplit/unbind_0": {"chunks": 2, "axis": 1}
-        # TODO: fix NNCFGraph tracing so valid reference below
-        # will be generated by NNCF
-        #'ModelForTestWithSplit/unbind_0': {'chunks': 20, 'axis': 1}
+        "ModelForTestWithSplit/unbind_0": {"chunks": 20, "axis": 1},
     }
 
     for node in graph.get_all_nodes():
@@ -460,6 +458,70 @@ def test_getitem_attributes(input_shape):
                 assert getitem_nodes_with_attributes[node.node_name] is None
 
 
+class ParallelEdgesModel(nn.Module):
+    def forward(self, x):
+        mm_res = torch.mm(x, x)
+        return mm_res, x + mm_res
+
+
+def test_parallel_edges_in_nncf_graph():
+    input_shape = (3, 3)
+    model = ParallelEdgesModel()
+    input_info = ModelInputInfo(input_shape)
+    graph_builder = GraphBuilder(
+        create_dummy_forward_fn(
+            [
+                input_info,
+            ],
+            with_input_tracing=True,
+            with_output_tracing=True,
+        )
+    )
+
+    nncf_graph = graph_builder.build_graph(model)
+
+    input_node = nncf_graph.get_node_by_name("/nncf_model_input_0")
+    mm_node = nncf_graph.get_node_by_name("ParallelEdgesModel/mm_0")
+    ref_input_edges = {
+        NNCFGraphEdge(
+            input_node,
+            mm_node,
+            input_port_id=0,
+            output_port_id=0,
+            tensor_shape=(3, 3),
+            dtype=Dtype.FLOAT,
+            parallel_input_port_ids=[],
+        ),
+        NNCFGraphEdge(
+            input_node,
+            mm_node,
+            input_port_id=1,
+            output_port_id=0,
+            tensor_shape=(3, 3),
+            dtype=Dtype.FLOAT,
+            parallel_input_port_ids=[],
+        ),
+    }
+    mm_node_input_edges = nncf_graph.get_input_edges(mm_node)
+    assert set(mm_node_input_edges) == ref_input_edges
+    ref_output_edges = ref_input_edges.copy()
+
+    add_node = nncf_graph.get_node_by_name("ParallelEdgesModel/__add___0")
+    ref_output_edges.add(
+        NNCFGraphEdge(
+            input_node,
+            add_node,
+            input_port_id=0,
+            output_port_id=0,
+            tensor_shape=(3, 3),
+            dtype=Dtype.FLOAT,
+            parallel_input_port_ids=[],
+        )
+    )
+    input_node_output_edges = nncf_graph.get_output_edges(input_node)
+    assert set(input_node_output_edges) == ref_output_edges
+
+
 TEST_KEYWORD_1 = "keyword1"
 TEST_KEYWORD_2 = "keyword2"
 INPUT_INFO_CONFIG_VS_FORWARD_ARGS = [