Add AllocateTempIntermediateTensor() function call (#1041)

tensorflow/lite/micro/micro_context.cc

@@ -80,6 +80,16 @@ TfLiteTensor* MicroContext::AllocateTempOutputTensor(const TfLiteNode* node,
   return AllocateTempTfLiteTensor(tensor_index);
 }
 
+TfLiteTensor* MicroContext::AllocateTempIntermediateTensor(
+    const TfLiteNode* node, int index) {
+  const int tensor_index = GetTensorIndex(index, node->intermediates->size,
+                                          node->intermediates->data);
+  if (tensor_index < 0) {
+    return nullptr;
+  }
+  return AllocateTempTfLiteTensor(tensor_index);
+}
+
 void MicroContext::DeallocateTempTfLiteTensor(TfLiteTensor* tensor) {
   return allocator_.DeallocateTempTfLiteTensor(tensor);
 }

tensorflow/lite/micro/micro_context.h

@@ -73,6 +73,13 @@ class MicroContext {
   virtual TfLiteTensor* AllocateTempOutputTensor(const TfLiteNode* node,
                                                  int index);
 
+  // Returns a temporary TfLiteTensor struct for the specified intermediate
+  // tensor of a given mode. This is the recommended API over the deprecated
+  // GetIntermediates/GetIntermediatesSafe to get a temp intermediate tensor.
+  // The returned tensor shall be freed via calling DeallocateTempTfLiteTensor.
+  virtual TfLiteTensor* AllocateTempIntermediateTensor(const TfLiteNode* node,
+                                                       int index);
+
   // Deallocates a temp TfLiteTensor.
   // Virtual so that it can be faked for kernel tests.
   virtual void DeallocateTempTfLiteTensor(TfLiteTensor* tensor);

tensorflow/lite/micro/micro_context_test.cc

@@ -137,4 +137,20 @@ TF_LITE_MICRO_TEST(TestGetTempOutputTensor) {
   TF_LITE_MICRO_EXPECT_TRUE(invalid_output == nullptr);
 }
 
+TF_LITE_MICRO_TEST(TestGetTempIntermediateTensor) {
+  tflite::MicroContext micro_context = tflite::CreateMicroContext();
+
+  TfLiteNode node;
+  int intermediate_data[] = {1, 0};
+  node.intermediates = IntArrayFromInts(intermediate_data);
+
+  TfLiteTensor* output = micro_context.AllocateTempIntermediateTensor(&node, 0);
+  TF_LITE_MICRO_EXPECT_TRUE(output != nullptr);
+  micro_context.DeallocateTempTfLiteTensor(output);
+
+  TfLiteTensor* invalid_output =
+      micro_context.AllocateTempIntermediateTensor(&node, 1);
+  TF_LITE_MICRO_EXPECT_TRUE(invalid_output == nullptr);
+}
+
 TF_LITE_MICRO_TESTS_END