IVGCVSW-7785 Extend support for 3D tensors BATCH_TO_SPACE and SPACE_TO_BATCH in CpuRef

* Both layers were assuming 4D tensors, now 3D is supported too.
* Remove some unnecessary includes
* Add Unit Tests

Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Change-Id: I7bdd11e4936a27cd97ec65fd915e6ccaa1494cff

1 files changed, 27 insertions, 0 deletions

diff --git a/src/backends/backendsCommon/test/layerTests/BatchToSpaceNdTestImpl.hpp b/src/backends/backendsCommon/test/layerTests/BatchToSpaceNdTestImpl.hpp
index b3007771c9..4f8b7d0193 100644
--- a/src/backends/backendsCommon/test/layerTests/BatchToSpaceNdTestImpl.hpp
+++ b/src/backends/backendsCommon/test/layerTests/BatchToSpaceNdTestImpl.hpp
@@ -279,6 +279,33 @@ LayerTestResult<T, 4> BatchToSpaceNdNhwcTest7(
 }
 
 template<armnn::DataType ArmnnType, typename T = armnn::ResolveType<ArmnnType>>
+LayerTestResult<T, 3> BatchToSpaceNdNhwcTest8(
+        armnn::IWorkloadFactory& workloadFactory,
+        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager,
+        const armnn::ITensorHandleFactory& tensorHandleFactory)
+{
+    const unsigned int inputShape[] = {4, 2, 1};
+    const unsigned int outputShape[] = {1, 8, 1};
+
+    std::vector<float> input({
+                                  1.0f,  2.0f,  3.0f,  4.0f,
+                                  5.0f,  6.0f,  7.0f,  8.0f
+                             });
+
+    std::vector<float> expectedOutput({
+                                           1.0f,  3.0f,  5.0f,  7.0f,
+                                           2.0f,  4.0f,  6.0f,  8.0f
+                                      });
+
+    std::vector<unsigned int> blockShape {4};
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}};
+
+    return BatchToSpaceNdHelper<ArmnnType, 3, 3>(workloadFactory, memoryManager, tensorHandleFactory,
+                                                 armnn::DataLayout::NHWC, inputShape, input, blockShape,
+                                                 crops, outputShape, expectedOutput);
+}
+
+template<armnn::DataType ArmnnType, typename T = armnn::ResolveType<ArmnnType>>
 LayerTestResult<T, 4> BatchToSpaceNdNchwTest1(
         armnn::IWorkloadFactory &workloadFactory,
         const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager,