From 53911a91f67018eb2f3987fb658e8fb956df61f2 Mon Sep 17 00:00:00 2001 From: Roman Kazantsev Date: Wed, 11 Nov 2020 14:47:00 +0300 Subject: [PATCH] Fix GatherND to output dynamic rank tensor (#3073) Signed-off-by: Roman Kazantsev --- ngraph/core/src/op/gather_nd.cpp | 2 +- ngraph/test/type_prop/gather_nd.cpp | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/ngraph/core/src/op/gather_nd.cpp b/ngraph/core/src/op/gather_nd.cpp index bee1bc8976b..9ad86145abb 100644 --- a/ngraph/core/src/op/gather_nd.cpp +++ b/ngraph/core/src/op/gather_nd.cpp @@ -142,7 +142,7 @@ void op::v5::GatherND::validate_and_infer_types() } else { - set_output_type(0, data_type, PartialShape{Dimension::dynamic()}); + set_output_type(0, data_type, PartialShape::dynamic()); } } diff --git a/ngraph/test/type_prop/gather_nd.cpp b/ngraph/test/type_prop/gather_nd.cpp index b0c9b14d514..ea4c6d23070 100644 --- a/ngraph/test/type_prop/gather_nd.cpp +++ b/ngraph/test/type_prop/gather_nd.cpp @@ -96,6 +96,18 @@ TEST(type_prop, gather_nd_batch_dim2_with_dyn_dim3) ASSERT_TRUE(G5->get_output_partial_shape(0).same_scheme(out_shape)); } +TEST(type_prop, gather_nd_batch_dim0_with_dyn_ind_dim) +{ + PartialShape params_shape{ + 7, Dimension::dynamic(), Dimension::dynamic(), 12, Dimension::dynamic()}; + PartialShape indices_shape{7, 5, 3, Dimension::dynamic()}; + auto P = make_shared(element::f32, params_shape); + auto I = make_shared(element::i32, indices_shape); + auto G5 = make_shared(P, I, 0); + ASSERT_EQ(G5->get_element_type(), element::f32); + ASSERT_TRUE(G5->get_output_partial_shape(0).same_scheme(PartialShape::dynamic())); +} + TEST(type_prop, gather_nd_fail_batch_dims_greater_indices_rank) { Shape params_shape{2, 3, 4, 5};