From 36f3bce53bbb025e4ad6179efe82d13dd8c34c7e Mon Sep 17 00:00:00 2001 From: ddavis-2015 Date: Thu, 10 Jul 2025 17:51:42 -0700 Subject: [PATCH] Add const-tensor checks @tensorflow/micro Add consistent const-tensor checks and error messages across these kernels: TRANSPOSE STRIDED_SLICE FILL BROADCAST_TO EXPAND_DIMS Modify associated kernel unit tests. bug=fixes #3140 --- tensorflow/lite/micro/kernels/broadcast_to.cc | 5 ++++- .../lite/micro/kernels/broadcast_to_test.cc | 4 +++- tensorflow/lite/micro/kernels/expand_dims.cc | 8 +++---- .../lite/micro/kernels/expand_dims_test.cc | 4 +++- tensorflow/lite/micro/kernels/fill.cc | 13 ++++++------ tensorflow/lite/micro/kernels/fill_test.cc | 21 ++++++++++++------- .../micro/kernels/strided_slice_common.cc | 8 ++++++- .../lite/micro/kernels/strided_slice_test.cc | 16 +++++++++++++- .../lite/micro/kernels/transpose_common.cc | 2 ++ .../lite/micro/kernels/transpose_test.cc | 3 +++ 10 files changed, 59 insertions(+), 25 deletions(-) diff --git a/tensorflow/lite/micro/kernels/broadcast_to.cc b/tensorflow/lite/micro/kernels/broadcast_to.cc index 61deaa31b8f..d4fabd94996 100644 --- a/tensorflow/lite/micro/kernels/broadcast_to.cc +++ b/tensorflow/lite/micro/kernels/broadcast_to.cc @@ -1,4 +1,4 @@ -/* Copyright 2022 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2025 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -95,6 +95,9 @@ TfLiteStatus BroadcastToPrepare(TfLiteContext* context, TfLiteNode* node) { // the same as TFLite. TF_LITE_ENSURE(context, input->type != kTfLiteString); + TF_LITE_ENSURE_MSG(context, IsConstantTensor(shape), + "Non-constant >shape< tensor is not supported"); + TF_LITE_ENSURE_STATUS(ValidateOutputTensor(context, input, shape, output)); micro_context->DeallocateTempTfLiteTensor(input); micro_context->DeallocateTempTfLiteTensor(shape); diff --git a/tensorflow/lite/micro/kernels/broadcast_to_test.cc b/tensorflow/lite/micro/kernels/broadcast_to_test.cc index 09a97567720..a3d1dc93711 100644 --- a/tensorflow/lite/micro/kernels/broadcast_to_test.cc +++ b/tensorflow/lite/micro/kernels/broadcast_to_test.cc @@ -1,4 +1,4 @@ -/* Copyright 2022 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2025 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -44,6 +44,8 @@ tflite::micro::KernelRunner CreateBroadcastToTestRunner( tensors[0] = CreateTensor(input_data, IntArrayFromInts(input_shape)); tensors[1] = CreateTensor(dims_data, IntArrayFromInts(dims_shape)); + // shape must be a const tensor + tensors[1].allocation_type = kTfLiteMmapRo; tensors[2] = CreateTensor(output_data, IntArrayFromInts(output_shape)); // The output type matches the value type. diff --git a/tensorflow/lite/micro/kernels/expand_dims.cc b/tensorflow/lite/micro/kernels/expand_dims.cc index d47b42cbe0c..55eee67d8f4 100644 --- a/tensorflow/lite/micro/kernels/expand_dims.cc +++ b/tensorflow/lite/micro/kernels/expand_dims.cc @@ -1,4 +1,4 @@ -/* Copyright 2021 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2025 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -99,10 +99,8 @@ TfLiteStatus ExpandDimsPrepare(TfLiteContext* context, TfLiteNode* node) { micro_context->AllocateTempOutputTensor(node, kOutputTensor); TF_LITE_ENSURE(context, output != nullptr); output->type = input->type; - if (IsDynamicTensor(axis)) { - MicroPrintf("DynamicTensor is not yet supported by Expand_Dims."); - return kTfLiteError; - } + TF_LITE_ENSURE_MSG(context, IsConstantTensor(axis), + "Non-constant >axis< tensor is not supported"); TF_LITE_ENSURE_OK(context, VerifyTensorDim(context, input, axis, output)); micro_context->DeallocateTempTfLiteTensor(input); diff --git a/tensorflow/lite/micro/kernels/expand_dims_test.cc b/tensorflow/lite/micro/kernels/expand_dims_test.cc index 39a83b57471..c147011637c 100644 --- a/tensorflow/lite/micro/kernels/expand_dims_test.cc +++ b/tensorflow/lite/micro/kernels/expand_dims_test.cc @@ -1,4 +1,4 @@ -/* Copyright 2021 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2025 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -57,6 +57,8 @@ micro::KernelRunner CreateExpandDimsKernelRunner( tensors[kDimsTensorIndex] = CreateTensor(input_data, in_dims); tensors[kAxisTensorIndex] = CreateTensor(axis_data, ax_dims); + // axis must be a const tensor + tensors[kAxisTensorIndex].allocation_type = kTfLiteMmapRo; tensors[kOutputTensorIndex] = CreateTensor(output_data, out_dims, true); TfLiteIntArray* inputs_array = diff --git a/tensorflow/lite/micro/kernels/fill.cc b/tensorflow/lite/micro/kernels/fill.cc index 1486fcb8ee7..e07c4609648 100644 --- a/tensorflow/lite/micro/kernels/fill.cc +++ b/tensorflow/lite/micro/kernels/fill.cc @@ -1,4 +1,4 @@ -/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2025 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -87,12 +87,11 @@ TfLiteStatus FillPrepare(TfLiteContext* context, TfLiteNode* node) { // The dimension of the output tensor is known in model already. TFLITE_DCHECK(output->dims != nullptr); - if (dims->data.data != nullptr) { - // When the dims tensor is specified in model already (i.e. is not an - // activation tensor), the dims tensor must match the output tensor shape. - // As a byproduct, ensures the dims tensor is of an integer type. - TF_LITE_ENSURE_OK(context, EnsureEq(context, output->dims, dims)); - } + TF_LITE_ENSURE_MSG(context, IsConstantTensor(dims), + "Non-constant >dims< tensor is not supported"); + // The dims tensor must match the output tensor shape. + // As a byproduct, ensures the dims tensor is of an integer type. + TF_LITE_ENSURE_OK(context, EnsureEq(context, output->dims, dims)); micro_context->DeallocateTempTfLiteTensor(dims); micro_context->DeallocateTempTfLiteTensor(value); diff --git a/tensorflow/lite/micro/kernels/fill_test.cc b/tensorflow/lite/micro/kernels/fill_test.cc index 703554024b1..1df370350d1 100644 --- a/tensorflow/lite/micro/kernels/fill_test.cc +++ b/tensorflow/lite/micro/kernels/fill_test.cc @@ -1,4 +1,4 @@ -/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2025 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -44,6 +44,10 @@ tflite::micro::KernelRunner CreateFillTestRunner( static TfLiteTensor tensors[3]; tensors[0] = CreateTensor(dims_data, IntArrayFromInts(dims_shape)); + if (dims_data != nullptr) { + // dims must be a const tensor + tensors[0].allocation_type = kTfLiteMmapRo; + } tensors[1] = CreateTensor(value_data, IntArrayFromInts(value_shape)); tensors[2] = CreateTensor(output_data, IntArrayFromInts(output_shape)); @@ -154,15 +158,13 @@ TF_LITE_MICRO_TEST(FillInt8Int32Dims) { output_data); } -// Verify the FILL still works when the input dims tensor is an activation -// tensor (i.e. has not prepopulated value). Fill a 2x2x2 tensor with a int8 -// scalar value. -TF_LITE_MICRO_TEST(FillInt8NoInputDimsData) { +TF_LITE_MICRO_TEST(FillInt8NonConstDimsTensorFail) { constexpr int kDim1 = 2; constexpr int kDim2 = 2; constexpr int kDim3 = 2; - // The dims tensor with unknown data. Note that shape is always known. + // Simulate the dims tensor with dynamic data. Note that shape is always + // known. int dims_shape[] = {1, 3}; int32_t* dims_data = nullptr; @@ -172,8 +174,11 @@ TF_LITE_MICRO_TEST(FillInt8NoInputDimsData) { int output_shape[] = {3, kDim1, kDim2, kDim3}; int8_t output_data[kDim1 * kDim2 * kDim3]; - TestFill(dims_shape, dims_data, value_shape, value_data, output_shape, - output_data); + tflite::micro::KernelRunner runner = + CreateFillTestRunner(dims_shape, dims_data, value_shape, value_data, + output_shape, output_data); + + TF_LITE_MICRO_EXPECT_EQ(runner.InitAndPrepare(), kTfLiteError); } TF_LITE_MICRO_TEST(FillFloatInt32Dims) { diff --git a/tensorflow/lite/micro/kernels/strided_slice_common.cc b/tensorflow/lite/micro/kernels/strided_slice_common.cc index 165e1f39039..3d2860b637c 100644 --- a/tensorflow/lite/micro/kernels/strided_slice_common.cc +++ b/tensorflow/lite/micro/kernels/strided_slice_common.cc @@ -1,4 +1,4 @@ -/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2025 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -141,6 +141,12 @@ TfLiteStatus StridedSlicePrepare(TfLiteContext* context, TfLiteNode* node) { StridedSliceContext op_context(context, node); TF_LITE_ENSURE_MSG(context, op_context.dims <= kMaxDim, "input dim should not exceed 4"); + TF_LITE_ENSURE_MSG(context, IsConstantTensor(op_context.begin), + "Non-constant >begin< tensor is not supported"); + TF_LITE_ENSURE_MSG(context, IsConstantTensor(op_context.end), + "Non-constant >end< tensor is not supported"); + TF_LITE_ENSURE_MSG(context, IsConstantTensor(op_context.strides), + "Non-constant >strides< tensor is not supported"); auto params = BuildStridedSliceParams(&op_context); memcpy(op_params, ¶ms, sizeof(StridedSliceParams)); return CheckOutputSize(context, &op_context); diff --git a/tensorflow/lite/micro/kernels/strided_slice_test.cc b/tensorflow/lite/micro/kernels/strided_slice_test.cc index 16c3d9c852e..5c799feff3b 100644 --- a/tensorflow/lite/micro/kernels/strided_slice_test.cc +++ b/tensorflow/lite/micro/kernels/strided_slice_test.cc @@ -1,4 +1,4 @@ -/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2025 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -12,6 +12,8 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "tensorflow/lite/micro/kernels/strided_slice.h" + #include #include "tensorflow/lite/c/builtin_op_data.h" @@ -82,6 +84,12 @@ void TestStridedSliceFloat(int* input_shape, int* begin_shape, int* end_shape, CreateTensor(strides_data, strides_dims), CreateTensor(output_data, output_dims), }; + // begin must be a const tensor + tensors[kStridedSliceBeginTensor].allocation_type = kTfLiteMmapRo; + // end must be a const tensor + tensors[kStridedSliceEndTensor].allocation_type = kTfLiteMmapRo; + // strides must be a const tensor + tensors[kStridedSliceStridesTensor].allocation_type = kTfLiteMmapRo; ValidateStridedSliceGoldens(tensors, tensors_size, expected_output, output_data, ElementCount(*output_dims), @@ -116,6 +124,12 @@ void TestStridedSliceQuantized(int* input_shape, int* begin_shape, CreateTensor(strides_data, strides_dims), CreateQuantizedTensor(output_data, output_dims, 1.0, zero_point), }; + // begin must be a const tensor + tensors[kStridedSliceBeginTensor].allocation_type = kTfLiteMmapRo; + // end must be a const tensor + tensors[kStridedSliceEndTensor].allocation_type = kTfLiteMmapRo; + // strides must be a const tensor + tensors[kStridedSliceStridesTensor].allocation_type = kTfLiteMmapRo; ValidateStridedSliceGoldens(tensors, tensors_size, expected_output, output_data, ElementCount(*output_dims), diff --git a/tensorflow/lite/micro/kernels/transpose_common.cc b/tensorflow/lite/micro/kernels/transpose_common.cc index 1a99f3b8b2c..cb5dc5e4a99 100644 --- a/tensorflow/lite/micro/kernels/transpose_common.cc +++ b/tensorflow/lite/micro/kernels/transpose_common.cc @@ -34,6 +34,8 @@ TfLiteStatus TransposePrepare(TfLiteContext* context, TfLiteNode* node) { "Transpose op only supports 1D-5D input arrays."); TF_LITE_ENSURE_TYPES_EQ(context, op_context.input->type, op_context.output->type); + TF_LITE_ENSURE_MSG(context, IsConstantTensor(op_context.perm), + "Non-constant >perm< tensor is not supported"); int dims = NumDimensions(op_context.input); const int32_t* perm_data = GetTensorData(op_context.perm); diff --git a/tensorflow/lite/micro/kernels/transpose_test.cc b/tensorflow/lite/micro/kernels/transpose_test.cc index 14f433795e9..28b10238351 100644 --- a/tensorflow/lite/micro/kernels/transpose_test.cc +++ b/tensorflow/lite/micro/kernels/transpose_test.cc @@ -18,6 +18,7 @@ limitations under the License. #include "tensorflow/lite/c/builtin_op_data.h" #include "tensorflow/lite/c/common.h" #include "tensorflow/lite/micro/kernels/kernel_runner.h" +#include "tensorflow/lite/micro/kernels/transpose.h" #include "tensorflow/lite/micro/micro_utils.h" #include "tensorflow/lite/micro/test_helpers.h" #include "tensorflow/lite/micro/testing/micro_test.h" @@ -123,6 +124,8 @@ void TestTranspose(int* input_dims_data, T* input_data, int* output_dims_data, CreateTensor(params->perm, perm_dims), CreateTensor(output_data, output_dims), }; + // perm must be a const tensor + tensors[kTransposePermTensor].allocation_type = kTfLiteMmapRo; TF_LITE_MICRO_EXPECT_EQ( kTfLiteOk, ValidateTranspose(tensors, tensors_size, expected_output_data,