Microfix

daniil-lyakhov · Nov 17, 2023 · 60f60ff · 60f60ff
1 parent f3da495
commit 60f60ff
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/nncf/openvino/graph/model_utils.py b/nncf/openvino/graph/model_utils.py
@@ -59,4 +59,4 @@ def get_input_nodes(nncf_graph: NNCFGraph) -> List[NNCFNode]:
     :param nncf_graph: NNCFGraph to work with.
     :return: Target NNCFGraph input nodes.
     """
-    return list(set(nncf_graph.get_input_nodes()).upadte(nncf_graph.get_nodes_by_metatypes([OVReadValueMetatype])))
+    return nncf_graph.get_input_nodes() + nncf_graph.get_nodes_by_metatypes([OVReadValueMetatype])
diff --git a/nncf/openvino/graph/nncf_graph_builder.py b/nncf/openvino/graph/nncf_graph_builder.py
@@ -199,8 +199,10 @@ def _set_non_weighted_layer_attributes(node: ov.Node, metatype: OVOpMetatype, nn
         """
         if metatype == OVConcatMetatype:
             nncf_node = nncf_graph.get_node_by_name(node.get_friendly_name())
+            axis = axis = node.get_attributes()["axis"]
+            num_inputs = len(node.inputs())
             nncf_node.layer_attributes = OVLayerAttributes(
-                {}, MultipleInputLayerAttributes(axis=node.get_axis(), num_inputs=len(node.inputs()))
+                {}, MultipleInputLayerAttributes(axis=axis, num_inputs=num_inputs)
             )
 
     def _set_weighted_layer_attributes(

diff --git a/tests/openvino/native/test_layer_attributes.py b/tests/openvino/native/test_layer_attributes.py
@@ -160,13 +160,21 @@ def get_lstm(node_name, input_shape):
     return [input_1], [output]
 
 
-def get_concat_node(op_name, input_shape):
+def get_concat_node(op_name, input_shape, axis):
     num_inputs = 3
     inputs = [opset.parameter(input_shape, name=f"Input{i}") for i in range(num_inputs)]
-    output = opset.concat(inputs, axis=1, name=op_name)
+    output = opset.concat(inputs, axis=axis, name=op_name)
     return inputs, [output]
 
 
+def get_concat_axis_default(op_name, input_shape):
+    return get_concat_node(op_name, input_shape, axis=0)
+
+
+def get_concat_axis_custom(op_name, input_shape):
+    return get_concat_node(op_name, input_shape, axis=1)
+
+
 def get_one_layer_model(op_name: str, node_creator, input_shape):
     inputs, outputs = node_creator(op_name, input_shape)
     results = [opset.result(o, name=f"Result{i}") for i, o in enumerate(outputs)]
@@ -446,7 +454,10 @@ class LayerAttributesTestCase:
         None,
     ),
     LayerAttributesTestCase(
-        get_concat_node, (1, 2, 3, 3), OVLayerAttributes({}, MultipleInputLayerAttributes(1, 3)), None
+        get_concat_axis_default, (1, 2, 3, 3), OVLayerAttributes({}, MultipleInputLayerAttributes(0, 3)), None
+    ),
+    LayerAttributesTestCase(
+        get_concat_axis_custom, (1, 2, 3, 3), OVLayerAttributes({}, MultipleInputLayerAttributes(1, 3)), None
     ),
 ]