jax-ml/jax · error · BufferError

__dlpack__ only supported for unsharded arrays.

Error message

__dlpack__ only supported for unsharded arrays.

What it means

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.

Source

Thrown at jax/_src/array.py:438

    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":
      return DLDeviceType.kDLCPU, 0

    elif self.platform() == "gpu":
      platform_version = _get_device(self).client.platform_version
      if "cuda" in platform_version:
        if self.sharding.memory_kind == "pinned_host":
          dl_device_type = DLDeviceType.kDLCUDAHost
        else:
          dl_device_type = DLDeviceType.kDLCUDA
      elif "rocm" in platform_version:
        if self.sharding.memory_kind == "pinned_host":
          dl_device_type = DLDeviceType.kDLROCMHost
        else:
          dl_device_type = DLDeviceType.kDLROCM

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Consolidate to a single device: x = jax.device_get(x) then re-upload, or jax.device_put(x, jax.devices()[0])
2. Force full replication before export using jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P()))
3. Convert through the host: np.asarray(x) then torch.as_tensor(np.asarray(x))

Example fix

// before
t = torch.from_dlpack(sharded_x)  # BufferError: only supported for unsharded arrays
// after
import numpy as np
t = torch.as_tensor(np.asarray(sharded_x))
# or: sharded_x = jax.device_put(sharded_x, jax.devices()[0])

def is_unsharded(x) -> bool:
    # public proxy: single addressable buffer on one device
    return x.is_fully_addressable and len(x.sharding.device_set) == 1

if not is_unsharded(x):
    x = jax.device_put(x, jax.devices()[0])

def is_dlpack_ready(x) -> bool:
    return (hasattr(x, '__dlpack_device__')
            and x.is_fully_addressable
            and len(x.sharding.device_set) == 1)

try:
    t = torch.from_dlpack(x)
except BufferError:
    x = jax.device_put(x, jax.devices()[0])
    t = torch.from_dlpack(x)

When it happens

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

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

Related errors

• to_dlpack can only pack a dlpack tensor from an array on a s
• __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
• The device specification passed to to_dlpack contains an uns

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