/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/tf2xla/functionalize_control_flow.h"

#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/ops/control_flow_ops_internal.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/resource_variable_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/tf2xla/cc/ops/functional_ops.h"
#include "tensorflow/compiler/tf2xla/test_util.h"
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/graph/graph_constructor.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/util/equal_graph_def.h"

namespace tensorflow {
namespace {

// Returns the names of the "cond" and "body" functions for the While node
// in a graph.
Status FindWhileCondAndBody(const GraphDef& graph, NameAttrList* cond,
                            NameAttrList* body) {
  for (const NodeDef& node : graph.node()) {
    if (node.op() == "XlaWhile") {
      const NameAttrList* result;
      TF_RETURN_IF_ERROR(GetNodeAttr(node, "cond", &result));
      *cond = *result;
      TF_RETURN_IF_ERROR(GetNodeAttr(node, "body", &result));
      *body = *result;
      return Status::OK();
    }
  }
  return errors::NotFound("No XlaWhile node found in graph");
}

// Graph:
// x = array_ops.placeholder(dtypes.int32)
// y = control_flow_ops.while_loop(lambda i: i < 10, lambda i: i + 1, [x])
TEST(FunctionalizeControlFlow, OneLoopVar) {
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

    auto dummy = ops::Placeholder(scope.WithOpName("Dummy"), DT_INT32);

    auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32);
    auto enter =
        ops::internal::Enter(scope.WithOpName("while/Enter"), source, "aloop");
    // Add an unused Enter node. These should be ignored.
    auto enter2 =
        ops::internal::Enter(scope.WithOpName("while/Enter2"), source, "aloop");
    auto merge = ops::Merge(scope.WithOpName("while/Merge"),
                            std::initializer_list<Input>{enter, dummy});
    auto ten = ops::Const<int32>(
        scope.WithOpName("while/Less/y").WithControlDependencies(merge.output),
        10);
    auto less = ops::Less(scope.WithOpName("while/Less"), merge.output, ten);
    auto loop_cond = ops::LoopCond(scope.WithOpName("while/LoopCond"), less);
    auto switch_ =
        ops::Switch(scope.WithOpName("while/Switch"), merge.output, loop_cond);
    auto exit = ops::internal::Exit(scope.WithOpName("while/Exit"),
                                    switch_.output_false);
    auto identity =
        ops::Identity(scope.WithOpName("while/Identity"), switch_.output_true);
    auto one = ops::Const<int32>(
        scope.WithOpName("while/add/y").WithControlDependencies(identity), 1);
    auto add = ops::Add(scope.WithOpName("while/add"), identity, one);
    auto next_iteration =
        ops::NextIteration(scope.WithOpName("while/NextIteration"), add);

    auto sink = ops::Identity(scope.WithOpName("sink"), exit);

    // Remove the dummy node and add the loop backedge.
    scope.graph()->RemoveNode(dummy.node());
    scope.graph()->AddEdge(next_iteration.node(), 0, merge.output.node(), 1);

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

  FunctionLibraryDefinition library(OpRegistry::Global(), {});
  TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library));

  GraphDef graph_def;
  graph.ToGraphDef(&graph_def);

  NameAttrList cond_fn, body_fn;
  TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn));

  // Outer graph
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32);
    auto while_op =
        ops::XlaWhile(scope.WithOpName("while/LoopCond"),
                      std::initializer_list<Input>{source}, cond_fn, body_fn);
    auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]);
    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));
    TF_EXPECT_GRAPH_EQ(expected, graph_def);
  }

  // Condition graph
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0);
    auto ten = ops::Const<int32>(
        scope.WithOpName("while/Less/y").WithControlDependencies(arg), 10);
    auto less = ops::Less(scope.WithOpName("while/Less"), arg, ten);
    auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0);

    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));

    InstantiationResultForTest result;
    TF_EXPECT_OK(InstantiateFunctionForTest(cond_fn.name(), library, &result));

    EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types);
    EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types);
    TF_EXPECT_GRAPH_EQ(expected, result.gdef);
  }

  // Body graph.
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0);
    auto identity = ops::Identity(scope.WithOpName("while/Identity"), arg);
    auto one = ops::Const<int32>(
        scope.WithOpName("while/add/y").WithControlDependencies(identity), 1);
    auto add = ops::Add(scope.WithOpName("while/add"), identity, one);
    auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0);

    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));

    InstantiationResultForTest result;
    TF_EXPECT_OK(InstantiateFunctionForTest(body_fn.name(), library, &result));

    EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types);
    EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types);
    TF_EXPECT_GRAPH_EQ(expected, result.gdef);
  }
}

// Graph:
// x = array_ops.placeholder(dtypes.int32)
// y = array_ops.placeholder(dtypes.int32)
// cond = lambda (i, j): i + 3 < 10
// body = lambda (i, j): (i < 10, j * 2)
// z = control_flow_ops.while_loop(cond, body, [x, y])
TEST(FunctionalizeControlFlow, TwoLoopVars) {
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

    auto dummy = ops::Placeholder(scope.WithOpName("Dummy"), DT_INT32);

    auto x = ops::Placeholder(scope.WithOpName("Placeholder/x"), DT_INT32);
    auto y = ops::Placeholder(scope.WithOpName("Placeholder/y"), DT_INT32);
    auto enter_x =
        ops::internal::Enter(scope.WithOpName("while/Enter/x"), x, "aloop");
    auto enter_y =
        ops::internal::Enter(scope.WithOpName("while/Enter/y"), y, "aloop");
    auto merge_x = ops::Merge(scope.WithOpName("while/Merge/x"),
                              std::initializer_list<Input>{enter_x, dummy});
    auto merge_y = ops::Merge(scope.WithOpName("while/Merge/y"),
                              std::initializer_list<Input>{enter_y, dummy});

    // Loop condition
    auto three = ops::Const<int32>(scope.WithOpName("while/cond/three")
                                       .WithControlDependencies(merge_x.output),
                                   3);
    auto cond_add =
        ops::Add(scope.WithOpName("while/cond/Add"), merge_x.output, three);
    auto ten = ops::Const<int32>(scope.WithOpName("while/cond/ten")
                                     .WithControlDependencies(merge_x.output),
                                 10);
    auto less = ops::Less(scope.WithOpName("while/cond/Less"), cond_add, ten);
    auto loop_cond = ops::LoopCond(scope.WithOpName("while/LoopCond"), less);

    auto switch_x = ops::Switch(scope.WithOpName("while/Switch/x"),
                                merge_x.output, loop_cond);
    auto switch_y = ops::Switch(scope.WithOpName("while/Switch/y"),
                                merge_y.output, loop_cond);

    auto exit_x = ops::internal::Exit(scope.WithOpName("while/Exit/x"),
                                      switch_x.output_false);
    auto exit_y = ops::internal::Exit(scope.WithOpName("while/Exit/y"),
                                      switch_y.output_false);

    auto identity_x = ops::Identity(scope.WithOpName("while/Identity/x"),
                                    switch_x.output_true);
    auto identity_y = ops::Identity(scope.WithOpName("while/Identity/y"),
                                    switch_y.output_true);

    auto one = ops::Const<int32>(
        scope.WithOpName("while/add/one").WithControlDependencies(identity_x),
        1);
    auto two = ops::Const<int32>(
        scope.WithOpName("while/mul/two").WithControlDependencies(identity_x),
        2);

    auto add = ops::Add(scope.WithOpName("while/add"), identity_x, one);
    auto mul = ops::Add(scope.WithOpName("while/mul"), identity_y, two);
    auto next_iteration_x =
        ops::NextIteration(scope.WithOpName("while/NextIteration/x"), add);
    auto next_iteration_y =
        ops::NextIteration(scope.WithOpName("while/NextIteration/y"), mul);

    auto sink_x = ops::Identity(scope.WithOpName("sink_x"), exit_x);
    auto sink_y = ops::Identity(scope.WithOpName("sink_y"), exit_y);

    // Remove the dummy node and add the loop backedges.
    scope.graph()->RemoveNode(dummy.node());
    scope.graph()->AddEdge(next_iteration_x.node(), 0, merge_x.output.node(),
                           1);
    scope.graph()->AddEdge(next_iteration_y.node(), 0, merge_y.output.node(),
                           1);

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

  FunctionLibraryDefinition library(OpRegistry::Global(), {});
  TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library));

  GraphDef graph_def;
  graph.ToGraphDef(&graph_def);

  NameAttrList cond_fn, body_fn;
  TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn));

  // Outer graph.
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto x = ops::Placeholder(scope.WithOpName("Placeholder/x"), DT_INT32);
    auto y = ops::Placeholder(scope.WithOpName("Placeholder/y"), DT_INT32);
    auto while_op =
        ops::XlaWhile(scope.WithOpName("while/LoopCond"),
                      std::initializer_list<Input>{x, y}, cond_fn, body_fn);
    auto sink_x = ops::Identity(scope.WithOpName("sink_x"), while_op[0]);
    auto sink_y = ops::Identity(scope.WithOpName("sink_y"), while_op[1]);
    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));
    TF_EXPECT_GRAPH_EQ(expected, graph_def);
  }

  // Condition graph.
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0);
    auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1);
    auto three = ops::Const<int32>(scope.WithOpName("while/cond/three")
                                       .WithControlDependencies(arg0.output),
                                   3);
    auto cond_add =
        ops::Add(scope.WithOpName("while/cond/Add"), arg0.output, three);
    auto ten = ops::Const<int32>(
        scope.WithOpName("while/cond/ten").WithControlDependencies(arg0.output),
        10);
    auto less = ops::Less(scope.WithOpName("while/cond/Less"), cond_add, ten);
    auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0);

    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));

    InstantiationResultForTest result;
    TF_EXPECT_OK(InstantiateFunctionForTest(cond_fn.name(), library, &result));

    EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32}), result.arg_types);
    EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types);
    TF_EXPECT_GRAPH_EQ(expected, result.gdef);
  }

  // Body graph.
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0);
    auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1);

    auto identity_x = ops::Identity(scope.WithOpName("while/Identity/x"), arg0);
    auto identity_y = ops::Identity(scope.WithOpName("while/Identity/y"), arg1);

    auto one = ops::Const<int32>(
        scope.WithOpName("while/add/one").WithControlDependencies(identity_x),
        1);
    auto two = ops::Const<int32>(
        scope.WithOpName("while/mul/two").WithControlDependencies(identity_x),
        2);

    auto add = ops::Add(scope.WithOpName("while/add"), identity_x, one);
    auto mul = ops::Add(scope.WithOpName("while/mul"), identity_y, two);
    auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0);
    auto retval1 = ops::_Retval(scope.WithOpName("_retval1_RetVal"), mul, 1);

    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));

    InstantiationResultForTest result;
    TF_EXPECT_OK(InstantiateFunctionForTest(body_fn.name(), library, &result));

    EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32}), result.arg_types);
    EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32}), result.ret_types);
    TF_EXPECT_GRAPH_EQ(expected, result.gdef);
  }
}

// Example with nesting, loop-invariant arguments, and resource variables.
//
// accum = resource_variable_ops.ResourceVariable(1)
// x = array_ops.placeholder(2, dtype=dtypes.int32)
// y = 3 + x
//
// def inner_body(j, k):
//   add = state_ops.assign_add(accum, k * j + x)
//   with ops.control_dependencies([add]):
//     return [j + 1, k]
//
// def body(i):
//   m = control_flow_ops.while_loop(lambda j, k: j < 5, inner_body,
//                                   [1, y], name="inner")
//   with ops.control_dependencies(m):
//     return [i + 1]
//
// z = control_flow_ops.while_loop(lambda i: i < 10, body, [0], name="outer")
TEST(FunctionalizeControlFlow, Complex) {
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

    auto dummy = ops::Placeholder(scope.WithOpName("Dummy"), DT_INT32);

    auto x = ops::Placeholder(scope.WithOpName("x"), DT_INT32);
    auto three = ops::Const<int32>(scope.WithOpName("three"), 3);
    auto y = ops::Add(scope.WithOpName("y"), x, three);

    auto var = ops::VarHandleOp(scope.WithOpName("Variable"), DT_INT32,
                                TensorShape({}));

    // Outer loop
    auto zero = ops::Const<int32>(scope.WithOpName("outer/Const"), 0);
    auto enter_i =
        ops::internal::Enter(scope.WithOpName("outer/Enter_i"), zero, "outer");
    auto merge_i = ops::Merge(scope.WithOpName("outer/Merge_i"),
                              std::initializer_list<Input>{enter_i, dummy});
    auto ten = ops::Const<int32>(scope.WithOpName("outer/Less/y")
                                     .WithControlDependencies(merge_i.output),
                                 10);
    auto less_i =
        ops::Less(scope.WithOpName("outer/Less_i"), merge_i.output, ten);
    auto outer_loop_cond =
        ops::LoopCond(scope.WithOpName("outer/LoopCond"), less_i);
    auto switch_i = ops::Switch(scope.WithOpName("outer/Switch"),
                                merge_i.output, outer_loop_cond);
    auto exit_i = ops::internal::Exit(scope.WithOpName("outer/Exit"),
                                      switch_i.output_false);
    auto identity_i =
        ops::Identity(scope.WithOpName("outer/Identity"), switch_i.output_true);

    auto enter_x_outer =
        ops::internal::Enter(scope.WithOpName("outer/Enter_x"), x, "outer",
                             ops::internal::Enter::Attrs().IsConstant(true));
    auto enter_k_outer =
        ops::internal::Enter(scope.WithOpName("outer/Enter_k"), y, "outer",
                             ops::internal::Enter::Attrs().IsConstant(true));
    auto enter_var_outer =
        ops::internal::Enter(scope.WithOpName("outer/Enter_var"), var, "outer",
                             ops::internal::Enter::Attrs().IsConstant(true));

    // Inner loop
    auto one_j = ops::Const<int32>(
        scope.WithOpName("outer/j").WithControlDependencies(identity_i), 1);
    auto enter_j = ops::internal::Enter(scope.WithOpName("outer/inner/Enter_j"),
                                        one_j, "inner");
    auto enter_k =
        ops::internal::Enter(scope.WithOpName("outer/inner/Enter_k")
                                 .WithControlDependencies(identity_i),
                             enter_k_outer, "inner");
    auto enter_x = ops::internal::Enter(
        scope.WithOpName("outer/inner/Enter_x"), enter_x_outer, "inner",
        ops::internal::Enter::Attrs().IsConstant(true));
    auto enter_var = ops::internal::Enter(
        scope.WithOpName("outer/inner/Enter_var"), enter_var_outer, "inner",
        ops::internal::Enter::Attrs().IsConstant(true));

    auto merge_j = ops::Merge(scope.WithOpName("outer/inner/Merge_j"),
                              std::initializer_list<Input>{enter_j, dummy});
    auto merge_k = ops::Merge(scope.WithOpName("outer/inner/Merge_k"),
                              std::initializer_list<Input>{enter_k, dummy});

    auto five = ops::Const<int32>(scope.WithOpName("outer/inner/Five")
                                      .WithControlDependencies(merge_j.output),
                                  5);
    auto less_j =
        ops::Less(scope.WithOpName("outer/inner/Less_j"), merge_j.output, five);
    auto loop_cond = ops::LoopCond(scope.WithOpName("outer/LoopCond"), less_j);

    auto switch_j = ops::Switch(scope.WithOpName("outer/inner/Switch_j"),
                                merge_j.output, loop_cond);
    auto switch_k = ops::Switch(scope.WithOpName("outer/inner/Switch_k"),
                                merge_k.output, loop_cond);
    auto exit_j = ops::internal::Exit(scope.WithOpName("outer/inner/Exit_j"),
                                      switch_j.output_false);
    auto exit_k = ops::internal::Exit(scope.WithOpName("outer/inner/Exit_k"),
                                      switch_k.output_false);
    auto identity_j = ops::Identity(scope.WithOpName("outer/inner/Identity_j"),
                                    switch_j.output_true);
    auto identity_k = ops::Identity(scope.WithOpName("outer/inner/Identity_k"),
                                    switch_k.output_true);

    // Variable update
    auto mul_jk =
        ops::Mul(scope.WithOpName("outer/inner/mul"), identity_j, identity_k);
    auto add_jkx =
        ops::Add(scope.WithOpName("outer/inner/add"), mul_jk, enter_x);
    auto assign = ops::AssignAddVariableOp(
        scope.WithOpName("outer/inner/assign_add"), enter_var, add_jkx);

    auto one =
        ops::Const<int32>(scope.WithOpName("outer/inner/One")
                              .WithControlDependencies(
                                  gtl::ArraySlice<Operation>{assign.operation}),
                          1);
    auto add_j =
        ops::Add(scope.WithOpName("outer/inner/add_j"), identity_j, one);

    auto next_iteration_j = ops::NextIteration(
        scope.WithOpName("outer/inner/NextIteration_j"), add_j);
    auto next_iteration_k = ops::NextIteration(
        scope.WithOpName("outer/inner/NextIteration_k"), identity_k);

    // Body and backedge for outer loop.
    auto one_outer = ops::Const<int32>(
        scope.WithOpName("outer/add/y").WithControlDependencies(identity_i), 1);
    auto add_i =
        ops::Add(scope.WithOpName("outer/add")
                     .WithControlDependencies(gtl::ArraySlice<Operation>{
                         exit_j.output.op(), exit_k.output.op()}),
                 identity_i, one_outer);
    auto next_iteration_i =
        ops::NextIteration(scope.WithOpName("outer/NextIteration"), add_i);

    auto sink = ops::Identity(scope.WithOpName("sink"), exit_i);

    // Remove the dummy node and add the loop backedge.
    scope.graph()->RemoveNode(dummy.node());
    scope.graph()->AddEdge(next_iteration_i.node(), 0, merge_i.output.node(),
                           1);
    scope.graph()->AddEdge(next_iteration_j.node(), 0, merge_j.output.node(),
                           1);
    scope.graph()->AddEdge(next_iteration_k.node(), 0, merge_k.output.node(),
                           1);

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

  FunctionLibraryDefinition library(OpRegistry::Global(), {});
  TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library));

  GraphDef graph_def;
  graph.ToGraphDef(&graph_def);

  NameAttrList outer_cond_fn, outer_body_fn;
  TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &outer_cond_fn, &outer_body_fn));

  // Outer graph.
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto x = ops::Placeholder(scope.WithOpName("x"), DT_INT32);
    auto three = ops::Const<int32>(scope.WithOpName("three"), 3);
    auto y = ops::Add(scope.WithOpName("y"), x, three);

    auto var = ops::VarHandleOp(scope.WithOpName("Variable"), DT_INT32,
                                TensorShape({}));

    auto zero = ops::Const<int32>(scope.WithOpName("outer/Const"), 0);

    auto while_op = ops::XlaWhile(scope.WithOpName("outer/LoopCond"),
                                  std::initializer_list<Input>{zero, y, x, var},
                                  outer_cond_fn, outer_body_fn);
    auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]);
    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));
    TF_EXPECT_GRAPH_EQ(expected, graph_def);
  }

  // Outer condition graph.
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0);
    auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1);
    auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2);
    auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3);

    auto ten = ops::Const<int32>(
        scope.WithOpName("outer/Less/y").WithControlDependencies(arg0.output),
        10);
    auto less = ops::Less(scope.WithOpName("outer/Less_i"), arg0, ten);
    auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0);

    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));

    InstantiationResultForTest result;
    TF_EXPECT_OK(
        InstantiateFunctionForTest(outer_cond_fn.name(), library, &result));

    EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}),
              result.arg_types);
    EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types);
    TF_EXPECT_GRAPH_EQ(expected, result.gdef);
  }

  // Outer body graph.
  NameAttrList inner_cond_fn, inner_body_fn;
  {
    InstantiationResultForTest result;
    TF_EXPECT_OK(
        InstantiateFunctionForTest(outer_body_fn.name(), library, &result));

    // Find the inner condition and body names.
    TF_EXPECT_OK(
        FindWhileCondAndBody(result.gdef, &inner_cond_fn, &inner_body_fn));

    Scope scope = Scope::NewRootScope().ExitOnError();
    auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0);
    auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1);
    auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2);
    auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3);

    auto identity_i = ops::Identity(scope.WithOpName("outer/Identity"), arg0);
    auto one_j = ops::Const<int32>(
        scope.WithOpName("outer/j").WithControlDependencies(identity_i), 1);
    auto while_op =
        ops::XlaWhile(scope.WithOpName("outer/LoopCond_1"),
                      std::initializer_list<Input>{one_j, arg1, arg2, arg3},
                      inner_cond_fn, inner_body_fn);

    auto one_outer = ops::Const<int32>(
        scope.WithOpName("outer/add/y").WithControlDependencies(identity_i), 1);
    auto add_i =
        ops::Add(scope.WithOpName("outer/add")
                     .WithControlDependencies(gtl::ArraySlice<Operation>{
                         while_op[0].op(), while_op[1].op()}),
                 identity_i, one_outer);

    auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add_i, 0);
    auto retval1 = ops::_Retval(scope.WithOpName("_retval1_RetVal"), arg1, 1);
    auto retval2 = ops::_Retval(scope.WithOpName("_retval2_RetVal"), arg2, 2);

    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));

    EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}),
              result.arg_types);
    EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32}), result.ret_types);
    TF_EXPECT_GRAPH_EQ(expected, result.gdef);
  }

  // Inner condition graph.
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0);
    auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1);
    auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2);
    auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3);

    auto five = ops::Const<int32>(
        scope.WithOpName("outer/inner/Five").WithControlDependencies(arg0), 5);
    auto less_j = ops::Less(scope.WithOpName("outer/inner/Less_j"), arg0, five);
    auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less_j, 0);

    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));

    InstantiationResultForTest result;
    TF_EXPECT_OK(
        InstantiateFunctionForTest(inner_cond_fn.name(), library, &result));

    EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}),
              result.arg_types);
    EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types);
    TF_EXPECT_GRAPH_EQ(expected, result.gdef);
  }

  // Inner body graph.
  {
    Scope scope = Scope::NewRootScope().ExitOnError();
    auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0);
    auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1);
    auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2);
    auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3);

    auto identity_j =
        ops::Identity(scope.WithOpName("outer/inner/Identity_j"), arg0);
    auto identity_k =
        ops::Identity(scope.WithOpName("outer/inner/Identity_k"), arg1);

    auto mul_jk =
        ops::Mul(scope.WithOpName("outer/inner/mul"), identity_j, identity_k);
    auto add_jkx = ops::Add(scope.WithOpName("outer/inner/add"), mul_jk, arg2);
    auto assign = ops::AssignAddVariableOp(
        scope.WithOpName("outer/inner/assign_add"), arg3, add_jkx);

    auto one =
        ops::Const<int32>(scope.WithOpName("outer/inner/One")
                              .WithControlDependencies(
                                  gtl::ArraySlice<Operation>{assign.operation}),
                          1);
    auto add_j =
        ops::Add(scope.WithOpName("outer/inner/add_j"), identity_j, one);

    auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add_j, 0);
    auto retval1 =
        ops::_Retval(scope.WithOpName("_retval1_RetVal"), identity_k, 1);
    auto retval2 = ops::_Retval(scope.WithOpName("_retval2_RetVal"), arg2, 2);

    GraphDef expected;
    TF_EXPECT_OK(scope.ToGraphDef(&expected));

    InstantiationResultForTest result;
    TF_EXPECT_OK(
        InstantiateFunctionForTest(inner_body_fn.name(), library, &result));

    EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}),
              result.arg_types);
    EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32}), result.ret_types);
    TF_EXPECT_GRAPH_EQ(expected, result.gdef);
  }
}

}  // namespace
}  // namespace tensorflow
