Fix of RollWithEmptyAxesReplacer transformation. (#5937)

* Fixed Roll output shape bug in RollWithEmptyAxesReplacer transformations.

* Added check in unit test.
This commit is contained in:
Anastasia Popova authored and GitHub committed 2021-06-03 14:45:05 +03:00
1 parent 9384cfefc6
commit 3825e90ae1
2 files changed
+6 -1

No files matched your search

@@ -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))
@@ -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')