def toposort(dag, sort_fn):
  """Topologically sorts the given dag.

  Arguments:
    - dag: a dict mapping from a node to its incoming nodes.
    - sort_fn: a function for sorting nodes in the same batch.

  Returns:
    A list containing topologically sorted nodes.
  """

  # Returns the next batch of nodes without incoming edges
  def get_next_batch(dag):
    while True:
      no_prev_nodes = set(node for node, prev in dag.items() if not prev)
      if not no_prev_nodes:
        break
      yield sorted(no_prev_nodes, key=sort_fn)
      dag = {
          node: (prev - no_prev_nodes)
          for node, prev in dag.items()
          if node not in no_prev_nodes
      }
    assert not dag, 'found cyclic dependency'

  sorted_nodes = []
  for batch in get_next_batch(dag):
    sorted_nodes.extend(batch)

  return sorted_nodes