@max-sixty Thanks! I'll really appreciate your help. I've tracked the possible source down to a dimension problem. I've tried to create a minimal example as follows using the current idxmax implementation from above. I copied only the dask related lines of code:

```python

create dask backed 3d array

darray = da.from_array(np.random.RandomState(0).randn(102030).reshape(10, 20, 30), chunks=(10, 20, 30), name='data_arr') array = xr.DataArray(darray, dims=["x", "y", 'z']) array = array.assign_coords({'x': (['x'], np.arange(10)), 'y': (['y'], np.arange(20)), 'z': (['z'], np.arange(30)), })

func=lambda x, args, kwargs: x.argmax(args, **kwargs) indx = func(array, dim='z', axis=None, keep_attrs=True, skipna=False) coordarray = array['z'] res = indx.copy( data=indx.data.map_blocks( lambda ind, coord: coord[(ind,)], coordarray, dtype=coordarray.dtype ) ) print(res)

the following line break breaks

print(res.compute())

using only 2dim array everything works as intended

array2d = array.sel(y=0, drop=True) indx = func(array2d, dim='z', axis=None, keep_attrs=True, skipna=False) coordarray = array['z'] res = indx.copy( data=indx.data.map_blocks( lambda ind, coord: coord[(ind,)], coordarray, dtype=coordarray.dtype ) ) print(res)

this works for two dim data

print(res.compute()) ```