parameters: (dag, sort_fn)
return_statement: return sorted_nodes
docstring: 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.
docstring_summary: Topologically sorts the given dag.
docstring_tokens: Topologically sorts the given dag .
function: 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
function_tokens: def toposort ( dag , sort_fn ) : # 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