namespace mlir {
namespace TF {

namespace {

const char kDeviceAttr[] = "device";
const char kDeviceGpu[] = "GPU";

std::optional<std::string> GetOpDevice(mlir::Operation *op) {
  mlir::StringAttr device = op->getAttrOfType<mlir::StringAttr>(kDeviceAttr);
  if (!device || device.getValue().empty()) {
    return std::nullopt;
  }
  tensorflow::DeviceNameUtils::ParsedName parsed_name;

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

                             StringPiece tensor_name, Device* device,
                             Tensor* cpu_tensor, StatusCallback done) override;

  void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
                              Tensor* output_tensor,
                              StatusCallback done) const override {
    done(errors::Internal("device->device copy not implemented."));
  }

 private:

                                   TF_Status* status) {
    params->device->hardware_name = hardware_name;
    params->device->device_vendor = vendor;
    params->device->pci_bus_id = pci_bus_id;
  };

  device_fns_.get_numa_node = [](const SP_Device* device) { return 123; };
  device_fns_.get_memory_bandwidth = [](const SP_Device* device) -> int64_t {
    return 54;
  };

    Device* device = dynamic_cast<Device*>(ctx->device());
    if (device == nullptr) {
      return errors::Internal("DeviceBase was not a Device.");
    }
    ctx->op_device_context()->CopyCPUTensorToDevice(
        &const_tensor, device, output_tensor,
        [&](Status status) { TF_CHECK_OK(status); });

    if (device->device_type() == DEVICE_GPU) {

  }

  void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
                              Tensor* output_tensor,
                              StatusCallback done) const override {
    done(errors::Internal("device->device copy not implemented."));
  }

 private:
  se::Stream* stream_;                        // Not owned.

                        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,

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