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

DLPack is supported for PjRt-compatible backends only.

Error message

DLPack is supported for PjRt-compatible backends only.

What it means

Importing a DLPack tensor into jax requires creating a PjRt buffer on the target device; if the IFRT device is not a PjRtDevice, jaxlib cannot construct a buffer and throws.

Source

Thrown at jaxlib/dlpack.cc:364

                      }
                      // PyCapsule_GetPointer may have raised. Restore the
                      // previous exception if there was one.
                      PyErr_SetRaisedException(exc);
                    }));
  if (!capsule.ptr()) {
    throw nb::python_error();
  }
  return capsule;
}

absl::StatusOr<nb::object> DLPackManagedTensorToBuffer(
    const nb::capsule& tensor, ifrt::Device* ifrt_device,
    nb_class_ptr<PyClient> client, std::optional<std::intptr_t> stream,
    std::optional<bool> copy, std::optional<DLDeviceType> dl_device_type) {
  ifrt::PjRtDevice* device =
      xla::ifrt::dyn_cast_or_null<ifrt::PjRtDevice>(ifrt_device);
  if (device == nullptr) {
    throw xla::XlaRuntimeError(
        "DLPack is supported for PjRt-compatible backends only.");
  }
  if (!device->IsAddressable()) {
    throw xla::XlaRuntimeError(
        "DLPack is only supported for devices addressable by the current "
        "process.");
  }
  if (std::string_view(tensor.name()) != kDlTensorCapsuleName) {
    return xla::InvalidArgument(
        "DLPack tensor must be a capsule with name \"dltensor\", got \"%s\". "
        "Note that a DLPack tensor may be consumed at most once.",
        std::string_view(tensor.name()));
  }
  DLManagedTensor* dlmt = static_cast<DLManagedTensor*>(tensor.data());
  if (dlmt->dl_tensor.ndim < 0) {
    return xla::InvalidArgument(
        "Number of dimensions in DLManagedTensor must be nonnegative, got %d",
        dlmt->dl_tensor.ndim);

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a device from a PjRt-compatible backend (default cpu/gpu devices)
2. Omit the device argument and let jax place the buffer on its default device

Example fix

# before
jax.dlpack.from_dlpack(t, device=ifrt_device)
# after
jax.dlpack.from_dlpack(t)  # or device=jax.devices('gpu')[0]

assert device.platform in ('cpu','gpu','cuda','rocm','tpu')

When it happens

Trigger: Calling dlpack.from_dlpack(tensor, device=non_pjrt_ifrt_device) or jax.dlpack.from_dlpack with a device from a non-PjRt-compatible backend.

Common situations: torch/cupy -> jax interop targeting an experimental IFRT-native backend device.

Related errors

• This operation is implemented for a PjRt-compatible backend
• to_dlpack can only pack a dlpack tensor from an array on a s
• __dlpack__ only supported for unsharded arrays.
• __dlpack__ device only supported for TPU pinned host memory
• __dlpack__ device only supported for CPU, GPU and TPU pinned

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