Revise Unsqueeze op - op class (#5555)

* revise unsqueeze op class

* Added checks and tests for second input incorrect type and shape

* Remove axes type constraints to keep backward compatibility
This commit is contained in:
Bartosz Lesniewski 2021-05-13 06:45:29 +02:00 committed by GitHub
parent f928f7fc56
commit 39fde540d2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 59 additions and 4 deletions

View File

@ -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<int64_t>();
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<int64_t> axes(begin(normalized_axes), end(normalized_axes));
vector<Dimension> output_shape{data_partial_shape};
for (auto axis : axes)
{

View File

@ -37,3 +37,54 @@ TEST(type_prop, unsqueeze_dynamic)
Dimension::dynamic(),
Dimension::dynamic()}));
}
TEST(type_prop, unsqueeze_incorrect_axes_shape)
{
auto param = make_shared<op::Parameter>(element::f32, Shape{4, 1, 4, 1, 8});
auto axes_node =
make_shared<ngraph::op::Constant>(element::u64, Shape{1, 1, 1}, vector<int64_t>{1});
try
{
auto unsqueeze = make_shared<op::v0::Unsqueeze>(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<op::Parameter>(element::f32, Shape{4, 1, 4, 1, 8});
auto axes_node = make_shared<ngraph::op::Constant>(element::u64, Shape{0}, vector<int64_t>{});
try
{
auto unsqueeze = make_shared<op::v0::Unsqueeze>(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<op::Parameter>(element::f32, Shape{4, 1, 4, 1, 8});
auto axes_node = make_shared<ngraph::op::Parameter>(element::u64, PartialShape::dynamic());
auto unsqueeze = make_shared<op::v0::Unsqueeze>(param, axes_node);
ASSERT_EQ(unsqueeze->get_element_type(), element::f32);
ASSERT_EQ(unsqueeze->get_output_partial_shape(0), PartialShape::dynamic());
}