%18 = "tf.Identity"(%17) {device = ""} : (tensor<i1>) -> tensor<i1>

// CHECK: [[z:%.*]] = mlrt.cond [[cond]] @then @else
// CHECK: return [[z]]
func.func @while_cond_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> (tensor<i1>) {
  %r = "tf.If"(%cond, %x, %y, %z) {then_branch = @then, else_branch = @else, is_stateless = true} : (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>) -> tensor<i1>
  return %r : tensor<i1>
}

// CHECK-LABEL: func @while_body_if

    } else if (auto if_op = llvm::dyn_cast<TF::IfOp>(use.getOwner())) {
      auto then_branch = if_op.then_function();
      auto else_branch = if_op.else_function();
      assert(then_branch && else_branch);
      auto type_from_then = GetElementTypeFromAccess(
          then_branch.getArgument(use.getOperandNumber() - 1), module,
          infer_from_op);

    llvm::StringMap<PartitionedCallDecompositionInfo>*
        decomposed_partitioned_call_callees) {
  // 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(

    TF_Int32Tensor:$max_num_elements
  );
}

def TF_IfOp : TF_Op<"If", [DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
  let summary = "output = cond ? then_branch(input) : else_branch(input)";

  let description = [{
output = cond ? then_branch(input) : else_branch(input)

cond: A Tensor. If the tensor is a scalar of non-boolean type, the
    scalar is converted to a boolean according to the

  // CHECK: %[[GVAR2:.*]] = "tf.MlirLocalVarOp"() : () -> tensor<!tf_type.resource<tensor<5x3xf32>>>
  // 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

  // CHECK: %[[CLUSTER:.*]]:2 = "tf_device.cluster"()
  %2 = "tf_device.cluster"() ({
    // CHECK: %[[IF:.*]]:2 = "tf.If"(%[[ARG0]], %[[READ0]], %[[READ1]])
    %3:2 = "tf.If"(%arg0, %0, %1) {then_branch = @if_then, else_branch = @if_else,
        is_stateless = false}
      : (tensor<i1>, tensor<*x!tf_type.resource<tensor<4xf32>>>, tensor<*x!tf_type.resource<tensor<4xf32>>>)

    // and else_branch, and it has `recv_if_pred_node` as cond input.
    Node *if_oc_node = node_name_index["oc_if_if"];
    EXPECT_NE(if_oc_node, nullptr);
    Node *if_oc_node_cond_input;
    TF_CHECK_OK(if_oc_node->input_node(0, &if_oc_node_cond_input));
    EXPECT_EQ(if_oc_node_cond_input, recv_if_pred_node);

    // Check that then_branch outside compilation has node "identity_true_fn".

#include "tensorflow/core/public/version.h"

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();