// Rewrite the branches.
  Region& then_branch = if_op.getThenBranch();
  Region& else_branch = if_op.getElseBranch();
  if (failed(DecomposeTensorListOpsInternal(
          &then_branch.front(), module, buffer_to_size,
          decomposed_partitioned_call_callees)))
    return failure();
  if (failed(DecomposeTensorListOpsInternal(
          &else_branch.front(), module, buffer_to_size,
          decomposed_partitioned_call_callees)))

  // Build the XLA computation func.
  // "const0" (bool)
  // "const1" (int32)
  // "if0" (pred = "const0", input = "const1", then_branch = "true_fn",
  //        else_branch = "false_fn")
  FunctionDefLibrary fdl;
  {
    tensorflow::Scope s = tensorflow::Scope::NewRootScope();
    Output arg = ops::_Arg(s.WithOpName("arg"), DT_INT32, 0);

namespace tensorflow {

TEST(RearrangeFunctionArgumentForFunctionTest, Basic) {
  FunctionDefLibrary fdl;
  {
    // Function for StatefulPartitionedCall's "f", If's
    // "then_branch"/"else_branch".
    // "arg0" (T=DT_RESOURCE), "arg1" (T=DT_BOOL)
    // "ret0" = "arg1"
    // "ret1" = "arg0"
    tensorflow::Scope s = tensorflow::Scope::NewRootScope();

  if (!s.ok()) {
    *uncompilable_reason = s.message();
    return false;
  }
  return true;
}

// Tests whether 'if_node' is compilable. Every operator in the then_branch and
// else_branch functions must be compilable for 'if_node' to be compilable.
bool RecursiveCompilabilityChecker::IsCompilableIf(
    const Node& if_node, FunctionLibraryRuntime* lib_runtime,
    std::vector<StackFrameView>* stack_trace,

    %1 = "tf.If"(%arg0, %arg1, %handle) {
      then_branch = @testIfThen, else_branch = @testIfElse, is_stateless = false
    } : (tensor<i1>, tensor<2xf32>, tensor<!tf_type.resource<tensor<0xf32>>>) -> tensor<2xf32>

    func.return %1 : tensor<2xf32>
  }
}

// -----

  %0 = tf_executor.graph {
    %output, %control = tf_executor.island {
      // CHECK: {{%.+}}, [[IF_CONTROL:%.+]] = tf_executor.island wraps "tf.If"
      %1 = "tf.If"(%arg0, %arg1) {device = "", else_branch = @stateful_composite_op_control_else, is_stateless = false, then_branch = @stateful_composite_op_control_then} : (tensor<i1>, tensor<*x!tf_type.resource<tensor<i32>>>) -> tensor<i32>

// CHECK:  [[SHAPE_1:%.*]] = "tf.Shape"([[INPUT]]) : (tensor<3x10xf32>) -> tensor<?xi32>
// CHECK:  [[RESULT:%.*]] = "tf.If"([[DIFF_RES]], [[INPUT]], [[SHAPE_1]], [[SIZE_DIFF]], [[SIZE]]) <{else_branch = @cond_false, is_stateless = true, then_branch = @cond_true}> : (tensor<i1>, tensor<3x10xf32>, tensor<?xi32>, tensor<i32>, tensor<i32>) -> tensor<?x10xf32>
// CHECK:  return [[RESULT]] : tensor<?x10xf32>
}

  // CHECK-NOT: tf.Stack
  %stack = "tf.StackV2"(%max_size) {elem_type = f32, stack_name = "s"} : (tensor<i32>) -> tensor<!tf_type.resource>
  %if_op = "tf.If"(%arg0, %stack) {then_branch = @if_then, else_branch = @if_else, is_stateless = false}
    : (tensor<i1>, tensor<!tf_type.resource>) -> tensor<!tf_type.resource>
  // CHECK: "tf.Slice"
  %pop = "tf.StackPopV2"(%if_op) : (tensor<!tf_type.resource>) -> tensor<f32>

      NodeBuilder(kFunctionalIfNodeName, "If", &root.graph()->flib_def())
          .Input(predicate.node())
          .Input(if_inputs)
          .Attr("then_branch", uncompilable_fn1_attr)
          .Attr("else_branch", uncompilable_fn2_attr)
          .Attr("Tout", {DT_INT32})
          .Finalize(root.graph(), &if_node));
  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  TF_ASSERT_OK(root.ToGraph(graph.get()));

      %outputs_65,...