IVGCVSW-3944 Add ArgMinMax output shape validation

Signed-off-by: James Conroy <james.conroy@arm.com>
Change-Id: I469895da158b062cd19248832525fa21527f7d41

4 files changed, 136 insertions, 0 deletions

diff --git a/src/armnn/layers/ArgMinMaxLayer.cpp b/src/armnn/layers/ArgMinMaxLayer.cpp
index aad95eb0cf..bfd71d519b 100644
--- a/src/armnn/layers/ArgMinMaxLayer.cpp
+++ b/src/armnn/layers/ArgMinMaxLayer.cpp
@@ -6,6 +6,8 @@
 
 #include "LayerCloneBase.hpp"
 
+#include <TensorUtils.hpp>
+
 #include <armnn/TypesUtils.hpp>
 #include <backendsCommon/WorkloadData.hpp>
 #include <backendsCommon/WorkloadFactory.hpp>
@@ -30,6 +32,43 @@ ArgMinMaxLayer* ArgMinMaxLayer::Clone(Graph& graph) const
     return CloneBase<ArgMinMaxLayer>(graph, m_Param, GetName());
 }
 
+std::vector<TensorShape> ArgMinMaxLayer::InferOutputShapes(const std::vector<TensorShape>& inputShapes) const
+{
+    BOOST_ASSERT(inputShapes.size() == 1);
+
+    TensorShape inputShape = inputShapes[0];
+    auto inputNumDimensions = inputShape.GetNumDimensions();
+
+    auto axis = m_Param.m_Axis;
+    auto unsignedAxis = armnnUtils::GetUnsignedAxis(inputNumDimensions, axis);
+
+    BOOST_ASSERT(unsignedAxis <= inputNumDimensions);
+
+    // 1D input shape results in scalar output
+    if (inputShape.GetNumDimensions() == 1)
+    {
+        std::vector<unsigned int> tensorDimensions(1, 1);
+        TensorShape outputShape(1, tensorDimensions.data());
+
+        return std::vector<TensorShape>({ outputShape });
+    }
+
+    std::vector<unsigned int> tensorDimensions(inputNumDimensions - 1, 0);
+    for (unsigned int i = 0; i < unsignedAxis; ++i)
+    {
+        tensorDimensions[i] = inputShape[i];
+    }
+
+    for (unsigned int i = unsignedAxis + 1; i < inputNumDimensions; ++i)
+    {
+        tensorDimensions[i - 1] = inputShape[i];
+    }
+
+    TensorShape outputShape = TensorShape(inputNumDimensions - 1, tensorDimensions.data());
+
+    return std::vector<TensorShape>({ outputShape });
+}
+
 void ArgMinMaxLayer::ValidateTensorShapesFromInputs()
 {
     VerifyLayerConnections(1, CHECK_LOCATION());
diff --git a/src/armnn/layers/ArgMinMaxLayer.hpp b/src/armnn/layers/ArgMinMaxLayer.hpp
index ca1337f065..43ea056c9e 100644
--- a/src/armnn/layers/ArgMinMaxLayer.hpp
+++ b/src/armnn/layers/ArgMinMaxLayer.hpp
@@ -25,6 +25,11 @@ public:
     /// @param [in] graph The graph into which this layer is being cloned.
     ArgMinMaxLayer* Clone(Graph& graph) const override;
 
+    /// Infers the output shape from a given input shape and axis parameter.
+    /// @param [in] inputShapes The vector of input shapes for ArgMinMax.
+    /// @return A vector of inferred output shapes.
+    std::vector<TensorShape> InferOutputShapes(const std::vector<TensorShape>& inputShapes) const override;
+
     /// Check if the input tensor shape(s)
     /// will lead to a valid configuration of @ref ArgMinMaxLayer.
     void ValidateTensorShapesFromInputs() override;
diff --git a/src/armnn/test/InferOutputTests.cpp b/src/armnn/test/InferOutputTests.cpp
index 8606745623..3293cef0f7 100644
--- a/src/armnn/test/InferOutputTests.cpp
+++ b/src/armnn/test/InferOutputTests.cpp
@@ -11,6 +11,12 @@
 
 BOOST_AUTO_TEST_SUITE(LayerValidateOutput)
 
+// ArgMinMax
+ARMNN_SIMPLE_TEST_CASE(ArgMinMaxInferOutputShape4d, ArgMinMaxInferOutputShape4dTest)
+ARMNN_SIMPLE_TEST_CASE(ArgMinMaxInferOutputShape3d, ArgMinMaxInferOutputShape3dTest)
+ARMNN_SIMPLE_TEST_CASE(ArgMinMaxInferOutputShape2d, ArgMinMaxInferOutputShape2dTest)
+ARMNN_SIMPLE_TEST_CASE(ArgMinMaxInferOutputShape1d, ArgMinMaxInferOutputShape1dTest)
+
 // BatchToSpace
 ARMNN_SIMPLE_TEST_CASE(BatchToSpaceInferOutputShape, BatchToSpaceInferOutputShapeTest)
 
diff --git a/src/armnn/test/InferOutputTests.hpp b/src/armnn/test/InferOutputTests.hpp
index c428a9db61..feb2125656 100644
--- a/src/armnn/test/InferOutputTests.hpp
+++ b/src/armnn/test/InferOutputTests.hpp
@@ -10,6 +10,7 @@
 #include <armnn/ArmNN.hpp>
 
 #include <Graph.hpp>
+#include <layers/ArgMinMaxLayer.hpp>
 #include <layers/BatchToSpaceNdLayer.hpp>
 #include <layers/SpaceToDepthLayer.hpp>
 #include <layers/PreluLayer.hpp>
@@ -18,6 +19,91 @@
 #include <boost/algorithm/string.hpp>
 #include <boost/test/unit_test.hpp>
 
+void ArgMinMaxInferOutputShapeImpl(const armnn::ArgMinMaxDescriptor       descriptor,
+                                   const std::vector<armnn::TensorShape>& inputShapes,
+                                   std::vector<armnn::TensorShape>&       outputShapes)
+{
+    armnn::Graph graph;
+    auto argMinMaxLayer = graph.AddLayer<armnn::ArgMinMaxLayer>(descriptor, "argMinMax");
+    outputShapes = argMinMaxLayer->InferOutputShapes(inputShapes);
+}
+
+void ArgMinMaxInferOutputShape4dTest()
+{
+    armnn::Graph graph;
+    armnn::ArgMinMaxDescriptor descriptor;
+    descriptor.m_Axis = 2;
+
+    const std::vector<armnn::TensorShape> inputShapes
+    {
+        { 1, 3, 2, 4 }
+    };
+
+    std::vector<armnn::TensorShape> outputShapes;
+    BOOST_CHECK_NO_THROW(ArgMinMaxInferOutputShapeImpl(descriptor, inputShapes, outputShapes));
+
+    armnn::TensorShape expectedOutputShape( { 1, 3, 4 } );
+    BOOST_CHECK(outputShapes.size() == 1);
+    BOOST_CHECK(outputShapes[0] == expectedOutputShape);
+}
+
+void ArgMinMaxInferOutputShape3dTest()
+{
+    armnn::Graph graph;
+    armnn::ArgMinMaxDescriptor descriptor;
+    descriptor.m_Axis = 0;
+
+    const std::vector<armnn::TensorShape> inputShapes
+    {
+        { 1, 3, 2 }
+    };
+
+    std::vector<armnn::TensorShape> outputShapes;
+    BOOST_CHECK_NO_THROW(ArgMinMaxInferOutputShapeImpl(descriptor, inputShapes, outputShapes));
+
+    armnn::TensorShape expectedOutputShape( { 3, 2 } );
+    BOOST_CHECK(outputShapes.size() == 1);
+    BOOST_CHECK(outputShapes[0] == expectedOutputShape);
+}
+
+void ArgMinMaxInferOutputShape2dTest()
+{
+    armnn::Graph graph;
+    armnn::ArgMinMaxDescriptor descriptor;
+    descriptor.m_Axis = 1;
+
+    const std::vector<armnn::TensorShape> inputShapes
+    {
+        { 3, 2 }
+    };
+
+    std::vector<armnn::TensorShape> outputShapes;
+    BOOST_CHECK_NO_THROW(ArgMinMaxInferOutputShapeImpl(descriptor, inputShapes, outputShapes));
+
+    armnn::TensorShape expectedOutputShape( { 3 } );
+    BOOST_CHECK(outputShapes.size() == 1);
+    BOOST_CHECK(outputShapes[0] == expectedOutputShape);
+}
+
+void ArgMinMaxInferOutputShape1dTest()
+{
+    armnn::Graph graph;
+    armnn::ArgMinMaxDescriptor descriptor;
+    descriptor.m_Axis = 0;
+
+    const std::vector<armnn::TensorShape> inputShapes
+    {
+        { 5 }
+    };
+
+    std::vector<armnn::TensorShape> outputShapes;
+    BOOST_CHECK_NO_THROW(ArgMinMaxInferOutputShapeImpl(descriptor, inputShapes, outputShapes));
+
+    armnn::TensorShape expectedOutputShape( { 1 } );
+    BOOST_CHECK(outputShapes.size() == 1);
+    BOOST_CHECK(outputShapes[0] == expectedOutputShape);
+}
+
 void BatchToSpaceInferOutputShapeTest()
 {
     armnn::Graph graph;