}

  return absl::OkStatus();
}

// Provide SendDeviceMemoryFunction for XLA host callbacks.  This callback
// handles transferring from device to host.
xla::SendDeviceMemoryFunction GetSendDeviceMemoryFunction(
    OpKernelContext* ctx, const std::string& program_key) {
  return
      [ctx, program_key](
          int64_t channel_id, se::Stream* stream, const xla::Shape& shape,