jax-ml/jax · error · BufferError

Unknown GPU platform for dlpack

Error message

Unknown GPU platform for __dlpack__: {platform_version}

What it means

When exporting via DLPack, JAX maps the CUDA platform_version string to a known GPU backend (CUDA, ROCm, OneAPI). If the runtime's platform_version string contains none of the recognized markers (e.g. 'cuda', 'rocm', 'oneapi'), JAX cannot determine which DLPack device type to advertise and raises BufferError.

Solutions

Upgrade jax and jaxlib to matching versions that recognize the platform string
Inspect x.platform() and the extension's platform_version to confirm which backend is active
Route the transfer through host memory: np.asarray(x) then move to the target framework
Report the platform_version string upstream so it can be whitelisted

Example fix

// before
t = torch.from_dlpack(gpu_x)  # BufferError: Unknown GPU platform
// after
import numpy as np
t = torch.as_tensor(np.asarray(gpu_x)).to('cuda')

Defensive patterns

Strategy: fallback

Validate before calling

x_platform = x.platform()
version_ok = any(m in jax._src.xla_bridge.get_backend().platform_version
                 for m in ('cuda', 'rocm', 'oneapi')) if x_platform != 'cpu' else True
if not version_ok:
    x = np.asarray(x)  # go via host instead of DLPack

Try / catch

try:
    t = torch.from_dlpack(x)
except BufferError:
    # unknown/unsupported platform: round-trip through host
    t = torch.as_tensor(np.asarray(x)).to(x.platform())

Prevention

Pin jax and jaxlib to matching versions from the same release
Run a smoke-test from_dlpack transfer in CI for your exact GPU stack
Check x.platform() and backend platform_version when interop behaves oddly

When it happens

Trigger: Calling __dlpack__/from_dlpack on a GPU array when jax's cuda/rocm extension reports an unrecognized platform_version string — e.g. development builds, patched drivers, or new backend variants not known to the installed JAX version.

Common situations: Newer GPU runtime or fork (e.g. a different HIP build) with an older JAX that doesn't recognize its version string; mismatched jax/jaxlib versions; source-built jaxlib with custom platform branding.

Related errors

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

Appendix: source

Thrown at jax/_src/array.py:460

    if self.platform() == "cpu":
      return DLDeviceType.kDLCPU, 0

    elif self.platform() == "gpu":
      platform_version = _get_device(self).client.platform_version
      if "cuda" in platform_version:
        if self.sharding.memory_kind == "pinned_host":
          dl_device_type = DLDeviceType.kDLCUDAHost
        else:
          dl_device_type = DLDeviceType.kDLCUDA
      elif "rocm" in platform_version:
        if self.sharding.memory_kind == "pinned_host":
          dl_device_type = DLDeviceType.kDLROCMHost
        else:
          dl_device_type = DLDeviceType.kDLROCM
      elif "oneapi" in platform_version:
        dl_device_type = DLDeviceType.kDLOneAPI
      else:
        raise BufferError("Unknown GPU platform for __dlpack__: "
                         f"{platform_version}")

      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

    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:

View on GitHub (pinned to 1e1c6a8fc0)