jax-ml/jax · error · BufferError

Couldn't get local_hardware_id for __dlpack__

Error message

Couldn't get local_hardware_id for __dlpack__

What it means

During DLPack export of a GPU array, JAX needs the device's local_hardware_id (the physical device index) to fill the DLPack device tuple. If the backend device object returns None for local_hardware_id, the export cannot proceed and BufferError is raised.

Source

Thrown at jax/_src/array.py:465

      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:
        raise BufferError("Couldn't get local_hardware_id for __dlpack__")

      return dl_device_type, local_hardware_id

    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Verify devices initialize correctly: print(jax.devices()) and check local_hardware_id on each
  2. Upgrade/downgrade jaxlib to match the installed JAX and CUDA/ROCm stack
  3. Avoid device filtering (CUDA_VISIBLE_DEVICES quirks) or container device cgroup restrictions and retry
  4. Fall back to host transfer: np.asarray(x)

Example fix

// before
t = torch.from_dlpack(x)  # BufferError: Couldn't get local_hardware_id
// after
for d in jax.devices():
    assert d.local_hardware_id is not None, d
# if that fails, fix backend init; otherwise transfer via host:
t = torch.as_tensor(np.asarray(x)).to('cuda')
Defensive patterns

Strategy: validation

Validate before calling

devices_ok = all(d.local_hardware_id is not None for d in jax.devices())
if not devices_ok:
    raise RuntimeError('GPU backend not fully initialized; fix before DLPack use')

Try / catch

try:
    t = torch.from_dlpack(x)
except BufferError:
    t = torch.as_tensor(np.asarray(x)).to('cuda')

Prevention

When it happens

Trigger: Calling __dlpack__/__dlpack_device__ on a CUDA/ROCm/OneAPI array where _get_device(self).local_hardware_id is None — typically a mis-initialized or partially initialized backend, or an unusual device abstraction from plugins.

Common situations: Custom XLA backend plugins or device proxies that don't populate local_hardware_id; jaxlib/device driver version mismatch; GPU visible-device configurations that strip hardware IDs (e.g. cgroup/device filtering in containers).

Related errors


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