| 113 | } |
| 114 | |
| 115 | TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node, bool is_arg_max) { |
| 116 | const TfLiteTensor* input = GetInput(context, node, kInputTensor); |
| 117 | const TfLiteTensor* axis = GetInput(context, node, kAxis); |
| 118 | TfLiteTensor* output = GetOutput(context, node, kOutputTensor); |
| 119 | if (IsDynamicTensor(output)) { |
| 120 | TF_LITE_ENSURE_STATUS(ResizeOutput(context, input, axis, output)); |
| 121 | } |
| 122 | |
| 123 | #define TF_LITE_ARG_MIN_MAX(data_type, axis_type, output_type) \ |
| 124 | optimized_ops::ArgMinMax( \ |
| 125 | GetTensorShape(input), GetTensorData<data_type>(input), \ |
| 126 | GetTensorData<axis_type>(axis), GetTensorShape(output), \ |
| 127 | GetTensorData<output_type>(output), \ |
| 128 | GetComparefunction<data_type>(is_arg_max)) |
| 129 | if (axis->type == kTfLiteInt32) { |
| 130 | switch (output->type) { |
| 131 | case kTfLiteInt32: { |
| 132 | switch (input->type) { |
| 133 | case kTfLiteFloat32: |
| 134 | TF_LITE_ARG_MIN_MAX(float, int32_t, int32_t); |
| 135 | break; |
| 136 | case kTfLiteUInt8: |
| 137 | TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int32_t); |
| 138 | break; |
| 139 | case kTfLiteInt8: |
| 140 | TF_LITE_ARG_MIN_MAX(int8_t, int32_t, int32_t); |
| 141 | break; |
| 142 | case kTfLiteInt32: |
| 143 | TF_LITE_ARG_MIN_MAX(int32_t, int32_t, int32_t); |
| 144 | break; |
| 145 | default: |
| 146 | context->ReportError(context, |
| 147 | "Only float32, uint8, int8 and int32 are " |
| 148 | "supported currently, got %s.", |
| 149 | TfLiteTypeGetName(input->type)); |
| 150 | return kTfLiteError; |
| 151 | } |
| 152 | } break; |
| 153 | case kTfLiteInt64: { |
| 154 | switch (input->type) { |
| 155 | case kTfLiteFloat32: |
| 156 | TF_LITE_ARG_MIN_MAX(float, int32_t, int64_t); |
| 157 | break; |
| 158 | case kTfLiteUInt8: |
| 159 | TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int64_t); |
| 160 | break; |
| 161 | case kTfLiteInt8: |
| 162 | TF_LITE_ARG_MIN_MAX(int8_t, int32_t, int64_t); |
| 163 | break; |
| 164 | case kTfLiteInt32: |
| 165 | TF_LITE_ARG_MIN_MAX(int32_t, int32_t, int64_t); |
| 166 | break; |
| 167 | default: |
| 168 | context->ReportError(context, |
| 169 | "Only float32, uint8, int8 and int32 are " |
| 170 | "supported currently, got %s.", |
| 171 | TfLiteTypeGetName(input->type)); |
| 172 | return kTfLiteError; |
no test coverage detected