2 // Copyright © 2017 Arm Ltd. All rights reserved.
3 // SPDX-License-Identifier: MIT
5 #include "ArgMinMaxLayer.hpp"
7 #include "LayerCloneBase.hpp"
9 #include <armnn/TypesUtils.hpp>
10 #include <backendsCommon/WorkloadData.hpp>
11 #include <backendsCommon/WorkloadFactory.hpp>
16 ArgMinMaxLayer::ArgMinMaxLayer(const ArgMinMaxDescriptor& param, const char* name)
17 : LayerWithParameters(1, 1, LayerType::ArgMinMax, param, name)
21 std::unique_ptr<IWorkload> ArgMinMaxLayer::CreateWorkload(const Graph& graph,
22 const IWorkloadFactory& factory) const
24 ArgMinMaxQueueDescriptor descriptor;
25 return factory.CreateArgMinMax(descriptor, PrepInfoAndDesc(descriptor, graph));
28 ArgMinMaxLayer* ArgMinMaxLayer::Clone(Graph& graph) const
30 return CloneBase<ArgMinMaxLayer>(graph, m_Param, GetName());
33 void ArgMinMaxLayer::ValidateTensorShapesFromInputs()
35 VerifyLayerConnections(1, CHECK_LOCATION());
37 auto inferredShapes = InferOutputShapes({ GetInputSlot(0).GetConnection()->GetTensorInfo().GetShape() });
39 BOOST_ASSERT(inferredShapes.size() == 1);
41 ConditionalThrowIfNotEqual<LayerValidationException>(
42 "ArgMinMaxLayer: TensorShape set on OutputSlot does not match the inferred shape.",
43 GetOutputSlot(0).GetTensorInfo().GetShape(),
47 void ArgMinMaxLayer::Accept(ILayerVisitor& visitor) const
49 visitor.VisitArgMinMaxLayer(this, GetParameters(), GetName());