tensors, whose types are the same as what else_branch returns.
else_branch: A function that takes 'inputs' and returns a list of
    tensors.  whose types are the same as what then_branch returns.
  }];

  let arguments = (ins
    TF_Tensor:$cond,
    Variadic<TF_Tensor>:$input,

    FlatSymbolRefAttr:$then_branch,
    FlatSymbolRefAttr:$else_branch,

  // CHECK: "tf.If"(%[[COND]], %[[VAR]], %[[GVAR1]], %[[GVAR2]])
  %1 = "tf.If"(%cond, %ta#0) {
    then_branch = @then_branch, else_branch = @else_branch, device = "", is_stateless = false}
       : (tensor<i1>, tensor<!tf_type.resource>) -> tensor<!tf_type.resource>

  // CHECK: [[cond_res:%.*]]:2 = tfrt.cond [[cond]]
  // CHECK-SAME: @cond_true @cond_false(%arg0, %arg2) : (!tfrt.chain, !tfrt_fallback.tf_tensor)
  %2 = "tf.If"(%arg0, %arg1) {else_branch = @cond_false, then_branch = @cond_true, is_stateless = true} : (tensor<i1>, tensor<i32>) -> tensor<i32>
  // CHECK: [[out_ch:%.*]] = tfrt.merge.chains [[cond_res]]#0, %arg0 : !tfrt.chain, !tfrt.chain

  %tmp1 = "tfl.add"(%const, %const) { fused_activation_function = "NONE" } : (tensor<1xf32>,tensor<1xf32>) -> (tensor<1xf32>)
  %tmp2 = "tf.If"(%tmp1, %tmp1) { then_branch=@id, else_branch=@id, is_stateless = true } : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
  %tmp3 = "tfl.add"(%const, %tmp2) { fused_activation_function = "NONE" } : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)

    // CHECK: tf_device.return [[TRUE]] :
    tf_device.return %13 : tensor<i1>
  }) : () -> tensor<i1>
  func.return
}

// CHECK-LABEL: func @case_branch
func.func @case_branch(%arg0: !tf_res) -> tensor<i1> {
  // CHECK: [[TRUE:%.+]] = "tf.Const"
  // CHECK-SAME: value = dense<true> : tensor<i1>
  %0 = "tf.VarIsInitializedOp"(%arg0) : (!tf_res) -> tensor<i1>

  // 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>
  }
}

// -----