Repository: dask/dask

Issue:
dask.array.argmin fails on 3D input
```
In [29]: da.from_array(np.random.randn(2, 3, 4), chunks=(2, 3, 4)).argmin(axis=0).compute()
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-29-b25ac117010d> in <module>()
----> 1 da.from_array(np.random.randn(2, 3, 4), chunks=(2, 3, 4)).argmin(axis=0).compute()

/Users/shoyer/conda/envs/xarray-dev/lib/python3.5/site-packages/dask/base.py in compute(self, **kwargs)
     35
     36     def compute(self, **kwargs):
---> 37         return compute(self, **kwargs)[0]
     38
     39     @classmethod

/Users/shoyer/conda/envs/xarray-dev/lib/python3.5/site-packages/dask/base.py in compute(*args, **kwargs)
    108                 for opt, val in groups.items()])
    109     keys = [var._keys() for var in variables]
--> 110     results = get(dsk, keys, **kwargs)
    111
    112     results_iter = iter(results)

/Users/shoyer/conda/envs/xarray-dev/lib/python3.5/site-packages/dask/threaded.py in get(dsk, result, cache, num_workers, **kwargs)
     55     results = get_async(pool.apply_async, len(pool._pool), dsk, result,
     56                         cache=cache, queue=queue, get_id=_thread_get_id,
---> 57                         **kwargs)
     58
     59     return results

/Users/shoyer/conda/envs/xarray-dev/lib/python3.5/site-packages/dask/async.py in get_async(apply_async, num_workers, dsk, result, cache, queue, get_id, raise_on_exception, rerun_exceptions_locally, callbacks, **kwargs)
    479                 _execute_task(task, data)  # Re-execute locally
    480             else:
--> 481                 raise(remote_exception(res, tb))
    482         state['cache'][key] = res
    483         finish_task(dsk, key, state, results, keyorder.get)

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (3,4) (4,1) (1,3)

Traceback
---------
  File "/Users/shoyer/conda/envs/xarray-dev/lib/python3.5/site-packages/dask/async.py", line 264, in execute_task
    result = _execute_task(task, data)
  File "/Users/shoyer/conda/envs/xarray-dev/lib/python3.5/site-packages/dask/async.py", line 246, in _execute_task
    return func(*args2)
  File "/Users/shoyer/conda/envs/xarray-dev/lib/python3.5/site-packages/toolz/functoolz.py", line 381, in __call__
    ret = f(ret)
  File "/Users/shoyer/conda/envs/xarray-dev/lib/python3.5/site-packages/dask/array/reductions.py", line 450, in arg_agg
    return _arg_combine(data, axis, argfunc)[0]
  File "/Users/shoyer/conda/envs/xarray-dev/lib/python3.5/site-packages/dask/array/reductions.py", line 416, in _arg_combine
    arg = (arg + offsets)[tuple(inds)]
```
This was reported in xarray: https://github.com/pydata/xarray/issues/759

Discussion:

Patch:
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index 3f05e1628..a29003285 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -1,12 +1,12 @@
 from __future__ import absolute_import, division, print_function
 
 from functools import partial, wraps
-from itertools import product
+from itertools import product, repeat
 from math import factorial, log, ceil
 import operator
 
 import numpy as np
-from toolz import compose, partition_all, merge, get
+from toolz import compose, partition_all, merge, get, accumulate, pluck
 
 from . import chunk
 from .core import _concatenate2, Array, atop, sqrt, lol_tuples
@@ -34,6 +34,22 @@ def reduction(x, chunk, aggregate, axis=None, keepdims=None, dtype=None,
     if dtype and 'dtype' in getargspec(aggregate).args:
         aggregate = partial(aggregate, dtype=dtype)
 
+    # Map chunk across all blocks
+    inds = tuple(range(x.ndim))
+    tmp = atop(partial(chunk, axis=axis, keepdims=True), inds, x, inds)
+    tmp._chunks = tuple((1,)*len(c) if i in axis else c for (i, c)
+                        in enumerate(tmp.chunks))
+
+    return _tree_reduce(tmp, aggregate, axis, keepdims, dtype, split_every,
+                       combine)
+
+
+def _tree_reduce(x, aggregate, axis, keepdims, dtype, split_every=None,
+                combine=None):
+    """Perform the tree reduction step of a reduction.
+
+    Lower level, users should use ``reduction`` or ``arg_reduction`` directly.
+    """
     # Normalize split_every
     split_every = split_every or _globals.get('split_every', 4)
     if isinstance(split_every, dict):
@@ -44,24 +60,18 @@ def reduction(x, chunk, aggregate, axis=None, keepdims=None, dtype=None,
     else:
         split_every = dict((k, v) for (k, v) in enumerate(x.numblocks) if k in axis)
 
-    # Map chunk across all blocks
-    inds = tuple(range(x.ndim))
-    tmp = atop(partial(chunk, axis=axis, keepdims=True), inds, x, inds)
-    tmp._chunks = tuple((1,)*len(c) if i in axis else c for (i, c)
-                        in enumerate(tmp.chunks))
-
     # Reduce across intermediates
     depth = 1
-    for i, n in enumerate(tmp.numblocks):
+    for i, n in enumerate(x.numblocks):
         if i in split_every and split_every[i] != 1:
             depth = int(builtins.max(depth, ceil(log(n, split_every[i]))))
     func = compose(partial(combine or aggregate, axis=axis, keepdims=True),
                    partial(_concatenate2, axes=axis))
     for i in range(depth - 1):
-        tmp = partial_reduce(func, tmp, split_every, True, None)
+        x = partial_reduce(func, x, split_every, True, None)
     func = compose(partial(aggregate, axis=axis, keepdims=keepdims),
                    partial(_concatenate2, axes=axis))
-    return partial_reduce(func, tmp, split_every, keepdims=keepdims,
+    return partial_reduce(func, x, split_every, keepdims=keepdims,
                           dtype=dtype)
 
 
@@ -403,71 +413,130 @@ def vnorm(a, ord=None, axis=None, dtype=None, keepdims=False, split_every=None):
                    split_every=split_every)**(1./ord)
 
 
-def _arg_combine(data, axis, argfunc):
+def _arg_combine(data, axis, argfunc, keepdims=False):
     """Merge intermediate results from ``arg_*`` functions"""
+    axis = None if len(axis) == data.ndim or data.ndim == 1 else axis[0]
     vals = data['vals']
     arg = data['arg']
-    ns = data['n']
-    args = argfunc(vals, axis=axis)
-    offsets = np.roll(np.cumsum(ns, axis=axis), 1, axis)
-    offsets[tuple(slice(None) if i != axis else 0 for i in range(ns.ndim))] = 0
-    inds = list(reversed(np.meshgrid(*map(np.arange, args.shape), sparse=True)))
-    inds.insert(axis, args)
-
-    arg = (arg + offsets)[tuple(inds)]
-    vals = vals[tuple(inds)]
-    n = ns.sum(axis=axis).take(0, 0)
-    return arg, vals, n
-
-
-def arg_chunk(func, argfunc, x, axis=None, **kwargs):
-    axis = axis[0] if isinstance(axis, tuple) else axis
-    vals = func(x, axis=axis, keepdims=True)
-    arg = argfunc(x, axis=axis, keepdims=True)
+    if axis is None:
+        local_args = argfunc(vals, axis=axis, keepdims=keepdims)
+        vals = vals.ravel()[local_args]
+        arg = arg.ravel()[local_args]
+    else:
+        local_args = argfunc(vals, axis=axis)
+        inds = np.ogrid[tuple(map(slice, local_args.shape))]
+        inds.insert(axis, local_args)
+        vals = vals[inds]
+        arg = arg[inds]
+        if keepdims:
+            vals = np.expand_dims(vals, axis)
+            arg = np.expand_dims(arg, axis)
+    return arg, vals
+
+
+def arg_chunk(func, argfunc, x, axis, offset_info):
+    arg_axis = None if len(axis) == x.ndim or x.ndim == 1 else axis[0]
+    vals = func(x, axis=arg_axis, keepdims=True)
+    arg = argfunc(x, axis=arg_axis, keepdims=True)
+    if arg_axis is None:
+        offset, total_shape = offset_info
+        ind = np.unravel_index(arg.ravel()[0], x.shape)
+        total_ind = tuple(o + i for (o, i) in zip(offset, ind))
+        arg[:] = np.ravel_multi_index(total_ind, total_shape)
+    else:
+        arg += offset_info
+
     result = np.empty(shape=vals.shape, dtype=[('vals', vals.dtype),
-                                               ('arg', arg.dtype),
-                                               ('n', 'i8')])
+                                               ('arg', arg.dtype)])
     result['vals'] = vals
     result['arg'] = arg
-    result['n'] = x.shape[axis]
     return result
 
 
 def arg_combine(func, argfunc, data, axis=None, **kwargs):
-    axis = axis[0] if isinstance(axis, tuple) else axis
-    arg, vals, n = _arg_combine(data, axis, argfunc)
-    shape = tuple(s if i != axis else 1 for (i, s) in enumerate(data.shape))
-    result = np.empty(shape=shape, dtype=[('vals', vals.dtype),
-                                          ('arg', arg.dtype),
-                                          ('n', 'i8')])
-    result['vals'] = vals.reshape(shape)
-    result['arg'] = arg.reshape(shape)
-    result['n'] = n
+    arg, vals = _arg_combine(data, axis, argfunc, keepdims=True)
+    result = np.empty(shape=vals.shape, dtype=[('vals', vals.dtype),
+                                               ('arg', arg.dtype)])
+    result['vals'] = vals
+    result['arg'] = arg
     return result
 
 
 def arg_agg(func, argfunc, data, axis=None, **kwargs):
-    axis = axis[0] if isinstance(axis, tuple) else axis
-    return _arg_combine(data, axis, argfunc)[0]
+    return _arg_combine(data, axis, argfunc, keepdims=False)[0]
+
+
+def arg_reduction(x, chunk, combine, agg, axis=None, split_every=None):
+    """Generic function for argreduction.
+
+    Parameters
+    ----------
+    x : Array
+    chunk : callable
+        Partialed ``arg_chunk``.
+    combine : callable
+        Partialed ``arg_combine``.
+    agg : callable
+        Partialed ``arg_agg``.
+    axis : int, optional
+    split_every : int or dict, optional
+    """
+    if axis is None:
+        axis = tuple(range(x.ndim))
+        ravel = True
+    elif isinstance(axis, int):
+        if axis < 0:
+            axis += x.ndim
+        if axis < 0 or axis >= x.ndim:
+            raise ValueError("axis entry is out of bounds")
+        axis = (axis,)
+        ravel = x.ndim == 1
+    else:
+        raise TypeError("axis must be either `None` or int, "
+                        "got '{0}'".format(axis))
+
+    # Map chunk across all blocks
+    name = 'arg-reduce-chunk-{0}'.format(tokenize(chunk, axis))
+    old = x.name
+    keys = list(product(*map(range, x.numblocks)))
+    offsets = list(product(*(accumulate(operator.add, bd[:-1], 0)
+                             for bd in x.chunks)))
+    if ravel:
+        offset_info = zip(offsets, repeat(x.shape))
+    else:
+        offset_info = pluck(axis[0], offsets)
+
+    chunks = tuple((1,)*len(c) if i in axis else c for (i, c)
+                   in enumerate(x.chunks))
+    dsk = dict(((name,) + k, (chunk, (old,) + k, axis, off)) for (k, off)
+               in zip(keys, offset_info))
+    tmp = Array(merge(dsk, x.dask), name, chunks)
+    return _tree_reduce(tmp, agg, axis, False, np.int64, split_every, combine)
+
 
+def make_arg_reduction(func, argfunc):
+    """Create a argreduction callable.
 
-def arg_reduction(func, argfunc):
+    Parameters
+    ----------
+    func : callable
+        The reduction (e.g. ``min``)
+    argfunc : callable
+        The argreduction (e.g. ``argmin``)
+    """
     chunk = partial(arg_chunk, func, argfunc)
-    agg = partial(arg_agg, func, argfunc)
     combine = partial(arg_combine, func, argfunc)
+    agg = partial(arg_agg, func, argfunc)
     @wraps(argfunc)
-    def _(a, axis=None, split_every=None):
-        if axis < 0:
-            axis = a.ndim + axis
-        return reduction(a, chunk, agg, axis=axis, dtype='i8',
-                         split_every=split_every, combine=combine)
+    def _(x, axis=None, split_every=None):
+        return arg_reduction(x, chunk, combine, agg, axis, split_every)
     return _
 
 
-argmin = arg_reduction(chunk.min, chunk.argmin)
-argmax = arg_reduction(chunk.max, chunk.argmax)
-nanargmin = arg_reduction(chunk.nanmin, chunk.nanargmin)
-nanargmax = arg_reduction(chunk.nanmax, chunk.nanargmax)
+argmin = make_arg_reduction(chunk.min, chunk.argmin)
+argmax = make_arg_reduction(chunk.max, chunk.argmax)
+nanargmin = make_arg_reduction(chunk.nanmin, chunk.nanargmin)
+nanargmax = make_arg_reduction(chunk.nanmax, chunk.nanargmax)
 
 
 def cumreduction(func, binop, ident, x, axis, dtype=None):


Test patch:
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
index 9b13a98dc..7b734416f 100644
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -73,16 +73,6 @@ def test_reductions_1D(dtype):
     reduction_1d_test(da.nanmin, a, np.nanmin, x, False)
     reduction_1d_test(da.nanmax, a, np.nanmax, x, False)
 
-    assert eq(da.argmax(a, axis=0), np.argmax(x, axis=0))
-    assert eq(da.argmin(a, axis=0), np.argmin(x, axis=0))
-    assert eq(da.nanargmax(a, axis=0), np.nanargmax(x, axis=0))
-    assert eq(da.nanargmin(a, axis=0), np.nanargmin(x, axis=0))
-
-    assert eq(da.argmax(a, axis=0, split_every=2), np.argmax(x, axis=0))
-    assert eq(da.argmin(a, axis=0, split_every=2), np.argmin(x, axis=0))
-    assert eq(da.nanargmax(a, axis=0, split_every=2), np.nanargmax(x, axis=0))
-    assert eq(da.nanargmin(a, axis=0, split_every=2), np.nanargmin(x, axis=0))
-
 
 def reduction_2d_test(da_func, darr, np_func, narr, use_dtype=True,
                       split_every=True):
@@ -144,23 +134,32 @@ def test_reductions_2D(dtype):
     reduction_2d_test(da.nanmin, a, np.nanmin, x, False)
     reduction_2d_test(da.nanmax, a, np.nanmax, x, False)
 
-    assert eq(da.argmax(a, axis=0), np.argmax(x, axis=0))
-    assert eq(da.argmin(a, axis=0), np.argmin(x, axis=0))
-    assert eq(da.nanargmax(a, axis=0), np.nanargmax(x, axis=0))
-    assert eq(da.nanargmin(a, axis=0), np.nanargmin(x, axis=0))
-    assert eq(da.argmax(a, axis=1), np.argmax(x, axis=1))
-    assert eq(da.argmin(a, axis=1), np.argmin(x, axis=1))
-    assert eq(da.nanargmax(a, axis=1), np.nanargmax(x, axis=1))
-    assert eq(da.nanargmin(a, axis=1), np.nanargmin(x, axis=1))
-
-    assert eq(da.argmax(a, axis=0, split_every=2), np.argmax(x, axis=0))
-    assert eq(da.argmin(a, axis=0, split_every=2), np.argmin(x, axis=0))
-    assert eq(da.nanargmax(a, axis=0, split_every=2), np.nanargmax(x, axis=0))
-    assert eq(da.nanargmin(a, axis=0, split_every=2), np.nanargmin(x, axis=0))
-    assert eq(da.argmax(a, axis=1, split_every=2), np.argmax(x, axis=1))
-    assert eq(da.argmin(a, axis=1, split_every=2), np.argmin(x, axis=1))
-    assert eq(da.nanargmax(a, axis=1, split_every=2), np.nanargmax(x, axis=1))
-    assert eq(da.nanargmin(a, axis=1, split_every=2), np.nanargmin(x, axis=1))
+
+@pytest.mark.parametrize(['dfunc', 'func'],
+        [(da.argmin, np.argmin), (da.argmax, np.argmax),
+         (da.nanargmin, np.nanargmin), (da.nanargmax, np.nanargmax)])
+def test_arg_reductions(dfunc, func):
+    x = np.random.random((10, 10, 10))
+    a = da.from_array(x, chunks=(3, 4, 5))
+
+    assert eq(dfunc(a), func(x))
+    assert eq(dfunc(a, 0), func(x, 0))
+    assert eq(dfunc(a, 1), func(x, 1))
+    assert eq(dfunc(a, 2), func(x, 2))
+    with set_options(split_every=2):
+        assert eq(dfunc(a), func(x))
+        assert eq(dfunc(a, 0), func(x, 0))
+        assert eq(dfunc(a, 1), func(x, 1))
+        assert eq(dfunc(a, 2), func(x, 2))
+
+    pytest.raises(ValueError, lambda: dfunc(a, 3))
+    pytest.raises(TypeError, lambda: dfunc(a, (0, 1)))
+
+    x2 = np.arange(10)
+    a2 = da.from_array(x2, chunks=3)
+    assert eq(dfunc(a2), func(x2))
+    assert eq(dfunc(a2, 0), func(x2, 0))
+    assert eq(dfunc(a2, 0, split_every=2), func(x2, 0))
 
 
 def test_reductions_2D_nans():