Repository: dask/dask

Issue:
Blockwise optimization bug when collections appear multiple times
Okay! This is a weird one, and I'm still investigating what is going on, so bear with me. I was working on #7417 and ran into some surprising behavior with blockwise optimization. It appears to happen (at least) when 
1. a collection with the same name is used more than once in a high level graph,
2. when there are at least two blockwise layers that might be combined, and
3. when there are multiple partitions.

A minimal example: the following succeeds

```python
import dask.array as da
import numpy as np

x = np.array([[1,1],[2,2]])
xx = da.from_array(x, chunks=1)  # create a multi-partition dask array

# Two blockwise operations in a row!
xx = xx * 2
z= da.matmul(xx,xx)

# optimize_graph=False turns off other optimizations, but I've manually verified
# that optimize_blockwise it the culprit
z.compute(optimize_graph=False) 

# produces the right answer
# [[12,12], [24,24]]
```

But if I turn on graph optimization (and `optimize_blockwise` does seem to be the important one), I get the wrong answer:

```python
import dask.array as da
import numpy as np

x = np.array([[1,1],[2,2]])
xx = da.from_array(x, chunks=1)  # create a multi-partition dask array

# Two blockwise operations in a row!
xx = xx * 2
z= da.matmul(xx,xx)

z.compute(optimize_graph=True) 

# produces the wrong answer
# [[8,8], [32,32]]
```

It does seem to be related to `xx` appearing twice. If I create a second array with a *different* name, but which is otherwise identical, the snippet succeeds. I'm a bit at a loss as to what might be happening right now, so this is a bit of an in-progress investigation that I'm putting out in case someone with fresher eyes has an idea. @gjoseph92 especially might be interested.

**Environment**:

- Dask version: `main`
- Python version: `3.93
- Operating System: ubuntu
- Install method (conda, pip, source): source


Discussion:
Hmm - Thanks for raising this. I can certainly reproduce

>It does seem to be related to xx appearing twice. If I create a second array with a different name, but which is otherwise identical, the snippet succeeds.

When you say a different name, do you mean that you modify `xx._name`? I still get the error when a distinct variable name is used (because the collection name is still the same):

```
xx = da.from_array(np.array([[1,1],[2,2]]), chunks=1)
yy = da.from_array(np.array([[1,1],[2,2]]), chunks=1)
assert xx._name == yy._name  # True
```
> When you say a different name, do you mean that you modify `xx._name`? I still get the error when a distinct variable name used (because the collection name is still the same):

Exactly -- I don't need to modify the name manually, I can pass it in as a kwarg, but if I give them different names as opposed to the deterministic tokenization, I get the right answer again:


```python
xx = da.from_array(np.array([[1,1],[2,2]]), name="xx", chunks=1)
yy = da.from_array(np.array([[1,1],[2,2]]), name="yy", chunks=1)
```


Also, thanks for taking a look @rjzamora 
This indeed seems to be a logic bug in optimize_blockwise. It looks like the matmul layer (`Blockwise<(('mul-bd4187c86bf5b3023c3b8ab1f2397129', ('.0', '.1')), ('mul-bd4187c86bf5b3023c3b8ab1f2397129', ('.1', '.2'))) -> matmul-3e94acba0bdfaf326a7220437b4448e8>`) originally targets two distinct `indices`, but these distinct indices become a single result in the fused graph... still need to figure out how to fix it.

Resolving patch:
diff --git a/dask/blockwise.py b/dask/blockwise.py
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -1300,6 +1300,11 @@ def _optimize_blockwise(full_graph, keys=()):
     return HighLevelGraph(out, dependencies)
 
 
+def _unique_dep(dep, ind):
+    # Append blockwise index information to dependency name
+    return dep + "_" + "_".join(str(i) for i in list(ind))
+
+
 def rewrite_blockwise(inputs):
     """Rewrite a stack of Blockwise expressions into a single blockwise expression
 
@@ -1355,9 +1360,15 @@ def rewrite_blockwise(inputs):
 
             changed = True
 
+            # Change dep name to avoid fusing the same dep
+            # (in different iteration orders) into a single
+            # subgraph key/dependency
+            # (see: https://github.com/dask/dask/issues/8535)
+            local_dep = dep if dep == root else _unique_dep(dep, ind)
+
             # Replace _n with dep name in existing tasks
             # (inc, _0) -> (inc, 'b')
-            dsk = {k: subs(v, {blockwise_token(i): dep}) for k, v in dsk.items()}
+            dsk = {k: subs(v, {blockwise_token(i): local_dep}) for k, v in dsk.items()}
 
             # Remove current input from input indices
             # [('a', 'i'), ('b', 'i')] -> [('a', 'i')]
@@ -1401,6 +1412,10 @@ def rewrite_blockwise(inputs):
                     indices.append(index)
             new_dsk = subs(inputs[dep].dsk, sub)
 
+            # Change new_dsk key to match local_dep
+            if dep != local_dep and dep in new_dsk:
+                new_dsk[local_dep] = new_dsk.pop(dep)
+
             # indices.extend(new_indices)
             dsk.update(new_dsk)
 


Test patch:
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -10,6 +10,7 @@
 from dask.blockwise import (
     _BLOCKWISE_DEFAULT_PREFIX,
     Blockwise,
+    _unique_dep,
     index_subs,
     optimize_blockwise,
     rewrite_blockwise,
@@ -35,21 +36,40 @@
                 (c, "i", {c: (dec, _0)}, [(a, "i")]),
                 (d, "i", {d: (add, _0, _1, _2)}, [(a, "i"), (b, "i"), (c, "i")]),
             ],
-            (d, "i", {b: (inc, _0), c: (dec, _0), d: (add, _0, b, c)}, [(a, "i")]),
+            (
+                d,
+                "i",
+                {
+                    _unique_dep(b, "i"): (inc, _0),
+                    _unique_dep(c, "i"): (dec, _0),
+                    d: (add, _0, _unique_dep(b, "i"), _unique_dep(c, "i")),
+                },
+                [(a, "i")],
+            ),
         ],
         [
             [
                 (b, "i", {b: (inc, _0)}, [(a, "i")]),
                 (c, "j", {c: (inc, _0)}, [(b, "j")]),
             ],
-            (c, "j", {b: (inc, _0), c: (inc, b)}, [(a, "j")]),
+            (
+                c,
+                "j",
+                {_unique_dep(b, "j"): (inc, _0), c: (inc, _unique_dep(b, "j"))},
+                [(a, "j")],
+            ),
         ],
         [
             [
                 (b, "i", {b: (sum, _0)}, [(a, "ij")]),
                 (c, "k", {c: (inc, _0)}, [(b, "k")]),
             ],
-            (c, "k", {b: (sum, _0), c: (inc, b)}, [(a, "kA")]),
+            (
+                c,
+                "k",
+                {_unique_dep(b, "k"): (sum, _0), c: (inc, _unique_dep(b, "k"))},
+                [(a, "kA")],
+            ),
         ],
         [
             [
@@ -60,7 +80,11 @@
             (
                 g,
                 "ij",
-                {g: (add, c, d), c: (inc, _0), d: (inc, _1)},
+                {
+                    g: (add, _unique_dep(c, "i"), _unique_dep(d, "j")),
+                    _unique_dep(c, "i"): (inc, _0),
+                    _unique_dep(d, "j"): (inc, _1),
+                },
                 [(a, "i"), (b, "j")],
             ),
         ],
@@ -69,14 +93,27 @@
                 (b, "ji", {b: (np.transpose, _0)}, [(a, "ij")]),
                 (c, "ij", {c: (add, _0, _1)}, [(a, "ij"), (b, "ij")]),
             ],
-            (c, "ij", {c: (add, _0, b), b: (np.transpose, _1)}, [(a, "ij"), (a, "ji")]),
+            (
+                c,
+                "ij",
+                {
+                    c: (add, _0, _unique_dep(b, "ij")),
+                    _unique_dep(b, "ij"): (np.transpose, _1),
+                },
+                [(a, "ij"), (a, "ji")],
+            ),
         ],
         [
             [
                 (c, "i", {c: (add, _0, _1)}, [(a, "i"), (b, "i")]),
                 (d, "i", {d: (inc, _0)}, [(c, "i")]),
             ],
-            (d, "i", {d: (inc, c), c: (add, _0, _1)}, [(a, "i"), (b, "i")]),
+            (
+                d,
+                "i",
+                {d: (inc, _unique_dep(c, "i")), _unique_dep(c, "i"): (add, _0, _1)},
+                [(a, "i"), (b, "i")],
+            ),
         ],
         [
             [
@@ -86,7 +123,10 @@
             (
                 d,
                 "ij",
-                {d: (np.dot, b, _0), b: (np.transpose, _1)},
+                {
+                    d: (np.dot, _unique_dep(b, "ik"), _0),
+                    _unique_dep(b, "ik"): (np.transpose, _1),
+                },
                 [(c, "kj"), (a, "ki")],
             ),
         ],
@@ -99,7 +139,11 @@
             (
                 g,
                 "i",
-                {g: (add, c, f), f: (add, _2, _3), c: (add, _0, _1)},
+                {
+                    g: (add, _unique_dep(c, "i"), _unique_dep(f, "i")),
+                    _unique_dep(f, "i"): (add, _2, _3),
+                    _unique_dep(c, "i"): (add, _0, _1),
+                },
                 [(a, i), (b, i), (d, i), (e, i)],
             ),
         ],
@@ -112,7 +156,11 @@
             (
                 g,
                 "i",
-                {g: (add, c, f), f: (add, _0, _2), c: (add, _0, _1)},
+                {
+                    g: (add, _unique_dep(c, "i"), _unique_dep(f, "i")),
+                    _unique_dep(f, "i"): (add, _0, _2),
+                    _unique_dep(c, "i"): (add, _0, _1),
+                },
                 [(a, "i"), (b, "i"), (e, "i")],
             ),
         ],
@@ -121,14 +169,24 @@
                 (b, "i", {b: (sum, _0)}, [(a, "ij")]),
                 (c, "i", {c: (inc, _0)}, [(b, "i")]),
             ],
-            (c, "i", {c: (inc, b), b: (sum, _0)}, [(a, "iA")]),
+            (
+                c,
+                "i",
+                {c: (inc, _unique_dep(b, "i")), _unique_dep(b, "i"): (sum, _0)},
+                [(a, "iA")],
+            ),
         ],
         [
             [
                 (c, "i", {c: (inc, _0)}, [(b, "i")]),
                 (d, "i", {d: (add, _0, _1, _2)}, [(a, "i"), (b, "i"), (c, "i")]),
             ],
-            (d, "i", {d: (add, _0, _1, c), c: (inc, _1)}, [(a, "i"), (b, "i")]),
+            (
+                d,
+                "i",
+                {d: (add, _0, _1, _unique_dep(c, "i")), _unique_dep(c, "i"): (inc, _1)},
+                [(a, "i"), (b, "i")],
+            ),
         ],
         # Include literals
         [
@@ -143,7 +201,7 @@
             (
                 c,
                 "j",
-                {b: (add, _1, _2), c: (add, b, _0)},
+                {_unique_dep(b, "j"): (add, _1, _2), c: (add, _unique_dep(b, "j"), _0)},
                 [(456, None), (a, "j"), (123, None)],
             ),
         ],
@@ -156,7 +214,7 @@
             (
                 c,
                 "j",
-                {b: (add, _1, _2), c: (add, b, _0)},
+                {_unique_dep(b, "j"): (add, _1, _2), c: (add, _unique_dep(b, "j"), _0)},
                 [(False, None), (a, "j"), (0, None)],
             ),
         ],
@@ -166,7 +224,65 @@
                 (b, "i", {b: (add, _0, _1)}, [(a, "i"), (123, None)]),
                 (c, "j", {c: (add, _0, _1)}, [(b, "j"), (123, None)]),
             ],
-            (c, "j", {b: (add, _1, _0), c: (add, b, _0)}, [(123, None), (a, "j")]),
+            (
+                c,
+                "j",
+                {_unique_dep(b, "j"): (add, _1, _0), c: (add, _unique_dep(b, "j"), _0)},
+                [(123, None), (a, "j")],
+            ),
+        ],
+        # Check cases where two distinct indices are used
+        # for the same dependency name, and where the same
+        # dependency-index combination is repeated
+        # (See: https://github.com/dask/dask/issues/8535)
+        [
+            [
+                (b, "jk", {b: (add, _0, _1)}, [(a, "jk"), (2, None)]),
+                (c, "ijk", {c: (add, _0, _1)}, [(b, "ij"), (b, "jk")]),
+                (d, "ijk", {d: (inc, _0, _1)}, [(c, "ijk"), (123, None)]),
+            ],
+            (
+                "d",
+                "ijk",
+                {
+                    "d": (inc, _unique_dep(c, "ijk"), _0),
+                    _unique_dep(c, "ijk"): (
+                        add,
+                        _unique_dep(b, "ij"),
+                        _unique_dep(b, "jk"),
+                    ),
+                    _unique_dep(b, "ij"): (add, _1, _2),
+                    _unique_dep(b, "jk"): (add, _3, _2),
+                },
+                [(123, None), (a, "ij"), (2, None), (a, "jk")],
+            ),
+        ],
+        [
+            [
+                (b, "jk", {b: (add, _0, _1)}, [(a, "jk"), (2, None)]),
+                (c, "ijk", {c: (add, _0, _1)}, [(b, "ij"), (b, "jk")]),
+                (d, "ijk", {d: (add, _0, _1, _2)}, [(b, "ij"), (c, "ij"), (b, "ij")]),
+            ],
+            (
+                "d",
+                "ijk",
+                {
+                    "d": (
+                        add,
+                        _unique_dep(b, "ij"),
+                        _unique_dep(c, "ij"),
+                        _unique_dep(b, "ij"),
+                    ),
+                    _unique_dep(c, "ij"): (
+                        add,
+                        _unique_dep(b, "ij"),
+                        _unique_dep(b, "jk"),
+                    ),
+                    _unique_dep(b, "ij"): (add, _0, _1),
+                    _unique_dep(b, "jk"): (add, _2, _1),
+                },
+                [(a, "ij"), (2, None), (a, "jk")],
+            ),
         ],
     ],
 )
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -455,3 +455,16 @@ def test_fuse_roots_annotations():
     assert {"foo": "bar"} in [l.annotations for l in hlg.layers.values()]
     za = da.Array(hlg, z.name, z.chunks, z.dtype)
     assert_eq(za, z)
+
+
+@pytest.mark.parametrize("optimize_graph", [True, False])
+def test_optimize_blockwise_duplicate_dependency(optimize_graph):
+    # Two blockwise operations in a row with duplicate name
+    # (See: https://github.com/dask/dask/issues/8535)
+    xx = da.from_array(np.array([[1, 1], [2, 2]]), chunks=1)
+    xx = xx * 2
+    z = da.matmul(xx, xx)
+
+    # Compare to known answer
+    result = z.compute(optimize_graph=optimize_graph)
+    assert assert_eq(result, [[12, 12], [24, 24]])