Variable Tensor API for TF Lite.

PiperOrigin-RevId: 200457602

1 files changed, 12 insertions, 1 deletions

diff --git a/tensorflow/contrib/lite/arena_planner_test.cc b/tensorflow/contrib/lite/arena_planner_test.cc
index 16171df10a..f0fd35216f 100644
--- a/tensorflow/contrib/lite/arena_planner_test.cc
+++ b/tensorflow/contrib/lite/arena_planner_test.cc
@@ -100,12 +100,18 @@ class TestGraph {
   std::vector<TfLiteTensor>* tensors() { return &tensors_; }
   const std::vector<int>& inputs() { return inputs_; }
   const std::vector<int>& outputs() { return outputs_; }
+  const std::vector<int>& variables() { return variables_; }
+
+  void SetVariables(const std::vector<int>& variables) {
+    variables_ = variables;
+  }
 
  private:
   std::vector<TfLiteNode> nodes_;
   std::vector<TfLiteTensor> tensors_;
   std::vector<int> inputs_;
   std::vector<int> outputs_;
+  std::vector<int> variables_;
 };
 
 // The GraphInfo for a TestGraph.
@@ -123,6 +129,9 @@ class TestGraphInfo : public GraphInfo {
   }
   const std::vector<int>& inputs() const override { return graph_->inputs(); }
   const std::vector<int>& outputs() const override { return graph_->outputs(); }
+  const std::vector<int>& variables() const override {
+    return graph_->variables();
+  }
 
  private:
   TestGraph* graph_;
@@ -306,13 +315,15 @@ TEST_F(ArenaPlannerTest, SimpleGraphWithPersistentTensor) {
                   {
                       /* in, out, tmp */
                       {{0, 1}, {2}, {}},   // First op
-                      {{2, 0}, {4}, {5}},  // Second op, with temporary
+                      {{2, 0}, {4}, {5}},  // Second op, with persistent
                       {{4, -1}, {3}, {}}   // Third op, with optional
                   },
                   {3});
 
   // Make #1 persistent so it goes into its own arena.
   (*graph.tensors())[1].allocation_type = kTfLiteArenaRwPersistent;
+  // The only use case for kTfLiteArenaRwPersistent is variable tensor now.
+  graph.SetVariables({1});
 
   SetGraph(&graph);
   Execute(0, 10);