issue_comments

Would using obj.fillna(0) not mess with argmax if for instance all the data is negative? Could fill with the min value instead?

Ah yes true. I was slightly anticipating e.g. filling with NaT if the dim was time-like, though time types are not something I am familiar with.

Ah yes both ways are working now, thanks. Just had a little play around with timings, and this seems like a reasonably quick way to achieve correct NaN behaviour:

```python def xr_idxmax(obj, dim): sig = ([(dim,), (dim,)], [()]) kwargs = {'axis': -1}

allna = obj.isnull().all(dim)

return apply_ufunc(gufunc_idxmax, obj.fillna(-np.inf), obj[dim],
                   signature=sig, kwargs=kwargs,
                   dask_array='allowed').where(~allna).fillna(np.nan)

```

i.e. originally replace all NaN values with -Inf, use the usual argmax, and remask the all-NaN values afterwards.

So I thought take was just the functional equivalent of fancy indexing - I spotted it in the dask api and assumed it would work but having tried it does indeed just raise a 'not implemented error'. Just as a note, with the map_blocks approach above take is working for some cases where x[inds, ] is not -- related to #1237?

Regarding edge cases: multiple maxes is presumably fine as long as user is aware it just takes the first. However, nanargmax is probably the actual desired function here, but looks like it will raise on all-nan slices. Would dropping these and then re-aligning be too much overhead?

Just as I am interested in having this functionality, and the new apply_ufunc is available, would something along these lines suffice?

```python from wherever import argmax, take # numpy or dask

def gufunc_idxmax(x, y, axis=None): indx = argmax(x, axis) return take(y, indx)

def idxmax(obj, dim): sig = ([(dim,), (dim,)], [()]) kwargs = {'axis': -1} return apply_ufunc(gufunc_idxmin, obj, obj[dim], signature=sig, kwargs=kwargs, dask_array='allowed') ```