int XlaDevice::Metadata::device_ordinal() const { return device_ordinal_; }

se::Platform* XlaDevice::Metadata::platform() const { return platform_; }

xla::LocalClient* XlaDevice::Metadata::client() const {
  auto client = xla::ClientLibrary::GetOrCreateLocalClient(platform_);
  return client.value();
}

const DeviceType& XlaDevice::Metadata::jit_device_type() const {

==============================================================================*/

// The XlaDevice executes a TensorFlow graph using the XLA linear algebra
// runtime.
//
// Operators assigned to an XlaDevice are compiled into XLA computations.
// Tensors on an XlaDevice are thin wrappers around XLA ScopedShapedBuffers.
//
// XlaDevice is instantiated separately for each XLA backend (e.g., CPU or GPU),

TEST_F(XlaPlatformInfoTest, BuildXlaDeviceCompilerXlaDeviceMetadata) {
  device_setup_.AddDevicesAndSetUp({DEVICE_XLA_GPU});

  Device* device = device_setup_.GetDevice(DEVICE_XLA_GPU);
  const XlaDevice::Metadata* metadata = nullptr;
  TF_CHECK_OK(XlaDevice::GetMetadataFromDevice(device, &metadata));
  XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(device);

  TF_ASSERT_OK_AND_ASSIGN(
      DeviceType compilation_device_type,

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Registers the XLA_GPU device, which is an XlaDevice instantiation that runs
// operators using XLA via the XLA "CUDA" or "ROCM" (GPU) backend.

#include <array>
#include <set>

#include "absl/memory/memory.h"
#include "absl/strings/numbers.h"

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Registers the XLA_CPU device, which is an XlaDevice instantiation that runs
// operators using XLA via the XLA "Host" (CPU) backend.

#include <array>

#include "absl/memory/memory.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/flags.h"

 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,
      std::shared_ptr<se::DeviceMemoryAllocator> device_allocator)
      : device_type_(device_type),
        platform_id_(platform_id),

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

#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:

  VLOG(1) << "Creating " << device_count << " TPU devices";
  for (int i = 0; i < device_count; ++i) {
    TF_RETURN_IF_ERROR(tpu::TpuNodeContext::Initialize(i));

    XlaDevice::Options options;
    options.platform = platform;
    options.device_name_prefix = name_prefix;
    options.device_name = DEVICE_TPU_NODE;
    options.device_ordinal = i;

  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