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_;
  tensorflow::Allocator* device_allocator_;
};

          std::move(se_event));
  XlaHostRecvDeviceContext* device_context =
      new XlaHostRecvDeviceContext(stream.get(), gpu_dst, shape, done_event);
  TF_ASSERT_OK(device_context->CopyDeviceTensorToCPUSync(
      &device_tensor, "", device_.get(), &dest_cpu_tensor));

  tensorflow::test::ExpectClose(origin_cpu_tensor, dest_cpu_tensor);
  device_context->Unref();
}

//  done_event->Init();
//
//  XlaHostSendDeviceContext device_context(&stream, &gpu_dst,
//    shape, done_event);
//  device_context.CopyCPUTensorToDeviceSync(
//    &cpu_tensor, &device, &device_tensor);

class XlaHostSendDeviceContext : public DeviceContext {
 public:
  XlaHostSendDeviceContext(
      se::Stream* stream, se::DeviceMemoryBase* device_memory_base,

//  done_event->Init();
//  Tensor dest_cpu_tensor;
//
//  XlaHostRecvDeviceContext device_context(&stream, gpu_dst,
//    shape, done_event);
//  device_context.CopyDeviceTensorToCPUSync(
//    &device_tensor, "", &device, &dest_cpu_tensor);

class XlaHostRecvDeviceContext : public DeviceContext {
 public:
  XlaHostRecvDeviceContext(
      se::Stream* stream, const se::DeviceMemoryBase& device_memory_base,

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

 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,