CVE-2020-15211: Out-of-bounds Write (CWE-787)
value: In TensorFlow Lite before versions 1.15.4, 2.0.3, 2.1.2, 2.2.1 and 2.3.1, saved models in the flatbuffer format use a double indexing scheme: a model has a set of subgraphs, each subgraph has a set of operators and each operator has a set of input/output tensors. The flatbuffer format uses indices for the tensors, indexing into an array of tensors that is owned by the subgraph. This results in a pattern of double array indexing when trying to get the data of each tensor. However, some operators can have some tensors be optional. To handle this scenario, the flatbuffer model uses a negative `-1` value as index for these tensors. This results in special casing during validation at model loading time. Unfortunately, this means that the `-1` index is a valid tensor index for any operator, including those that don't expect optional inputs and including for output tensors. Thus, this allows writing and reading from outside the bounds of heap allocated arrays, although only at a specific offset from the start of these arrays. This results in both read and write gadgets, albeit very limited in scope. The issue is patched in several commits (46d5b0852, 00302787b7, e11f5558, cd31fd0ce, 1970c21, and fff2c83), and is released in TensorFlow versions 1.15.4, 2.0.3, 2.1.2, 2.2.1, or 2.3.1. A potential workaround would be to add a custom `Verifier` to the model loading code to ensure that only operators which accept optional inputs use the `-1` special value and only for the tensors that they expect to be optional. Since this allow-list type approach is erro-prone, we advise upgrading to the patched code.
CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 4.8
Published: 2020-09-25T19:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/lite/delegates/delegate_test.cc, tensorflow/lite/delegates/gpu/common/model_builder.cc, tensorflow/lite/experimental/delegates/coreml/builders/convolution_op_builder.cc, tensorflow/lite/experimental/delegates/coreml/builders/fully_connected_op_builder.cc, tensorflow/lite/experimental/delegates/coreml/builders/pad_op_builder.cc, tensorflow/lite/experimental/delegates/coreml/builders/reshape_op_builder.cc, tensorflow/lite/experimental/kernels/ctc_beam_search_decoder.cc, tensorflow/lite/experimental/kernels/unidirectional_sequence_gru.cc, tensorflow/lite/experimental/microfrontend/audio_microfrontend.cc, tensorflow/lite/interpreter_test.cc, tensorflow/lite/java/src/test/native/interpreter_test_jni.cc, tensorflow/lite/kernels/hashtable/hashtable.cc, tensorflow/lite/kernels/hashtable/hashtable_find.cc, tensorflow/lite/kernels/hashtable/hashtable_import.cc, tensorflow/lite/kernels/hashtable/hashtable_size.cc, tensorflow/lite/kernels/kernel_util.cc, tensorflow/lite/kernels/kernel_util.h, tensorflow/lite/profiling/profile_summarizer_test.cc, tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
Commit date: 2020-09-18 13:53:51 -0700
Commit message: [tflite]: Insert `nullptr` checks when obtaining tensors.

As part of ongoing refactoring, `tflite::GetInput`, `tflite::GetOutput`, `tflite::GetTemporary` and `tflite::GetIntermediates` will return `nullptr` in some cases. Hence, we insert the `nullptr` checks on all usages.

We also insert `nullptr` checks on usages of `tflite::GetVariableInput` and `tflite::GetOptionalInputTensor` but only in the cases where there is no obvious check that `nullptr` is acceptable (that is, we only insert the check for the output of these two functions if the tensor is accessed as if it is always not `nullptr`).

PiperOrigin-RevId: 332517854
Change-Id: Ic27221dd1f0fbe302f311c2fe5a846ed8ff02016

Diff: diff --git a/tensorflow/lite/delegates/delegate_test.cc b/tensorflow/lite/delegates/delegate_test.cc
index aac44b0ff51..c5a31374887 100644
--- a/tensorflow/lite/delegates/delegate_test.cc
+++ b/tensorflow/lite/delegates/delegate_test.cc
@@ -40,13 +40,16 @@ TfLiteRegistration AddOpRegistration() {
   reg.custom_name = "my_add";
   reg.builtin_code = tflite::BuiltinOperator_CUSTOM;
 
   reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
     // Set output size to input size
-    const TfLiteTensor* input1 = GetInput(context, node, 0);
-    const TfLiteTensor* input2 = GetInput(context, node, 1);
-    TfLiteTensor* output = GetOutput(context, node, 0);
+    const TfLiteTensor* input1;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input1));
+    const TfLiteTensor* input2;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &input2));
+    TfLiteTensor* output;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
 
     TF_LITE_ENSURE_EQ(context, input1->dims->size, input2->dims->size);
     for (int i = 0; i < input1->dims->size; ++i) {
       TF_LITE_ENSURE_EQ(context, input1->dims->data[i], input2->dims->data[i]);
     }
@@ -56,17 +59,20 @@ TfLiteRegistration AddOpRegistration() {
     return kTfLiteOk;
   };
 
   reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
     // Copy input data to output data.
-    const TfLiteTensor* a0 = GetInput(context, node, 0);
+    const TfLiteTensor* a0;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
     TF_LITE_ENSURE(context, a0);
     TF_LITE_ENSURE(context, a0->data.f);
-    const TfLiteTensor* a1 = GetInput(context, node, 1);
+    const TfLiteTensor* a1;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &a1));
     TF_LITE_ENSURE(context, a1);
     TF_LITE_ENSURE(context, a1->data.f);
-    TfLiteTensor* out = GetOutput(context, node, 0);
+    TfLiteTensor* out;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out));
     TF_LITE_ENSURE(context, out);
     TF_LITE_ENSURE(context, out->data.f);
     int num = a0->dims->data[0];
     for (int i = 0; i < num; i++) {
       out->data.f[i] = a0->data.f[i] + a1->data.f[i];
@@ -265,11 +271,12 @@ class TestDelegate : public ::testing::Test {
             a1 = GetInput(context, node, 1);
           } else {
             a0 = GetInput(context, node, 0);
             a1 = a0;
           }
-          TfLiteTensor* out = GetOutput(context, node, 0);
+          TfLiteTensor* out;
+          TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out));
           int num = 1;
           for (int i = 0; i < a0->dims->size; ++i) {
             num *= a0->dims->data[i];
           }
           for (int i = 0; i < num; i++) {
@@ -287,12 +294,14 @@ class TestDelegate : public ::testing::Test {
       // testing parameters.
       if (automatic_shape_propagation_) {
         reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
           // Shapes should already by propagated by the runtime, just need to
           // check.
-          const TfLiteTensor* input1 = GetInput(context, node, 0);
-          TfLiteTensor* output = GetOutput(context, node, 0);
+          const TfLiteTensor* input1;
+          TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input1));
+          TfLiteTensor* output;
+          TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
           const int input_dims_size = input1->dims->size;
           TF_LITE_ENSURE(context, output->dims->size == input_dims_size);
           for (int i = 0; i < input_dims_size; ++i) {
             TF_LITE_ENSURE(context,
                            output->dims->data[i] == input1->dims->data[i]);
@@ -313,11 +322,12 @@ class TestDelegate : public ::testing::Test {
             input2 = GetInput(context, node, 1);
           } else {
             input1 = GetInput(context, node, 0);
             input2 = input1;
           }
-          TfLiteTensor* output = GetOutput(context, node, 0);
+          TfLiteTensor* output;
+          TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
 
           TF_LITE_ENSURE_STATUS(context->ResizeTensor(
               context, output, TfLiteIntArrayCopy(input1->dims)));
           return kTfLiteOk;
         };
@@ -1167,15 +1177,18 @@ class TestDelegateWithDynamicTensors : public ::testing::Test {
   static TfLiteRegistration DynamicCopyOpRegistration() {
     TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
 
     reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
       // Output 0 is dynamic
-      TfLiteTensor* output0 = GetOutput(context, node, 0);
+      TfLiteTensor* output0;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output0));
       SetTensorToDynamic(output0);
       // Output 1 has the same shape as input.
-      const TfLiteTensor* input = GetInput(context, node, 0);
-      TfLiteTensor* output1 = GetOutput(context, node, 1);
+      const TfLiteTensor* input;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
+      TfLiteTensor* output1;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 1, &output1));
       TF_LITE_ENSURE_STATUS(context->ResizeTensor(
           context, output1, TfLiteIntArrayCopy(input->dims)));
       return kTfLiteOk;
     };
 
@@ -1191,15 +1204,18 @@ class TestDelegateWithDynamicTensors : public ::testing::Test {
 
     reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
       // If tensors are resized, the runtime should propagate shapes
       // automatically if correct flag is set. Ensure values are correct.
       // Output 0 should be dynamic.
-      TfLiteTensor* output0 = GetOutput(context, node, 0);
+      TfLiteTensor* output0;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output0));
       TF_LITE_ENSURE(context, IsDynamicTensor(output0));
       // Output 1 has the same shape as input.
-      const TfLiteTensor* input = GetInput(context, node, 0);
-      TfLiteTensor* output1 = GetOutput(context, node, 1);
+      const TfLiteTensor* input;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
+      TfLiteTensor* output1;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 1, &output1));
       TF_LITE_ENSURE(context, input->dims->size == output1->dims->size);
       TF_LITE_ENSURE(context, input->dims->data[0] == output1->dims->data[0]);
       return kTfLiteOk;
     };
 
diff --git a/tensorflow/lite/delegates/gpu/common/model_builder.cc b/tensorflow/lite/delegates/gpu/common/model_builder.cc
index 9a394f8f6c5..e56afa30497 100644
--- a/tensorflow/lite/delegates/gpu/common/model_builder.cc
+++ b/tensorflow/lite/delegates/gpu/common/model_builder.cc
@@ -1164,10 +1164,13 @@ class MulOperationParser : public TFLiteOperationParser {
     if (tflite_node->inputs->size != 2) {
       return absl::UnimplementedError("MUL requires two input tensors.");
     }
     auto input0 = tflite::GetInput(context, tflite_node, 0);
     auto input1 = tflite::GetInput(context, tflite_node, 1);
+    if (input0 == nullptr || input1 == nullptr) {
+      return absl::InvalidArgumentError("At least one input tensor is null");
+    }
     if (input0->dims->size == input1->dims->size) {
       // this code checks that at least one input of Mul not smaller in all
       // dimensions. Sometimes Mul used for matrix-vector multiplication that we
       // currently don't support. For example input0 HWC(1, 256, 1), input1
       // HWC(1, 1, 256) -> output HWC (1, 256, 256). In this case it can be
@@ -1378,11 +1381,14 @@ class PadOperationParser : public TFLiteOperationParser {
     }
     RETURN_IF_ERROR(CheckMaxSupportedOpVersion(registration, 2));
     RETURN_IF_ERROR(CheckInputsOutputs(context, tflite_node,
                                        /*runtime_inputs=*/1, /*outputs=*/1));
     RETURN_IF_ERROR(CheckTensorIsAvailable(context, tflite_node, 1));
-    auto pad_tensor = tflite::GetInput(context, tflite_node, 1);
+    const TfLiteTensor* pad_tensor = tflite::GetInput(context, tflite_node, 1);
+    if (pad_tensor == nullptr) {
+      return absl::InvalidArgumentError("Padding tensor was null");
+    }
     if (pad_tensor->dims->size != 2) {
       return absl::InvalidArgumentError(absl::StrCat(
           "Invalid paddings tensor dimension: expected 2 dim, got ",
           pad_tensor->dims->size, " dim"));
     }
diff --git a/tensorflow/lite/experimental/delegates/coreml/builders/convolution_op_builder.cc b/tensorflow/lite/experimental/delegates/coreml/builders/convolution_op_builder.cc
index 8dbc18722f1..e6c3f892c3b 100644
--- a/tensorflow/lite/experimental/delegates/coreml/builders/convolution_op_builder.cc
+++ b/tensorflow/lite/experimental/delegates/coreml/builders/convolution_op_builder.cc
@@ -326,11 +326,13 @@ bool IsConvolutionOpSupported(const TfLiteRegistration* registration,
   }
 
   const int kOutputShapeTensor = 0;  // Only used for TransposeConv
   const int kWeightTensor = 1;
   const int kBiasTensor = 2;  // Only used for non-TransposeConv
-  const TfLiteTensor* weights = GetInput(context, node, kWeightTensor);
+  const TfLiteTensor* weights;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kWeightTensor, &weights));
   const int max_kernel_size = 16384;
   if (!IsConstantTensor(weights)) {
     return false;
   }
   if (weights->dims->data[1] > max_kernel_size ||
diff --git a/tensorflow/lite/experimental/delegates/coreml/builders/fully_connected_op_builder.cc b/tensorflow/lite/experimental/delegates/coreml/builders/fully_connected_op_builder.cc
index 376830922c9..4179ed0445e 100644
--- a/tensorflow/lite/experimental/delegates/coreml/builders/fully_connected_op_builder.cc
+++ b/tensorflow/lite/experimental/delegates/coreml/builders/fully_connected_op_builder.cc
@@ -151,12 +151,14 @@ bool IsFullyConnectedOpSupported(const TfLiteRegistration* registration,
   const int kBias = 2;
 
   if (fc_params->weights_format != kTfLiteFullyConnectedWeightsFormatDefault) {
     return false;
   }
-  const TfLiteTensor* input = GetInput(context, node, kInput);
-  const TfLiteTensor* weights = GetInput(context, node, kWeights);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
+  const TfLiteTensor* weights;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kWeights, &weights));
 
   if (!IsFloatType(input->type)) {
     return false;
   }
   if (!IsFloatType(weights->type) || !IsConstantTensor(weights)) {
@@ -167,11 +169,12 @@ bool IsFullyConnectedOpSupported(const TfLiteRegistration* registration,
   if (input->dims->data[input->dims->size - 1] != NumElements(input)) {
     return false;
   }
 
   if (node->inputs->size > 2) {
-    const TfLiteTensor* bias = GetInput(context, node, kBias);
+    const TfLiteTensor* bias;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBias, &bias));
     if (!IsFloatType(bias->type) || !IsConstantTensor(bias)) {
       return false;
     }
   }
 
diff --git a/tensorflow/lite/experimental/delegates/coreml/builders/pad_op_builder.cc b/tensorflow/lite/experimental/delegates/coreml/builders/pad_op_builder.cc
index d8ef4f61ddb..36c50d816df 100644
--- a/tensorflow/lite/experimental/delegates/coreml/builders/pad_op_builder.cc
+++ b/tensorflow/lite/experimental/delegates/coreml/builders/pad_op_builder.cc
@@ -95,11 +95,12 @@ OpBuilder* CreateMirrorPadOpBuilder(GraphBuilder* graph_builder) {
 }
 
 bool IsPadOpSupported(const TfLiteRegistration* registration,
                       const TfLiteNode* node, TfLiteContext* context) {
   // padding is d x 2 tensor, where d is the dimension of input.
-  const TfLiteTensor* padding = GetInput(context, node, 1);
+  const TfLiteTensor* padding;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &padding));
   if (!IsConstantTensor(padding)) {
     TF_LITE_KERNEL_LOG(context,
                        "%s: Only constant padding is supported for PAD.",
                        padding->name);
     return false;
diff --git a/tensorflow/lite/experimental/delegates/coreml/builders/reshape_op_builder.cc b/tensorflow/lite/experimental/delegates/coreml/builders/reshape_op_builder.cc
index b7b78653d36..39b15e9349a 100644
--- a/tensorflow/lite/experimental/delegates/coreml/builders/reshape_op_builder.cc
+++ b/tensorflow/lite/experimental/delegates/coreml/builders/reshape_op_builder.cc
@@ -124,11 +124,12 @@ bool IsReshapeOpSupported(const TfLiteRegistration* registration,
         reinterpret_cast<TfLiteReshapeParams*>(node->builtin_data);
     return params->num_dimensions == 3 || params->num_dimensions == 4;
   }
 
   const int kShapeTensor = 1;
-  const auto* shape = GetInput(context, node, kShapeTensor);
+  const TfLiteTensor* shape;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kShapeTensor, &shape));
   if (shape->allocation_type != kTfLiteMmapRo) {
     TF_LITE_KERNEL_LOG(context, "Reshape has non-const shape.");
     return false;
   }
   const bool is_shape_tensor =
diff --git a/tensorflow/lite/experimental/kernels/ctc_beam_search_decoder.cc b/tensorflow/lite/experimental/kernels/ctc_beam_search_decoder.cc
index c5e019fc2ee..9b7d731c5f8 100644
--- a/tensorflow/lite/experimental/kernels/ctc_beam_search_decoder.cc
+++ b/tensorflow/lite/experimental/kernels/ctc_beam_search_decoder.cc
@@ -60,37 +60,46 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE(context, option->beam_width >= top_paths);
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
   // The outputs should be top_paths * 3 + 1.
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 3 * top_paths + 1);
 
-  const TfLiteTensor* inputs = GetInput(context, node, kInputsTensor);
+  const TfLiteTensor* inputs;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputsTensor, &inputs));
   TF_LITE_ENSURE_EQ(context, NumDimensions(inputs), 3);
   // TensorFlow only supports float.
   TF_LITE_ENSURE_EQ(context, inputs->type, kTfLiteFloat32);
   const int batch_size = SizeOfDimension(inputs, 1);
 
-  const TfLiteTensor* sequence_length =
-      GetInput(context, node, kSequenceLengthTensor);
+  const TfLiteTensor* sequence_length;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kSequenceLengthTensor,
+                                          &sequence_length));
   TF_LITE_ENSURE_EQ(context, NumDimensions(sequence_length), 1);
   TF_LITE_ENSURE_EQ(context, NumElements(sequence_length), batch_size);
   // TensorFlow only supports int32.
   TF_LITE_ENSURE_EQ(context, sequence_length->type, kTfLiteInt32);
 
   // Resize decoded outputs.
   // Do not resize indices & values cause we don't know the values yet.
   for (int i = 0; i < top_paths; ++i) {
-    TfLiteTensor* indices = GetOutput(context, node, i);
+    TfLiteTensor* indices;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, i, &indices));
     SetTensorToDynamic(indices);
-    TfLiteTensor* values = GetOutput(context, node, i + top_paths);
+    TfLiteTensor* values;
+    TF_LITE_ENSURE_OK(context,
+                      GetOutputSafe(context, node, i + top_paths, &values));
     SetTensorToDynamic(values);
-    TfLiteTensor* output_shape = GetOutput(context, node, i + 2 * top_paths);
+    TfLiteTensor* output_shape;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, i + 2 * top_paths,
+                                             &output_shape));
     SetTensorToDynamic(output_shape);
   }
 
   // Resize log probability outputs.
-  TfLiteTensor* log_probability_output =
-      GetOutput(context, node, top_paths * 3);
+  TfLiteTensor* log_probability_output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, top_paths * 3,
+                                           &log_probability_output));
   TfLiteIntArray* log_probability_output_shape_array = TfLiteIntArrayCreate(2);
   log_probability_output_shape_array->data[0] = batch_size;
   log_probability_output_shape_array->data[1] = top_paths;
   return context->ResizeTensor(context, log_probability_output,
                                log_probability_output_shape_array);
@@ -125,17 +134,22 @@ TfLiteStatus StoreAllDecodedSequences(
 
   for (int p = 0; p < top_paths; ++p) {
     const int32_t p_num = num_entries[p];
 
     // Resize the decoded outputs.
-    TfLiteTensor* indices = GetOutput(context, node, p);
+    TfLiteTensor* indices;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, p, &indices));
     TF_LITE_ENSURE_OK(context, Resize(context, {p_num, 2}, indices));
 
-    TfLiteTensor* values = GetOutput(context, node, p + top_paths);
+    TfLiteTensor* values;
+    TF_LITE_ENSURE_OK(context,
+                      GetOutputSafe(context, node, p + top_paths, &values));
     TF_LITE_ENSURE_OK(context, Resize(context, {p_num}, values));
 
-    TfLiteTensor* decoded_shape = GetOutput(context, node, p + 2 * top_paths);
+    TfLiteTensor* decoded_shape;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, p + 2 * top_paths,
+                                             &decoded_shape));
     TF_LITE_ENSURE_OK(context, Resize(context, {2}, decoded_shape));
 
     int32_t max_decoded = 0;
     int32_t offset = 0;
 
@@ -159,13 +173,16 @@ TfLiteStatus StoreAllDecodedSequences(
   }
   return kTfLiteOk;
 }
 
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
-  const TfLiteTensor* inputs = GetInput(context, node, kInputsTensor);
-  const TfLiteTensor* sequence_length =
-      GetInput(context, node, kSequenceLengthTensor);
+  const TfLiteTensor* inputs;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputsTensor, &inputs));
+  const TfLiteTensor* sequence_length;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kSequenceLengthTensor,
+                                          &sequence_length));
   const CTCBeamSearchDecoderParams* option =
       reinterpret_cast<CTCBeamSearchDecoderParams*>(node->user_data);
 
   const int max_time = SizeOfDimension(inputs, 0);
   const int batch_size = SizeOfDimension(inputs, 1);
@@ -205,11 +222,13 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   optimized_ops::TTypes<float>::Flat input_chip_t(input_chip_t_data, dims);
 
   std::vector<std::vector<std::vector<int>>> best_paths(batch_size);
   std::vector<float> log_probs;
 
-  TfLiteTensor* log_probabilities = GetOutput(context, node, 3 * top_paths);
+  TfLiteTensor* log_probabilities;
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node, 3 * top_paths, &log_probabilities));
   float* log_probabilities_output = GetTensorData<float>(log_probabilities);
 
   // Assumption: the blank index is num_classes - 1
   for (int b = 0; b < batch_size; ++b) {
     auto& best_paths_b = best_paths[b];
diff --git a/tensorflow/lite/experimental/kernels/unidirectional_sequence_gru.cc b/tensorflow/lite/experimental/kernels/unidirectional_sequence_gru.cc
index 84f54b8d5aa..1fd6dabb52f 100644
--- a/tensorflow/lite/experimental/kernels/unidirectional_sequence_gru.cc
+++ b/tensorflow/lite/experimental/kernels/unidirectional_sequence_gru.cc
@@ -125,78 +125,94 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   TF_LITE_ENSURE_EQ(context, node->inputs->size, kInputNum);
   TF_LITE_ENSURE_EQ(context, node->outputs->size, kOutputNum);
 
   // input's dim = [n_time, n_batch, n_input]
-  const TfLiteTensor* input = GetInput(context, node, kInput);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
   TF_LITE_ENSURE_EQ(context, input->dims->size, 3);
   const int n_time = input->dims->data[0];
   const int n_batch = input->dims->data[1];
   const int n_input = input->dims->data[2];
 
   // input_state's dim = [n_batch, n_output]
-  const TfLiteTensor* input_state = GetInput(context, node, kInputState);
+  const TfLiteTensor* input_state;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputState, &input_state));
   TF_LITE_ENSURE_EQ(context, input_state->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, input_state->dims->data[0], n_batch);
   const int n_output = input_state->dims->data[1];
 
   // gate_weight' dim = [2 * n_output, n_input + n_output]
-  const TfLiteTensor* gate_weight = GetInput(context, node, kGateWeight);
+  const TfLiteTensor* gate_weight;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kGateWeight, &gate_weight));
   TF_LITE_ENSURE_EQ(context, gate_weight->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, gate_weight->dims->data[0], 2 * n_output);
   TF_LITE_ENSURE_EQ(context, gate_weight->dims->data[1], n_input + n_output);
 
   // gate_bias' dim = [2 * n_output]
-  const TfLiteTensor* gate_bias = GetInput(context, node, kGateBias);
+  const TfLiteTensor* gate_bias;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kGateBias, &gate_bias));
   TF_LITE_ENSURE_EQ(context, gate_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, gate_bias->dims->data[0], 2 * n_output);
 
   // candidate_weight' dim = [n_output, n_input + n_output]
-  const TfLiteTensor* candidate_weight =
-      GetInput(context, node, kCandidateWeight);
+  const TfLiteTensor* candidate_weight;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kCandidateWeight,
+                                          &candidate_weight));
   TF_LITE_ENSURE_EQ(context, candidate_weight->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, candidate_weight->dims->data[0], n_output);
   TF_LITE_ENSURE_EQ(context, candidate_weight->dims->data[1],
                     n_input + n_output);
 
   // candidate_bias' dim = [n_output]
-  const TfLiteTensor* candidate_bias = GetInput(context, node, kCandidateBias);
+  const TfLiteTensor* candidate_bias;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node, kCandidateBias, &candidate_bias));
   TF_LITE_ENSURE_EQ(context, candidate_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, candidate_bias->dims->data[0], n_output);
 
   // output's dim = [n_time, n_batch, n_output]
-  TfLiteTensor* output = GetOutput(context, node, kOutput);
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kOutput, &output));
   TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
   output_size->data[0] = n_time;
   output_size->data[1] = n_batch;
   output_size->data[2] = n_output;
   TF_LITE_ENSURE_OK(context,
                     context->ResizeTensor(context, output, output_size));
 
   // output_state's dim = [n_batch, n_output]
-  TfLiteTensor* output_state = GetOutput(context, node, kOutputState);
+  TfLiteTensor* output_state;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputState, &output_state));
   TF_LITE_ENSURE_OK(
       context, context->ResizeTensor(context, output_state,
                                      TfLiteIntArrayCopy(input_state->dims)));
 
   TfLiteIntArrayFree(node->temporaries);
   node->temporaries = TfLiteIntArrayCreate(kTemporaryNum);
 
   // activation's dim = [n_batch, 2 * n_output]
   node->temporaries->data[kActivation] = *scratch_tensor_index;
-  TfLiteTensor* activation = GetTemporary(context, node, kActivation);
+  TfLiteTensor* activation;
+  TF_LITE_ENSURE_OK(context,
+                    GetTemporarySafe(context, node, kActivation, &activation));
   activation->type = input->type;
   activation->allocation_type = kTfLiteArenaRw;
   TfLiteIntArray* activation_size = TfLiteIntArrayCreate(2);
   activation_size->data[0] = n_batch;
   activation_size->data[1] = 2 * n_output;
   TF_LITE_ENSURE_OK(
       context, context->ResizeTensor(context, activation, activation_size));
 
   // concat's dim  = [n_batch, n_input + n_output]
   node->temporaries->data[kConcat] = (*scratch_tensor_index) + kConcat;
-  TfLiteTensor* concat = GetTemporary(context, node, kConcat);
+  TfLiteTensor* concat;
+  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kConcat, &concat));
   concat->type = input->type;
   concat->allocation_type = kTfLiteArenaRw;
   TfLiteIntArray* concat_size = TfLiteIntArrayCreate(2);
   concat_size->data[0] = n_batch;
   concat_size->data[1] = n_input + n_output;
@@ -205,21 +221,37 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   return kTfLiteOk;
 }
 
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
-  const TfLiteTensor* input = GetInput(context, node, kInput);
-  const TfLiteTensor* input_state = GetInput(context, node, kInputState);
-  const TfLiteTensor* gate_weight = GetInput(context, node, kGateWeight);
-  const TfLiteTensor* gate_bias = GetInput(context, node, kGateBias);
-  const TfLiteTensor* candidate_weight =
-      GetInput(context, node, kCandidateWeight);
-  const TfLiteTensor* candidate_bias = GetInput(context, node, kCandidateBias);
-  TfLiteTensor* output = GetOutput(context, node, kOutput);
-  TfLiteTensor* output_state = GetOutput(context, node, kOutputState);
-  TfLiteTensor* activation = GetTemporary(context, node, kActivation);
-  TfLiteTensor* concat = GetTemporary(context, node, kConcat);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
+  const TfLiteTensor* input_state;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputState, &input_state));
+  const TfLiteTensor* gate_weight;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kGateWeight, &gate_weight));
+  const TfLiteTensor* gate_bias;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kGateBias, &gate_bias));
+  const TfLiteTensor* candidate_weight;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kCandidateWeight,
+                                          &candidate_weight));
+  const TfLiteTensor* candidate_bias;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node, kCandidateBias, &candidate_bias));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kOutput, &output));
+  TfLiteTensor* output_state;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputState, &output_state));
+  TfLiteTensor* activation;
+  TF_LITE_ENSURE_OK(context,
+                    GetTemporarySafe(context, node, kActivation, &activation));
+  TfLiteTensor* concat;
+  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kConcat, &concat));
   auto cpu_backend_context = CpuBackendContext::GetFromContext(context);
 
   if (gate_weight->type == kTfLiteFloat32) {
     GruImpl(input, input_state, gate_weight, gate_bias, candidate_weight,
             candidate_bias, output, output_state, activation, concat,
diff --git a/tensorflow/lite/experimental/microfrontend/audio_microfrontend.cc b/tensorflow/lite/experimental/microfrontend/audio_microfrontend.cc
index b8d89f69b26..090871da2ef 100644
--- a/tensorflow/lite/experimental/microfrontend/audio_microfrontend.cc
+++ b/tensorflow/lite/experimental/microfrontend/audio_microfrontend.cc
@@ -89,12 +89,15 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
       reinterpret_cast<TfLiteAudioMicrofrontendParams*>(node->user_data);
 
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
-  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
 
   TF_LITE_ENSURE_EQ(context, NumDimensions(input), 1);
 
   TF_LITE_ENSURE_EQ(context, input->type, kTfLiteInt16);
   output->type = kTfLiteInt32;
@@ -178,12 +181,15 @@ void GenerateFeatures(TfLiteAudioMicrofrontendParams* data,
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   auto* data =
       reinterpret_cast<TfLiteAudioMicrofrontendParams*>(node->user_data);
   FrontendReset(data->state);
 
-  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
 
   if (data->out_float) {
     GenerateFeatures<float>(data, input, output);
   } else {
     GenerateFeatures<int32>(data, input, output);
diff --git a/tensorflow/lite/interpreter_test.cc b/tensorflow/lite/interpreter_test.cc
index bd0f724a7bf..66728ea89e9 100644
--- a/tensorflow/lite/interpreter_test.cc
+++ b/tensorflow/lite/interpreter_test.cc
@@ -619,12 +619,14 @@ TfLiteRegistration GetPassthroughOpRegistration() {
     delete static_cast<int*>(buffer);
   };
   reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
     auto* first_new_tensor = static_cast<int*>(node->user_data);
 
-    const TfLiteTensor* tensor0 = GetInput(context, node, 0);
-    TfLiteTensor* tensor1 = GetOutput(context, node, 0);
+    const TfLiteTensor* tensor0;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &tensor0));
+    TfLiteTensor* tensor1;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &tensor1));
 
     TfLiteIntArray* newSize = TfLiteIntArrayCopy(tensor0->dims);
     TF_LITE_ENSURE_STATUS(context->ResizeTensor(context, tensor1, newSize));
 
     TfLiteIntArrayFree(node->temporaries);
@@ -644,11 +646,12 @@ TfLiteRegistration GetPassthroughOpRegistration() {
     TF_LITE_ENSURE_STATUS(setup_temporary(node->temporaries->data[1]));
 
     return kTfLiteOk;
   };
   reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
-    const TfLiteTensor* a0 = GetInput(context, node, 0);
+    const TfLiteTensor* a0;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
 
     auto populate = [&](int id) {
       TfLiteTensor* t = &context->tensors[id];
       int num = a0->dims->data[0];
       for (int i = 0; i < num; i++) {
@@ -778,30 +781,35 @@ TEST(BasicInterpreter, ThreeStepAllocate) {
             kTfLiteOk);
 
   // String-in String-out node.
   TfLiteRegistration reg_copy = {nullptr, nullptr, nullptr, nullptr};
   reg_copy.invoke = [](TfLiteContext* context, TfLiteNode* node) {
-    const TfLiteTensor* input = GetInput(context, node, 0);
-    TfLiteTensor* output = GetOutput(context, node, 0);
+    const TfLiteTensor* input;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
+    TfLiteTensor* output;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
     DynamicBuffer buf;
     StringRef str_ref = GetString(input, 0);
     buf.AddString(str_ref);
     buf.WriteToTensorAsVector(output);
     return kTfLiteOk;
   };
 
   // String-in Int-out node.
   TfLiteRegistration reg_len = {nullptr, nullptr, nullptr, nullptr};
   reg_len.prepare = [](TfLiteContext* context, TfLiteNode* node) {
-    TfLiteTensor* output = GetOutput(context, node, 0);
+    TfLiteTensor* output;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
     TfLiteIntArray* outputSize = TfLiteIntArrayCreate(1);
     outputSize->data[0] = 1;
     return context->ResizeTensor(context, output, outputSize);
   };
   reg_len.invoke = [](TfLiteContext* context, TfLiteNode* node) {
-    const TfLiteTensor* a0 = GetInput(context, node, 0);
-    TfLiteTensor* a1 = GetOutput(context, node, 0);
+    const TfLiteTensor* a0;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
+    TfLiteTensor* a1;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &a1));
     a1->data.i32[0] = a0->bytes;
     return kTfLiteOk;
   };
 
   ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr,
@@ -846,18 +854,22 @@ TEST(BasicInterpreter, AllocateTwice) {
                                                      quantized),
             kTfLiteOk);
 
   TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
   reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
-    const TfLiteTensor* tensor0 = GetInput(context, node, 0);
-    TfLiteTensor* tensor1 = GetOutput(context, node, 0);
+    const TfLiteTensor* tensor0;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &tensor0));
+    TfLiteTensor* tensor1;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &tensor1));
     TfLiteIntArray* newSize = TfLiteIntArrayCopy(tensor0->dims);
     return context->ResizeTensor(context, tensor1, newSize);
   };
   reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
-    const TfLiteTensor* a0 = GetInput(context, node, 0);
-    TfLiteTensor* a1 = GetOutput(context, node, 0);
+    const TfLiteTensor* a0;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
+    TfLiteTensor* a1;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &a1));
     int num = a0->dims->data[0];
     for (int i = 0; i < num; i++) {
       a1->data.f[i] = a0->data.f[i];
     }
     return kTfLiteOk;
@@ -1203,22 +1215,26 @@ class TestExecutionPlan : public ::testing::Test {
   TfLiteRegistration CopyOpRegistration() {
     TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
 
     reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
       // Set output size to input size
-      const TfLiteTensor* tensor0 = GetInput(context, node, 0);
-      TfLiteTensor* tensor1 = GetOutput(context, node, 0);
+      const TfLiteTensor* tensor0;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &tensor0));
+      TfLiteTensor* tensor1;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &tensor1));
       TfLiteIntArray* newSize = TfLiteIntArrayCopy(tensor0->dims);
       return context->ResizeTensor(context, tensor1, newSize);
     };
 
     reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
       CallReporting* call_reporting =
           static_cast<CallReporting*>(node->builtin_data);
       // Copy input data to output data.
-      const TfLiteTensor* a0 = GetInput(context, node, 0);
-      TfLiteTensor* a1 = GetOutput(context, node, 0);
+      const TfLiteTensor* a0;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
+      TfLiteTensor* a1;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &a1));
       int num = a0->dims->data[0];
       for (int i = 0; i < num; i++) {
         a1->data.f[i] = a0->data.f[i];
       }
       call_reporting->Record();
@@ -1401,12 +1417,14 @@ class CancellationTest : public ::testing::Test {
     TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
 
     // Set output size to the input size in CancelOp::Prepare(). Code exists to
     // have a framework in Prepare. The input and output tensors are not used.
     reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
-      const TfLiteTensor* in_tensor = GetInput(context, node, 0);
-      TfLiteTensor* out_tensor = GetOutput(context, node, 0);
+      const TfLiteTensor* in_tensor;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &in_tensor));
+      TfLiteTensor* out_tensor;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out_tensor));
       TfLiteIntArray* new_size = TfLiteIntArrayCopy(in_tensor->dims);
       return context->ResizeTensor(context, out_tensor, new_size);
     };
 
     reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
@@ -1421,12 +1439,14 @@ class CancellationTest : public ::testing::Test {
     TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
 
     // Set output size to the input size in OkOp::Prepare(). Code exists to have
     // a framework in Prepare. The input and output tensors are not used.
     reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
-      const TfLiteTensor* in_tensor = GetInput(context, node, 0);
-      TfLiteTensor* out_tensor = GetOutput(context, node, 0);
+      const TfLiteTensor* in_tensor;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &in_tensor));
+      TfLiteTensor* out_tensor;
+      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out_tensor));
       TfLiteIntArray* new_size = TfLiteIntArrayCopy(in_tensor->dims);
       return context->ResizeTensor(context, out_tensor, new_size);
     };
 
     reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
diff --git a/tensorflow/lite/java/src/test/native/interpreter_test_jni.cc b/tensorflow/lite/java/src/test/native/interpreter_test_jni.cc
index 76bb7b7d4cd..7981a8bb8f5 100644
--- a/tensorflow/lite/java/src/test/native/interpreter_test_jni.cc
+++ b/tensorflow/lite/java/src/test/native/interpreter_test_jni.cc
@@ -31,19 +31,25 @@ Java_org_tensorflow_lite_InterpreterTest_getNativeHandleForDelegate(
   static TfLiteRegistration registration = {
       .init = nullptr,
       .free = nullptr,
       .prepare =
           [](TfLiteContext* context, TfLiteNode* node) {
-            const TfLiteTensor* input = tflite::GetInput(context, node, 0);
-            TfLiteTensor* output = tflite::GetOutput(context, node, 0);
+            const TfLiteTensor* input;
+            TF_LITE_ENSURE_OK(context,
+                              tflite::GetInputSafe(context, node, 0, &input));
+            TfLiteTensor* output;
+            TF_LITE_ENSURE_OK(context,
+                              tflite::GetOutputSafe(context, node, 0, &output));
             TfLiteIntArray* output_dims = TfLiteIntArrayCopy(input->dims);
             output->type = kTfLiteFloat32;
             return context->ResizeTensor(context, output, output_dims);
           },
       .invoke =
           [](TfLiteContext* context, TfLiteNode* node) {
-            TfLiteTensor* output = tflite::GetOutput(context, node, 0);
+            TfLiteTensor* output;
+            TF_LITE_ENSURE_OK(context,
+                              tflite::GetOutputSafe(context, node, 0, &output));
             std::fill(output->data.f,
                       output->data.f + tflite::NumElements(output), 7.0f);
             return kTfLiteOk;
           },
       .profiling_string = nullptr,
diff --git a/tensorflow/lite/kernels/hashtable/hashtable.cc b/tensorflow/lite/kernels/hashtable/hashtable.cc
index d1f9551ddf0..bca8f2c9208 100644
--- a/tensorflow/lite/kernels/hashtable/hashtable.cc
+++ b/tensorflow/lite/kernels/hashtable/hashtable.cc
@@ -78,13 +78,13 @@ TfLiteStatus PrepareHashtable(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE(context, (params->key_dtype == kTfLiteInt64 &&
                            params->value_dtype == kTfLiteString) ||
                               (params->key_dtype == kTfLiteString &&
                                params->value_dtype == kTfLiteInt64));
 
-  TfLiteTensor* resource_handle_tensor =
-      GetOutput(context, node, kResourceHandleTensor);
-  TF_LITE_ENSURE(context, resource_handle_tensor != nullptr);
+  TfLiteTensor* resource_handle_tensor;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kResourceHandleTensor,
+                                           &resource_handle_tensor));
   TF_LITE_ENSURE_EQ(context, resource_handle_tensor->type, kTfLiteInt32);
   TfLiteIntArray* outputSize = TfLiteIntArrayCreate(1);
   outputSize->data[0] = 1;
   return context->ResizeTensor(context, resource_handle_tensor, outputSize);
 }
@@ -95,12 +95,13 @@ TfLiteStatus EvalHashtable(TfLiteContext* context, TfLiteNode* node) {
       reinterpret_cast<const TfLiteHashtableParams*>(node->user_data);
 
   // The resource id is generated based on the given table name.
   const int resource_id = std::hash<std::string>{}(params->table_name);
 
-  TfLiteTensor* resource_handle_tensor =
-      GetOutput(context, node, kResourceHandleTensor);
+  TfLiteTensor* resource_handle_tensor;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kResourceHandleTensor,
+                                           &resource_handle_tensor));
   auto* resource_handle_data =
       GetTensorData<std::int32_t>(resource_handle_tensor);
   resource_handle_data[0] = resource_id;
 
   Subgraph* subgraph = reinterpret_cast<Subgraph*>(context->impl_);
diff --git a/tensorflow/lite/kernels/hashtable/hashtable_find.cc b/tensorflow/lite/kernels/hashtable/hashtable_find.cc
index 10236cfce07..f26fe824b4c 100644
--- a/tensorflow/lite/kernels/hashtable/hashtable_find.cc
+++ b/tensorflow/lite/kernels/hashtable/hashtable_find.cc
@@ -32,39 +32,50 @@ constexpr int kOutputTensor = 0;
 
 TfLiteStatus PrepareHashtableFind(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
-  const TfLiteTensor* input_resource_id_tensor =
-      GetInput(context, node, kInputResourceIdTensor);
+  const TfLiteTensor* input_resource_id_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
+                                          &input_resource_id_tensor));
   TF_LITE_ENSURE_EQ(context, input_resource_id_tensor->type, kTfLiteInt32);
   TF_LITE_ENSURE_EQ(context, NumDimensions(input_resource_id_tensor), 1);
   TF_LITE_ENSURE_EQ(context, SizeOfDimension(input_resource_id_tensor, 0), 1);
 
-  const TfLiteTensor* default_value_tensor =
-      GetInput(context, node, kDefaultValueTensor);
+  const TfLiteTensor* default_value_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kDefaultValueTensor,
+                                          &default_value_tensor));
 
-  const TfLiteTensor* key_tensor = GetInput(context, node, kKeyTensor);
-  TfLiteTensor* output_tensor = GetOutput(context, node, kOutputTensor);
+  const TfLiteTensor* key_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kKeyTensor, &key_tensor));
+  TfLiteTensor* output_tensor;
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node, kOutputTensor, &output_tensor));
   TF_LITE_ENSURE_EQ(context, default_value_tensor->type, output_tensor->type);
   TF_LITE_ENSURE(context, (key_tensor->type == kTfLiteInt64 &&
                            output_tensor->type == kTfLiteString) ||
                               (key_tensor->type == kTfLiteString &&
                                output_tensor->type == kTfLiteInt64));
   return context->ResizeTensor(context, output_tensor,
                                TfLiteIntArrayCopy(key_tensor->dims));
 }
 
 TfLiteStatus EvalHashtableFind(TfLiteContext* context, TfLiteNode* node) {
-  const TfLiteTensor* input_resource_id_tensor =
-      GetInput(context, node, kInputResourceIdTensor);
+  const TfLiteTensor* input_resource_id_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
+                                          &input_resource_id_tensor));
   int resource_id = input_resource_id_tensor->data.i32[0];
 
-  const TfLiteTensor* key_tensor = GetInput(context, node, kKeyTensor);
-  const TfLiteTensor* default_value_tensor =
-      GetInput(context, node, kDefaultValueTensor);
-  TfLiteTensor* output_tensor = GetOutput(context, node, 0);
+  const TfLiteTensor* key_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kKeyTensor, &key_tensor));
+  const TfLiteTensor* default_value_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kDefaultValueTensor,
+                                          &default_value_tensor));
+  TfLiteTensor* output_tensor;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output_tensor));
 
   Subgraph* subgraph = reinterpret_cast<Subgraph*>(context->impl_);
   auto& resources = subgraph->resources();
   auto* lookup = resource::GetHashtableResource(&resources, resource_id);
   TF_LITE_ENSURE(context, lookup != nullptr);
diff --git a/tensorflow/lite/kernels/hashtable/hashtable_import.cc b/tensorflow/lite/kernels/hashtable/hashtable_import.cc
index 1b5c0424526..fad9345a1a4 100644
--- a/tensorflow/lite/kernels/hashtable/hashtable_import.cc
+++ b/tensorflow/lite/kernels/hashtable/hashtable_import.cc
@@ -31,18 +31,23 @@ constexpr int kValueTensor = 2;
 
 TfLiteStatus PrepareHashtableImport(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 0);
 
-  const TfLiteTensor* input_resource_id_tensor =
-      GetInput(context, node, kInputResourceIdTensor);
+  const TfLiteTensor* input_resource_id_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
+                                          &input_resource_id_tensor));
   TF_LITE_ENSURE_EQ(context, input_resource_id_tensor->type, kTfLiteInt32);
   TF_LITE_ENSURE_EQ(context, NumDimensions(input_resource_id_tensor), 1);
   TF_LITE_ENSURE_EQ(context, SizeOfDimension(input_resource_id_tensor, 0), 1);
 
-  const TfLiteTensor* key_tensor = GetInput(context, node, kKeyTensor);
-  const TfLiteTensor* value_tensor = GetInput(context, node, kValueTensor);
+  const TfLiteTensor* key_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kKeyTensor, &key_tensor));
+  const TfLiteTensor* value_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kValueTensor, &value_tensor));
   TF_LITE_ENSURE(context, (key_tensor->type == kTfLiteInt64 &&
                            value_tensor->type == kTfLiteString) ||
                               (key_tensor->type == kTfLiteString &&
                                value_tensor->type == kTfLiteInt64));
   // TODO(b/144731295): Tensorflow lookup ops support 1-D vector in storing
@@ -50,16 +55,21 @@ TfLiteStatus PrepareHashtableImport(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE(context, HaveSameShapes(key_tensor, value_tensor));
   return kTfLiteOk;
 }
 
 TfLiteStatus EvalHashtableImport(TfLiteContext* context, TfLiteNode* node) {
-  const TfLiteTensor* input_resource_id_tensor =
-      GetInput(context, node, kInputResourceIdTensor);
+  const TfLiteTensor* input_resource_id_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
+                                          &input_resource_id_tensor));
   const int resource_id = input_resource_id_tensor->data.i32[0];
 
-  const TfLiteTensor* key_tensor = GetInput(context, node, kKeyTensor);
-  const TfLiteTensor* value_tensor = GetInput(context, node, kValueTensor);
+  const TfLiteTensor* key_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kKeyTensor, &key_tensor));
+  const TfLiteTensor* value_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kValueTensor, &value_tensor));
 
   Subgraph* subgraph = reinterpret_cast<Subgraph*>(context->impl_);
   auto& resources = subgraph->resources();
   auto* lookup = resource::GetHashtableResource(&resources, resource_id);
   TF_LITE_ENSURE(context, lookup != nullptr);
diff --git a/tensorflow/lite/kernels/hashtable/hashtable_size.cc b/tensorflow/lite/kernels/hashtable/hashtable_size.cc
index 48029795ae0..34a8031a3c0 100644
--- a/tensorflow/lite/kernels/hashtable/hashtable_size.cc
+++ b/tensorflow/lite/kernels/hashtable/hashtable_size.cc
@@ -30,30 +30,35 @@ constexpr int kOutputTensor = 0;
 
 TfLiteStatus PrepareHashtableSize(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
-  const TfLiteTensor* input_resource_id_tensor =
-      GetInput(context, node, kInputResourceIdTensor);
+  const TfLiteTensor* input_resource_id_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
+                                          &input_resource_id_tensor));
   TF_LITE_ENSURE_EQ(context, input_resource_id_tensor->type, kTfLiteInt32);
   TF_LITE_ENSURE_EQ(context, NumDimensions(input_resource_id_tensor), 1);
   TF_LITE_ENSURE_EQ(context, SizeOfDimension(input_resource_id_tensor, 0), 1);
 
-  TfLiteTensor* output_tensor = GetOutput(context, node, kOutputTensor);
-  TF_LITE_ENSURE(context, output_tensor != nullptr);
+  TfLiteTensor* output_tensor;
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node, kOutputTensor, &output_tensor));
   TF_LITE_ENSURE_EQ(context, output_tensor->type, kTfLiteInt64);
   TfLiteIntArray* outputSize = TfLiteIntArrayCreate(1);
   outputSize->data[0] = 1;
   return context->ResizeTensor(context, output_tensor, outputSize);
 }
 
 TfLiteStatus EvalHashtableSize(TfLiteContext* context, TfLiteNode* node) {
-  const TfLiteTensor* input_resource_id_tensor =
-      GetInput(context, node, kInputResourceIdTensor);
+  const TfLiteTensor* input_resource_id_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
+                                          &input_resource_id_tensor));
   int resource_id = input_resource_id_tensor->data.i32[0];
 
-  TfLiteTensor* output_tensor = GetOutput(context, node, kOutputTensor);
+  TfLiteTensor* output_tensor;
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node, kOutputTensor, &output_tensor));
   auto* output_data = GetTensorData<std::int64_t>(output_tensor);
 
   Subgraph* subgraph = reinterpret_cast<Subgraph*>(context->impl_);
   auto& resources = subgraph->resources();
   auto* lookup = resource::GetHashtableResource(&resources, resource_id);
diff --git a/tensorflow/lite/kernels/kernel_util.cc b/tensorflow/lite/kernels/kernel_util.cc
index cd243335c9c..a834d8ab913 100644
--- a/tensorflow/lite/kernels/kernel_util.cc
+++ b/tensorflow/lite/kernels/kernel_util.cc
@@ -28,58 +28,163 @@ limitations under the License.
 
 namespace tflite {
 
 namespace {
 
-inline TfLiteTensor* GetMutableInput(const TfLiteContext* context,
-                                     const TfLiteNode* node, int index) {
-  if (index >= 0 && index < node->inputs->size) {
-    const int tensor_index = node->inputs->data[index];
+// Assumes tensor_index is a valid index (in bounds)
+inline TfLiteTensor* GetTensorAtIndex(const TfLiteContext* context,
+                                      int tensor_index) {
+  if (context->tensors != nullptr) {
+    return &context->tensors[tensor_index];
+  } else {
+    return context->GetTensor(context, tensor_index);
+  }
+}
+
+// Validate in a single place to reduce binary size
+inline TfLiteStatus ValidateTensorIndexingSafe(const TfLiteContext* context,
+                                               int index, int max_size,
+                                               const int* tensor_indices,
+                                               int* tensor_index) {
+  if (index < 0 || index >= max_size) {
+    TF_LITE_KERNEL_LOG(const_cast<TfLiteContext*>(context),
+                       "Invalid tensor index %d (not in [0, %d))\n", index,
+                       max_size);
+    return kTfLiteError;
+  }
+  if (tensor_indices[index] == kTfLiteOptionalTensor) {
+    TF_LITE_KERNEL_LOG(const_cast<TfLiteContext*>(context),
+                       "Tensor at index %d was optional but was expected\n",
+                       index);
+    return kTfLiteError;
+  }
+
+  *tensor_index = tensor_indices[index];
+  return kTfLiteOk;
+}
+
+// Same as above but returns -1 for invalid inputs instead of status + logging
+// error.
+inline int ValidateTensorIndexing(const TfLiteContext* context, int index,
+                                  int max_size, const int* tensor_indices) {
+  if (index >= 0 && index < max_size) {
+    const int tensor_index = tensor_indices[index];
     if (tensor_index != kTfLiteOptionalTensor) {
-      if (context->tensors != nullptr) {
-        return &context->tensors[tensor_index];
-      } else {
-        return context->GetTensor(context, tensor_index);
-      }
+      return tensor_index;
     }
   }
-  return nullptr;
+  return -1;
+}
+
+inline TfLiteTensor* GetMutableInput(const TfLiteContext* context,
+                                     const TfLiteNode* node, int index) {
+  const int tensor_index = ValidateTensorIndexing(
+      context, index, node->inputs->size, node->inputs->data);
+  if (tensor_index < 0) {
+    return nullptr;
+  }
+  return GetTensorAtIndex(context, tensor_index);
+}
+
+inline TfLiteStatus GetMutableInputSafe(const TfLiteContext* context,
+                                        const TfLiteNode* node, int index,
+                                        const TfLiteTensor** tensor) {
+  int tensor_index;
+  TF_LITE_ENSURE_OK(
+      context, ValidateTensorIndexingSafe(context, index, node->inputs->size,
+                                          node->inputs->data, &tensor_index));
+  *tensor = GetTensorAtIndex(context, tensor_index);
+  return kTfLiteOk;
 }
 
 }  // anonymous namespace.
 
 const TfLiteTensor* GetInput(const TfLiteContext* context,
                              const TfLiteNode* node, int index) {
   return GetMutableInput(context, node, index);
 }
 
+TfLiteStatus GetInputSafe(const TfLiteContext* context, const TfLiteNode* node,
+                          int index, const TfLiteTensor** tensor) {
+  return GetMutableInputSafe(context, node, index, tensor);
+}
+
 TfLiteTensor* GetVariableInput(TfLiteContext* context, const TfLiteNode* node,
                                int index) {
   TfLiteTensor* tensor = GetMutableInput(context, node, index);
   return tensor->is_variable ? tensor : nullptr;
 }
 
 TfLiteTensor* GetOutput(TfLiteContext* context, const TfLiteNode* node,
                         int index) {
-  if (index >= 0 && index < node->outputs->size) {
-    const int tensor_index = node->outputs->data[index];
-    if (tensor_index != kTfLiteOptionalTensor) {
-      if (context->tensors != nullptr) {
-        return &context->tensors[tensor_index];
-      } else {
-        return context->GetTensor(context, tensor_index);
-      }
-    }
+  const int tensor_index = ValidateTensorIndexing(
+      context, index, node->outputs->size, node->outputs->data);
+  if (tensor_index < 0) {
+    return nullptr;
   }
-  return nullptr;
+  return GetTensorAtIndex(context, tensor_index);
+}
+
+TfLiteStatus GetOutputSafe(const TfLiteContext* context, const TfLiteNode* node,
+                           int index, TfLiteTensor** tensor) {
+  int tensor_index;
+  TF_LITE_ENSURE_OK(
+      context, ValidateTensorIndexingSafe(context, index, node->outputs->size,
+                                          node->outputs->data, &tensor_index));
+  *tensor = GetTensorAtIndex(context, tensor_index);
+  return kTfLiteOk;
 }
 
 const TfLiteTensor* GetOptionalInputTensor(const TfLiteContext* context,
                                            const TfLiteNode* node, int index) {
   return GetInput(context, node, index);
 }
 
+#ifndef TF_LITE_STATIC_MEMORY
+TfLiteTensor* GetTemporary(TfLiteContext* context, const TfLiteNode* node,
+                           int index) {
+  const int tensor_index = ValidateTensorIndexing(
+      context, index, node->temporaries->size, node->temporaries->data);
+  if (tensor_index < 0) {
+    return nullptr;
+  }
+  return GetTensorAtIndex(context, tensor_index);
+}
+
+TfLiteStatus GetTemporarySafe(const TfLiteContext* context,
+                              const TfLiteNode* node, int index,
+                              TfLiteTensor** tensor) {
+  int tensor_index;
+  TF_LITE_ENSURE_OK(context, ValidateTensorIndexingSafe(
+                                 context, index, node->temporaries->size,
+                                 node->temporaries->data, &tensor_index));
+  *tensor = GetTensorAtIndex(context, tensor_index);
+  return kTfLiteOk;
+}
+
+const TfLiteTensor* GetIntermediates(TfLiteContext* context,
+                                     const TfLiteNode* node, int index) {
+  const int tensor_index = ValidateTensorIndexing(
+      context, index, node->intermediates->size, node->intermediates->data);
+  if (tensor_index < 0) {
+    return nullptr;
+  }
+  return GetTensorAtIndex(context, tensor_index);
+}
+
+TfLiteStatus GetIntermediatesSafe(const TfLiteContext* context,
+                                  const TfLiteNode* node, int index,
+                                  TfLiteTensor** tensor) {
+  int tensor_index;
+  TF_LITE_ENSURE_OK(context, ValidateTensorIndexingSafe(
+                                 context, index, node->intermediates->size,
+                                 node->intermediates->data, &tensor_index));
+  *tensor = GetTensorAtIndex(context, tensor_index);
+  return kTfLiteOk;
+}
+#endif  // TF_LITE_STATIC_MEMORY
+
 // Per-axis
 TfLiteStatus PopulateConvolutionQuantizationParams(
     TfLiteContext* context, const TfLiteTensor* input,
     const TfLiteTensor* filter, const TfLiteTensor* bias, TfLiteTensor* output,
     const TfLiteFusedActivation& activation, int32_t* multiplier, int* shift,
diff --git a/tensorflow/lite/kernels/kernel_util.h b/tensorflow/lite/kernels/kernel_util.h
index 5950effca0e..06f24b8e7d1 100644
--- a/tensorflow/lite/kernels/kernel_util.h
+++ b/tensorflow/lite/kernels/kernel_util.h
@@ -38,10 +38,21 @@ namespace tflite {
 // This is because the index might point to the optional tensor constant
 // (kTfLiteOptionalTensor) in which case there is no tensor to return.
 const TfLiteTensor* GetInput(const TfLiteContext* context,
                              const TfLiteNode* node, int index);
 
+// Same as `GetInput` but returns boolean and uses output argument for tensor.
+//
+//   TfLiteTensor* my_tensor;
+//   TF_LITE_ENSURE_OK(context,
+//                     GetInputSafe(context, node, kMyTensorIdx, &my_tensor));
+//   // can use my_tensor directly from here onwards, it is not nullptr
+//
+// Should be used in cases where the binary size is too large.
+TfLiteStatus GetInputSafe(const TfLiteContext* context, const TfLiteNode* node,
+                          int index, const TfLiteTensor** tensor);
+
 // Note: You must check if result is not null:
 //
 //   TfLiteTensor* my_tensor = GetVariableInput(context, node, kMyTensorIdx);
 //   TF_LITE_ENSURE(context, my_tensor != nullptr);
 //
@@ -58,10 +69,21 @@ TfLiteTensor* GetVariableInput(TfLiteContext* context, const TfLiteNode* node,
 // This is because the index might point to the optional tensor constant
 // (kTfLiteOptionalTensor) in which case there is no tensor to return.
 TfLiteTensor* GetOutput(TfLiteContext* context, const TfLiteNode* node,
                         int index);
 
+// Same as `GetOutput` but returns boolean and uses output argument for tensor.
+//
+//   TfLiteTensor* my_tensor;
+//   TF_LITE_ENSURE_OK(context,
+//                     GetOutputSafe(context, node, kMyTensorIdx, &my_tensor));
+//   // can use my_tensor directly from here onwards, it is not nullptr
+//
+// Should be used in cases where the binary size is too large.
+TfLiteStatus GetOutputSafe(const TfLiteContext* context, const TfLiteNode* node,
+                           int index, TfLiteTensor** tensor);
+
 // Note: You must check if result is not null:
 //
 //   TfLiteTensor* my_tensor = GetOptionalInputTensor(context, node, kIdx);
 //   TF_LITE_ENSURE(context, my_tensor != nullptr);
 //
@@ -70,62 +92,73 @@ TfLiteTensor* GetOutput(TfLiteContext* context, const TfLiteNode* node,
 //
 // Deprecated. GetInput has the same functionality.
 const TfLiteTensor* GetOptionalInputTensor(const TfLiteContext* context,
                                            const TfLiteNode* node, int index);
 
-inline int NumDimensions(const TfLiteTensor* t) { return t->dims->size; }
-inline int SizeOfDimension(const TfLiteTensor* t, int dim) {
-  return t->dims->data[dim];
-}
-
 #ifndef TF_LITE_STATIC_MEMORY
 // Note: You must check if result is not null:
 //
 //   TfLiteTensor* my_tensor = GetTemporary(context, node, kMyTensorIdx);
 //   TF_LITE_ENSURE(context, my_tensor != nullptr);
 //
 // This is because the index might point to the optional tensor constant
 // (kTfLiteOptionalTensor) in which case there is no tensor to return.
-inline TfLiteTensor* GetTemporary(TfLiteContext* context,
-                                  const TfLiteNode* node, int index) {
-  if (index >= 0 && index < node->temporaries->size) {
-    const int tensor_index = node->temporaries->data[index];
-    if (tensor_index != kTfLiteOptionalTensor) {
-      if (context->tensors != nullptr) {
-        return &context->tensors[tensor_index];
-      }
-    }
-  }
-  return nullptr;
-}
+TfLiteTensor* GetTemporary(TfLiteContext* context, const TfLiteNode* node,
+                           int index);
+
+// Same as `GetTemporary` but returns boolean and uses output argument for
+// tensor.
+//
+//   TfLiteTensor* my_tensor;
+//   TF_LITE_ENSURE_OK(context,
+//                     GetTemporarySafe(context, node, kMyTensorIdx,
+//                     &my_tensor));
+//   // can use my_tensor directly from here onwards, it is not nullptr
+//
+// Should be used in cases where the binary size is too large.
+TfLiteStatus GetTemporarySafe(const TfLiteContext* context,
+                              const TfLiteNode* node, int index,
+                              TfLiteTensor** tensor);
 
 // Note: You must check if result is not null:
 //
 //   TfLiteTensor* my_tensor = GetIntermediates(context, node, kMyTensorIdx);
 //   TF_LITE_ENSURE(context, my_tensor != nullptr);
 //
 // This is because the index might point to the optional tensor constant
 // (kTfLiteOptionalTensor) in which case there is no tensor to return.
-inline const TfLiteTensor* GetIntermediates(TfLiteContext* context,
-                                            const TfLiteNode* node, int index) {
-  if (index >= 0 && index < node->intermediates->size) {
-    const int tensor_index = node->intermediates->data[index];
-    if (tensor_index != kTfLiteOptionalTensor) {
-      if (context->tensors != nullptr) {
-        return &context->tensors[tensor_index];
-      }
-    }
-  }
-  return nullptr;
+const TfLiteTensor* GetIntermediates(TfLiteContext* context,
+                                     const TfLiteNode* node, int index);
+
+// Same as `GetIntermediates` but returns boolean and uses output argument for
+// tensor.
+//
+//   TfLiteTensor* my_tensor;
+//   TF_LITE_ENSURE_OK(context,
+//                     GetIntermediatesSafe(context, node, kMyTensorIdx,
+//                     &my_tensor));
+//   // can use my_tensor directly from here onwards, it is not nullptr
+//
+// Should be used in cases where the binary size is too large.
+TfLiteStatus GetIntermediatesSafe(const TfLiteContext* context,
+                                  const TfLiteNode* node, int index,
+                                  TfLiteTensor** tensor);
+#endif  // TF_LITE_STATIC_MEMORY
+
+inline int NumDimensions(const TfLiteTensor* t) { return t->dims->size; }
+inline int SizeOfDimension(const TfLiteTensor* t, int dim) {
+  return t->dims->data[dim];
 }
 
+inline int NumInputs(const TfLiteNode* node) { return node->inputs->size; }
+inline int NumOutputs(const TfLiteNode* node) { return node->outputs->size; }
+
+#ifndef TF_LITE_STATIC_MEMORY
 inline int NumIntermediates(const TfLiteNode* node) {
   return node->intermediates->size;
 }
 #endif  // TF_LITE_STATIC_MEMORY
-inline int NumInputs(const TfLiteNode* node) { return node->inputs->size; }
-inline int NumOutputs(const TfLiteNode* node) { return node->outputs->size; }
 
 inline int64_t NumElements(const TfLiteIntArray* dims) {
   int64_t count = 1;
   for (int i = 0; i < dims->size; ++i) {
     count *= dims->data[i];
diff --git a/tensorflow/lite/profiling/profile_summarizer_test.cc b/tensorflow/lite/profiling/profile_summarizer_test.cc
index fd81c00e603..7eccc2d3843 100644
--- a/tensorflow/lite/profiling/profile_summarizer_test.cc
+++ b/tensorflow/lite/profiling/profile_summarizer_test.cc
@@ -34,14 +34,18 @@ namespace profiling {
 namespace {
 
 const char* kOpName = "SimpleOpEval";
 
 TfLiteStatus SimpleOpEval(TfLiteContext* context, TfLiteNode* node) {
-  const TfLiteTensor* input1 = tflite::GetInput(context, node, /*index=*/0);
-  const TfLiteTensor* input2 = tflite::GetInput(context, node, /*index=*/1);
-
-  TfLiteTensor* output = GetOutput(context, node, /*index=*/0);
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, /*index=*/0, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, /*index=*/1, &input2));
+
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, /*index=*/0, &output));
 
   int32_t* output_data = output->data.i32;
   *output_data = *(input1->data.i32) + *(input2->data.i32);
   return kTfLiteOk;
 }
diff --git a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
index 1ac996abe87..6249649f4b7 100644
--- a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
+++ b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
@@ -464,30 +464,55 @@ struct OpData {
 // tensors match each other.
 TfLiteStatus lstm_eval(TfLiteContext* context, TfLiteNode* node, Logger* logger,
                        ErrorReporter* error_reporter) {
   const auto* params = static_cast<TfLiteLSTMParams*>(node->builtin_data);
 
-  const TfLiteTensor* input =
-      GetInput(context, node, ops::builtin::lstm::full::kInputTensor);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node,
+                            ops::builtin::lstm::full::kInputTensor, &input));
 
   const TfLiteTensor* input_to_input_weights = GetOptionalInputTensor(
       context, node, ops::builtin::lstm::full::kInputToInputWeightsTensor);
-  const TfLiteTensor* input_to_forget_weights = GetInput(
-      context, node, ops::builtin::lstm::full::kInputToForgetWeightsTensor);
-  const TfLiteTensor* input_to_cell_weights = GetInput(
-      context, node, ops::builtin::lstm::full::kInputToCellWeightsTensor);
-  const TfLiteTensor* input_to_output_weights = GetInput(
-      context, node, ops::builtin::lstm::full::kInputToOutputWeightsTensor);
+  const TfLiteTensor* input_to_forget_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kInputToForgetWeightsTensor,
+                   &input_to_forget_weights));
+  const TfLiteTensor* input_to_cell_weights;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node,
+                            ops::builtin::lstm::full::kInputToCellWeightsTensor,
+                            &input_to_cell_weights));
+  const TfLiteTensor* input_to_output_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kInputToOutputWeightsTensor,
+                   &input_to_output_weights));
 
   const TfLiteTensor* recurrent_to_input_weights = GetOptionalInputTensor(
       context, node, ops::builtin::lstm::full::kRecurrentToInputWeightsTensor);
-  const TfLiteTensor* recurrent_to_forget_weights = GetInput(
-      context, node, ops::builtin::lstm::full::kRecurrentToForgetWeightsTensor);
-  const TfLiteTensor* recurrent_to_cell_weights = GetInput(
-      context, node, ops::builtin::lstm::full::kRecurrentToCellWeightsTensor);
-  const TfLiteTensor* recurrent_to_output_weights = GetInput(
-      context, node, ops::builtin::lstm::full::kRecurrentToOutputWeightsTensor);
+  const TfLiteTensor* recurrent_to_forget_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kRecurrentToForgetWeightsTensor,
+                   &recurrent_to_forget_weights));
+  const TfLiteTensor* recurrent_to_cell_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kRecurrentToCellWeightsTensor,
+                   &recurrent_to_cell_weights));
+  const TfLiteTensor* recurrent_to_output_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kRecurrentToOutputWeightsTensor,
+                   &recurrent_to_output_weights));
 
   const TfLiteTensor* cell_to_input_weights = GetOptionalInputTensor(
       context, node, ops::builtin::lstm::full::kCellToInputWeightsTensor);
   const TfLiteTensor* cell_to_forget_weights = GetOptionalInputTensor(
       context, node, ops::builtin::lstm::full::kCellToForgetWeightsTensor);
@@ -507,34 +532,47 @@ TfLiteStatus lstm_eval(TfLiteContext* context, TfLiteNode* node, Logger* logger,
       context, node,
       ops::builtin::lstm::full::kOutputLayerNormCoefficientsTensor);
 
   const TfLiteTensor* input_gate_bias = GetOptionalInputTensor(
       context, node, ops::builtin::lstm::full::kInputGateBiasTensor);
-  const TfLiteTensor* forget_gate_bias =
-      GetInput(context, node, ops::builtin::lstm::full::kForgetGateBiasTensor);
-  const TfLiteTensor* cell_gate_bias =
-      GetInput(context, node, ops::builtin::lstm::full::kCellGateBiasTensor);
-  const TfLiteTensor* output_gate_bias =
-      GetInput(context, node, ops::builtin::lstm::full::kOutputGateBiasTensor);
+  const TfLiteTensor* forget_gate_bias;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node,
+                            ops::builtin::lstm::full::kForgetGateBiasTensor,
+                            &forget_gate_bias));
+  const TfLiteTensor* cell_gate_bias;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node, ops::builtin::lstm::full::kCellGateBiasTensor,
+                   &cell_gate_bias));
+  const TfLiteTensor* output_gate_bias;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node,
+                            ops::builtin::lstm::full::kOutputGateBiasTensor,
+                            &output_gate_bias));
 
   const TfLiteTensor* projection_weights = GetOptionalInputTensor(
       context, node, ops::builtin::lstm::full::kProjectionWeightsTensor);
   const TfLiteTensor* projection_bias = GetOptionalInputTensor(
       context, node, ops::builtin::lstm::full::kProjectionBiasTensor);
 
   // Index the scratch buffers pointers to the global scratch buffer.
-  TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
+  TfLiteTensor* scratch_buffer;
+  TF_LITE_ENSURE_OK(
+      context, GetTemporarySafe(context, node, /*index=*/0, &scratch_buffer));
 
   TfLiteTensor* output_state = GetVariableInput(
       context, node, ops::builtin::lstm::full::kOutputStateTensor);
   TF_LITE_ENSURE(context, output_state != nullptr);
   TfLiteTensor* cell_state = GetVariableInput(
       context, node, ops::builtin::lstm::full::kCellStateTensor);
   TF_LITE_ENSURE(context, cell_state != nullptr);
 
-  TfLiteTensor* output =
-      GetOutput(context, node, ops::builtin::lstm::full::kOutputTensor);
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node,
+                             ops::builtin::lstm::full::kOutputTensor, &output));
 
   std::vector<int> intermediate_tensor_indexes(node->intermediates->size);
   for (int i = 0; i < node->intermediates->size; ++i) {
     intermediate_tensor_indexes[i] = node->intermediates->data[i];
   }

Vulnerable code:     const TfLiteTensor* input1 = GetInput(context, node, 0);
    const TfLiteTensor* input2 = GetInput(context, node, 1);
    TfLiteTensor* output = GetOutput(context, node, 0);
    const TfLiteTensor* a0 = GetInput(context, node, 0);
    const TfLiteTensor* a1 = GetInput(context, node, 1);
    TfLiteTensor* out = GetOutput(context, node, 0);
          TfLiteTensor* out = GetOutput(context, node, 0);
          const TfLiteTensor* input1 = GetInput(context, node, 0);
          TfLiteTensor* output = GetOutput(context, node, 0);
          TfLiteTensor* output = GetOutput(context, node, 0);
      TfLiteTensor* output0 = GetOutput(context, node, 0);
      const TfLiteTensor* input = GetInput(context, node, 0);
      TfLiteTensor* output1 = GetOutput(context, node, 1);
      TfLiteTensor* output0 = GetOutput(context, node, 0);
      const TfLiteTensor* input = GetInput(context, node, 0);
      TfLiteTensor* output1 = GetOutput(context, node, 1);
    auto pad_tensor = tflite::GetInput(context, tflite_node, 1);
  const TfLiteTensor* weights = GetInput(context, node, kWeightTensor);
  const TfLiteTensor* input = GetInput(context, node, kInput);
  const TfLiteTensor* weights = GetInput(context, node, kWeights);
    const TfLiteTensor* bias = GetInput(context, node, kBias);
  const TfLiteTensor* padding = GetInput(context, node, 1);
  const auto* shape = GetInput(context, node, kShapeTensor);
  const TfLiteTensor* inputs = GetInput(context, node, kInputsTensor);
  const TfLiteTensor* sequence_length =
      GetInput(context, node, kSequenceLengthTensor);
    TfLiteTensor* indices = GetOutput(context, node, i);
    TfLiteTensor* values = GetOutput(context, node, i + top_paths);
    TfLiteTensor* output_shape = GetOutput(context, node, i + 2 * top_paths);
  TfLiteTensor* log_probability_output =
      GetOutput(context, node, top_paths * 3);
    TfLiteTensor* indices = GetOutput(context, node, p);
    TfLiteTensor* values = GetOutput(context, node, p + top_paths);
    TfLiteTensor* decoded_shape = GetOutput(context, node, p + 2 * top_paths);
  const TfLiteTensor* inputs = GetInput(context, node, kInputsTensor);
  const TfLiteTensor* sequence_length =
      GetInput(context, node, kSequenceLengthTensor);
  TfLiteTensor* log_probabilities = GetOutput(context, node, 3 * top_paths);
  const TfLiteTensor* input = GetInput(context, node, kInput);
  const TfLiteTensor* input_state = GetInput(context, node, kInputState);
  const TfLiteTensor* gate_weight = GetInput(context, node, kGateWeight);
  const TfLiteTensor* gate_bias = GetInput(context, node, kGateBias);
  const TfLiteTensor* candidate_weight =
      GetInput(context, node, kCandidateWeight);
  const TfLiteTensor* candidate_bias = GetInput(context, node, kCandidateBias);
  TfLiteTensor* output = GetOutput(context, node, kOutput);
  TfLiteTensor* output_state = GetOutput(context, node, kOutputState);
  TfLiteTensor* activation = GetTemporary(context, node, kActivation);
  TfLiteTensor* concat = GetTemporary(context, node, kConcat);
  const TfLiteTensor* input = GetInput(context, node, kInput);
  const TfLiteTensor* input_state = GetInput(context, node, kInputState);
  const TfLiteTensor* gate_weight = GetInput(context, node, kGateWeight);
  const TfLiteTensor* gate_bias = GetInput(context, node, kGateBias);
  const TfLiteTensor* candidate_weight =
      GetInput(context, node, kCandidateWeight);
  const TfLiteTensor* candidate_bias = GetInput(context, node, kCandidateBias);
  TfLiteTensor* output = GetOutput(context, node, kOutput);
  TfLiteTensor* output_state = GetOutput(context, node, kOutputState);
  TfLiteTensor* activation = GetTemporary(context, node, kActivation);
  TfLiteTensor* concat = GetTemporary(context, node, kConcat);
  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
    const TfLiteTensor* tensor0 = GetInput(context, node, 0);
    TfLiteTensor* tensor1 = GetOutput(context, node, 0);
    const TfLiteTensor* a0 = GetInput(context, node, 0);
    const TfLiteTensor* input = GetInput(context, node, 0);
    TfLiteTensor* output = GetOutput(context, node, 0);
    TfLiteTensor* output = GetOutput(context, node, 0);
    const TfLiteTensor* a0 = GetInput(context, node, 0);
    TfLiteTensor* a1 = GetOutput(context, node, 0);
    const TfLiteTensor* tensor0 = GetInput(context, node, 0);
    TfLiteTensor* tensor1 = GetOutput(context, node, 0);
    const TfLiteTensor* a0 = GetInput(context, node, 0);
    TfLiteTensor* a1 = GetOutput(context, node, 0);
      const TfLiteTensor* tensor0 = GetInput(context, node, 0);
      TfLiteTensor* tensor1 = GetOutput(context, node, 0);
      const TfLiteTensor* a0 = GetInput(context, node, 0);
      TfLiteTensor* a1 = GetOutput(context, node, 0);
      const TfLiteTensor* in_tensor = GetInput(context, node, 0);
      TfLiteTensor* out_tensor = GetOutput(context, node, 0);
      const TfLiteTensor* in_tensor = GetInput(context, node, 0);
      TfLiteTensor* out_tensor = GetOutput(context, node, 0);
            const TfLiteTensor* input = tflite::GetInput(context, node, 0);
            TfLiteTensor* output = tflite::GetOutput(context, node, 0);
            TfLiteTensor* output = tflite::GetOutput(context, node, 0);
  TfLiteTensor* resource_handle_tensor =
      GetOutput(context, node, kResourceHandleTensor);
  TF_LITE_ENSURE(context, resource_handle_tensor != nullptr);
  TfLiteTensor* resource_handle_tensor =
      GetOutput(context, node, kResourceHandleTensor);
  const TfLiteTensor* input_resource_id_tensor =
      GetInput(context, node, kInputResourceIdTensor);
  const TfLiteTensor* default_value_tensor =
      GetInput(context, node, kDefaultValueTensor);
  const TfLiteTensor* key_tensor = GetInput(context, node, kKeyTensor);
  TfLiteTensor* output_tensor = GetOutput(context, node, kOutputTensor);
  const TfLiteTensor* input_resource_id_tensor =
      GetInput(context, node, kInputResourceIdTensor);
  const TfLiteTensor* key_tensor = GetInput(context, node, kKeyTensor);
  const TfLiteTensor* default_value_tensor =
      GetInput(context, node, kDefaultValueTensor);
  TfLiteTensor* output_tensor = GetOutput(context, node, 0);
  const TfLiteTensor* input_resource_id_tensor =
      GetInput(context, node, kInputResourceIdTensor);
  const TfLiteTensor* key_tensor = GetInput(context, node, kKeyTensor);
  const TfLiteTensor* value_tensor = GetInput(context, node, kValueTensor);
  const TfLiteTensor* input_resource_id_tensor =
      GetInput(context, node, kInputResourceIdTensor);
  const TfLiteTensor* key_tensor = GetInput(context, node, kKeyTensor);
  const TfLiteTensor* value_tensor = GetInput(context, node, kValueTensor);
  const TfLiteTensor* input_resource_id_tensor =
      GetInput(context, node, kInputResourceIdTensor);
  TfLiteTensor* output_tensor = GetOutput(context, node, kOutputTensor);
  TF_LITE_ENSURE(context, output_tensor != nullptr);
  const TfLiteTensor* input_resource_id_tensor =
      GetInput(context, node, kInputResourceIdTensor);
  TfLiteTensor* output_tensor = GetOutput(context, node, kOutputTensor);
inline TfLiteTensor* GetMutableInput(const TfLiteContext* context,
                                     const TfLiteNode* node, int index) {
  if (index >= 0 && index < node->inputs->size) {
    const int tensor_index = node->inputs->data[index];
      if (context->tensors != nullptr) {
        return &context->tensors[tensor_index];
      } else {
        return context->GetTensor(context, tensor_index);
      }
  return nullptr;
  if (index >= 0 && index < node->outputs->size) {
    const int tensor_index = node->outputs->data[index];
    if (tensor_index != kTfLiteOptionalTensor) {
      if (context->tensors != nullptr) {
        return &context->tensors[tensor_index];
      } else {
        return context->GetTensor(context, tensor_index);
      }
    }
  return nullptr;
inline int NumDimensions(const TfLiteTensor* t) { return t->dims->size; }
inline int SizeOfDimension(const TfLiteTensor* t, int dim) {
  return t->dims->data[dim];
}

inline TfLiteTensor* GetTemporary(TfLiteContext* context,
                                  const TfLiteNode* node, int index) {
  if (index >= 0 && index < node->temporaries->size) {
    const int tensor_index = node->temporaries->data[index];
    if (tensor_index != kTfLiteOptionalTensor) {
      if (context->tensors != nullptr) {
        return &context->tensors[tensor_index];
      }
    }
  }
  return nullptr;
}
inline const TfLiteTensor* GetIntermediates(TfLiteContext* context,
                                            const TfLiteNode* node, int index) {
  if (index >= 0 && index < node->intermediates->size) {
    const int tensor_index = node->intermediates->data[index];
    if (tensor_index != kTfLiteOptionalTensor) {
      if (context->tensors != nullptr) {
        return &context->tensors[tensor_index];
      }
    }
  }
  return nullptr;
inline int NumInputs(const TfLiteNode* node) { return node->inputs->size; }
inline int NumOutputs(const TfLiteNode* node) { return node->outputs->size; }
  const TfLiteTensor* input1 = tflite::GetInput(context, node, /*index=*/0);
  const TfLiteTensor* input2 = tflite::GetInput(context, node, /*index=*/1);

  TfLiteTensor* output = GetOutput(context, node, /*index=*/0);
  const TfLiteTensor* input =
      GetInput(context, node, ops::builtin::lstm::full::kInputTensor);
  const TfLiteTensor* input_to_forget_weights = GetInput(
      context, node, ops::builtin::lstm::full::kInputToForgetWeightsTensor);
  const TfLiteTensor* input_to_cell_weights = GetInput(
      context, node, ops::builtin::lstm::full::kInputToCellWeightsTensor);
  const TfLiteTensor* input_to_output_weights = GetInput(
      context, node, ops::builtin::lstm::full::kInputToOutputWeightsTensor);
  const TfLiteTensor* recurrent_to_forget_weights = GetInput(
      context, node, ops::builtin::lstm::full::kRecurrentToForgetWeightsTensor);
  const TfLiteTensor* recurrent_to_cell_weights = GetInput(
      context, node, ops::builtin::lstm::full::kRecurrentToCellWeightsTensor);
  const TfLiteTensor* recurrent_to_output_weights = GetInput(
      context, node, ops::builtin::lstm::full::kRecurrentToOutputWeightsTensor);
  const TfLiteTensor* forget_gate_bias =
      GetInput(context, node, ops::builtin::lstm::full::kForgetGateBiasTensor);
  const TfLiteTensor* cell_gate_bias =
      GetInput(context, node, ops::builtin::lstm::full::kCellGateBiasTensor);
  const TfLiteTensor* output_gate_bias =
      GetInput(context, node, ops::builtin::lstm::full::kOutputGateBiasTensor);
  TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
  TfLiteTensor* output =
      GetOutput(context, node, ops::builtin::lstm::full::kOutputTensor);
Fixed code:     const TfLiteTensor* input1;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input1));
    const TfLiteTensor* input2;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &input2));
    TfLiteTensor* output;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
    const TfLiteTensor* a0;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
    const TfLiteTensor* a1;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &a1));
    TfLiteTensor* out;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out));
          TfLiteTensor* out;
          TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out));
          const TfLiteTensor* input1;
          TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input1));
          TfLiteTensor* output;
          TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
          TfLiteTensor* output;
          TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
      TfLiteTensor* output0;
      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output0));
      const TfLiteTensor* input;
      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
      TfLiteTensor* output1;
      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 1, &output1));
      TfLiteTensor* output0;
      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output0));
      const TfLiteTensor* input;
      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
      TfLiteTensor* output1;
      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 1, &output1));
    if (input0 == nullptr || input1 == nullptr) {
      return absl::InvalidArgumentError("At least one input tensor is null");
    }
    const TfLiteTensor* pad_tensor = tflite::GetInput(context, tflite_node, 1);
    if (pad_tensor == nullptr) {
      return absl::InvalidArgumentError("Padding tensor was null");
    }
  const TfLiteTensor* weights;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kWeightTensor, &weights));
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
  const TfLiteTensor* weights;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kWeights, &weights));
    const TfLiteTensor* bias;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBias, &bias));
  const TfLiteTensor* padding;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &padding));
  const TfLiteTensor* shape;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kShapeTensor, &shape));
  const TfLiteTensor* inputs;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputsTensor, &inputs));
  const TfLiteTensor* sequence_length;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kSequenceLengthTensor,
                                          &sequence_length));
    TfLiteTensor* indices;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, i, &indices));
    TfLiteTensor* values;
    TF_LITE_ENSURE_OK(context,
                      GetOutputSafe(context, node, i + top_paths, &values));
    TfLiteTensor* output_shape;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, i + 2 * top_paths,
                                             &output_shape));
  TfLiteTensor* log_probability_output;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, top_paths * 3,
                                           &log_probability_output));
    TfLiteTensor* indices;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, p, &indices));
    TfLiteTensor* values;
    TF_LITE_ENSURE_OK(context,
                      GetOutputSafe(context, node, p + top_paths, &values));
    TfLiteTensor* decoded_shape;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, p + 2 * top_paths,
                                             &decoded_shape));
  const TfLiteTensor* inputs;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputsTensor, &inputs));
  const TfLiteTensor* sequence_length;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kSequenceLengthTensor,
                                          &sequence_length));
  TfLiteTensor* log_probabilities;
  TF_LITE_ENSURE_OK(
      context, GetOutputSafe(context, node, 3 * top_paths, &log_probabilities));
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
  const TfLiteTensor* input_state;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputState, &input_state));
  const TfLiteTensor* gate_weight;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateWeight, &gate_weight));
  const TfLiteTensor* gate_bias;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateBias, &gate_bias));
  const TfLiteTensor* candidate_weight;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kCandidateWeight,
                                          &candidate_weight));
  const TfLiteTensor* candidate_bias;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node, kCandidateBias, &candidate_bias));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kOutput, &output));
  TfLiteTensor* output_state;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputState, &output_state));
  TfLiteTensor* activation;
  TF_LITE_ENSURE_OK(context,
                    GetTemporarySafe(context, node, kActivation, &activation));
  TfLiteTensor* concat;
  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kConcat, &concat));
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
  const TfLiteTensor* input_state;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputState, &input_state));
  const TfLiteTensor* gate_weight;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateWeight, &gate_weight));
  const TfLiteTensor* gate_bias;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateBias, &gate_bias));
  const TfLiteTensor* candidate_weight;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kCandidateWeight,
                                          &candidate_weight));
  const TfLiteTensor* candidate_bias;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node, kCandidateBias, &candidate_bias));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kOutput, &output));
  TfLiteTensor* output_state;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputState, &output_state));
  TfLiteTensor* activation;
  TF_LITE_ENSURE_OK(context,
                    GetTemporarySafe(context, node, kActivation, &activation));
  TfLiteTensor* concat;
  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kConcat, &concat));
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensor, &output));
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensor, &output));
    const TfLiteTensor* tensor0;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &tensor0));
    TfLiteTensor* tensor1;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &tensor1));
    const TfLiteTensor* a0;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
    const TfLiteTensor* input;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
    TfLiteTensor* output;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
    TfLiteTensor* output;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
    const TfLiteTensor* a0;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
    TfLiteTensor* a1;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &a1));
    const TfLiteTensor* tensor0;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &tensor0));
    TfLiteTensor* tensor1;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &tensor1));
    const TfLiteTensor* a0;
    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
    TfLiteTensor* a1;
    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &a1));
      const TfLiteTensor* tensor0;
      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &tensor0));
      TfLiteTensor* tensor1;
      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &tensor1));
      const TfLiteTensor* a0;
      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
      TfLiteTensor* a1;
      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &a1));
      const TfLiteTensor* in_tensor;
      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &in_tensor));
      TfLiteTensor* out_tensor;
      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out_tensor));
      const TfLiteTensor* in_tensor;
      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &in_tensor));
      TfLiteTensor* out_tensor;
      TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out_tensor));
            const TfLiteTensor* input;
            TF_LITE_ENSURE_OK(context,
                              tflite::GetInputSafe(context, node, 0, &input));
            TfLiteTensor* output;
            TF_LITE_ENSURE_OK(context,
                              tflite::GetOutputSafe(context, node, 0, &output));
            TfLiteTensor* output;
            TF_LITE_ENSURE_OK(context,
                              tflite::GetOutputSafe(context, node, 0, &output));
  TfLiteTensor* resource_handle_tensor;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kResourceHandleTensor,
                                           &resource_handle_tensor));
  TfLiteTensor* resource_handle_tensor;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kResourceHandleTensor,
                                           &resource_handle_tensor));
  const TfLiteTensor* input_resource_id_tensor;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
                                          &input_resource_id_tensor));
  const TfLiteTensor* default_value_tensor;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kDefaultValueTensor,
                                          &default_value_tensor));
  const TfLiteTensor* key_tensor;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kKeyTensor, &key_tensor));
  TfLiteTensor* output_tensor;
  TF_LITE_ENSURE_OK(
      context, GetOutputSafe(context, node, kOutputTensor, &output_tensor));
  const TfLiteTensor* input_resource_id_tensor;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
                                          &input_resource_id_tensor));
  const TfLiteTensor* key_tensor;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kKeyTensor, &key_tensor));
  const TfLiteTensor* default_value_tensor;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kDefaultValueTensor,
                                          &default_value_tensor));
  TfLiteTensor* output_tensor;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output_tensor));
  const TfLiteTensor* input_resource_id_tensor;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
                                          &input_resource_id_tensor));
  const TfLiteTensor* key_tensor;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kKeyTensor, &key_tensor));
  const TfLiteTensor* value_tensor;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kValueTensor, &value_tensor));
  const TfLiteTensor* input_resource_id_tensor;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
                                          &input_resource_id_tensor));
  const TfLiteTensor* key_tensor;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kKeyTensor, &key_tensor));
  const TfLiteTensor* value_tensor;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kValueTensor, &value_tensor));
  const TfLiteTensor* input_resource_id_tensor;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
                                          &input_resource_id_tensor));
  TfLiteTensor* output_tensor;
  TF_LITE_ENSURE_OK(
      context, GetOutputSafe(context, node, kOutputTensor, &output_tensor));
  const TfLiteTensor* input_resource_id_tensor;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputResourceIdTensor,
                                          &input_resource_id_tensor));
  TfLiteTensor* output_tensor;
  TF_LITE_ENSURE_OK(
      context, GetOutputSafe(context, node, kOutputTensor, &output_tensor));
// Assumes tensor_index is a valid index (in bounds)
inline TfLiteTensor* GetTensorAtIndex(const TfLiteContext* context,
                                      int tensor_index) {
  if (context->tensors != nullptr) {
    return &context->tensors[tensor_index];
  } else {
    return context->GetTensor(context, tensor_index);
  }
}

// Validate in a single place to reduce binary size
inline TfLiteStatus ValidateTensorIndexingSafe(const TfLiteContext* context,
                                               int index, int max_size,
                                               const int* tensor_indices,
                                               int* tensor_index) {
  if (index < 0 || index >= max_size) {
    TF_LITE_KERNEL_LOG(const_cast<TfLiteContext*>(context),
                       "Invalid tensor index %d (not in [0, %d))\n", index,
                       max_size);
    return kTfLiteError;
  }
  if (tensor_indices[index] == kTfLiteOptionalTensor) {
    TF_LITE_KERNEL_LOG(const_cast<TfLiteContext*>(context),
                       "Tensor at index %d was optional but was expected\n",
                       index);
    return kTfLiteError;
  }

  *tensor_index = tensor_indices[index];
  return kTfLiteOk;
}

// Same as above but returns -1 for invalid inputs instead of status + logging
// error.
inline int ValidateTensorIndexing(const TfLiteContext* context, int index,
                                  int max_size, const int* tensor_indices) {
  if (index >= 0 && index < max_size) {
    const int tensor_index = tensor_indices[index];
      return tensor_index;
  return -1;
}

inline TfLiteTensor* GetMutableInput(const TfLiteContext* context,
                                     const TfLiteNode* node, int index) {
  const int tensor_index = ValidateTensorIndexing(
      context, index, node->inputs->size, node->inputs->data);
  if (tensor_index < 0) {
    return nullptr;
  }
  return GetTensorAtIndex(context, tensor_index);
}

inline TfLiteStatus GetMutableInputSafe(const TfLiteContext* context,
                                        const TfLiteNode* node, int index,
                                        const TfLiteTensor** tensor) {
  int tensor_index;
  TF_LITE_ENSURE_OK(
      context, ValidateTensorIndexingSafe(context, index, node->inputs->size,
                                          node->inputs->data, &tensor_index));
  *tensor = GetTensorAtIndex(context, tensor_index);
  return kTfLiteOk;
TfLiteStatus GetInputSafe(const TfLiteContext* context, const TfLiteNode* node,
                          int index, const TfLiteTensor** tensor) {
  return GetMutableInputSafe(context, node, index, tensor);
}

  const int tensor_index = ValidateTensorIndexing(
      context, index, node->outputs->size, node->outputs->data);
  if (tensor_index < 0) {
    return nullptr;
  return GetTensorAtIndex(context, tensor_index);
}

TfLiteStatus GetOutputSafe(const TfLiteContext* context, const TfLiteNode* node,
                           int index, TfLiteTensor** tensor) {
  int tensor_index;
  TF_LITE_ENSURE_OK(
      context, ValidateTensorIndexingSafe(context, index, node->outputs->size,
                                          node->outputs->data, &tensor_index));
  *tensor = GetTensorAtIndex(context, tensor_index);
  return kTfLiteOk;
#ifndef TF_LITE_STATIC_MEMORY
TfLiteTensor* GetTemporary(TfLiteContext* context, const TfLiteNode* node,
                           int index) {
  const int tensor_index = ValidateTensorIndexing(
      context, index, node->temporaries->size, node->temporaries->data);
  if (tensor_index < 0) {
    return nullptr;
  }
  return GetTensorAtIndex(context, tensor_index);
}

TfLiteStatus GetTemporarySafe(const TfLiteContext* context,
                              const TfLiteNode* node, int index,
                              TfLiteTensor** tensor) {
  int tensor_index;
  TF_LITE_ENSURE_OK(context, ValidateTensorIndexingSafe(
                                 context, index, node->temporaries->size,
                                 node->temporaries->data, &tensor_index));
  *tensor = GetTensorAtIndex(context, tensor_index);
  return kTfLiteOk;
}

const TfLiteTensor* GetIntermediates(TfLiteContext* context,
                                     const TfLiteNode* node, int index) {
  const int tensor_index = ValidateTensorIndexing(
      context, index, node->intermediates->size, node->intermediates->data);
  if (tensor_index < 0) {
    return nullptr;
  }
  return GetTensorAtIndex(context, tensor_index);
}

TfLiteStatus GetIntermediatesSafe(const TfLiteContext* context,
                                  const TfLiteNode* node, int index,
                                  TfLiteTensor** tensor) {
  int tensor_index;
  TF_LITE_ENSURE_OK(context, ValidateTensorIndexingSafe(
                                 context, index, node->intermediates->size,
                                 node->intermediates->data, &tensor_index));
  *tensor = GetTensorAtIndex(context, tensor_index);
  return kTfLiteOk;
}
#endif  // TF_LITE_STATIC_MEMORY

// Same as `GetInput` but returns boolean and uses output argument for tensor.
//
//   TfLiteTensor* my_tensor;
//   TF_LITE_ENSURE_OK(context,
//                     GetInputSafe(context, node, kMyTensorIdx, &my_tensor));
//   // can use my_tensor directly from here onwards, it is not nullptr
//
// Should be used in cases where the binary size is too large.
TfLiteStatus GetInputSafe(const TfLiteContext* context, const TfLiteNode* node,
                          int index, const TfLiteTensor** tensor);

// Same as `GetOutput` but returns boolean and uses output argument for tensor.
//
//   TfLiteTensor* my_tensor;
//   TF_LITE_ENSURE_OK(context,
//                     GetOutputSafe(context, node, kMyTensorIdx, &my_tensor));
//   // can use my_tensor directly from here onwards, it is not nullptr
//
// Should be used in cases where the binary size is too large.
TfLiteStatus GetOutputSafe(const TfLiteContext* context, const TfLiteNode* node,
                           int index, TfLiteTensor** tensor);

TfLiteTensor* GetTemporary(TfLiteContext* context, const TfLiteNode* node,
                           int index);

// Same as `GetTemporary` but returns boolean and uses output argument for
// tensor.
//
//   TfLiteTensor* my_tensor;
//   TF_LITE_ENSURE_OK(context,
//                     GetTemporarySafe(context, node, kMyTensorIdx,
//                     &my_tensor));
//   // can use my_tensor directly from here onwards, it is not nullptr
//
// Should be used in cases where the binary size is too large.
TfLiteStatus GetTemporarySafe(const TfLiteContext* context,
                              const TfLiteNode* node, int index,
                              TfLiteTensor** tensor);
const TfLiteTensor* GetIntermediates(TfLiteContext* context,
                                     const TfLiteNode* node, int index);

// Same as `GetIntermediates` but returns boolean and uses output argument for
// tensor.
//
//   TfLiteTensor* my_tensor;
//   TF_LITE_ENSURE_OK(context,
//                     GetIntermediatesSafe(context, node, kMyTensorIdx,
//                     &my_tensor));
//   // can use my_tensor directly from here onwards, it is not nullptr
//
// Should be used in cases where the binary size is too large.
TfLiteStatus GetIntermediatesSafe(const TfLiteContext* context,
                                  const TfLiteNode* node, int index,
                                  TfLiteTensor** tensor);
#endif  // TF_LITE_STATIC_MEMORY

inline int NumDimensions(const TfLiteTensor* t) { return t->dims->size; }
inline int SizeOfDimension(const TfLiteTensor* t, int dim) {
  return t->dims->data[dim];
inline int NumInputs(const TfLiteNode* node) { return node->inputs->size; }
inline int NumOutputs(const TfLiteNode* node) { return node->outputs->size; }

#ifndef TF_LITE_STATIC_MEMORY
  const TfLiteTensor* input1;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, /*index=*/0, &input1));
  const TfLiteTensor* input2;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, /*index=*/1, &input2));

  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, /*index=*/0, &output));
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node,
                            ops::builtin::lstm::full::kInputTensor, &input));
  const TfLiteTensor* input_to_forget_weights;
  TF_LITE_ENSURE_OK(
      context,
      GetInputSafe(context, node,
                   ops::builtin::lstm::full::kInputToForgetWeightsTensor,
                   &input_to_forget_weights));
  const TfLiteTensor* input_to_cell_weights;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node,
                            ops::builtin::lstm::full::kInputToCellWeightsTensor,
                            &input_to_cell_weights));
  const TfLiteTensor* input_to_output_weights;
  TF_LITE_ENSURE_OK(
      context,
      GetInputSafe(context, node,
                   ops::builtin::lstm::full::kInputToOutputWeightsTensor,
                   &input_to_output_weights));
  const TfLiteTensor* recurrent_to_forget_weights;
  TF_LITE_ENSURE_OK(
      context,
      GetInputSafe(context, node,
                   ops::builtin::lstm::full::kRecurrentToForgetWeightsTensor,
                   &recurrent_to_forget_weights));
  const TfLiteTensor* recurrent_to_cell_weights;
  TF_LITE_ENSURE_OK(
      context,
      GetInputSafe(context, node,
                   ops::builtin::lstm::full::kRecurrentToCellWeightsTensor,
                   &recurrent_to_cell_weights));
  const TfLiteTensor* recurrent_to_output_weights;
  TF_LITE_ENSURE_OK(
      context,
      GetInputSafe(context, node,
                   ops::builtin::lstm::full::kRecurrentToOutputWeightsTensor,
                   &recurrent_to_output_weights));
  const TfLiteTensor* forget_gate_bias;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node,
                            ops::builtin::lstm::full::kForgetGateBiasTensor,
                            &forget_gate_bias));
  const TfLiteTensor* cell_gate_bias;
  TF_LITE_ENSURE_OK(
      context,
      GetInputSafe(context, node, ops::builtin::lstm::full::kCellGateBiasTensor,
                   &cell_gate_bias));
  const TfLiteTensor* output_gate_bias;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node,
                            ops::builtin::lstm::full::kOutputGateBiasTensor,
                            &output_gate_bias));
  TfLiteTensor* scratch_buffer;
  TF_LITE_ENSURE_OK(
      context, GetTemporarySafe(context, node, /*index=*/0, &scratch_buffer));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(
      context, GetOutputSafe(context, node,
                             ops::builtin::lstm::full::kOutputTensor, &output));