Internal cleanup of NumElements() function.
PiperOrigin-RevId: 281604730
Change-Id: I34ae9018779abb6ff759b43364855fc623ad2490
diff --git a/tensorflow/lite/tools/optimize/quantization_utils.cc b/tensorflow/lite/tools/optimize/quantization_utils.cc
index 52c3d77..8588f2f 100644
--- a/tensorflow/lite/tools/optimize/quantization_utils.cc
+++ b/tensorflow/lite/tools/optimize/quantization_utils.cc
@@ -41,10 +41,6 @@
} // namespace
TfLiteStatus NumElements(const TensorT& tensor, uint64_t* num_elements) {
- if (tensor.shape.empty()) {
- *num_elements = 0;
- return kTfLiteOk;
- }
*num_elements = 1;
for (const int64_t dim : tensor.shape) {
if (dim <= 0 || *num_elements > UINT64_MAX / static_cast<uint64_t>(dim)) {
diff --git a/tensorflow/lite/tools/optimize/quantization_utils_test.cc b/tensorflow/lite/tools/optimize/quantization_utils_test.cc
index 9487939..ece0123 100644
--- a/tensorflow/lite/tools/optimize/quantization_utils_test.cc
+++ b/tensorflow/lite/tools/optimize/quantization_utils_test.cc
@@ -65,7 +65,8 @@
tensor.shape = {};
EXPECT_EQ(kTfLiteOk, NumElements(tensor, &num_elements));
- EXPECT_EQ(num_elements, 0);
+ // Scalars with empty shape have 1 element.
+ EXPECT_EQ(num_elements, 1);
tensor.shape = {1, 2, 3, -1};
EXPECT_EQ(kTfLiteError, NumElements(tensor, &num_elements));
