BUG: astype(..., copy=True) doesn't copy on dask

Author: lithomas1

array_api_compat/dask/array/_aliases.py

@@ -39,7 +39,14 @@
 
 isdtype = get_xp(np)(_aliases.isdtype)
 unstack = get_xp(da)(_aliases.unstack)
-astype = _aliases.astype
+
+def astype(x: Array, dtype: Dtype, /, *, copy: bool = True) -> Array:
+    if not copy and dtype == x.dtype:
+        return x
+    # dask astype doesn't respect copy=True so copy
+    # manually via numpy
+    x = np.array(x, dtype=dtype, copy=copy)
+    return da.from_array(x)
 
 # Common aliases
 

tests/test_array_namespace.py

@@ -9,6 +9,7 @@
 
 import array_api_compat
 from array_api_compat import array_namespace
+import array_api_compat.numpy
 
 from ._helpers import import_, all_libraries, wrapped_libraries
 
@@ -22,6 +23,7 @@ def test_array_namespace(library, api_version, use_compat):
     if use_compat is True and library in {'array_api_strict', 'jax.numpy', 'sparse'}:
         pytest.raises(ValueError, lambda: array_namespace(array, use_compat=use_compat))
         return
+    print(use_compat)
     namespace = array_api_compat.array_namespace(array, api_version=api_version, use_compat=use_compat)
 
     if use_compat is False or use_compat is None and library not in wrapped_libraries:
@@ -36,6 +38,17 @@ def test_array_namespace(library, api_version, use_compat):
                 assert namespace == jax.experimental.array_api
         else:
             assert namespace == xp
+    elif use_compat is None:
+        if library == "dask.array":
+            # dask should always return wrapped version
+            # since dask.array is not array API compatible
+            assert namespace == array_api_compat.dask.array
+        elif library == "numpy":
+            assert namespace == array_api_compat.numpy
+        elif library == "torch":
+            assert namespace == array_api_compat.torch
+        else:
+            assert namespace == xp
     else:
         if library == "dask.array":
             assert namespace == array_api_compat.dask.array

tests/test_common.py

@@ -272,3 +272,16 @@ def test_asarray_copy(library):
         assert all(b[0] == 1.0)
     else:
         assert all(b[0] == 0.0)
+
+@pytest.mark.parametrize("library", wrapped_libraries)
+def test_astype_copy(library):
+    # array-api-tests currently doesn't check copy=True
+    # makes a copy when dtypes are the same
+    # so we check that here
+    xp = import_(library, wrapper=True)
+    a = xp.asarray([1])
+    b = xp.astype(a, a.dtype, copy=True)
+
+    a[0] = 10
+
+    assert b[0] == 1