// CHECK:       ^bb0
// CHECK:         (%[[CARG:[a-zA-Z0-9_]+]]: tensor<*xf32>)
// CHECK:       tf.OpA
// CHECK:       ^bb0
// CHECK:         (%[[BARG:[a-zA-Z0-9_]+]]: tensor<*xf32>)
// CHECK-NOT:   tf.OpB
// CHECK:       tf.OpC
// CHECK-SAME:    %[[BARG]]
func.func @remove_first_result(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  %0:2 = "tf.WhileRegion"(%arg0, %arg1) ({
  ^bb0(%arg2: tensor<*xf32>, %arg3: tensor<*xf32>):

   // CHECK: OpA
   // CHECK: OpB
   // CHECK: OpC
   "tf.OpA"() : () -> ()
   %status_x, %program_x = "tf._TPUCompileMlir"() { metadata = "X", mlir_module = "..." } : () -> (tensor<!tf_type.string>, tensor<!tf_type.string>)
   "tf.OpB"() : () -> ()
   %status_y, %program_y = "tf._TPUCompileMlir"() { metadata = "Y", mlir_module = "..." } : () -> (tensor<!tf_type.string>, tensor<!tf_type.string>)
   "tf.OpC"() : () -> ()
}

        %b:2 = "tf.opB"(%a) : (tensor<i1>) -> (tensor<i32>, tensor<f32>)
        tf_device.return %b#1, %b#0 : tensor<f32>, tensor<i32>
      }) {device = "CPU:0"} : () -> (tensor<f32>, tensor<i32>)
      %c = "tf.opC"() : () -> tensor<i1>
      tf_executor.yield %a, %launch#0, %launch#1, %c : tensor<i1>, tensor<f32>, tensor<i32>, tensor<i1>
    }
    tf_executor.fetch
  }
  func.return
}

// CHECK-NOT: tf_executor.island

  "tf.opC"(%4) {_xla_compile_device_type = "TPU", _replication_info = "replicate", device = "/device:TPU:0", name = "name", is_stateless = true} : (tensor<i1>) -> ()
  %5 = "tf.TPUReplicatedInput"(%arg5) {index = 4 : i64, is_packed = true} : (tensor<i1>) -> tensor<i1>

  %0:3 = tf_executor.graph {
    %1:4 = tf_executor.island {
      %3 = "tf.opA"(%arg0) : (tensor<i1>) -> tensor<i1>
      %4 = "tf.opB"(%3) : (tensor<i1>) -> tensor<i1>
      %5 = "tf.opC"(%4) : (tensor<i1>) -> tensor<i1>
      tf_executor.yield %3, %5, %4 : tensor<i1>, tensor<i1>, tensor<i1>
    }
    tf_executor.fetch %1#1, %1#0, %1#2 : tensor<i1>, tensor<i1>, tensor<i1>
  }

      %c = "tf.opC"() : () -> tensor<i1>
      tf_executor.yield %a, %launch#0, %launch#1, %c : tensor<i1>, tensor<f32>, tensor<i32>, tensor<i1>
    }
    tf_executor.fetch
  }
  func.return
}

// CHECK:      %[[A:.*]] = "tf.opA"
// CHECK:      %[[B:.*]]:2 = "tf.opB"(%[[A]])
// CHECK-SAME: device = "CPU:0"
// CHECK:      %[[C:.*]] = "tf.opC"
// CHECK-NOT:  "tf_device.launch"

//    "tf_device.parallel_execute"() ({
//      "tf.OpC"
//    }, {
//      "tf.OpD"
//    })
//  }
//
// The above IR will be flattened after `ReplicateToIslandPass` and
// `ParallelExecuteToIslandsPass` as follows:
//
// tf_executor.island wraps "tf.OpA"
// tf_executor.island {_parallel_execution_ids=r0:0} wraps "tf.OpB"
// tf_executor.island {_parallel_execution_ids=r0:0,p0:0} wraps "tf.OpC"

    }
    %3:3 = tf_executor.island() {
      %4:2 = "tf_device.parallel_execute"() ({
        %5 = "tf.opB"(%1#0) : (tensor<i1>) -> tensor<i1>
        tf_device.return %5 : tensor<i1>
      }, {
        %5 = "tf.opC"(%1#0) : (tensor<i1>) -> tensor<i32>
        tf_device.return %5 : tensor<i32>
      }) {} : () -> (tensor<i1>, tensor<i32>)
      tf_executor.yield %4#0, %4#1 : tensor<i1>, tensor<i32>
    }

// possible that live-out values don't dominate the region. For example:
//
// ```
// %0 = "tf.OpA"()
// %1 = "tf.OpB"(%0)
// %2 = "tf.OpC"(%0)
// ```
//
// Assuming `tf.OpA` and `tf.OpC` are clustered together, the region will be
// inserted right after `tf.OpC`. The live-out `%0`, however, is used by
// `tf.OpB`, which won't dominate the region. This function reorders all users

      %3 = "tf.opB"(%2, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
      tf_executor.yield %3 : tensor<i32>
    }
    %4:2 = tf_executor.island {
      %5 = "tf.opC"(%1#0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
      %6 = "tf.opD"(%5, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
      tf_executor.yield %6 : tensor<i32>
    }
    tf_executor.fetch %4#0 : tensor<i32>