jax-ml/jax · error · xla::XlaRuntimeError

This operation is implemented for a PjRt-compatible backend

Error message

This operation is implemented for a PjRt-compatible backend only.

What it means

Exporting an array to DLPack requires reaching the underlying PjRtBuffer; if the IFRT array does not implement PjRtCompatibleArray, jaxlib cannot get a device pointer for the tensor and throws.

Source

Thrown at jaxlib/dlpack.cc:283

          data, element_type, dimensions, byte_strides,
          xla::PjRtClient::HostBufferSemantics::kMutableZeroCopy,
          on_delete_callback, memory_space, /*device_layout=*/nullptr));
  return std::make_pair(std::move(buffer), true);
}

}  // namespace

absl::StatusOr<nb::capsule> BufferToDLPackManagedTensor(
    nb::handle py_buffer, std::optional<std::intptr_t> stream) {
  ifrt::Array* ifrt_array = nb::cast<PyArray>(py_buffer).ifrt_array();
  if (ifrt_array == nullptr) {
    return xla::Unimplemented(
        "BufferToDLPackManagedTensor called on deleted array.");
  }
  auto* arr =
      xla::ifrt::dyn_cast_or_null<ifrt::PjRtCompatibleArray>(ifrt_array);
  if (arr == nullptr) {
    throw xla::XlaRuntimeError(
        "This operation is implemented for a PjRt-compatible backend only.");
  }
  xla::PjRtBuffer* pjrt_buffer = arr->pjrt_buffers().front().get();

  if (pjrt_buffer->IsTuple()) {
    return xla::Unimplemented(
        "BufferToDLPackManagedTensor is not implemented for tuple "
        "buffers.");
  }
  if (pjrt_buffer->has_dynamic_dimensions()) {
    return xla::Unimplemented("DynamicShape is not implemented in DLPack.");
  }

  auto pack = std::make_unique<DLPackTensor>();
  DLTensor& dt = pack->tensor.dl_tensor;
  {
    // AcquireExternalReference may block; there are no API guarantees.
    GlobalPyRefManager()->CollectGarbage();

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Copy the array to a PjRt-compatible backend device (e.g. jax.device_put(arr, jax.devices('cpu')[0])) before exporting
Avoid DLPack interop on non-PjRt IFRT backends; use numpy round-trip as a fallback

Example fix

# before
torch.from_dlpack(ifrt_array)
# after
import jax
a = jax.device_put(ifrt_array, jax.devices('cpu')[0])
torch.from_dlpack(a)

Defensive patterns

Strategy: fallback

Validate before calling

from jax.extend import backend as jeb
assert jeb.get_backend_c_api() is not None, 'backend not PjRt-compatible; DLPack export unsupported'

Try / catch

try:
    capsule = arr.__dlpack__()
except Exception:
    capsule = None
    np_arr = np.asarray(arr)

Prevention

device_put arrays onto a standard backend before interop
Prefer numpy round-trip for exotic backends

When it happens

Trigger: Calling dlpack.to_dlpack(device_array) / array.__dlpack__() on an array living on a non-PjRt-compatible IFRT backend.

Common situations: Interop (torch, cupy) with arrays on experimental IFRT-native backends.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/e523e000192f9caf. Report an issue: GitHub.