jax-ml/jax · error · BufferError

dlpack device only supported for TPU pinned host memory

Error message

__dlpack__ device only supported for TPU pinned host memory

What it means

For TPU arrays, JAX only supports DLPack device export when the array lives in 'pinned_host' memory, because the DLPack TPU device type (kDLTPUHost) corresponds to pinned host buffers. Any other TPU memory kind is rejected with BufferError.

Source

Thrown at jax/_src/array.py:473

        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:
      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}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Transfer through host memory: x_host = jax.device_get(x) then np.asarray / torch.as_tensor
Place the array in pinned host memory first: jax.device_put(x, jax.sharding.SingleDeviceSharding(dev, memory_kind='pinned_host')) (if supported by your setup)
Avoid DLPack on TPU; use .tolist()/np.asarray for interop

Example fix

// before
t = torch.from_dlpack(tpu_x)  # BufferError: TPU pinned host memory required
// after
import numpy as np
t = torch.as_tensor(np.asarray(tpu_x))

Defensive patterns

Strategy: fallback

Validate before calling

if x.platform() == 'tpu' and x.sharding.memory_kind != 'pinned_host':
    x = jax.device_get(x)  # export via host instead of DLPack

Try / catch

try:
    t = torch.from_dlpack(x)
except BufferError:
    t = torch.as_tensor(np.asarray(x))  # TPU: host round-trip

Prevention

Don't rely on DLPack for TPU interop; default to device_get/np.asarray
If pinned-host export is needed, explicitly device_put with memory_kind='pinned_host'
Document platform-specific interop limits in shared utility code

When it happens

Trigger: Calling __dlpack_device__/from_dlpack on a jax array on a TPU whose sharding.memory_kind is not 'pinned_host' — i.e. ordinary on-device TPU memory, the default for device_put/jit outputs.

Common situations: Trying to move TPU results directly to another framework via DLPack; assuming DLPack works uniformly across platforms; migrating GPU interop code to TPU without accounting for the pinned-host restriction.

Related errors

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