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

    host_alloc_attr.set_on_host(true);
    host_allocator_ = device_->GetAllocator(host_alloc_attr);

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

    tensorflow::DeviceContext* device_context;
    auto status = device_->TryGetDeviceContext(&device_context);
    TF_EXPECT_OK(status);

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

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

  XlaPlatformInfo platform_info(
      compilation_device_type, /*platform_id=*/nullptr,
      /*xla_device_metadata=*/nullptr,
      /*pjrt_device_metadata=*/pjrt_device_metadata.get(),
      /*device_allocator=*/nullptr);

  XlaCompiler::Options options = GenerateCompilerOptionsForPjRt(
      *device_setup_.flr(), device, platform_info,
      /*pjrt_device_compiler=*/nullptr);

                                /*xla_device_metadata=*/nullptr,
                                /*pjrt_device_metadata=*/nullptr,
                                /*device_allocator=*/nullptr);

  XlaDeviceCompiler* xla_device_compiler = nullptr;
  TF_EXPECT_OK(BuildXlaDeviceCompiler(device, nullptr, platform_info,
                                      compilation_device_type,

  } else if (default_device_ordinal != -1) {
    build_options.set_device_ordinal(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);

//
// 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 =