Repository: dask/dask

Issue:
Dask array with non-highlevelgraph leads to KeyError
I have some [older code](https://github.com/ska-sa/katsdpcal/blob/5caa806649de33324977727a5ca1b5b76d61ba08/katsdpcal/scan.py#L86-L92) that assigns a plain dict directly into `array.dask`, as a workaround for an old dask bug. If that's no longer supported then feel free to close this bug (although it would be nice to add one or two sanity checks). Since dask 2.8.0 it raises a KeyError inside the optimizer.

Here's a minimum reproducing example:
```python
#!/usr/bin/env python3
import dask
import dask.array as da
import numpy as np

a = da.from_array(np.ones((4, 4), np.float64), chunks=(2, 4))
a.dask = dict(a.dask)
b = da.from_array(np.zeros((4, 4), np.float64), chunks=(2, 4))

x = a + b
da.compute(x)
```

```
Traceback (most recent call last):
  File "./bw.py", line 11, in <module>
    da.compute(x)
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/base.py", line 433, in compute
    dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/base.py", line 219, in collections_to_dsk
    [opt(dsk, keys, **kwargs) for opt, (dsk, keys) in groups.items()],
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/base.py", line 219, in <listcomp>
    [opt(dsk, keys, **kwargs) for opt, (dsk, keys) in groups.items()],
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/array/optimization.py", line 43, in optimize
    dsk = fuse_roots(dsk, keys=keys)
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/blockwise.py", line 822, in fuse_roots
    and not any(dependencies[dep] for dep in deps)  # no need to fuse if 0 or 1
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/blockwise.py", line 822, in <genexpr>
    and not any(dependencies[dep] for dep in deps)  # no need to fuse if 0 or 1
KeyError: 'array-d6b7f7b33e2b14cd720de87a517147e8'
```
(run against dask 2.10.0, Python 3.6)

From what I can tell, the trouble starts when HighLevelGraph.from_collections is given a dependency that doesn't use a HighLevelGraph, and it names the resulting layer just using `id`: [here](https://github.com/dask/dask/blob/0274a5c6db618e7af8444562149ecc5774ae22a3/dask/highlevelgraph.py#L134). However, the Blockwise for the addition keys its `indices` based on the array names. Then `optimize_blockwise` uses those `indices` keys to construct the new dependencies, [here](https://github.com/dask/dask/blob/0274a5c6db618e7af8444562149ecc5774ae22a3/dask/blockwise.py#L552). That means the new dependencies don't line up with the layers. The final explosion happens in `fuse_roots`, but it could presumably happen anywhere that expected the dependencies to be consistent.

Discussion:
bmerry: There is a test failure (at least when it was tested on my branch), but it also appears to be failing on master.

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index f6154ffd8..e3c47ff1b 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -4,7 +4,8 @@ repos:
     hooks:
     - id: black
       language_version: python3.7
--   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v2.3.0
+-   repo: https://gitlab.com/pycqa/flake8
+    rev: 3.7.9
     hooks:
     - id: flake8
+      language_version: python3.7
diff --git a/dask/array/core.py b/dask/array/core.py
index 15218b53c..84fcbf620 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -182,13 +182,12 @@ def slices_from_chunks(chunks):
       (slice(2, 4, None), slice(3, 6, None)),
       (slice(2, 4, None), slice(6, 9, None))]
     """
-    cumdims = [list(accumulate(add, (0,) + bds[:-1])) for bds in chunks]
-    shapes = product(*chunks)
-    starts = product(*cumdims)
-    return [
-        tuple(slice(s, s + dim) for s, dim in zip(start, shape))
-        for start, shape in zip(starts, shapes)
+    cumdims = [cached_cumsum(bds, initial_zero=True) for bds in chunks]
+    slices = [
+        [slice(s, s + dim) for s, dim in zip(starts, shapes)]
+        for starts, shapes in zip(cumdims, chunks)
     ]
+    return list(product(*slices))
 
 
 def getem(
diff --git a/dask/blockwise.py b/dask/blockwise.py
index 5fc7cf5d6..28d575a4f 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -581,6 +581,10 @@ def rewrite_blockwise(inputs):
     --------
     optimize_blockwise
     """
+    if len(inputs) == 1:
+        # Fast path: if there's only one input we can just use it as-is.
+        return inputs[0]
+
     inputs = {inp.output: inp for inp in inputs}
     dependencies = {
         inp.output: {d for d, v in inp.indices if v is not None and d in inputs}
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 8c16b8326..1a9ce204f 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -131,7 +131,10 @@ class HighLevelGraph(Mapping):
                     with ignoring(AttributeError):
                         deps[name] |= set(collection.__dask_layers__())
                 else:
-                    key = id(graph)
+                    try:
+                        [key] = collection.__dask_layers__()
+                    except AttributeError:
+                        key = id(graph)
                     layers[key] = graph
                     deps[name].add(key)
                     deps[key] = set()
diff --git a/docs/source/array-design.rst b/docs/source/array-design.rst
index bc10c3390..4cd33add1 100644
--- a/docs/source/array-design.rst
+++ b/docs/source/array-design.rst
@@ -33,7 +33,10 @@ Example
    >>> x.name
    'arange-539766a'
 
-   >>> x.dask  # somewhat simplified
+   >>> x.__dask_graph__()
+   <dask.highlevelgraph.HighLevelGraph at 0x7f9f6f686d68>
+
+   >>> dict(x.__dask_graph__())  # somewhat simplified
    {('arange-539766a', 0): (np.arange, 0, 5),
     ('arange-539766a', 1): (np.arange, 5, 10),
     ('arange-539766a', 2): (np.arange, 10, 15)}
@@ -51,8 +54,10 @@ Keys of the Dask graph
 By special convention, we refer to each block of the array with a tuple of the
 form ``(name, i, j, k)``, with ``i, j, k`` being the indices of the block
 ranging from ``0`` to the number of blocks in that dimension.  The Dask graph
-must hold key-value pairs referring to these keys.  Moreover, it likely also 
-holds other key-value pairs required to eventually compute the desired values:
+must hold key-value pairs referring to these keys.  Moreover, it likely also
+holds other key-value pairs required to eventually compute the desired values
+(usually organised in a :doc:`HighLevelGraph <high-level-graphs>`, but shown
+in a flattened form here for illustration):
 
 .. code-block:: python
 
@@ -103,10 +108,11 @@ Some ways in which ``chunks`` reflects properties of our array:
 Create an Array Object
 ----------------------
 
-In order to create an ``da.Array`` object we need a dictionary with these special
+In order to create an ``da.Array`` object we need a graph with these special
 keys::
 
-    dsk = {('x', 0, 0): ...}
+    layer = {('x', 0, 0): ...}
+    dsk = HighLevelGraph.from_collections('x', layer, dependencies=())
 
 a name specifying which keys this array refers to::
 
@@ -127,7 +133,7 @@ shapes.
 Example - ``eye`` function
 --------------------------
 
-As an example, lets build the ``np.eye`` function for ``dask.array`` to make the
+As an example, let's build the ``np.eye`` function for ``dask.array`` to make the
 identity matrix:
 
 .. code-block:: python
@@ -138,11 +144,12 @@ identity matrix:
 
        name = 'eye' + next(tokens)  # unique identifier
 
-       dsk = {(name, i, j): (np.eye, blocksize)
-                            if i == j else
-                            (np.zeros, (blocksize, blocksize))
+       layer = {(name, i, j): (np.eye, blocksize)
+                              if i == j else
+                              (np.zeros, (blocksize, blocksize))
                 for i in range(n // blocksize)
                 for j in range(n // blocksize)}
+       dsk = dask.highlevelgraph.HighLevelGraph.from_collections(name, layer, dependencies=())
 
        dtype = np.eye(0).dtype  # take dtype default from numpy
 


Test patch:
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
index 1e22396d9..dbac96c26 100644
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -587,3 +587,12 @@ def test_atop_legacy():
     z = da.blockwise(inc, "i", x, "i", dtype=x.dtype)
     assert_eq(y, z)
     assert y.name == z.name
+
+
+def test_non_hlg():
+    # Regression test for https://github.com/dask/dask/issues/5850
+    a = da.from_array(np.ones(1, np.float64), chunks=(1,))
+    a.dask = dict(a.dask)  # Convert from HighLevelGraph to plain dict
+    b = da.from_array(np.zeros(1, np.float64), chunks=(1,))
+    x = a + b
+    assert_eq(x, a)