1 files changed, 41 insertions, 23 deletions

diff --git a/OutputShapeUtils.cpp b/OutputShapeUtils.cpp
index a1a49d70..739038ac 100644
--- a/OutputShapeUtils.cpp
+++ b/OutputShapeUtils.cpp
@@ -8,6 +8,42 @@
 #include <algorithm>
 #include <vector>
 
+namespace
+{
+
+using namespace armnn;
+
+TensorShape CalculateMaxShape(const TensorShape& inShape0, const TensorShape& inShape1)
+{
+    // NOTE: The inferred output size will be the maximum size along each dimension
+    // of inShape0 and inShape1, starting with the trailing dimensions, and working its way forward.
+    //
+    // Example: inShape0={4, 1, 2}, inShape1={5, 4, 3, 1} => outputShape={5, 4, 3, 2}
+
+    const unsigned int numInput0Dims = inShape0.GetNumDimensions();
+    const unsigned int numInput1Dims = inShape1.GetNumDimensions();
+
+    const unsigned int maxNumDims = std::max(numInput0Dims, numInput1Dims);
+
+    TensorShape outputShape = TensorShape(maxNumDims);
+    for (unsigned int reverseIdx = 1u; reverseIdx <= maxNumDims; ++reverseIdx)
+    {
+        const int input0Idx = numInput0Dims - reverseIdx;
+        const int input1Idx = numInput1Dims - reverseIdx;
+
+        const unsigned int input0DimSize = input0Idx >= 0 ? inShape0[input0Idx] : 0u;
+        const unsigned int input1DimSize = input1Idx >= 0 ? inShape1[input1Idx] : 0u;
+
+        const unsigned int outputIdx = maxNumDims - reverseIdx;
+        outputShape[outputIdx] = std::max(input0DimSize, input1DimSize);
+    }
+
+    return outputShape;
+}
+
+} // namespace annonymous
+
+
 namespace armnn_driver
 {
 
@@ -38,30 +74,12 @@ TensorShape InferPadOutputShape(const TensorShape& inputShape,
 
 TensorShape InferPreluOutputShape(const TensorShape& inputShape, const TensorShape& alphaShape)
 {
-    // NOTE: The inferred PReLU output size will be the maximum size along each dimension
-    // of input and alpha, starting with the trailing dimensions, and working its way forward.
-    //
-    // Example: inputShape={4, 1, 2}, alphaShape={5, 4, 3, 1} => outputShape={5, 4, 3, 2}
-
-    const unsigned int numInputDims = inputShape.GetNumDimensions();
-    const unsigned int numAlphaDims = alphaShape.GetNumDimensions();
-
-    const unsigned int maxNumDims = std::max(numInputDims, numAlphaDims);
-
-    TensorShape outputShape = TensorShape(maxNumDims);
-    for (unsigned int reverseIdx = 1u; reverseIdx <= maxNumDims; ++reverseIdx)
-    {
-        const int inputIdx = numInputDims - reverseIdx;
-        const int alphaIdx = numAlphaDims - reverseIdx;
-
-        const unsigned int inputDimSize = inputIdx >= 0 ? inputShape[inputIdx] : 0u;
-        const unsigned int alphaDimSize = alphaIdx >= 0 ? alphaShape[alphaIdx] : 0u;
-
-        const unsigned int outputIdx = maxNumDims - reverseIdx;
-        outputShape[outputIdx] = std::max(inputDimSize, alphaDimSize);
-    }
+    return CalculateMaxShape(inputShape, alphaShape);
+}
 
-    return outputShape;
+TensorShape InferSubOutputShape(const TensorShape& input0Shape, const TensorShape& input1Shape)
+{
+    return CalculateMaxShape(input0Shape, input1Shape);
 }
 
 } // namespace armnn_driver
\ No newline at end of file