Repository: dask/dask

Issue:
Regression in 0.18.2: argtopk(split_every=2) broken
```
import dask.array as da 
a = da.from_array([[1, 4, 2, 5, 3], 
                   [7, 1, 8, 0, 2]], chunks=1) 
a.argtopk(-3, split_every=2).compute()


AttributeError: 'tuple' object has no attribute 'shape'
```

I'm still in the middle of investigating how it happened and, most importantly, why the regression tests didn't spot it.


Discussion:
There's a second one.
```
b = da.random.random((10, ), chunks=c)
b.argtopk(5, split_every=2).compute()
```
c=1 -> AttributeError: 'tuple' object has no attribute 'shape'
c=2 -> AttributeError: 'tuple' object has no attribute 'shape'
c=3 -> ValueError: all the input arrays must have same number of dimensions
c=4 -> works
Third :angry:

```
>>> b = da.random.random((10, ), chunks=1)
>>> b.argtopk(5).compute()
ValueError: too many values to unpack (expected 2)
```

Resolving patch:
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
--- a/dask/array/chunk.py
+++ b/dask/array/chunk.py
@@ -9,6 +9,7 @@
 from . import numpy_compat as npcompat
 
 from ..compatibility import getargspec
+from ..core import flatten
 from ..utils import ignoring
 
 try:
@@ -233,6 +234,7 @@ def argtopk(a_plus_idx, k, axis, keepdims):
     axis = axis[0]
 
     if isinstance(a_plus_idx, list):
+        a_plus_idx = list(flatten(a_plus_idx))
         a = np.concatenate([ai for ai, _ in a_plus_idx], axis)
         idx = np.concatenate([broadcast_to(idxi, ai.shape)
                               for ai, idxi in a_plus_idx], axis)


Test patch:
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -460,7 +460,26 @@ def test_topk_argtopk1(npfunc, daskfunc, split_every):
         daskfunc(b, -k, axis=3, split_every=split_every)
 
 
-def test_topk_argtopk2():
+@pytest.mark.parametrize('npfunc,daskfunc', [
+    (np.sort, da.topk),
+    (np.argsort, da.argtopk),
+])
+@pytest.mark.parametrize('split_every', [None, 2, 3, 4])
+@pytest.mark.parametrize('chunksize', [1, 2, 3, 4, 5, 10])
+def test_topk_argtopk2(npfunc, daskfunc, split_every, chunksize):
+    """Fine test use cases when k is larger than chunk size"""
+    a = da.random.random((10, ), chunks=chunksize)
+    k = 5
+
+    # top 5 elements, sorted descending
+    assert_eq(npfunc(a)[-k:][::-1],
+              daskfunc(a, k, split_every=split_every))
+    # bottom 5 elements, sorted ascending
+    assert_eq(npfunc(a)[:k],
+              daskfunc(a, -k, split_every=split_every))
+
+
+def test_topk_argtopk3():
     a = da.random.random((10, 20, 30), chunks=(4, 8, 8))
 
     # Support for deprecated API for topk
@@ -468,5 +487,7 @@ def test_topk_argtopk2():
         assert_eq(da.topk(a, 5), da.topk(5, a))
 
     # As Array methods
-    assert_eq(a.topk(5, axis=1, split_every=2), da.topk(a, 5, axis=1, split_every=2))
-    assert_eq(a.argtopk(5, axis=1, split_every=2), da.argtopk(a, 5, axis=1, split_every=2))
+    assert_eq(a.topk(5, axis=1, split_every=2),
+              da.topk(a, 5, axis=1, split_every=2))
+    assert_eq(a.argtopk(5, axis=1, split_every=2),
+              da.argtopk(a, 5, axis=1, split_every=2))