class XlaDeviceContext : public DeviceContext {
 public:
  explicit XlaDeviceContext(
      std::shared_ptr<se::Stream> compute_stream,
      std::shared_ptr<se::Stream> host_to_device_stream,
      std::shared_ptr<se::Stream> device_to_host_stream,
      std::vector<std::shared_ptr<se::Stream>> device_to_device_streams,
      xla::LocalClient* client,

    return false;
  }
  return stream_executor_->ClearAllocatorStats();
}

XlaDeviceContext::XlaDeviceContext(
    std::shared_ptr<se::Stream> compute_stream,
    std::shared_ptr<se::Stream> host_to_device_stream,
    std::shared_ptr<se::Stream> device_to_host_stream,
    std::vector<std::shared_ptr<se::Stream>> device_to_device_streams,
    xla::LocalClient* client,
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,

      colocation_constraints_(),
      disable_shape_inference_(disable_shape_inference) {}

Scope::Impl::Impl(const std::shared_ptr<Graph>& graph,
                  const std::shared_ptr<Status>& status,
                  const std::shared_ptr<NameMap>& name_map,
                  const std::shared_ptr<ShapeRefiner>& refiner)
    : graph_(graph),
      status_(status),
      name_map_(name_map),
      refiner_(refiner),

#include "tensorflow/core/lib/core/status.h"
#include "tsl/platform/errors.h"

namespace tensorflow {

bool XlaKernelCreator::CanCreateKernel(
    const FunctionLibraryRuntime& flr,
    const std::shared_ptr<const NodeProperties>& props) const {
  return CanCreateXlaKernel(props->node_def) &&
         !XlaOpRegistry::IsCompilationDevice(flr.device()->device_type());
}

static Status CreateXlaKernel(FunctionLibraryRuntime* flr,

    TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "stream", &stream_,
                                            &need_new_device_context));
  }

  std::shared_ptr<se::Stream> host_to_device_stream;
  std::shared_ptr<se::Stream> device_to_host_stream;
  std::vector<std::shared_ptr<se::Stream>> device_to_device_streams;
  if (use_multiple_streams_) {
    TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "host_to_device_stream",

  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),
        xla_device_metadata_(xla_device_metadata),
        pjrt_device_metadata_(pjrt_device_metadata),

class TaggedValueHash;
using List = std::vector<TaggedValue>;
using ListPtr = std::shared_ptr<List>;
using Dict =
    absl::flat_hash_map<TaggedValue, TaggedValue, TaggedValueHash<TaggedValue>>;
using DictPtr = std::shared_ptr<Dict>;
using TuplePtr = std::shared_ptr<Tuple>;
using Func =
    std::function<absl::StatusOr<TaggedValue>(TaggedValue, TaggedValue)>;

  // If use_multiple_streams_, host to device transfers are performed using this
  // stream.
  std::shared_ptr<se::Stream> host_to_device_stream_ TF_GUARDED_BY(mu_);
  // If use_multiple_streams_, transfers between different devices are performed
  // using these streams.
  std::vector<std::shared_ptr<se::Stream>> device_to_device_streams_
      TF_GUARDED_BY(mu_);

  // See comments in options.

  // streams.
  // It is legal to reset the definition event of a tensor when overwriting the
  // tensor's value (at which point, it is effectively a new tensor once again.)
  void ResetDefinitionEvent(std::shared_ptr<se::Event> event,
                            se::Stream* stream);

  // Refresh the status of streams_defined_on_. Return the first not-OK stream's
  // status or OK.
  Status RefreshStatusOfStreams();

  DeviceType compilation_device_type = DeviceType(DEVICE_GPU_XLA_JIT);

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