import dask.array as da
from colorize import visualize_colors

A, B = 20, 99 
x = da.random.normal(size=(A, B), chunks=(1, None))
for _ in range(2):
    y = (x[:, None, :] * x[:, :, None]).cumsum(axis=0)
    x = x.cumsum(axis=0)
w = (y * x[:, None]).sum(axis=(1, 2))

visualize_colors(w.__dask_graph__(), filename='dask.pdf')