jax-ml/jax · error · BufferError

to_dlpack can only pack a dlpack tensor from an array on a s

Error message

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.

What it means

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.

Source

Thrown at jax/_src/array.py:424

    """
    return self.sharding.is_fully_addressable

  def __array__(self, dtype: np.dtype | None = None,
                context: None = None, copy: bool | None = None):
    del context  # unused
    # copy argument is supported by np.asarray starting in numpy 2.0
    kwds = {} if copy is None else {'copy': copy}
    return np.asarray(self._value, dtype=dtype, **kwds)  # pyrefly: ignore[no-matching-overload]

  def __dlpack__(self, *, stream: int | Any | None = None,
                 max_version: tuple[int, int] | None = None,
                 dl_device: tuple[DLDeviceType, int] | None = None,
                 copy: bool | None = None):
    from jax._src.dlpack import to_dlpack  # pyrefly: ignore[missing-import]

    device_set = self.sharding.device_set
    if len(device_set) > 1:
      raise BufferError(
        "to_dlpack can only pack a dlpack tensor from an array on a singular "
        f"device, but an array with a Sharding over {len(device_set)} devices "
        "was provided."
      )
    device, = device_set
    return to_dlpack(self, stream=stream,
                     max_version=max_version,
                     src_device=device,
                     dl_device=dl_device,
                     copy=copy)

  def __dlpack_device__(self) -> tuple[enum.Enum, int]:
    if len(self._arrays) != 1:
      raise BufferError("__dlpack__ only supported for unsharded arrays.")

    from jax._src.dlpack import DLDeviceType  # pyrefly: ignore[missing-import]

    if self.platform() == "cpu":

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Gather the array to one device first: x = jax.device_get(x) (to host) or jax.device_put(x, jax.devices()[0])
2. If using NamedSharding, use jax.lax.with_sharding_constraint or GSPMD to make it fully replicated before export
3. Use x.addressable_data(0) only if it is fully replicated; otherwise copy via np.asarray(x)

Example fix

// before
x = pjit_fn(inputs)  # sharded over 8 devices
t = torch.from_dlpack(x)  # BufferError
// after
x = jax.device_put(x, jax.devices()[0])
t = torch.from_dlpack(x)

def ensure_single_device(x, dev=None):
    dev = dev or jax.devices()[0]
    if len(x.sharding.device_set) > 1:
        return jax.device_put(x, dev)
    return x

x = ensure_single_device(x)
t = torch.from_dlpack(x)

def is_dlpack_exportable(x) -> bool:
    return (hasattr(x, '__dlpack__')
            and len(getattr(x, 'sharding').device_set) == 1)

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• __dlpack__ only supported for unsharded arrays.
• from_dlpack can only unpack a dlpack tensor onto a singular
• __dlpack__ device only supported for TPU pinned host memory
• __dlpack__ device only supported for CPU, GPU and TPU pinned
• __cuda_array_interface__() is supported only for unsharded a

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