}

// -----

func.func private @then_branch() -> ()
func.func private @else_branch() -> ()

// Test tf.If device is preserved.
// CHECK-LABEL: func @testIfDevice
func.func @testIfDevice(%arg0: tensor<i1>) {
  "tf.If"(%arg0) {then_branch = @then_branch, else_branch = @else_branch, is_stateless = false, device = "/device:CPU:0"} : (tensor<i1>) -> ()

  // CHECK: "tf.IfRegion"

  // CHECK-NEXT: tf.If
  // CHECK-SAME: <{else_branch = @multiple_uses_merged_if_0_0_else, is_stateless = true, then_branch = @multiple_uses_merged_if_0_0_then}>
  %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>)

  %0 = "tf.Const"() {value = dense<10> : tensor<i32> } : () -> tensor<i32>
  // CHECK: [[res:%.*]] = "tf.If"([[cond]], [[x]])
  // CHECK-SAME: {else_branch = @else_removed_const_args_0, is_stateless = false, then_branch = @then_removed_const_args_0}
  // CHECK-NEXT: return [[res]]
  %1 = "tf.If"(%cond, %x, %0) {else_branch = @else, then_branch = @then, is_stateless = false} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %1 : tensor<i32>
}

  %unused = "tf.If"(%cond, %cond) {device = "/job:localhost/replica:0/task:0/device:CPU:0", else_branch = @else_branch, is_stateless = false, then_branch = @then_branch} : (tensor<i1>, tensor<i1>) -> tensor<i1>

# CHECK-DAG: constant dense<[5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00]>
# CHECK:      "tf.If"{{.+}}else_branch = @cond_false_10{{.+}}is_stateless = true{{.+}}then_branch = @cond_true_10
# CHECK:      "tf.If"{{.+}}else_branch = @cond_false0{{.+}}is_stateless = false{{.+}}then_branch = @cond_true0
# CHECK:      func private @cond_false_10
# CHECK-NEXT: tfl.div

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

  %0 = "tf.While"(%arg) { cond = @no_side_effect_cond, body = @no_side_effect_body, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
  // CHECK: tf.If
  // CHECK-SAME: is_stateless = true
  %1 = "tf.If"(%cond, %arg) {else_branch = @no_side_effect_body, then_branch = @no_side_effect_body2, is_stateless = false} : (tensor<i1>, tensor<i32>) -> tensor<i32>
  func.return %0, %1 : tensor<i32>, tensor<i32>
}

// CHECK-LABEL: func @nested_set_stateless

  %3 = "tf.If"(%2, %0, %1) {else_branch = @sub, then_branch = @addormul, is_stateless = true} : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
  func.return %3 : tensor<f32>
}

func.func private @addormul(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32>  {
  %0 = arith.constant dense<false> : tensor<i1>

# CHECK:  func {{.*}} tf_saved_model.exported_names = ["key_1"]
# CHECK: "tf.If"
# CHECK-SAME: else_branch = @[[else:[a-zA-Z_0-9]+]]
# CHECK-SAME: then_branch = @[[then:[a-zA-Z_0-9]+]]

# CHECK:  func {{.*}} tf_saved_model.exported_names = ["key_2"]
# CHECK: "tf.If"
# CHECK-SAME: else_branch = @[[else]]
# CHECK-SAME: then_branch = @[[then]]

# CHECK: func private @[[else]](
# CHECK: func private @[[then]](

func.func @basic(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
  // 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>
}

// -----