LSTMCell tensor names fix (#5696)

* Added handling of debug information in create_node().

* Code refactoring.

* Checks fixed.

* Added comments, added unit test.

* Renamed unit test class.

* Fixed port number in unit test.

model-optimizer/mo/graph/port.py

@@ -101,7 +101,7 @@ class Port:
             else:
                 data_node = self.node.out_node(self.idx, control_flow=self.control_flow)
                 assert data_node.value is None or \
-                       np.array_equal(data_node.soft_get('force_shape', data_node.shape),  int64_array(shape))
+                       np.array_equal(data_node.soft_get('force_shape', data_node.shape), int64_array(shape))
                 self.node.out_node(self.idx, control_flow=self.control_flow).shape = int64_array(shape)
 
     def _get_value(self):
@@ -263,25 +263,37 @@ class Port:
         return consumer_ports
 
     def get_tensor_names(self, port_renumber: bool = False):
-        def get_tensor_names_list(attrs):
-            tensor_names_list = []
+        """
+        Gets sorted tensor names list.
+        :param port_renumber: defines whether data node index should be calculated considering port renumbering.
+        """
+        tensor_debug_info = self.get_tensor_debug_info(port_renumber)
+        tensor_names_list = []
+        for attr in tensor_debug_info:
+            if attr is not None and len(attr) >= 2:
+                tensor_name = attr[1]
+                if tensor_name is not None and len(tensor_name) > 0:
+                    tensor_names_list.append(tensor_name.replace(',', '\\,'))
+        return sorted(tensor_names_list)
+
+    def get_tensor_debug_info(self, port_renumber: bool = False):
+        """
+        Gets tensor debug info attribute.
+        :param port_renumber: defines whether data node index should be calculated considering port renumbering.
+        """
+        def get_tensor_debug_info_from_attrs(attrs):
             if 'fw_tensor_debug_info' in attrs:
-                if attrs['fw_tensor_debug_info'] is None:
-                    return tensor_names_list
-                for attr in attrs['fw_tensor_debug_info']:
-                    if attr is not None and len(attr) >= 2:
-                        tensor_name = attr[1]
-                        if tensor_name is not None and len(tensor_name) > 0:
-                            tensor_names_list.append(tensor_name.replace(',', '\\,'))
-            return tensor_names_list
+                if attrs['fw_tensor_debug_info'] is not None:
+                    return attrs['fw_tensor_debug_info']
+            return []
 
-        assert self.type != 'in', "Can't get tensor names for input port at {} node".format(self.node.name)
+        assert self.type != 'in', "Can't get tensor debug info for input port at {} node".format(self.node.name)
 
-        fw_names = []
+        fw_debug_info = []
         if self.node.graph.stage == 'front':
             if self.idx in self.node.out_edges():
                 out_edge = self.node.out_edge(self.idx)
-                fw_names += get_tensor_names_list(out_edge)
+                fw_debug_info += get_tensor_debug_info_from_attrs(out_edge)
         else:
             # before port renumbering we use sequential numbering
             node_idx = self.idx
@@ -293,8 +305,9 @@ class Port:
 
             if node_idx in self.node.out_nodes():
                 out_node = self.node.out_node(node_idx)
-                fw_names += get_tensor_names_list(out_node.attrs())
-        return sorted(fw_names)
+                fw_debug_info += get_tensor_debug_info_from_attrs(out_node.attrs())
+        return fw_debug_info
+
 
     def disconnect(self):
         if self.type == 'out':

model-optimizer/mo/ops/op.py

@@ -121,13 +121,20 @@ class Op(object):
         if attrs is None:
             attrs = dict()
         new_node = self.add_node(attrs)
-        # Missed careful handling of debug information
         for i, inp in enumerate(inputs):
             edge_attr = {'in': i, 'out': inp[1],
                          'in_attrs': ['in', 'permutation'],
                          'out_attrs': ['out', 'permutation'],
                          'data_attrs': []} if not inp[0].has_valid('kind') or inp[0].kind == 'op' \
                 else {'in': i, 'in_attrs': ['in', 'permutation']}
+
+            # handling of debug information
+            if inp[0].has_port('out', inp[1]):
+                debug_info = inp[0].out_port(inp[1]).get_tensor_debug_info()
+                if debug_info is not None and len(debug_info) > 0:
+                    edge_attr.update({'fw_tensor_debug_info': debug_info})
+                    edge_attr['data_attrs'].append('fw_tensor_debug_info')
+
             if edge_attrs is not None:
                 edge_attr.update(edge_attrs)
             new_node.add_input_port(i, skip_if_exist=True)

model-optimizer/unit_tests/mo/ops/op_test.py

@@ -0,0 +1,30 @@
+# Copyright (C) 2018-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+from extensions.ops.lstm_cell import LSTMCell
+from mo.graph.graph import Node
+from unit_tests.utils.graph import build_graph, regular_op
+
+nodes = {
+    **regular_op('Op1', {'type': 'Op1', 'kind': 'op', 'op': 'Op1'}),
+    **regular_op('Op2', {'type': 'Op2', 'kind': 'op', 'op': 'Op2'}),
+    **regular_op('Op3', {'type': 'Op3', 'kind': 'op', 'op': 'Op3'}),
+}
+
+
+class TestOp(unittest.TestCase):
+    def test_create_node(self):
+        graph = build_graph(nodes, [('Op1', 'Op3', {'in': 0, 'out': 0, 'fw_tensor_debug_info': [('Op1', 'Op1')]}),
+                                    ('Op2', 'Op3', {'in': 1, 'out': 0, 'fw_tensor_debug_info': [('Op2', 'Op2')]})])
+        graph.stage = 'front'
+        input1 = Node(graph, 'Op1')
+        input2 = Node(graph, 'Op2')
+        inputs = [(input1, 0), (input2, 0)]
+
+        lstm_op = LSTMCell(graph, dict(name='LSTMCell'))
+        _ = lstm_op.create_node(inputs)
+
+        self.assertTrue(input1.out_edge(0)['fw_tensor_debug_info'] == [('Op1', 'Op1')])
+        self.assertTrue(input2.out_edge(0)['fw_tensor_debug_info'] == [('Op2', 'Op2')])