Sync from upstream TF.

tensorflow · Oct 10, 2024 · fca5ed3 · fca5ed3
1 parent e86d97b
commit fca5ed3
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/tensorflow/lite/kernels/internal/reference/batch_matmul.h b/tensorflow/lite/kernels/internal/reference/batch_matmul.h
@@ -111,7 +111,8 @@ inline void BatchMatMul(const RuntimeShape& lhs_shape, const int8_t* lhs_data,
  const float* scaling_factors,
  const int32_t* input_offset, int32_t* row_sums,
  const RuntimeShape& output_shape, float* output_data,
- bool* compute_row_sums) {
+ bool* compute_row_sums,
+ const float* per_channel_scales) {
  const RuntimeShape extended_lhs_shape =
  RuntimeShape::ExtendedShape(5, lhs_shape);
  const RuntimeShape extended_rhs_shape =
@@ -188,7 +189,11 @@ inline void BatchMatMul(const RuntimeShape& lhs_shape, const int8_t* lhs_data,
  int32_t row_sum = woff_ptr2[i];
  total -= row_sum * batch_offset;
  int idx = lhs_rows * j + i;
- out_ptr[idx] += batch_scaling_factor * total;
+ float ss = batch_scaling_factor;
+ if (per_channel_scales) {
+ ss *= per_channel_scales[j];
+ }
+ out_ptr[idx] += ss * total;
  }
  }
  }