#include "tensorflow/compiler/jit/kernels/xla_ops.h"
#include "tensorflow/compiler/jit/xla_compile_on_demand_op.h"
#include "tensorflow/compiler/jit/xla_device.h"
#include "tensorflow/compiler/jit/xla_device_ops.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/stream_executor/platform_manager.h"
#include "tensorflow/core/common_runtime/device_factory.h"

#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/kernels/xla_ops.h"
#include "tensorflow/compiler/jit/xla_device.h"
#include "tensorflow/compiler/jit/xla_device_ops.h"
#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/stream_executor/gpu/gpu_init.h"

#include <memory>
#include <optional>
#include <string>

#include "tensorflow/compiler/jit/device_compiler.h"
#include "tensorflow/compiler/jit/pjrt_base_device.h"
#include "tensorflow/compiler/jit/xla_device.h"
#include "xla/stream_executor/integrations/tf_allocator_adapter.h"
#include "tensorflow/core/framework/op_kernel.h"

namespace tensorflow {

// Holds some information about the platform on which an

#ifndef TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_
#define TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_

#include <atomic>

#include "tensorflow/compiler/jit/device_compiler.h"
#include "tensorflow/compiler/jit/xla_device.h"
#include "tensorflow/compiler/jit/xla_launch_util.h"
#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "xla/stream_executor/integrations/tf_allocator_adapter.h"
#include "tensorflow/core/framework/allocator.h"

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