from dask.highlevelgraph import HighLevelGraph
import dask
from dask.threaded import get
def add(x, y):
  return x + y
 
def square(x):
  return x ** 2
 
 
layers = {
  'add': {
    ('add', 0): (add, 1, 2),
    ('add', 1): (add, 3, 4)
  },
  'square': {
    ('square', 0): (square, ('add', 0)),
    ('square', 1): (square, ('add', 1))
  },
  'agg': {
    'agg': (add, ('square', 0), ('square', 1))}
}
 
dependencies = {
  'add': set(),
  'square': {'add'},
  'agg': {'square'}
}
 
g = HighLevelGraph(layers, dependencies)
 
get(g, 'agg')
58
g = HighLevelGraph(layers, dependencies)
dask.visualize(g)

png

dask.visualize(graph)

png