def parse_head_pruning_descriptors(
    descriptors,
    reverse_descriptors=False,
    n_heads=None
):
    """Returns a dictionary mapping layers to the set of heads to prune in
    this layer (for each kind of attention)"""
    to_prune = {
        "E": {},
        "A": {},
        "D": {},
    }
    for descriptor in descriptors:
        attn_type, layer, heads = descriptor.split(":")
        layer = int(layer) - 1
        heads = set(int(head) - 1 for head in heads.split(","))
        if layer not in to_prune[attn_type]:
            to_prune[attn_type][layer] = set()
        to_prune[attn_type][layer].update(heads)
    # Reverse
    if reverse_descriptors:
        if n_heads is None:
            raise ValueError("You need to specify the total number of heads")
        for attn_type in to_prune:
            for layer, heads in to_prune[attn_type].items():
                to_prune[attn_type][layer] = set([head for head in range(n_heads)
                                                  if head not in heads])
    return to_prune