aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/lite/arena_planner_test.cc
diff options
context:
space:
mode:
authorGravatar Yu-Cheng Ling <ycling@google.com>2018-06-13 14:48:22 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-06-13 14:50:40 -0700
commita3273e090f7ea8401ea283ad052350aeffa5fdc1 (patch)
tree9497eab175c8d8ed52a466e5869d33433e6333da /tensorflow/contrib/lite/arena_planner_test.cc
parent2f7f04a7a03003e8fe345667ddf0b088032f0e03 (diff)
Variable Tensor API for TF Lite.
PiperOrigin-RevId: 200457602
Diffstat (limited to 'tensorflow/contrib/lite/arena_planner_test.cc')
-rw-r--r--tensorflow/contrib/lite/arena_planner_test.cc13
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);