IVGCVSW-3258 Add front end support for new SpaceToDepth layer

Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
Change-Id: Id677e29a734f2b36483d939ad370079bdc11551e

1 files changed, 26 insertions, 3 deletions

diff --git a/src/armnn/test/LayerValidateOutputTest.cpp b/src/armnn/test/LayerValidateOutputTest.cpp
index 999844e252..acefd51110 100644
--- a/src/armnn/test/LayerValidateOutputTest.cpp
+++ b/src/armnn/test/LayerValidateOutputTest.cpp
@@ -4,11 +4,12 @@
 //
 #include <armnn/ArmNN.hpp>
 
-#include <boost/algorithm/string.hpp>
-#include <boost/test/unit_test.hpp>
-#include <layers/BatchToSpaceNdLayer.hpp>
 #include <Graph.hpp>
+#include <layers/BatchToSpaceNdLayer.hpp>
+#include <layers/SpaceToDepthLayer.hpp>
 
+#include <boost/algorithm/string.hpp>
+#include <boost/test/unit_test.hpp>
 
 BOOST_AUTO_TEST_SUITE(LayerValidateOutput)
 
@@ -35,4 +36,26 @@ BOOST_AUTO_TEST_CASE(TestBatchToSpaceInferOutputShape)
     BOOST_CHECK(expectedShape == batchToSpaceLayer->InferOutputShapes(shapes).at(0));
 }
 
+BOOST_AUTO_TEST_CASE(TestSpaceToDepthInferOutputShape)
+{
+    armnn::Graph graph;
+
+    armnn::SpaceToDepthDescriptor descriptor;
+    descriptor.m_BlockSize  = 2;
+    descriptor.m_DataLayout = armnn::DataLayout::NHWC;
+
+    armnn::SpaceToDepthLayer* const spaceToDepthLayer =
+        graph.AddLayer<armnn::SpaceToDepthLayer>(descriptor, "spaceToDepth");
+
+    std::vector<armnn::TensorShape> shapes;
+    const std::vector<unsigned int> dimSizes{ 1, 16, 8, 3 };
+    armnn::TensorShape shape(4, dimSizes.data());
+    shapes.push_back(shape);
+
+    const std::vector<unsigned int> expectedDimSizes{ 1, 8, 4, 12 };
+    armnn::TensorShape expectedShape(4, expectedDimSizes.data());
+
+    BOOST_CHECK(expectedShape == spaceToDepthLayer->InferOutputShapes(shapes).at(0));
+}
+
 BOOST_AUTO_TEST_SUITE_END()