// Copies 'cpu_tensor' to `device_memory_base_` with `shape_`.
  // `device_tensor` is unused.
  void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,
                             Tensor* device_tensor, StatusCallback done,
                             bool sync_dst_compute) const override;

  void CopyDeviceTensorToCPU(const Tensor* device_tensor,
                             StringPiece tensor_name, Device* device,

TEST_F(XlaHostSendRecvDeviceContextTest, CopyDeviceTensorToCPU) {
  SetDevice("GPU");
  Tensor origin_cpu_tensor(host_allocator_, DT_FLOAT, TensorShape({2, 2}));
  test::FillValues<float>(&origin_cpu_tensor, {1.2, 2.3, 3.4, 4.5});
  Tensor device_tensor(device_allocator_, DT_FLOAT, TensorShape({2, 2}));
  Tensor dest_cpu_tensor(host_allocator_, DT_FLOAT, TensorShape({2, 2}));

  stream_executor::Platform* platform =

                             Tensor* device_tensor, StatusCallback done,
                             bool sync_dst_compute) const override {
    done(errors::Internal("host->device copy not implemented."));
  }

  // Copies `device_memory_base_` with `shape_` into `cpu_tensor`.
  // `device_tensor` is unused.
  void CopyDeviceTensorToCPU(const Tensor* device_tensor,

  TF_ASSERT_OK(device_context_->CopyCPUTensorToDeviceSync(
      &origin_cpu_tensor, device_.get(), &device_tensor));
  TF_ASSERT_OK(device_context_->CopyDeviceTensorToCPUSync(
      &device_tensor, "", device_.get(), &dest_cpu_tensor));
  LOG(INFO) << "H2D - D2H roundtrip completes. tensor: "
            << dest_cpu_tensor.DebugString(4);

  tensorflow::test::ExpectClose(origin_cpu_tensor, dest_cpu_tensor);
}
#endif

  void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,
                             Tensor* device_tensor, StatusCallback done,
                             bool sync_dst_compute) const override;
  void CopyDeviceTensorToCPU(const Tensor* device_tensor,
                             absl::string_view tensor_name, Device* device,
                             Tensor* cpu_tensor, StatusCallback done) override;

// GPU MIG.
def TransferToDeviceOp: Gpu_Op<"transfer_to_device"> {
  let summary = "Transfer a CPU tensor to device.";

  let description = [{
    Transfer a CPU tensor to device.

    Example:
      %device_tensor = gpurt.transfer_to_device %cpu_tensor
  }];

  let arguments = (ins TFTensorType);
  let results = (outs TFTensorType);
  let assemblyFormat = "operands attr-dict";
}

      thread::ThreadPool* thread_pool);

  void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,
                             Tensor* device_tensor, StatusCallback done,
                             bool sync_dst_compute) const override;
  void CopyDeviceTensorToCPU(const Tensor* device_tensor,
                             absl::string_view tensor_name, Device* device,
                             Tensor* cpu_tensor, StatusCallback done) override;

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

namespace tensorflow {

void XlaHostSendDeviceContext::CopyCPUTensorToDevice(
    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()) {

#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) {
  DataType dtype = EncodePrimitiveTypeAsDataType(shape_.element_type()).value();
  TensorShape tensor_shape;