// CHECK-NOT: Optional
  %0 = builtin.unrealized_conversion_cast to tensor<5xi1>
  %1 = "tf.OptionalNone"() : () -> tensor<!tf_type.variant<tensor<f32>>>
  %2 = "tf.If"(%0, %1) <{else_branch = @false, is_stateless = false, then_branch = @true}>
      : (tensor<5xi1>, tensor<!tf_type.variant<tensor<f32>>>) -> (tensor<!tf_type.variant<tensor<f32>>>)
  return
}

// CHECK: func private @tf.IfRegion_else
// CHECK-NEXT: "tf.Neg"
// CHECK: func private @tf.IfRegion_then
// CHECK-NEXT: "tf.Abs"

func.func @testNested(%arg0: tensor<i1>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  // CHECK: "tf.If"({{.+}}) <{else_branch = @tf.IfRegion1_else, {{.+}} then_branch = @tf.IfRegion1_then}
  %0 = "tf.IfRegion"(%arg0) ({

                              /*operands=*/else_returns);

  // Then branch gets the data and calls the reduce_function.
  auto& then_branch = dataset_if.getThenBranch();
  then_branch.push_back(new Block);
  builder.setInsertionPointToEnd(&then_branch.front());
  // Add iterator operational data access inside if.
  auto get_value = builder.create<TF::OptionalGetValueOp>(loc, dataset_types,

    bool* has_outside_compilation) {
  // Instantiate "then_branch" and "else_branch".
  NameAttrList then_branch, else_branch;
  TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "then_branch", &then_branch));
  TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "else_branch", &else_branch));

  // Extract outside compilation for then_branch and else_branch.
  bool then_branch_has_outside_compilation = false;

    %4 = "tf.If"(%3, %3, %arg0, %cst) {_tpu_replicate = "cluster", _xla_outside_compilation = "0", device = "", else_branch = @if_else_branch, is_stateless = false, then_branch = @if_then_branch} : (tensor<*xi1>, tensor<*xi1>, tensor<*xi64>, tensor<i64>) -> tensor<*xi1>
    func.return %4 : tensor<*xi1>
  }

  %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