dask visualize() KeyError - graphviz

Question

I have been successfully visualizing dask objects for a few days now, but I just restarted my Jupyter notebook kernel and am running into a KeyError that I can't seem to debug.
Before I restarted the kernel, the following code worked fine:
def sigmoid(x):
'''Sigmoid function of x.'''
return 1/(1+da.exp(-x))
X = da.from_array(np.random.random((4,2)), chunks=2)
beta = np.ones((2,1))
##
p = (X.dot(beta)).map_blocks(sigmoid)
W = da.diag((p*(1-p))[:,0])
hessian = da.dot(X.T.dot(W),X)
hessian.visualize()
But now when I run, I get the following traceback:
KeyError Traceback (most recent call last)
<ipython-input-62-660a47cb4654> in <module>()
7 W = da.diag((p*(1-p))[:,0])
8 hessian = dot(X.T.dot(W),X)
----> 9 hessian.visualize()
.../anaconda3/anaconda/lib/python3.5/site-packages/dask/base.py in visualize(self, filename, format, optimize_graph, **kwargs)
59 """
60 return visualize(self, filename=filename, format=format,
---> 61 optimize_graph=optimize_graph, **kwargs)
62
63 def compute(self, **kwargs):
.../anaconda3/anaconda/lib/python3.5/site-packages/dask/base.py in visualize(*args, **kwargs)
234 dsk = merge(dsks)
235
--> 236 return dot_graph(dsk, filename=filename, **kwargs)
237
238
.../anaconda3/anaconda/lib/python3.5/site-packages/dask/dot.py in dot_graph(dsk, filename, format, **kwargs)
217 dask.dot.to_graphviz
218 """
--> 219 g = to_graphviz(dsk, **kwargs)
220
221 fmts = ['.png', '.pdf', '.dot', '.svg', '.jpeg', '.jpg']
.../anaconda3/anaconda/lib/python3.5/site-packages/dask/dot.py in to_graphviz(dsk, data_attributes, function_attributes, rankdir, graph_attr, node_attr, edge_attr, **kwargs)
130 if func_name not in seen:
131 seen.add(func_name)
--> 132 g.node(func_name, label=task_label(v), shape='circle',
133 **function_attributes.get(k, {}))
134 g.edge(func_name, k_name)
.../anaconda3/anaconda/lib/python3.5/site-packages/dask/dot.py in task_label(task)
29 if hasattr(func, 'funcs'):
30 if len(func.funcs) > 1:
---> 31 return '{0}(...)'.format(funcname(func.funcs[0]))
32 else:
33 head = funcname(func.funcs[0])
KeyError: 0
It appears there is some label missing for one of the tasks or something like that?
EDIT: Actually, to reproduce the error you need to modify the sigmoid function as follows:
from multipledispatch import dispatch
#dispatch(da.Array)
def sigmoid(x):
'''Sigmoid function of x.'''
return 1/(1+da.exp(-x))
#dispatch(np.ndarray)
def sigmoid(x):
'''Sigmoid function of x.'''
return 1/(1+np.exp(-x))

Answer 1

Thanks for the reproducible example.
Unfortunately, this runs fine for me:
In [1]: import numpy as np
In [2]: import dask.array as da
In [3]: from multipledispatch import dispatch
...:
...: #dispatch(da.Array)
...: def sigmoid(x):
...: '''Sigmoid function of x.'''
...: return 1/(1+da.exp(-x))
...:
...: #dispatch(np.ndarray)
...: def sigmoid(x):
...: '''Sigmoid function of x.'''
...: return 1/(1+np.exp(-x))
...:
In [4]: pdb
Automatic pdb calling has been turned ON
In [5]: def sigmoid(x):
...: '''Sigmoid function of x.'''
...: return 1/(1+da.exp(-x))
...:
...: X = da.from_array(np.random.random((4,2)), chunks=2)
...: beta = np.ones((2,1))
...:
...: ##
...: p = (X.dot(beta)).map_blocks(sigmoid)
...: W = da.diag((p*(1-p))[:,0])
...: hessian = da.dot(X.T.dot(W),X)
...: hessian.visualize()
...:
Out[5]: <IPython.core.display.Image object>
The fact that it worked until you restarted something is a bit odd. Perhaps something is up with your environment?