inline int ArgMaxVector(const int8_t* input_data, int size) {
  int32_t max_index = 0;
  int8_t max_value = input_data[0];
  int32_t i = 0;
#ifdef USE_NEON
  constexpr int VECTOR_SIZE = 16;
  if (size >= VECTOR_SIZE) {
    int8x16_t max_value_s8x16;
    for (; i <= size - VECTOR_SIZE; i += VECTOR_SIZE) {
      max_value_s8x16 = vld1q_s8(input_data + i);
      int8_t max_from_vec;
#ifdef __aarch64__
      max_from_vec = vmaxvq_s8(max_value_s8x16);
#else   // 32 bit
      int8x8_t max_val_s8x8 =
          vpmax_s8(vget_low_s8(max_value_s8x16), vget_high_s8(max_value_s8x16));
      max_val_s8x8 = vpmax_s8(max_val_s8x8, max_val_s8x8);
      max_val_s8x8 = vpmax_s8(max_val_s8x8, max_val_s8x8);
      max_val_s8x8 = vpmax_s8(max_val_s8x8, max_val_s8x8);
      max_from_vec = vget_lane_s8(max_val_s8x8, 0);
#endif  // __aarch64__
      if (max_from_vec > max_value) {
        max_value = max_from_vec;
        max_index = i;
      }
    }
  }
  for (int start_idx = max_index; start_idx < max_index + VECTOR_SIZE;
       start_idx++) {
    if (input_data[start_idx] == max_value) {
      max_index = start_idx;
      break;
    }
  }

#endif  // USE_NEON
  // Leftover loop.
  for (; i < size; ++i) {
    const int8_t curr_value = input_data[i];
    if (curr_value > max_value) {
      max_value = curr_value;
      max_index = i;
    }
  }

  return max_index;
}
inline int ArgMaxVector(const int8_t* input_data, int size) {
  int32_t max_index = 0;
  int8_t max_value = input_data[0];
  int32_t i = 0;
#ifdef USE_NEON
  constexpr int VECTOR_SIZE = 16;
  if (size >= VECTOR_SIZE) {
    int8x16_t max_value_s8x16;
    for (; i <= size - VECTOR_SIZE; i += VECTOR_SIZE) {
      max_value_s8x16 = vld1q_s8(input_data + i);
      int8_t max_from_vec;
#ifdef __aarch64__
      max_from_vec = vmaxvq_s8(max_value_s8x16);
#else    
      int8x8_t max_val_s8x8 =
          vpmax_s8(vget_low_s8(max_value_s8x16), vget_high_s8(max_value_s8x16));
      max_val_s8x8 = vpmax_s8(max_val_s8x8, max_val_s8x8);
      max_val_s8x8 = vpmax_s8(max_val_s8x8, max_val_s8x8);
      max_val_s8x8 = vpmax_s8(max_val_s8x8, max_val_s8x8);
      max_from_vec = vget_lane_s8(max_val_s8x8, 0);
#endif   
      if (max_from_vec > max_value) {
        max_value = max_from_vec;
        max_index = i;
      }
    }
  }
  for (int start_idx = max_index; start_idx < max_index + VECTOR_SIZE;
       start_idx++) {
    if (input_data[start_idx] == max_value) {
      max_index = start_idx;
      break;
    }
  }

#endif   
   
  for (; i < size; ++i) {
    const int8_t curr_value = input_data[i];
    if (curr_value > max_value) {
      max_value = curr_value;
      max_index = i;
    }
  }

  return max_index;
}