Propagate input rank for dynamic output shape in BatchToSpace,DepthTo… (#4588)

* Propagate input rank for dynamic output shape in BatchToSpace,DepthToSpace,Pad,SpaceToBatch,SpaceToDepth

* Add tests for SpaceToDepth, DepthToSpace dynamic input

ngraph/core/src/op/batch_to_space.cpp

@@ -124,7 +124,7 @@ void op::v1::BatchToSpace::validate_and_infer_types()
     }
     else
     {
-        set_output_type(0, data_type, PartialShape::dynamic());
+        set_output_type(0, data_type, PartialShape::dynamic(data_pshape.rank()));
     }
 }
 

ngraph/core/src/op/depth_to_space.cpp

@@ -112,7 +112,7 @@ void op::DepthToSpace::validate_and_infer_types()
     }
     else
     {
-        set_output_type(0, data_type, PartialShape::dynamic());
+        set_output_type(0, data_type, PartialShape::dynamic(data_pshape.rank()));
     }
 }
 

ngraph/core/src/op/pad.cpp

@@ -192,7 +192,7 @@ void op::v1::Pad::validate_and_infer_types()
     }
     else
     {
-        set_output_type(0, get_input_element_type(0), PartialShape::dynamic());
+        set_output_type(0, get_input_element_type(0), PartialShape::dynamic(arg_shape_rank));
     }
 }
 

ngraph/core/src/op/space_to_batch.cpp

@@ -124,7 +124,7 @@ void op::v1::SpaceToBatch::validate_and_infer_types()
     }
     else
     {
-        set_output_type(0, data_type, PartialShape::dynamic());
+        set_output_type(0, data_type, PartialShape::dynamic(data_pshape.rank()));
     }
 }
 

ngraph/core/src/op/space_to_depth.cpp

@@ -109,7 +109,7 @@ void ngraph::op::v0::SpaceToDepth::validate_and_infer_types()
     }
     else
     {
-        set_output_type(0, data_type, PartialShape::dynamic());
+        set_output_type(0, data_type, PartialShape::dynamic(data_pshape.rank()));
     }
 }
 

ngraph/frontend/onnx_import/src/op/depth_to_space.cpp

@@ -28,7 +28,9 @@ namespace ngraph
                 OutputVector depth_to_space(const Node& node)
                 {
                     auto data = node.get_ng_inputs().at(0);
-                    NGRAPH_CHECK(data.get_shape().size() == 4, "Input must be 4-dimensional");
+                    const auto& shape = data.get_partial_shape();
+                    NGRAPH_CHECK(shape.rank().is_static() && shape.rank().get_length() == 4,
+                                 "Input must be 4-dimensional");
 
                     const auto mode = node.get_attribute_value<std::string>("mode", "DCR");
                     default_opset::DepthToSpace::DepthToSpaceMode ngraph_mode;

ngraph/frontend/onnx_import/src/op/space_to_depth.cpp

@@ -28,7 +28,9 @@ namespace ngraph
                 OutputVector space_to_depth(const Node& node)
                 {
                     auto data = node.get_ng_inputs().at(0);
-                    NGRAPH_CHECK(data.get_shape().size() == 4, "Input must be 4-dimensional");
+                    const auto& shape = data.get_partial_shape();
+                    NGRAPH_CHECK(shape.rank().is_static() && shape.rank().get_length() == 4,
+                                 "Input must be 4-dimensional");
                     std::size_t block_size = node.get_attribute_value<std::int64_t>("blocksize");
                     const auto mode = default_opset::SpaceToDepth::SpaceToDepthMode::BLOCKS_FIRST;
                     return OutputVector{

ngraph/test/models/onnx/dynamic_shapes/depth_to_space.prototxt

@@ -0,0 +1,59 @@
+ir_version: 3
+producer_name: "nGraph ONNX Importer"
+graph {
+  node {
+    input: "A"
+    output: "B"
+    op_type: "DepthToSpace"
+    attribute {
+      name: "blocksize"
+      i: 2
+      type: INT
+    }
+    attribute {
+      name: "mode"
+      s: "DCR"
+      type: STRING
+    }
+  }
+  name: "compute_graph"
+  input {
+    name: "A"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+          }
+          dim {
+          }
+          dim {
+          }
+          dim {
+          }
+        }
+      }
+    }
+  }
+  output {
+    name: "B"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+          }
+          dim {
+          }
+          dim {
+          }
+          dim {
+          }
+        }
+      }
+    }
+  }
+}
+opset_import {
+  version: 11
+}

ngraph/test/models/onnx/dynamic_shapes/space_to_depth.prototxt

@@ -0,0 +1,54 @@
+ir_version: 3
+producer_name: "nGraph ONNX Importer"
+graph {
+  node {
+    input: "A"
+    output: "B"
+    op_type: "SpaceToDepth"
+    attribute {
+      name: "blocksize"
+      i: 2
+      type: INT
+    }
+  }
+  name: "compute_graph"
+  input {
+    name: "A"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+          }
+          dim {
+          }
+          dim {
+          }
+          dim {
+          }
+        }
+      }
+    }
+  }
+  output {
+    name: "B"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+          }
+          dim {
+          }
+          dim {
+          }
+          dim {
+          }
+        }
+      }
+    }
+  }
+}
+opset_import {
+  version: 1
+}

ngraph/test/onnx/onnx_import_dyn_shapes.in.cpp

@@ -1337,3 +1337,40 @@ NGRAPH_TEST(${BACKEND_NAME}, onnx_model_max_pool_dyn_rank_without_default_attrs)
     test_case.add_expected_output<float>(Shape{1, 1, 3, 3}, {5, 6, 7, 9, 10, 11, 13, 14, 15});
     test_case.run();
 }
+
+NGRAPH_TEST(${BACKEND_NAME}, onnx_model_depth_to_space_dynamic_input)
+{
+    auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/dynamic_shapes/depth_to_space.prototxt"));
+
+    std::vector<float> input(32);
+    std::iota(input.begin(), input.end(), 0);
+
+    std::vector<float> expected_output{
+        0.f, 8.f,  1.f, 9.f,  16.f, 24.f, 17.f, 25.f, 2.f, 10.f, 3.f, 11.f, 18.f, 26.f, 19.f, 27.f,
+        4.f, 12.f, 5.f, 13.f, 20.f, 28.f, 21.f, 29.f, 6.f, 14.f, 7.f, 15.f, 22.f, 30.f, 23.f, 31.f};
+
+    auto test_case = test::TestCase<TestEngine, TestCaseType::DYNAMIC>(function);
+    test_case.add_input(Shape{1, 8, 2, 2}, input);
+    test_case.add_expected_output(Shape{1, 2, 4, 4}, expected_output);
+    test_case.run();
+}
+
+NGRAPH_TEST(${BACKEND_NAME}, onnx_model_space_to_depth_dynamic_input)
+{
+    auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/dynamic_shapes/space_to_depth.prototxt"));
+
+    std::vector<float> input(32);
+    std::iota(input.begin(), input.end(), 0);
+
+    std::vector<float> expected_output{
+        0.f, 2.f, 8.f,  10.f, 16.f, 18.f, 24.f, 26.f, 1.f, 3.f, 9.f,  11.f, 17.f, 19.f, 25.f, 27.f,
+        4.f, 6.f, 12.f, 14.f, 20.f, 22.f, 28.f, 30.f, 5.f, 7.f, 13.f, 15.f, 21.f, 23.f, 29.f, 31.f,
+    };
+
+    auto test_case = test::TestCase<TestEngine, TestCaseType::DYNAMIC>(function);
+    test_case.add_input(Shape{1, 2, 4, 4}, input);
+    test_case.add_expected_output(Shape{1, 8, 2, 2}, expected_output);
+    test_case.run();
+}

ngraph/test/runtime/ie/unit_test.manifest

@@ -242,6 +242,9 @@ IE_CPU.nothing_to_reverse
 onnx_size_dyn_op
 onnx_model_gru_defaults_fwd_const_dynamic
 onnx_model_rnn_defaults_fwd_const_dynamic
+onnx_model_depth_to_space_dynamic_input
+onnx_model_space_to_depth_dynamic_input
+
 
 # Constant network
 # MKLDNNGraph::CreateGraph: No inputs for the topology

ngraph/test/type_prop/batch_to_space.cpp

@@ -90,3 +90,35 @@ TEST(type_prop, batch_to_space_and_space_to_batch)
     ASSERT_EQ(space_to_batch->get_element_type(), element::f32);
     ASSERT_EQ(space_to_batch->get_shape(), (Shape{4800, 9, 11, 2}));
 }
+
+TEST(type_prop, batch_to_space_dynamic_shape_static_rank)
+{
+    auto data = make_shared<op::Parameter>(element::f32, PartialShape::dynamic(4));
+    auto block_shape =
+        make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{1, 10, 5, 1});
+    auto pads_begin =
+        make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{0, 3, 1, 0});
+    auto pads_end = make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{0, 3, 0, 0});
+
+    auto batch_to_space =
+        make_shared<op::v1::BatchToSpace>(data, block_shape, pads_begin, pads_end);
+
+    ASSERT_EQ(batch_to_space->get_element_type(), element::f32);
+    ASSERT_EQ(batch_to_space->get_output_partial_shape(0), PartialShape::dynamic(4));
+}
+
+TEST(type_prop, batch_to_space_dynamic_shape_dynamic_rank)
+{
+    auto data = make_shared<op::Parameter>(element::f32, PartialShape::dynamic());
+    auto block_shape =
+        make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{1, 10, 5, 1});
+    auto pads_begin =
+        make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{0, 3, 1, 0});
+    auto pads_end = make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{0, 3, 0, 0});
+
+    auto batch_to_space =
+        make_shared<op::v1::BatchToSpace>(data, block_shape, pads_begin, pads_end);
+
+    ASSERT_EQ(batch_to_space->get_element_type(), element::f32);
+    ASSERT_EQ(batch_to_space->get_output_partial_shape(0), PartialShape::dynamic());
+}

ngraph/test/type_prop/depth_to_space.cpp

@@ -112,3 +112,23 @@ TEST(type_prop, depth_to_space_blocksize_not_matched)
         FAIL() << "DepthToSpace decomposition failed for unexpected reason";
     }
 }
+
+TEST(type_prop, depth_to_space_dynamic_shape_static_rank)
+{
+    auto A = make_shared<op::Parameter>(element::f32, PartialShape::dynamic(4));
+    auto space_to_depth =
+        make_shared<op::DepthToSpace>(A, op::DepthToSpace::DepthToSpaceMode::DEPTH_FIRST, 2);
+
+    ASSERT_EQ(space_to_depth->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_depth->get_output_partial_shape(0), PartialShape::dynamic(4));
+}
+
+TEST(type_prop, depth_to_space_dynamic_shape_dynamic_rank)
+{
+    auto A = make_shared<op::Parameter>(element::f32, PartialShape::dynamic());
+    auto space_to_depth =
+        make_shared<op::DepthToSpace>(A, op::DepthToSpace::DepthToSpaceMode::DEPTH_FIRST, 2);
+
+    ASSERT_EQ(space_to_depth->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_depth->get_output_partial_shape(0), PartialShape::dynamic());
+}

ngraph/test/type_prop/pad.cpp

@@ -281,3 +281,27 @@ TEST(type_prop, pad_v1_deduce_too_small_for_reflect)
         FAIL() << "Deduced type check failed for unexpected reason";
     }
 }
+
+TEST(type_prop, pad_v1_dynamic_output_with_dynamic_rank)
+{
+    auto arg = make_shared<op::Parameter>(element::f32, PartialShape::dynamic());
+    auto pads_begin = make_shared<op::Parameter>(element::i32, Shape{1});
+    auto pads_end = make_shared<op::Parameter>(element::i32, Shape{1});
+    auto arg_pad_value = op::Constant::create(element::f32, Shape{}, {0});
+
+    auto pad =
+        make_shared<op::v1::Pad>(arg, pads_begin, pads_end, arg_pad_value, op::PadMode::CONSTANT);
+    ASSERT_EQ(pad->get_output_partial_shape(0), PartialShape::dynamic());
+}
+
+TEST(type_prop, pad_v1_dynamic_output_with_static_rank)
+{
+    auto arg = make_shared<op::Parameter>(element::f32, Shape{1, 2, 3});
+    auto pads_begin = make_shared<op::Parameter>(element::i32, Shape{1});
+    auto pads_end = make_shared<op::Parameter>(element::i32, Shape{1});
+    auto arg_pad_value = op::Constant::create(element::f32, Shape{}, {0});
+
+    auto pad =
+        make_shared<op::v1::Pad>(arg, pads_begin, pads_end, arg_pad_value, op::PadMode::CONSTANT);
+    ASSERT_EQ(pad->get_output_partial_shape(0), PartialShape::dynamic(3));
+}

ngraph/test/type_prop/space_to_batch.cpp

@@ -90,3 +90,35 @@ TEST(type_prop, space_to_batch_and_batch_to_space)
     ASSERT_EQ(batch_to_space->get_element_type(), element::f32);
     ASSERT_EQ(batch_to_space->get_shape(), (Shape{2, 100, 1024, 3}));
 }
+
+TEST(type_prop, space_to_batch_dynamic_shape_static_rank)
+{
+    auto data = make_shared<op::Parameter>(element::f32, PartialShape::dynamic(4));
+    auto block_shape =
+        make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{1, 10, 5, 1});
+    auto pads_begin =
+        make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{0, 3, 1, 0});
+    auto pads_end = make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{0, 3, 0, 0});
+
+    auto space_to_batch =
+        make_shared<op::v1::SpaceToBatch>(data, block_shape, pads_begin, pads_end);
+
+    ASSERT_EQ(space_to_batch->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_batch->get_output_partial_shape(0), PartialShape::dynamic(4));
+}
+
+TEST(type_prop, space_to_batch_dynamic_shape_dynamic_rank)
+{
+    auto data = make_shared<op::Parameter>(element::f32, PartialShape::dynamic());
+    auto block_shape =
+        make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{1, 10, 5, 1});
+    auto pads_begin =
+        make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{0, 3, 1, 0});
+    auto pads_end = make_shared<op::Constant>(element::i64, Shape{4}, vector<int64_t>{0, 3, 0, 0});
+
+    auto space_to_batch =
+        make_shared<op::v1::SpaceToBatch>(data, block_shape, pads_begin, pads_end);
+
+    ASSERT_EQ(space_to_batch->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_batch->get_output_partial_shape(0), PartialShape::dynamic());
+}

ngraph/test/type_prop/space_to_depth.cpp

@@ -61,6 +61,26 @@ TEST(type_prop, space_to_depth_output_shape_depth_first_5D)
     ASSERT_EQ(space_to_depth->get_shape(), (Shape{1, 12 * 8, 4 / 2, 1080 / 2, 1616 / 2}));
 }
 
+TEST(type_prop, space_to_depth_dynamic_shape_static_rank)
+{
+    auto A = make_shared<op::Parameter>(element::f32, PartialShape::dynamic(4));
+    const auto mode = ngraph::op::SpaceToDepth::SpaceToDepthMode::BLOCKS_FIRST;
+    auto space_to_depth = make_shared<op::SpaceToDepth>(A, mode, 8);
+
+    ASSERT_EQ(space_to_depth->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_depth->get_output_partial_shape(0), PartialShape::dynamic(4));
+}
+
+TEST(type_prop, space_to_depth_dynamic_shape_dynamic_rank)
+{
+    auto A = make_shared<op::Parameter>(element::f32, PartialShape::dynamic());
+    const auto mode = ngraph::op::SpaceToDepth::SpaceToDepthMode::BLOCKS_FIRST;
+    auto space_to_depth = make_shared<op::SpaceToDepth>(A, mode, 8);
+
+    ASSERT_EQ(space_to_depth->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_depth->get_output_partial_shape(0), PartialShape::dynamic());
+}
+
 TEST(type_prop, space_to_depth_input_rank_not_supported)
 {
     auto A = make_shared<op::Parameter>(element::f32, Shape{1, 8});