TF_ASSIGN_OR_RETURN(auto device_context, GetDeviceContextDefault());
  device_context->Ref();
  *out_context = device_context;
  return absl::OkStatus();
}

// Warn about XLA_CPU/XLA_GPU exactly once.
static void ShowXlaDeviceDeprecationWarning(
    absl::string_view compilation_device_name) {
  static absl::once_flag once;
  if (absl::StrContains(compilation_device_name, "CPU") ||

    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,

  // TODO(cheshire): Avoid duplication with framework/op_kernel.h
  DeviceContext* device_context = nullptr;
  if (device != nullptr) {
    TF_RETURN_IF_ERROR(device->TryGetDeviceContext(&device_context));
    bool using_default_context = false;
    auto cleanup = absl::MakeCleanup([&] {
      if (device_context != nullptr && !using_default_context) {
        device_context->Unref();
      }
    });
    if (device_context == nullptr) {

  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,

            tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
                std::move(event));

        Rendezvous::Args args;
        // Rendezvous::Args owns the device context pointer.
        args.device_context = new XlaHostRecvDeviceContext(
            stream, device_memory_base, shape, done_event);

        Tensor host_tensor;
        TF_RETURN_IF_ERROR(

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

        device_type,
        xla::GetTfrtCpuClient(/*asynchronous=*/true, /*cpu_device_count=*/1)
            .value()));

    // device_context_ should be a PjRtDeviceContext.
    TF_CHECK_OK(device_->TryGetDeviceContext(&device_context_));

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