From 3825e90ae1dc3724c28b3ff02f0b237ad97add6e Mon Sep 17 00:00:00 2001 From: Anastasia Popova Date: Thu, 3 Jun 2021 14:45:05 +0300 Subject: [PATCH] Fix of RollWithEmptyAxesReplacer transformation. (#5937) * Fixed Roll output shape bug in RollWithEmptyAxesReplacer transformations. * Added check in unit test. --- .../extensions/front/RollWithEmptyAxesReplacer.py | 2 +- .../extensions/front/RollWithEmptyAxesReplacer_test.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/model-optimizer/extensions/front/RollWithEmptyAxesReplacer.py b/model-optimizer/extensions/front/RollWithEmptyAxesReplacer.py index 236a9328958..1ebdb434252 100644 --- a/model-optimizer/extensions/front/RollWithEmptyAxesReplacer.py +++ b/model-optimizer/extensions/front/RollWithEmptyAxesReplacer.py @@ -41,7 +41,7 @@ class RollWithEmptyAxesReplacer(FrontReplacementPattern): # reshape to original shape shape_of = Shape(graph, {'name': node_name + '/shape_of'}).create_node() - roll_node.in_port(0).get_connection().add_destination(shape_of.in_port(0)) + reshape_to_1d.in_port(0).get_connection().add_destination(shape_of.in_port(0)) reshape_to_orig_shape = Reshape(graph, {}).create_node() rename_nodes([(roll_node, node_name + '/roll'), (reshape_to_orig_shape, node_name)]) shape_of.out_port(0).connect(reshape_to_orig_shape.in_port(1)) diff --git a/model-optimizer/unit_tests/extensions/front/RollWithEmptyAxesReplacer_test.py b/model-optimizer/unit_tests/extensions/front/RollWithEmptyAxesReplacer_test.py index e172737f15e..8e60ad3ac4b 100644 --- a/model-optimizer/unit_tests/extensions/front/RollWithEmptyAxesReplacer_test.py +++ b/model-optimizer/unit_tests/extensions/front/RollWithEmptyAxesReplacer_test.py @@ -50,3 +50,8 @@ class RollWithEmptyAxesReplacerTest(unittest.TestCase): (flag, resp) = compare_graphs(graph, graph_ref, 'result', check_op_attrs=True) self.assertTrue(flag, resp) + + shape_of_nodes = graph.get_op_nodes(type='ShapeOf') + self.assertTrue(len(shape_of_nodes) == 1) + shape_of = shape_of_nodes[0] + self.assertTrue(shape_of.in_node().soft_get('name') == 'placeholder')