xlstm_jax.utils.model_param_handling.convert_state_dict

xlstm_jax.utils.model_param_handling.convert_state_dict#

Attributes#

LOGGER

Functions#

find_parameter_match_key(from_key, to_keys, match_dict)

Finds the matching parameter key for the target state dict.

create_full_state_dict_key_mapping(from_state_dict, ...)

Creates a full state dict key mapping from the source state dict to the target state dict.

move_state_dict_params_(from_state_dict, ...)

Move the params of a model from one state dict to another state dict.

move_safetensors_state_dict_params_(...)

Move the params of a model from one state dict to another state dict.

convert_state_dict_keys_(state_dict, full_key_mapping)

Converts the keys of the state dict according to the key mapping.

apply_weight_transforms_(state_dict, ...)

Applies weight transforms to the weights of the state dict.

Module Contents#

xlstm_jax.utils.model_param_handling.convert_state_dict.LOGGER#
xlstm_jax.utils.model_param_handling.convert_state_dict.find_parameter_match_key(from_key, to_keys, match_dict)#

Finds the matching parameter key for the target state dict.

Parameters:

  • from_key (str) – The key of the source state dict.

  • to_keys (list[str]) – The keys of the target state dict.

  • match_dict (dict[str, str]) – The dict that maps the source state dict keys to the target state dict keys. Should contain unique substrings of the source state dict keys as keys and the corresponding target state dict keys substrings as values.

Returns:

The target state dict key that matches the source state dict key.

Return type:

str

xlstm_jax.utils.model_param_handling.convert_state_dict.create_full_state_dict_key_mapping(from_state_dict, to_state_dict, match_dict)#

Creates a full state dict key mapping from the source state dict to the target state dict.

Parameters:

  • from_state_dict (dict) – The source state dict.

  • to_state_dict (dict) – The target state dict.

  • match_dict (dict) – The dict that maps the source state dict keys to the target state dict keys.

Returns:

The full state dict key mapping from the source state dict to the target state dict.

Return type:

dict

xlstm_jax.utils.model_param_handling.convert_state_dict.move_state_dict_params_(from_state_dict, to_state_dict, match_dict)#

Move the params of a model from one state dict to another state dict. Modifies the to_state_dict in place.

Parameters:

  • from_state_dict (dict[str, Any]) – The source state dict.

  • to_state_dict (dict[str, Any]) – The target state dict.

  • match_dict (dict[str, str]) – The dict that maps the source state dict keys to the target state dict keys. Should contain unique substrings of the source state dict keys as keys and the corresponding target state dict keys substrings as values.

Returns:

The target (modified to_state_dict) state dict with the converted parameters.

Return type:

dict[str, Any]

xlstm_jax.utils.model_param_handling.convert_state_dict.move_safetensors_state_dict_params_(from_state_dict_path, to_state_dict, match_dict)#

Move the params of a model from one state dict to another state dict. It loads the from_state_dict from a file on-the-fly. This means only the to_state_dict is in memory. Modifies the to_state_dict in place.

Parameters:

  • from_state_dict (Path) – The path to the source state dict.

  • to_state_dict (dict[str, Any]) – The target state dict.

  • match_dict (dict[str, str]) – The dict that maps the source state dict keys to the target state dict keys. Should contain unique substrings of the source state dict keys as keys and the corresponding target state dict keys substrings as values.

  • from_state_dict_path (dict[str, Any])

Returns:

The target (modified to_state_dict) state dict with the converted parameters.

Return type:

dict[str, Any]

xlstm_jax.utils.model_param_handling.convert_state_dict.convert_state_dict_keys_(state_dict, full_key_mapping)#

Converts the keys of the state dict according to the key mapping.

Parameters:

  • state_dict (dict[str, Any]) – The state dict to convert.

  • full_key_mapping (dict[str, str]) – The key mapping that maps the old keys to the new keys.

Returns:

The converted state dict with the new keys.

Return type:

dict[str, Any]

xlstm_jax.utils.model_param_handling.convert_state_dict.apply_weight_transforms_(state_dict, apply_transforms_to_keys)#

Applies weight transforms to the weights of the state dict.

There are currently these transforms supported:

  • “transpose”: Transposes the weight tensor. Accepts only 2D tensors.

  • “squeeze-XXX”: Squeezes the XXX dimension of the weight tensor.

    If XXX is not given, squeezes all dimensions of size 1.

  • “flatten”: Flattens the weight tensor.

If possible the transforms are applied in-place on tensors. Also the state_dict is modified in-place.

Parameters:

  • state_dict (dict[str, torch.Tensor]) – The state dict with the weights.

  • apply_transforms_to_keys (dict[str, list[str]]) – The dict that maps the transform to the keys of the state dict.

Returns:

The state dict with the transformed weights

Return type:

dict[str, torch.Tensor]