se::Stream* dst_compute_stream = dst_xla_context->stream();
    auto* const dst_compute_stream_impl =
        static_cast<tpu::TpuStreamInterface*>(dst_compute_stream);

    if (src_compute_stream_impl->IsSameSharedMemoryLocation(
            dst_compute_stream_impl)) {
      // Surprisingly, this path does get triggered in practice.
      *output = *input;
      done(absl::OkStatus());
      return absl::OkStatus();
    }