jax-ml/jax · error · BufferError

Unknown GPU platform for dlpack: {platform_version}

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.

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)

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.