Building off autograd's reverse toposort.

toposort.py

import operator

def toposort(start_node, children=operator.attrgetter('children')):
  """Generate nodes in DAG's topological order. This guarantees
  for any edge U -> V, we always visit U before visiting V.

  This lets us play the tape via the "push" dataflow model.
  https://stackoverflow.com/questions/981027/what-are-pro-cons-of-push-pull-data-flow-models
  """
  parent_counts = {}
  stack = [start_node]
  while stack:
    node = stack.pop()
    if node in parent_counts:
      parent_counts[node] += 1
    else:
      parent_counts[node] = 1
      stack.extend(children(node))

  parentless_nodes = [start_node]
  while parentless_nodes:
    node = parentless_nodes.pop()
    yield node
    for child in children(node):
      if parent_counts[child] == 1:
        parentless_nodes.append(child)
      else:
        parent_counts[child] -= 1

def reverse_toposort(end_node, parents=operator.attrgetter('parents')):
  """Generate nodes in DAG's reverse topological order. This guarantees
  for any edge U -> V, we always visit V before visiting U.

  This lets us play the tape via the "pull" dataflow model.
  https://stackoverflow.com/questions/981027/what-are-pro-cons-of-push-pull-data-flow-models
  """
  child_counts = {}
  stack = [end_node]
  while stack:
    node = stack.pop()
    if node in child_counts:
      child_counts[node] += 1
    else:
      child_counts[node] = 1
      stack.extend(parents(node))

  childless_nodes = [end_node]
  while childless_nodes:
    node = childless_nodes.pop()
    yield node
    for parent in parents(node):
      if child_counts[parent] == 1:
        childless_nodes.append(parent)
      else:
        child_counts[parent] -= 1