Repository: dask/dask

Issue:
Dask array reshape bug for arrays with fully chunked leading axes
xref: https://github.com/pystatgen/sgkit/issues/456

https://github.com/dask/dask/pull/6748 is a special case optimisation:

> When the slow-moving (early) axes in .reshape are all size 1

One of our tests in https://github.com/pystatgen/sgkit/issues/456 falls into that category. And afaiu https://github.com/dask/dask/pull/6748 introduced a bug:

```python
> # c4038add1a087ba3a82207a557cdcad9753b689d is the https://github.com/dask/dask/pull/6748
> dask git:(c4038add) g ck c4038add1a087ba3a82207a557cdcad9753b689d
HEAD is now at c4038add Avoid rechunking in reshape with chunksize=1 (#6748)
> dask git:(c4038add) python
>>> import numpy as np
>>> import dask.array as da
>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (4,)))
>>> a.reshape(6,4)
dask.array<reshape, shape=(6, 4), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> # merging dimensions at the front works fine, now let's try the last two (which is our use case)
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 24), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> # NOTICE: the shape is (2, 24) NOT (2,12)!, now let's try to compute this:
>>> a.reshape(2,12).compute()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/rav/projects/dask/dask/base.py", line 167, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/Users/rav/projects/dask/dask/base.py", line 454, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/Users/rav/projects/dask/dask/threaded.py", line 76, in get
    results = get_async(
  File "/Users/rav/projects/dask/dask/local.py", line 503, in get_async
    return nested_get(result, state["cache"])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in nested_get
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in <listcomp>
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in nested_get
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in <listcomp>
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in nested_get
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in <listcomp>
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 301, in nested_get
    return coll[ind]
KeyError: ('reshape-aa29de8b0f6be5be25495836ed047c4a', 1, 0)
```

Notice the invalid shape after `a.reshape(2,12)`.

The same code works fine without https://github.com/dask/dask/pull/6748:

```python
> dask git:(c4038add) g ck head~1
Previous HEAD position was c4038add Avoid rechunking in reshape with chunksize=1 (#6748)
HEAD is now at 94bdd4e3 Try to make categoricals work on join (#6205)
> dask git:(94bdd4e3) python
>>> import numpy as np
>>> import dask.array as da
>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (4,)))
>>> a.reshape(6,4)
dask.array<reshape, shape=(6, 4), dtype=int64, chunksize=(3, 4), chunktype=numpy.ndarray>
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 12), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> a.reshape(2,12).compute()
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]])
```

This issue seems to be affected by the chunking of the array (`a`), in the case above (and in our reginie case) the lower axis are completely chunked, see:

```python
>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (4,)))
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 24), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> # BAD

>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (2,2)))
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 24), dtype=int64, chunksize=(1, 2), chunktype=numpy.ndarray>
>>> # BAD

>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (1, 1, 1, 1)))
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 24), dtype=int64, chunksize=(1, 1), chunktype=numpy.ndarray>
>>> # BAD

>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 2), (2, 2)))
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 12), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> # OK
```

**Environment**:

- Dask version: introduced in 2020.12.0
- Python version: 3.8
- Operating System: mac
- Install method: conda/pip/source

Dask array reshape bug for arrays with fully chunked leading axes
xref: https://github.com/pystatgen/sgkit/issues/456

https://github.com/dask/dask/pull/6748 is a special case optimisation:

> When the slow-moving (early) axes in .reshape are all size 1

One of our tests in https://github.com/pystatgen/sgkit/issues/456 falls into that category. And afaiu https://github.com/dask/dask/pull/6748 introduced a bug:

```python
> # c4038add1a087ba3a82207a557cdcad9753b689d is the https://github.com/dask/dask/pull/6748
> dask git:(c4038add) g ck c4038add1a087ba3a82207a557cdcad9753b689d
HEAD is now at c4038add Avoid rechunking in reshape with chunksize=1 (#6748)
> dask git:(c4038add) python
>>> import numpy as np
>>> import dask.array as da
>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (4,)))
>>> a.reshape(6,4)
dask.array<reshape, shape=(6, 4), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> # merging dimensions at the front works fine, now let's try the last two (which is our use case)
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 24), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> # NOTICE: the shape is (2, 24) NOT (2,12)!, now let's try to compute this:
>>> a.reshape(2,12).compute()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/rav/projects/dask/dask/base.py", line 167, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/Users/rav/projects/dask/dask/base.py", line 454, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/Users/rav/projects/dask/dask/threaded.py", line 76, in get
    results = get_async(
  File "/Users/rav/projects/dask/dask/local.py", line 503, in get_async
    return nested_get(result, state["cache"])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in nested_get
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in <listcomp>
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in nested_get
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in <listcomp>
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in nested_get
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 299, in <listcomp>
    return tuple([nested_get(i, coll) for i in ind])
  File "/Users/rav/projects/dask/dask/local.py", line 301, in nested_get
    return coll[ind]
KeyError: ('reshape-aa29de8b0f6be5be25495836ed047c4a', 1, 0)
```

Notice the invalid shape after `a.reshape(2,12)`.

The same code works fine without https://github.com/dask/dask/pull/6748:

```python
> dask git:(c4038add) g ck head~1
Previous HEAD position was c4038add Avoid rechunking in reshape with chunksize=1 (#6748)
HEAD is now at 94bdd4e3 Try to make categoricals work on join (#6205)
> dask git:(94bdd4e3) python
>>> import numpy as np
>>> import dask.array as da
>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (4,)))
>>> a.reshape(6,4)
dask.array<reshape, shape=(6, 4), dtype=int64, chunksize=(3, 4), chunktype=numpy.ndarray>
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 12), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> a.reshape(2,12).compute()
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]])
```

This issue seems to be affected by the chunking of the array (`a`), in the case above (and in our reginie case) the lower axis are completely chunked, see:

```python
>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (4,)))
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 24), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> # BAD

>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (2,2)))
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 24), dtype=int64, chunksize=(1, 2), chunktype=numpy.ndarray>
>>> # BAD

>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 1, 1), (1, 1, 1, 1)))
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 24), dtype=int64, chunksize=(1, 1), chunktype=numpy.ndarray>
>>> # BAD

>>> a = da.from_array(np.arange(24).reshape(2, 3, 4), chunks=((1, 1), (1, 2), (2, 2)))
>>> a.reshape(2,12)
dask.array<reshape, shape=(2, 12), dtype=int64, chunksize=(1, 4), chunktype=numpy.ndarray>
>>> # OK
```

**Environment**:

- Dask version: introduced in 2020.12.0
- Python version: 3.8
- Operating System: mac
- Install method: conda/pip/source



Discussion:
Hmm thanks for the report and apologies for the bug. Are you interested in digging into this? It'll be a bit before I can take a look.

The good news is that the bug should be localized to just these lines: https://github.com/dask/dask/blob/2632bbce82d471b7fef426729faec9a4e78ee137/dask/array/reshape.py#L49-L54. They're a bit dense though.
@TomAugspurger np about the bug, and thanks for working on making dask better. In sgkit, we will mark this case as xfail, I might take a look over the weekend but if I don't come back with a fix early next week, chances are slim I will have time to work on this in the future.
cc'ing the @dask/array team in case anyone wants to dig into what is hopefully a small bug in the reshape code (beware, the reshape code is not for wimps :) )
> Hmm thanks for the report and apologies for the bug. Are you interested in digging into this? It'll be a bit before I can take a look.
> 
> The good news is that the bug should be localized to just these lines:
> 
> https://github.com/dask/dask/blob/2632bbce82d471b7fef426729faec9a4e78ee137/dask/array/reshape.py#L49-L54
> . They're a bit dense though.

I may be mistaken (as @mrocklin mentioned, this is pretty dense), but I believe line 52-54 should be:
```python
result_outchunks[oi] = inchunks[i] * np.prod(list(map(len, inchunks[ileft:ii])))
```
The issue at present seems to be that the product is taken over all the input dimensions up to `ii` when in fact we only want to do it over the dimensions which are being combined. I am about to head off for the long weekend, otherwise I would do some more rigorous testing/submit a PR. If there is no further movement before I get back, I will take another look. 
Hmm thanks for the report and apologies for the bug. Are you interested in digging into this? It'll be a bit before I can take a look.

The good news is that the bug should be localized to just these lines: https://github.com/dask/dask/blob/2632bbce82d471b7fef426729faec9a4e78ee137/dask/array/reshape.py#L49-L54. They're a bit dense though.
@TomAugspurger np about the bug, and thanks for working on making dask better. In sgkit, we will mark this case as xfail, I might take a look over the weekend but if I don't come back with a fix early next week, chances are slim I will have time to work on this in the future.
cc'ing the @dask/array team in case anyone wants to dig into what is hopefully a small bug in the reshape code (beware, the reshape code is not for wimps :) )
> Hmm thanks for the report and apologies for the bug. Are you interested in digging into this? It'll be a bit before I can take a look.
> 
> The good news is that the bug should be localized to just these lines:
> 
> https://github.com/dask/dask/blob/2632bbce82d471b7fef426729faec9a4e78ee137/dask/array/reshape.py#L49-L54
> . They're a bit dense though.

I may be mistaken (as @mrocklin mentioned, this is pretty dense), but I believe line 52-54 should be:
```python
result_outchunks[oi] = inchunks[i] * np.prod(list(map(len, inchunks[ileft:ii])))
```
The issue at present seems to be that the product is taken over all the input dimensions up to `ii` when in fact we only want to do it over the dimensions which are being combined. I am about to head off for the long weekend, otherwise I would do some more rigorous testing/submit a PR. If there is no further movement before I get back, I will take another look. 

Resolving patch:
diff --git a/dask/array/reshape.py b/dask/array/reshape.py
--- a/dask/array/reshape.py
+++ b/dask/array/reshape.py
@@ -49,8 +49,8 @@ def reshape_rechunk(inshape, outshape, inchunks):
             if all(len(inchunks[i]) == inshape[i] for i in range(ii)):
                 for i in range(ii + 1):
                     result_inchunks[i] = inchunks[i]
-                result_outchunks[oi] = inchunks[i] * np.prod(
-                    list(map(len, inchunks[:i]))
+                result_outchunks[oi] = inchunks[ii] * np.prod(
+                    list(map(len, inchunks[ileft:ii]))
                 )
             else:
                 for i in range(ileft + 1, ii + 1):  # need single-shape dimensions


Test patch:
diff --git a/dask/array/tests/test_reshape.py b/dask/array/tests/test_reshape.py
--- a/dask/array/tests/test_reshape.py
+++ b/dask/array/tests/test_reshape.py
@@ -107,6 +107,8 @@ def test_reshape_unknown_sizes():
         ((2, 2, 3, 4), ((1, 1), (2,), (1, 2), (4,)), (4, 3, 4), ((2, 2), (1, 2), (4,))),
         # (2, 3, 4) -> (24,).
         ((2, 3, 4), ((1, 1), (1, 1, 1), (2, 2)), (24,), ((2,) * 12,)),
+        # (2, 3, 4) -> (2, 12)
+        ((2, 3, 4), ((1, 1), (1, 1, 1), (4,)), (2, 12), ((1, 1), (4,) * 3)),
     ],
 )
 def test_reshape_all_chunked_no_merge(inshape, inchunks, outshape, outchunks):