Using apply_ufunc with dask='allowed' for a function with xarray operations

[tlsw231]

I have a function that has xarray-based operations (e.g., rolling) and I am trying to call it with apply_ufunc. I get this error:
'Array' object has no attribute 'rolling'. Here is the code:

import xarray as xr               
import cftime
import numpy as np

times=xr.cftime_range(start='2022-01',end='2023-12',freq='1MS')
(nt,ny,nx)=(len(times),2,3)
dummy = xr.DataArray(np.random.random([nt,ny,nx]),dims=['time','y','x'],
                    coords=[times,range(ny),range(nx)] )

def rolling_max(array):
    return array.rolling(time=2).mean().groupby('time.year').max()

Here is the call to apply_ufunc that throws the error:

xr.apply_ufunc(rolling_max,dummy.chunk(dict(time=-1)),input_core_dims=[['time']],
               output_core_dims=[['year']],dask='allowed',exclude_dims=set(("time",)) )

[alessioarena]

Hi, if you are applying an xarray operation you should not need to apply it using apply_ufunc. You can assume that all xarray operations will work with dask.
In that specific case:

dummy.rolling(time=2).mean().groupby('time.year').max()

Should work fine.

The only change you should make to have that being a dask operation is declaring your dataset as

import dask.array as da
dummy = xr.DataArray(da.random.random([nt,ny,nx]),dims=['time','y','x'],
                    coords=[times,range(ny),range(nx)] )

[dcherian]

@tlsw231 See https://docs.xarray.dev/en/stable/user-guide/dask.html#automatic-parallelization-with-apply-ufunc-and-map-blocks for more.