namespace tensorflow {

void XlaHostRecvDeviceContext::CopyDeviceTensorToCPU(
    const Tensor* device_tensor, StringPiece tensor_name, Device* device,
    Tensor* cpu_tensor, StatusCallback done) {
  DataType dtype = EncodePrimitiveTypeAsDataType(shape_.element_type()).value();
  TensorShape tensor_shape;
  Status status = XLAShapeToTensorShape(shape_, &tensor_shape);
  if (!status.ok()) {
    done(status);
    return;
  }

          const xla::Shape& xla_shape) -> absl::StatusOr<xla::Shape> {
    TensorShape shape;
    TF_RETURN_IF_ERROR(XLAShapeToTensorShape(xla_shape, &shape));
    TF_ASSIGN_OR_RETURN(DataType dtype, EncodePrimitiveTypeAsDataType(
                                            xla_shape.element_type()));
    auto layout_preference = shape_determination_fns.layout_preference_fn(
        shape, dtype, std::nullopt);

      const xla::Literal* const xla_literal = arguments[i];
      tensorflow::TensorShape shape;
      TF_ASSIGN_OR_RETURN(auto data_type,
                          tensorflow::EncodePrimitiveTypeAsDataType(
                              xla_literal->shape().element_type()));
      TF_RETURN_IF_ERROR(
          tensorflow::XLAShapeToTensorShape(xla_literal->shape(), &shape));