fix ConvertGroupedStridedSlice.py for XLNet (#10496)

This commit is contained in:
Pavel Esir authored and GitHub committed 2022-02-18 22:40:28 +03:00
1 parent 6b22d0d109
commit fb6359586d
2 files changed
+43

No files matched your search

@@ -91,6 +91,10 @@ class ConvertGroupedStridedSlice(MiddleReplacementPattern):
sorted_out_nodes = sorted(out_nodes, key=lambda n: list(n.slices))
out_nodes = unique_by(sorted_out_nodes, strided_slices_equality)
# if there is only one StridedSlice out_node with unique 'slices',
# there is nothing to optimize, continue to the next data node
if len(out_nodes) <= 1:
continue
for node in out_nodes:
if len(node.slices) != len(out_nodes[0].slices):
@@ -806,6 +806,45 @@ class ConvertGroupedStridedSliceTests(unittest.TestCase):
(flag, resp) = compare_graphs(graph, graph_ref, 'concat_1_data', check_op_attrs=True)
self.assertTrue(flag, resp)
# one unuque StridedSlice
def test_12(self):
graph = build_graph(nodes_attributes,
[('placeholder_1', 'placeholder_1_data'),
('placeholder_1_data', 'sslice_1'),
('sslice_1', 'sslice_1_data'),
('placeholder_1_data', 'sslice_2'),
('sslice_2', 'sslice_2_data'),
],
{'placeholder_1_data': {'shape': np.array([1, 511])},
'sslice_1': {'slices': np.array([slice(0, 1, 1), slice(0, 1, 1)]),
'begin_mask': np.array([0, 1, 0]),
'end_mask': np.array([0, 1, 0]),
'new_axis_mask': np.array([0, 0, 0]),
'shrink_axis_mask': np.array([0, 0, 0]),
'ellipsis_mask': np.array([0, 0, 0])},
'sslice_1_data': {'shape': np.array([1, 1, 511])},
'sslice_2': {'slices': np.array([slice(0, 1, 1), slice(0, 1, 1)]),
'begin_mask': np.array([0, 1, 0]),
'end_mask': np.array([0, 1, 0]),
'new_axis_mask': np.array([0, 0, 0]),
'shrink_axis_mask': np.array([0, 0, 0]),
'ellipsis_mask': np.array([0, 0, 0])},
'sslice_2_data': {'shape': np.array([1, 1, 511])},
})
graph.graph['layout'] = 'NHWC'
graph_ref = graph.copy()
pattern = ConvertGroupedStridedSlice()
pattern.find_and_replace_pattern(graph)
(flag, resp) = compare_graphs(graph, graph_ref, 'sslice_1_data', check_op_attrs=True)
self.assertTrue(flag, resp)
(flag, resp) = compare_graphs(graph, graph_ref, 'sslice_2_data', check_op_attrs=True)
self.assertTrue(flag, resp)
class AddReshapeAfterStridedSliceTests(unittest.TestCase):
def test_ss_1_shrink_last(self):