%arg1: tensor<*x!tf_type.resource<tensor<64xf32>>> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"},
  %arg2: tensor<*x!tf_type.resource<tensor<16xf32>>> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) {
  // CHECK-NEXT: %[[ID_0:.*]] = "tf.IdentityN"(%[[ARG_0]])

  LoggingDevice* device = new LoggingDevice;
  device->arrived_flag = arrived_flag;
  device->executed_flag = executed_flag;
  device->device_name = name;
  device->underlying_device = "/job:localhost/replica:0/task:0/device:CPU:0";
  device->strict_scope_placement = strict_scope_placement;
  TFE_RegisterCustomDevice(context, custom_device, name, device, status);
}

    %3:3 = tf_executor.Merge %2#0, %0#0 : tensor<*xi32> {device = "", N = 2, T = "tfdtype$DT_INT32"} loc("while/Merge")
    %4:2 = tf_executor.island(%3#2) wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_INT32", value = dense<10> : tensor<i32>} : () -> tensor<i32> loc("while/Less/y")

class DeviceInfoCache {
 public:
  bool IsGpu(DeviceId device) const { return is_gpu_[device.id()]; }
  bool IsCpu(DeviceId device) const { return is_cpu_[device.id()]; }

  absl::string_view GetNameFor(DeviceId device) const {
    return names_[device.id()];
  }

  absl::StatusOr<DeviceId> GetIdFor(absl::string_view name);

    Op<TFRT_GPU_Dialect, mnemonic, traits> {
}

// TODO(b/260267885): We may add a device argument when we want to support
// GPU MIG.
def TransferToDeviceOp: Gpu_Op<"transfer_to_device"> {
  let summary = "Transfer a CPU tensor to device.";

  let description = [{
    Transfer a CPU tensor to device.

    Example:
      %device_tensor = gpurt.transfer_to_device %cpu_tensor
  }];

  };

  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"

    %0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
    %1 = tfl.mul %0, %arg2 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
    func.return %1 : tensor<1xf32>