def _get_op_to_stage_map(stages):
    """Given a list of stages, returns a map from each op in each
    stage to the encompassing stage."""
    op_to_stage = {}
    for stage in stages:
        for op in stage.ops:
            op_to_stage[op] = stage
    return op_to_stage