};

  devices.ForEach([&](jit::DeviceId device) {
    if (device_info_cache.IsGpu(device)) {
      if (maybe_gpu_device) {
        multiple_gpu_devices = is_multiple_devices(device, &maybe_gpu_device);
        if (multiple_gpu_devices) return false;
      } else {
        maybe_gpu_device = device;
      }
    } else if (device_info_cache.IsCpu(device)) {
      if (maybe_cpu_device) {

    %1 = "tf.Rank"(%0) {_tpu_replicate = "cluster", device = ""} : (tensor<*xi1>) -> tensor<*xi32>
    %2 = "tf.Range"(%cst_0, %1, %cst_1) {_tpu_replicate = "cluster", _xla_outside_compilation = "0", device = ""} : (tensor<i32>, tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
    %3 = "tf.All"(%0, %2) {_tpu_replicate = "cluster", _xla_outside_compilation = "0", device = "", keep_dims = false} : (tensor<*xi1>, tensor<*xi32>) -> tensor<*xi1>

     %3 = builtin.unrealized_conversion_cast to tensor<i32>
     // CHECK: TPUCopyWithDynamicShape{{.*}}device = "foobar"
     %4, %5 = "tf.TPUCopyWithDynamicShape"(%0, %1, %2, %3) {operandSegmentSizes = array<i32: 2, 2>} : (tensor<2048xi32>, tensor<2048xi32>, tensor<i32>, tensor<i32>) -> (tensor<2048xi32>, tensor<2048xi32>)
    "tf.TPUExecute"(%4, %arg0) {device = "foobar"} : (tensor<2048xi32>, tensor<!tf_type.string>) -> ()
    return
}

// -----

   // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
   // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
  %1 = "tfl.add"(%arg0, %0) {fused_activation_function = "RELU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
   // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"

      // CHECK-SAME: device = "/device:test_device:0"
      %3 = "tf_device.launch_func"(%2) {device = "/device:test_device:0", func = @_func} : (tensor<?xf32>) -> tensor<?xf32>

      // CHECK: %[[CALL_OUTPUT_1:[0-9]*]] = "tf.PartitionedCall"(%[[CALL_OUTPUT_0]])
      // CHECK-SAME: f = @_func
      // CHECK-SAME: device = "/device:test_device:1"

// RUN: tf-opt -split-input-file -verify-diagnostics %s -tf-replica-id-to-device-ordinal | FileCheck %s


// Tests device ordinal is set correctly for multiple devices.
// CHECK-LABEL: func @device_ordinal_attr_added_multiple_devices
module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0", "/job:worker/replica:0/task:0/device:TPU:1"]} {

// RUN: tf-opt -tf-simple-device-assignment='default-device=gpu' %s | FileCheck %s

// CHECK-LABEL: func @device_test
func.func @device_test(%arg0: tensor<3x1xf32>) -> (tensor<3x3xf32>) {

  // CHECK: device = "gpu"
  %0 = "tf.Const"() {value = dense<[[1.0, 2.0, 3.0]]> : tensor<1x3xf32>} : () -> tensor<1x3xf32>
  // CHECK: device = "gpu"

def TFXLADeviceSpecificTransforms : Pass<"tfxla-device-specific-transforms",
                                            "mlir::func::FuncOp"> {
  let summary = "Transforms ops that require device context into device independent TF Ops.";

  let description = [{"Transforms device specific ops into device independent"
                    "ops."}];

  let options = [
    Option<"device_type_", "device-type", "std::string",

  // CHECK-NEXT: [[ch2:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key(1) cost({{.*}}) device("/device:CPU:0") "tf.AssignVariableOp"
  // CHECK-NEXT: [[out_ch:%.*]] = tfrt.merge.chains [[ch1]], [[ch2]]
  // CHECK-NEXT: tfrt.return [[out_ch]]

// CHECK-LABEL: func @device_test
func.func @device_test(%arg0: tensor<3x1xf32>) -> (tensor<3x3xf32>) attributes {tf.device = "xpu"} {

  // CHECK: device = "xpu"
  %0 = "tf.Const"() {value = dense<[[1.0, 2.0, 3.0]]> : tensor<1x3xf32>} : () -> tensor<1x3xf32>
  // CHECK: device = "xpu"