From ca44ffc43cbe6972f8cfc2a803626418554195c2 Mon Sep 17 00:00:00 2001
From: Teresa Charlin <teresa.charlinreyes@arm.com>
Date: Tue, 11 Aug 2020 23:00:18 +0100
Subject: IVGCVSW-5207 Add to PadLayer the functions to validate and infer
 shapes

Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Change-Id: I63a6da418862b0a843637d73cbe178086f23bc84
---
 src/armnn/layers/PadLayer.cpp | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

(limited to 'src/armnn/layers/PadLayer.cpp')

diff --git a/src/armnn/layers/PadLayer.cpp b/src/armnn/layers/PadLayer.cpp
index 4fcbc77c7b..a8c749c570 100644
--- a/src/armnn/layers/PadLayer.cpp
+++ b/src/armnn/layers/PadLayer.cpp
@@ -36,10 +36,39 @@ PadLayer* PadLayer::Clone(Graph& graph) const
     return std::move(layer);
 }
 
+std::vector<TensorShape> PadLayer::InferOutputShapes(const std::vector<TensorShape>& inputShapes) const
+{
+    ARMNN_ASSERT(inputShapes.size() == 1);
+    const TensorShape& inputShape = inputShapes[0];
+
+    unsigned int rank = inputShape.GetNumDimensions();
+    ARMNN_ASSERT(m_Param.m_PadList.size() == rank);
+    ARMNN_ASSERT(rank != 0);
+
+    std::vector<unsigned int> outputDimensionSizes;
+    outputDimensionSizes.reserve(rank);
+    for (unsigned int i = 0; i < rank; ++i)
+    {
+        outputDimensionSizes[i] = inputShape[i] + m_Param.m_PadList[i].first + m_Param.m_PadList[i].second;
+    }
+
+    TensorShape tensorShape = TensorShape( rank, outputDimensionSizes.data());
+    return std::vector<TensorShape>({ tensorShape });
+}
+
 void PadLayer::ValidateTensorShapesFromInputs()
 {
+    VerifyLayerConnections(1, CHECK_LOCATION());
+
+    const TensorShape& outputShape = GetOutputSlot(0).GetTensorInfo().GetShape();
+
+    VerifyShapeInferenceType(outputShape, m_ShapeInferenceMethod);
+
+    auto inferredShapes = InferOutputShapes({ GetInputSlot(0).GetConnection()->GetTensorInfo().GetShape() });
+
+    ARMNN_ASSERT(inferredShapes.size() == 1);
 
-    return;
+    ValidateAndCopyShape(outputShape, inferredShapes[0], m_ShapeInferenceMethod, "PadLayer");
 }
 
 void PadLayer::Accept(ILayerVisitor& visitor) const
-- 
cgit v1.2.1