from typing import OrderedDict


def map_layers(weight):
    """ Since the pre-trained weights provided for bls17 by us were trained with
        different layer names, we map the layer names in the state dictionaries
        to the new names using the following function map_layers().
    """
    return OrderedDict([(k.replace('z', 'w'), v) if 'z' in k else (k, v) for k, v in weight.items()])