Tidying up multiple issues

* Fixed error in InferOutputShape implementation
* Added better error checking to the BatchToSpace implementation.
* Added defaults to the batchToSpace descriptors.
* Changed crops to be a vector of pairs to align with the SpaceToBatch implementation


Change-Id: Ib1c16d871f0898a1caeb6629c1fee6380a773e14

5 files changed, 29 insertions, 34 deletions

diff --git a/include/armnn/Descriptors.hpp b/include/armnn/Descriptors.hpp
index bda8cf7396..32ac959808 100644
--- a/include/armnn/Descriptors.hpp
+++ b/include/armnn/Descriptors.hpp
@@ -299,19 +299,20 @@ struct BatchNormalizationDescriptor
 struct BatchToSpaceNdDescriptor
 {
     BatchToSpaceNdDescriptor()
-        : m_BlockShape()
-        , m_Crops()
+        : m_BlockShape({0, 0})
+        , m_Crops({{0, 0}, {0, 0}})
         , m_DataLayout(DataLayout::NCHW)
     {}
 
-    BatchToSpaceNdDescriptor(std::vector<unsigned int> blockShape, std::vector<std::vector<unsigned int>> crops)
+    BatchToSpaceNdDescriptor(std::vector<unsigned int> blockShape,
+                             std::vector<std::pair<unsigned int, unsigned int>> crops)
         : m_BlockShape(blockShape)
         , m_Crops(crops)
         , m_DataLayout(DataLayout::NCHW)
     {}
 
     std::vector<unsigned int> m_BlockShape;
-    std::vector<std::vector<unsigned int>> m_Crops;
+    std::vector<std::pair<unsigned int, unsigned int>> m_Crops;
     DataLayoutIndexed m_DataLayout;
 };
 
diff --git a/src/armnn/layers/BatchToSpaceNdLayer.cpp b/src/armnn/layers/BatchToSpaceNdLayer.cpp
index 595ce4a7fe..9366a8710b 100644
--- a/src/armnn/layers/BatchToSpaceNdLayer.cpp
+++ b/src/armnn/layers/BatchToSpaceNdLayer.cpp
@@ -57,23 +57,19 @@ std::vector<TensorShape> BatchToSpaceNdLayer::InferOutputShapes(const std::vecto
 
     std::vector<unsigned int> theBlockShape = m_Param.m_BlockShape;
 
-    unsigned int overallSize = inBatchSize;
+    unsigned int overallSize = inBatchSize * inputShape[dataLayout.GetHeightIndex()]
+                               * inputShape[dataLayout.GetWidthIndex()];
 
-    for (unsigned int i = 0; i < theBlockShape.size(); ++i)
-    {
-        overallSize = overallSize * theBlockShape.at(i);
-    }
-
-    std::vector<std::vector<unsigned int>> crops = m_Param.m_Crops;
+    std::vector<std::pair<unsigned int, unsigned int>> crops = m_Param.m_Crops;
 
-    std::vector<unsigned int> yCrops = crops[0];
-    std::vector<unsigned int> xCrops = crops[1];
+    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 = theBlockShape.at(0) * (inputHeight - (yCrops[0] + yCrops[1]));
+    unsigned int outputHeight = theBlockShape.at(0) * (inputHeight - (yCrops.first + yCrops.second));
 
     unsigned int inputWidth = inputShape[dataLayout.GetWidthIndex()];
-    unsigned int outputWidth = theBlockShape.at(1) * (inputWidth - (xCrops[0] + xCrops[1]));
+    unsigned int outputWidth = theBlockShape.at(1) * (inputWidth - (xCrops.first + xCrops.second));
 
     unsigned int outputBatchSize = overallSize / (outputHeight * outputWidth);
 
diff --git a/src/backends/backendsCommon/test/LayerTests.cpp b/src/backends/backendsCommon/test/LayerTests.cpp
index 4a003036ca..85b3e1b2b7 100755
--- a/src/backends/backendsCommon/test/LayerTests.cpp
+++ b/src/backends/backendsCommon/test/LayerTests.cpp
@@ -6178,7 +6178,7 @@ LayerTestResult<T, OutputDim> BatchToSpaceNdHelper(armnn::IWorkloadFactory &work
                                                    const unsigned int *inputShape,
                                                    const std::vector<T> &inputData,
                                                    const std::vector<unsigned int> &blockShape,
-                                                   const std::vector<std::vector<unsigned int>> &crops,
+                                                   const std::vector<std::pair<unsigned int, unsigned int>> &crops,
                                                    const unsigned int *outputShape,
                                                    const std::vector<T> &outputData,
                                                    float scale = 1.0f,
@@ -6266,7 +6266,7 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test1(armnn::IWorkloadFactory
     });
 
     std::vector<unsigned int> blockShape {2, 2};
-    std::vector<std::vector<unsigned int>> crops = {{0, 0}, {0, 0}};
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
 
     return BatchToSpaceNdHelper<float, 4, 4>(workloadFactory, armnn::DataLayout::NHWC, inputShape, input, blockShape,
             crops, outputShape, expectedOutput);
@@ -6286,7 +6286,7 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test2(armnn::IWorkloadFactory
     std::vector<float> expectedOutput({1.0f,   2.0f,  3.0f,  4.0f});
 
     std::vector<unsigned int> blockShape({2, 2});
-    std::vector<std::vector<unsigned int>> crops = {{0, 0}, {0, 0}};
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
 
     return BatchToSpaceNdHelper<float, 4, 4>(workloadFactory, armnn::DataLayout::NHWC, inputShape, input, blockShape,
                                              crops, outputShape, expectedOutput);
@@ -6302,7 +6302,7 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test3(armnn::IWorkloadFactory
     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::vector<unsigned int>> crops = {{0, 0}, {0, 0}};
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
 
     return BatchToSpaceNdHelper<float, 4, 4>(workloadFactory, armnn::DataLayout::NHWC, inputShape, input, blockShape,
                                              crops, outputShape, expectedOutput);
@@ -6331,7 +6331,7 @@ LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test1(armnn::IWorkloadFactory
     });
 
     std::vector<unsigned int> blockShape({2, 2});
-    std::vector<std::vector<unsigned int>> crops = {{0, 0}, {0, 0}};
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
 
     return BatchToSpaceNdHelper<float, 4, 4>(workloadFactory, armnn::DataLayout::NCHW, inputShape, input, blockShape,
                                              crops, outputShape, expectedOutput);
diff --git a/src/backends/reference/workloads/BatchToSpaceNd.cpp b/src/backends/reference/workloads/BatchToSpaceNd.cpp
index bedf8418ef..4313085ba5 100644
--- a/src/backends/reference/workloads/BatchToSpaceNd.cpp
+++ b/src/backends/reference/workloads/BatchToSpaceNd.cpp
@@ -34,23 +34,17 @@ void BatchToSpaceNd(const DataLayoutIndexed& dataLayout,
                     const TensorInfo& inputTensorInfo,
                     const TensorInfo& outputTensorInfo,
                     const std::vector<unsigned int>& blockShape,
-                    const std::vector<std::vector<unsigned int>>& cropsData,
+                    const std::vector<std::pair<unsigned int, unsigned int>>& cropsData,
                     const float* inputData,
                     float* outputData)
 {
     TensorShape inputShape = inputTensorInfo.GetShape();
-    unsigned int inputNumDims = inputShape.GetNumDimensions();
-    if (inputNumDims != 4)
-    {
-        throw armnn::InvalidArgumentException("Expected Input with 4 Dimensions");
-    }
+
+    BOOST_ASSERT_MSG(inputShape.GetNumDimensions() == 4, "Expected Input with 4 Dimensions");
 
     TensorShape outputShape = outputTensorInfo.GetShape();
-    unsigned int outputNumDims = outputShape.GetNumDimensions();
-    if (outputNumDims != 4)
-    {
-        throw armnn::InvalidArgumentException("Expected Output with 4 Dimensions");
-    }
+
+    BOOST_ASSERT_MSG(outputShape.GetNumDimensions() == 4, "Expected Output with 4 Dimensions");
 
     const unsigned int inputBatchSize = inputShape[0];
     const unsigned int channels = inputShape[dataLayout.GetChannelsIndex()];
@@ -59,11 +53,15 @@ void BatchToSpaceNd(const DataLayoutIndexed& dataLayout,
     const unsigned int outputHeight = outputShape[dataLayout.GetHeightIndex()];
     const unsigned int outputWidth = outputShape[dataLayout.GetWidthIndex()];
 
+    BOOST_ASSERT_MSG(blockShape.size() > 0, "BlockShape must contain 1 or more entries");
+
     const unsigned int blockShapeHeight = blockShape[0];
     const unsigned int blockShapeWidth = blockShape[1];
 
-    const unsigned int cropsTop = cropsData[0][0];
-    const unsigned int cropsLeft = cropsData[1][0];
+    BOOST_ASSERT_MSG(cropsData.size() > 0, "Crops must contain 1 or more entries");
+
+    const unsigned int cropsTop = cropsData[0].first;
+    const unsigned int cropsLeft = cropsData[1].first;
 
     for (unsigned int inBatch = 0; inBatch < inputBatchSize; ++inBatch)
     {
diff --git a/src/backends/reference/workloads/BatchToSpaceNd.hpp b/src/backends/reference/workloads/BatchToSpaceNd.hpp
index 7923ceadd0..091d092777 100644
--- a/src/backends/reference/workloads/BatchToSpaceNd.hpp
+++ b/src/backends/reference/workloads/BatchToSpaceNd.hpp
@@ -16,7 +16,7 @@ void BatchToSpaceNd(const DataLayoutIndexed& dataLayout,
                     const TensorInfo& inputTensorInfo,
                     const TensorInfo& outputTensorInfo,
                     const std::vector<unsigned int>& blockShape,
-                    const std::vector<std::vector<unsigned int>>& cropsData,
+                    const std::vector<std::pair<unsigned int, unsigned int>>& cropsData,
                     const float* inputData,
                     float* outputData);
 } // namespace armnn
\ No newline at end of file