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

  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =
      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
          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));

#include "tensorflow/compiler/jit/xla_host_recv_device_context.h"

#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/compiler/tf2xla/type_util.h"

namespace tensorflow {

void XlaHostRecvDeviceContext::CopyDeviceTensorToCPU(
    const Tensor* device_tensor, StringPiece tensor_name, Device* device,
    Tensor* cpu_tensor, StatusCallback done) {