diff --git a/ngraph/core/src/op/unsqueeze.cpp b/ngraph/core/src/op/unsqueeze.cpp index dd593d0920e..cc1cfa9da0a 100644 --- a/ngraph/core/src/op/unsqueeze.cpp +++ b/ngraph/core/src/op/unsqueeze.cpp @@ -34,23 +34,27 @@ void op::v0::Unsqueeze::validate_and_infer_types() const auto data_rank = data_partial_shape.rank(); const auto axes_constant = get_constant_from_source(input_value(1)); + auto axes_pshape = get_input_partial_shape(1); + + NODE_VALIDATION_CHECK(this, + axes_pshape.rank().compatible(0) || axes_pshape.rank().compatible(1), + "Second input (axes) should not be of rank higher than 1. Got: ", + axes_pshape.rank().get_length()); + if (data_rank.is_dynamic() || !axes_constant) { set_output_type(0, get_input_element_type(0), PartialShape::dynamic()); return; } - uint64_t data_rank_value = data_partial_shape.rank().get_length(); - - // Get value of axes from Constant const auto axes_values = axes_constant->cast_vector(); + uint64_t data_rank_value = data_partial_shape.rank().get_length(); const int64_t expanded_rank = data_rank_value + axes_values.size(); NODE_VALIDATION_CHECK(this, !axes_values.empty(), "'axes' input is mandatory"); auto normalized_axes = normalize_axes(this->description(), axes_values, expanded_rank); set axes(begin(normalized_axes), end(normalized_axes)); - vector output_shape{data_partial_shape}; for (auto axis : axes) { diff --git a/ngraph/test/type_prop/unsqueeze.cpp b/ngraph/test/type_prop/unsqueeze.cpp index 252e3139b73..c38ed797437 100644 --- a/ngraph/test/type_prop/unsqueeze.cpp +++ b/ngraph/test/type_prop/unsqueeze.cpp @@ -37,3 +37,54 @@ TEST(type_prop, unsqueeze_dynamic) Dimension::dynamic(), Dimension::dynamic()})); } + +TEST(type_prop, unsqueeze_incorrect_axes_shape) +{ + auto param = make_shared(element::f32, Shape{4, 1, 4, 1, 8}); + auto axes_node = + make_shared(element::u64, Shape{1, 1, 1}, vector{1}); + + try + { + auto unsqueeze = make_shared(param, axes_node); + FAIL() << "Unsqueeze axes invalid rank not detected"; + } + catch (const NodeValidationFailure& error) + { + EXPECT_HAS_SUBSTRING(error.what(), + "Second input (axes) should not be of rank higher than 1"); + } + catch (...) + { + FAIL() << "Deduced type check failed for unexpected reason"; + } +} + +TEST(type_prop, unsqueeze_empty_axes) +{ + auto param = make_shared(element::f32, Shape{4, 1, 4, 1, 8}); + auto axes_node = make_shared(element::u64, Shape{0}, vector{}); + try + { + auto unsqueeze = make_shared(param, axes_node); + FAIL() << "Unsqueeze axes empty not detected"; + } + catch (const NodeValidationFailure& error) + { + EXPECT_HAS_SUBSTRING(error.what(), "'axes' input is mandatory"); + } + catch (...) + { + FAIL() << "Deduced type check failed for unexpected reason"; + } +} + +TEST(type_prop, unsqueeze_dynamic_axes) +{ + auto param = make_shared(element::f32, Shape{4, 1, 4, 1, 8}); + auto axes_node = make_shared(element::u64, PartialShape::dynamic()); + + auto unsqueeze = make_shared(param, axes_node); + ASSERT_EQ(unsqueeze->get_element_type(), element::f32); + ASSERT_EQ(unsqueeze->get_output_partial_shape(0), PartialShape::dynamic()); +} \ No newline at end of file