def convert_mlstm_layer_weights_(
        state_dict: dict[str, torch.Tensor], block_keys: list[str]
    ):
        """Modifies the state dict in place."""

        # fuse qkv o weights
        state_dict = concat_weights_and_biases(
            state_dict,
            ".*(q|k|v|ogate_preact).weight",
            ".q.weight",
            ".qkv_opreact.weight",
        )

        # fuse qkv o biases
        state_dict = concat_weights_and_biases(
            state_dict,
            ".*(q|k|v|ogate_preact).bias",
            ".q.bias",
            ".qkv_opreact.bias",
        )

        # fuse if gate weights
        state_dict = concat_weights_and_biases(
            state_dict,
            ".*(i|f)gate_preact.weight",
            ".igate_preact.weight",
            ".ifgate_preact.weight",
        )

        # fuse if gate biases
        state_dict = concat_weights_and_biases(
            state_dict,
            ".*(i|f)gate_preact.bias",
            ".igate_preact.bias",
            ".ifgate_preact.bias",
        )

        return state_dict
Modifies the state dict in place.
xlstm/xlstm_large/utils.py