Function TEST

tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc:65–180

Graph: x = array_ops.placeholder(dtypes.int32) y = array_ops.placeholder(dtypes.int32) z = control_flow_ops.cond( math_ops.less(y, x), lambda: math_ops.multiply(y, 17), lambda: math_ops.add(x, 23))

//     math_ops.less(y, x), lambda: math_ops.multiply(y, 17),
//     lambda: math_ops.add(x, 23))
TEST(FunctionalizeControlFlow, Conditional) {
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

    auto x = ops::Placeholder(scope.WithOpName("x"), DT_INT32);
    auto y = ops::Placeholder(scope.WithOpName("y"), DT_INT32);
    auto less = ops::Less(scope.WithOpName("cond/Less"), y, x);
    auto switch_1 = ops::Switch(scope.WithOpName("cond/Switch"), less, less);

    auto identity_t =
        ops::Identity(scope.WithOpName("cond/Identity"), switch_1.output_true);
    auto seventeen = ops::Const<int32>(
        scope.WithOpName("cond").WithControlDependencies(identity_t), 17);
    auto switch_2 = ops::Switch(scope.WithOpName("cond/Switch"), y, less);
    auto mul = ops::Multiply(scope.WithOpName("cond/Mul"), switch_2.output_true,
                             seventeen);

    auto identity_f =
        ops::Identity(scope.WithOpName("cond/Identity"), switch_1.output_false);
    auto twenty_three = ops::Const<int32>(
        scope.WithOpName("cond").WithControlDependencies(identity_f), 23);
    auto switch_3 = ops::Switch(scope.WithOpName("cond/Switch"), x, less);
    auto add = ops::Add(scope.WithOpName("cond/false/add"),
                        switch_3.output_false, twenty_three);

    auto merge = ops::Merge(scope.WithOpName("cond/Merge"),
                            std::initializer_list<Input>{add, mul});

    TF_EXPECT_OK(scope.ToGraph(&graph));
  }

  FunctionLibraryDefinition library(OpRegistry::Global(), {});
  GraphDef optimized_graph_def;
  graph.ToGraphDef(&optimized_graph_def);
  TF_ASSERT_OK(
      FunctionalizeControlFlowForGraphDef(&optimized_graph_def, &library));
  TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library));
  GraphDef converted_graph_def;
  graph.ToGraphDef(&converted_graph_def);

  for (const GraphDef& graph_def : {optimized_graph_def, converted_graph_def}) {
    string op_name;
    NameAttrList then_fn;
    NameAttrList else_fn;
    TF_EXPECT_OK(FindIfThenAndElse(graph_def, &op_name, &then_fn, &else_fn));
    InstantiationResultForTest else_result;
    TF_EXPECT_OK(
        InstantiateFunctionForTest(else_fn.name(), library, &else_result));

    // Outer graph
    {
      Scope scope = Scope::NewRootScope().ExitOnError();
      auto y = ops::Placeholder(scope.WithOpName("y"), DT_INT32);
      auto x = ops::Placeholder(scope.WithOpName("x"), DT_INT32);
      auto less = ops::Less(scope.WithOpName("cond/Less"), y, x);
      auto if_op = ops::If(scope.WithOpName(op_name), less,
                           std::initializer_list<Input>{less, y, x}, {DT_INT32},