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 10 commits into
base: main
Choose a base branch
from
+14 −6
Draft

Support CuPy-backed arrays in DaskManager #11383

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
10 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
3980fcb
Merge branch 'main' into dask_with_cupy
keewis Jun 18, 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 typing import Any

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) -> Any:
def _get_single_namespace(x):
if hasattr(x, "__array_namespace__"):
return x.__array_namespace__()
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
6 changes: 4 additions & 2 deletions xarray/namedarray/daskmanager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import numpy as np

from xarray.compat.array_api_compat import get_array_namespace
from xarray.core.indexing import ImplicitToExplicitIndexingAdapter
from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint, T_ChunkedArray
from xarray.namedarray.utils import is_duck_dask_array, module_available
Expand Down Expand Up @@ -68,8 +69,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 = get_array_namespace(data.get_duck_array())

@keewis keewis Jun 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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)

@weiji14 weiji14 Jun 15, 2026

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.

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...

@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.

i went down this road a while ago. the problem is that we need our lazy arrays coerced to an in-memory type at some point. If they advertise __array_namespace__ they can treated as an in-memory type (deep within dask), and nothing works. Did you also run in to this problem?

@weiji14 weiji14 Jun 17, 2026

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.

I think I've went pretty deep, but not deep into dask internals (yet?). If it ends up needing changes in dask, I'm just gonna push on findind a way to remove dask entirely - #9038 (comment).

Right now I think I've solved the repl issues by changing some of the logic in formatting.py. Still need to work my way through some other logic that have been basing their logic around duck arrays / __array_namespace__ but shouldn't...

@weiji14 weiji14 Jun 17, 2026

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.

Actually:

[ ] Ensure calling .compute() on an xarray.Dataset backed by CuPy arrays don't get coerced to NumPy arrays.

I did realize that even with the cupy meta fixes here, the Dask-backed arrays still get loaded into NumPy instead. So maybe that's what you're referring to? That there's no way to have a fully CuPy-only array pipeline without changing some internals of Dask?

@dcherian dcherian Jun 18, 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.

Is there an example I can run without an nvidia gpu?

@weiji14 weiji14 Jun 18, 2026

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.

Unfortunately no, I only know how to reproduce this on this cupy-xarray PR - https://github.com/xarray-contrib/cupy-xarray/pull/81/changes#diff-b866be1141ec4295c6a2ef9b8effab4a58d6dacd262918d0db4f56d6d7927818:

ds: xr.Dataset = xr.open_mfdataset(
    paths=[
        "https://github.com/developmentseed/titiler/raw/1.2.0/src/titiler/mosaic/tests/fixtures/B01.tif",
        "https://github.com/developmentseed/titiler/raw/1.2.0/src/titiler/mosaic/tests/fixtures/B09.tif",
    ],
    engine="cog3pio",
    concat_dim="band",
    combine="nested",
    device_id=None,
)
ds.raster.load()  # inplace load, requires https://github.com/pydata/xarray/pull/11381
# assert isinstance(
#     ds.raster, cp.ndarray  # TODO wait for https://github.com/pydata/xarray/pull/11383 ?
# )

But I should figure out a good duck array test somehow, and will let you know if it can be reproduced without a GPU.

@dcherian dcherian Jun 18, 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.

I am pretty sure the get_array_namespace problem will reproduce with any duck array honestly.

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