From a41d2ecb5af8554cbaa3690329e5d15e1d7af1f0 Mon Sep 17 00:00:00 2001 From: TFLM-bot Date: Wed, 28 Jul 2021 09:13:42 -0700 Subject: [PATCH] Automated sync from github.com/tensorflow/tensorflow (#328) Co-authored-by: Advait Jain --- .../lite/kernels/internal/reference/reduce.h | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/tensorflow/lite/kernels/internal/reference/reduce.h b/tensorflow/lite/kernels/internal/reference/reduce.h index 5b83eef3..348e170e 100644 --- a/tensorflow/lite/kernels/internal/reference/reduce.h +++ b/tensorflow/lite/kernels/internal/reference/reduce.h @@ -160,7 +160,8 @@ inline bool InitTensorDataForReduce(const int* dims, const int num_dims, for (int idx = 0; idx < num_dims; ++idx) { size_t current = static_cast(dims[idx]); // Overflow prevention. - if (num_elements > std::numeric_limits::max() / current) { + if (current > 0 && + num_elements > std::numeric_limits::max() / current) { return false; } num_elements *= current; @@ -181,17 +182,20 @@ inline bool ReduceGeneric(const T* input_data, const int* input_dims, bool keep_dims, int* temp_index, int* resolved_axis, T init_value, T reducer(const T current, const T in)) { - // Return early when input shape has zero dim. - for (int i = 0; i < input_num_dims; ++i) { - if (input_dims[i] == 0) return true; - } - // Reset output data. if (!InitTensorDataForReduce(output_dims, output_num_dims, init_value, output_data)) { return false; } + // Return early when input shape has zero dim. This is done after initializing + // data for output tensor because there are cases that the input tensor is + // empty but output tensor is not. In that case, output tensor should be + // filled with init_value. + for (int i = 0; i < input_num_dims; ++i) { + if (input_dims[i] == 0) return true; + } + // Resolve axis. int num_resolved_axis = 0; if (!ResolveAxis(input_num_dims, axis, num_axis_dimensions, resolved_axis, @@ -402,6 +406,14 @@ inline bool QuantizedMeanOrSum(const T* input_data, int32_t input_zero_point, temp_sum[idx] = U(); } + // Return early when input shape has zero dim. This is done after initializing + // data for output tensor because there are cases that the input tensor is + // empty but output tensor is not. In that case, output tensor should be + // filled with init_value. + for (int i = 0; i < input_num_dims; ++i) { + if (input_dims[i] == 0) return true; + } + // Resolve axis. int num_resolved_axis = 0; if (!ResolveAxis(input_num_dims, axis, num_axis_dimensions, resolved_axis, -- GitLab