Add dask.array specific implementation of asarray()

This also fixes several bugs in the implementation, and updates some tests.

Author: asmeurer

array_api_compat/common/_aliases.py

@@ -6,18 +6,18 @@
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    import numpy as np
     from typing import Optional, Sequence, Tuple, Union
     from ._typing import ndarray, Device, Dtype, NestedSequence, SupportsBufferProtocol
 
 from typing import NamedTuple
-from types import ModuleType
 import inspect
 
-from ._helpers import _check_device, is_numpy_array, array_namespace
+from ._helpers import _check_device
 
 # These functions are modified from the NumPy versions.
 
+# Creation functions add the device keyword (which does nothing for NumPy)
+
 def arange(
     start: Union[int, float],
     /,
@@ -268,92 +268,6 @@ def var(
 def permute_dims(x: ndarray, /, axes: Tuple[int, ...], xp) -> ndarray:
     return xp.transpose(x, axes)
 
-# Creation functions add the device keyword (which does nothing for NumPy)
-
-# asarray also adds the copy keyword
-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,
-    namespace = 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 namespace is None:
-        try:
-            xp = array_namespace(obj, _use_compat=False)
-        except ValueError:
-            # TODO: What about lists of arrays?
-            raise ValueError("A namespace must be specified for asarray() with non-array input")
-    elif isinstance(namespace, ModuleType):
-        xp = namespace
-    elif namespace == 'numpy':
-        import numpy as xp
-    elif namespace == 'cupy':
-        import cupy as xp
-    elif namespace == 'dask.array':
-        import dask.array as xp
-    else:
-        raise ValueError("Unrecognized namespace argument to asarray()")
-
-    _check_device(xp, device)
-    if is_numpy_array(obj):
-        import numpy as np
-        if hasattr(np, '_CopyMode'):
-            # Not present in older NumPys
-            COPY_FALSE = (False, np._CopyMode.IF_NEEDED)
-            COPY_TRUE = (True, np._CopyMode.ALWAYS)
-        else:
-            COPY_FALSE = (False,)
-            COPY_TRUE = (True,)
-    else:
-        COPY_FALSE = (False,)
-        COPY_TRUE = (True,)
-    if copy in COPY_FALSE and namespace != "dask.array":
-        # copy=False is not yet implemented in xp.asarray
-        raise NotImplementedError("copy=False is not yet implemented")
-    if (hasattr(xp, "ndarray") and isinstance(obj, xp.ndarray)):
-        if dtype is not None and obj.dtype != dtype:
-            copy = True
-        if copy in COPY_TRUE:
-            return xp.array(obj, copy=True, dtype=dtype)
-        return obj
-    elif namespace == "dask.array":
-        if copy in COPY_TRUE:
-            if dtype is None:
-                return obj.copy()
-            # Go through numpy, since dask copy is no-op by default
-            import numpy as np
-            obj = np.array(obj, dtype=dtype, copy=True)
-            return xp.array(obj, dtype=dtype)
-        else:
-            import dask.array as da
-            import numpy as np
-            if not isinstance(obj, da.Array):
-                obj = np.asarray(obj, dtype=dtype)
-                return da.from_array(obj)
-            return obj
-
-    return xp.asarray(obj, dtype=dtype, **kwargs)
-
 # np.reshape calls the keyword argument 'newshape' instead of 'shape'
 def reshape(x: ndarray,
             /,

array_api_compat/dask/array/_aliases.py

@@ -37,7 +37,7 @@
 if TYPE_CHECKING:
     from typing import Optional, Union
 
-    from ...common._typing import Device, Dtype, Array
+    from ...common._typing import Device, Dtype, Array, NestedSequence, SupportsBufferProtocol
 
 import dask.array as da
 
@@ -76,10 +76,6 @@ def _dask_arange(
 arange = get_xp(da)(_dask_arange)
 eye = get_xp(da)(_aliases.eye)
 
-from functools import partial
-asarray = partial(_aliases._asarray, namespace='dask.array')
-asarray.__doc__ = _aliases._asarray.__doc__
-
 linspace = get_xp(da)(_aliases.linspace)
 eye = get_xp(da)(_aliases.eye)
 UniqueAllResult = get_xp(da)(_aliases.UniqueAllResult)
@@ -113,6 +109,47 @@ def _dask_arange(
 matmul = get_xp(np)(_aliases.matmul)
 tensordot = get_xp(np)(_aliases.tensordot)
 
+
+# asarray also adds the copy keyword, which is not present in numpy 1.0.
+def asarray(
+    obj: Union[
+        Array,
+        bool,
+        int,
+        float,
+        NestedSequence[bool | int | float],
+        SupportsBufferProtocol,
+    ],
+    /,
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+    copy: "Optional[Union[bool, np._CopyMode]]" = None,
+    **kwargs,
+) -> Array:
+    """
+    Array API compatibility wrapper for asarray().
+
+    See the corresponding documentation in the array library and/or the array API
+    specification for more details.
+    """
+    if copy is False:
+        # copy=False is not yet implemented in dask
+        raise NotImplementedError("copy=False is not yet implemented")
+    elif copy is True:
+        if isinstance(obj, da.Array) and dtype is None:
+            return obj.copy()
+        # Go through numpy, since dask copy is no-op by default
+        obj = np.array(obj, dtype=dtype, copy=True)
+        return da.array(obj, dtype=dtype)
+    else:
+        if not isinstance(obj, da.Array) or dtype is not None and obj.dtype != dtype:
+            obj = np.asarray(obj, dtype=dtype)
+            return da.from_array(obj)
+        return obj
+
+    return da.asarray(obj, dtype=dtype, **kwargs)
+
 from dask.array import (
     # Element wise aliases
     arccos as acos,

tests/test_common.py

@@ -94,11 +94,11 @@ def test_asarray_copy(library):
     xp = import_(library, wrapper=True)
     asarray = xp.asarray
     is_lib_func = globals()[is_functions[library]]
-    all = xp.all
+    all = xp.all if library != 'dask.array' else lambda x: xp.all(x).compute()
 
     if library == 'numpy' and xp.__version__[0] < '2' and not hasattr(xp, '_CopyMode') :
         supports_copy_false = False
-    elif library == 'cupy':
+    elif library in ['cupy', 'dask.array']:
         supports_copy_false = False
     else:
         supports_copy_false = True
@@ -133,14 +133,18 @@ def test_asarray_copy(library):
     assert all(b[0] == 0)
 
     a = asarray([1.0], dtype=xp.float32)
+    assert a.dtype == xp.float32
     b = asarray(a, dtype=xp.float64, copy=None)
     assert is_lib_func(b)
+    assert b.dtype == xp.float64
     a[0] = 0.0
     assert all(b[0] == 1.0)
 
     a = asarray([1.0], dtype=xp.float64)
+    assert a.dtype == xp.float64
     b = asarray(a, dtype=xp.float64, copy=None)
     assert is_lib_func(b)
+    assert b.dtype == xp.float64
     a[0] = 0.0
     assert all(b[0] == 0.0)