def pattern():
return dict(nodes=[('transpose_const',
dict(kind='op',
type='Const',
value=lambda v: v is not None and np.
array_equal(v, int64_array([0, 2, 3, 1])))),
('transpose_const_data', dict(kind='data')),
('transpose', dict(kind='op', type='Transpose')),
('transpose_data', dict(kind='data')),
('reduce_const',
dict(kind='op',
type='Const',
value=lambda v: v is not None and np.
array_equal(v, int64_array([1, 2])))),
('reduce_const_data', dict(kind='data')),
('reduce',
dict(kind='op',
type=lambda t: t in reduce_map.keys(),
keep_dims=False))],
edges=[('transpose_const', 'transpose_const_data'),
('transpose_const_data', 'transpose', {
'in': 1
}), ('transpose', 'transpose_data'),
('transpose_data', 'reduce', {
'in': 0
}), ('reduce_const', 'reduce_const_data'),
('reduce_const_data', 'reduce', {
'in': 1
})])
def find_and_replace_pattern(self, graph: Graph):
rsorted_nodes = graph.pseudo_topological_sort(reverse=True)
for reduce_type in reduce_map.keys():
reduces_of_type = [
n for n in rsorted_nodes
if n.id in graph and n.soft_get('type') == reduce_type
]
for second_reduce_node in reduces_of_type:
if second_reduce_node.id not in graph:
continue
first_reduce_node = second_reduce_node.in_port(
0).get_source().node
if first_reduce_node.soft_get('type', None) == reduce_type:
ReduceMerge.fuse_reduces(first_reduce=first_reduce_node,
second_reduce=second_reduce_node)