{"record":{"id":"a91b2e9a7ace201f","repo":"jax-ml/jax","slug":"to-dlpack-can-only-pack-a-dlpack-tensor-from-an-ar","errorCode":null,"errorMessage":"to_dlpack can only pack a dlpack tensor from an array on a singular device, but an array with a Sharding over {len(device_set)} devices was provided.","messagePattern":"to_dlpack can only pack a dlpack tensor from an array on a singular device, but an array with a Sharding over (.+?) devices was provided\\.","errorType":"exception","errorClass":"BufferError","httpStatus":null,"severity":"error","filePath":"jax/_src/array.py","lineNumber":424,"sourceCode":"    \"\"\"\n    return self.sharding.is_fully_addressable\n\n  def __array__(self, dtype: np.dtype | None = None,\n                context: None = None, copy: bool | None = None):\n    del context  # unused\n    # copy argument is supported by np.asarray starting in numpy 2.0\n    kwds = {} if copy is None else {'copy': copy}\n    return np.asarray(self._value, dtype=dtype, **kwds)  # pyrefly: ignore[no-matching-overload]\n\n  def __dlpack__(self, *, stream: int | Any | None = None,\n                 max_version: tuple[int, int] | None = None,\n                 dl_device: tuple[DLDeviceType, int] | None = None,\n                 copy: bool | None = None):\n    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\":","sourceCodeStart":406,"sourceCodeEnd":442,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/array.py#L406-L442","documentation":"jax.Array.__dlpack__ refuses to export an array whose Sharding spans more than one device, because a single dlpack capsule can only describe one device's memory. This occurs when the array is sharded across multiple GPUs/TPUs under jax.jit with multi-device sharding or is not fully materialized on one device.","triggerScenarios":"Calling torch.from_numpy-like interop via DLPack (e.g. torch.from_dlpack(x), np.from_dlpack(x), cupy.from_dlpack(x)) on an ArrayImpl/Array whose sharding.device_set has >1 device, e.g. arrays returned from pjit/shard_map or jax.device_put with a multi-device Sharding.","commonSituations":"Moving sharded TPU/GPU results to PyTorch or NumPy in multi-device pipelines; arrays produced under jax.jit with PartitionSpec sharding; forgetting to gather before export; TPU pod slices where every op returns multi-device arrays.","solutions":["Gather the array to one device first: x = jax.device_get(x) (to host) or jax.device_put(x, jax.devices()[0])","If using NamedSharding, use jax.lax.with_sharding_constraint or GSPMD to make it fully replicated before export","Use x.addressable_data(0) only if it is fully replicated; otherwise copy via np.asarray(x)"],"exampleFix":"// before\nx = pjit_fn(inputs)  # sharded over 8 devices\nt = torch.from_dlpack(x)  # BufferError\n// after\nx = jax.device_put(x, jax.devices()[0])\nt = torch.from_dlpack(x)","handlingStrategy":"validation","validationCode":"def ensure_single_device(x, dev=None):\n    dev = dev or jax.devices()[0]\n    if len(x.sharding.device_set) > 1:\n        return jax.device_put(x, dev)\n    return x\n\nx = ensure_single_device(x)\nt = torch.from_dlpack(x)","typeGuard":"def is_dlpack_exportable(x) -> bool:\n    return (hasattr(x, '__dlpack__')\n            and len(getattr(x, 'sharding').device_set) == 1)","tryCatchPattern":"try:\n    t = torch.from_dlpack(x)\nexcept BufferError:\n    t = torch.as_tensor(np.asarray(x))  # host round-trip fallback","preventionTips":["Gather/replicate sharded results before any cross-framework handoff","Wrap interop boundaries in a to_torch(x) helper that handles sharding centrally","Log x.sharding when interop fails on multi-device jobs to catch surprise sharding"],"tags":["jax","dlpack","interop","sharding","multi-device"],"backgroundTag":"dlpack-export-sharded-array","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}