[SPEC] Add ND input data support for DepthToSpace and SpaceToDepth (#4003)

* First debug version of support ND input

* Added type_prop tests

* Implemented ND support for SpaceToDepth

* Added type_prop tests for SpaceToDepth

* Added fused op tests

* code refactor

* Fixed DepthToSpaceMode::BLOCKS_FIRST

* Code review remarks introduced

* Code review ramarks. Part.2

* Using NODE_VALIDATION_CHECK introduced

src/ngraph/op/fused/depth_to_space.cpp

@@ -13,6 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //*****************************************************************************
+#include <cmath>
 #include <cstddef>
 #include <memory>
 
@@ -46,59 +47,101 @@ op::DepthToSpace::DepthToSpace(const Output<Node>& data,
 NodeVector op::DepthToSpace::decompose_op() const
 {
     auto data = input_value(0);
-    const Shape& data_shape = data.get_shape();
+    auto data_shape = data.get_shape();
 
-    // Set default values to each dimension to be able to work with both 3D or 4D data.
-    size_t n{1}, c{1}, h{1}, w{1};
+    NODE_VALIDATION_CHECK(this,
+                          (data_shape.size() >= 3),
+                          "The input tensor with rank lower than 3 is not supported (input rank: ",
+                          data_shape.size(),
+                          ")");
 
-    NGRAPH_CHECK((data_shape.size() == 3 || data_shape.size() == 4),
-                 "The provided tensor shape: ",
-                 data_shape,
-                 " is not supported.");
-
-    // Assume NCHW data layout
-    if (data_shape.size() == 4)
-    {
-        n = data_shape.at(0);
-        c = data_shape.at(1);
-        h = data_shape.at(2);
-        w = data_shape.at(3);
-    }
-    // Without batch.
-    else if (data_shape.size() == 3)
+    if (data_shape.size() == 3)
     {
-        c = data_shape.at(0);
-        h = data_shape.at(1);
-        w = data_shape.at(2);
+        // Insert batch axis
+        data_shape.insert(data_shape.begin(), 1);
+        data = builder::reshape(data, data_shape);
     }
-
-    NGRAPH_CHECK((c % (m_blocksize * m_blocksize) == 0 && m_blocksize > 0),
-                 "SpaceToDepth: The depth axis size must be a multiple of ",
-                 "squared block_size attribute value.");
+    const size_t n_dim = data_shape.at(0);
+    const size_t c_dim = data_shape.at(1);
+    const size_t spatial_dim_index = 2;
+    const size_t spatial_dims = data_shape.size() - spatial_dim_index;
+    const auto c_dim_divider = static_cast<int>(std::pow(m_blocksize, spatial_dims));
+
+    NODE_VALIDATION_CHECK(this,
+                          m_blocksize > 0 && c_dim % c_dim_divider == 0,
+                          "DepthToSpace: The input data's 'channels' axis size: ",
+                          c_dim,
+                          " must be a equivalent to ",
+                          "'block_size'^'spatial_dims': ",
+                          c_dim_divider);
 
     auto bs = static_cast<size_t>(m_blocksize);
-    size_t c_flat = c / (bs * bs);
+    size_t c_flat = c_dim / c_dim_divider;
 
     // First we have to disperse the data from depth channel, then rearrange them
     // so as appropriate chunks of data where close to their destination place.
     // Finally squeeze data from respective dimensions.
     shared_ptr<Node> flat_node;
+    Shape dispersed_shape{n_dim};
+    for (int i = 0; i < spatial_dims; ++i)
+    {
+        dispersed_shape.push_back(bs);
+    }
+    for (int i = 0; i < spatial_dims; ++i)
+    {
+        dispersed_shape.push_back(data_shape.at(spatial_dim_index + i));
+    }
+    vector<size_t> axes_order{0};
     switch (m_mode)
     {
+    // x' = reshape(data, [N, C / (block_size ^ K), block_size, block_size, ..., block_size, D1, D2,
+    // ..., DK])
+    // x'' = transpose(x', [0,  1,  K + 2, 2, K + 3, 3, K + 4, 4, ..., K + (K + 1), K + 1])
+    // y = reshape(x'', [N, C / (block_size ^ K), D1 * block_size, D2 * block_size, D3 * block_size,
+    // ..., DK * block_size])
     case DepthToSpaceMode::DEPTH_FIRST:
     {
-        flat_node = builder::reshape(data, Shape{n, c_flat, bs, bs, h, w});
-        flat_node = builder::reorder_axes(flat_node, {0, 1, 4, 2, 5, 3});
+        dispersed_shape.insert(dispersed_shape.begin() + 1, c_flat);
+        flat_node = builder::reshape(data, dispersed_shape);
+
+        axes_order.push_back(1);
+        for (int i = spatial_dim_index; i < data_shape.size(); ++i)
+        {
+            axes_order.push_back(spatial_dims + i);
+            axes_order.push_back(i);
+        }
+
+        flat_node = builder::reorder_axes(flat_node, axes_order);
         break;
     }
+    // x' = reshape(data, [N, block_size, block_size, ..., block_size, C / (block_size ^ K), D1, D2,
+    // ..., DK])
+    // x'' = transpose(x', [0,  K + 1,  K + 2, 1, K + 3, 2, K + 4, 3, ..., K + (K + 1), K])
+    // y = reshape(x'', [N, C / (block_size ^ K), D1 * block_size, D2 * block_size, D3 * block_size,
+    // ..., DK * block_size])
     case DepthToSpaceMode::BLOCKS_FIRST:
     default:
     {
-        flat_node = builder::reshape(data, Shape{n, bs, bs, c_flat, h, w});
-        flat_node = builder::reorder_axes(flat_node, {0, 3, 4, 1, 5, 2});
+        dispersed_shape.insert(dispersed_shape.begin() + spatial_dims + 1, c_flat);
+        flat_node = builder::reshape(data, dispersed_shape);
+
+        axes_order.push_back(spatial_dims + 1);
+        for (int i = 2; i < data_shape.size(); ++i)
+        {
+            axes_order.push_back(spatial_dims + i);
+            axes_order.push_back(i - 1);
+        }
+        flat_node = builder::reorder_axes(flat_node, axes_order);
     }
     }
-    return NodeVector{builder::reshape(flat_node, Shape{n, c_flat, h * bs, w * bs})};
+    Shape squeezed_shape{n_dim, c_flat};
+    for (int i = spatial_dim_index; i < data_shape.size(); ++i)
+    {
+        squeezed_shape.push_back(data_shape.at(i) * bs);
+    }
+    flat_node = builder::reshape(flat_node, squeezed_shape);
+
+    return NodeVector{flat_node};
 }
 
 shared_ptr<Node> op::DepthToSpace::copy_with_new_args(const NodeVector& new_args) const

src/ngraph/op/fused/space_to_depth.cpp

@@ -13,6 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //*****************************************************************************
+#include <cmath>
 #include <cstddef>
 #include <memory>
 
@@ -43,57 +44,93 @@ op::SpaceToDepth::SpaceToDepth(const Output<Node>& data, const std::string& mode
 NodeVector op::SpaceToDepth::decompose_op() const
 {
     auto data = input_value(0);
-    const Shape& data_shape = data.get_shape();
+    auto data_shape = data.get_shape();
 
-    // Set default values to each dimension to be able to work with both 3D or 4D data.
-    size_t n{1}, c{1}, h{1}, w{1};
+    NODE_VALIDATION_CHECK(this,
+                          (data_shape.size() >= 3),
+                          "The input tensor with rank lower than 3 is not supported (input rank: ",
+                          data_shape.size(),
+                          ")");
 
-    NGRAPH_CHECK((data_shape.size() == 3 || data_shape.size() == 4),
-                 "The provided tensor shape: ",
-                 data_shape,
-                 " is not supported.");
+    NODE_VALIDATION_CHECK(this, m_blocksize > 0, "m_blocksize must be greater than 0");
 
-    // Assume NCHW data layout
-    if (data_shape.size() == 4)
+    if (data_shape.size() == 3)
     {
-        n = data_shape.at(0);
-        c = data_shape.at(1);
-        h = data_shape.at(2);
-        w = data_shape.at(3);
-    }
-    // Without batch.
-    else if (data_shape.size() == 3)
-    {
-        c = data_shape.at(0);
-        h = data_shape.at(1);
-        w = data_shape.at(2);
+        // Insert batch axis
+        data_shape.insert(data_shape.begin(), 1);
+        data = builder::reshape(data, data_shape);
     }
 
-    NGRAPH_CHECK((h % m_blocksize == 0 && w % m_blocksize == 0 && m_blocksize > 0),
-                 "SpaceToDepth: The width and height axes size must be a multiple of ",
-                 "squared block_size attribute value");
+    const size_t n_dim = data_shape.at(0);
+    const size_t c_dim = data_shape.at(1);
+    const size_t spatial_dim_index = 2;
+    const size_t spatial_dims = data_shape.size() - spatial_dim_index;
 
-    size_t bs = static_cast<size_t>(m_blocksize);
-    size_t w_flat = w / bs;
-    size_t h_flat = h / bs;
-    size_t c_high = c * bs * bs;
+    for (int i = spatial_dim_index; i < data_shape.size(); ++i)
+    {
+        NODE_VALIDATION_CHECK(this,
+                              m_blocksize > 0 && data_shape.at(i) % m_blocksize == 0,
+                              "The dimension on position: ",
+                              i,
+                              " equal to: ",
+                              data_shape.at(i),
+                              " must be a multiple of m_blocksize: ",
+                              m_blocksize);
+    }
 
-    // First we have to disperse the data from height and width channels, then
+    // First we have to disperse the data from spatial dimensions, then
     // rearrange them so as appropriate chunks of data where close to their
     // destination place. Finally squeeze data from respective dimensions.
-    Output<Node> flat_node = builder::reshape(data, Shape{n, c, h_flat, bs, w_flat, bs});
+    Shape dispersed_shape{n_dim, c_dim};
+    for (int i = 0; i < spatial_dims; ++i)
+    {
+        dispersed_shape.push_back(data_shape.at(i + spatial_dim_index) / m_blocksize);
+        dispersed_shape.push_back(m_blocksize);
+    }
+    auto flat_node = builder::reshape(data, dispersed_shape);
+    // calculate axes to transpose
+    // [0, 3, 5, ..., spatial_dims + (spatial_dims + 1), 2, 4, ..., K + K])
+    vector<size_t> axes_order{0};
+    for (size_t i = 0, j = 3; i < spatial_dims; ++i, j += 2)
+    {
+        axes_order.push_back(j);
+    }
+    for (size_t i = 0, j = 2; i < spatial_dims; ++i, j += 2)
+    {
+        axes_order.push_back(j);
+    }
+
     switch (m_mode)
     {
+    // x' = reshape(data, [N, C, D1/block_size, block_size, D2/block_size, block_size, ...,
+    // DK/block_size, block_size])
+    // x'' = transpose(x', [0,  1, 3, 5, ..., K + (K + 1),  2, 4, ..., K + K])
+    // y = reshape(x'', [N, C * (block_size ^ K), D1 / block_size, D2 / block_size, ..., DK /
+    // block_size])
     case SpaceToDepthMode::DEPTH_FIRST:
     {
-        flat_node = builder::reorder_axes(flat_node, {0, 1, 3, 5, 2, 4});
+        axes_order.insert(axes_order.begin() + 1, 1);
         break;
     }
+    // x' = reshape(data, [N, C, D1/block_size, block_size, D2/block_size, block_size, ... ,
+    // DK/block_size, block_size])
+    // x'' = transpose(x',  [0,  3, 5, ..., K + (K + 1), 1,  2, 4, ..., K + K])
+    // y = reshape(x'', [N, C * (block_size ^ K), D1 / block_size, D2 / block_size, ..., DK /
+    // block_size])
     case SpaceToDepthMode::BLOCKS_FIRST:
-    default: { flat_node = builder::reorder_axes(flat_node, {0, 3, 5, 1, 2, 4});
+    default: { axes_order.insert(axes_order.begin() + spatial_dims + 1, 1);
     }
     }
-    return NodeVector{builder::reshape(flat_node, Shape{n, c_high, h_flat, w_flat})};
+    flat_node = builder::reorder_axes(flat_node, axes_order);
+    Shape squeezed_shape{n_dim};
+    for (int i = 0; i < spatial_dims; ++i)
+    {
+        squeezed_shape.push_back(data_shape.at(spatial_dim_index + i) / m_blocksize);
+    }
+    squeezed_shape.insert(squeezed_shape.begin() + 1, c_dim * std::pow(m_blocksize, spatial_dims));
+    flat_node = builder::reshape(flat_node, squeezed_shape);
+
+    return NodeVector{flat_node};
 }
 
 shared_ptr<Node> op::SpaceToDepth::copy_with_new_args(const NodeVector& new_args) const

test/backend/fused_op.in.cpp

@@ -2649,3 +2649,83 @@ NGRAPH_TEST(${BACKEND_NAME}, cross_entropy_with_one_hot)
     auto result = read_vector<float>(result0);
     EXPECT_TRUE(test::all_close_f(result, expected, 23));
 }
+
+NGRAPH_TEST(${BACKEND_NAME}, depth_to_space_space_to_depth_block_first)
+{
+    auto backend = runtime::Backend::create("${BACKEND_NAME}");
+    Shape dts_input_shape{2, 32, 2, 4, 2, 4};
+    size_t block_size = 2;
+
+    auto dts_input = make_shared<op::Parameter>(element::f32, dts_input_shape);
+    auto depth_to_space = make_shared<op::DepthToSpace>(
+        dts_input, op::DepthToSpace::DepthToSpaceMode::BLOCKS_FIRST, block_size);
+    auto dts_func = make_shared<Function>(NodeVector{depth_to_space}, ParameterVector{dts_input});
+
+    auto dts_input_tensor = backend->create_tensor(element::f32, dts_input_shape);
+    const auto data_size = shape_size(dts_input_shape);
+    vector<float> data(data_size);
+    std::iota(data.begin(), data.end(), 0);
+    copy_data(dts_input_tensor, data);
+    const auto dts_output_shape = depth_to_space->get_output_shape(0);
+    auto dts_output_tensor = backend->create_tensor(element::f32, dts_output_shape);
+    auto handle = backend->compile(dts_func);
+    handle->call_with_validate({dts_output_tensor}, {dts_input_tensor});
+    auto dts_result = read_vector<float>(dts_output_tensor);
+
+    // use depth_to_space output as space_to_depth input
+    auto std_input = make_shared<op::Parameter>(element::f32, dts_output_shape);
+    auto space_to_depth = make_shared<op::SpaceToDepth>(
+        std_input, op::SpaceToDepth::SpaceToDepthMode::BLOCKS_FIRST, block_size);
+    auto std_func = make_shared<Function>(NodeVector{space_to_depth}, ParameterVector{std_input});
+
+    auto std_input_tensor = backend->create_tensor(element::f32, dts_output_shape);
+    copy_data(std_input_tensor, dts_result);
+    auto std_output_tensor = backend->create_tensor(element::f32, dts_input_shape);
+    handle = backend->compile(std_func);
+    handle->call_with_validate({std_output_tensor}, {std_input_tensor});
+    auto std_result = read_vector<float>(std_output_tensor);
+
+    // expected output of space_to_depth is input of depth_to_space
+    ASSERT_EQ(dts_input_shape, space_to_depth->get_output_shape(0));
+    EXPECT_TRUE(test::all_close_f(std_result, data, data_size));
+}
+
+NGRAPH_TEST(${BACKEND_NAME}, depth_to_space_space_to_depth_depth_first)
+{
+    auto backend = runtime::Backend::create("${BACKEND_NAME}");
+    Shape dts_input_shape{2, 32, 2, 4, 2, 4};
+    size_t block_size = 2;
+
+    auto dts_input = make_shared<op::Parameter>(element::f32, dts_input_shape);
+    auto depth_to_space = make_shared<op::DepthToSpace>(
+        dts_input, op::DepthToSpace::DepthToSpaceMode::DEPTH_FIRST, block_size);
+    auto dts_func = make_shared<Function>(NodeVector{depth_to_space}, ParameterVector{dts_input});
+
+    auto dts_input_tensor = backend->create_tensor(element::f32, dts_input_shape);
+    const auto data_size = shape_size(dts_input_shape);
+    vector<float> data(data_size);
+    std::iota(data.begin(), data.end(), 0);
+    copy_data(dts_input_tensor, data);
+    const auto dts_output_shape = depth_to_space->get_output_shape(0);
+    auto dts_output_tensor = backend->create_tensor(element::f32, dts_output_shape);
+    auto handle = backend->compile(dts_func);
+    handle->call_with_validate({dts_output_tensor}, {dts_input_tensor});
+    auto dts_result = read_vector<float>(dts_output_tensor);
+
+    // use depth_to_space output as space_to_depth input
+    auto std_input = make_shared<op::Parameter>(element::f32, dts_output_shape);
+    auto space_to_depth = make_shared<op::SpaceToDepth>(
+        std_input, op::SpaceToDepth::SpaceToDepthMode::DEPTH_FIRST, block_size);
+    auto std_func = make_shared<Function>(NodeVector{space_to_depth}, ParameterVector{std_input});
+
+    auto std_input_tensor = backend->create_tensor(element::f32, dts_output_shape);
+    copy_data(std_input_tensor, dts_result);
+    auto std_output_tensor = backend->create_tensor(element::f32, dts_input_shape);
+    handle = backend->compile(std_func);
+    handle->call_with_validate({std_output_tensor}, {std_input_tensor});
+    auto std_result = read_vector<float>(std_output_tensor);
+
+    // expected output of space_to_depth is input of depth_to_space
+    ASSERT_EQ(dts_input_shape, space_to_depth->get_output_shape(0));
+    EXPECT_TRUE(test::all_close_f(std_result, data, data_size));
+}

test/type_prop/depth_to_space.cpp

@@ -21,7 +21,7 @@
 using namespace std;
 using namespace ngraph;
 
-TEST(type_prop, depth_to_space)
+TEST(type_prop, depth_to_space_output_shape_block_first_4D)
 {
     auto A = make_shared<op::Parameter>(element::f32, Shape{1, 128, 8, 8});
     auto space_to_depth =
@@ -31,9 +31,49 @@ TEST(type_prop, depth_to_space)
     ASSERT_EQ(space_to_depth->get_shape(), (Shape{1, 2, 64, 64}));
 }
 
+TEST(type_prop, depth_to_space_output_shape_block_first_4D_2)
+{
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 12, 1080, 1616});
+    auto space_to_depth =
+        make_shared<op::DepthToSpace>(A, op::DepthToSpace::DepthToSpaceMode::BLOCKS_FIRST, 2);
+
+    ASSERT_EQ(space_to_depth->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_depth->get_shape(), (Shape{1, 3, 2 * 1080, 2 * 1616}));
+}
+
+TEST(type_prop, depth_to_space_output_shape_block_first_5D)
+{
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 16, 3, 1080, 1616});
+    auto space_to_depth =
+        make_shared<op::DepthToSpace>(A, op::DepthToSpace::DepthToSpaceMode::BLOCKS_FIRST, 2);
+
+    ASSERT_EQ(space_to_depth->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_depth->get_shape(), (Shape{1, 2, 2 * 3, 2 * 1080, 2 * 1616}));
+}
+
+TEST(type_prop, depth_to_space_output_shape_depth_first_4D)
+{
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 12, 1080, 1616});
+    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_shape(), (Shape{1, 3, 2 * 1080, 2 * 1616}));
+}
+
+TEST(type_prop, depth_to_space_output_shape_depth_first_5D)
+{
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 16, 3, 1080, 1616});
+    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_shape(), (Shape{1, 2, 2 * 3, 2 * 1080, 2 * 1616}));
+}
+
 TEST(type_prop, depth_to_space_input_rank_not_supported)
 {
-    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 8, 8, 8, 4});
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 8});
     try
     {
         auto space_to_depth =
@@ -42,7 +82,30 @@ TEST(type_prop, depth_to_space_input_rank_not_supported)
     }
     catch (const ngraph_error& error)
     {
-        EXPECT_HAS_SUBSTRING(error.what(), "The provided tensor shape: ");
+        EXPECT_HAS_SUBSTRING(
+            error.what(),
+            "The input tensor with rank lower than 3 is not supported (input rank: 2)");
+    }
+    catch (...)
+    {
+        FAIL() << "DepthToSpace decomposition failed for unexpected reason";
+    }
+}
+
+TEST(type_prop, depth_to_space_blocksize_not_matched)
+{
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 7, 4, 4});
+    try
+    {
+        auto space_to_depth =
+            make_shared<op::DepthToSpace>(A, op::DepthToSpace::DepthToSpaceMode::DEPTH_FIRST, 2);
+        FAIL() << "Not matched blocksize for DepthToSpace exception not thrown";
+    }
+    catch (const ngraph_error& error)
+    {
+        EXPECT_HAS_SUBSTRING(error.what(),
+                             "DepthToSpace: The input data's 'channels' axis size: 7"
+                             " must be a equivalent to 'block_size'^'spatial_dims': 4");
     }
     catch (...)
     {

test/type_prop/space_to_depth.cpp

@@ -21,7 +21,7 @@
 using namespace std;
 using namespace ngraph;
 
-TEST(type_prop, space_to_depth)
+TEST(type_prop, space_to_depth_output_shape_block_first_4D)
 {
     auto A = make_shared<op::Parameter>(element::f32, Shape{1, 2, 64, 64});
     const auto mode = ngraph::op::SpaceToDepth::SpaceToDepthMode::BLOCKS_FIRST;
@@ -31,18 +31,71 @@ TEST(type_prop, space_to_depth)
     ASSERT_EQ(space_to_depth->get_shape(), (Shape{1, 128, 8, 8}));
 }
 
+TEST(type_prop, space_to_depth_output_shape_block_first_4D_2)
+{
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 12, 1080, 1616});
+    const auto mode = ngraph::op::SpaceToDepth::SpaceToDepthMode::BLOCKS_FIRST;
+    auto space_to_depth = make_shared<op::SpaceToDepth>(A, mode, 2);
+
+    ASSERT_EQ(space_to_depth->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_depth->get_shape(), (Shape{1, 12 * 4, 1080 / 2, 1616 / 2}));
+}
+
+TEST(type_prop, space_to_depth_output_shape_depth_first_4D)
+{
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 12, 1080, 1616});
+    const auto mode = ngraph::op::SpaceToDepth::SpaceToDepthMode::DEPTH_FIRST;
+    auto space_to_depth = make_shared<op::SpaceToDepth>(A, mode, 2);
+
+    ASSERT_EQ(space_to_depth->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_depth->get_shape(), (Shape{1, 12 * 4, 1080 / 2, 1616 / 2}));
+}
+
+TEST(type_prop, space_to_depth_output_shape_depth_first_5D)
+{
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 12, 4, 1080, 1616});
+    const auto mode = ngraph::op::SpaceToDepth::SpaceToDepthMode::DEPTH_FIRST;
+    auto space_to_depth = make_shared<op::SpaceToDepth>(A, mode, 2);
+
+    ASSERT_EQ(space_to_depth->get_element_type(), element::f32);
+    ASSERT_EQ(space_to_depth->get_shape(), (Shape{1, 12 * 8, 4 / 2, 1080 / 2, 1616 / 2}));
+}
+
 TEST(type_prop, space_to_depth_input_rank_not_supported)
 {
-    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 8, 8, 8, 4});
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 8});
     try
     {
         auto space_to_depth =
-            make_shared<op::DepthToSpace>(A, op::DepthToSpace::DepthToSpaceMode::DEPTH_FIRST, 2);
+            make_shared<op::SpaceToDepth>(A, op::SpaceToDepth::SpaceToDepthMode::DEPTH_FIRST, 2);
         FAIL() << "Not supported input shape for SpaceToDepth exception not thrown";
     }
     catch (const ngraph_error& error)
     {
-        EXPECT_HAS_SUBSTRING(error.what(), "The provided tensor shape: ");
+        EXPECT_HAS_SUBSTRING(
+            error.what(),
+            "The input tensor with rank lower than 3 is not supported (input rank: 2)");
+    }
+    catch (...)
+    {
+        FAIL() << "SpaceToDepth decomposition failed for unexpected reason";
+    }
+}
+
+TEST(type_prop, space_to_depth_blocksize_not_matched)
+{
+    auto A = make_shared<op::Parameter>(element::f32, Shape{1, 3, 8, 7});
+    try
+    {
+        auto space_to_depth =
+            make_shared<op::SpaceToDepth>(A, op::SpaceToDepth::SpaceToDepthMode::DEPTH_FIRST, 4);
+        FAIL() << "Not matched blocksize SpaceToDepth exception not thrown";
+    }
+    catch (const ngraph_error& error)
+    {
+        EXPECT_HAS_SUBSTRING(
+            error.what(),
+            "The dimension on position: 3 equal to: 7 must be a multiple of m_blocksize: 4");
     }
     catch (...)
     {