->stream->parent()
                      ->GetPlatform()
                      ->id();
  } else if (XlaDevice::GetMetadataFromDevice(device_base, &xla_device_metadata)
                 .ok()) {
    // If we are on an XlaDevice, use the underlying XLA platform's allocator
    // directly. We could use the StreamExecutor's allocator which may

      layout_preference_fn, shape_representation_fn}};
}

std::unique_ptr<XlaDevice::Metadata> CreateXlaDeviceMetadata(
    DeviceType compilation_device_type) {
  return std::make_unique<XlaDevice::Metadata>(
      /*device_ordinal=*/0, /*platform=*/nullptr, compilation_device_type,
      GetShapeDeterminationFns(), XlaDevice::PaddedShapeFn(),
      /*use_multiple_streams=*/false);
}

#include "tensorflow/core/common_runtime/local_device.h"
#include "tensorflow/core/framework/device_base.h"

namespace tensorflow {

// tensorflow::PjRtBaseDevice replaces the deprecated tensorflow::XlaDevice.
// This accelerator agnostic device is mainly used to store metadata.
class PjRtBaseDevice : public LocalDevice {
 public:
  // Stores metadata about the PjRtBaseDevice.
  class Metadata {
   public:

#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/mutex.h"

namespace tensorflow {

// The implementation of a Tensor for an XlaDevice. All device tensors are
// actually one of these.
//
// To distinguish between "normal" device tensors and XlaTensors, the raw
// pointer data stored in the TensorBuffer is a tagged pointer.
class XlaTensor {
 public:

  if (debug_options_.dump_graphs) {
    DumpGraphToFile("before_mark_for_compilation", *graph_, flib_def_);
  }

  // Mark clusters for compilation that:
  // * are placed on a device that requires compilation (an XlaDevice),
  // * are explicitly marked for compilation (_XlaCompile=true), or
  // * have more than debug_options_.xla_min_cluster_size elements (applicable
  //   only if compilation is enabled, otherwise there will be no such