diff --git a/array_api_compat/common/_helpers.py b/array_api_compat/common/_helpers.py index d50e0d83..77175d0d 100644 --- a/array_api_compat/common/_helpers.py +++ b/array_api_compat/common/_helpers.py @@ -775,42 +775,28 @@ def _cupy_to_device( /, stream: int | Any | None = None, ) -> _CupyArray: - import cupy as cp # pyright: ignore[reportMissingTypeStubs] - from cupy.cuda import Device as _Device # pyright: ignore - from cupy.cuda import stream as stream_module # pyright: ignore - from cupy_backends.cuda.api import runtime # pyright: ignore + import cupy as cp - if device == x.device: - return x - elif device == "cpu": + if device == "cpu": # allowing us to use `to_device(x, "cpu")` # is useful for portable test swapping between # host and device backends return x.get() - elif not isinstance(device, _Device): - raise ValueError(f"Unsupported device {device!r}") - else: - # see cupy/cupy#5985 for the reason how we handle device/stream here - prev_device: Any = runtime.getDevice() # pyright: ignore[reportUnknownMemberType] - prev_stream = None - if stream is not None: - prev_stream: Any = stream_module.get_current_stream() # pyright: ignore - # stream can be an int as specified in __dlpack__, or a CuPy stream - if isinstance(stream, int): - stream = cp.cuda.ExternalStream(stream) # pyright: ignore - elif isinstance(stream, cp.cuda.Stream): # pyright: ignore[reportUnknownMemberType] - pass - else: - raise ValueError("the input stream is not recognized") - stream.use() # pyright: ignore[reportUnknownMemberType] - try: - runtime.setDevice(device.id) # pyright: ignore[reportUnknownMemberType] - arr = x.copy() - finally: - runtime.setDevice(prev_device) # pyright: ignore[reportUnknownMemberType] - if stream is not None: - prev_stream.use() - return arr + if not isinstance(device, cp.cuda.Device): + raise TypeError(f"Unsupported device type {device!r}") + + if stream is None: + with device: + return cp.asarray(x) + + # stream can be an int as specified in __dlpack__, or a CuPy stream + if isinstance(stream, int): + stream = cp.cuda.ExternalStream(stream) + elif not isinstance(stream, cp.cuda.Stream): + raise TypeError(f"Unsupported stream type {stream!r}") + + with device, stream: + return cp.asarray(x) def _torch_to_device( diff --git a/array_api_compat/cupy/_aliases.py b/array_api_compat/cupy/_aliases.py index fd1460ae..adb74bff 100644 --- a/array_api_compat/cupy/_aliases.py +++ b/array_api_compat/cupy/_aliases.py @@ -64,8 +64,6 @@ finfo = get_xp(cp)(_aliases.finfo) iinfo = get_xp(cp)(_aliases.iinfo) -_copy_default = object() - # asarray also adds the copy keyword, which is not present in numpy 1.0. def asarray( @@ -79,7 +77,7 @@ def asarray( *, dtype: Optional[DType] = None, device: Optional[Device] = None, - copy: Optional[bool] = _copy_default, + copy: Optional[bool] = None, **kwargs, ) -> Array: """ @@ -89,25 +87,13 @@ def asarray( specification for more details. """ with cp.cuda.Device(device): - # cupy is like NumPy 1.26 (except without _CopyMode). See the comments - # in asarray in numpy/_aliases.py. - if copy is not _copy_default: - # A future version of CuPy will change the meaning of copy=False - # to mean no-copy. We don't know for certain what version it will - # be yet, so to avoid breaking that version, we use a different - # default value for copy so asarray(obj) with no copy kwarg will - # always do the copy-if-needed behavior. - - # This will still need to be updated to remove the - # NotImplementedError for copy=False, but at least this won't - # break the default or existing behavior. - if copy is None: - copy = False - elif copy is False: - raise NotImplementedError("asarray(copy=False) is not yet supported in cupy") - kwargs['copy'] = copy - - return cp.array(obj, dtype=dtype, **kwargs) + if copy is None: + return cp.asarray(obj, dtype=dtype, **kwargs) + else: + res = cp.array(obj, dtype=dtype, copy=copy, **kwargs) + if not copy and res is not obj: + raise ValueError("Unable to avoid copy while creating an array as requested") + return res def astype( diff --git a/cupy-xfails.txt b/cupy-xfails.txt index a30572f8..df85d9ca 100644 --- a/cupy-xfails.txt +++ b/cupy-xfails.txt @@ -11,9 +11,6 @@ array_api_tests/test_array_object.py::test_scalar_casting[__index__(int64)] # testsuite bug (https://github.com/data-apis/array-api-tests/issues/172) array_api_tests/test_array_object.py::test_getitem -# copy=False is not yet implemented -array_api_tests/test_creation_functions.py::test_asarray_arrays - # attributes are np.float32 instead of float # (see also https://github.com/data-apis/array-api/issues/405) array_api_tests/test_data_type_functions.py::test_finfo[float32] diff --git a/tests/test_common.py b/tests/test_common.py index 6b1aa160..d1933899 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -17,6 +17,7 @@ from array_api_compat import ( device, is_array_api_obj, is_lazy_array, is_writeable_array, size, to_device ) +from array_api_compat.common._helpers import _DASK_DEVICE from ._helpers import all_libraries, import_, wrapped_libraries, xfail @@ -189,23 +190,26 @@ class C: @pytest.mark.parametrize("library", all_libraries) -def test_device(library, request): +def test_device_to_device(library, request): if library == "ndonnx": - xfail(request, reason="Needs ndonnx >=0.9.4") + xfail(request, reason="Stub raises ValueError") + if library == "sparse": + xfail(request, reason="No __array_namespace_info__()") xp = import_(library, wrapper=True) + devices = xp.__array_namespace_info__().devices() - # We can't test much for device() and to_device() other than that - # x.to_device(x.device) works. - + # Default device x = xp.asarray([1, 2, 3]) dev = device(x) - x2 = to_device(x, dev) - assert device(x2) == device(x) - - x3 = xp.asarray(x, device=dev) - assert device(x3) == device(x) + for dev in devices: + if dev is None: # JAX >=0.5.3 + continue + if dev is _DASK_DEVICE: # TODO this needs a better design + continue + y = to_device(x, dev) + assert device(y) == dev @pytest.mark.parametrize("library", wrapped_libraries) diff --git a/tests/test_cupy.py b/tests/test_cupy.py new file mode 100644 index 00000000..f8b4a4d8 --- /dev/null +++ b/tests/test_cupy.py @@ -0,0 +1,22 @@ +import pytest +from array_api_compat import device, to_device + +xp = pytest.importorskip("array_api_compat.cupy") +from cupy.cuda import Stream + + +def test_to_device_with_stream(): + devices = xp.__array_namespace_info__().devices() + streams = [ + Stream(), + Stream(non_blocking=True), + Stream(null=True), + Stream(ptds=True), + 123, # dlpack stream + ] + + a = xp.asarray([1, 2, 3]) + for dev in devices: + for stream in streams: + b = to_device(a, dev, stream=stream) + assert device(b) == dev