examples/research_projects/jax-projects/model_parallel/partitions.py/_replacement_rules
def _replacement_rules(rules):
    def replace(key, val):
        for rule, replacement in rules:
            if _match(rule, key):
                return replacement
        return val

    return replace