| 272 | |
| 273 | template <KernelType kernel_type> |
| 274 | TfLiteStatus EvalMean(TfLiteContext* context, TfLiteNode* node) { |
| 275 | OpContext op_context(context, node); |
| 276 | OpData* data = reinterpret_cast<OpData*>(node->user_data); |
| 277 | |
| 278 | int num_axis = static_cast<int>(NumElements(op_context.axis)); |
| 279 | TfLiteTensor* temp_index = GetTemporary(context, node, /*index=*/0); |
| 280 | TfLiteTensor* resolved_axis = GetTemporary(context, node, /*index=*/1); |
| 281 | TfLiteTensor* temp_sum = GetTemporary(context, node, /*index=*/2); |
| 282 | // Resize the output tensor if the output tensor is dynamic. |
| 283 | if (IsDynamicTensor(op_context.output)) { |
| 284 | TF_LITE_ENSURE_OK(context, |
| 285 | ResizeTempAxis(context, &op_context, resolved_axis)); |
| 286 | TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); |
| 287 | TF_LITE_ENSURE_OK(context, ResizeTempSum(context, &op_context, temp_sum)); |
| 288 | } |
| 289 | |
| 290 | if (kernel_type == kGenericOptimized) { |
| 291 | // Use optimized ops if available. |
| 292 | switch (op_context.input->type) { |
| 293 | case kTfLiteInt8: { |
| 294 | tflite::MeanParams op_params; |
| 295 | op_params.axis_count = num_axis; |
| 296 | ResolveAxis(GetTensorData<int>(op_context.axis), num_axis, &op_params); |
| 297 | const TfLiteTensor* input = op_context.input; |
| 298 | optimized_integer_ops::Mean( |
| 299 | op_params, GetTensorShape(input), GetTensorData<int8_t>(input), |
| 300 | op_context.input->params.zero_point, op_context.input->params.scale, |
| 301 | GetTensorShape(op_context.output), |
| 302 | GetTensorData<int8_t>(op_context.output), |
| 303 | op_context.output->params.zero_point, |
| 304 | op_context.output->params.scale, |
| 305 | CpuBackendContext::GetFromContext(context)); |
| 306 | return kTfLiteOk; |
| 307 | } break; |
| 308 | case kTfLiteUInt8: { |
| 309 | tflite::MeanParams op_params; |
| 310 | op_params.axis_count = num_axis; |
| 311 | ResolveAxis(GetTensorData<int>(op_context.axis), num_axis, &op_params); |
| 312 | const TfLiteTensor* input = op_context.input; |
| 313 | // TODO(b/13910232): Handle the below special case in the optimized |
| 314 | // method. |
| 315 | if (op_context.params->keep_dims && NumDimensions(input) == 4 && |
| 316 | op_params.axis_count == 2 && |
| 317 | ((op_params.axis[0] == 1 && op_params.axis[1] == 2) || |
| 318 | (op_params.axis[0] == 2 && op_params.axis[1] == 1))) { |
| 319 | optimized_ops::Mean( |
| 320 | op_params, GetTensorShape(input), GetTensorData<uint8_t>(input), |
| 321 | op_context.input->params.zero_point, |
| 322 | op_context.input->params.scale, GetTensorShape(op_context.output), |
| 323 | GetTensorData<uint8_t>(op_context.output), |
| 324 | op_context.output->params.zero_point, |
| 325 | op_context.output->params.scale, |
| 326 | CpuBackendContext::GetFromContext(context)); |
| 327 | return kTfLiteOk; |
| 328 | } |
| 329 | } break; |
| 330 | default: |
| 331 | break; |
nothing calls this directly
no test coverage detected