// In the newly cloned function, check that we have a _tf.If operation and capture the then and else branch.
// CHECK: func private @[[FUNCTIONALIZE_FUNC]]
// CHECK: "tf.If"
// CHECK-SAME:  else_branch = @[[ELSE_FUNC:[A-Za-z0-9_]*]]
// CHECK-SAME:  then_branch = @[[THEN_FUNC:[A-Za-z0-9_]*]]

// We expect the _tf.Add in the else func and the _tf.Mul in the then func

// CHECK: func private @[[ELSE_FUNC]]

from tensorflow.python.ops import control_flow_ops

# Tests V1 control flow is functionalized.

# CHECK-NOT: tf_executor.Switch
# CHECK-NOT: tf_executor.Merge
# CHECK: "tf.If"
# CHECK-SAME: else_branch = @"key/[[else:[a-zA-Z_0-9]+]]"
# CHECK-SAME: then_branch = @"key/[[then:[a-zA-Z_0-9]+]]"

# CHECK: func private @"key/[[else]]"(
# CHECK-SAME: tf._original_func_name
# CHECK: func private @"key/[[then]]"(

  %cond = "tf.Const"() {device = "/CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
  // CHECK: "tf.If"([[cond]], [[handle]])
  %x = "tf.If"(%cond, %handle) {then_branch = @some_other_func, else_branch = @some_other_func, is_stateless = false} : (tensor<i1>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  %r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>

  %16 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>

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

      if (type_from_body.has_value()) return type_from_body;
    } 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);

          then_branch = @if_then_and_else,
          else_branch = @if_then_and_else, is_stateless = false}
        : (tensor<i1>, !tf_res,
           !tf_res) -> ()
      "tf.If"(%arg1, %var_handle, %id0) {
      // expected-error@above {{Conflicting device assignment for resource}}
          then_branch = @if_then_and_else,
          else_branch = @if_then_and_else,
          is_stateless = false}

// CHECK-NEXT: }

func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  %0 = "tfl.less"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1>
  %1 = "tf.If"(%0, %arg0, %arg1) {else_branch = @cond_false, then_branch = @cond_true, is_stateless = false} : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  func.return %1 : tensor<1xf32>
}

                             decomposed_partitioned_call_callees) {
  auto then_branch = if_op.then_function();
  auto else_branch = if_op.else_function();
  auto grads = AccessedGradients({then_branch, else_branch}, module);
  auto ta_arg_buffer_type = [&](int64_t index) -> Type {
    auto it = stats->find(if_op.getOperand(index + 1));
    if (it == stats->end()) return nullptr;

      list {
        type: DT_INT32
        type: DT_INT32
      }
    }
  }
  attr {
    key: "_tpu_replicate"
    value {
      s: "cluster"
    }
  }
  attr {
    key: "else_branch"
    value {
      func {
        name: "cond_false"
      }
    }
  }
  attr {
    key: "then_branch"
    value {
      func {
        name: "cond_true"
        attr {