-
-
Notifications
You must be signed in to change notification settings - Fork 1.2k
coords.update() converts CuPy arrays to NumPy for dimension coordinates #11251
Copy link
Copy link
Open
Labels
needs triageIssue that has not been reviewed by xarray team memberIssue that has not been reviewed by xarray team member
Description
What is your issue?
When assigning CuPy arrays as coordinates to a dataset using DataSet.coords.update(), the resulting coordinate data is converted to a NumPy array.
As a result, GPU-native workflows that are dependent on coordinates (e.g. grid interpolation) become CPU-bound or require additional data transfers to the GPU.
Minimal example:
data_cp = cp.arange(6)
ds = xr.Dataset({
"var": (("x"), data_cp)
})
x_cp = cp.array([0,1,2,3,4,5])
print(type(x_cp)) # <class 'cupy.ndarray'>
ds.coords.update({"x": x_cp.copy()})
print(type(ds.coords["x"].data)) # <class 'numpy.ndarray'>xarray: 2026.2.0
Question
Is this conversion intended, or should CuPy-backed coordinates be supported?
Reactions are currently unavailable
Metadata
Metadata
Assignees
Labels
needs triageIssue that has not been reviewed by xarray team memberIssue that has not been reviewed by xarray team member