[Shape Inference] Select - shape inference review (#14708)

* Select shape_infer tests update

* Add Select type_prop tests

* Add evaluate_lower/upper for select

* Revert evaluate_lower/upper for Select

* Use get_node_input_partial_shapes

* Style and headers improvements

* Style apply

* Rename select shape infer file tests

* Use default ctor for output_shapes init

* Use helper for shape_labels init and add more  dim test cases

src/core/shape_inference/include/select_shape_inference.hpp

@@ -45,4 +45,4 @@ void shape_infer(const Select* op, const std::vector<T>& input_shapes, std::vect
 
 }  // namespace v1
 }  // namespace op
-}  // namespace ov
+}  // namespace ov

src/core/src/op/select.cpp

@@ -5,14 +5,12 @@
 #include "ngraph/op/select.hpp"
 
 #include <memory>
-#include <ngraph/validation_util.hpp>
-#include <select_shape_inference.hpp>
 
 #include "itt.hpp"
 #include "ngraph/attribute_visitor.hpp"
-#include "ngraph/op/convert.hpp"
-#include "ngraph/op/not.hpp"
 #include "ngraph/runtime/reference/select.hpp"
+#include "ngraph/validation_util.hpp"
+#include "select_shape_inference.hpp"
 
 using namespace std;
 using namespace ngraph;
@@ -41,10 +39,9 @@ void op::v1::Select::validate_and_infer_types() {
                           element::Type::merge(result_et, get_input_element_type(1), get_input_element_type(2)),
                           "Argument 1 and 2 element types must match.");
 
-    std::vector<ov::PartialShape> output_shapes = {ov::PartialShape{}};
-    const std::vector<ov::PartialShape> input_shapes = {get_input_partial_shape(0),
-                                                        get_input_partial_shape(1),
-                                                        get_input_partial_shape(2)};
+    const auto input_shapes = get_node_input_partial_shapes(*this);
+    auto output_shapes = std::vector<ov::PartialShape>(1);
+
     shape_infer(this, input_shapes, output_shapes);
     set_output_type(0, result_et, output_shapes[0]);
 }

src/core/tests/type_prop/select.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+#include "dimension_tracker.hpp"
 #include "gtest/gtest.h"
 #include "ngraph/ngraph.hpp"
 #include "util/type_prop.hpp"
@@ -20,6 +21,153 @@ TEST(type_prop, select_deduce) {
     ASSERT_EQ(bc->get_shape(), (Shape{2, 4}));
 }
 
+TEST(type_prop, select_default_constructor) {
+    auto cond_param = make_shared<op::Parameter>(element::boolean, Shape{2, 4});
+    auto then_param = make_shared<op::Parameter>(element::f32, Shape{2, 4});
+    auto else_param = make_shared<op::Parameter>(element::f32, Shape{2, 4});
+    auto op = make_shared<op::v1::Select>();
+
+    EXPECT_EQ(op->get_auto_broadcast().m_type, op::AutoBroadcastType::NUMPY);
+
+    op->set_argument(0, cond_param);
+    op->set_argument(1, then_param);
+    op->set_argument(2, else_param);
+
+    op->set_auto_broadcast(op::AutoBroadcastType::NONE);
+    EXPECT_EQ(op->get_auto_broadcast().m_type, op::AutoBroadcastType::NONE);
+
+    op->validate_and_infer_types();
+
+    EXPECT_EQ(op->get_element_type(), element::f32);
+    EXPECT_EQ(op->get_output_shape(0), (Shape{2, 4}));
+}
+
+TEST(type_prop, select_labels_cond_numpy) {
+    auto labeled_shape = PartialShape{{2, 8}, {3, 7}, {1, 10}, {1, 6}, {1, 10}};
+    set_shape_labels(labeled_shape, 10);
+    std::vector<size_t> expected_labels{10, 11, 12, ov::no_label, 14};
+
+    auto cond_param = make_shared<op::Parameter>(element::boolean, labeled_shape);
+    auto then_param = make_shared<op::Parameter>(element::f32, PartialShape::dynamic(5));
+    auto else_param = make_shared<op::Parameter>(element::f32, PartialShape({{1, 5}, {1, 11}, 5, {1, 8}}));
+    auto op = make_shared<op::v1::Select>(cond_param, then_param, else_param);
+
+    const auto& out_shape = op->get_output_partial_shape(0);
+
+    EXPECT_EQ(op->get_element_type(), element::f32);
+    EXPECT_EQ(out_shape, PartialShape({{2, 8}, {3, 7}, -1, 5, -1}));
+    EXPECT_EQ(get_shape_labels(out_shape), expected_labels);
+}
+
+TEST(type_prop, select_labels_then_numpy) {
+    auto labeled_shape = PartialShape::dynamic(5);
+    set_shape_labels(labeled_shape, 10);
+    std::vector<size_t> expected_labels{ov::no_label, ov::no_label, 12, ov::no_label, 14};
+
+    auto cond_param =
+        make_shared<op::Parameter>(element::boolean, PartialShape{{2, 8}, {3, 7}, {1, 10}, {1, 6}, {1, 10}});
+    auto then_param = make_shared<op::Parameter>(element::f32, labeled_shape);
+    auto else_param = make_shared<op::Parameter>(element::f32, PartialShape({{1, 5}, {1, 11}, 5, {1, 8}}));
+    auto op = make_shared<op::v1::Select>(cond_param, then_param, else_param);
+
+    const auto& out_shape = op->get_output_partial_shape(0);
+
+    EXPECT_EQ(op->get_element_type(), element::f32);
+    EXPECT_EQ(out_shape, PartialShape({{2, 8}, {3, 7}, -1, 5, -1}));
+    EXPECT_EQ(get_shape_labels(out_shape), expected_labels);
+}
+
+TEST(type_prop, select_labels_else_numpy) {
+    auto labeled_shape = PartialShape{{1, 5}, {1, 11}, 5, {1, 8}};
+    set_shape_labels(labeled_shape, 10);
+
+    std::vector<size_t> expected_labels{ov::no_label, ov::no_label, 11, 12, 13};
+
+    auto cond_param =
+        make_shared<op::Parameter>(element::boolean, PartialShape{{2, 8}, {3, 7}, {1, 10}, {1, 6}, {1, 10}});
+    auto then_param = make_shared<op::Parameter>(element::f32, PartialShape::dynamic(5));
+    auto else_param = make_shared<op::Parameter>(element::f32, labeled_shape);
+    auto op = make_shared<op::v1::Select>(cond_param, then_param, else_param);
+
+    const auto& out_shape = op->get_output_partial_shape(0);
+
+    EXPECT_EQ(op->get_element_type(), element::f32);
+    EXPECT_EQ(out_shape, PartialShape({{2, 8}, {3, 7}, -1, 5, -1}));
+    EXPECT_EQ(get_shape_labels(out_shape), expected_labels);
+}
+
+TEST(type_prop, select_labels_all_params_numpy) {
+    auto labeled_shape_cond = PartialShape{-1, 2, 1, 3, 1, {2, 5}, {1, 8}, {5, -1}, {-1, 5}};
+    auto labeled_shape_then = PartialShape{-1, 2, 4, 1, 1, {1, 5}, {2, 8}, {5, -1}, {-1, 5}};
+    auto labeled_shape_else = PartialShape{-1, 2, 1, 3, 5, {1, 7}, {1, 8}, {5, -1}, {-1, 5}};
+
+    set_shape_labels(labeled_shape_cond, 10);
+    set_shape_labels(labeled_shape_then, 20);
+    set_shape_labels(labeled_shape_else, 30);
+
+    std::vector<size_t> expected_labels{10, 11, 22, 13, 34, 15, 26, 17, 18};
+
+    auto cond_param = make_shared<op::Parameter>(element::boolean, labeled_shape_cond);
+    auto then_param = make_shared<op::Parameter>(element::f32, labeled_shape_then);
+    auto else_param = make_shared<op::Parameter>(element::f32, labeled_shape_else);
+    auto op = make_shared<op::v1::Select>(cond_param, then_param, else_param);
+
+    const auto& out_shape = op->get_output_partial_shape(0);
+
+    EXPECT_EQ(op->get_element_type(), element::f32);
+    EXPECT_EQ(out_shape, (PartialShape{-1, 2, 4, 3, 5, {2, 5}, {2, 8}, {5, -1}, {-1, 5}}));
+    EXPECT_EQ(get_shape_labels(out_shape), expected_labels);
+}
+
+TEST(type_prop, select_labels_all_params_none) {
+    auto labeled_shape_cond = PartialShape{-1, 2, 4, 3, 5, {2, 5}, {2, 8}, {5, -1}, {-1, 5}};
+    auto labeled_shape_then = PartialShape{-1, 2, 4, 3, 5, {2, 5}, {2, 8}, {5, -1}, {-1, 5}};
+    auto labeled_shape_else = PartialShape{-1, 2, 4, 3, 5, {2, 5}, {2, 8}, {5, -1}, {-1, 5}};
+
+    set_shape_labels(labeled_shape_cond, 10);
+    set_shape_labels(labeled_shape_then, 20);
+    set_shape_labels(labeled_shape_else, 30);
+
+    std::vector<size_t> expected_labels{10, 11, 12, 13, 14, 15, 16, 17, 18};
+
+    auto cond_param = make_shared<op::Parameter>(element::boolean, labeled_shape_cond);
+    auto then_param = make_shared<op::Parameter>(element::f32, labeled_shape_then);
+    auto else_param = make_shared<op::Parameter>(element::f32, labeled_shape_else);
+    auto op = make_shared<op::v1::Select>(cond_param, then_param, else_param, op::AutoBroadcastType::NONE);
+
+    const auto& out_shape = op->get_output_partial_shape(0);
+
+    EXPECT_EQ(op->get_element_type(), element::f32);
+    EXPECT_EQ(out_shape, (PartialShape{-1, 2, 4, 3, 5, {2, 5}, {2, 8}, {5, -1}, {-1, 5}}));
+    EXPECT_EQ(get_shape_labels(out_shape), expected_labels);
+}
+
+TEST(type_prop, select_labels_all_params_pdpd) {
+    auto labeled_shape_cond = PartialShape{-1, -1, -1, -1, -1, -1, -1};
+    auto labeled_shape_then = PartialShape{-1, 2, 4, 3, 5, {2, 5}, {2, 8}, {5, -1}, {-1, 5}};
+    auto labeled_shape_else = PartialShape{-1, 2, 4, 3, 5, {2, 5}, {2, 8}, {5, -1}, {-1, 5}};
+
+    set_shape_labels(labeled_shape_cond, 10);
+    set_shape_labels(labeled_shape_then, 20);
+    set_shape_labels(labeled_shape_else, 30);
+
+    std::vector<size_t> expected_labels{20, 21, 22, 23, 24, 25, 26, 27, 28};
+
+    auto cond_param = make_shared<op::Parameter>(element::boolean, labeled_shape_cond);
+    auto then_param = make_shared<op::Parameter>(element::f32, labeled_shape_then);
+    auto else_param = make_shared<op::Parameter>(element::f32, labeled_shape_else);
+    auto op = make_shared<op::v1::Select>(cond_param,
+                                          then_param,
+                                          else_param,
+                                          op::AutoBroadcastSpec{op::AutoBroadcastType::PDPD});
+
+    const auto& out_shape = op->get_output_partial_shape(0);
+
+    EXPECT_EQ(op->get_element_type(), element::f32);
+    EXPECT_EQ(out_shape, (PartialShape{-1, 2, 4, 3, 5, {2, 5}, {2, 8}, {5, -1}, {-1, 5}}));
+    EXPECT_EQ(get_shape_labels(out_shape), expected_labels);
+}
+
 TEST(type_prop, select_dynamic) {
     auto param_0 =
         make_shared<op::Parameter>(element::boolean, PartialShape({{2, 8}, {3, 7}, {1, 10}, {1, 6}, {1, 10}}));
@@ -250,9 +398,9 @@ INSTANTIATE_TEST_SUITE_P(
                       SelectParams({{2}, {2, 4}, {2}, {2, 4}},
                                    {element::boolean, element::f32, element::dynamic, element::f32},
                                    {op::AutoBroadcastType::PDPD, 0}),
-                      // TODO: Whats the right behavior here?
-                      // SelectParams({{2}, {2, 4}, {2}, {2, 4}}, {element::boolean, element::f32,
-                      // element::dynamic, element::f32}, {op::AutoBroadcastType::PDPD, 0}),
+                      SelectParams({{2}, {2, 4}, {2}, {2, 4}},
+                                   {element::boolean, element::f32, element::dynamic, element::f32},
+                                   {op::AutoBroadcastType::PDPD, 0}),
                       SelectParams({{4}, {2, 4}, {2, 4}, {2, 4}},
                                    {element::boolean, element::f32, element::f32, element::f32},
                                    {op::AutoBroadcastType::PDPD, 1}),

src/plugins/intel_cpu/tests/unit/shape_inference_test/select_shape_inference_test.cpp

@@ -14,35 +14,35 @@ using namespace ov;
 using namespace ov::intel_cpu;
 
 TEST(StaticShapeInferenceTest, SelectTestBCastModeNUMPY) {
-    auto cond = std::make_shared<op::v0::Parameter>(element::boolean, PartialShape{});
-    auto ptrue = std::make_shared<op::v0::Parameter>(element::f32, PartialShape{});
-    auto pfalse = std::make_shared<op::v0::Parameter>(element::f32, PartialShape{});
+    auto cond = std::make_shared<op::v0::Parameter>(element::boolean, PartialShape::dynamic());
+    auto ptrue = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
+    auto pfalse = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
     auto select = std::make_shared<op::v1::Select>(cond, ptrue, pfalse, op::AutoBroadcastType::NUMPY);
     {
         std::vector<StaticShape> static_input_shapes = {StaticShape{}, StaticShape{4}, StaticShape{2, 4}},
                                  static_output_shapes = {StaticShape{}};
         shape_inference(select.get(), static_input_shapes, static_output_shapes);
-        ASSERT_EQ(static_output_shapes[0], StaticShape({2, 4}));
+        EXPECT_EQ(static_output_shapes[0], StaticShape({2, 4}));
     }
 
     {
         std::vector<StaticShape> static_input_shapes = {StaticShape{}, StaticShape{2, 4}, StaticShape{2, 4}},
                                  static_output_shapes = {StaticShape{}};
         shape_inference(select.get(), static_input_shapes, static_output_shapes);
-        ASSERT_EQ(static_output_shapes[0], StaticShape({2, 4}));
+        EXPECT_EQ(static_output_shapes[0], StaticShape({2, 4}));
     }
 
     {
         std::vector<StaticShape> static_input_shapes = {StaticShape{4}, StaticShape{2, 4}, StaticShape{4}},
                                  static_output_shapes = {StaticShape{}};
         shape_inference(select.get(), static_input_shapes, static_output_shapes);
-        ASSERT_EQ(static_output_shapes[0], StaticShape({2, 4}));
+        EXPECT_EQ(static_output_shapes[0], StaticShape({2, 4}));
     }
 }
 TEST(StaticShapeInferenceTest, SelectTestBCastModePDPD) {
-    auto cond = std::make_shared<op::v0::Parameter>(element::boolean, PartialShape{});
-    auto ptrue = std::make_shared<op::v0::Parameter>(element::f32, PartialShape{});
-    auto pfalse = std::make_shared<op::v0::Parameter>(element::f32, PartialShape{});
+    auto cond = std::make_shared<op::v0::Parameter>(element::boolean, PartialShape::dynamic());
+    auto ptrue = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
+    auto pfalse = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
     auto select =
         std::make_shared<op::v1::Select>(cond, ptrue, pfalse, op::AutoBroadcastSpec{op::AutoBroadcastType::PDPD, 1});
     std::vector<StaticShape> static_input_shapes = {StaticShape{4}, StaticShape{2, 4}, StaticShape{4}},
@@ -52,9 +52,9 @@ TEST(StaticShapeInferenceTest, SelectTestBCastModePDPD) {
 }
 
 TEST(StaticShapeInferenceTest, SelectTestBCastModeNone) {
-    auto cond = std::make_shared<op::v0::Parameter>(element::boolean, PartialShape{});
-    auto ptrue = std::make_shared<op::v0::Parameter>(element::f32, PartialShape{});
-    auto pfalse = std::make_shared<op::v0::Parameter>(element::f32, PartialShape{});
+    auto cond = std::make_shared<op::v0::Parameter>(element::boolean, PartialShape::dynamic());
+    auto ptrue = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
+    auto pfalse = std::make_shared<op::v0::Parameter>(element::f32, PartialShape::dynamic());
     auto select = std::make_shared<op::v1::Select>(cond, ptrue, pfalse, op::AutoBroadcastType::NONE);
 
     std::vector<StaticShape> static_input_shapes = {StaticShape{6, 4}, StaticShape{6, 4}, StaticShape{6, 4}},