jax-ml/jax · error · BufferError

dlpack device only supported for CPU, GPU and TPU pinned

Error message

__dlpack__ device only supported for CPU, GPU and TPU pinned host, got platform: {self.platform()}

What it means

jax.Array.__dlpack_device__ supports only CPU, GPU, and TPU (pinned host) platforms. Arrays on any other backend (e.g. plugin backends, or future/alternate XLA devices) trigger BufferError with the offending platform name included.

Source

Thrown at jax/_src/array.py:484

      return dl_device_type, local_hardware_id

    elif self.platform() == "tpu":
      if self.sharding.memory_kind == "pinned_host":
        dl_device_type = DLDeviceType.kDLTPUHost
      else:
        raise BufferError(
            "__dlpack__ device only supported for TPU pinned host memory"
        )

      local_hardware_id = _get_device(self).local_hardware_id
      if local_hardware_id is None:
        raise BufferError("Couldn't get local_hardware_id for __dlpack__")

      return dl_device_type, local_hardware_id

    else:
      raise BufferError(
          "__dlpack__ device only supported for CPU, GPU and TPU pinned host,"
          f" got platform: {self.platform()}"
      )

  def __reduce__(self):
    fun, args, arr_state = self._value.__reduce__()
    aval_state = {'weak_type': self.aval.weak_type}
    return (_reconstruct_array, (fun, args, arr_state, aval_state))

  @use_cpp_method()
  def unsafe_buffer_pointer(self):
    if len(self._arrays) != 1:
      raise ValueError("unsafe_buffer_pointer() is supported only for unsharded"
                       " arrays.")
    return self._arrays[0].unsafe_buffer_pointer()

  @property
  @use_cpp_method()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Transfer through host memory: np.asarray(x) then move to the target framework
Check x.platform() before attempting DLPack and pick an interop path per backend
Request/await DLPack support for the backend upstream
Pin your workload to a supported backend (cpu/cuda/rocm/tpu) if interop is essential

Example fix

// before
t = torch.from_dlpack(x)  # BufferError: only CPU, GPU, TPU pinned host
// after
import numpy as np
t = torch.as_tensor(np.asarray(x))  # host round-trip works on any backend

Defensive patterns

Strategy: validation

Validate before calling

DLPACK_PLATFORMS = {'cpu', 'cuda', 'rocm', 'oneapi', 'tpu'}
if x.platform() not in DLPACK_PLATFORMS:
    x = np.asarray(x)  # unsupported backend: transfer via host

Type guard

def supports_dlpack(x) -> bool:
    return hasattr(x, '__dlpack__') and x.platform() in {
        'cpu', 'cuda', 'rocm', 'oneapi', 'tpu'}

Try / catch

try:
    t = torch.from_dlpack(x)
except BufferError:
    t = torch.as_tensor(np.asarray(x))  # backend-agnostic fallback

Prevention

Branch interop strategy on x.platform() explicitly
Keep a host-memory transfer path as the universal fallback
Track upstream DLPack support for any plugin backend you depend on

When it happens

Trigger: Calling from_dlpack consumers on an array whose x.platform() is not 'cpu', 'cuda'/'rocm'/'oneapi' variants, or 'tpu' — e.g. IPU/Metal/plugin backends or custom XLA devices.

Common situations: Using jax with third-party plugin backends (ipu, metal, etc.) and attempting DLPack interop with torch/numpy; experimental backends where DLPack was never implemented.

Related errors

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