Repository: dask/dask

Issue:
Mutually inconsistent `x_chunk` and `axis` argument-values during `dask.array.reduction()`
Hi all @dask

I continue to thank you guys for the wonderful program that `dask` is!

I'm not sure if the behavior described below is by design or an oversight.  But feel free to let me know what you think.

When "tree-reducing" a `dask` array `x` along multiple axes using 
```python
dask.array.reduction(x, chunk, aggregate, axis=axes, concatenate=concatenate, ...)
```
the values of the `x_chunk` and `axis` arguments passed to the callable `aggregate` when `concatenate=False` (or to the looked-up function `concatenate()` when `concatenate=True`) are in general not mutually consistent.  Here, the value of `x_chunk` is, in general, a nested list of lists of `chunk_type`-arrays.  The number of nesting levels of `x_chunk` reduces as each axis gets fully reduced, but the value of `axis` always remains the same.

This makes it difficult/impossible at each stage of recursion to determine within the callable `aggregate` or `concatenate()` which axes have already been fully reduced and those which are not yet fully reduced, and act accordingly.  And as a result, one could end up inadvertently reducing/concatenating along the wrong axis!  Thus I'm often forced to circumvent this issue by reducing one axis at a time using a for-loop for example:

```python
red_axes = (1, 2)
for axis in red_axes[::-1]:
   red = dask.array.reduction(red, chunk, aggregate, axis=axis, dtype=red.dtype, meta=meta)
```
since, by doing so, I'm always assured of a single nesting-level in each loop and there is no confusion.

There could be two ways to solve this problem permanently:
Either amend the source-code to
1.  at each stage of recursion, pass only those axes that have not yet been fully reduced into `aggregate` (or `concatenate()`) using the argument `axis`, or
2. fix the number of nesting levels of the list of lists `x_chunk` throughout the reduction procedure.  In this case, those axes that have already been fully reduced will have a corresponding nesting-level whose lists are of length 1.

I've already examined the source-code to determine the viability of both options above.  But before I go ahead and submit a PR, your thoughts would be much appreciated.  Thanks.


Discussion:
To illustrate the above issue, see the following example:

```python
import dask.array as da
import numpy as np
from threading import Lock
from more_itertools import collapse
from functools import reduce


print_lock = Lock()


def display(x_chunks, axis):
    x_chunks = [f'   {chunk}' if type(chunk) is list
                else f'   array({chunk})'
                for chunk in x_chunks]
    with print_lock:
        print()
        print(f'{axis=}\n'
              f'x_chunks = [')
        print(*x_chunks, sep=',\n', end='')
        print('\n]')


def chunk_func(x_chunks, axis, keepdims):
    display(x_chunks, axis)
    out = reduce(np.add, collapse(x_chunks, levels=1))
    if keepdims:
        return out
    else:
        return out[0][0]


a, b = 2, 4
x = da.from_array(np.arange(a*b).reshape(a, b), chunks=1)
x_0 = da.reduction(
    x,
    lambda x, axis, keepdims: x,
    chunk_func,
    concatenate=False,
    axis=(0, 1),
    dtype=x.dtype,
    meta=np.array([])
).compute()
```

The following is the output from the first stage of recursion during `reduction()` where we observe `x_chunks` is a list of lists (two-dimensional) as expected:
```python
axis=(0, 1)
x_chunks = [
   [array([[0]]), array([[1]])],
   [array([[4]]), array([[5]])]
]

axis=(0, 1)
x_chunks = [
   [array([[2]]), array([[3]])],
   [array([[6]]), array([[7]])]
]
```
Next is the output from the second stage of recursion where we observe `x_chunks` is just a list (one-dimensional) since one axis is already completely reduced but not the other.  But which one?
```python
axis=(0, 1)
x_chunks = [
   array([[10]]),
   array([[18]])
]
```
Well, obviously axis 0 was reduced first since it initially had only 2 blocks while axis 1 initially had 4 blocks.  But how can `chunk_func` determine this if `axis` is always `(0, 1)` and the function has no access to block information?
Thanks for this issue @ParticularMiner! @bryanwweber thanks for taking a look! 
Hi @ParticularMiner sorry, I haven't had a chance to look at this yet. Hope to this week!

Resolving patch:
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
--- a/dask/array/chunk.py
+++ b/dask/array/chunk.py
@@ -244,6 +244,7 @@ def argtopk_aggregate(a_plus_idx, k, axis, keepdims):
     and return the index only.
     """
     assert keepdims is True
+    a_plus_idx = a_plus_idx if len(a_plus_idx) > 1 else a_plus_idx[0]
     a, idx = argtopk(a_plus_idx, k, axis, keepdims)
     axis = axis[0]
 
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -230,7 +230,7 @@ def _tree_reduce(
             depth = int(builtins.max(depth, ceil(log(n, split_every[i]))))
     func = partial(combine or aggregate, axis=axis, keepdims=True)
     if concatenate:
-        func = compose(func, partial(_concatenate2, axes=axis))
+        func = compose(func, partial(_concatenate2, axes=sorted(axis)))
     for i in range(depth - 1):
         x = partial_reduce(
             func,
@@ -243,7 +243,7 @@ def _tree_reduce(
         )
     func = partial(aggregate, axis=axis, keepdims=keepdims)
     if concatenate:
-        func = compose(func, partial(_concatenate2, axes=axis))
+        func = compose(func, partial(_concatenate2, axes=sorted(axis)))
     return partial_reduce(
         func,
         x,
@@ -293,9 +293,11 @@ def partial_reduce(
         out_chunks = list(getter(out_chunks))
     dsk = {}
     for k, p in zip(keys, product(*parts)):
-        decided = {i: j[0] for (i, j) in enumerate(p) if len(j) == 1}
-        dummy = dict(i for i in enumerate(p) if i[0] not in decided)
-        g = lol_tuples((x.name,), range(x.ndim), decided, dummy)
+        free = {
+            i: j[0] for (i, j) in enumerate(p) if len(j) == 1 and i not in split_every
+        }
+        dummy = dict(i for i in enumerate(p) if i[0] in split_every)
+        g = lol_tuples((x.name,), range(x.ndim), free, dummy)
         dsk[(name,) + k] = (func, g)
     graph = HighLevelGraph.from_collections(name, dsk, dependencies=[x])
 


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
@@ -1,6 +1,6 @@
 import os
 import warnings
-from itertools import zip_longest
+from itertools import permutations, zip_longest
 
 import pytest
 
@@ -809,3 +809,27 @@ def test_nan_func_does_not_warn(func):
     with pytest.warns(None) as rec:
         getattr(da, func)(d).compute()
     assert not rec  # did not warn
+
+
+@pytest.mark.parametrize("chunks", list(permutations(((2, 1) * 8, (3,) * 8, (6,) * 4))))
+@pytest.mark.parametrize("split_every", [2, 4])
+@pytest.mark.parametrize(
+    "axes", list(permutations((0, 1, 2), 2)) + list(permutations((0, 1, 2)))
+)
+def test_chunk_structure_independence(axes, split_every, chunks):
+    # Reducing an array should not depend on its chunk-structure!!!
+    # See Issue #8541: https://github.com/dask/dask/issues/8541
+    shape = tuple(np.sum(s) for s in chunks)
+    np_array = np.arange(np.prod(shape)).reshape(*shape)
+    x = da.from_array(np_array, chunks=chunks)
+    reduced_x = da.reduction(
+        x,
+        lambda x, axis, keepdims: x,
+        lambda x, axis, keepdims: x,
+        keepdims=True,
+        axis=axes,
+        split_every=split_every,
+        dtype=x.dtype,
+        meta=x._meta,
+    )
+    _assert_eq(reduced_x, np_array, check_chunks=False, check_shape=False)