XlaHostSendDeviceContext(
      se::Stream* stream, se::DeviceMemoryBase* device_memory_base,
      const xla::Shape& shape,
      tsl::AsyncValueRef<std::unique_ptr<se::Event>>& done_event)
      : stream_(stream),
        device_memory_base_(device_memory_base),
        shape_(shape),
        done_event_(done_event) {}

  // Copies 'cpu_tensor' to `device_memory_base_` with `shape_`.
  // `device_tensor` is unused.

 public:
  XlaHostRecvDeviceContext(
      se::Stream* stream, const se::DeviceMemoryBase& device_memory_base,
      const xla::Shape& shape,
      tsl::AsyncValueRef<std::unique_ptr<se::Event>>& done_event)
      : stream_(stream),
        device_memory_base_(device_memory_base),
        shape_(shape),
        done_event_(done_event) {}

    const Tensor* cpu_tensor, Device* device, Tensor* device_tensor,
    StatusCallback done, bool sync_dst_compute) const {
  auto status = stream_->Memcpy(device_memory_base_, cpu_tensor->data(),
                                device_memory_base_->size());
  if (!status.ok()) {
    done(status);
    return;
  }
  status = stream_->RecordEvent(done_event_.get().get());
  if (!status.ok()) {
    done(status);

  if (!status.ok()) {
    done(status);
    return;
  }

  *cpu_tensor = Tensor(dtype, tensor_shape);

  status = stream_->Memcpy(cpu_tensor->data(), device_memory_base_,
                           device_memory_base_.size());
  if (!status.ok()) {
    done(status);
    return;
  }
  status = stream_->RecordEvent(done_event_.get().get());
  if (!status.ok()) {
    done(status);
    return;

  SP_DeviceMemoryBase device_memory_base{SP_DEVICE_MEMORY_BASE_STRUCT_SIZE};
  // `opaque` field inside SP_DeviceMemoryBase is not const.
  // Therefore, we need to cast away the constness before setting it.
  device_memory_base.opaque = const_cast<void*>(mem->opaque());
  device_memory_base.size = mem->size();
  device_memory_base.payload = mem->payload();
  return device_memory_base;
}

    OpKernelContext* ctx, const std::string& program_key) {
  return
      [ctx, program_key](
          int64_t channel_id, se::Stream* stream, const xla::Shape& shape,
          const se::DeviceMemoryBase& device_memory_base,
          const absl::flat_hash_map<std::string, std::string>& frontend_attrs)
          -> absl::StatusOr<tsl::AsyncValueRef<std::unique_ptr<se::Event>>> {