jax-ml/jax · error · TypeError

Array passed to from_dlpack is on unsupported device type

Error message

Array passed to from_dlpack is on unsupported device type (DLDeviceType: {dl_device_type}, array: {external_array}

What it means

from_dlpack queries the producer's __dlpack_device__ and only knows how to map known DLDeviceType values (CPU, CUDA) onto JAX devices. It fires when the external array reports an unsupported device type, so JAX cannot place the imported buffer.

Solutions

1. Move the source tensor to CPU or a supported device in the source framework before from_dlpack
2. Install the JAX build matching the accelerator (jax-rocm vs jax-cuda)
3. Fall back to jax.numpy.asarray(t) which copies through the host

Example fix

# before
jax.dlpack.from_dlpack(vulkan_tensor)

# after
jax.dlpack.from_dlpack(vulkan_tensor.cpu())  # or asarray fallback

from jax._src.dlpack import _DL_DEVICE_TO_PLATFORM
if external.__dlpack_device__()[0] not in _DL_DEVICE_TO_PLATFORM:
    external = external.cpu() if hasattr(external, 'cpu') else None
assert external is not None

try:
    jax.dlpack.from_dlpack(t)
except TypeError:
    jax.dlpack.from_dlpack(t.cpu())

When it happens

Trigger: from_dlpack on tensors living on devices like kDLOpenCL, kDLVulkan, kDLMetal (on unsupported builds), or any enum absent from JAX's _DL_DEVICE_TO_PLATFORM.

Common situations: Interop with Vulkan/OpenCL-backed arrays, or a JAX build lacking the matching accelerator backend (e.g. ROCm tensor with CUDA-only jax).

Related errors

• __dlpack__ device only supported for CPU, GPU and TPU…
• __dlpack__ device only supported for TPU pinned host memory
• DLPack is supported for PjRt-compatible backends only.
• __dlpack__ only supported for unsharded arrays.
• The device specification passed to to_dlpack contains an…

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

Appendix: source

Thrown at jax/_src/dlpack.py:250

  """
  if isinstance(device, Sharding):
    device_set = device.device_set
    if len(device_set) > 1:
      raise ValueError(
        "from_dlpack can only unpack a dlpack tensor onto a singular device, but "
        f"a Sharding with {len(device_set)} devices was provided."
      )
    device, = device_set
  if not hasattr(external_array, "__dlpack__") or not hasattr(external_array, "__dlpack_device__"):
    raise TypeError(
        "The array passed to from_dlpack must have __dlpack__ and __dlpack_device__ methods."
    )

  dl_device_type, device_id = external_array.__dlpack_device__()
  try:
    dl_device_platform = _DL_DEVICE_TO_PLATFORM[dl_device_type]
  except KeyError:
    raise TypeError(
        "Array passed to from_dlpack is on unsupported device type "
        f"(DLDeviceType: {dl_device_type}, array: {external_array}"
    ) from None

  backend = xla_bridge.get_backend(dl_device_platform)
  dlpack_device = backend.device_from_local_hardware_id(device_id)
  _check_device(device, dlpack_device, copy)
  if _is_tensorflow_tensor(external_array):
    # TensorFlow does not support stream=.
    stream = None
  elif dl_device_type in (
      DLDeviceType.kDLCUDAHost,
      DLDeviceType.kDLROCMHost,
      DLDeviceType.kDLTPUHost,
  ):
    # Some producers (e.g. torch.Tensor with is_pinned()) route pinned tensors
    # through their CPU __dlpack__, which rejects a non-None stream argument.
    stream = None

View on GitHub (pinned to 1e1c6a8fc0)