// Note: we currently assume the PjRtBuffer is a PjRtStreamExecutorBuffer.
  bool use_pjrt_tensor_buffer_;
};

void PjRtDeviceToDeviceCopy(DeviceContext* send_dev_context,
                            DeviceContext* recv_dev_context, Device* src,
                            Device* dst, AllocatorAttributes src_alloc_attr,
                            AllocatorAttributes dst_alloc_attr,

  Status Sync() override;

  Status TryGetDeviceContext(DeviceContext** out_context) override
      TF_LOCKS_EXCLUDED(mu_);

  Status MakeTensorFromProto(const TensorProto& tensor_proto,
                             const AllocatorAttributes alloc_attrs,
                             Tensor* tensor) override TF_LOCKS_EXCLUDED(mu_);

  Status MakeTensorFromProto(DeviceContext* device_context,

    device_allocator_ = device_->GetAllocator(device_alloc_attr);

    tensorflow::DeviceContext* device_context;
    auto status = device_->TryGetDeviceContext(&device_context);
    TF_EXPECT_OK(status);
    device_context_.reset(device_context);
  }

  std::unique_ptr<Device> device_;
  tensorflow::core::RefCountPtr<DeviceContext> device_context_;
  tensorflow::Allocator* host_allocator_;

#include "xla/stream_executor/stream.h"
#include "tensorflow/core/framework/device_base.h"
#include "tfrt/concurrency/async_value_ref.h"  // from @tf_runtime

namespace tensorflow {

// XlaHostSendDeviceContext is a DeviceContext that is intended to be
// used to transfer from host->device using Rendezvous. It transfers the
// content of `device_memory_base` with `shape` using `stream`. Only

#include "xla/stream_executor/stream.h"
#include "tensorflow/core/framework/device_base.h"
#include "tfrt/concurrency/async_value_ref.h"  // from @tf_runtime

namespace tensorflow {

// XlaHostRecvDeviceContext is a DeviceContext that is intended to be
// used to transfer from device->host using Rendezvous. It transfers the
// content of `device_memory_base` with `shape` using `stream`. Only

  }

  return device_contexts_;
}

absl::StatusOr<DeviceContext*> XlaDevice::GetDeviceContextWithIndex(int index) {
  mutex_lock lock(mu_);
  TF_ASSIGN_OR_RETURN(auto device_contexts, GetDeviceContextLocked());
  return device_contexts.at(index);
}

absl::StatusOr<DeviceContext*> XlaDevice::GetDeviceContextDefault() {
  return GetDeviceContextWithIndex(0);
}

    SetDevice(device_type, std::move(device));

    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_fns{
        UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
    device_context_ = core::RefCountPtr<DeviceContext>(
        new PjRtDeviceContext(shape_fns, /*use_pjrt_tensor_buffer=*/true));

    // Get the host allocator.
    AllocatorAttributes host_alloc_attr;
    host_alloc_attr.set_on_host(true);

      tensorflow::AsyncValueTensor::FromTensor(output_tensor);
  result_tensor->GetBuffer()->GetReadyFuture().OnReady(std::move(done));
}

void PjRtDeviceToDeviceCopy(DeviceContext* send_dev_context,
                            DeviceContext* recv_dev_context, Device* src,
                            Device* dst, AllocatorAttributes src_alloc_attr,
                            AllocatorAttributes dst_alloc_attr,

 private:
  // The stream executor of the device.
  se::StreamExecutor* stream_executor_;
};

// Helper class for managing data transfers between host and XLA devices.
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,

// strictly hierarchical, that is, transfers between devices on different hosts
// can only take place using the host network.
void TpuDeviceToDeviceCopy(DeviceContext* src_dev_context,
                           DeviceContext* dst_dev_context, Device* src,
                           Device* dst, AllocatorAttributes src_allocator_attrs,
                           AllocatorAttributes dst_allocator_attrs,