TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node, OpData* data,
                        const RuntimeShape& input_shape,
                        const TfLiteTensor* input,
                        const RuntimeShape& filter_shape,
                        const TfLiteTensor* filter,
                        TfLiteTensor* input_quantized,
                        TfLiteTensor* scaling_factors,
                        TfLiteTensor* accum_scratch, TfLiteTensor* row_sums,
                        TfLiteTensor* input_offsets, TfLiteTensor* output) {
  const int32_t num_input_dims = input_shape.DimensionsCount();

  // Input row/cols have been swapped at this point, so dims are
  // {input_size, num_batches}
  const int input_size = input_shape.Dims(num_input_dims - 2);
  const int batch_size = input_shape.Dims(num_input_dims - 1);

  int num_batches_to_quantize = batch_size;
  for (int i = 0; i < input_shape.DimensionsCount() - 2; ++i) {
    num_batches_to_quantize *= input_shape.Dims(i);
  }
  // Quantize input from float to uint8 + quantization params (scaling factor).
  float* scaling_factors_ptr = GetTensorData<float>(scaling_factors);
  int32_t* input_offset_ptr = nullptr;
  int32_t* row_sums_ptr = nullptr;
  // Only asymmetric quantization is supported.
  input_offset_ptr = GetTensorData<int32_t>(input_offsets);
  row_sums_ptr = GetTensorData<int32_t>(row_sums);
  int8_t* quant_data = GetTensorData<int8_t>(input_quantized);
  const int8_t* filter_data = GetTensorData<int8_t>(filter);
  const float* input_ptr = GetTensorData<float>(input);
  // Quantize each batch independently.
  for (int b = 0; b < num_batches_to_quantize; ++b) {
    const int offset = b * input_size;
    tensor_utils::AsymmetricQuantizeFloats(
        input_ptr + offset, input_size, quant_data + offset,
        &scaling_factors_ptr[b], &input_offset_ptr[b]);
    // Incorporate scaling of the filter.
    scaling_factors_ptr[b] *= filter->params.scale;
  }

  RuntimeShape output_shape = GetTensorShape(output);
  int output_size = 1;
  for (int i = 0; i < output_shape.DimensionsCount(); ++i) {
    output_size *= output_shape.Dims(i);
  }
  std::fill_n(GetTensorData<float>(output), output_size, 0.0f);
  if (kernel_type == kGenericOptimized) {
    optimized_ops::BatchMatMul(
        filter_shape, filter_data, input_shape, quant_data, scaling_factors_ptr,
        input_offset_ptr, row_sums_ptr, GetTensorShape(output),
        GetTensorData<int32_t>(accum_scratch), GetTensorData<float>(output),
        &(data->compute_row_sums), CpuBackendContext::GetFromContext(context));
  } else {
    reference_ops::BatchMatMul(
        filter_shape, filter_data, input_shape, quant_data, scaling_factors_ptr,
        input_offset_ptr, row_sums_ptr, GetTensorShape(output),
        GetTensorData<float>(output), &(data->compute_row_sums));
  }

  return kTfLiteOk;
}
TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node, OpData* data,
                        const RuntimeShape& input_shape,
                        const TfLiteTensor* input,
                        const RuntimeShape& filter_shape,
                        const TfLiteTensor* filter,
                        TfLiteTensor* input_quantized,
                        TfLiteTensor* scaling_factors,
                        TfLiteTensor* accum_scratch, TfLiteTensor* row_sums,
                        TfLiteTensor* input_offsets, TfLiteTensor* output) {
  const int32_t num_input_dims = input_shape.DimensionsCount();

   
   
  const int input_size = input_shape.Dims(num_input_dims - 2);
  const int batch_size = input_shape.Dims(num_input_dims - 1);

  int num_batches_to_quantize = batch_size;
  for (int i = 0; i < input_shape.DimensionsCount() - 2; ++i) {
    num_batches_to_quantize *= input_shape.Dims(i);
  }
   
  float* scaling_factors_ptr = GetTensorData<float>(scaling_factors);
  int32_t* input_offset_ptr = nullptr;
  int32_t* row_sums_ptr = nullptr;
   
  input_offset_ptr = GetTensorData<int32_t>(input_offsets);
  row_sums_ptr = GetTensorData<int32_t>(row_sums);
  int8_t* quant_data = GetTensorData<int8_t>(input_quantized);
  const int8_t* filter_data = GetTensorData<int8_t>(filter);
  const float* input_ptr = GetTensorData<float>(input);
   
  for (int b = 0; b < num_batches_to_quantize; ++b) {
    const int offset = b * input_size;
    tensor_utils::AsymmetricQuantizeFloats(
        input_ptr + offset, input_size, quant_data + offset,
        &scaling_factors_ptr[b], &input_offset_ptr[b]);
     
    scaling_factors_ptr[b] *= filter->params.scale;
  }

  RuntimeShape output_shape = GetTensorShape(output);
  int output_size = 1;
  for (int i = 0; i < output_shape.DimensionsCount(); ++i) {
    output_size *= output_shape.Dims(i);
  }
  std::fill_n(GetTensorData<float>(output), output_size, 0.0f);
  if (kernel_type == kGenericOptimized) {
    optimized_ops::BatchMatMul(
        filter_shape, filter_data, input_shape, quant_data, scaling_factors_ptr,
        input_offset_ptr, row_sums_ptr, GetTensorShape(output),
        GetTensorData<int32_t>(accum_scratch), GetTensorData<float>(output),
        &(data->compute_row_sums), CpuBackendContext::GetFromContext(context));
  } else {
    reference_ops::BatchMatMul(
        filter_shape, filter_data, input_shape, quant_data, scaling_factors_ptr,
        input_offset_ptr, row_sums_ptr, GetTensorShape(output),
        GetTensorData<float>(output), &(data->compute_row_sums));
  }

  return kTfLiteOk;
}