Function TEST

tensorflow/core/common_runtime/function_test.cc:1889–1932

}  // end namespace

TEST(OptimizationTest, RemoveDeadNodes) {
  auto T = DT_INT32;
  auto func = FDH::Define(
      // Name
      "F",
      // Args
      {"x: int32"},
      // Return values
      {"y: int32"},
      // Attrs
      {},
      // Nodes
      {// a = Square<T>(x)
       {{"a"}, "Square", {"x"}, {{"T", T}}},
       // 1
       FDH::Const("o", 1),
       // A bunch of extra arithmetic that y doesn't depend on
       {{"x1"}, "Add", {"o", "o"}, {{"T", T}}},
       {{"x2"}, "Mul", {"a", "x1"}, {{"T", T}}},
       {{"x3"}, "Mul", {"x1", "x2"}, {{"T", T}}},
       // A stateful node.
       {{"keep_me"}, "RandomUniform", {"o"}, {{"T", T}, {"dtype", DT_FLOAT}}},
       // y = Add<T>(a, o)
       {{"y"}, "Add", {"a", "o"}, {{"T", T}}}});

  GraphDef expected;
  {
    Scope s = Scope::DisabledShapeInferenceScope();
    auto x = ops::_Arg(s.WithOpName("x"), DT_INT32, 0);
    auto o = ops::Const(s.WithOpName("o"), 1);
    auto keep_me = ops::RandomUniform(s.WithOpName("keep_me"), {o}, DT_FLOAT);
    auto x1 = ops::Add(s.WithOpName("x1"), o, o);
    auto a = ops::Square(s.WithOpName("a"), x);
    auto y = ops::Add(s.WithOpName("y"), a, o);
    auto x2 = ops::Mul(s.WithOpName("x2"), a, x1);
    auto x3 = ops::Mul(s.WithOpName("x3"), x1, x2);
    auto ret = ops::_Retval(s.WithOpName("y_RetVal"), y, 0);
    TF_ASSERT_OK(s.ToGraphDef(&expected));
  }
  TF_EXPECT_GRAPH_EQ(expected, Optimize(DoNothing, func));

  // TODO(zhifengc): Comes up another test case.
  TF_EXPECT_GRAPH_EQ(expected, Optimize(::tensorflow::RemoveDeadNodes, func));
}

TEST(OptimizationTest, RemoveIdentityNodes_Ref) {
  auto T = DT_FLOAT;