Skip to content

Support CuPy-backed arrays in DaskManager #11383

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
weiji14 wants to merge 11 commits into
base: main
Choose a base branch
from
+39 −7
Draft

Support CuPy-backed arrays in DaskManager #11383

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
1a03e39
Support CuPy-backed arrays in DaskManager.from_array method
weiji14 Jun 13, 2026
4465ea4
Fix no-untyped-call
weiji14 Jun 13, 2026
16009ee
Merge branch 'main' into dask_with_cupy
weiji14 Jun 14, 2026
b77cc57
Implement __array_namespace__ method on NDArrayMixin
weiji14 Jun 15, 2026
558f171
Merge branch 'main' into dask_with_cupy
weiji14 Jun 17, 2026
286053f
Fix repls for LazilyIndexedArray, MemoryCachedArray and IndexingAdapter
weiji14 Jun 17, 2026
551abe8
Fix mypy unused-ignore
weiji14 Jun 17, 2026
ba652b8
Merge branch 'main' into dask_with_cupy
weiji14 Jun 18, 2026
7132f73
Put __array_namespace__ on ImplicitToExplicitIndexingAdapter instead
weiji14 Jun 18, 2026
76721fe
Merge branch 'main' into dask_with_cupy
weiji14 Jun 26, 2026
39909d5
Test to check duck_arrays that are chunked then computed return the same
weiji14 Jun 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion xarray/compat/array_api_compat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from types import ModuleType

import numpy as np

from xarray.namedarray.pycompat import array_type
Expand Down Expand Up @@ -46,7 +48,7 @@ def result_type(*arrays_and_dtypes, xp) -> np.dtype:
return _future_array_api_result_type(*arrays_and_dtypes, xp=xp)


def get_array_namespace(*values):
def get_array_namespace(*values) -> ModuleType:
def _get_single_namespace(x):
if hasattr(x, "__array_namespace__"):
return x.__array_namespace__()
Expand Down
6 changes: 5 additions & 1 deletion xarray/core/formatting.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
from xarray.core.indexing import (
BasicIndexer,
ExplicitlyIndexed,
IndexingAdapter,
LazilyIndexedArray,
MemoryCachedArray,
)
from xarray.core.options import OPTIONS, _get_boolean_with_default
Expand Down Expand Up @@ -700,7 +702,9 @@ def short_data_repr(array):

if isinstance(array, np.ndarray):
return short_array_repr(array)
elif is_duck_array(internal_data):
elif not isinstance(
internal_data, (LazilyIndexedArray, MemoryCachedArray, IndexingAdapter)

@dcherian dcherian Jun 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
internal_data, (LazilyIndexedArray, MemoryCachedArray, IndexingAdapter)
internal_data, ExplicitlyIndexed

is usually the way to do it

@weiji14 weiji14 Jun 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Excellent hint, but actually, not needed anymore. I talked to @keewis on the Pangeo call just now, and figured out the right-level of abstraction is to put the __array_namespace__ method in ImplicitToExplicitIndexingAdapter instead of NDArrayMixin (commit 7132f73). So we won't need to change this repr stuff anymore.

@weiji14 weiji14 Jun 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh wait, commit 7132f73 might not be correct either 😅 I'll actually need to use your suggestion here, and really should write a proper test for this first...

@dcherian dcherian Jun 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here github still has the version of the notebook:L https://github.com/dcherian/cupy-xarray/blob/kvikio-entrypoint/docs/source/kvikio.ipynb

Doesn't work: Chunk with dask is the section

And this is the issue:
https://github.com/dask/dask/blob/cbbac094376f03d4d3fb933a6f115e3416981263/dask/array/core.py#L120-L144

Today, that np.asarray is what triggers a concrete read in to memory. If you advertise array-namespace and thus satisfy is_arraylike that read isn't triggered. we could add a patch to dask to hardcode support for xarray classes and call .get_duck_array (i bet this is fine).

My bigger concern is that advertising array api compliance will break random downstream code in a similar way somehow (but maybe not).

@weiji14 weiji14 Jun 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Funny, I missed this comment, but stumbled on those exact same cached notebook cells just now while revisiting xarray-contrib/cupy-xarray#70 😆 That was the aha moment I needed to write a sane test afterwards :)

Today, that np.asarray is what triggers a concrete read in to memory. If you advertise array-namespace and thus satisfy is_arraylike that read isn't triggered. we could add a patch to dask to hardcode support for xarray classes and call .get_duck_array (i bet this is fine).

My bigger concern is that advertising array api compliance will break random downstream code in a similar way somehow (but maybe not).

I feel like we should lean towards array api compliance rather than patching dask, so that other similar chunking libraries (e.g. cubed?) can rely on the same standards rather than hardcoding? But yeah, not sure how much code this will break 😬

) and is_duck_array(internal_data):
return limit_lines(repr(array.data), limit=40)
elif getattr(array, "_in_memory", None):
return short_array_repr(array)
Expand Down
2 changes: 1 addition & 1 deletion xarray/core/indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -687,7 +687,7 @@ def __array__(
else:
return np.asarray(self.get_duck_array(), dtype=dtype)

def get_duck_array(self):
def get_duck_array(self) -> duckarray:
return self.array.get_duck_array()

def __getitem__(self, key: Any):
Expand Down
5 changes: 5 additions & 0 deletions xarray/core/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -704,6 +704,11 @@ def dtype(self: Any) -> np.dtype:
def shape(self: Any) -> tuple[int, ...]:
return self.array.shape

def __array_namespace__(self: Any) -> ModuleType:
from xarray.compat.array_api_compat import get_array_namespace

return get_array_namespace(self.array)

def __getitem__(self: Any, key):
return self.array[key]

Expand Down
2 changes: 1 addition & 1 deletion xarray/namedarray/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -843,7 +843,7 @@ def chunk(
# Using OuterIndexer is a pragmatic choice: dask does not yet handle
# different indexing types in an explicit way:
# https://github.com/dask/dask/issues/2883
ndata = ImplicitToExplicitIndexingAdapter(data_old, OuterIndexer) # type: ignore[assignment]
ndata = ImplicitToExplicitIndexingAdapter(data_old, OuterIndexer)

if is_dict_like(chunks):
chunks = tuple(starmap(chunks.get, enumerate(ndata.shape)))
Expand Down
5 changes: 3 additions & 2 deletions xarray/namedarray/daskmanager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,9 @@ def from_array(
import dask.array as da

if isinstance(data, ImplicitToExplicitIndexingAdapter):
# 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 = data.__array_namespace__()
kwargs["meta"] = xp.ndarray

return da.from_array(
data,
Expand Down
2 changes: 1 addition & 1 deletion xarray/namedarray/pycompat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ def to_duck_array(data: Any, **kwargs: dict[str, Any]) -> duckarray[_ShapeType,
return loaded_data

if isinstance(data, ExplicitlyIndexed | ImplicitToExplicitIndexingAdapter):
return data.get_duck_array() # type: ignore[no-untyped-call, no-any-return]
return data.get_duck_array()
elif is_duck_array(data):
return data
else:
Expand Down
Loading