`PjRtCApiBuffer::AcquireExternalReference` does not support `WaitUntilBufferReadyOnStream`

openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators

Apache License 2.0

2.68k stars 430 forks source link

`PjRtCApiBuffer::AcquireExternalReference` does not support `WaitUntilBufferReadyOnStream` #11535

Closed jjwilke closed 6 months ago

jjwilke commented 6 months ago

When wrapping a stream executor client with the C API, the returned external reference from PjRtCApiBuffer::AcquireExternalReference does not wrap a stream executor external reference and does not currently support WaitUntilBufferReadyOnStream. This prevents calls to, e.g.:

dlt_2 = xla_client._xla.buffer_to_dlpack_managed_tensor(buffer, stream.ptr)

jyingl3 commented 6 months ago

Work-in-progress PR: https://github.com/openxla/xla/pull/11629, https://github.com/google/jax/pull/20805

jyingl3 commented 6 months ago

The PRs have been merged. Closed this issue.