1 //
2 // Copyright © 2017 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5
6 #include "ArgMinMaxLayer.hpp"
7 #include "LayerCloneBase.hpp"
8
9 #include <armnn/TypesUtils.hpp>
10
11 #include <armnnUtils/TensorUtils.hpp>
12
13 #include <backendsCommon/WorkloadData.hpp>
14 #include <backendsCommon/WorkloadFactory.hpp>
15
16 namespace armnn
17 {
18
ArgMinMaxLayer(const ArgMinMaxDescriptor & param,const char * name)19 ArgMinMaxLayer::ArgMinMaxLayer(const ArgMinMaxDescriptor& param, const char* name)
20 : LayerWithParameters(1, 1, LayerType::ArgMinMax, param, name)
21 {
22 }
23
CreateWorkload(const IWorkloadFactory & factory) const24 std::unique_ptr<IWorkload> ArgMinMaxLayer::CreateWorkload(const IWorkloadFactory& factory) const
25 {
26 ArgMinMaxQueueDescriptor descriptor;
27 SetAdditionalInfo(descriptor);
28
29 return factory.CreateArgMinMax(descriptor, PrepInfoAndDesc(descriptor));
30 }
31
Clone(Graph & graph) const32 ArgMinMaxLayer* ArgMinMaxLayer::Clone(Graph& graph) const
33 {
34 return CloneBase<ArgMinMaxLayer>(graph, m_Param, GetName());
35 }
36
InferOutputShapes(const std::vector<TensorShape> & inputShapes) const37 std::vector<TensorShape> ArgMinMaxLayer::InferOutputShapes(const std::vector<TensorShape>& inputShapes) const
38 {
39 ARMNN_ASSERT(inputShapes.size() == 1);
40
41 TensorShape inputShape = inputShapes[0];
42 auto inputNumDimensions = inputShape.GetNumDimensions();
43
44 auto axis = m_Param.m_Axis;
45 auto unsignedAxis = armnnUtils::GetUnsignedAxis(inputNumDimensions, axis);
46
47 ARMNN_ASSERT(unsignedAxis <= inputNumDimensions);
48
49 // 1D input shape results in scalar output
50 if (inputShape.GetNumDimensions() == 1)
51 {
52 std::vector<unsigned int> tensorDimensions(1, 1);
53 TensorShape outputShape(1, tensorDimensions.data());
54
55 return std::vector<TensorShape>({ outputShape });
56 }
57
58 std::vector<unsigned int> tensorDimensions(inputNumDimensions - 1, 0);
59 for (unsigned int i = 0; i < unsignedAxis; ++i)
60 {
61 tensorDimensions[i] = inputShape[i];
62 }
63
64 for (unsigned int i = unsignedAxis + 1; i < inputNumDimensions; ++i)
65 {
66 tensorDimensions[i - 1] = inputShape[i];
67 }
68
69 TensorShape outputShape = TensorShape(inputNumDimensions - 1, tensorDimensions.data());
70
71 return std::vector<TensorShape>({ outputShape });
72 }
73
ValidateTensorShapesFromInputs()74 void ArgMinMaxLayer::ValidateTensorShapesFromInputs()
75 {
76 VerifyLayerConnections(1, CHECK_LOCATION());
77
78 const TensorShape& outputShape = GetOutputSlot(0).GetTensorInfo().GetShape();
79
80 VerifyShapeInferenceType(outputShape, m_ShapeInferenceMethod);
81
82 auto inferredShapes = InferOutputShapes({ GetInputSlot(0).GetConnection()->GetTensorInfo().GetShape() });
83
84 ARMNN_ASSERT(inferredShapes.size() == 1);
85
86 ValidateAndCopyShape(outputShape, inferredShapes[0], m_ShapeInferenceMethod, "ArgMinMaxLayer");
87 }
88
Accept(ILayerVisitor & visitor) const89 void ArgMinMaxLayer::Accept(ILayerVisitor& visitor) const
90 {
91 visitor.VisitArgMinMaxLayer(this, GetParameters(), GetName());
92 }
93
94 } // namespace armnn
95