func.return %1 : tensor<*xf32>
}

// -----

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

// 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) ({

func.func @testIf1Result(tensor<i1>, tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> {
^bb0(%arg0: tensor<i1>, %arg1: tensor<*xf32>, %arg2: tensor<*xf32>):
  %1 = "tf.If"(%arg0, %arg1, %arg2) {
    then_branch = @testIf1Then, else_branch = @testIf1Else, is_stateless = false
  } : (tensor<i1>, tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>

// CHECK:   [[TOBOOL:%.+]] = "tf.ToBool"(%arg0) : (tensor<i1>) -> tensor<i1>

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

  // CHECK: [[cond:%.*]] = "tf.Const"()
  %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>

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

        : () -> !tf_res
      "tf.If"(%arg1, %id0, %var_handle) {
          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,

                         llvm::StringMap<PartitionedCallTensorArrayOpsInfo>*
                             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));

    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(