Repository: dask/dask

Issue:
argtopk broken with uneven chunks
```
>>> a = da.ones(9, chunks=((4, 5), ))
>>> a.argtopk(2)
ValueError: Shapes do not align {0: {2, 3}}
```

The cause of the error is trivial; just need to change https://github.com/dask/dask/blob/669cde6b1a02e5fc7f0bcb924489c6ae8e697323/dask/array/reductions.py#L904 
to ``chunks=(a.chunks[axis], )``.

What baffles me however is the underlying reason why this went undetected so far. Can anybody tell me if there is any good use case for why this should not raise ValueError?!?
```
>>> normalize_chunks((4, 5), shape=(9, ))
((4, 4, 1),)
```


Discussion:

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2023,8 +2023,8 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None,
 
     Parameters
     ----------
-    chunks: tuple, int, or string
-        The chunks to be normalized.  See example below for more details
+    chunks: tuple, int, dict, or string
+        The chunks to be normalized.  See examples below for more details
     shape: Tuple[int]
         The shape of the array
     limit: int (optional)
@@ -2058,10 +2058,15 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None,
     >>> normalize_chunks(10, shape=(30, 5))
     ((10, 10, 10), (5,))
 
+    Expands dict inputs
+
+    >>> normalize_chunks({0: 2, 1: 3}, shape=(6, 6))
+    ((2, 2, 2), (3, 3))
+
     The value -1 gets mapped to full size
 
-    >>> normalize_chunks((-1,), shape=(10,))
-    ((10,),)
+    >>> normalize_chunks((5, -1), shape=(10, 10))
+    ((5, 5), (10,))
 
     Use the value "auto" to automatically determine chunk sizes along certain
     dimensions.  This uses the ``limit=`` and ``dtype=`` keywords to
@@ -2090,11 +2095,14 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None,
     if not chunks and shape and all(s == 0 for s in shape):
         chunks = ((0,),) * len(shape)
 
+    if (shape and len(shape) == 1 and len(chunks) > 1 and
+            all(isinstance(c, (Number, str)) for c in chunks)):
+        chunks = chunks,
+
     if shape and len(chunks) != len(shape):
-        if not (len(shape) == 1 and sum(chunks) == shape[0]):
-            raise ValueError(
-                "Chunks and shape must be of the same length/dimension. "
-                "Got chunks=%s, shape=%s" % (chunks, shape))
+        raise ValueError(
+            "Chunks and shape must be of the same length/dimension. "
+            "Got chunks=%s, shape=%s" % (chunks, shape))
     if -1 in chunks:
         chunks = tuple(s if c == -1 else c for c, s in zip(chunks, shape))
 
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -901,7 +901,7 @@ def argtopk(a, k, axis=-1, split_every=None):
     axis = validate_axis(a.ndim, axis)
 
     # Generate nodes where every chunk is a tuple of (a, original index of a)
-    idx = arange(a.shape[axis], chunks=a.chunks[axis], dtype=np.int64)
+    idx = arange(a.shape[axis], chunks=(a.chunks[axis], ), dtype=np.int64)
     idx = idx[tuple(slice(None) if i == axis else np.newaxis
                     for i in range(a.ndim))]
     a_plus_idx = a.map_blocks(chunk.argtopk_preprocess, idx,


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -2226,6 +2226,23 @@ def test_from_array_with_missing_chunks():
 
 def test_normalize_chunks():
     assert normalize_chunks(3, (4, 6)) == ((3, 1), (3, 3))
+    assert normalize_chunks(((3, 3), (8,)), (6, 8)) == ((3, 3), (8, ))
+    assert normalize_chunks((4, 5), (9,)) == ((4, 5), )
+    assert normalize_chunks((4, 5), (9, 9)) == ((4, 4, 1), (5, 4))
+    assert normalize_chunks(-1, (5, 5)) == ((5,), (5, ))
+    assert normalize_chunks((3, -1), (5, 5)) == ((3, 2), (5, ))
+    assert normalize_chunks({0: 3}, (5, 5)) == ((3, 2), (5,))
+    assert normalize_chunks([[2, 2], [3, 3]]) == ((2, 2), (3, 3))
+    assert normalize_chunks(10, (30, 5)), ((10, 10, 10), (5,))
+    assert normalize_chunks((), (0, 0)), ((0,), (0,))
+    assert normalize_chunks(-1, (0, 3)), ((0,), (3,))
+    assert normalize_chunks("auto", shape=(20,), limit=5, dtype='uint8') == \
+        ((5, 5, 5, 5),)
+
+    with pytest.raises(ValueError):
+        normalize_chunks(((10,), ), (11, ))
+    with pytest.raises(ValueError):
+        normalize_chunks(((5, ), (5, )), (5, ))
 
 
 def test_align_chunks_to_previous_chunks():
diff --git a/dask/array/tests/test_fft.py b/dask/array/tests/test_fft.py
--- a/dask/array/tests/test_fft.py
+++ b/dask/array/tests/test_fft.py
@@ -9,9 +9,7 @@
 from dask.array.fft import fft_wrap
 from dask.array.utils import assert_eq, same_keys
 
-from dask.array.core import (
-    normalize_chunks as _normalize_chunks,
-)
+from dask.array.core import normalize_chunks
 
 
 all_1d_funcnames = [
@@ -201,21 +199,21 @@ def test_fftfreq(n, d, c):
     r1 = np.fft.fftfreq(n, d)
     r2 = da.fft.fftfreq(n, d, chunks=c)
 
-    assert _normalize_chunks(c, r2.shape) == r2.chunks
+    assert normalize_chunks(c, r2.shape) == r2.chunks
 
     assert_eq(r1, r2)
 
 
 @pytest.mark.parametrize("n", [1, 2, 3, 6, 7])
 @pytest.mark.parametrize("d", [1.0, 0.5, 2 * np.pi])
-@pytest.mark.parametrize("c", [lambda m: m // 2 + 1, lambda m: (1, m // 2)])
+@pytest.mark.parametrize("c", [lambda m: (m // 2 + 1, ), lambda m: (1, m // 2)])
 def test_rfftfreq(n, d, c):
-    c = c(n)
+    c = [ci for ci in c(n) if ci != 0]
 
     r1 = np.fft.rfftfreq(n, d)
     r2 = da.fft.rfftfreq(n, d, chunks=c)
 
-    assert _normalize_chunks(c, r2.shape) == r2.chunks
+    assert normalize_chunks(c, r2.shape) == r2.chunks
 
     assert_eq(r1, r2)
 
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
@@ -426,7 +426,7 @@ def test_topk_argtopk1(npfunc, daskfunc, split_every):
     k = 5
     # Test at least 3 levels of aggregation when split_every=2
     # to stress the different chunk, combine, aggregate kernels
-    a = da.random.random(800, chunks=100)
+    a = da.random.random(800, chunks=((120, 80, 100, 200, 300), ))
     b = da.random.random((10, 20, 30), chunks=(4, 8, 8))
 
     # 1-dimensional arrays