jax-ml/jax · error · ValueError

from_dlpack can only unpack a dlpack tensor onto a singular

Error message

from_dlpack can only unpack a dlpack tensor onto a singular device, but a Sharding with {len(device_set)} devices was provided.

What it means

from_dlpack can place a DLPack buffer on exactly one device; passing a multi-device Sharding (e.g. a NamedSharding or PmapSharding over several devices) is ambiguous and rejected.

Source

Thrown at jax/_src/dlpack.py:236

      If ``copy=True`` then a copy is always performed, even if unpacked onto
      the same device. If ``copy=False`` then the copy is never performed and
      will raise an error if necessary. When ``copy=None`` then a copy may be
      performed if needed for a device transfer.

  Returns:
    A jax.Array

  Note:
    While JAX arrays are always immutable, dlpack buffers cannot be marked as
    immutable, and it is possible for processes external to JAX to mutate them
    in-place. If a jax Array is constructed from a dlpack buffer and the buffer
    is later modified in-place, it may lead to undefined behavior when using
    the associated JAX array.
  """
  if isinstance(device, Sharding):
    device_set = device.device_set
    if len(device_set) > 1:
      raise ValueError(
        "from_dlpack can only unpack a dlpack tensor onto a singular device, but "
        f"a Sharding with {len(device_set)} devices was provided."
      )
    device, = device_set
  if not hasattr(external_array, "__dlpack__") or not hasattr(external_array, "__dlpack_device__"):
    raise TypeError(
        "The array passed to from_dlpack must have __dlpack__ and __dlpack_device__ methods."
    )

  dl_device_type, device_id = external_array.__dlpack_device__()
  try:
    dl_device_platform = _DL_DEVICE_TO_PLATFORM[dl_device_type]
  except KeyError:
    raise TypeError(
        "Array passed to from_dlpack is on unsupported device type "
        f"(DLDeviceType: {dl_device_type}, array: {external_array}"
    ) from None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a single Device (e.g. jax.devices()[0]) and shard afterwards with jax.device_put(arr, sharding)
2. Use a single-device sharding (device_set of size 1)
3. Import the array first, then redistribute with device_put(arr, desired_sharding)

Example fix

# before
jax.dlpack.from_dlpack(t, device=multi_device_sharding)

# after
arr = jax.dlpack.from_dlpack(t, device=jax.devices()[0])
arr = jax.device_put(arr, multi_device_sharding)

if isinstance(device, jax.sharding.Sharding):
    assert len(device.device_set) == 1, 'use a single device for from_dlpack'
    device, = device.device_set

import jax.sharding as sh
def is_single_device_sharding(s) -> bool:
    return isinstance(s, sh.Sharding) and len(s.device_set) == 1

try:
    jax.dlpack.from_dlpack(t, device=dev)
except ValueError as e:
    if 'singular device' not in str(e): raise
    arr = jax.dlpack.from_dlpack(t); arr = jax.device_put(arr, dev)

When it happens

Trigger: jax.dlpack.from_dlpack(external, device=jax.sharding.NamedSharding(mesh, P)) where the sharding's device_set has more than one device.

Common situations: Reusing sharding objects from a jitted multi-device pipeline as the device argument for interop imports.

Related errors

• to_dlpack can only pack a dlpack tensor from an array on a s
• __dlpack__ only supported for unsharded arrays.
• Query, key and value should have same sharding.
• 0th dimension of all xs should be replicated. Got {}
• 0th dimension of leaf passed to `jax.lax.map` should be repl

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