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:
mrocklin: @crusaderky do you think that this is ready to merge?  
crusaderky: @mrocklin yes, I was just waiting for CI to give green light 

Patch:
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
index 498edfdf7..a36a4de7f 100644
--- a/dask/array/chunk.py
+++ b/dask/array/chunk.py
@@ -9,6 +9,7 @@ import numpy as np
 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
index c0e642847..83db8b57b 100644
--- 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))