Repository: dask/dask

Issue:
ShareDict seems to cause issues with HighLevelGraphs in optimize_atop
Following up on https://github.com/dask/dask/issues/4038#issuecomment-434338006, this script:

```python
from dask.sharedict import ShareDict
import numpy as np


def foo(A):
    return A[None, ...]


A = da.ones(shape=(10, 20, 4), chunks=(2, 5, 4))

name = 'D-' + dask.base.tokenize(A)

dsk = da.core.top(
        foo, name, ("nsrc", "ntime", "nbl", "npol"),
        A.name, ("ntime", "nbl", "npol"),
        new_axes={"nsrc": 1},
        numblocks={a.name: a.numblocks for a in (A,)})

array_dsk = ShareDict()
array_dsk.update(dsk)
array_dsk.update(A.__dask_graph__())

chunks = ((1,),) + A.chunks

D = da.Array(array_dsk, name, chunks, dtype=A.dtype)
print(D.compute().shape)    # Succeeds
D = D.sum(axis=0)
print(D.compute().shape)    # Fails
```

causes the following error:


```traceback
$ python dask_top_2.py 
(1, 10, 20, 4)
Traceback (most recent call last):
  File "dask_top_2.py", line 32, in <module>
    pprint(D.compute().shape)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 155, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 389, in compute
    dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 193, in collections_to_dsk
    for opt, (dsk, keys) in groups.items()]))
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/array/optimization.py", line 38, in optimize
    dsk = optimize_atop(dsk, keys=keys)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/array/top.py", line 594, in optimize_atop
    if (isinstance(layers[dep], TOP) and
KeyError: 'D-122ec2f0f994205583cdc43c8fba6f6d'
```


Discussion:
python 2.7.14
dask master branch
Thanks for raising this issue @sjperkins .  Checking in, is everything above necessary in order to reproduce your error?  If not then can I ask you to simplify things first please?  This will really help to isolate where the problem is.
@mrocklin  I've updated the example.
I think this is ShareDict related, the above example works with a normal dict.

Resolving patch:
diff --git a/dask/array/top.py b/dask/array/top.py
--- a/dask/array/top.py
+++ b/dask/array/top.py
@@ -591,7 +591,8 @@ def optimize_atop(full_graph, keys=()):
             deps = set(top_layers)
             while deps:  # we gather as many sub-layers as we can
                 dep = deps.pop()
-                if (isinstance(layers[dep], TOP) and
+                if (dep in layers and
+                        isinstance(layers[dep], TOP) and
                         not (dep != layer and dep in keep) and  # output layer
                         layers[dep].concatenate == layers[layer].concatenate):  # punt on mixed concatenate
                     top_layers.add(dep)


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
@@ -380,3 +380,28 @@ def f(data, x):
                 dtype=A.dtype)
 
     assert_eq(A, B)
+
+
+def test_gh_4176():
+    def foo(A):
+        return A[None, ...]
+
+    A = da.ones(shape=(10, 20, 4), chunks=(2, 5, 4))
+
+    name = 'D'
+
+    dsk = da.core.top(
+        foo, name, ("nsrc", "ntime", "nbl", "npol"),
+        A.name, ("ntime", "nbl", "npol"),
+        new_axes={"nsrc": 1},
+        numblocks={a.name: a.numblocks for a in (A,)}
+    )
+
+    array_dsk = dask.sharedict.ShareDict()
+    array_dsk.update(dsk)
+    array_dsk.update(A.__dask_graph__())
+
+    chunks = ((1,),) + A.chunks
+
+    D = da.Array(array_dsk, name, chunks, dtype=A.dtype)
+    D.sum(axis=0).compute()