IVGCVSW-3456 Add support for dynamic output shape in ConvertPrelu
Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
Change-Id: I8fc7a716455be3f51b51177f6896a73790a41fc3
diff --git a/OutputShapeUtils.cpp b/OutputShapeUtils.cpp
new file mode 100644
index 0000000..de27630
--- /dev/null
+++ b/OutputShapeUtils.cpp
@@ -0,0 +1,43 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "OutputShapeUtils.hpp"
+
+#include <algorithm>
+
+namespace armnn_driver
+{
+
+using namespace armnn;
+
+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 outputShape;
+}
+
+} // namespace armnn_driver
\ No newline at end of file