jax-ml/jax · error · ValueError

unsafe_buffer_pointer() is supported only for unsharded arra

Error message

unsafe_buffer_pointer() is supported only for unsharded arrays.

What it means

Array.unsafe_buffer_pointer() returns the raw device/host address of the array's underlying buffer and therefore requires the array to be backed by exactly one buffer (unsharded). Sharded or multi-buffer arrays have no single pointer, so ValueError is raised.

Source

Thrown at jax/_src/array.py:497

        raise BufferError("Couldn't get local_hardware_id for __dlpack__")

      return dl_device_type, local_hardware_id

    else:
      raise BufferError(
          "__dlpack__ device only supported for CPU, GPU and TPU pinned host,"
          f" got platform: {self.platform()}"
      )

  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Consolidate first: x = jax.device_put(x, jax.devices()[0]) (or device_get for host)
2. Guard: if len(x.addressable_data(0)._arrays) ... — practically, check x.is_fully_addressable and single-device sharding before calling
3. Use x.__cuda_array_interface__ after making it unsharded, which exposes the same pointer for CUDA consumers

Example fix

// before
ptr = sharded_x.unsafe_buffer_pointer()  # ValueError
// after
x = jax.device_put(sharded_x, jax.devices()[0])
ptr = x.unsafe_buffer_pointer()

def to_single_buffer(x):
    if len(x.sharding.device_set) > 1 or not x.is_fully_addressable:
        return jax.device_put(x, jax.devices()[0])
    return x

x = to_single_buffer(x)
ptr = x.unsafe_buffer_pointer()

def has_single_buffer(x) -> bool:
    return x.is_fully_addressable and len(x.sharding.device_set) == 1

try:
    ptr = x.unsafe_buffer_pointer()
except ValueError:
    x = jax.device_put(x, jax.devices()[0])
    ptr = x.unsafe_buffer_pointer()

When it happens

Trigger: Calling x.unsafe_buffer_pointer() (often indirectly via __cuda_array_interface__ or custom C++/CUDA kernel plumbing) on an array with len(x._arrays) != 1 — i.e. any sharded array on multi-device setups or non-addressable committed arrays.

Common situations: Passing jax buffers into custom CUDA kernels or profiling tools that need a raw pointer; TPU/multi-GPU pjit outputs fed to low-level code; assuming jit outputs are always single-buffer.

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• to_dlpack can only pack a dlpack tensor from an array on a s
• __dlpack__ only supported for unsharded arrays.
• __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/c71d50b2dfd2043e. Report an issue: GitHub.