if (device_type == DEVICE_CPU) {
    platform_id = se::host::kHostPlatformId;
  } else if (device_type == DEVICE_GPU) {
    auto device = static_cast<Device*>(device_base);
    platform_id = device->tensorflow_accelerator_device_info()
                      ->stream->parent()
                      ->GetPlatform()
                      ->id();

                        XlaRunOp);
REGISTER_KERNEL_BUILDER(
    Name("_XlaRun").Device(DEVICE_DEFAULT).HostMemory("key"), XlaRunOp);

REGISTER_KERNEL_BUILDER(Name("_XlaMerge").Device(DEVICE_CPU), XlaMergeOp);
REGISTER_KERNEL_BUILDER(Name("_XlaMerge").Device(DEVICE_GPU), XlaMergeOp);
REGISTER_KERNEL_BUILDER(Name("_XlaMerge").Device(DEVICE_DEFAULT), XlaMergeOp);

                                              Device* device,
                                              Tensor* output_tensor,
                                              StatusCallback done) const {
  done(errors::Unimplemented("XLA->XLA same-device copies not implemented."));
}

void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
                                             Device* device,

                      task)
            .str());
    for (int device = 0; device < num_devices_per_task; ++device)
      devices.push_back(
          llvm::formatv("/job:worker/replica:0/task:{0}/device:TPU:{1}", task,
                        device)
              .str());
  }

  return devices;
}

TEST(TPURewriteDeviceUtilTest,

    args: a list of fallback tensors.
    device: the name of the tensorflow device that is associated with the
      input fallback tensors.

    results: the converted corert tensorhandles.
  }];

  let arguments = (ins
    Variadic<TFTensorType>:$args,
    StrAttr:$device
  );

  let results = (outs
    Variadic<CoreRT_TensorHandleType>:$results
  );

  // Creates a device that represents a TPU distributed system.
  const DeviceAttributes attrs = Device::BuildDeviceAttributes(
      absl::StrCat(name_prefix, "/device:", DEVICE_TPU_SYSTEM, ":", 0),
      DeviceType(DEVICE_TPU_SYSTEM), Bytes(memory_limit), DeviceLocality(),
      absl::StrCat("device: ", DEVICE_TPU_SYSTEM, " device"));
  devices->push_back(std::make_unique<VirtualDevice>(options.env, attrs));

  %0 = "tf.Const"() {__op_key = 0: i32, device = "/device:CPU:0", value = dense<0> : tensor<i64>} : () -> tensor<i64>
  %1 = "tf.Const"() {__op_key = 1: i32, device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
  %2 = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>

                            void* device_info) = nullptr;

  // Pins the op to `device` based on inputs to `op`. Returns true
  // signifying to pin to the current custom device. Returns false
  // to pin to the physical device.
  //
  // This function is guaranteed to be called only when all of the custom-device
  // inputs are on this device.
  bool (*shall_pin_to_this_device)(const TFE_Op* op, TF_Status* s) = nullptr;

    // `enabled_for_gpu_` below.
    bool enabled_for_all_;

    // If true, enable Device API (PjRt) for TF GPU device. This is a helper
    // flag so that individual tests can turn on PjRt for GPU specifically.
    // Once the rollout to GPU is complete, this flag can be deprecated.
    bool enabled_for_gpu_;

   private:
    // Devices for which using Device API (PjRt) is allowed in the XlaLaunch op.

  let description = [{
    The ExecuteOp executes an operation on the specified device using a custom device.
  }];

  let arguments = (ins
    TFDeviceType:$device,
    Variadic<TFTensorType>:$args,
    StrAttr:$node_def,
    I32Attr:$op_key
  );

  let results = (outs
    Variadic<TFTensorType>:$results
  );