Add const-tensor checks (#3141)

@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

Author: ddavis-2015

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);

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.

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);

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 =

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);

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) {

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, &params, sizeof(StridedSliceParams));
   return CheckOutputSize(context, &op_context);

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 <cstdint>
 
 #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),

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<int32_t>(op_context.perm);

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,