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

DLPack is only supported for devices addressable by the curr

Error message

DLPack is only supported for devices addressable by the current process.

What it means

DLPack import needs to write the tensor data onto the target device; only devices addressable by the current process (i.e. local) can be written to, so remote/virtual devices are rejected.

Source

Thrown at jaxlib/dlpack.cc:368

                    }));
  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);
  }
  absl::Span<int64_t const> dimensions(
      reinterpret_cast<int64_t*>(dlmt->dl_tensor.shape), dlmt->dl_tensor.ndim);
  TF_ASSIGN_OR_RETURN(xla::PrimitiveType element_type,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Select an addressable local device (jax.local_devices()[0])
In multi-process setups, import the DLPack tensor on the process that owns the data

Example fix

# before
jax.dlpack.from_dlpack(t, device=jax.devices()[3])
# after
jax.dlpack.from_dlpack(t, device=jax.local_devices()[0])

Defensive patterns

Strategy: validation

Validate before calling

import jax
local = set(id(d) for d in jax.local_devices())
assert id(dev) in local, 'device not addressable by this process'

Type guard

def is_local(d) -> bool:
    import jax
    return d in jax.local_devices()

Prevention

Always import DLPack tensors onto jax.local_devices()[0]
In multi-host runs, import on the host that owns the data

When it happens

Trigger: jax.dlpack.from_dlpack(tensor, device=d) where d is a non-addressable device — e.g. a remote device in multi-process jax, or another process's device.

Common situations: Multi-host/multi-process jax runs where the device list includes remote devices and the caller picked devices[1] which lives on another node.

Related errors

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