{"record":{"id":"66800e65a26cbd2b","repo":"jax-ml/jax","slug":"dlpack-only-supported-for-unsharded-arrays","errorCode":null,"errorMessage":"__dlpack__ only supported for unsharded arrays.","messagePattern":"__dlpack__ only supported for unsharded arrays\\.","errorType":"exception","errorClass":"BufferError","httpStatus":null,"severity":"error","filePath":"jax/_src/array.py","lineNumber":438,"sourceCode":"    from jax._src.dlpack import to_dlpack  # pyrefly: ignore[missing-import]\n\n    device_set = self.sharding.device_set\n    if len(device_set) > 1:\n      raise BufferError(\n        \"to_dlpack can only pack a dlpack tensor from an array on a singular \"\n        f\"device, but an array with a Sharding over {len(device_set)} devices \"\n        \"was provided.\"\n      )\n    device, = device_set\n    return to_dlpack(self, stream=stream,\n                     max_version=max_version,\n                     src_device=device,\n                     dl_device=dl_device,\n                     copy=copy)\n\n  def __dlpack_device__(self) -> tuple[enum.Enum, int]:\n    if len(self._arrays) != 1:\n      raise BufferError(\"__dlpack__ only supported for unsharded arrays.\")\n\n    from jax._src.dlpack import DLDeviceType  # pyrefly: ignore[missing-import]\n\n    if self.platform() == \"cpu\":\n      return DLDeviceType.kDLCPU, 0\n\n    elif self.platform() == \"gpu\":\n      platform_version = _get_device(self).client.platform_version\n      if \"cuda\" in platform_version:\n        if self.sharding.memory_kind == \"pinned_host\":\n          dl_device_type = DLDeviceType.kDLCUDAHost\n        else:\n          dl_device_type = DLDeviceType.kDLCUDA\n      elif \"rocm\" in platform_version:\n        if self.sharding.memory_kind == \"pinned_host\":\n          dl_device_type = DLDeviceType.kDLROCMHost\n        else:\n          dl_device_type = DLDeviceType.kDLROCM","sourceCodeStart":420,"sourceCodeEnd":456,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/array.py#L420-L456","documentation":"jax.Array.__dlpack_device__ (called by consumers via from_dlpack to discover the exporting device) raises BufferError when the array is backed by more than one underlying buffer, i.e. it is sharded. The DLPack device-protocol needs exactly one device per capsule, so only unsharded arrays are supported.","triggerScenarios":"torch.from_dlpack(x), np.from_dlpack(x), or any consumer that queries __dlpack_device__ on a jax array with len(x._arrays) != 1 — typically an array sharded across devices or with committed multi-buffer state (e.g. outputs of pjit/shard_map on multi-device setups).","commonSituations":"PyTorch/CuPy interop on TPU or multi-GPU pods; arrays returned from sharded jitted functions; converting training metrics or logits to torch for a custom loss; version changes where previously single-device arrays became sharded by default sharding strategies.","solutions":["Consolidate to a single device: x = jax.device_get(x) then re-upload, or jax.device_put(x, jax.devices()[0])","Force full replication before export using jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P()))","Convert through the host: np.asarray(x) then torch.as_tensor(np.asarray(x))"],"exampleFix":"// before\nt = torch.from_dlpack(sharded_x)  # BufferError: only supported for unsharded arrays\n// after\nimport numpy as np\nt = torch.as_tensor(np.asarray(sharded_x))\n# or: sharded_x = jax.device_put(sharded_x, jax.devices()[0])","handlingStrategy":"fallback","validationCode":"def is_unsharded(x) -> bool:\n    # public proxy: single addressable buffer on one device\n    return x.is_fully_addressable and len(x.sharding.device_set) == 1\n\nif not is_unsharded(x):\n    x = jax.device_put(x, jax.devices()[0])","typeGuard":"def is_dlpack_ready(x) -> bool:\n    return (hasattr(x, '__dlpack_device__')\n            and x.is_fully_addressable\n            and len(x.sharding.device_set) == 1)","tryCatchPattern":"try:\n    t = torch.from_dlpack(x)\nexcept BufferError:\n    x = jax.device_put(x, jax.devices()[0])\n    t = torch.from_dlpack(x)","preventionTips":["Centralize framework transfers in one utility that gathers first","Test interop paths under the same device mesh as production training","Prefer device_get → host when correctness matters more than zero-copy"],"tags":["jax","dlpack","sharding","buffer-protocol","interop"],"backgroundTag":"dlpack-export-sharded-array","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}