jax-ml/jax · error · TypeError

Array passed to from_dlpack is on unsupported device type (D

Error message

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

What it means

The external array's __dlpack_device__ reports a DLDeviceType that JAX cannot map to a backend (no CPU/GPU match), so from_dlpack cannot construct a JAX array on that device.

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)

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
Defensive patterns

Strategy: fallback

Validate before calling

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 / catch

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

Prevention

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


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