{"record":{"id":"acecf8f05325c62a","repo":"jax-ml/jax","slug":"dlpack-device-only-supported-for-tpu-pinned-ho","errorCode":null,"errorMessage":"__dlpack__ device only supported for TPU pinned host memory","messagePattern":"__dlpack__ device only supported for TPU pinned host memory","errorType":"exception","errorClass":"BufferError","httpStatus":null,"severity":"error","filePath":"jax/_src/array.py","lineNumber":473,"sourceCode":"        else:\n          dl_device_type = DLDeviceType.kDLROCM\n      elif \"oneapi\" in platform_version:\n        dl_device_type = DLDeviceType.kDLOneAPI\n      else:\n        raise BufferError(\"Unknown GPU platform for __dlpack__: \"\n                         f\"{platform_version}\")\n\n      local_hardware_id = _get_device(self).local_hardware_id\n      if local_hardware_id is None:\n        raise BufferError(\"Couldn't get local_hardware_id for __dlpack__\")\n\n      return dl_device_type, local_hardware_id\n\n    elif self.platform() == \"tpu\":\n      if self.sharding.memory_kind == \"pinned_host\":\n        dl_device_type = DLDeviceType.kDLTPUHost\n      else:\n        raise BufferError(\n            \"__dlpack__ device only supported for TPU pinned host memory\"\n        )\n\n      local_hardware_id = _get_device(self).local_hardware_id\n      if local_hardware_id is None:\n        raise BufferError(\"Couldn't get local_hardware_id for __dlpack__\")\n\n      return dl_device_type, local_hardware_id\n\n    else:\n      raise BufferError(\n          \"__dlpack__ device only supported for CPU, GPU and TPU pinned host,\"\n          f\" got platform: {self.platform()}\"\n      )\n\n  def __reduce__(self):\n    fun, args, arr_state = self._value.__reduce__()\n    aval_state = {'weak_type': self.aval.weak_type}","sourceCodeStart":455,"sourceCodeEnd":491,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/array.py#L455-L491","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nt = torch.from_dlpack(tpu_x)  # BufferError: TPU pinned host memory required\n// after\nimport numpy as np\nt = torch.as_tensor(np.asarray(tpu_x))","handlingStrategy":"fallback","validationCode":"if x.platform() == 'tpu' and x.sharding.memory_kind != 'pinned_host':\n    x = jax.device_get(x)  # export via host instead of DLPack","typeGuard":null,"tryCatchPattern":"try:\n    t = torch.from_dlpack(x)\nexcept BufferError:\n    t = torch.as_tensor(np.asarray(x))  # TPU: host round-trip","preventionTips":["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"],"tags":["jax","dlpack","tpu","pinned-memory","interop"],"backgroundTag":"dlpack-tpu-memory-kind","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}