Function TEST

tensorflow/core/framework/function_test.cc:66–111

)doc");

TEST(TFunc, SquarePlusOne) {
  auto fdef = FDH::Create(
      // Name
      "SquarePlusOne",
      // Inputs
      {"x: T"},
      // Outputs
      {"y: T"},
      // Attrs
      {"T: {float, double, int32, int64}"},
      // Nodes
      {// a = Square<T>(x)
       {{"a"}, "Square", {"x"}, {{"T", "$T"}}},
       // o = One<T>()
       // NOTE: We can also have a Cast<Tin, Tout>(x) instead.
       {{"o"}, "One", {}, {{"T", "$T"}}},
       // y = Add<T>(a, o)
       {{"y"}, "Add", {"a:y", "o:y"}, {{"T", "$T"}}}},
      // Returns
      {{"y", "y:z:0"}});

  const char* e = R"P(
SquarePlusOne[T:{float, double, int32, int64}](x:T) -> (y:T) {
  a = Square[T=$T](x)
  o = One[T=$T]()
  y = Add[T=$T](a:y, o:y)
  return y = y:z:0
}
)P";
  EXPECT_EQ(DebugString(fdef), e);

  // Instantiate one with T=float
  InstantiationResult result;
  TF_ASSERT_OK(
      InstantiateFunction(fdef, Attrs({{"T", DT_FLOAT}}), GetOpSig, &result));
  const char* e2 = R"P(
(x:float) -> (y:float) {
  a = Square[T=float](x)
  o = One[T=float]()
  y = Add[T=float](a, o)
}
)P";
  EXPECT_EQ(result.arg_types, DataTypeVector({DT_FLOAT}));
  EXPECT_EQ(result.ret_types, DataTypeVector({DT_FLOAT}));
  EXPECT_EQ(DebugString(result.nodes), e2);
}

TEST(TFunc, ControlDep) {
  auto fdef = FDH::Create(