da.asarray should not materialize the graph

Author: crusaderky

array_api_compat/dask/array/_aliases.py

@@ -144,24 +144,23 @@ def asarray(
     See the corresponding documentation in the array library and/or the array API
     specification for more details.
     """
+    if isinstance(obj, da.Array):
+        if dtype is not None and dtype != obj.dtype:
+            if copy is False:
+                raise ValueError("Unable to avoid copy when changing dtype")
+            obj = obj.astype(dtype)
+        return obj.copy() if copy else obj
+
     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:
-            # copy=True to be uniform across dask < 2024.12 and >= 2024.12
-            # see https://github.com/dask/dask/pull/11524/
-            obj = np.array(obj, dtype=dtype, copy=True)
-            return da.from_array(obj)
-        return obj
-
-    return da.asarray(obj, dtype=dtype, **kwargs)
+        raise NotImplementedError(
+            "Unable to avoid copy when converting a non-dask object to dask"
+        )
+
+    # copy=None to be uniform across dask < 2024.12 and >= 2024.12
+    # see https://github.com/dask/dask/pull/11524/
+    obj = np.array(obj, dtype=dtype, copy=True)
+    return da.from_array(obj)
+
 
 from dask.array import (
     # Element wise aliases

tests/test_all.py

@@ -40,5 +40,7 @@ def test_all(library):
         all_names = module.__all__
 
         if set(dir_names) != set(all_names):
-            assert set(dir_names) - set(all_names) == set(), f"Some dir() names not included in __all__ for {mod_name}"
-            assert set(all_names) - set(dir_names) == set(), f"Some __all__ names not in dir() for {mod_name}"
+            extra_dir = set(dir_names) - set(all_names)
+            extra_all = set(all_names) - set(dir_names)
+            assert not extra_dir, f"Some dir() names not included in __all__ for {mod_name}: {extra_dir}"
+            assert not extra_all, f"Some __all__ names not in dir() for {mod_name}: {extra_all}"

tests/test_common.py

@@ -226,11 +226,17 @@ def test_asarray_copy(library):
     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 in ['cupy', 'dask.array']:
-        supports_copy_false = False
+        supports_copy_false_other_ns = False
+        supports_copy_false_same_ns = False
+    elif library == 'cupy':
+        supports_copy_false_other_ns = False
+        supports_copy_false_same_ns = False
+    elif library == 'dask.array':
+        supports_copy_false_other_ns = False
+        supports_copy_false_same_ns = True
     else:
-        supports_copy_false = True
+        supports_copy_false_other_ns = True
+        supports_copy_false_same_ns = True
 
     a = asarray([1])
     b = asarray(a, copy=True)
@@ -240,7 +246,7 @@ def test_asarray_copy(library):
     assert all(a[0] == 0)
 
     a = asarray([1])
-    if supports_copy_false:
+    if supports_copy_false_same_ns:
         b = asarray(a, copy=False)
         assert is_lib_func(b)
         a[0] = 0
@@ -249,7 +255,7 @@ def test_asarray_copy(library):
         pytest.raises(NotImplementedError, lambda: asarray(a, copy=False))
 
     a = asarray([1])
-    if supports_copy_false:
+    if supports_copy_false_same_ns:
         pytest.raises(ValueError, lambda: asarray(a, copy=False,
                                                   dtype=xp.float64))
     else:
@@ -281,7 +287,7 @@ def test_asarray_copy(library):
     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:
+        if supports_copy_false_other_ns:
             pytest.raises(ValueError, lambda: asarray(obj, copy=False))
         else:
             pytest.raises(NotImplementedError, lambda: asarray(obj, copy=False))
@@ -294,7 +300,7 @@ def test_asarray_copy(library):
     assert all(b[0] == 1.0)
 
     a = array.array('f', [1.0])
-    if supports_copy_false:
+    if supports_copy_false_other_ns:
         b = asarray(a, copy=False)
         assert is_lib_func(b)
         a[0] = 0.0

tests/test_dask.py

@@ -0,0 +1,57 @@
+import dask
+import numpy as np
+import pytest
+import dask.array as da
+
+from array_api_compat import array_namespace
+
+
+@pytest.fixture
+def no_compute():
+    """
+    Cause the test to raise if at any point anything calls compute() or persist(),
+    e.g. as it can be triggered implicitly by __bool__, __array__, etc.
+    """
+    def get(dsk, *args, **kwargs):
+        raise AssertionError("Called compute() or persist()")
+    
+    with dask.config.set(scheduler=get):
+        yield
+
+
+def test_no_compute(no_compute):
+    """Test the no_compute fixture"""
+    a = da.asarray(True)
+    with pytest.raises(AssertionError, match="Called compute"):
+        bool(a)
+
+
+def test_asarray_no_compute(no_compute):
+    a = da.arange(10)
+    xp = array_namespace(a)  # wrap in array_api_compat.dask.array
+
+    xp.asarray(a)
+    xp.asarray(a, dtype=np.int16)
+    xp.asarray(a, dtype=a.dtype)
+    xp.asarray(a, copy=True)
+    xp.asarray(a, copy=True, dtype=np.int16)
+    xp.asarray(a, copy=True, dtype=a.dtype)
+    xp.asarray(a, copy=False)
+    xp.asarray(a, copy=False, dtype=a.dtype)
+
+
+@pytest.mark.parametrize("copy", [True, False])
+def test_astype_no_compute(no_compute, copy):
+    a = da.arange(10)
+    xp = array_namespace(a)  # wrap in array_api_compat.dask.array
+    xp.astype(a, np.int16, copy=copy)
+    xp.astype(a, a.dtype, copy=copy)
+
+
+def test_clip_no_compute(no_compute):
+    a = da.arange(10)
+    xp = array_namespace(a)  # wrap in array_api_compat.dask.array
+
+    xp.clip(a)
+    xp.clip(a, 1)
+    xp.clip(a, 1, 8)