IVGCVSW-2763 Fix bug in BatchToSpaceNdLayer::InferOutputShapes

*Also added test case for non-zero crop and moved BatchToSpaceNd test 
functions together

Change-Id: I142ba356165618b2811a4ab650ca6ced35220d9c
Signed-off-by: Nattapat Chaimanowong <nattapat.chaimanowong@arm.com>

5 files changed, 285 insertions, 247 deletions

diff --git a/src/armnn/layers/BatchToSpaceNdLayer.cpp b/src/armnn/layers/BatchToSpaceNdLayer.cpp
index 9accf28137..a168fe8bbd 100644
--- a/src/armnn/layers/BatchToSpaceNdLayer.cpp
+++ b/src/armnn/layers/BatchToSpaceNdLayer.cpp
@@ -16,6 +16,8 @@
 
 #include <DataLayoutIndexed.hpp>
 
+#include <numeric>
+
 using namespace armnnUtils;
 
 namespace armnn
@@ -55,68 +57,40 @@ void BatchToSpaceNdLayer::ValidateTensorShapesFromInputs()
 
 std::vector<TensorShape> BatchToSpaceNdLayer::InferOutputShapes(const std::vector<TensorShape>& inputShapes) const
 {
-    const DataLayoutIndexed dataLayout = m_Param.m_DataLayout;
-    const TensorShape& inputShape = inputShapes[0];
-    unsigned int inBatchSize = inputShape[0];
-    unsigned int channelSize = inputShape[dataLayout.GetChannelsIndex()];
-
-    std::vector<unsigned int> theBlockShape = m_Param.m_BlockShape;
-
-    unsigned int overallSize = inBatchSize * inputShape[dataLayout.GetHeightIndex()]
-                               * inputShape[dataLayout.GetWidthIndex()];
+    BOOST_ASSERT(inputShapes.size() == 1);
 
-    std::vector<std::pair<unsigned int, unsigned int>> crops = m_Param.m_Crops;
-
-    std::pair<unsigned int, unsigned int> yCrops = crops[0];
-    std::pair<unsigned int, unsigned int> xCrops = crops[1];
-
-    unsigned int inputHeight = inputShape[dataLayout.GetHeightIndex()];
-    unsigned int outputHeight;
+    const TensorShape& inputShape = inputShapes[0];
+    TensorShape outputShape(inputShape);
 
-    unsigned int yCropsTotal = yCrops.first + yCrops.second;
+    unsigned int accumulatedBlockShape = std::accumulate(m_Param.m_BlockShape.begin(),
+                                                         m_Param.m_BlockShape.end(),
+                                                         1U,
+                                                         std::multiplies<>());
 
-    BOOST_ASSERT_MSG(yCropsTotal <= inputHeight,
-                     "BatchToSpaceLayer: Overall height crop should be less than or equal to the input height.");
+    BOOST_ASSERT(inputShape[0] % accumulatedBlockShape == 0);
 
-    unsigned int croppedHeight = inputHeight - yCropsTotal;
+    outputShape[0] = inputShape[0] / accumulatedBlockShape;
 
-    if (theBlockShape.at(0) > 0)
-    {
-        outputHeight = theBlockShape.at(0) * croppedHeight;
-    }
-    else
-    {
-        outputHeight = croppedHeight;
-    }
+    DataLayoutIndexed dimensionIndices = m_Param.m_DataLayout;
+    unsigned int heightIndex = dimensionIndices.GetHeightIndex();
+    unsigned int widthIndex = dimensionIndices.GetWidthIndex();
 
-    unsigned int outputWidth;
-    unsigned int inputWidth = inputShape[dataLayout.GetWidthIndex()];
+    unsigned int heightCrop = m_Param.m_Crops[0].first + m_Param.m_Crops[0].second;
+    unsigned int widthCrop = m_Param.m_Crops[1].first + m_Param.m_Crops[1].second;
 
-    unsigned int xCropsTotal = xCrops.first + xCrops.second;
+    unsigned int outputHeight = inputShape[heightIndex] * m_Param.m_BlockShape[0];
+    unsigned int outputWidth = inputShape[widthIndex] * m_Param.m_BlockShape[1];
 
-    BOOST_ASSERT_MSG(xCropsTotal <= inputWidth,
-                     "BatchToSpaceLayer: Overall width crop should be less than or equal to the input width.");
-    unsigned int croppedWidth = inputWidth - xCropsTotal;
+    BOOST_ASSERT_MSG(heightCrop <= outputHeight,
+        "BatchToSpaceLayer: Overall height crop should be less than or equal to the uncropped output height.");
 
-    if (theBlockShape.at(1) > 0)
-    {
-        outputWidth = theBlockShape.at(1) * croppedWidth;
-    }
-    else
-    {
-        outputWidth = croppedWidth;
-    }
+    BOOST_ASSERT_MSG(widthCrop <= outputWidth,
+        "BatchToSpaceLayer: Overall width crop should be less than or equal to the uncropped output width.");
 
-    unsigned int outputBatchSize = overallSize / (outputHeight * outputWidth);
+    outputShape[heightIndex] = outputHeight - heightCrop;
+    outputShape[widthIndex] = outputWidth - widthCrop;
 
-    if (dataLayout == DataLayout::NHWC)
-    {
-        return std::vector<TensorShape>({ TensorShape({ outputBatchSize, outputHeight, outputWidth, channelSize }) });
-    }
-    else
-    {
-        return std::vector<TensorShape>({ TensorShape({ outputBatchSize, channelSize, outputHeight, outputWidth }) });
-    }
+    return std::vector<TensorShape>({ outputShape });
 }
 
 void BatchToSpaceNdLayer::Accept(ILayerVisitor& visitor) const
diff --git a/src/armnn/test/LayerValidateOutputTest.cpp b/src/armnn/test/LayerValidateOutputTest.cpp
index 62b9c4a0d8..999844e252 100644
--- a/src/armnn/test/LayerValidateOutputTest.cpp
+++ b/src/armnn/test/LayerValidateOutputTest.cpp
@@ -17,22 +17,22 @@ BOOST_AUTO_TEST_CASE(TestBatchToSpaceInferOutputShape)
     armnn::Graph graph;
 
     armnn::BatchToSpaceNdDescriptor descriptor;
-    std::vector<unsigned int> theBlockShape = {2, 2};
-    descriptor.m_BlockShape = theBlockShape;
+    descriptor.m_BlockShape = {2, 2};
+    descriptor.m_Crops = {{0, 0}, {2, 0}};
     descriptor.m_DataLayout = armnn::DataLayout::NHWC;
 
     armnn::BatchToSpaceNdLayer* const batchToSpaceLayer =
         graph.AddLayer<armnn::BatchToSpaceNdLayer>(descriptor, "batchToSpace");
 
     std::vector<armnn::TensorShape> shapes;
-    const std::vector<unsigned int> theDimSizes = {4, 2, 2, 1};
+    const std::vector<unsigned int> theDimSizes = {8, 1, 3, 1};
     armnn::TensorShape shape(4, theDimSizes.data());
     shapes.push_back(shape);
 
-    const std::vector<unsigned int> expectedDimSizes = {1, 4, 4, 1};
+    const std::vector<unsigned int> expectedDimSizes = {2, 2, 4, 1};
     armnn::TensorShape expectedShape(4, expectedDimSizes.data());
 
     BOOST_CHECK(expectedShape == batchToSpaceLayer->InferOutputShapes(shapes).at(0));
 }
 
-BOOST_AUTO_TEST_SUITE_END()
\ No newline at end of file
+BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/backends/backendsCommon/test/LayerTests.cpp b/src/backends/backendsCommon/test/LayerTests.cpp
index a6b3b3d5bb..ce02fedb98 100644
--- a/src/backends/backendsCommon/test/LayerTests.cpp
+++ b/src/backends/backendsCommon/test/LayerTests.cpp
@@ -8342,10 +8342,9 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test1(
     const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
 {
     const unsigned int inputShape[] = {4, 2, 2, 1};
-    const unsigned int outputShape[] = {1, 4, 4, 1 };
+    const unsigned int outputShape[] = {1, 4, 4, 1};
 
-    std::vector<float> input
-    ({
+    std::vector<float> input({
         // Batch 0, Height 0, Width (2) x Channel (1)
         1.0f, 3.0f,
         // Batch 0, Height 1, Width (2) x Channel (1)
@@ -8369,8 +8368,7 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test1(
         14.0f, 16.0f
     });
 
-    std::vector<float> expectedOutput
-    ({
+    std::vector<float> expectedOutput({
         1.0f,   2.0f,  3.0f,  4.0f,
         5.0f,   6.0f,  7.0f,  8.0f,
         9.0f,  10.0f, 11.0f,  12.0f,
@@ -8392,13 +8390,12 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test2(
     const unsigned int inputShape[] = {4, 1, 1, 1};
     const unsigned int outputShape[] = {1, 2, 2, 1};
 
-    std::vector<float> input
-    ({
-         // Batch 0, Height 0, Width (2) x Channel (1)
-         1.0f, 2.0f, 3.0f, 4.0f
+    std::vector<float> input({
+        // Batch 0, Height 0, Width (2) x Channel (1)
+        1.0f, 2.0f, 3.0f, 4.0f
     });
 
-    std::vector<float> expectedOutput({1.0f,   2.0f,  3.0f,  4.0f});
+    std::vector<float> expectedOutput({1.0f, 2.0f, 3.0f, 4.0f});
 
     std::vector<unsigned int> blockShape({2, 2});
     std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
@@ -8415,9 +8412,9 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test3(
     const unsigned int inputShape[] = {4, 1, 1, 3};
     const unsigned int outputShape[] = {1, 2, 2, 3};
 
-    std::vector<float> input({ 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f });
+    std::vector<float> input({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f});
 
-    std::vector<float> expectedOutput({ 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f });
+    std::vector<float> expectedOutput({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f});
 
     std::vector<unsigned int> blockShape({2, 2});
     std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
@@ -8427,6 +8424,39 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test3(
         crops, outputShape, expectedOutput);
 }
 
+LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test4(
+    armnn::IWorkloadFactory& workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
+{
+    const unsigned int inputShape[] = {8, 1, 3, 1};
+    const unsigned int outputShape[] = {2, 2, 4, 1};
+
+    std::vector<float> input({
+        0.0f, 1.0f, 3.0f,
+        0.0f, 9.0f, 11.0f,
+        0.0f, 2.0f, 4.0f,
+        0.0f, 10.0f, 12.0f,
+        0.0f, 5.0f, 7.0f,
+        0.0f, 13.0f, 15.0f,
+        0.0f, 6.0f, 8.0f,
+        0.0f, 14.0f, 16.0f
+    });
+
+    std::vector<float> expectedOutput({
+        1.0f, 2.0f, 3.0f, 4.0f,
+        5.0f, 6.0f, 7.0f, 8.0f,
+        9.0f, 10.0f, 11.0f, 12.0f,
+        13.0f, 14.0f, 15.0f, 16.0f
+    });
+
+    std::vector<unsigned int> blockShape({2, 2});
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {2, 0}};
+
+    return BatchToSpaceNdHelper<float, 4, 4>(workloadFactory, memoryManager,
+        armnn::DataLayout::NHWC, inputShape, input, blockShape,
+        crops, outputShape, expectedOutput);
+}
+
 LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test1(
     armnn::IWorkloadFactory &workloadFactory,
     const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
@@ -8434,21 +8464,20 @@ LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test1(
     const unsigned int inputShape[] = {4, 3, 1, 1};
     const unsigned int outputShape[] = {1, 3, 2, 2};
 
-    std::vector<float> input({ 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f });
+    std::vector<float> input({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f});
 
-    std::vector<float> expectedOutput
-    ({
-         // Batch 0, Channel 0, Height (2) x Width (2)
-         1.0f,  4.0f,
-         7.0f, 10.0f,
+    std::vector<float> expectedOutput({
+        // Batch 0, Channel 0, Height (2) x Width (2)
+        1.0f,  4.0f,
+        7.0f, 10.0f,
 
-         // Batch 0, Channel 1, Height (2) x Width (2)
-         2.0f,  5.0f,
-         8.0f, 11.0f,
+        // Batch 0, Channel 1, Height (2) x Width (2)
+        2.0f,  5.0f,
+        8.0f, 11.0f,
 
-         // Batch 0, Channel 2, Height (2) x Width (2)
-         3.0f,  6.0f,
-         9.0f, 12.0f,
+        // Batch 0, Channel 2, Height (2) x Width (2)
+        3.0f,  6.0f,
+        9.0f, 12.0f,
     });
 
     std::vector<unsigned int> blockShape({2, 2});
@@ -8460,19 +8489,18 @@ LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test1(
 }
 
 LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test2(
-        armnn::IWorkloadFactory& workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
+    armnn::IWorkloadFactory& workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
 {
     const unsigned int inputShape[] = {4, 1, 1, 1};
     const unsigned int outputShape[] = {1, 1, 2, 2};
 
-    std::vector<float> input
-            ({
-                     // Batch 0, Height 0, Width (2) x Channel (1)
-                     1.0f, 2.0f, 3.0f, 4.0f
-             });
+    std::vector<float> input({
+        // Batch 0, Height 0, Width (2) x Channel (1)
+        1.0f, 2.0f, 3.0f, 4.0f
+    });
 
-    std::vector<float> expectedOutput({1.0f,   2.0f,  3.0f,  4.0f});
+    std::vector<float> expectedOutput({1.0f, 2.0f, 3.0f, 4.0f});
 
     std::vector<unsigned int> blockShape({2, 2});
     std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
@@ -8483,28 +8511,27 @@ LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test2(
 }
 
 LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test3(
-        armnn::IWorkloadFactory& workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
+    armnn::IWorkloadFactory& workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
 {
     const unsigned int inputShape[] = {4, 3, 1, 1};
     const unsigned int outputShape[] = {1, 3, 2, 2};
 
-    std::vector<float> input({ 1.0f, 3.0f, 5.0f, 7.0f, 9.0f, 11.0f, 2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f });
+    std::vector<float> input({1.0f, 3.0f, 5.0f, 7.0f, 9.0f, 11.0f, 2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f});
 
-    std::vector<float> expectedOutput
-            ({
-                     // Batch 0, Channel 0, Height (2) x Width (2)
-                     1.0f,  7.0f,
-                     2.0f,  8.0f,
+    std::vector<float> expectedOutput({
+        // Batch 0, Channel 0, Height (2) x Width (2)
+        1.0f,  7.0f,
+        2.0f,  8.0f,
 
-                     // Batch 0, Channel 1, Height (2) x Width (2)
-                     3.0f,  9.0f,
-                     4.0f, 10.0f,
+        // Batch 0, Channel 1, Height (2) x Width (2)
+        3.0f,  9.0f,
+        4.0f, 10.0f,
 
-                     // Batch 0, Channel 2, Height (2) x Width (2)
-                     5.0f, 11.0f,
-                     6.0f, 12.0f,
-             });
+        // Batch 0, Channel 2, Height (2) x Width (2)
+        5.0f, 11.0f,
+        6.0f, 12.0f,
+    });
 
     std::vector<unsigned int> blockShape({2, 2});
     std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
@@ -8521,8 +8548,8 @@ LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest1(
     const unsigned int inputShape[] = {4, 2, 2, 1};
     const unsigned int outputShape[] = {1, 4, 4, 1};
 
-    std::vector<uint8_t> input({ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16 });
-    std::vector<uint8_t> expectedOutput({ 1, 5, 2, 6, 9, 13, 10, 14, 3, 7, 4, 8, 11, 15, 12, 16});
+    std::vector<uint8_t> input({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+    std::vector<uint8_t> expectedOutput({1, 5, 2, 6, 9, 13, 10, 14, 3, 7, 4, 8, 11, 15, 12, 16});
 
     std::vector<unsigned int> blockShape({2, 2});
     std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
@@ -8531,6 +8558,161 @@ LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest1(
                                                input, blockShape, crops, outputShape, expectedOutput);
 }
 
+LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest2(
+    armnn::IWorkloadFactory& workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
+{
+    const unsigned int inputShape[] = {4, 1, 1, 1};
+    const unsigned int outputShape[] = {1, 2, 2, 1};
+
+    std::vector<uint8_t> input({
+        // Batch 0, Height 0, Width (2) x Channel (1)
+        1, 2, 3, 4
+    });
+
+    std::vector<uint8_t> expectedOutput({1, 2, 3, 4});
+
+    std::vector<unsigned int> blockShape({2, 2});
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
+
+    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
+                                               armnn::DataLayout::NHWC, inputShape, input, blockShape,
+                                               crops, outputShape, expectedOutput);
+}
+
+LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest3(
+    armnn::IWorkloadFactory& workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
+{
+    const unsigned int inputShape[] = {4, 1, 1, 3};
+    const unsigned int outputShape[] = {1, 2, 2, 3};
+
+    std::vector<uint8_t> input({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+
+    std::vector<uint8_t> expectedOutput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+
+    std::vector<unsigned int> blockShape({2, 2});
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
+
+    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
+                                               armnn::DataLayout::NHWC, inputShape, input, blockShape,
+                                               crops, outputShape, expectedOutput);
+}
+
+
+LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest1(
+    armnn::IWorkloadFactory &workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
+{
+    const unsigned int inputShape[] = {4, 3, 1, 1};
+    const unsigned int outputShape[] = {1, 3, 2, 2};
+
+    std::vector<uint8_t> input({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+
+    std::vector<uint8_t> expectedOutput({
+        // Batch 0, Channel 0, Height (2) x Width (2)
+        1,  4,
+        7, 10,
+
+        // Batch 0, Channel 1, Height (2) x Width (2)
+        2,  5,
+        8, 11,
+
+        // Batch 0, Channel 2, Height (2) x Width (2)
+        3,  6,
+        9, 12,
+    });
+
+    std::vector<unsigned int> blockShape({2, 2});
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
+
+    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
+                                               armnn::DataLayout::NCHW, inputShape, input, blockShape,
+                                               crops, outputShape, expectedOutput);
+}
+
+LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest2(
+    armnn::IWorkloadFactory& workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
+{
+    const unsigned int inputShape[] = {4, 1, 1, 1};
+    const unsigned int outputShape[] = {1, 1, 2, 2};
+
+    std::vector<uint8_t> input({
+        // Batch 0, Height 0, Width (2) x Channel (1)
+        1, 2, 3, 4
+    });
+
+    std::vector<uint8_t> expectedOutput({1, 2, 3, 4});
+
+    std::vector<unsigned int> blockShape({2, 2});
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
+
+    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
+                                               armnn::DataLayout::NCHW, inputShape, input, blockShape,
+                                               crops, outputShape, expectedOutput);
+}
+
+LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest3(
+    armnn::IWorkloadFactory& workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
+{
+    const unsigned int inputShape[] = {4, 3, 1, 1};
+    const unsigned int outputShape[] = {1, 3, 2, 2};
+
+    std::vector<uint8_t> input({1, 3, 5, 7, 9, 11, 2, 4, 6, 8, 10, 12});
+
+    std::vector<uint8_t> expectedOutput({
+        // Batch 0, Channel 0, Height (2) x Width (2)
+        1,  7,
+        2,  8,
+
+        // Batch 0, Channel 1, Height (2) x Width (2)
+        3,  9,
+        4, 10,
+
+        // Batch 0, Channel 2, Height (2) x Width (2)
+        5, 11,
+        6, 12,
+    });
+
+    std::vector<unsigned int> blockShape({2, 2});
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
+
+    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
+                                               armnn::DataLayout::NCHW, inputShape, input, blockShape,
+                                               crops, outputShape, expectedOutput);
+}
+
+LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest4(
+        armnn::IWorkloadFactory& workloadFactory,
+        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
+{
+    const unsigned int inputShape[] = {8, 1, 1, 3};
+    const unsigned int outputShape[] = {2, 1, 2, 4};
+
+    std::vector<uint8_t> input({
+        0, 1, 3, 0,  9, 11,
+        0, 2, 4, 0, 10, 12,
+        0, 5, 7, 0, 13, 15,
+        0, 6, 8, 0, 14, 16
+    });
+
+    std::vector<uint8_t> expectedOutput({
+        1,  2,  3,  4,
+        5,  6,  7,  8,
+        9, 10, 11, 12,
+       13, 14, 15, 16
+    });
+
+    std::vector<unsigned int> blockShape({2, 2});
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {2, 0}};
+
+    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
+        armnn::DataLayout::NCHW, inputShape, input, blockShape,
+        crops, outputShape, expectedOutput);
+}
+
 LayerTestResult<float, 4> StridedSlice4DFloat32Test(
     armnn::IWorkloadFactory& workloadFactory,
     const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
@@ -8656,134 +8838,6 @@ LayerTestResult<uint8_t, 2> StridedSlice2DReverseUint8Test(
 {
     return StridedSlice2DReverseTest<armnn::DataType::QuantisedAsymm8>(workloadFactory, memoryManager);
 }
-LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest2(
-        armnn::IWorkloadFactory& workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
-{
-    const unsigned int inputShape[] = {4, 1, 1, 1};
-    const unsigned int outputShape[] = {1, 2, 2, 1};
-
-    std::vector<uint8_t> input
-            ({
-                     // Batch 0, Height 0, Width (2) x Channel (1)
-                     1, 2, 3, 4
-             });
-
-    std::vector<uint8_t> expectedOutput({1, 2, 3, 4});
-
-    std::vector<unsigned int> blockShape({2, 2});
-    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
-
-    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
-                                               armnn::DataLayout::NHWC, inputShape, input, blockShape,
-                                               crops, outputShape, expectedOutput);
-}
-
-LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest3(
-        armnn::IWorkloadFactory& workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
-{
-    const unsigned int inputShape[] = {4, 1, 1, 3};
-    const unsigned int outputShape[] = {1, 2, 2, 3};
-
-    std::vector<uint8_t> input({ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 });
-
-    std::vector<uint8_t> expectedOutput({ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 });
-
-    std::vector<unsigned int> blockShape({2, 2});
-    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
-
-    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
-                                               armnn::DataLayout::NHWC, inputShape, input, blockShape,
-                                               crops, outputShape, expectedOutput);
-}
-
-
-LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest1(
-        armnn::IWorkloadFactory &workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
-{
-    const unsigned int inputShape[] = {4, 3, 1, 1};
-    const unsigned int outputShape[] = {1, 3, 2, 2};
-
-    std::vector<uint8_t> input({ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 });
-
-    std::vector<uint8_t> expectedOutput
-            ({
-                     // Batch 0, Channel 0, Height (2) x Width (2)
-                     1,  4,
-                     7, 10,
-
-                     // Batch 0, Channel 1, Height (2) x Width (2)
-                     2,  5,
-                     8, 11,
-
-                     // Batch 0, Channel 2, Height (2) x Width (2)
-                     3,  6,
-                     9, 12,
-             });
-
-    std::vector<unsigned int> blockShape({2, 2});
-    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
-
-    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
-                                               armnn::DataLayout::NCHW, inputShape, input, blockShape,
-                                               crops, outputShape, expectedOutput);
-}
-
-LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest2(
-        armnn::IWorkloadFactory& workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
-{
-    const unsigned int inputShape[] = {4, 1, 1, 1};
-    const unsigned int outputShape[] = {1, 1, 2, 2};
-
-    std::vector<uint8_t> input
-            ({
-                     // Batch 0, Height 0, Width (2) x Channel (1)
-                     1, 2, 3, 4
-             });
-
-    std::vector<uint8_t> expectedOutput({1, 2, 3, 4});
-
-    std::vector<unsigned int> blockShape({2, 2});
-    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
-
-    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
-                                             armnn::DataLayout::NCHW, inputShape, input, blockShape,
-                                             crops, outputShape, expectedOutput);
-}
-
-LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest3(
-        armnn::IWorkloadFactory& workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager)
-{
-    const unsigned int inputShape[] = {4, 3, 1, 1};
-    const unsigned int outputShape[] = {1, 3, 2, 2};
-
-    std::vector<uint8_t> input({ 1, 3, 5, 7, 9, 11, 2, 4, 6, 8, 10, 12 });
-
-    std::vector<uint8_t> expectedOutput
-            ({
-                     // Batch 0, Channel 0, Height (2) x Width (2)
-                     1,  7,
-                     2,  8,
-
-                     // Batch 0, Channel 1, Height (2) x Width (2)
-                     3,  9,
-                     4, 10,
-
-                     // Batch 0, Channel 2, Height (2) x Width (2)
-                     5, 11,
-                     6, 12,
-             });
-    std::vector<unsigned int> blockShape({2, 2});
-    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
-
-    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, memoryManager,
-                                               armnn::DataLayout::NCHW, inputShape, input, blockShape,
-                                               crops, outputShape, expectedOutput);
-}
 
 LayerTestResult<float, 4> Debug4DFloat32Test(
     armnn::IWorkloadFactory& workloadFactory,
diff --git a/src/backends/backendsCommon/test/LayerTests.hpp b/src/backends/backendsCommon/test/LayerTests.hpp
index 93385f00ce..587ffe9a3e 100644
--- a/src/backends/backendsCommon/test/LayerTests.hpp
+++ b/src/backends/backendsCommon/test/LayerTests.hpp
@@ -1196,41 +1196,49 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test3(
     armnn::IWorkloadFactory& workloadFactory,
     const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
+LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test4(
+    armnn::IWorkloadFactory& workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
+
 LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test1(
     armnn::IWorkloadFactory &workloadFactory,
     const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
 LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test2(
-        armnn::IWorkloadFactory &workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
+    armnn::IWorkloadFactory &workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
 LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test3(
-        armnn::IWorkloadFactory &workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
+    armnn::IWorkloadFactory &workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
 LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest1(
     armnn::IWorkloadFactory &workloadFactory,
     const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
 LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest2(
-        armnn::IWorkloadFactory &workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
+    armnn::IWorkloadFactory &workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
 LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest3(
-        armnn::IWorkloadFactory &workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
+    armnn::IWorkloadFactory &workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
 LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest1(
-        armnn::IWorkloadFactory &workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
+    armnn::IWorkloadFactory &workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
 LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest2(
-        armnn::IWorkloadFactory &workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
+    armnn::IWorkloadFactory &workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
 LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest3(
-        armnn::IWorkloadFactory &workloadFactory,
-        const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
+    armnn::IWorkloadFactory &workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
+
+LayerTestResult<uint8_t, 4> BatchToSpaceNdNchwUintTest4(
+    armnn::IWorkloadFactory &workloadFactory,
+    const armnn::IBackendInternal::IMemoryManagerSharedPtr& memoryManager);
 
 LayerTestResult<float, 4> StridedSlice4DFloat32Test(
     armnn::IWorkloadFactory& workloadFactory,
diff --git a/src/backends/reference/test/RefLayerTests.cpp b/src/backends/reference/test/RefLayerTests.cpp
index cf8e6a2e25..4be917f721 100644
--- a/src/backends/reference/test/RefLayerTests.cpp
+++ b/src/backends/reference/test/RefLayerTests.cpp
@@ -449,6 +449,7 @@ ARMNN_AUTO_TEST_CASE(SpaceToBatchNdPaddingNHWCUint8, SpaceToBatchNdPaddingNHWCUi
 ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNhwcFloat321, BatchToSpaceNdNhwcFloat32Test1)
 ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNhwcFloat322, BatchToSpaceNdNhwcFloat32Test2)
 ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNhwcFloat323, BatchToSpaceNdNhwcFloat32Test3)
+ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNhwcFloat324, BatchToSpaceNdNhwcFloat32Test4)
 
 ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNchwFloat321, BatchToSpaceNdNchwFloat32Test1)
 ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNchwFloat322, BatchToSpaceNdNchwFloat32Test2)
@@ -461,6 +462,7 @@ ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNhwcUint3, BatchToSpaceNdNhwcUintTest3)
 ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNchwUint1, BatchToSpaceNdNchwUintTest1)
 ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNchwUint2, BatchToSpaceNdNchwUintTest2)
 ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNchwUint3, BatchToSpaceNdNchwUintTest3)
+ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNchwUint4, BatchToSpaceNdNchwUintTest4)
 
 // Strided Slice
 ARMNN_AUTO_TEST_CASE(StridedSlice4DFloat32, StridedSlice4DFloat32Test)