Mark variable tensors as known in SubgraphWriter::CheckInputOutput.

PiperOrigin-RevId: 348859590
Change-Id: I234b9c170c24863e6221167b27a7dcccd9bf126c

Author: Abdurrahman Akkas

tensorflow/lite/tools/serialization/BUILD

@@ -36,6 +36,7 @@ cc_library(
         "//tensorflow/lite/schema:schema_fbs_with_reflection",
         "//tensorflow/lite/schema:schema_utils",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
     ],
 )
 

tensorflow/lite/tools/serialization/writer_lib.cc

@@ -19,6 +19,7 @@ limitations under the License.
 #include <unordered_map>
 #include <unordered_set>
 
+#include "absl/container/flat_hash_set.h"
 #include "tensorflow/lite/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/context_util.h"
@@ -326,7 +327,9 @@ TfLiteStatus SubgraphWriter::RegisterCustomWriter(
 TfLiteStatus SubgraphWriter::CheckInputOutput(
     const std::vector<int>& inputs, const std::vector<int>& outputs,
     const std::vector<int>& execution_plan) {
-  std::unordered_set<int> known_tensors(inputs.begin(), inputs.end());
+  absl::flat_hash_set<int> known_tensors(inputs.begin(), inputs.end());
+  known_tensors.insert(subgraph_->variables().begin(),
+                       subgraph_->variables().end());
   // Scan execution plan and confirm input tensors are known before each node
   // executes. Then append output tensors to known tensors.
   for (int op_index : execution_plan) {

tensorflow/lite/tools/serialization/writer_lib_test.cc

@@ -22,6 +22,7 @@ limitations under the License.
 #include <tuple>
 
 #include <gtest/gtest.h>
+#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/register.h"
@@ -230,6 +231,50 @@ TEST_P(SingleSubgraphTest, CustomInputOutputErrorCasesTest) {
             kTfLiteOk);
 }
 
+// Tests if SetCustomInputOutput handles variable tensors correctly.
+TEST_P(SingleSubgraphTest, CustomInputOutputVariableTensorTest) {
+  Interpreter interpreter;
+  tflite::ops::builtin::BuiltinOpResolver resolver;
+
+  // Create tensors.
+  interpreter.AddTensors(3);
+  interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "a", {3},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadWrite(1, kTfLiteFloat32, "b", {3},
+                                           TfLiteQuantization(),
+                                           /*is_variable=*/true);
+  interpreter.SetTensorParametersReadWrite(2, kTfLiteFloat32, "c", {3},
+                                           TfLiteQuantization());
+  interpreter.SetInputs({0});
+  interpreter.SetOutputs({2});
+  interpreter.SetVariables({1});
+
+  // Create an Add node.
+  TfLiteAddParams* builtin_data =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  builtin_data->activation = kTfLiteActNone;
+  builtin_data->pot_scale_int16 = false;
+  interpreter.AddNodeWithParameters({0, 1}, {2}, nullptr, 0,
+                                    reinterpret_cast<void*>(builtin_data),
+                                    resolver.FindOp(BuiltinOperator_ADD, 1));
+
+  // Write model to file.
+  const std::string test_file = CreateFilePath("test_variables.tflite");
+  SubgraphWriter writer(&interpreter.primary_subgraph());
+  EXPECT_EQ(writer.SetCustomInputOutput(/*inputs=*/{0}, /*outputs=*/{2},
+                                        /*execution_plan=*/{0}),
+            kTfLiteOk);
+  writer.Write(test_file);
+
+  // Read model and test.
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(test_file.c_str());
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+  CHECK_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+}
+
 TEST_P(SingleSubgraphTest, PerTensorQuantizedModelTest) {
   Interpreter interpreter;
   interpreter.AddTensors(3);