def test_avoid_upwards_branching(abcde):
        """
             a1
             |
             a2
             |
             a3    d1
            /  \  /
          b1    c1
          |     |
          b2    c2
                |
                c3
    
        Prefer b1 over c1 because it won't stick around waiting for d1 to complete
        """
        a, b, c, d, e = abcde
        dsk = {(a, 1): (f, (a, 2)),
               (a, 2): (f, (a, 3)),
               (a, 3): (f, (b, 1), (c, 1)),
               (b, 1): (f, (b, 2)),
               (c, 1): (f, (c, 2)),
               (c, 2): (f, (c, 3)),
               (d, 1): (f, (c, 1))}
    
        o = order(dsk)    

import matplotlib.pyplot as plt

import dask.array as da
n = 10
x  = da.random.normal(size=(n, 100), chunks=(1, 100))
y  = da.random.normal(size=(n,), chunks=(1,))
xy = (x * y[:, None]).cumsum(axis=0)
xx = (x[:, None, :] * x[:, :, None]).cumsum(axis=0)
beta = da.stack([da.linalg.solve(xx[i], xy[i]) for i in range(xx.shape[0])],
                axis=0)
ey = (x * beta).sum(axis=1)
ey.visualize('dask.png', color='order', node_attr=dict(penwidth='8'),
             optimize_graph=True, cmap=plt.cm.RdBu)


A, B = 10, 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))

w.visualize('gh3055.png', color='order', node_attr=dict(penwidth='8'),
            cmap=plt.cm.RdBu)

%%time
n = 100
x  = da.random.normal(size=(n, 100), chunks=(1, 100))
y  = da.random.normal(size=(n,), chunks=(1,))
xy = (x * y[:, None]).cumsum(axis=0)
xx = (x[:, None, :] * x[:, :, None]).cumsum(axis=0)
beta = da.stack([da.linalg.solve(xx[i], xy[i]) for i in range(xx.shape[0])],
                axis=0)
ey = (x * beta).sum(axis=1)
CPU times: user 239 ms, sys: 4.21 ms, total: 243 ms
Wall time: 233 ms

%time dsk = dict(ey.__dask_graph__())
CPU times: user 74.1 ms, sys: 0 ns, total: 74.1 ms
Wall time: 73.2 ms

%timeit _ = order(dsk)
10 loops, best of 3: 25.9 ms per loop

%timeit _ = order(dsk)
10 loops, best of 3: 39.3 ms per loop