pimm.export.remap_state_dict_keys#

remap_state_dict_keys(state_dict: Dict[str, Any], key_mapping: Dict[str, str], *, keep_unmapped: bool = True) → Dict[str, Any][source]#

Remap state dict keys according to mapping.

Parameters:

  • state_dict – Source state dictionary.

  • key_mapping – Dictionary mapping old keys or prefixes to new keys or prefixes. Exact key matches take priority. Prefix mappings are applied when a state-dict key starts with the source string.

  • keep_unmapped – If True, preserve keys that do not match any mapping.

Returns:

Remapped state dictionary.