[WIP] ENH: dask+cupy, dask+sparse etc. namespaces

crusaderky · crusaderky · commit 7f39b039aab4 · 2025-03-06T00:32:18.000Z
diff --git a/array_api_compat/common/_helpers.py b/array_api_compat/common/_helpers.py
@@ -368,7 +368,11 @@ def is_dask_namespace(xp) -> bool:
     is_pydata_sparse_namespace
     is_array_api_strict_namespace
     """
-    return xp.__name__ in {'dask.array', _compat_module_name() + '.dask.array'}
+    names = {'dask.array', _compat_module_name() + '.dask.array'}
+    return (
+        xp.__name__ in names
+        or any(xp.__name__.startswith(name + '.') for name in names)
+    )
 
 
 def is_jax_namespace(xp) -> bool:
@@ -541,8 +545,10 @@ def your_function(x, y):
         elif is_dask_array(x):
             if _use_compat:
                 _check_api_version(api_version)
-                from ..dask import array as dask_namespace
-                namespaces.add(dask_namespace)
+                from ..dask.array import wrap_namespace
+                mxp = array_namespace(x._meta, use_compat=False)
+                xp = wrap_namespace(mxp)
+                namespaces.add(xp)
             else:
                 import dask.array as da
                 namespaces.add(da)
diff --git a/array_api_compat/dask/array/__init__.py b/array_api_compat/dask/array/__init__.py
@@ -2,6 +2,7 @@
 
 # These imports may overwrite names from the import * above.
 from ._aliases import * # noqa: F403
+from ._meta import wrap_namespace # noqa: F401
 
 __array_api_version__ = '2024.12'
 
diff --git a/array_api_compat/dask/array/_aliases.py b/array_api_compat/dask/array/_aliases.py
@@ -157,6 +157,7 @@ def asarray(
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
     copy: Optional[Union[bool, np._CopyMode]] = None,
+    like: Optional[Array] = None,
     **kwargs,
 ) -> Array:
     """
@@ -172,7 +173,11 @@ def asarray(
             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:
+            obj = obj.copy()
+        if like is not None:
+            obj = da.asarray(obj, like=like)
+        return obj
 
     if copy is False:
         raise NotImplementedError(
@@ -181,7 +186,11 @@ def asarray(
 
     # 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)
+    if like is not None:
+        mxp = array_namespace(like)
+        obj = mxp.asarray(obj, dtype=dtype, copy=True)
+    else:
+        obj = np.array(obj, dtype=dtype, copy=True)
     return da.from_array(obj)
 
 
diff --git a/array_api_compat/dask/array/_meta.py b/array_api_compat/dask/array/_meta.py
@@ -0,0 +1,49 @@
+import functools
+import sys
+import types
+
+from ...common import is_numpy_namespace
+
+__all__ = ['wrap_namespace']
+_all_ignore = ['functools', 'sys', 'types', 'is_numpy_namespace']
+
+
+def wrap_namespace(xp):
+    """Create a bespoke Dask namespace that wraps around another namespace.
+
+    Parameters
+    ----------
+    xp : namespace
+        Namespace to be wrapped by Dask
+
+    Returns
+    -------
+    namespace :
+        A module object that duplicates array_api_compat.dask.array, with the
+        difference that all creation functions will create an array with the same
+        meta namespace as the input.
+    """
+    from .. import array as da_compat
+
+    if is_numpy_namespace(xp):
+        return da_compat
+
+    mod_name = f'{da_compat.__name__}.{xp.__name__}'
+    try:
+        return sys.modules[mod_name]
+    except KeyError:
+        pass
+
+    mod = types.ModuleType(mod_name)
+    sys.modules[mod_name] = mod
+
+    meta = xp.empty(())
+    for name, v in da_compat.__dict__.items():
+        if name.startswith('_'):
+            continue
+        if name in {'arange',  'asarray', 'empty', 'eye', 'from_dlpack',
+                    'full', 'linspace', 'ones', 'zeros'}:
+            v = functools.wraps(v)(functools.partial(v, like=meta))
+        setattr(mod, name, v)
+
+    return mod
