jax-ml/jax · error · NotImplementedError

stream argument of array.to_device()

Error message

stream argument of array.to_device()

What it means

`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.

Source

Thrown at jax/_src/numpy/array_methods.py:470

  Refer to :func:`jax.numpy.take` for full documentation.
  """
  return indexing.take(self, indices, axis=axis, out=out, mode=mode, unique_indices=unique_indices,
                       indices_are_sorted=indices_are_sorted, fill_value=fill_value)

def _to_device(self: Array, device: xc.Device | Sharding, *,
               stream: int | Any | None = None):
  """Return a copy of the array on the specified device

  Args:
    device: :class:`~jax.Device` or :class:`~jax.sharding.Sharding`
      to which the created array will be committed.
    stream: not implemented, passing a non-None value will lead to an error.
  Returns:
    copy of array placed on the specified device or devices.
  """
  if stream is not None:
    raise NotImplementedError("stream argument of array.to_device()")
  return api.device_put(self, device)


def _trace(self: Array, offset: int | ArrayLike = 0, axis1: int = 0, axis2: int = 1,
           dtype: DTypeLike | None = None, out: None = None) -> Array:
  """Return the sum along the diagonal.

  Refer to :func:`jax.numpy.trace` for full documentation.
  """
  return lax_numpy.trace(self, offset=offset, axis1=axis1, axis2=axis2, dtype=dtype, out=out)

def _transpose(self: Array, *args: Any) -> Array:
  """Returns a copy of the array with axes transposed.

  Refer to :func:`jax.numpy.transpose` for full documentation.
  """
  if not args:
    axis = None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Call to_device without the stream argument
2. If ordering with another library's work matters, use explicit synchronization (e.g. torch.cuda.synchronize() or jax.block_until_ready) around the copy

Example fix

# before
arr.to_device(jax.devices('gpu')[0], stream=cupy_stream)

# after
arr.to_device(jax.devices('gpu')[0])

def to_device(arr, device, stream=None):
    if stream is not None:
        raise NotImplementedError('JAX ignores streams; synchronize explicitly')
    return arr.to_device(device)

When it happens

Trigger: Calling `arr.to_device(dev, stream=some_cuda_stream)` with stream not None.

Common situations: Interop code migrating from CuPy/PyTorch where copying with an explicit stream is standard; copy-pasting a non-JAX to_device signature.

Related errors

• array ref with memory space only works inside of a `jit`.
• pinned array ref only works inside of a `jit`.
• __pow__ modulo not implemented
• __rpow__ modulo not implemented
• for grad support, subclass {type(self)} must implement `vjp_

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