jax-ml/jax · error · ValueError

__cuda_array_interface__() is supported only for unsharded a

Error message

__cuda_array_interface__() is supported only for unsharded arrays.

What it means

The __cuda_array_interface__ property exposes the CUDA array interface dict (needed by CuPy, Numba, PyTorch zero-copy paths) and only exists for unsharded arrays backed by a single CUDA buffer. Sharded arrays raise ValueError because there is no single device pointer/shape tuple to expose.

Source

Thrown at jax/_src/array.py:505

      )

  def __reduce__(self):
    fun, args, arr_state = self._value.__reduce__()
    aval_state = {'weak_type': self.aval.weak_type}
    return (_reconstruct_array, (fun, args, arr_state, aval_state))

  @use_cpp_method()
  def unsafe_buffer_pointer(self):
    if len(self._arrays) != 1:
      raise ValueError("unsafe_buffer_pointer() is supported only for unsharded"
                       " arrays.")
    return self._arrays[0].unsafe_buffer_pointer()

  @property
  @use_cpp_method()
  def __cuda_array_interface__(self):
    if len(self._arrays) != 1:
      raise ValueError("__cuda_array_interface__() is supported only for "
                       "unsharded arrays.")
    return self._arrays[0].__cuda_array_interface__  # bind-properties

  @use_cpp_method()
  def on_device_size_in_bytes(self):
    """Returns the total global on-device size of the array in bytes."""
    arr = self._arrays[0]
    per_shard_size = arr.on_device_size_in_bytes()
    return per_shard_size * self.sharding.num_devices

  def devices(self) -> set[Device]:
    self._check_if_deleted()
    return self.sharding.device_set

  @property
  def device_buffer(self):
    raise AttributeError(
      "arr.device_buffer has been deprecated. Use arr.addressable_data(0)")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Copy to a single device: x = jax.device_put(x, jax.devices()[0]) then retry
2. Gather to host and re-upload in the target library: cp.asarray(np.asarray(x))
3. Check sharding before interop: assert x.is_fully_replicated or single-device sharding
4. Use jax's own device_get when a host copy is acceptable

Example fix

// before
import cupy as cp
c = cp.asarray(sharded_x)  # ValueError: only for unsharded arrays
// after
x = jax.device_put(sharded_x, jax.devices()[0])
c = cp.asarray(x)

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

import cupy as cp
c = cp.asarray(to_unsharded_cuda(x))

def is_zero_copy_cupy_ready(x) -> bool:
    return (x.platform() in ('cuda', 'rocm')
            and x.is_fully_addressable
            and len(x.sharding.device_set) == 1)

try:
    c = cp.asarray(x)
except ValueError:
    c = cp.asarray(np.asarray(x))  # host round-trip

When it happens

Trigger: CuPy.asarray(x), numba.cuda.jit kernels taking x, torch tensors constructed from jax arrays, or any consumer of __cuda_array_interface__ when len(x._arrays) != 1 (multi-GPU/TPU sharded arrays).

Common situations: Zero-copy handoff of jax arrays to CuPy/Numba/torch on multi-GPU jobs; outputs of pjit/shard_map passed to custom kernels; migration from single-GPU code to sharded pipelines without gathering.

Related errors

• to_dlpack can only pack a dlpack tensor from an array on a s
• __dlpack__ only supported for unsharded arrays.
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• __dlpack__ device only supported for TPU pinned host memory

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