`coords.update()` converts CuPy arrays to NumPy for dimension coordinates

[Dobbermann2]

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?

[welcome]

Thanks for opening your first issue here at xarray! Be sure to follow the issue template!
If you have an idea for a solution, we would really welcome a Pull Request with proposed changes.
See the Contributing Guide for more.
It may take us a while to respond here, but we really value your contribution. Contributors like you help make xarray better.
Thank you!

[headtr1ck]

Problem seems to be here: https://github.com/pydata/xarray/blob/main/xarray/core/coordinates.py#L1280
Didn't dig deeper than that though.

[Dobbermann2]

This appears to be a requirement for using pandas indexing in xarray. Coordinate data arrays are explicitly converted to NumPy arrays before being passed to the pandas Index constructor, since pandas expects strict NumPy arrays.

xarray/xarray/core/indexes.py

Line 543 in 0646e8d

values = to_numpy(array)

[VeckoTheGecko]

since pandas expects strict NumPy arrays.

I don't think this is the case - looking the Pandas API docs for Index it mentions that the input needs to be 'array-like' and it also exposes the .array variable

.array
The ExtensionArray of the data backing this Index.

This property provides direct access to the underlying array data of an Index without requiring conversion to a NumPy array. It returns an ExtensionArray, which is the native storage format for pandas extension dtypes.

This might just be part of the codebase in xarray that hasn't been updated to enable different types of arrays (which, from my POVm, would make sense since we haven't had a lot of testing/users in this area)

[Dobbermann2]

looking the Pandas API docs for Index it mentions that the input needs to be 'array-like'

You can indeed pass an 'array-like' to the pd.Index constructor, but then it does an explicit conversion to NumPy
https://github.com/pandas-dev/pandas/blob/1bba48f900275af485e457b056f3e4fc82b38f94/pandas/core/indexes/base.py#L537

This np.asarray() call actually raises an exception when a CuPy array is used as input, which is probably why xarray converts to NumPy beforehand.

[keewis]

if pandas indexes are the problem, you can disable that using the Coordinates constructor:

ds.coords.update(xr.Coordinates({"x": x_cp.copy()}, indexes={}))

(passing in a bare dict is the equivalent of xr.Coordinates(new_coords), which will create indexes on coordinates named like a dimension)

[Dobbermann2]

if pandas indexes are the problem, you can disable that using the Coordinates constructor:

ds.coords.update(xr.Coordinates({"x": x_cp.copy()}, indexes={}))
(passing in a bare dict is the equivalent of xr.Coordinates(new_coords), which will create indexes on coordinates named like a dimension)

That seems to be working! Thank you so much!

[keewis]

it might be interesting to see if cudf has a data structure similar to pandas.Index that we could reuse. If so, that discussion should be continued on cupy-xarray's issue tracker (cc @weiji14), because that's where the index implementation would most likely end up living.

[weiji14]

Stepping back a bit, is it necessary to have the coordinates in GPU memory? I'm assuming the grid interpolation itself is working on the matrix data, and the coordinates don't necessarily need to be in GPU memory?

Related: #10079

[keewis]

depends on what you do. If it's value selection / indexing or something similar it's probably fine to have coordinates in host memory, but if you're trying to compute the interpolation itself on the GPU you do need the coordinates in device memory (since in that case the coordinate values interact with the values of the data variables).