Support CuPy-backed arrays in DaskManager#11383
Conversation
The "meta" argument passed to dask.array.from_array should not be hardcoded to just `numpy.ndarray`, but allow for `cupy.ndarray` too.
Not sure how to type-hint np | cp | ??, so just use Any for output of get_array_namespace.
| # lazily loaded backend array classes should use NumPy array operations. | ||
| kwargs["meta"] = np.ndarray | ||
| # lazily loaded backend array classes should use NumPy or CuPy array operations. | ||
| xp = get_array_namespace(data.get_duck_array()) |
There was a problem hiding this comment.
we probably need to add some API allow getting the underlying array type / library without actually fetching data. Something like a data.get_array_namespace() or data.get_meta()? Not sure how easy it would be to implement that, though.
(I think this is what causes the tests to fail)
There was a problem hiding this comment.
Ideally we could just call xp = data.__array_namespace__() following the Array API spec - https://data-apis.org/array-api/2025.12/API_specification/generated/array_api.array.__array_namespace__.html, and it would propagate through all the subclassed layers to get the underlying array namespace (numpy or cupy). I thought this would work by putting it into the NDArrayMixin (b77cc57), but that breaks a lot of the lazy repls...
Might need to think this through a bit more. Wondering if there needs to be a cached .__cached_array_namespace__ attribute of some sort to work with the lazy objects...
Centralize retrieving of the __array_namespace__ through several subclassed layers, to avoid having to go through `.get_duck_array()`. Need to put `from xarray.compat.array_api_compat import get_array_namespace` import within the method to avoid circular import. Also type-hinted output of `get_array_namespace` as ModuleType following numpy/numpy#20719.
2 participants
Description
The default Dask ChunkManagerEntrypoint appears to be hardcoded to return NumPy arrays by default, even if the underlying arrays are CuPy arrays
TODO:
dask.array.from_arraythat was hardcoded tonp.ndarray(since 2019 in Ensure explicitly indexed arrays are preserved #3027), allowingcupy.ndarrayinstead when detected.compute()on an xarray.Dataset backed by CuPy arrays don't get coerced to NumPy arrays.Probably needs #11381 to be merged first. Part of resolving xarray-contrib/cupy-xarray#81 (comment).
Checklist
whats-new.rstapi.rstAI Disclosure