TfLiteStatus EvalMeanReferenceOps(TfLiteContext* context,
                                  const OpContext& op_context, int num_axis,
                                  OpData* data, TfLiteTensor* temp_index,
                                  TfLiteTensor* resolved_axis,
                                  TfLiteTensor* temp_sum) {
  tflite::MeanParams op_params;
  op_params.axis_count = num_axis;
  ResolveAxis(GetTensorData<int>(op_context.axis), num_axis, &op_params);
  const TfLiteTensor* input = op_context.input;
  // TODO(b/139102329): Handle all the cases in the combined reference
  // method.
  if (op_context.params->keep_dims && NumDimensions(input) == 4 &&
      op_params.axis_count == 2 &&
      ((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
       (op_params.axis[0] == 2 && op_params.axis[1] == 1))) {
    if (std::is_same<integer_type, uint8_t>::value) {
      reference_ops::Mean(op_params, GetTensorShape(op_context.input),
                          GetTensorData<uint8_t>(op_context.input),
                          op_context.input->params.zero_point,
                          op_context.input->params.scale,
                          GetTensorShape(op_context.output),
                          GetTensorData<uint8_t>(op_context.output),
                          op_context.output->params.zero_point,
                          op_context.output->params.scale);
    } else {
      reference_integer_ops::Mean(
          op_params, data->multiplier, data->shift, GetTensorShape(input),
          GetTensorData<integer_type>(input),
          op_context.input->params.zero_point,
          GetTensorShape(op_context.output),
          GetTensorData<integer_type>(op_context.output),
          op_context.output->params.zero_point);
    }
  } else if (input->params.zero_point == op_context.output->params.zero_point &&
             input->params.scale == op_context.output->params.scale) {
    TF_LITE_ENSURE(
        context,
        reference_ops::Mean(
            GetTensorData<integer_type>(input), input->dims->data,
            input->dims->size, GetTensorData<integer_type>(op_context.output),
            op_context.output->dims->data, op_context.output->dims->size,
            GetTensorData<int>(op_context.axis), num_axis,
            op_context.params->keep_dims, GetTensorData<int>(temp_index),
            GetTensorData<int>(resolved_axis), GetTensorData<int>(temp_sum)));
  } else {
    TF_LITE_ENSURE(
        context,
        reference_ops::QuantizedMeanOrSum<>(
            GetTensorData<integer_type>(input), input->params.zero_point,
            input->params.scale, input->dims->data, input->dims->size,
            GetTensorData<integer_type>(op_context.output),
            op_context.output->params.zero_point,
            op_context.output->params.scale, op_context.output->dims->data,
            op_context.output->dims->size, GetTensorData<int>(op_context.axis),
            num_axis, op_context.params->keep_dims,
            GetTensorData<int>(temp_index), GetTensorData<int>(resolved_axis),
            GetTensorData<int>(temp_sum),
            /*compute_sum=*/false));
  }
  return kTfLiteOk;
}
TfLiteStatus EvalMeanReferenceOps(TfLiteContext* context,
                                  const OpContext& op_context, int num_axis,
                                  OpData* data, TfLiteTensor* temp_index,
                                  TfLiteTensor* resolved_axis,
                                  TfLiteTensor* temp_sum) {
  tflite::MeanParams op_params;
  op_params.axis_count = num_axis;
  ResolveAxis(GetTensorData<int>(op_context.axis), num_axis, &op_params);
  const TfLiteTensor* input = op_context.input;
   
   
  if (op_context.params->keep_dims && NumDimensions(input) == 4 &&
      op_params.axis_count == 2 &&
      ((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
       (op_params.axis[0] == 2 && op_params.axis[1] == 1))) {
    if (std::is_same<integer_type, uint8_t>::value) {
      reference_ops::Mean(op_params, GetTensorShape(op_context.input),
                          GetTensorData<uint8_t>(op_context.input),
                          op_context.input->params.zero_point,
                          op_context.input->params.scale,
                          GetTensorShape(op_context.output),
                          GetTensorData<uint8_t>(op_context.output),
                          op_context.output->params.zero_point,
                          op_context.output->params.scale);
    } else {
      reference_integer_ops::Mean(
          op_params, data->multiplier, data->shift, GetTensorShape(input),
          GetTensorData<integer_type>(input),
          op_context.input->params.zero_point,
          GetTensorShape(op_context.output),
          GetTensorData<integer_type>(op_context.output),
          op_context.output->params.zero_point);
    }
  } else if (input->params.zero_point == op_context.output->params.zero_point &&
             input->params.scale == op_context.output->params.scale) {
    TF_LITE_ENSURE(
        context,
        reference_ops::Mean(
            GetTensorData<integer_type>(input), input->dims->data,
            input->dims->size, GetTensorData<integer_type>(op_context.output),
            op_context.output->dims->data, op_context.output->dims->size,
            GetTensorData<int>(op_context.axis), num_axis,
            op_context.params->keep_dims, GetTensorData<int>(temp_index),
            GetTensorData<int>(resolved_axis), GetTensorData<int>(temp_sum)));
  } else {
    TF_LITE_ENSURE(
        context,
        reference_ops::QuantizedMeanOrSum<>(
            GetTensorData<integer_type>(input), input->params.zero_point,
            input->params.scale, input->dims->data, input->dims->size,
            GetTensorData<integer_type>(op_context.output),
            op_context.output->params.zero_point,
            op_context.output->params.scale, op_context.output->dims->data,
            op_context.output->dims->size, GetTensorData<int>(op_context.axis),
            num_axis, op_context.params->keep_dims,
            GetTensorData<int>(temp_index), GetTensorData<int>(resolved_axis),
            GetTensorData<int>(temp_sum),
             false));
  }
  return kTfLiteOk;
}