Function TEST

tensorflow/core/grappler/optimizers/data/inject_prefetch_test.cc:31–55

namespace {

TEST(MakeStateless, ParallelMap) {
  using test::function::NDef;
  GrapplerItem item;
  item.graph = test::function::GDef(
      {NDef("start", "Const", {}, {{"value", 0}, {"dtype", DT_INT32}}),
       NDef("stop", "Const", {}, {{"value", 10}, {"dtype", DT_INT32}}),
       NDef("step", "Const", {}, {{"value", 1}, {"dtype", DT_INT32}}),
       NDef("range", "RangeDataset", {"start", "stop", "step"}, {}),
       NDef("num_parallel_calls", "Const", {},
            {{"value", 1}, {"dtype", DT_INT32}}),
       graph_tests_utils::MakeParallelMapNode("map", "range",
                                              "num_parallel_calls", "XTimesTwo",
                                              /*sloppy=*/false)},
      // FunctionLib
      {
          test::function::XTimesTwo(),
      });

  InjectPrefetch optimizer;
  GraphDef output;
  TF_ASSERT_OK(optimizer.Optimize(nullptr, item, &output));
  EXPECT_TRUE(graph_utils::ContainsNodeWithOp("PrefetchDataset", output));
  int index = graph_utils::FindGraphNodeWithOp("PrefetchDataset", output);
  EXPECT_FALSE(output.node(index).attr().at("legacy_autotune").b());
}

TEST(MakeStateless, ParallelInterleave) {
  using test::function::NDef;