// then device_allocator_ is the xla::Backend's memory allocator.  If the op
  // is placed on a regular CPU or GPU device then device_allocator_ is null.
  // The allocator is of unknown provenance; keep it in a shared pointer to
  // set an artificial refcount of one.
  std::shared_ptr<se::DeviceMemoryAllocator> device_allocator_;

  XlaPlatformInfo(const XlaPlatformInfo&) = delete;

    AllocatorAttributes device_alloc_attr;
    device_alloc_attr.set_on_host(false);
    device_allocator_ = device_->GetAllocator(device_alloc_attr);
  }

 protected:
  std::unique_ptr<Device> device_;
  Allocator* host_allocator_;
  Allocator* device_allocator_;
};

TEST_F(XlaHostSendRecvDeviceContextTest, CopyDeviceTensorToCPU) {
  SetDevice("GPU");

    return host_tensor;
  }

  // Creates a Tensor on device using the device_allocator_
  template <typename T>
  Tensor* CreateDeviceTensor(const TensorShape& shape,
                             const gtl::ArraySlice<T> data) {
    Tensor* host_tensor = CreateHostTensor<T>(shape, data);
    Tensor* device_tensor =
        new Tensor(device_allocator_, DataTypeToEnum<T>::v(), shape);

//
// Example usage:
//
//  Device device;
//  stream_executor::Stream stream(executor);
//  Tensor cpu_tensor(host_allocator, DT_FLOAT, TensorShape({2, 2}));
//  Tensor device_tensor(device_allocator, DT_FLOAT, TensorShape({2, 2}));
//  se::DeviceMemoryBase gpu_dst{device_tensor.data(), 4 * sizeof(float)};
//  xla::Shape shape(xla::F32, {2, 2}, {}, {})
//  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =

// Concrete once transfer is completed.
//
// Example usage:
//
//  Device device;
//  stream_executor::Stream stream(executor);
//  Tensor device_tensor(device_allocator, DT_FLOAT, TensorShape({2, 2}));
//  se::DeviceMemoryBase gpu_dst{device_tensor.data(), 4 * sizeof(float)};
//  xla::Shape shape(xla::F32, {2, 2}, {}, {})
//  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =

  options.graph_def_version = function_library.graph_def_version();
  options.allow_cpu_custom_calls =
      (platform_info.platform_id() == se::host::kHostPlatformId);
  options.device_allocator = GetAllocator(device, stream, platform_info);
  if (platform_info.xla_device_metadata()) {
    options.shape_determination_fns =
        platform_info.xla_device_metadata()->default_shape_determination_fns();
  }

                                   : local_client->default_device_ordinal());
  build_options.set_result_layout(result.xla_output_shape);
  build_options.set_device_allocator(options.device_allocator.get());
  build_options.set_alias_passthrough_params(options.alias_passthrough_params);
  build_options.mutable_debug_options()->set_xla_detailed_logging(
      options.detailed_logging);