Function TEST_F

tensorflow/core/grappler/optimizers/function_optimizer_test.cc:38–82

class FunctionOptimizerTest : public GrapplerTest {};

TEST_F(FunctionOptimizerTest, InlineFunction_SimpleFunction) {
  using test::function::NDef;

  FunctionOptimizer optimizer(RewriterConfig::DEFAULT, true);

  // Build a graph to compute y = XTimesTwo(x)
  GrapplerItem item;
  item.graph = test::function::GDef(
      {NDef("x", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
       NDef("y", "XTimesTwo", {"x"}, {{"T", DT_FLOAT}}, kDevice),
       NDef("z", "Identity", {"y"}, {{"T", DT_FLOAT}}, kDevice)},
      // FunctionLib
      {
          test::function::XTimesTwo(),
      });

  GraphDef output;
  TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));

  const string arg0 = "Func/y/input/_0";
  const string ret0 = "Func/y/output/_1";

  const Tensor kTwo = test::AsScalar<int64>(2);
  GraphDef expected = test::function::GDef(
      {NDef("x", "Placeholder", {}, {{"dtype", DT_FLOAT}}),
       NDef(arg0, "Identity", {"x"}, {{"T", DT_FLOAT}}),
       NDef("y/two", "Const", {}, {{"dtype", DT_INT64}, {"value", kTwo}}),
       NDef("y/scale", "Cast", {"y/two"},
            {{"DstT", DT_FLOAT}, {"SrcT", DT_INT64}}),
       NDef("y/y", "Mul", {arg0, "y/scale"}, {{"T", DT_FLOAT}}),
       NDef(ret0, "Identity", {"y/y"}, {{"T", DT_FLOAT}}),
       NDef("z", "Identity", {ret0}, {{"T", DT_FLOAT}})},
      {});
  for (NodeDef& node : *expected.mutable_node()) node.set_device(kDevice);

  CompareGraphs(expected, output);

  Tensor pi = test::AsScalar<float>(3.14f);
  item.fetch = {"z"};
  item.feed.emplace_back("x", pi);
  auto tensors_expected = EvaluateFetchNodes(item);
  GrapplerItem optimized = item.WithGraph(std::move(output));
  auto tensors = EvaluateFetchNodes(optimized);
  test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}

TEST_F(FunctionOptimizerTest, InlineFunction_FixedTypeFunction) {
  using test::function::NDef;