diff --git a/runtime/bindings/python/src/openvino/opset8/__init__.py b/runtime/bindings/python/src/openvino/opset8/__init__.py index 80711acc2f5..c53da9026b8 100644 --- a/runtime/bindings/python/src/openvino/opset8/__init__.py +++ b/runtime/bindings/python/src/openvino/opset8/__init__.py @@ -55,7 +55,7 @@ from openvino.opset1.ops import floor from openvino.opset1.ops import floor_mod from openvino.opset8.ops import gather from openvino.opset6.ops import gather_elements -from openvino.opset5.ops import gather_nd +from openvino.opset8.ops import gather_nd from openvino.opset1.ops import gather_tree from openvino.opset7.ops import gelu from openvino.opset1.ops import greater diff --git a/runtime/bindings/python/src/openvino/opset8/ops.py b/runtime/bindings/python/src/openvino/opset8/ops.py index c4a30fce7cf..ef29f36cf08 100644 --- a/runtime/bindings/python/src/openvino/opset8/ops.py +++ b/runtime/bindings/python/src/openvino/opset8/ops.py @@ -391,3 +391,26 @@ def slice( inputs = as_nodes(data, start, stop, step, axes) return _get_node_factory_opset8().create("Slice", inputs) + + +@nameable_op +def gather_nd( + data: NodeInput, + indices: NodeInput, + batch_dims: Optional[int] = 0, + name: Optional[str] = None, +) -> Node: + """Return a node which performs GatherND. + + @param data: N-D tensor with data for gathering + @param indices: K-D tensor of tuples with indices by which data is gathered + @param batch_dims: Scalar value of batch dimensions + @return: The new node which performs GatherND + """ + inputs = as_nodes(data, indices) + + attributes = { + "batch_dims": batch_dims + } + + return _get_node_factory_opset8().create("GatherND", inputs, attributes) diff --git a/runtime/bindings/python/tests/test_ngraph/test_data_movement.py b/runtime/bindings/python/tests/test_ngraph/test_data_movement.py index 624427210a8..6bdb57481c5 100644 --- a/runtime/bindings/python/tests/test_ngraph/test_data_movement.py +++ b/runtime/bindings/python/tests/test_ngraph/test_data_movement.py @@ -188,13 +188,13 @@ def test_select(): assert np.allclose(result, excepted) -def test_gather_nd(): +def test_gather_v8_nd(): indices_type = np.int32 data_dtype = np.float32 data = ov.parameter([2, 10, 80, 30, 50], dtype=data_dtype, name="data") indices = ov.parameter([2, 10, 30, 40, 2], dtype=indices_type, name="indices") batch_dims = 2 - expected_shape = [20, 30, 40, 50] + expected_shape = [2, 10, 30, 40, 50] node = ov.gather_nd(data, indices, batch_dims) assert node.get_type_name() == "GatherND"