%0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_0, then_branch = @no_side_effect_then_0, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)

// CHECK:   %{{.*}} = "tf.If"(%{{.*}}, %{{.*}}, %{{.*}}) <{else_branch = @cond_false, is_stateless = false, then_branch = @cond_true}> : (tensor<1xi1>, tensor<1xf32>, 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>

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

// -----

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

func.func @tf_if_gets_control_node(%arg0: tensor<1xi1>)->() {
 "tf.If"(%arg0) {_lower_using_switch_merge = true, _read_only_resource_inputs = [], device = "", else_branch = @noop, is_stateless = false, output_shapes = [#tf_type.shape<>], then_branch = @noop} : (tensor<1xi1>) -> ()
 func.return
}
// CHECK-NEXT: %[[CONTROL:.*]] = tfl.control_node controls "tf.If"
// CHECK-NEXT: return

// CHECK-LABEL: @tfl_if_gets_control_node

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

# 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]]"(
# CHECK-SAME: tf._original_func_name


def Test():

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

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

func.func @cond_true(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {

  %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>)