Open
Description
The first line works, the second raises an exception
import numpy as np
import xarray as xr
import cupy_xarray
xr.DataArray([1, 2, np.nan]).chunk(dim_0=1).as_cupy().sum().compute()
xr.DataArray([1, 2, np.nan]).chunk(dim_0=1).as_cupy().sum(min_count=1).compute()
xarray.DataArray'asarray-75d4a7ce4023e88c4c5563214cb235b4'
array(3.)
Coordinates: (0)
Indexes: (0)
Attributes: (0)
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[5], line 6
3 import cupy_xarray
5 xr.DataArray([1, 2, np.nan]).chunk(dim_0=1).as_cupy().sum().compute()
----> 6 xr.DataArray([1, 2, np.nan]).chunk(dim_0=1).as_cupy().sum(min_count=1).compute()
File [~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/xarray/core/dataarray.py:1179](http://localhost:8888/lab/tree/icec/seaice/nb/~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/xarray/core/dataarray.py#line=1178), in DataArray.compute(self, **kwargs)
1154 """Manually trigger loading of this array's data from disk or a
1155 remote source into memory and return a new array.
1156
(...)
1176 dask.compute
1177 """
1178 new = self.copy(deep=False)
-> 1179 return new.load(**kwargs)
File [~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/xarray/core/dataarray.py:1147](http://localhost:8888/lab/tree/icec/seaice/nb/~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/xarray/core/dataarray.py#line=1146), in DataArray.load(self, **kwargs)
1127 def load(self, **kwargs) -> Self:
1128 """Manually trigger loading of this array's data from disk or a
1129 remote source into memory and return this array.
1130
(...)
1145 dask.compute
1146 """
-> 1147 ds = self._to_temp_dataset().load(**kwargs)
1148 new = self._from_temp_dataset(ds)
1149 self._variable = new._variable
File [~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/xarray/core/dataset.py:863](http://localhost:8888/lab/tree/icec/seaice/nb/~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/xarray/core/dataset.py#line=862), in Dataset.load(self, **kwargs)
860 chunkmanager = get_chunked_array_type(*lazy_data.values())
862 # evaluate all the chunked arrays simultaneously
--> 863 evaluated_data: tuple[np.ndarray[Any, Any], ...] = chunkmanager.compute(
864 *lazy_data.values(), **kwargs
865 )
867 for k, data in zip(lazy_data, evaluated_data):
868 self.variables[k].data = data
File [~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/xarray/namedarray/daskmanager.py:86](http://localhost:8888/lab/tree/icec/seaice/nb/~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/xarray/namedarray/daskmanager.py#line=85), in DaskManager.compute(self, *data, **kwargs)
81 def compute(
82 self, *data: Any, **kwargs: Any
83 ) -> tuple[np.ndarray[Any, _DType_co], ...]:
84 from dask.array import compute
---> 86 return compute(*data, **kwargs)
File [~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/dask/base.py:662](http://localhost:8888/lab/tree/icec/seaice/nb/~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/dask/base.py#line=661), in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
659 postcomputes.append(x.__dask_postcompute__())
661 with shorten_traceback():
--> 662 results = schedule(dsk, keys, **kwargs)
664 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
File cupy[/_core/core.pyx:1717](http://localhost:8888/_core/core.pyx#line=1716), in cupy._core.core._ndarray_base.__array_function__()
File [~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/cupy/_sorting/search.py:211](http://localhost:8888/lab/tree/icec/seaice/nb/~/mambaforge/envs/cupy-seaice/lib/python3.12/site-packages/cupy/_sorting/search.py#line=210), in where(condition, x, y)
209 if fusion._is_fusing():
210 return fusion._call_ufunc(_where_ufunc, condition, x, y)
--> 211 return _where_ufunc(condition.astype('?'), x, y)
File cupy[/_core/_kernel.pyx:1286](http://localhost:8888/_core/_kernel.pyx#line=1285), in cupy._core._kernel.ufunc.__call__()
File cupy[/_core/_kernel.pyx:159](http://localhost:8888/_core/_kernel.pyx#line=158), in cupy._core._kernel._preprocess_args()
File cupy[/_core/_kernel.pyx:145](http://localhost:8888/_core/_kernel.pyx#line=144), in cupy._core._kernel._preprocess_arg()
TypeError: Unsupported type <class 'numpy.ndarray'>
Versions:
xr.__version__
np.__version__
cupy_xarray.__version__
'2024.6.0'
'1.26.4'
'0.1.3+9.g7fc3df5'
Same thing with numpy 2.0.0