Fix GatherND to output dynamic rank tensor (#3073)

Signed-off-by: Roman Kazantsev <roman.kazantsev@intel.com>

diff --git a/ngraph/core/src/op/gather_nd.cpp b/ngraph/core/src/op/gather_nd.cpp
index bee1bc8..9ad8614 100644 (file)
--- 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 b0c9b14..ea4c6d2 100644 (file)
--- 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<op::Parameter>(element::f32, params_shape);
+    auto I = make_shared<op::Parameter>(element::i32, indices_shape);
+    auto G5 = make_shared<op::v5::GatherND>(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};