ONE - On-device Neural Engine
Loading...
Searching...
No Matches
onert::backend::cpu::ops::ArgMinMaxLayer Class Reference

#include <ArgMinMaxLayer.h>

Collaboration diagram for onert::backend::cpu::ops::ArgMinMaxLayer:

Public Member Functions

 ArgMinMaxLayer ()
 
void configure (const IPortableTensor *indices, IPortableTensor *output, const IPortableTensor *axis, bool is_arg_max)
 
void run () override
 
- Public Member Functions inherited from onert::exec::IFunction
virtual ~IFunction ()=default
 
virtual void prepare ()
 

Detailed Description

Definition at line 27 of file ArgMinMaxLayer.h.

Constructor & Destructor Documentation

◆ ArgMinMaxLayer()

onert::backend::cpu::ops::ArgMinMaxLayer::ArgMinMaxLayer ( )
inline

Definition at line 30 of file ArgMinMaxLayer.h.

30: _input(nullptr), _output(nullptr), _axis(nullptr), _is_arg_max(true) {}

Member Function Documentation

◆ configure()

void onert::backend::cpu::ops::ArgMinMaxLayer::configure ( const IPortableTensor indices,
IPortableTensor output,
const IPortableTensor axis,
bool  is_arg_max 
)

Definition at line 67 of file ArgMinMaxLayer.cc.

69{
70 _input = input;
71 _output = output;
72 _axis = axis;
73 _is_arg_max = is_arg_max;
74}

◆ run()

void onert::backend::cpu::ops::ArgMinMaxLayer::run ( )
overridevirtual

Implements onert::exec::IFunction.

Definition at line 76 of file ArgMinMaxLayer.cc.

77{
78 if (_axis->total_size() != sizeof(int32_t))
79 {
80 throw std::runtime_error("ArgMinMax: wrong shape of axis");
81 }
82 auto axis = *getBuffer<int32_t>(_axis);
83 if (axis < 0)
84 {
85 axis += _input->getShape().rank();
86 }
87#define TF_LITE_ARG_MIN_MAX(input_type, axis_type, output_type) \
88 ArgMinMax(getShape(_input), getBuffer<input_type>(_input), getShape(_output), \
89 getBuffer<output_type>(_output), axis, GetComparefunction<input_type>(_is_arg_max));
90 if (_output->data_type() == ir::DataType::INT32)
91 {
92 switch (_input->data_type())
93 {
94 case ir::DataType::FLOAT32:
95 TF_LITE_ARG_MIN_MAX(float, int32_t, int32_t);
96 break;
97 case ir::DataType::QUANT_UINT8_ASYMM:
98 case ir::DataType::UINT8:
99 TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int32_t);
100 break;
101 case ir::DataType::QUANT_INT8_ASYMM:
102 TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int32_t);
103 break;
104 case ir::DataType::INT32:
105 TF_LITE_ARG_MIN_MAX(int32_t, int32_t, int32_t);
106 break;
107 default:
108 throw std::runtime_error("ArgMinMax: unsupported data type");
109 }
110 }
111 else if (_output->data_type() == ir::DataType::INT64)
112 {
113 switch (_input->data_type())
114 {
115 case ir::DataType::FLOAT32:
116 TF_LITE_ARG_MIN_MAX(float, int32_t, int64_t);
117 break;
118 case ir::DataType::QUANT_UINT8_ASYMM:
119 case ir::DataType::UINT8:
120 TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int64_t);
121 break;
122 case ir::DataType::QUANT_INT8_ASYMM:
123 TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int64_t);
124 break;
125 case ir::DataType::INT32:
126 TF_LITE_ARG_MIN_MAX(int32_t, int32_t, int64_t);
127 break;
128 default:
129 throw std::runtime_error("ArgMinMax: unsupported data type");
130 }
131 }
132 else
133 {
134 throw std::runtime_error("ArgMinMax: unsupported data type");
135 }
136
137#undef TF_LITE_ARG_MIN_MAX
138}
#define TF_LITE_ARG_MIN_MAX(input_type, axis_type, output_type)
size_t total_size() const override final
ir::DataType data_type() const override final
ir::Shape getShape() const override final
Get ir::Shape of tensor.

References onert::backend::IPortableTensor::data_type(), onert::backend::IPortableTensor::getShape(), TF_LITE_ARG_MIN_MAX, and onert::backend::IPortableTensor::total_size().


The documentation for this class was generated from the following files: