Merge branch 'main' into count_nonzero_torch_tuple_axis_

Author: ev-br

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(

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(
@@ -138,6 +124,11 @@ def count_nonzero(
    return result
 
 
+# take_along_axis: axis defaults to -1 but in cupy (and numpy) axis is a required arg
+def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
+    return cp.take_along_axis(x, indices, axis=axis)
+
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
 if hasattr(cp, 'vecdot'):
@@ -159,6 +150,7 @@ def count_nonzero(
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'bool', 'concat', 'count_nonzero', 'pow', 'sign']
+                              'bool', 'concat', 'count_nonzero', 'pow', 'sign',
+                              'take_along_axis']
 
 _all_ignore = ['cp', 'get_xp']

array_api_compat/numpy/_aliases.py

@@ -140,6 +140,11 @@ def count_nonzero(
     return result
 
 
+# take_along_axis: axis defaults to -1 but in numpy axis is a required arg
+def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
+    return np.take_along_axis(x, indices, axis=axis)
+
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
 if hasattr(np, "vecdot"):
@@ -175,6 +180,7 @@ def count_nonzero(
     "concat",
     "count_nonzero",
     "pow",
+    "take_along_axis"
 ]
 __all__ += _aliases.__all__
 _all_ignore = ["np", "get_xp"]

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]
@@ -37,6 +34,11 @@ array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__sub
 array_api_tests/test_operators_and_elementwise_functions.py::test_add[__add__(x, s)]
 # floating point inaccuracy
 array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[remainder(x1, x2)]
+# incomplete NEP50 support in CuPy 13.x (fixed in 14.0.0a1)
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[pow]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[logaddexp]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[nextafter]
+array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[hypot]
 
 # cupy (arg)min/max wrong with infinities
 # https://github.com/cupy/cupy/issues/7424
@@ -185,7 +187,6 @@ array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
 array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
 
 # 2024.12 support
-array_api_tests/test_signatures.py::test_func_signature[count_nonzero]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_and]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_left_shift]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_or]

dask-xfails.txt

@@ -24,6 +24,9 @@ array_api_tests/test_creation_functions.py::test_linspace
 # Shape mismatch
 array_api_tests/test_indexing_functions.py::test_take
 
+# missing `take_along_axis`, https://github.com/dask/dask/issues/3663
+array_api_tests/test_indexing_functions.py::test_take_along_axis
+
 # Array methods and attributes not already on da.Array cannot be wrapped
 array_api_tests/test_has_names.py::test_has_names[array_method-__array_namespace__]
 array_api_tests/test_has_names.py::test_has_names[array_method-to_device]

numpy-1-26-xfails.txt

@@ -50,6 +50,8 @@ array_api_tests/test_signatures.py::test_func_signature[bitwise_right_shift]
 array_api_tests/test_signatures.py::test_func_signature[bitwise_xor]
 array_api_tests/test_data_type_functions.py::TestResultType::test_with_scalars
 
+array_api_tests/test_operators_and_elementwise_functions.py::test_where_with_scalars
+
 # Stubs have a comment: (**note**: libraries may return ``NaN`` to match Python behavior.); Apparently, NumPy does just that
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]

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)

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