Support the copy keyword in asarray

array_api_compat/common/_aliases.py

@@ -291,9 +291,11 @@ def _asarray(
     """
     Array API compatibility wrapper for asarray().
 
-    See the corresponding documentation in NumPy/CuPy and/or the array API
+    See the corresponding documentation in the array library and/or the array API
     specification for more details.
 
+    'namespace' may be an array module namespace. This is needed to support
+    conversion of sequences of Python scalars.
     """
     if namespace is None:
         try:

array_api_compat/cupy/_aliases.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-from functools import partial
-
 import cupy as cp
 
 from ..common import _aliases
 from .._internal import get_xp
 
-asarray = asarray_cupy = partial(_aliases._asarray, namespace='cupy')
-asarray.__doc__ = _aliases._asarray.__doc__
-del partial
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from typing import Optional, Union
+    from ._typing import ndarray, Device, Dtype, NestedSequence, SupportsBufferProtocol
 
 bool = cp.bool_
 
@@ -62,6 +61,42 @@
 matrix_transpose = get_xp(cp)(_aliases.matrix_transpose)
 tensordot = get_xp(cp)(_aliases.tensordot)
 
+# asarray also adds the copy keyword, which is not present in numpy 1.0.
+def asarray(
+    obj: Union[
+        ndarray,
+        bool,
+        int,
+        float,
+        NestedSequence[bool | int | float],
+        SupportsBufferProtocol,
+    ],
+    /,
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+    copy: Optional[bool] = None,
+    **kwargs,
+) -> ndarray:
+    """
+    Array API compatibility wrapper for asarray().
+
+    See the corresponding documentation in the array library and/or the array API
+    specification for more details.
+
+    'namespace' may be an array module namespace. This is needed to support
+    conversion of sequences of Python scalars.
+    """
+    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 None:
+            copy = False
+        elif copy is False:
+            raise NotImplementedError("asarray(copy=False) is not yet supported in cupy")
+
+        return cp.array(obj, copy=copy, dtype=dtype, **kwargs)
+
 # 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'):
@@ -73,7 +108,7 @@
 else:
     isdtype = get_xp(cp)(_aliases.isdtype)
 
-__all__ = _aliases.__all__ + ['asarray', 'asarray_cupy', 'bool', 'acos',
+__all__ = _aliases.__all__ + ['asarray', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',
                               'atanh', 'bitwise_left_shift', 'bitwise_invert',
                               'bitwise_right_shift', 'concat', 'pow']

array_api_compat/numpy/__init__.py

@@ -21,4 +21,10 @@
 
 from ..common._helpers import * # noqa: F403
 
+try:
+    # Used in asarray(). Not present in older versions.
+    from numpy import _CopyMode
+except ImportError:
+    pass
+
 __array_api_version__ = '2022.12'

array_api_compat/numpy/_aliases.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
-from functools import partial
-
 from ..common import _aliases
 
 from .._internal import get_xp
 
-asarray = asarray_numpy = partial(_aliases._asarray, namespace='numpy')
-asarray.__doc__ = _aliases._asarray.__doc__
-del partial
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from typing import Optional, Union
+    from ._typing import ndarray, Device, Dtype, NestedSequence, SupportsBufferProtocol
 
 import numpy as np
 bool = np.bool_
@@ -62,6 +61,65 @@
 matrix_transpose = get_xp(np)(_aliases.matrix_transpose)
 tensordot = get_xp(np)(_aliases.tensordot)
 
+def _supports_buffer_protocol(obj):
+    try:
+        memoryview(obj)
+    except TypeError:
+        return False
+    return True
+
+# asarray also adds the copy keyword, which is not present in numpy 1.0.
+# asarray() is different enough between numpy, cupy, and dask, the logic
+# complicated enough that it's easier to define it separately for each module
+# rather than trying to combine everything into one function in common/
+def asarray(
+    obj: Union[
+        ndarray,
+        bool,
+        int,
+        float,
+        NestedSequence[bool | int | float],
+        SupportsBufferProtocol,
+    ],
+    /,
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+    copy: "Optional[Union[bool, np._CopyMode]]" = None,
+    **kwargs,
+) -> ndarray:
+    """
+    Array API compatibility wrapper for asarray().
+
+    See the corresponding documentation in the array library and/or the array API
+    specification for more details.
+
+    'namespace' may be an array module namespace. This is needed to support
+    conversion of sequences of Python scalars.
+    """
+    if np.__version__[0] >= '2':
+        # NumPy 2.0 asarray() is completely array API compatible. No need for
+        # the complicated logic below
+        return np.asarray(obj, dtype=dtype, device=device, copy=copy, **kwargs)
+
+    if device not in ["cpu", None]:
+        raise ValueError(f"Unsupported device for NumPy: {device!r}")
+
+    if hasattr(np, '_CopyMode'):
+        if copy is None:
+            copy = np._CopyMode.IF_NEEDED
+        elif copy is False:
+            copy = np._CopyMode.NEVER
+        elif copy is True:
+            copy = np._CopyMode.ALWAYS
+    else:
+        # Not present in older NumPys. In this case, we cannot really support
+        # copy=False.
+        if copy is False:
+            raise NotImplementedError("asarray(copy=False) requires a newer version of NumPy.")
+
+    return np.array(obj, copy=copy, dtype=dtype, **kwargs)
+
 # 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'):
@@ -73,7 +131,7 @@
 else:
     isdtype = get_xp(np)(_aliases.isdtype)
 
-__all__ = _aliases.__all__ + ['asarray', 'asarray_numpy', 'bool', 'acos',
+__all__ = _aliases.__all__ + ['asarray', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',
                               'atanh', 'bitwise_left_shift', 'bitwise_invert',
                               'bitwise_right_shift', 'concat', 'pow']

tests/_helpers.py

@@ -5,6 +5,9 @@
 import pytest
 
 
+wrapped_libraries = ["numpy", "cupy", "torch", "dask.array"]
+all_libraries = wrapped_libraries + ["jax.numpy"]
+
 def import_(library, wrapper=False):
     if 'jax' in library and sys.version_info < (3, 9):
         pytest.skip('JAX array API support does not support Python 3.8')

tests/test_all.py

@@ -12,11 +12,11 @@
 
 import sys
 
-from ._helpers import import_
+from ._helpers import import_, wrapped_libraries
 
 import pytest
 
-@pytest.mark.parametrize("library", ["common", "cupy", "numpy", "torch", "dask.array"])
+@pytest.mark.parametrize("library", ["common"] + wrapped_libraries)
 def test_all(library):
     import_(library, wrapper=True)
 

tests/test_array_namespace.py

@@ -9,9 +9,9 @@
 import array_api_compat
 from array_api_compat import array_namespace
 
-from ._helpers import import_
+from ._helpers import import_, all_libraries
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+@pytest.mark.parametrize("library", all_libraries)
 @pytest.mark.parametrize("api_version", [None, "2021.12"])
 def test_array_namespace(library, api_version):
     xp = import_(library)

tests/test_common.py

@@ -3,10 +3,11 @@
 
 from array_api_compat import is_array_api_obj, device, to_device
 
-from ._helpers import import_
+from ._helpers import import_, wrapped_libraries, all_libraries
 
 import pytest
 import numpy as np
+import array
 from numpy.testing import assert_allclose
 
 is_functions = {
@@ -29,7 +30,7 @@ def test_is_xp_array(library, func):
 
     assert is_array_api_obj(x)
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+@pytest.mark.parametrize("library", all_libraries)
 def test_device(library):
     xp = import_(library, wrapper=True)
 
@@ -43,7 +44,7 @@ def test_device(library):
     assert device(x) == device(x2)
 
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array"])
+@pytest.mark.parametrize("library", wrapped_libraries)
 def test_to_device_host(library):
     # different libraries have different semantics
     # for DtoH transfers; ensure that we support a portable
@@ -64,7 +65,7 @@ def test_to_device_host(library):
 
 @pytest.mark.parametrize("target_library", is_functions.keys())
 @pytest.mark.parametrize("source_library", is_functions.keys())
-def test_asarray(source_library, target_library, request):
+def test_asarray_cross_library(source_library, target_library, request):
     if source_library == "dask.array" and target_library == "torch":
         # Allow rest of test to execute instead of immediately xfailing
         # xref https://github.com/pandas-dev/pandas/issues/38902
@@ -83,3 +84,97 @@ def test_asarray(source_library, target_library, request):
     b = tgt_lib.asarray(a)
 
     assert is_tgt_type(b), f"Expected {b} to be a {tgt_lib.ndarray}, but was {type(b)}"
+
+@pytest.mark.parametrize("library", wrapped_libraries)
+def test_asarray_copy(library):
+    # Note, we have this test here because the test suite currently doesn't
+    # test the copy flag to asarray() very rigorously. Once
+    # https://github.com/data-apis/array-api-tests/issues/241 is fixed we
+    # should be able to delete this.
+    xp = import_(library, wrapper=True)
+    asarray = xp.asarray
+    is_lib_func = globals()[is_functions[library]]
+    all = xp.all
+
+    if library == 'numpy' and xp.__version__[0] < '2' and not hasattr(xp, '_CopyMode') :
+        supports_copy_false = False
+    elif library == 'cupy':
+        supports_copy_false = False
+    else:
+        supports_copy_false = True
+
+    a = asarray([1])
+    b = asarray(a, copy=True)
+    assert is_lib_func(b)
+    a[0] = 0
+    assert all(b[0] == 1)
+    assert all(a[0] == 0)
+
+    a = asarray([1])
+    if supports_copy_false:
+        b = asarray(a, copy=False)
+        assert is_lib_func(b)
+        a[0] = 0
+        assert all(b[0] == 0)
+    else:
+        pytest.raises(NotImplementedError, lambda: asarray(a, copy=False))
+
+    a = asarray([1])
+    if supports_copy_false:
+        pytest.raises(ValueError, lambda: asarray(a, copy=False,
+                                                  dtype=xp.float64))
+    else:
+        pytest.raises(NotImplementedError, lambda: asarray(a, copy=False, dtype=xp.float64))
+
+    a = asarray([1])
+    b = asarray(a, copy=None)
+    assert is_lib_func(b)
+    a[0] = 0
+    assert all(b[0] == 0)
+
+    a = asarray([1.0], dtype=xp.float32)
+    b = asarray(a, dtype=xp.float64, copy=None)
+    assert is_lib_func(b)
+    a[0] = 0.0
+    assert all(b[0] == 1.0)
+
+    a = asarray([1.0], dtype=xp.float64)
+    b = asarray(a, dtype=xp.float64, copy=None)
+    assert is_lib_func(b)
+    a[0] = 0.0
+    assert all(b[0] == 0.0)
+
+    # Python built-in types
+    for obj in [True, 0, 0.0, 0j, [0], [[0]]]:
+        asarray(obj, copy=True) # No error
+        asarray(obj, copy=None) # No error
+        if supports_copy_false:
+            pytest.raises(ValueError, lambda: asarray(obj, copy=False))
+        else:
+            pytest.raises(NotImplementedError, lambda: asarray(obj, copy=False))
+
+    # Use the standard library array to test the buffer protocol
+    a = array.array('f', [1.0])
+    b = asarray(a, copy=True)
+    assert is_lib_func(b)
+    a[0] = 0.0
+    assert all(b[0] == 1.0)
+
+    a = array.array('f', [1.0])
+    if supports_copy_false:
+        b = asarray(a, copy=False)
+        assert is_lib_func(b)
+        a[0] = 0.0
+        assert all(b[0] == 0.0)
+    else:
+        pytest.raises(NotImplementedError, lambda: asarray(a, copy=False))
+
+    a = array.array('f', [1.0])
+    b = asarray(a, copy=None)
+    assert is_lib_func(b)
+    a[0] = 0.0
+    if library == 'cupy':
+        # A copy is required for libraries where the default device is not CPU
+        assert all(b[0] == 1.0)
+    else:
+        assert all(b[0] == 0.0)

tests/test_isdtype.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from ._helpers import import_
+from ._helpers import import_, wrapped_libraries
 
 # Check the known dtypes by their string names
 
@@ -64,7 +64,7 @@ def isdtype_(dtype_, kind):
     assert type(res) is bool #  noqa: E721
     return res
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+@pytest.mark.parametrize("library", wrapped_libraries)
 def test_isdtype_spec_dtypes(library):
     xp = import_(library, wrapper=True)
 
@@ -98,7 +98,7 @@ def test_isdtype_spec_dtypes(library):
     'bfloat16',
 ]
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+@pytest.mark.parametrize("library", wrapped_libraries)
 @pytest.mark.parametrize("dtype_", additional_dtypes)
 def test_isdtype_additional_dtypes(library, dtype_):
     xp = import_(library, wrapper=True)