Extend MO for operation Einsum-7 (#5401)

* Extend MO for operation Einsum-7 Signed-off-by: Roman Kazantsev <roman.kazantsev@intel.com> * Add extractor for einsum and optimize code based on review feedback Signed-off-by: Roman Kazantsev <roman.kazantsev@intel.com> * Fix the code based on the review: correct code, tests and comments; move insert_transpose Signed-off-by: Roman Kazantsev <roman.kazantsev@intel.com> * Fix LayoutChangeForEinsum transformation condition Signed-off-by: Roman Kazantsev <roman.kazantsev@intel.com> * Update third-party dependencies Signed-off-by: Roman Kazantsev <roman.kazantsev@intel.com>
2021-05-11 21:36:04 +03:00
parent 9db7f849df
commit dc22c177d5
10 changed files with 543 additions and 2 deletions
@@ -0,0 +1,82 @@
+# Copyright (C) 2018-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+import numpy as np
+
+from extensions.back.LayoutChangeForEinsum import LayoutChangeForEinsum
+from mo.front.common.partial_infer.utils import int64_array
+from mo.utils.ir_engine.compare_graphs import compare_graphs
+from unit_tests.utils.graph import build_graph, result, regular_op_with_shaped_data, valued_const_with_data, connect
+
+nodes_attributes = {
+    # Parameter layers
+    **regular_op_with_shaped_data('placeholder_1', None, {'type': 'Parameter', 'op': 'Parameter'}),
+    **regular_op_with_shaped_data('placeholder_2', None, {'type': 'Parameter', 'op': 'Parameter'}),
+    **regular_op_with_shaped_data('placeholder_3', None, {'type': 'Parameter', 'op': 'Parameter'}),
+
+    # Einsum layer
+    **regular_op_with_shaped_data('einsum', None, {'type': 'Einsum', 'op': 'Einsum'}),
+
+    # Result layer
+    **result(),
+
+    # Transpose layers
+    **regular_op_with_shaped_data('transpose_1', None,
+                                  {'type': 'Transpose', 'op': 'Transpose', 'need_shape_inference': True}),
+    **regular_op_with_shaped_data('transpose_3', None,
+                                  {'type': 'Transpose', 'op': 'Transpose', 'need_shape_inference': True}),
+
+    # Const layers
+    **valued_const_with_data('axis_1_const', int64_array([0, 2, 3, 1])),
+    **valued_const_with_data('axis_3_const', int64_array([0, 4, 1, 2, 3])),
+}
+
+
+class LayoutChangeForEinsumTests(unittest.TestCase):
+    def test_layout_change_einsum(self):
+        graph = build_graph(nodes_attributes,
+                            [*connect('placeholder_1', '0:einsum'),
+                             *connect('placeholder_2', '1:einsum'),
+                             *connect('placeholder_3', '2:einsum'),
+                             *connect('einsum', 'output')],
+                            {  # this input stays as is since it is of a rank equal to 3
+                                'placeholder_1_d': {'shape': np.array([2, 3, 5])},
+                                # [3, 5, 7, 8] - NHWC, [3, 8, 5, 7] - NCHW
+                                # this input does not require additional transpose
+                                # since the corresponding subscript can be adjusted
+                                'placeholder_2_d': {'shape': np.array([3, 8, 5, 7])},
+                                # [3, 5, 10, 12] - NHWC, [3, 12, 5, 10] - NCHW
+                                # the third input must be transposed to NHWC layout
+                                # since ellipsis covers multiple dimensions in the end
+                                # the corresponding subscript is not changed
+                                'placeholder_3_d': {'shape': np.array([3, 12, 8, 10])},
+                                # equation is still for NHWC layout
+                                'einsum': {'equation': "abc,bcde,bc...->ade..."},
+                                # [2, 7, 8, 10, 12] - NHWC, [2, 12, 7, 8, 10] - NCHW
+                                # the output is in NCHW layout but its shape will be re-inferred since
+                                # the output stays in NHWC layout due to ellipsis in the end
+                                # and additional transpose to NCHW will be inserted
+                                'einsum_d': {'shape': np.array([2, 12, 7, 8, 10])},
+                            }, nodes_with_edges_only=True)
+        graph.graph['fw'] = 'tf'
+
+        graph_ref = build_graph(nodes_attributes,
+                                [*connect('placeholder_3', '0:transpose_1'),
+                                 *connect('axis_1_const', '1:transpose_1'),
+                                 *connect('placeholder_1', '0:einsum'),
+                                 *connect('placeholder_2', '1:einsum'),
+                                 *connect('transpose_1', '2:einsum'),
+                                 *connect('einsum', '0:transpose_3'),
+                                 *connect('axis_3_const', '1:transpose_3'),
+                                 *connect('transpose_3', 'output')],
+                                {'placeholder_1_d': {'shape': np.array([2, 3, 5])},
+                                 'placeholder_2_d': {'shape': np.array([3, 8, 5, 7])},
+                                 'einsum': {'equation': "abc,becd,bc...->ade..."},
+                                 'einsum_d': {'shape': np.array([2, 12, 7, 8, 10])}
+                                 })
+
+        LayoutChangeForEinsum().find_and_replace_pattern(graph)
+        (flag, resp) = compare_graphs(graph, graph_ref, 'output', check_op_attrs=True)
+        self.assertTrue(flag, resp)
@@ -0,0 +1,90 @@
+# Copyright (C) 2018-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+import numpy as np
+from generator import generator, generate
+
+from extensions.ops.einsum import Einsum
+from mo.front.common.partial_infer.utils import int64_array
+from mo.graph.graph import Graph
+from mo.graph.graph import Node
+from unit_tests.utils.graph import build_graph, regular_op_with_shaped_data, result, connect
+
+
+def create_einsum_graph(input_shapes: list, equation: str) -> Graph:
+    num_inputs = len(input_shapes)
+    assert num_inputs > 0, "Einsum node must have at least one input"
+    nodes = {}
+    edges = []
+    for input_ind in range(num_inputs):
+        input_name = 'input' + str(input_ind)
+        parameter_op = regular_op_with_shaped_data(input_name, input_shapes[input_ind],
+                                                   {'op': 'Parameter', 'type': 'Parameter'})
+        nodes.update(parameter_op)
+        edges += connect(input_name, str(input_ind) + ":einsum_node")
+    einsum_op = regular_op_with_shaped_data('einsum_node', None,
+                                            {'op': 'Einsum', 'type': 'Einsum', 'equation': equation})
+    nodes.update(einsum_op)
+    result_op = result('output')
+    nodes.update(result_op)
+    edges += connect('einsum_node', 'output')
+
+    graph = build_graph(nodes, edges, nodes_with_edges_only=True)
+    return graph
+
+
+@generator
+class TestEinsum(unittest.TestCase):
+    @generate(*[
+        # dot product
+        ([int64_array([10]), int64_array([10])], "i,i->", int64_array([])),
+        # matrix multiplication
+        ([int64_array([2, 3]), int64_array([3, 4])], "ab,bc->ac", int64_array([2, 4])),
+        # trace per batch
+        ([int64_array([2, 3, 3])], "kii->k", int64_array([2])),
+        # diagonal extraction
+        ([int64_array([6, 5, 5])], "kii->ki", int64_array([6, 5])),
+        # transpose
+        ([int64_array([1, 2, 3])], "ijk->kij", int64_array([3, 1, 2])),
+        # multiple matrix multiplication
+        ([int64_array([2, 5]), int64_array([5, 3, 6]), int64_array([5, 3])], "ab,bcd,bc->ca", int64_array([3, 2])),
+        # ellipsis for one operand
+        ([int64_array([5, 3, 4])], "a...->...", int64_array([3, 4])),
+        # ellipsis for multiple operands
+        ([int64_array([3, 5]), int64_array([1])], "a...,...->a...", int64_array([3, 5])),
+        # ellipsis with broadcasting
+        ([int64_array([9, 1, 4, 3]), int64_array([3, 11, 7, 1])], "a...b,b...->a...", int64_array([9, 11, 7, 4])),
+        # mixed case letters in equation
+        ([int64_array([1, 3, 5])], "AbC", int64_array([1, 5, 3])),
+        # mixed case letters and equation in implicit mode
+        ([int64_array([3, 11, 1, 5]), int64_array([1, 3, 1, 7])], "a...b,B...", int64_array([3, 11, 7, 1, 3, 5])),
+    ])
+    def test_einsum(self, input_shapes, equation, ref_output_shape):
+        graph = create_einsum_graph(input_shapes, equation)
+        einsum_node = Node(graph, 'einsum_node')
+        Einsum.infer(einsum_node)
+
+        # get the result
+        res_output_shape = graph.node['einsum_node_d']['shape']
+
+        self.assertTrue(np.array_equal(ref_output_shape, res_output_shape),
+                        'shape does not match expected: {} and given: {}'.format(ref_output_shape, res_output_shape))
+
+    @generate(*[
+        # incorrect subscript numbers or inputs
+        ([int64_array([3, 11]), int64_array([11, 4])], "ab,bc,cd->ac", None),
+        # invalid labels
+        ([int64_array([3, 11]), int64_array([11, 4])], "a$,Bc->ac", None),
+        # incompatible shapes
+        ([int64_array([3, 11]), int64_array([12, 4])], "ab,bc->ac", None),
+        # not broadcastable shapes
+        ([int64_array([11, 1, 4, 3]), int64_array([3, 11, 7, 5])], "a...b,b...->a...", None),
+        # missed ellipsis
+        ([int64_array([11, 1, 4, 3]), int64_array([3, 11, 7, 4])], "a...b,b...->a", None),
+    ])
+    def test_invalid_cases(self, input_shapes, equation, ref_output_shape):
+        graph = create_einsum_graph(input_shapes, equation)
+        einsum_node = Node(graph, 'einsum_node')
+        self.assertRaises(AssertionError, Einsum.infer, einsum_node)
@@ -14,7 +14,7 @@ nodes_attributes = {'data': {'kind': 'op'},
                    'data_data': {'shape': None, 'value': None, 'kind': 'data'},
                    'indices': {'kind': 'op'},
                    'indices_data': {'shape': None, 'value': None, 'kind': 'data'},
-                    'gathernd_node': {'op': 'ScatterNDUpdate', 'kind': 'op', 'batch_dims': 0},
+                    'gathernd_node': {'op': 'GatherNDUpdate', 'kind': 'op', 'batch_dims': 0},
                    'output': {'shape': None, 'value': None, 'kind': 'data'}}

 # graph 1
@@ -118,7 +118,7 @@ inputs_inv2 = {'data_data': {'shape': int64_array([10, 40, 20]), 'value': None},
 inputs_inv3 = {'data_data': {'shape': int64_array([10, 40, 20, 10, 2]), 'value': None},
               'indices_data': {'shape': int64_array([10, 40, 4]), 'value': None}}

-class TestScatterNDUpdate(unittest.TestCase):
+class TestGatherNDUpdate(unittest.TestCase):
    def setUp(self):
        nodes_attributes['gathernd_node']['batch_dims'] = 0