def _split_stages(node, duplicates=None, aliases=None, stages=None, parents=None):

    # init data structures
    if duplicates is None:
        duplicates = dict()
        aliases = set()
        stages = list()
        parents = list()

    if type(node) is list:

        # don't count assignment node as a parent,
        # it will always exist in the final stage which is processed outside of
        # this function
        if node[0][0] != "assign":
            parents.append(node)

        # post order traversal (pulls the stages deepest in the tree first)
        if len(node) > 3:
            _split_stages(node[3], duplicates, aliases, stages, parents)
        if len(node) > 4:
            _split_stages(node[4], duplicates, aliases, stages, parents)

        if len(parents) > 0:
            parents.pop()

        if node[0][0] in _reduction_ops:

            red_stack = _process_node(node, aliases, duplicates)
            if red_stack:
                # add this reduction stack to the stages
                stages.append(("reduction", red_stack))

            # decrement reduction count for all parents
            for parent in parents:
                parent[2] -= 1

            # walk up the parent list
            # TODO: potentially do this iteratively to find longest common set
            # of operations
            scalar_parent = None
            for parent in parents[::-1]:
                # find the highest parent that is both scalar and has no other
                # child reductions
                if parent[1] and parent[2] == 0:
                    scalar_parent = parent
                else:
                    break

            # if there are any scalar operations over this reduction, remove
            # them from the tree as well
            if scalar_parent is not None:

                scalar_stack = _process_node(
                    scalar_parent, aliases, duplicates)
                if scalar_stack:
                    # add this scalar stack to the stages
                    stages.append(("scalar", scalar_stack))

    return stages
Split out all reductions and post reduction scalar operations into seperate
dataset/ETHPy150Open NervanaSystems/neon/neon/backends/float_ew.py/_split_stages