Return kTfLiteError if case of error

Author: veblush

tensorflow/lite/micro/micro_interpreter_graph.cc

@@ -307,7 +307,11 @@ TfLiteStatus MicroInterpreterGraph::ResetVariableTensors() {
        subgraph_idx++) {
     const SubGraph* subgraph = (*subgraphs_)[subgraph_idx];
     for (size_t i = 0; i < subgraph->tensors()->size(); ++i) {
-      TF_LITE_ENSURE_STATUS(ResetVariableTensor(i, subgraph_idx));
+      auto* tensor = subgraph->tensors()->Get(i);
+      if (tensor->is_variable()) {
+        TF_LITE_ENSURE_STATUS(ResetTensorData(
+            tensor, &subgraph_allocations_[subgraph_idx].tensors[i]));
+      }
     }
   }
   if (resource_variables_ != nullptr) {
@@ -319,21 +323,43 @@ TfLiteStatus MicroInterpreterGraph::ResetVariableTensors() {
 
 TfLiteStatus MicroInterpreterGraph::ResetVariableTensor(int tensor_index,
                                                         int subgraph_index) {
+  if (static_cast<size_t>(subgraph_index) >= subgraphs_->size()) {
+    MicroPrintf("Accessing subgraph %d but only %d subgraphs found",
+                subgraph_index, subgraphs_->size());
+    return kTfLiteError;
+  }
   const SubGraph* subgraph = (*subgraphs_)[subgraph_index];
+  if (subgraph->tensors() == nullptr ||
+      static_cast<size_t>(tensor_index) >= subgraph->tensors()->size()) {
+    MicroPrintf(
+        "Accessing tensor %d but only %d tensors found in subgraph %d",
+        tensor_index,
+        (subgraph->tensors() != nullptr ? subgraph->tensors()->size() : 0),
+        subgraph_index);
+    return kTfLiteError;
+  }
   auto* tensor = subgraph->tensors()->Get(tensor_index);
-  if (tensor->is_variable()) {
-    size_t buffer_size;
-    TF_LITE_ENSURE_STATUS(TfLiteEvalTensorByteLength(
-        &subgraph_allocations_[subgraph_index].tensors[tensor_index],
-        &buffer_size));
-
-    int value = 0;
-    if (tensor->type() == tflite::TensorType_INT8) {
-      value = tensor->quantization()->zero_point()->Get(0);
-    }
-    memset(subgraph_allocations_[subgraph_index].tensors[tensor_index].data.raw,
-           value, buffer_size);
+  if (!tensor->is_variable()) {
+    MicroPrintf("Accessing tensor %d in subgraph %d which is not a variable",
+                tensor_index, subgraph_index);
+    return kTfLiteError;
   }
+
+  return ResetTensorData(
+      tensor, &subgraph_allocations_[subgraph_index].tensors[tensor_index]);
+}
+
+TfLiteStatus MicroInterpreterGraph::ResetTensorData(
+    const tflite::Tensor* tensor, TfLiteEvalTensor* eval_tensor) {
+  size_t buffer_size;
+  TF_LITE_ENSURE_STATUS(TfLiteEvalTensorByteLength(eval_tensor, &buffer_size));
+
+  int value = 0;
+  if (tensor->type() == tflite::TensorType_INT8) {
+    value = tensor->quantization()->zero_point()->Get(0);
+  }
+  memset(eval_tensor->data.raw, value, buffer_size);
+
   return kTfLiteOk;
 }
 

tensorflow/lite/micro/micro_interpreter_graph.h

@@ -103,6 +103,9 @@ class MicroInterpreterGraph : public MicroGraph {
   MicroResourceVariables* GetResourceVariables() { return resource_variables_; }
 
  private:
+  TfLiteStatus ResetTensorData(const tflite::Tensor* tensor,
+                               TfLiteEvalTensor* eval_tensor);
+
   TfLiteContext* context_;
   const Model* model_;
   MicroAllocator* allocator_;

tensorflow/lite/micro/micro_interpreter_graph_test.cc

@@ -75,12 +75,20 @@ TF_LITE_MICRO_TEST(TestResetVariableTensor) {
     for (size_t i = 0; i < buffer_size; ++i) {
       non_variable_tensor->data.uint8[i] = 0xBB;
     }
-    TF_LITE_MICRO_EXPECT_EQ(
-        kTfLiteOk, interpreter.ResetVariableTensor(non_variable_tensor_idx, 0));
+    TF_LITE_MICRO_EXPECT_EQ(kTfLiteError, interpreter.ResetVariableTensor(
+                                              non_variable_tensor_idx, 0));
     for (size_t i = 0; i < buffer_size; ++i) {
       TF_LITE_MICRO_EXPECT_EQ(0xBB, non_variable_tensor->data.uint8[i]);
     }
   }
+
+  // Test invalid tensor index.
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteError,
+                          interpreter.ResetVariableTensor(100, 0));
+
+  // Test invalid subgraph index.
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteError,
+                          interpreter.ResetVariableTensor(1, 100));
 }
 
 TF_LITE_MICRO_TESTS_END