Function TEST_F

tensorflow/core/grappler/costs/graph_memory_test.cc:45–71

};

TEST_F(GraphMemoryTest, Basic) {
  TrivialTestGraphInputYielder fake_input(4, 1, 10, false, {"/CPU:0"});
  GrapplerItem item;
  CHECK(fake_input.NextItem(&item));
  item.feed.clear();

  GraphMemory memory(item);
  Status s = memory.InferStatically(devices_);
  TF_CHECK_OK(s);
  const GraphMemory::MemoryUsage& mem_usage =
      memory.GetPeakMemoryUsage("/CPU:0");
  EXPECT_EQ(120, mem_usage.used_memory);

  std::set<string> tensors;
  for (const auto& t : mem_usage.live_tensors) {
    tensors.insert(strings::StrCat(t.node, ":", t.output_id));
  }
  // When the execution of the 'Square' node completes, TF can start executing
  // 'Square_1' and release the memory used by 'x'. Since we can't be sure of
  // the order in which this takes place, in the worst case the 3 tensors are in
  // memory.
  std::set<string> expected;
  expected.insert("Square:0");
  expected.insert("Square_1:0");
  expected.insert("x:0");
  EXPECT_EQ(expected, tensors);
}

TEST_F(GraphMemoryTest, UnknownBatchSize) {
  TrivialTestGraphInputYielder fake_input(4, 1, -1, false, {"/CPU:0"});