finalized MO implementation

model-optimizer/extensions/front/onnx/gatherelements_ext.py

@@ -1,5 +1,5 @@
 """
- Copyright (C) 2020 Intel Corporation
+ Copyright (C) 2017-2020 Intel Corporation
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

model-optimizer/extensions/ops/gatherelements.py

@@ -14,8 +14,10 @@
  limitations under the License.
 """
 
+import numpy as np
+
 from mo.graph.graph import Node, Graph
-from mo.ops.op import Op
+from mo.ops.op import Op, PermuteAttrs
 
 
 class GatherElements(Op):
@@ -37,7 +39,33 @@ def backend_attrs(self):
 
     @staticmethod
     def infer(node: Node):
+        data_shape = node.in_port(0).data.get_shape()
         indices_shape = node.in_port(1).data.get_shape()
+        axis = node.axis
+        data_rank = len(data_shape)
+
+        assert data_rank == len(indices_shape), f'data and indices inputs for node {node.name} must be of the ' \
+                                                f'same rank. Got {data_rank} and {len(indices_shape)}'
+        assert -data_rank < axis < data_rank, f'axis for node {node.name} must be within interval ' \
+                                              f'[{-data_rank},  {data_rank - 1}]. Instead axis={axis}'
+        if data_rank < 0:
+            axis += data_rank
+
+        for idx, (data_sz, ind_sz) in enumerate(zip(data_shape, indices_shape)):
+            if idx != axis and data_sz != ind_sz:
+                raise ValueError(f'Sizes along axis {idx} for node {node.name} do not match. '
+                                 f'data and indices must have equal size along all axes except for axis {axis}')
+
         node.out_port(0).data.set_shape(indices_shape)
 
-        # todo: add value_inference
+        data = node.in_port(0).data.get_value()
+        indices = node.in_port(1).data.get_value()
+        if data is not None and indices is not None:
+            out_value = np.empty(indices_shape, dtype=data.dtype)
+            for idx in np.ndindex(*indices_shape):
+                data_idx = list(idx)
+                data_idx[node.axis] = indices[idx]
+                out_value[idx] = data[tuple(data_idx)]
+            node.out_port(0).data.set_value(out_value)
+
+        PermuteAttrs.create_permute_attrs(node, attrs=[('axis', 'input:0')])

model-optimizer/extensions/ops/gatherelements_test.py

@@ -0,0 +1,115 @@
+"""
+ Copyright (C) 2017-2020 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+import unittest
+
+import numpy as np
+from generator import generator, generate
+
+from extensions.ops.gatherelements import GatherElements
+from mo.front.common.partial_infer.utils import int64_array
+from mo.graph.graph import Node
+from mo.utils.unittest.graph import build_graph, regular_op_with_empty_data, result, connect, \
+    valued_const_with_data
+
+
+@generator
+class GatherElementsInferTest(unittest.TestCase):
+    @generate(*[
+        ([[1, 2],
+          [3, 4]],
+         [[0, 1],
+          [0, 0]],
+         0,  # axis
+         [[1, 4],  # ref_res
+          [1, 2]]),
+
+        ([[1, 2],
+          [3, 4]],
+         [[0, 1],
+          [0, 0]],
+         1,  # axis
+         [[1, 2],  # ref_res
+          [3, 3]]),
+
+        ([[1, 2, 3],
+          [4, 5, 6],
+          [7, 8, 9]],
+         [[1, 2, 0],
+          [2, 0, 0]],
+         0,  # axis
+         [[4, 8, 3],  # ref_res
+          [7, 2, 3]]),
+
+        ([[1, 2],
+          [3, 4]],
+         [[0, 1],
+          [0, 0]],
+         -1,  # axis
+         [[1, 2],  # ref_res
+          [3, 3]]),
+
+        ([  # 3D case
+          [[1, 2],
+           [3, 4]],
+          [[5, 6],
+           [7, 8]],
+          [[9, 10],
+           [11, 12]]
+        ],
+         [
+          [[1, 0],
+           [0, 1]],
+          [[1, 1],
+           [1, 0]],
+          [[0, 0],
+           [1, 1]]
+         ],
+         -1,  # axis
+         [
+          [[2, 1],
+           [3, 4]],
+          [[6, 6],
+           [8, 7]],
+          [[9, 9],
+           [12, 12]]
+         ]),
+    ])
+
+    def test_gatherelements_value_infer(self, data, indices, axis, ref_res):
+        nodes = {
+            **valued_const_with_data('data', int64_array(data)),
+            **valued_const_with_data('indices', int64_array(indices)),
+            **regular_op_with_empty_data('gather_elements', {'op': 'GatherElements', 'axis': axis}),
+            **result()
+        }
+
+        graph = build_graph(nodes_attrs=nodes, edges=[
+            *connect('data', '0:gather_elements'),
+            *connect('indices', '1:gather_elements'),
+            *connect('gather_elements', 'output')
+        ], nodes_with_edges_only=True)
+        graph.stage = 'middle'
+
+        gather_el_node = Node(graph, 'gather_elements')
+        GatherElements.infer(gather_el_node)
+
+        res_output_shape = gather_el_node.out_node().shape
+        self.assertTrue(np.array_equal(int64_array(ref_res).shape, res_output_shape))
+
+        res_output_value = gather_el_node.out_node().value
+        if res_output_value is not None:
+            self.assertTrue(np.array_equal(int64_array(ref_res), res_output_value))