Add gather 7 to nG Python API and IE IR reader test (#5276)

* nG Python api for Gather and IE IR reader tests

* also included old IE IR reader test for Gather1

* added a blank line before Gather

* style corrections

* applied review comments

* removed blank line

* removed xfailed for Gather7 with batch_dims since CPU was moved to nGraph (...allows to fallback node execution on ngraph evaluate() method in case if optimized implementation is absent)

* added new visitor tests

inference-engine/tests/functional/inference_engine/ngraph_reader/gather_tests.cpp

@@ -4,7 +4,8 @@
 
 #include <string>
 #include "ngraph_reader_tests.hpp"
-TEST_F(NGraphReaderTests, ReadGatherNetwork) {
+
+TEST_F(NGraphReaderTests, Read_Gather1_Network) {
     std::string model = R"V0G0N(
 <net name="saved_model" version="10">
 	<layers>
@@ -117,7 +118,133 @@ TEST_F(NGraphReaderTests, ReadGatherNetwork) {
 </net>
 )V0G0N";
     compareIRs(model, modelV7, 16, [](Blob::Ptr& weights) {
-                auto* buffer = weights->buffer().as<int64_t*>();
-                buffer[0] = 0;
-            });
-}
+        auto* buffer = weights->buffer().as<int64_t*>();
+        buffer[0] = 0;
+    });
+}
+
+TEST_F(NGraphReaderTests, Read_Gather7_Network) {
+    std::string model = R"V0G0N(
+<net name="saved_model" version="10">
+	<layers>
+		<layer id="0" name="data" type="Parameter" version="opset1">
+			<data shape="2,3,4" element_type="f32"/>
+			<output>
+				<port id="0" precision="FP32">
+					<dim>2</dim>
+					<dim>3</dim>
+					<dim>4</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="1" name="indices" type="Parameter" version="opset1">
+			<data shape="8,16" element_type="i32"/>
+			<output>
+				<port id="0" precision="I32">
+					<dim>8</dim>
+					<dim>16</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="2" name="GatherV2/axis" type="Const" version="opset1">
+			<data offset="0" size="8" shape="" element_type="i64"/>
+			<output>
+				<port id="1" precision="I64"/>
+			</output>
+		</layer>
+		<layer id="3" name="Gather" type="Gather" version="opset7">
+            <data batch_dims="0"/>
+			<input>
+				<port id="0">
+					<dim>2</dim>
+					<dim>3</dim>
+					<dim>4</dim>
+				</port>
+				<port id="1">
+					<dim>8</dim>
+					<dim>16</dim>
+				</port>
+				<port id="2"/>
+			</input>
+			<output>
+				<port id="3" precision="FP32">
+					<dim>2</dim>
+					<dim>8</dim>
+					<dim>16</dim>
+					<dim>4</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="4" name="Identity/sink_port_0" type="Result" version="opset1">
+			<input>
+				<port id="0">
+					<dim>2</dim>
+					<dim>8</dim>
+					<dim>16</dim>
+					<dim>4</dim>
+				</port>
+			</input>
+		</layer>
+	</layers>
+	<edges>
+		<edge from-layer="0" from-port="0" to-layer="3" to-port="0"/>
+		<edge from-layer="1" from-port="0" to-layer="3" to-port="1"/>
+		<edge from-layer="2" from-port="1" to-layer="3" to-port="2"/>
+		<edge from-layer="3" from-port="3" to-layer="4" to-port="0"/>
+	</edges>
+</net>
+)V0G0N";
+    std::string modelV7 = R"V0G0N(
+<net name="saved_model" version="7">
+	<layers>
+		<layer id="0" name="data" type="Input" version="opset1">
+			<output>
+				<port id="0" precision="FP32">
+					<dim>2</dim>
+					<dim>3</dim>
+					<dim>4</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="1" name="indices" type="Input" version="opset1">
+			<output>
+				<port id="0" precision="I32">
+					<dim>8</dim>
+					<dim>16</dim>
+				</port>
+			</output>
+		</layer>
+		<layer id="3" name="Gather" type="Gather">
+			<data axis="1"/>
+			<input>
+				<port id="0">
+					<dim>2</dim>
+					<dim>3</dim>
+					<dim>4</dim>
+				</port>
+				<port id="1">
+					<dim>8</dim>
+					<dim>16</dim>
+				</port>
+			</input>
+			<output>
+				<port id="3" precision="FP32">
+					<dim>2</dim>
+					<dim>8</dim>
+					<dim>16</dim>
+					<dim>4</dim>
+				</port>
+			</output>
+		</layer>
+	</layers>
+	<edges>
+		<edge from-layer="0" from-port="0" to-layer="3" to-port="0"/>
+		<edge from-layer="1" from-port="0" to-layer="3" to-port="1"/>
+	</edges>
+</net>
+)V0G0N";
+    compareIRs(model, modelV7, 8, [](Blob::Ptr& weights) {
+        auto* buffer = weights->buffer().as<int64_t*>();
+        buffer[0] = 1;
+    });
+}

ngraph/python/src/ngraph/opset6/ops.py

@@ -84,7 +84,7 @@ def gather_elements(
     axis: Optional[int] = 0,
     name: Optional[str] = None,
 ) -> Node:
-    """Return a node which performs GatherND.
+    """Return a node which performs GatherElements.
 
     @param data:       N-D tensor with data for gathering
     @param indices:    N-D tensor with indices by which data is gathered

ngraph/python/src/ngraph/opset7/__init__.py

@@ -51,7 +51,7 @@ from ngraph.opset1.ops import exp
 from ngraph.opset1.ops import fake_quantize
 from ngraph.opset1.ops import floor
 from ngraph.opset1.ops import floor_mod
-from ngraph.opset1.ops import gather
+from ngraph.opset7.ops import gather
 from ngraph.opset6.ops import gather_elements
 from ngraph.opset5.ops import gather_nd
 from ngraph.opset1.ops import gather_tree

ngraph/python/src/ngraph/opset7/ops.py

@@ -104,6 +104,27 @@ def roll(
 
 
 @nameable_op
+def gather(
+        data: NodeInput,
+        indices: NodeInput,
+        axis: NodeInput,
+        batch_dims: Optional[int] = 0,
+) -> Node:
+    """Return a node which performs Gather.
+
+    @param data:         N-D tensor with data for gathering
+    @param indices:      N-D tensor with indices by which data is gathered
+    @param axis:         axis along which elements are gathered
+    @param batch_dims:   number of batch dimensions
+    @return:             The new node which performs Gather
+    """
+    inputs = as_nodes(data, indices, axis)
+    attributes = {
+        "batch_dims": batch_dims
+    }
+    return _get_node_factory_opset7().create("Gather", inputs, attributes)
+
+
 def dft(
         data: NodeInput,
         axes: NodeInput,

ngraph/python/tests/__init__.py

@@ -157,7 +157,6 @@ xfail_issue_52463 = xfail_test(reason="test_operator_add_size1_singleton_broadca
                                       "Not equal to tolerance")
 
 xfail_issue_49391 = xfail_test(reason="Roll is not implemented in CPU plugin.")
-
 xfail_issue_49359 = xfail_test(reason="DFT is not implemented in CPU plugin")
 xfail_issue_49375 = xfail_test(reason="IDFT is not implemented in CPU plugin")
 xfail_issue_45432 = xfail_test(reason="Einsum is not implemented in CPU plugin.")

ngraph/python/tests/test_ngraph/test_gather.py

@@ -0,0 +1,54 @@
+# Copyright (C) 2018-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import ngraph as ng
+import numpy as np
+
+from tests.test_ngraph.util import run_op_node
+
+
+def test_gather():
+    input_data = np.array(
+        [1.0, 1.1, 1.2, 2.0, 2.1, 2.2, 3.0, 3.1, 3.2], np.float32
+    ).reshape((3, 3))
+    input_indices = np.array([0, 2], np.int32).reshape(1, 2)
+    input_axis = np.array([1], np.int32)
+
+    expected = np.array([1.0, 1.2, 2.0, 2.2, 3.0, 3.2], dtype=np.float32).reshape(
+        (3, 1, 2)
+    )
+
+    result = run_op_node([input_data], ng.gather, input_indices, input_axis)
+    assert np.allclose(result, expected)
+
+
+def test_gather_with_scalar_axis():
+    input_data = np.array(
+        [1.0, 1.1, 1.2, 2.0, 2.1, 2.2, 3.0, 3.1, 3.2], np.float32
+    ).reshape((3, 3))
+    input_indices = np.array([0, 2], np.int32).reshape(1, 2)
+    input_axis = np.array(1, np.int32)
+
+    expected = np.array([1.0, 1.2, 2.0, 2.2, 3.0, 3.2], dtype=np.float32).reshape(
+        (3, 1, 2)
+    )
+
+    result = run_op_node([input_data], ng.gather, input_indices, input_axis)
+    assert np.allclose(result, expected)
+
+
+def test_gather_batch_dims_1():
+
+    input_data = np.array([[1, 2, 3, 4, 5],
+                           [6, 7, 8, 9, 10]], np.float32)
+
+    input_indices = np.array([[0, 0, 4],
+                              [4, 0, 0]], np.int32)
+    input_axis = np.array([1], np.int32)
+    batch_dims = 1
+
+    expected = np.array([[1, 1, 5],
+                         [10, 6, 6]], np.float32)
+
+    result = run_op_node([input_data], ng.gather, input_indices, input_axis, batch_dims)
+    assert np.allclose(result, expected)

ngraph/python/tests/test_ngraph/test_ops_reshape.py

@@ -1,10 +1,10 @@
 # Copyright (C) 2018-2021 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+import ngraph as ng
 import numpy as np
 import pytest
 
-import ngraph as ng
 from tests.runtime import get_runtime
 from tests.test_ngraph.util import run_op_node, run_op_numeric_data
 
@@ -120,21 +120,6 @@ def test_broadcast_bidirectional():
     assert node.get_output_size() == 1
 
 
-def test_gather():
-    input_data = np.array(
-        [1.0, 1.1, 1.2, 2.0, 2.1, 2.2, 3.0, 3.1, 3.2], np.float32
-    ).reshape((3, 3))
-    input_indices = np.array([0, 2], np.int32).reshape(1, 2)
-    input_axes = np.array([1], np.int32)
-
-    expected = np.array([1.0, 1.2, 2.0, 2.2, 3.0, 3.2], dtype=np.float32).reshape(
-        (3, 1, 2)
-    )
-
-    result = run_op_node([input_data], ng.gather, input_indices, input_axes)
-    assert np.allclose(result, expected)
-
-
 def test_transpose():
     input_tensor = np.arange(3 * 3 * 224 * 224, dtype=np.int32).reshape(
         (3, 3, 224, 224)

ngraph/test/CMakeLists.txt

@@ -219,6 +219,7 @@ set(SRC
     visitors/op/elu.cpp
     visitors/op/extractimagepatches.cpp
     visitors/op/fake_quantize.cpp
+    visitors/op/gather.cpp
     visitors/op/grn.cpp
     visitors/op/group_conv.cpp
     visitors/op/interpolate.cpp

ngraph/test/visitors/op/gather.cpp

@@ -0,0 +1,31 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "gtest/gtest.h"
+
+#include "ngraph/ngraph.hpp"
+#include "ngraph/opsets/opset1.hpp"
+#include "ngraph/opsets/opset7.hpp"
+
+#include "util/visitor.hpp"
+
+using namespace std;
+using namespace ngraph;
+using ngraph::test::NodeBuilder;
+using ngraph::test::ValueMap;
+
+TEST(attributes, gather_v7_op)
+{
+    NodeBuilder::get_ops().register_factory<opset7::Gather>();
+    auto data = make_shared<opset1::Parameter>(element::i32, Shape{2, 3, 4});
+    auto indices = make_shared<opset1::Parameter>(element::i32, Shape{2});
+    auto axis = make_shared<opset1::Constant>(element::i32, Shape{}, 2);
+    int64_t batch_dims = 1;
+
+    auto gather = make_shared<opset7::Gather>(data, indices, axis, batch_dims);
+    NodeBuilder builder(gather);
+    auto g_gather = as_type_ptr<opset7::Gather>(builder.create());
+
+    EXPECT_EQ(g_gather->get_batch_dims(), gather->get_batch_dims());
+}