{"record":{"id":"93f6d47118c8fc5e","repo":"jax-ml/jax","slug":"cuda-array-interface-is-supported-only-for-u","errorCode":null,"errorMessage":"__cuda_array_interface__() is supported only for unsharded arrays.","messagePattern":"__cuda_array_interface__\\(\\) is supported only for unsharded arrays\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/array.py","lineNumber":505,"sourceCode":"      )\n\n  def __reduce__(self):\n    fun, args, arr_state = self._value.__reduce__()\n    aval_state = {'weak_type': self.aval.weak_type}\n    return (_reconstruct_array, (fun, args, arr_state, aval_state))\n\n  @use_cpp_method()\n  def unsafe_buffer_pointer(self):\n    if len(self._arrays) != 1:\n      raise ValueError(\"unsafe_buffer_pointer() is supported only for unsharded\"\n                       \" arrays.\")\n    return self._arrays[0].unsafe_buffer_pointer()\n\n  @property\n  @use_cpp_method()\n  def __cuda_array_interface__(self):\n    if len(self._arrays) != 1:\n      raise ValueError(\"__cuda_array_interface__() is supported only for \"\n                       \"unsharded arrays.\")\n    return self._arrays[0].__cuda_array_interface__  # bind-properties\n\n  @use_cpp_method()\n  def on_device_size_in_bytes(self):\n    \"\"\"Returns the total global on-device size of the array in bytes.\"\"\"\n    arr = self._arrays[0]\n    per_shard_size = arr.on_device_size_in_bytes()\n    return per_shard_size * self.sharding.num_devices\n\n  def devices(self) -> set[Device]:\n    self._check_if_deleted()\n    return self.sharding.device_set\n\n  @property\n  def device_buffer(self):\n    raise AttributeError(\n      \"arr.device_buffer has been deprecated. Use arr.addressable_data(0)\")","sourceCodeStart":487,"sourceCodeEnd":523,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/array.py#L487-L523","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Copy to a single device: x = jax.device_put(x, jax.devices()[0]) then retry","Gather to host and re-upload in the target library: cp.asarray(np.asarray(x))","Check sharding before interop: assert x.is_fully_replicated or single-device sharding","Use jax's own device_get when a host copy is acceptable"],"exampleFix":"// before\nimport cupy as cp\nc = cp.asarray(sharded_x)  # ValueError: only for unsharded arrays\n// after\nx = jax.device_put(sharded_x, jax.devices()[0])\nc = cp.asarray(x)","handlingStrategy":"validation","validationCode":"def to_unsharded_cuda(x):\n    if len(x.sharding.device_set) > 1:\n        x = jax.device_put(x, jax.devices()[0])\n    return x\n\nimport cupy as cp\nc = cp.asarray(to_unsharded_cuda(x))","typeGuard":"def is_zero_copy_cupy_ready(x) -> bool:\n    return (x.platform() in ('cuda', 'rocm')\n            and x.is_fully_addressable\n            and len(x.sharding.device_set) == 1)","tryCatchPattern":"try:\n    c = cp.asarray(x)\nexcept ValueError:\n    c = cp.asarray(np.asarray(x))  # host round-trip","preventionTips":["Gather sharded arrays before handing them to CuPy/Numba/torch","Write one as_cupy/as_torch utility that handles sharding for the whole codebase","Add multi-GPU CI tests for interop paths so sharding regressions surface early"],"tags":["jax","cuda-array-interface","cupy","sharding","interop"],"backgroundTag":"cuda-array-interface-sharded-array","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}