Skip to content

Commit 9adbc8b

Browse files
TFLM-botTFLM-bot
committed
Sync from upstream TF.
1 parent 8f1f3b2 commit 9adbc8b

8 files changed

Lines changed: 73 additions & 13 deletions

File tree

‎tensorflow/lite/core/c/common.cc‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ limitations under the License.
2020
#endif // TF_LITE_STATIC_MEMORY
2121

2222
#include <cstring>
23+
#include <limits>
2324
#include <new>
2425
#include <type_traits>
2526
#include <utility>
@@ -460,15 +461,23 @@ TfLiteStatus TfLiteTensorCopy(const TfLiteTensor* src, TfLiteTensor* dst) {
460461

461462
TfLiteStatus TfLiteTensorResizeMaybeCopy(size_t num_bytes, TfLiteTensor* tensor,
462463
bool preserve_data) {
464+
if (tensor == nullptr) {
465+
return kTfLiteError;
466+
}
463467
if (tensor->allocation_type != kTfLiteDynamic &&
464468
tensor->allocation_type != kTfLitePersistentRo) {
465469
return kTfLiteOk;
466470
}
471+
// Guard against integer overflow: num_bytes + XNN_EXTRA_BYTES must not wrap.
472+
constexpr size_t kXnnExtraBytes = 16;
473+
if (num_bytes > std::numeric_limits<size_t>::max() - kXnnExtraBytes) {
474+
return kTfLiteError;
475+
}
467476
#ifdef TF_LITE_TENSORFLOW_PROFILER
468477
tflite::PauseHeapMonitoring(/*pause=*/true);
469478
#endif
470479
// This buffer may be consumed by XNNPack.
471-
size_t alloc_bytes = num_bytes + /*XNN_EXTRA_BYTES=*/16;
480+
size_t alloc_bytes = num_bytes + kXnnExtraBytes;
472481
// TODO(b/145340303): Tensor data should be aligned.
473482
if (!tensor->data.data) {
474483
tensor->data.data = (char*)malloc(alloc_bytes);

‎tensorflow/lite/kernels/internal/reference/arg_min_max.h‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,23 +15,50 @@ limitations under the License.
1515
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ARG_MIN_MAX_H_
1616
#define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ARG_MIN_MAX_H_
1717

18+
#include <cmath>
1819
#include <functional>
20+
#include <type_traits>
1921

2022
#include "tensorflow/lite/kernels/internal/types.h"
2123

2224
namespace tflite {
2325

2426
namespace reference_ops {
2527

28+
// Default comparator for non-floating-point types (no NaN possible).
2629
template <typename T>
27-
std::function<bool(T, T)> GetComparefunction(bool is_arg_max) {
30+
typename std::enable_if<!std::is_floating_point<T>::value,
31+
std::function<bool(T, T)>>::type
32+
GetComparefunction(bool is_arg_max) {
2833
if (is_arg_max) {
2934
return std::greater<T>();
3035
} else {
3136
return std::less<T>();
3237
}
3338
}
3439

40+
// NaN-aware comparator for floating-point types.
41+
// Matches TensorFlow eager semantics: NaN is treated as "less than any finite
42+
// value" for ArgMax and "greater than any finite value" for ArgMin. A NaN
43+
// candidate never replaces anything; a finite candidate always replaces a NaN
44+
// accumulator. For all-NaN inputs the first index is returned.
45+
template <typename T>
46+
typename std::enable_if<std::is_floating_point<T>::value,
47+
std::function<bool(T, T)>>::type
48+
GetComparefunction(bool is_arg_max) {
49+
if (is_arg_max) {
50+
return [](T candidate, T current) {
51+
return !std::isnan(candidate) &&
52+
(std::isnan(current) || candidate > current);
53+
};
54+
} else {
55+
return [](T candidate, T current) {
56+
return !std::isnan(candidate) &&
57+
(std::isnan(current) || candidate < current);
58+
};
59+
}
60+
}
61+
3562
template <typename T1, typename T2, typename T3, typename Cmp>
3663
void ArgMinMax(const RuntimeShape& input1_shape, const T1* input1_data,
3764
const T3* input2_data, const RuntimeShape& output_shape,

‎tensorflow/lite/kernels/internal/reference/select.h‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -237,10 +237,10 @@ void BroadcastSelect5DSlow(const RuntimeShape& input_condition_shape,
237237
const T* input_y_data,
238238
const RuntimeShape& output_shape, T* output_data) {
239239
ruy::profiler::ScopeLabel label("Select/BroadcastSelectSlow");
240-
TFLITE_DCHECK_LE(input_condition_shape.DimensionsCount(), 5);
241-
TFLITE_DCHECK_LE(input_x_shape.DimensionsCount(), 5);
242-
TFLITE_DCHECK_LE(input_y_shape.DimensionsCount(), 5);
243-
TFLITE_DCHECK_LE(output_shape.DimensionsCount(), 5);
240+
TFLITE_DCHECK_LE(input_condition_shape.DimensionsCount(), 8);
241+
TFLITE_DCHECK_LE(input_x_shape.DimensionsCount(), 8);
242+
TFLITE_DCHECK_LE(input_y_shape.DimensionsCount(), 8);
243+
TFLITE_DCHECK_LE(output_shape.DimensionsCount(), 8);
244244

245245
BroadcastSelectSimple(input_condition_shape, input_condition_data,
246246
input_x_shape, input_x_data, input_y_shape,

‎tensorflow/lite/kernels/kernel_util.cc‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -599,8 +599,18 @@ bool HasUnspecifiedDimension(const TfLiteTensor* tensor) {
599599
TfLiteStatus CheckedShapeProduct(TfLiteContext* context,
600600
std::initializer_list<int> dims,
601601
const char* error_message, size_t& product) {
602+
return CheckedShapeProduct(context, dims.begin(), dims.size(), error_message,
603+
product);
604+
}
605+
606+
TfLiteStatus CheckedShapeProduct(TfLiteContext* context, const int* dims,
607+
int count, const char* error_message,
608+
size_t& product) {
609+
TF_LITE_ENSURE(context, count >= 0);
610+
TF_LITE_ENSURE(context, dims != nullptr || count == 0);
602611
size_t checked_count = 1;
603-
for (const int d : dims) {
612+
for (int i = 0; i < count; ++i) {
613+
const int d = dims[i];
604614
TF_LITE_ENSURE_MSG(context, d >= 0, "Encountered a negative dimension.");
605615
TF_LITE_ENSURE_MSG(
606616
context,

‎tensorflow/lite/kernels/kernel_util.h‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -360,6 +360,20 @@ TfLiteStatus CheckedShapeProduct(TfLiteContext* context,
360360
* the dimensions is negative or if the product overflows.
361361
* @param context The context to use for error reporting.
362362
* @param dims The dimensions to multiply.
363+
* @param count The length of the dims array.
364+
* @param error_message The error message to use if an error is encountered.
365+
* @param product The output parameter to store the product.
366+
*/
367+
TfLiteStatus CheckedShapeProduct(TfLiteContext* context, const int* dims,
368+
int count, const char* error_message,
369+
size_t& product);
370+
371+
/**
372+
* Calculates the product of the given dimensions. Returns an error if any of
373+
* the dimensions is negative or if the product overflows. (Same as above
374+
* function with dims built on the fly)
375+
* @param context The context to use for error reporting.
376+
* @param dims The dimensions to multiply.
363377
* @param error_message The error message to use if an error is encountered.
364378
* @param product The output parameter to store the product.
365379
*/

‎tensorflow/lite/tools/flatbuffer_utils_test.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,9 @@
1818
import subprocess
1919
import sys
2020

21-
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema # pylint:disable=g-direct-tensorflow-import
22-
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import flatbuffer_utils
23-
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils
21+
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema # pylint:disable=g-direct-tensorflow-import
22+
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import flatbuffer_utils
23+
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils
2424
from tensorflow.python.framework import test_util
2525
from tensorflow.python.platform import test
2626

‎tensorflow/lite/tools/test_utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
"""
1919

2020
import flatbuffers
21-
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema_fb
21+
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema_fb
2222

2323
TFLITE_SCHEMA_VERSION = 3
2424

‎tensorflow/lite/tools/visualize_test.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616
import os
1717
import re
1818

19-
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils
20-
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import visualize
19+
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils
20+
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import visualize
2121
from tensorflow.python.framework import test_util
2222
from tensorflow.python.platform import test
2323

0 commit comments

Comments
 (0)