// XLA device.
  const XlaDevice::Metadata* xla_device_metadata() const {
    return xla_device_metadata_;
  }
  bool is_on_xla_device() const { return xla_device_metadata() != nullptr; }

  const PjRtBaseDevice::Metadata* pjrt_device_metadata() const {
    return pjrt_device_metadata_;
  }

 private:
  DeviceType device_type_;
  se::Platform::Id platform_id_;

  // xla_device_metadata_ lives in the tensorflow::DeviceBase in which the

  DeviceType device_type = platform_info.device_type();

  if (platform_info.xla_device_metadata()) {
    VLOG(2) << "Building PjRtDeviceCompiler using "
               "platform_info.xla_device_metadata().";

    *compilation_device_type =
        platform_info.xla_device_metadata()->jit_device_type();
    TF_ASSIGN_OR_RETURN(*pjrt_client, GetOrCreatePjRtClient(device_type));

  se::Platform::Id platform_id = nullptr;
  auto xla_device_metadata = CreateXlaDeviceMetadata(compilation_device_type);
  std::shared_ptr<se::DeviceMemoryAllocator> custom_allocator;
  XlaPlatformInfo platform_info(
      compilation_device_type, platform_id, xla_device_metadata.get(),
      /*pjrt_device_metadata=*/nullptr, custom_allocator);

      (platform_info.platform_id() == se::host::kHostPlatformId);
  options.device_allocator = GetAllocator(device, stream, platform_info);
  if (platform_info.xla_device_metadata()) {
    options.shape_determination_fns =
        platform_info.xla_device_metadata()->default_shape_determination_fns();
  }
  // If reference variables are not present in the graph, we can safely alias
  // passthrough parameters without performing a copy.

  // XlaDeviceCompiler.
  Device* device = nullptr;
  XlaPlatformInfo platform_info(DeviceType(DEVICE_TPU), /*platform_id=*/nullptr,
                                /*xla_device_metadata=*/nullptr,
                                /*pjrt_device_metadata=*/nullptr,
                                /*device_allocator=*/nullptr);

  XlaDeviceCompiler* xla_device_compiler = nullptr;

  XlaComputationLaunchContext launch_context(
      client, allocator, client->default_device_ordinal(),
      /*allocate_xla_tensors=*/platform_info_.xla_device_metadata() != nullptr,
      platform_info_.xla_device_metadata()
          ? platform_info_.xla_device_metadata()->UseMultipleStreams()
          : false);

  std::map<int, const Tensor*> snapshot_ptrs;
  for (auto& p : variable_args) {

  if (platform_name == DEVICE_CPU) {
    platform_id = se::host::kHostPlatformId;
  }
  XlaPlatformInfo platform_info(DeviceType(platform_name), platform_id,
                                /*xla_device_metadata=*/nullptr,
                                /*pjrt_device_metadata=*/nullptr,
                                /*device_allocator=*/nullptr);
  DeviceType compilation_device_type = platform_info.device_type();