Use of `da.where` before `da.einsum` on a chunked array produces incorrectly sized result

What happened:

Using dask.array.where prior to dask.array.einsum on some chunked arrays results in an incorrectly sized array. This seems to only occur if the array is chunked in the dimensions that are retained by einsum. Using einsum without first using ‘where’ produces the correct result. Note that this occurs even when ‘where’ does not change any values within the array.

What you expected to happen:

Use of dask.array.where shouldn’t affect the result shape of dask.array.einsum.

Minimal Complete Verifiable Example:

import numpy as np
import dask.array as da
a = da.asarray(
    [
        [1. , 0. ],
        [0. , 1. ],
        [0.5, 0.5],
    ], 
    chunks=((2,1), (2,))
)
b = da.where(a == -1, a, a)
assert a.chunks == b.chunks
np.testing.assert_array_equal(a, b)

da.einsum('xj,yj->xy', a, a).compute()

array([[1. , 0. , 0.5],
       [0. , 1. , 0.5],
       [0.5, 0.5, 0.5]])

da.einsum('xj,yj->xy', b, b).compute()

array([[1. , 0. , 1. , 0. ],
       [0. , 1. , 0. , 1. ],
       [0.5, 0.5, 0.5, 0.5],
       [0.5, 0.5, 0.5, 0.5]])

Environment:

• Dask version: 2021.9.0 (pypi)
• Python version: 3.9.6 (conda)
• Operating System: Ubuntu 21.04
• Install method: pip within conda env

[jcrist]

I’m not sure what PR fixed this (there have been a few touching this code recently), but the test example passes on main now:

import numpy as np
import dask.array as da
a = da.asarray(
    [
        [1. , 0. ],
        [0. , 1. ],
        [0.5, 0.5],
    ], 
    chunks=((2,1), (2,))
)
b = da.where(a == -1, a, a)
assert a.chunks == b.chunks
np.testing.assert_array_equal(a, b)

The issue seems to be resolved, closing.

[timothymillar]

Just for reference, the fix appears to be in dask 2022.01.0

Edit: digging a little deeper it was fixed in #8542 and seems to have been a symptom of duplicate dependency names in rewrite_blockwise, though I can’t quite grok the details.