usrbinenv python codingutf8 2021 the google research s and the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license utilities for constructing pytrees of partitionspecs utils adapted from https github comgoogleresearchgoogleresearchblobmasterflaxmodelst5xpartitions py import re from flax core frozendict import freeze from flax traverseutil import flattendict unflattendict from jax experimental import partitionspec as p sentinels unmatched object for specifying empty leaf dict emptydict object def matchqs ks compile regexes and force complete match partitionspec for gptneo replicate the hidden dim and shard feedforward and head dim embeddings atention mlp layer norms usr bin env python coding utf 8 2021 the google research s and the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license utilities for constructing pytrees of partitionspecs utils adapted from https github com google research google research blob master flax_models t5x partitions py sentinels for specifying empty leaf dict return true if regexes in qs match any window of strings in tuple ks compile regexes and force complete match partitionspec for gptneo replicate the hidden dim and shard feed forward and head dim embeddings atention mlp layer norms
import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P



_unmatched = object()


empty_dict = object()


def _match(qs, ks):
    
    
    qts = tuple((re.compile(x + "$") for x in qs))
    for i in range(len(ks) - len(qs) + 1):
        matches = [x.match(y) for x, y in zip(qts, ks[i:])]
        if matches and all(matches):
            return True
    return False


def _replacement_rules(rules):
    def replace(key, val):
        for rule, replacement in rules:
            if _match(rule, key):
                return replacement
        return val

    return replace




def _get_partition_rules():
    return [
        
        (("transformer", "wpe", "embedding"), P("mp", None)),
        (("transformer", "wte", "embedding"), P("mp", None)),
        
        (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(None, "mp")),
        (("attention", "out_proj", "kernel"), P("mp", None)),
        (("attention", "out_proj", "bias"), None),
        
        (("mlp", "c_fc", "kernel"), P(None, "mp")),
        (("mlp", "c_fc", "bias"), P("mp")),
        (("mlp", "c_proj", "kernel"), P("mp", None)),
        (("mlp", "c_proj", "bias"), None),
        
        ((r"ln_\d+", "bias"), None),
        ((r"\d+", r"ln_\d+", "scale"), None),
        (("ln_f", "bias"), None),
        (("ln_f", "scale"), None),
    ]


def set_partitions(in_dict):
    rules = _get_partition_rules()
    replace = _replacement_rules(rules)
    initd = {k: _unmatched for k in flatten_dict(in_dict)}
    result = {k: replace(k, v) for k, v in initd.items()}
    assert _unmatched not in result.values(), "Incomplete partition spec."
    return freeze(unflatten_dict(result))