jax-ml/jax · error · ValueError

Specified {device=} which requires a copy since the source d

Error message

Specified {device=} which requires a copy since the source device is {repr(src_device)}, however copy=False. Set copy=True or copy=None to perform the requested operation.

What it means

to_dlpack was asked to place the resulting DLPack buffer on a device different from the array's source device, which requires a copy, but copy=False forbids copies.

Source

Thrown at jax/_src/dlpack.py:70

def is_supported_dtype(dtype: DTypeLike) -> bool:
  """Check if dtype is supported by jax.dlpack."""
  if dtype is None:
    # NumPy will silently cast this to float64, which may be surprising.
    raise TypeError(f"Expected a string or dtype-like object; got {dtype=}")
  return np.dtype(dtype) in SUPPORTED_DTYPES_SET


def _to_dlpack(x: Array, stream: int | Any | None,
               src_device: _jax.Device | None = None,
               device: _jax.Device | None = None,
               copy: bool | None = None):

  if src_device is None:
    src_device, = x.devices()
  if device and (src_device is None or device != src_device):
    if copy is not None and not copy:
      raise ValueError(
        f"Specified {device=} which requires a copy since the source device "
        f"is {repr(src_device)}, however copy=False. Set copy=True or "
        "copy=None to perform the requested operation."
      )
    else:
      arr = device_put(x, device)
  else:
    arr = _array_copy(x) if copy else x
  return _jax.buffer_to_dlpack_managed_tensor(
    arr.addressable_data(0), stream=stream
  )


_DL_DEVICE_TO_PLATFORM = {
    DLDeviceType.kDLCPU: "cpu",
    DLDeviceType.kDLCUDA: "cuda",
    DLDeviceType.kDLCUDAHost: "cuda",
    DLDeviceType.kDLROCM: "rocm",

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Set copy=True (or copy=None, the default, which copies only when needed)
  2. First move the array with jax.device_put(x, device) and then to_dlpack with matching device
  3. Omit the device argument if you want the buffer on the array's current device

Example fix

# before
jax.dlpack.to_dlpack(x, device=dev1, copy=False)

# after
jax.dlpack.to_dlpack(x, device=dev1, copy=True)
Defensive patterns

Strategy: validation

Validate before calling

src, = x.devices()
if device is not None and device != src:
    assert copy is not False, 'copy=False with cross-device target'

Try / catch

try:
    jax.dlpack.to_dlpack(x, device=dev, copy=copy)
except ValueError:
    jax.dlpack.to_dlpack(x, device=dev, copy=True)

Prevention

When it happens

Trigger: Calling jax.dlpack.to_dlpack(x, device=other_device, copy=False) where other_device != x.devices()[0].

Common situations: Zero-copy interop with PyTorch/TensorFlow where the target GPU differs from the JAX array's GPU; multi-GPU pipelines assuming cross-device zero-copy is possible.

Related errors


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