Fix QuantizedMeanOrSum float static_cast overflow.
PiperOrigin-RevId: 277650994
Change-Id: I547624832fbc339e146debe1e3fd0ab62d661f64
diff --git a/tensorflow/lite/kernels/internal/reference/reference_ops.h b/tensorflow/lite/kernels/internal/reference/reference_ops.h
index ed5d2b2..2a69faa 100644
--- a/tensorflow/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/lite/kernels/internal/reference/reference_ops.h
@@ -2924,11 +2924,12 @@
for (size_t idx = 0; idx < num_outputs; ++idx) {
float float_mean = static_cast<float>(temp_sum[idx]) /
static_cast<float>(num_elements_in_axis);
-
- // Convert to float value.
- output_data[idx] =
- static_cast<T>(std::round(float_mean * scale + bias)) +
- output_zero_point;
+ float result =
+ std::min(std::round(float_mean * scale + bias) + output_zero_point,
+ static_cast<float>(std::numeric_limits<T>::max()));
+ result =
+ std::max(result, static_cast<float>(std::numeric_limits<T>::min()));
+ output_data[idx] = static_cast<T>(result);
}
}
}
