syncing with google3 change (#1301)

tensorflow/lite/micro/kernels/kernel_util.cc

@@ -40,9 +40,10 @@ int ValidateTensorIndexing(const TfLiteContext* context, int index,
 TfLiteRegistration RegisterOp(
     void* (*init)(TfLiteContext* context, const char* buffer, size_t length),
     TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node),
-    TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node)) {
+    TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node),
+    void (*free)(TfLiteContext* context, void* buffer)) {
   return {/*init=*/init,
-          /*free=*/nullptr,
+          /*free=*/free,
           /*prepare=*/prepare,
           /*invoke=*/invoke,
           /*profiling_string=*/nullptr,

tensorflow/lite/micro/kernels/kernel_util.h

@@ -31,7 +31,8 @@ namespace micro {
 TfLiteRegistration RegisterOp(
     void* (*init)(TfLiteContext* context, const char* buffer, size_t length),
     TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node),
-    TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node));
+    TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node),
+    void (*free)(TfLiteContext* context, void* buffer) = nullptr);
 
 // Prints out n bytes in a int8_t buffer as hex
 void PrintNBytes(const int8_t* tensor_data, int n_bytes,

tensorflow/lite/micro/micro_interpreter.cc

@@ -317,7 +317,17 @@ TfLiteTensor* MicroInterpreter::output(size_t index) {
   }
   return output_tensors_[index];
 }
+// Repurposing free subgraphs to reset state for some ops for now
+// will reset api is made. See b/220940833#comment25 for more context.
+TfLiteStatus MicroInterpreter::Reset() {
+  TfLiteStatus status = graph_.FreeSubgraphs();
+  if (status != kTfLiteOk) {
+    return status;
+  }
+  return graph_.ResetVariableTensors();
+}
 
+// TODO: remove this API completely in favor of MicroInterpreter::Reset
 TfLiteStatus MicroInterpreter::ResetVariableTensors() {
   return graph_.ResetVariableTensors();
 }

tensorflow/lite/micro/micro_interpreter.h

@@ -31,7 +31,7 @@ limitations under the License.
 #include "tensorflow/lite/portable_type_to_tflitetype.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
-// Copied from tensorflow/lite/version.h to avoid a dependency chain into
+/// Copied from tensorflow/lite/version.h to avoid a dependency chain into
 // tensorflow/core.
 #define TFLITE_SCHEMA_VERSION (3)
 
@@ -116,6 +116,11 @@ class MicroInterpreter {
     return nullptr;
   }
 
+  // Reset the state to be what you would expect when the interpreter is first
+  // created. i.e. after Init and Prepare is called for the very first time.
+  TfLiteStatus Reset();
+
+  // TODO: remove this in favor of Reset()
   // Reset all variable tensors to the default value.
   TfLiteStatus ResetVariableTensors();