{"record":{"id":"ab2babc723816e7f","repo":"jax-ml/jax","slug":"stream-argument-of-array-to-device","errorCode":null,"errorMessage":"stream argument of array.to_device()","messagePattern":"stream argument of array\\.to_device\\(\\)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/array_methods.py","lineNumber":470,"sourceCode":"\n  Refer to :func:`jax.numpy.take` for full documentation.\n  \"\"\"\n  return indexing.take(self, indices, axis=axis, out=out, mode=mode, unique_indices=unique_indices,\n                       indices_are_sorted=indices_are_sorted, fill_value=fill_value)\n\ndef _to_device(self: Array, device: xc.Device | Sharding, *,\n               stream: int | Any | None = None):\n  \"\"\"Return a copy of the array on the specified device\n\n  Args:\n    device: :class:`~jax.Device` or :class:`~jax.sharding.Sharding`\n      to which the created array will be committed.\n    stream: not implemented, passing a non-None value will lead to an error.\n  Returns:\n    copy of array placed on the specified device or devices.\n  \"\"\"\n  if stream is not None:\n    raise NotImplementedError(\"stream argument of array.to_device()\")\n  return api.device_put(self, device)\n\n\ndef _trace(self: Array, offset: int | ArrayLike = 0, axis1: int = 0, axis2: int = 1,\n           dtype: DTypeLike | None = None, out: None = None) -> Array:\n  \"\"\"Return the sum along the diagonal.\n\n  Refer to :func:`jax.numpy.trace` for full documentation.\n  \"\"\"\n  return lax_numpy.trace(self, offset=offset, axis1=axis1, axis2=axis2, dtype=dtype, out=out)\n\ndef _transpose(self: Array, *args: Any) -> Array:\n  \"\"\"Returns a copy of the array with axes transposed.\n\n  Refer to :func:`jax.numpy.transpose` for full documentation.\n  \"\"\"\n  if not args:\n    axis = None","sourceCodeStart":452,"sourceCodeEnd":488,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/array_methods.py#L452-L488","documentation":"`Array.to_device(device, stream)` accepts a stream parameter only for API compatibility with dlpack/NumPy; JAX manages streams internally and any non-None stream raises NotImplementedError. Passing a CUDA stream from another library cannot be honored.","triggerScenarios":"Calling `arr.to_device(dev, stream=some_cuda_stream)` with stream not None.","commonSituations":"Interop code migrating from CuPy/PyTorch where copying with an explicit stream is standard; copy-pasting a non-JAX to_device signature.","solutions":["Call to_device without the stream argument","If ordering with another library's work matters, use explicit synchronization (e.g. torch.cuda.synchronize() or jax.block_until_ready) around the copy"],"exampleFix":"# before\narr.to_device(jax.devices('gpu')[0], stream=cupy_stream)\n\n# after\narr.to_device(jax.devices('gpu')[0])","handlingStrategy":"validation","validationCode":"def to_device(arr, device, stream=None):\n    if stream is not None:\n        raise NotImplementedError('JAX ignores streams; synchronize explicitly')\n    return arr.to_device(device)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never forward stream kwargs from CuPy/PyTorch code to JAX","Use jax.block_until_ready / other-lib synchronize for ordering"],"tags":["jax","device-transfer","stream","not-implemented"],"backgroundTag":"unsupported-argument-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}