def create_colocation_group_to_ops_map(op_graph):
    """Generate a dict that maps a colocation group to its op id list."""
    retval = {}

    for op_id, op_data in op_graph.nodes().items():
        # assume there is only one group
        group = op_data['colocation_group']
        if group in retval:
            retval[group].append(op_id)
        else:
            retval[group] = [op_id]

    return retval