2 * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
3 * Copyright 2019 The TensorFlow Authors. All Rights Reserved.
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
9 * http://www.apache.org/licenses/LICENSE-2.0
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
18 #include "kernels/ArgMax.h"
19 #include "kernels/Utils.h"
20 #include "PALArgMax.h"
22 namespace luci_interpreter
27 ArgMax::ArgMax(const Tensor *input, const Tensor *axis, Tensor *output, const ArgMaxParams ¶ms)
28 : KernelWithParams<ArgMaxParams>({input, axis}, {output}, params)
32 void ArgMax::configure()
34 assert(axis()->element_type() == DataType::S32 || axis()->element_type() == DataType::S64);
35 assert(input()->shape().num_dims() >= 1);
36 const Shape &input_shape = input()->shape();
37 const int num_dims = input_shape.num_dims();
38 Shape output_shape(num_dims - 1);
40 // If axis value is negative, then update by adding input_shape's num_dims.
41 // If updated value also negative, then assert.
42 assert(axis()->shape().num_elements() == 1);
43 int axis_value = getTensorData<int32_t>(axis())[0];
45 axis_value = axis_value + num_dims;
46 assert(axis_value >= 0);
49 for (int i = 0; i < num_dims; i++)
53 output_shape.dim(j++) = input_shape.dim(i);
56 assert(output()->element_type() == _params.output_type);
58 // TODO: enable it only if kernel with dynamic shapes
59 output()->resize(output_shape);
62 void ArgMax::execute() const
65 #define TF_LITE_ARG_MAX(data_type, axis_type, output_type) \
66 luci_interpreter_pal::ArgMinMax(getTensorShape(input()), getTensorData<data_type>(input()), \
67 getTensorData<axis_type>(axis()), getTensorShape(output()), \
68 getTensorData<output_type>(output()), std::greater<data_type>())
69 if (axis()->element_type() == DataType::S32)
71 switch (_params.output_type)
74 switch (input()->element_type())
76 case DataType::FLOAT32:
77 TF_LITE_ARG_MAX(float, int32_t, int32_t);
80 TF_LITE_ARG_MAX(uint8_t, int32_t, int32_t);
83 assert(false && "Unsupported input type.");
87 switch (input()->element_type())
89 case DataType::FLOAT32:
90 TF_LITE_ARG_MAX(float, int32_t, int64_t);
93 TF_LITE_ARG_MAX(uint8_t, int32_t, int64_t);
96 assert(false && "Unsupported input type.");
100 assert(false && "Unsupported output type.");
105 switch (_params.output_type)
108 switch (input()->element_type())
110 case DataType::FLOAT32:
111 TF_LITE_ARG_MAX(float, int64_t, int32_t);
114 TF_LITE_ARG_MAX(uint8_t, int64_t, int32_t);
117 assert(false && "Unsupported input type.");
121 switch (input()->element_type())
123 case DataType::FLOAT32:
124 TF_LITE_ARG_MAX(float, int64_t, int64_t);
127 TF_LITE_ARG_MAX(uint8_t, int64_t, int64_t);
130 assert(false && "Unsupported input type.");
134 assert(false && "Unsupported output type.");
137 #undef TF_LITE_ARG_MAX
140 } // namespace kernels
141 } // namespace luci_interpreter