inline void ArgMinMax(const RuntimeShape& input1_shape, const T1* input1_data,
                      const T3* input2_data, const RuntimeShape& output_shape,
                      T2* output_data, const bool is_arg_max) {
  ruy::profiler::ScopeLabel label("ArgMinMax");

  TFLITE_DCHECK_GT(input1_shape.DimensionsCount(), 0);
  TFLITE_DCHECK_EQ(input1_shape.DimensionsCount() - 1,
                   output_shape.DimensionsCount());
  int axis = input2_data[0];
  if (axis < 0) {
    axis += input1_shape.DimensionsCount();
  }
  const int axis_size = input1_shape.Dims(axis);

  int outer_size = 1;
  for (int i = 0; i < axis; ++i) {
    TFLITE_DCHECK_EQ(input1_shape.Dims(i), output_shape.Dims(i));
    outer_size *= input1_shape.Dims(i);
  }

  int inner_size = 1;
  const int dims_count = input1_shape.DimensionsCount();
  for (int i = axis + 1; i < dims_count; ++i) {
    TFLITE_DCHECK_EQ(input1_shape.Dims(i), output_shape.Dims(i - 1));
    inner_size *= input1_shape.Dims(i);
  }

  // Call specialized function when axis=dims-1. So far, only float32 is
  // optimized so reroute to specialized function only when T1 is float32.
  if (inner_size == 1 &&
      (std::is_same<T1, float>::value || std::is_same<T1, int8_t>::value ||
       std::is_same<T1, uint8_t>::value)) {
    if (is_arg_max) {
      ArgMinMaxLastAxis<T1, T2, /*is_arg_max=*/true>(
          {outer_size, axis_size}, input1_data, {outer_size}, output_data);
    } else {
      ArgMinMaxLastAxis<T1, T2, /*is_arg_max=*/false>(
          {outer_size, axis_size}, input1_data, {outer_size}, output_data);
    }
    return;
  }

  reference_ops::ArgMinMax(input1_shape, input1_data, input2_data, output_shape,
                           output_data, is_arg_max);
}
inline void ArgMinMax(const RuntimeShape& input1_shape, const T1* input1_data,
                      const T3* input2_data, const RuntimeShape& output_shape,
                      T2* output_data, const bool is_arg_max) {
  ruy::profiler::ScopeLabel label("ArgMinMax");

  TFLITE_DCHECK_GT(input1_shape.DimensionsCount(), 0);
  TFLITE_DCHECK_EQ(input1_shape.DimensionsCount() - 1,
                   output_shape.DimensionsCount());
  int axis = input2_data[0];
  if (axis < 0) {
    axis += input1_shape.DimensionsCount();
  }
  const int axis_size = input1_shape.Dims(axis);

  int outer_size = 1;
  for (int i = 0; i < axis; ++i) {
    TFLITE_DCHECK_EQ(input1_shape.Dims(i), output_shape.Dims(i));
    outer_size *= input1_shape.Dims(i);
  }

  int inner_size = 1;
  const int dims_count = input1_shape.DimensionsCount();
  for (int i = axis + 1; i < dims_count; ++i) {
    TFLITE_DCHECK_EQ(input1_shape.Dims(i), output_shape.Dims(i - 1));
    inner_size *= input1_shape.Dims(i);
  }

   
   
  if (inner_size == 1 &&
      (std::is_same<T1, float>::value || std::is_same<T1, int8_t>::value ||
       std::is_same<T1, uint8_t>::value)) {
    if (is_arg_max) {
      ArgMinMaxLastAxis<T1, T2,  true>(
          {outer_size, axis_size}, input1_data, {outer_size}, output_data);
    } else {
      ArgMinMaxLastAxis<T1, T2,  false>(
          {outer_size, axis_size}, input1_data, {outer_size}, output_data);
    }
    return;
  }

  reference_ops::ArgMinMax(input1_shape, input1_data, input2_data, output_shape,
                           output_data, is_arg_max);
}