Repository: dask/dask
Pull request: Result of reducing an array should not depend on its chunk-structure!!!
- [x] Closes #8541
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`

It's taken me a while to find a good example to motivate the seriousness of #8541 and prove that it is in fact a bug.  I think I might have found one now.  

As the source-code currently stands, running the following code to reduce an array gives different results depending on the chunk-structure of the array `x`:

```python
import dask.array as da
import numpy as np


a, b = 6, 6
array = np.arange(a*b).reshape(a, b)
# Reducing an array should not depend on its chunking pattern!!!
for chunks in [
    ((3, 3), (2, 1, 2, 1)),
    ((2, 1, 2, 1), (3, 3)),
]:
    x = da.from_array(array, chunks=chunks)
    x_0 = da.reduction(
        x,
        lambda x, axis, keepdims: x,
        lambda x, axis, keepdims: x,
        keepdims=True,
        axis=(0, 1),
        dtype=x.dtype,
        meta=np.array([[]])
    ).compute()
    print(f"{x_0=}\n")
```

Output:
```python
x_0=array(
      [[ 0,  1,  2],
       [ 6,  7,  8],
       [12, 13, 14],
       [18, 19, 20],
       [24, 25, 26],
       [30, 31, 32],
       [ 3,  4,  5],
       [ 9, 10, 11],
       [15, 16, 17],
       [21, 22, 23],
       [27, 28, 29],
       [33, 34, 35]]
)

x_0=array(
      [[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23],
       [24, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35]]
)
```

This PR fixes this bug yielding the correct output:
```python
x_0=array(
      [[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23],
       [24, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35]]
)

x_0=array(
      [[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23],
       [24, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35]]
)
```

I think the reason why `dask` has gotten away with this bug for so long is that in most cases implemented in `dask/array/reductions.py` (for example, `sum`, `mean`, etc.) the reduction operation is *axis-independent*.  The bug only shows up in "axis-dependent reductions". 

Problem statement:
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.

Hints from issue 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!
Hi @bryanwweber, @scharlottej13 

Pls [follow this link](https://github.com/ParticularMiner/dask/pull/2#issue-1110345797) to see a fix in my `dask` fork in which I chose to fix the number of dimensions of the list of lists.  (This seemed to be the easiest and most backward-compatible option.)

There, the code above now gives at the 1st stage of recursion:
```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]])]
]
```
and at the 2nd stage of recursion:
```python
axis=(0, 1)
x_chunks = [
   [array([[10]]), array([[18]])]
]
```

Fortunately, all CI tests there have also passed except for `parquet`-related tests in python 3.7.  I don't yet know if the errors raised by those failed tests are related to this issue or not.  If you find that they are, do let me know.  

And of course, if you feel this fix is worth merging upstream, let me know, and I'll proceed to submit a PR here.


@bryanwweber, @scharlottej13 

Moreover, exchanging the values of `a` and `b`:
```python
a, b = 4, 2
```
yields:
1st stage of recursion:
```python
axis=(0, 1)
x_chunks = [
   [array([[0]]), array([[1]])],
   [array([[2]]), array([[3]])]
]

axis=(0, 1)
x_chunks = [
   [array([[4]]), array([[5]])],
   [array([[6]]), array([[7]])]
]
```
2nd stage of recursion:
```python
axis=(0, 1)
x_chunks = [
   [array([[6]])],
   [array([[22]])]
]
```
which is what one would expect.

Commit message: fixed more linting errors in test_chunk_structure_dependence()

Patch to review:
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
index b8b2bb7f5b3..1a7bb648d7a 100644
--- 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
index c2f1316e1ca..490827915d2 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -294,7 +294,7 @@ def partial_reduce(
     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)
+        dummy = dict(i for i in enumerate(p) if i[0] in split_every)
         g = lol_tuples((x.name,), range(x.ndim), decided, dummy)
         dsk[(name,) + k] = (func, g)
     graph = HighLevelGraph.from_collections(name, dsk, dependencies=[x])
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
index 6afb175df08..eb305f4f814 100644
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -809,3 +809,37 @@ 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", [((3, 3), (2, 1, 2, 1)), ((2, 1, 2, 1), (3, 3))])
+def test_chunk_structure_dependence(chunks):
+    # Reducing an array should not depend on its chunk-structure!!!
+    # See Issue #8541: https://github.com/dask/dask/issues/8541
+
+    def chunk_func(x_chunks, axis, keepdims):
+        return da.core._concatenate2(x_chunks, axes=axis)
+
+    a, b = 6, 6
+    array = np.arange(a * b).reshape(a, b)
+    x = da.from_array(array, chunks=chunks)
+    reduced_x = da.reduction(
+        x,
+        lambda x, axis, keepdims: x,
+        chunk_func,
+        keepdims=True,
+        concatenate=False,
+        axis=(0, 1),
+        dtype=x.dtype,
+        meta=np.array([[]]),
+    ).compute()
+    expected = np.array(
+        [
+            [0, 1, 2, 3, 4, 5],
+            [6, 7, 8, 9, 10, 11],
+            [12, 13, 14, 15, 16, 17],
+            [18, 19, 20, 21, 22, 23],
+            [24, 25, 26, 27, 28, 29],
+            [30, 31, 32, 33, 34, 35],
+        ]
+    )
+    np.testing.assert_equal(reduced_x, expected)


Code review comments:
diff_hunk: @@ -809,3 +809,37 @@ 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", [((3, 3), (2, 1, 2, 1)), ((2, 1, 2, 1), (3, 3))])
+def test_chunk_structure_dependence(chunks):
path: dask/array/tests/test_reductions.py
text: @user1:
It'd be good to have a test that includes more than 2 dimensions. It'd also be good to check with an array that has enough chunks that `split_every` could be tested too.

@author:
I completely agree.  See the code for a new test implementation (and below for a corresponding example-code).  In fact, after following your suggestion, a few `ValueError` exceptions were raised indicating that the reduction axes in the parameter `axis` had to be supplied in ascending order.  But I've now amended the code in such a way that this requirement no longer matters:

In `_tree_reduce()`, I added the line:
```python
    axis = sorted(axis)
```
just before passing `axis` to `_concatenate2`.

UPDATE:
The above "correction" was flawed and brought-up errors in other test units.  The correct thing to do was to ensure `_concatenate2()` alone (and not the callable user functions `combine()` or `aggregate`) received the sorted `axis` tuple:

```python
    if concatenate:
        func = compose(func, partial(_concatenate2, axes=sorted(axis)))
```

diff_hunk: @@ -809,3 +809,37 @@ 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", [((3, 3), (2, 1, 2, 1)), ((2, 1, 2, 1), (3, 3))])
+def test_chunk_structure_dependence(chunks):
+    # Reducing an array should not depend on its chunk-structure!!!
+    # See Issue #8541: https://github.com/dask/dask/issues/8541
+
+    def chunk_func(x_chunks, axis, keepdims):
+        return da.core._concatenate2(x_chunks, axes=axis)
+
+    a, b = 6, 6
+    array = np.arange(a * b).reshape(a, b)
+    x = da.from_array(array, chunks=chunks)
+    reduced_x = da.reduction(
+        x,
+        lambda x, axis, keepdims: x,
+        chunk_func,
+        keepdims=True,
+        concatenate=False,
path: dask/array/tests/test_reductions.py
text: @user1:
In a comment you gave a simpler test case that had `concatenate=True` - would you mind switching to that test case (if you think it's sufficient)?

@author:
Sure.  No problem.

diff_hunk: @@ -294,7 +294,7 @@ def partial_reduce(
     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)
+        dummy = dict(i for i in enumerate(p) if i[0] in split_every)
path: dask/array/reductions.py
text: @user1:
This isn't quite right (but I'm not sure what it should be). The keys in `decided` and `dummy` should never intersect (but `assert set(decided).union(dummy) == set(range(ndim))` should always be true). I suspect you probably want to change how `decided` is created, and leave the creation of `dummy` as it was before.

@author:
I agree with your latter assertion:
```python
assert set(decided).union(dummy) == set(range(ndim))
```
but I differ on your claim that "The keys in `decided` and `dummy` should never intersect".

The following was my reasoning (pardon me if I'm needlessly going over what you already know):

* I discovered that `lol_tuples()` (on the next line after `dummy`) outputs a single array-like list, or list-of-lists (lol), `g` whose number of dimensions always equals the number of keys in `dummy`.  
   https://github.com/dask/dask/blob/6824e5dd7807971270d5828a61618ccd81592199/dask/array/reductions.py#L294-L299


* On the next line, this lol `g` is then passed on to `_concatenate2()` together with the previously passed `axis` tuple (from `_tree_reduce()`) which never changes throughout the recursion.  https://github.com/dask/dask/blob/6824e5dd7807971270d5828a61618ccd81592199/dask/array/reductions.py#L232-L233

* Now since `_concatenate2()` needs a corresponding lol-dimension for each value of `axis`, then the length of `dummy` also needs to remain fixed even if its individual key-values change with each recursion.

* The keys of `decided`, on the other hand, are all those axes whose number of chunks (`.numblocks`) is equal to 1, either because the axis in question has been fully reduced, or because it is not taking part in the reduction, as can be deduced from following code-extracts: https://github.com/dask/dask/blob/6824e5dd7807971270d5828a61618ccd81592199/dask/array/reductions.py#L280-L283
https://github.com/dask/dask/blob/6824e5dd7807971270d5828a61618ccd81592199/dask/array/reductions.py#L295-L296
To confirm this, the following code snippet:
```python
from tlz import partition_all

x_numblocks = (16, 8, 4)
split_every = {0: 2, 1: 2}	# keys here are the reduction axes
parts = [
    list(partition_all(split_every.get(i, 1), range(n)))
    for (i, n) in enumerate(x_numblocks)
]
```
yields:
```python
parts = 
[[(0, 1), (2, 3), (4, 5), (6, 7), (8, 9), (10, 11), (12, 13), (14, 15)],
 [(0, 1), (2, 3), (4, 5), (6, 7)],
 [(0,), (1,), (2,), (3,)]]
```

* So I concluded that while `dummy` ought to remain fixed in length, `decided` will increase in length (starting from the number of axes not taking part in the reduction), thus intersecting and eventually subsuming `dummy`.

@user1:
> but I differ on your claim that "The keys in decided and dummy should never intersect".

Sorry, I think I wasn't clear here. I'm not saying that how you're handling generating `dummy` is wrong, all I'm saying is that the following assertion should never trigger:

```
assert not set(dummy).intersection(decided)
```

Otherwise this code is relying on an implementation detail of `lol_tuples`, which checks `dummies` first before checking `decided`. If this implemention were to change at all then the existing code would be wrong. However you want to handle this is up to you, but that assertion should pass (note, I don't think you should add that assertion to the code, just that you should think about it in your code above). 

@author:
Oh I see.  You're right about that.  Sorry for the misunderstanding.  Hmm ...  I'll need to think about this some more.

@author:
Have you any suggestion on how I could go about this?

@user1:
If you think your code is correct as is, then it might be as simple as just removing the intersecting ones from `dummy`. Not elegant, but functional. No strong thoughts though.

@author:
Thanks @user1 .  This was a good idea, and quite important too on deeper reflection.   I'm guessing that you meant to say, "remove the intersecting ones from `decided`," instead of `dummy`, right? 

I also renamed `decided` to `free`, which appears to be the designated antonym for `dummy` in the context of indices.  I hope it improves readability of the code. 

Well then, I suppose my latest commit completes this PR's review.  Would you agree?

diff_hunk: @@ -809,3 +809,37 @@ 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", [((3, 3), (2, 1, 2, 1)), ((2, 1, 2, 1), (3, 3))])
+def test_chunk_structure_dependence(chunks):
+    # Reducing an array should not depend on its chunk-structure!!!
+    # See Issue #8541: https://github.com/dask/dask/issues/8541
+
+    def chunk_func(x_chunks, axis, keepdims):
+        return da.core._concatenate2(x_chunks, axes=axis)
+
+    a, b = 6, 6
+    array = np.arange(a * b).reshape(a, b)
+    x = da.from_array(array, chunks=chunks)
+    reduced_x = da.reduction(
+        x,
+        lambda x, axis, keepdims: x,
+        chunk_func,
+        keepdims=True,
+        concatenate=False,
+        axis=(0, 1),
+        dtype=x.dtype,
+        meta=np.array([[]]),
+    ).compute()
+    expected = np.array(
+        [
+            [0, 1, 2, 3, 4, 5],
+            [6, 7, 8, 9, 10, 11],
+            [12, 13, 14, 15, 16, 17],
+            [18, 19, 20, 21, 22, 23],
+            [24, 25, 26, 27, 28, 29],
+            [30, 31, 32, 33, 34, 35],
+        ]
+    )
+    np.testing.assert_equal(reduced_x, expected)
path: dask/array/tests/test_reductions.py
text: @user1:
It's better to use `assert_eq` here, as it also checks that the resulting metadata matches the computed structure.

```suggestion
    reduced_x = da.reduction(
        x,
        lambda x, axis, keepdims: x,
        chunk_func,
        keepdims=True,
        concatenate=False,
        axis=(0, 1),
        dtype=x.dtype,
        meta=np.array([[]]),
    )
    expected = np.array(
        [
            [0, 1, 2, 3, 4, 5],
            [6, 7, 8, 9, 10, 11],
            [12, 13, 14, 15, 16, 17],
            [18, 19, 20, 21, 22, 23],
            [24, 25, 26, 27, 28, 29],
            [30, 31, 32, 33, 34, 35],
        ]
    )
    assert_eq(reduced_x, expected)
```

@author:
Done.  I had to add `reduced_x.compute_chunk_sizes()` before `assert_eq()` though.