VLOG(1) << "Created PJRT base device " << options.compilation_device_name
          << " device_name: " << name();
}

/*static*/ absl::StatusOr<const PjRtBaseDevice::Metadata*>
PjRtBaseDevice::GetMetadataFromDevice(DeviceBase* device) {
  PjRtBaseDevice* pjrt_device =
      dynamic_cast<PjRtBaseDevice*>(device->UnderlyingDevice());
  if (pjrt_device == nullptr) {
    return errors::Internal(

#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:
    Metadata(const DeviceType& jit_device_type,

// abstraction for normal, XLA devices and devices inheriting from
// PjRtBaseDevice.
class XlaPlatformInfo {
 public:
  XlaPlatformInfo() : device_type_("") {}
  XlaPlatformInfo(XlaPlatformInfo&&) = default;
  explicit XlaPlatformInfo(
      const DeviceType device_type, se::Platform::Id platform_id,
      const XlaDevice::Metadata* xla_device_metadata,
      const PjRtBaseDevice::Metadata* pjrt_device_metadata,

      GetShapeDeterminationFns(), XlaDevice::PaddedShapeFn(),
      /*use_multiple_streams=*/false);
}

std::unique_ptr<PjRtBaseDevice::Metadata> CreatePjRtDeviceMetadata(
    DeviceType compilation_device_type) {
  return std::make_unique<PjRtBaseDevice::Metadata>(compilation_device_type,
                                                    GetShapeDeterminationFns());
}

}

XlaPlatformInfo XlaPlatformInfoFromDevice(DeviceBase* device_base) {
  se::Platform::Id platform_id = nullptr;
  const XlaDevice::Metadata* xla_device_metadata = nullptr;
  const PjRtBaseDevice::Metadata* pjrt_device_metadata = nullptr;
  std::shared_ptr<se::DeviceMemoryAllocator> custom_allocator;

  const std::string& device_type = device_base->device_type();
  if (device_type == DEVICE_CPU) {