def convert_feedforward_weights(
        state_dict: dict[str, torch.Tensor], block_keys: list[str]
    ):
        """Modifies the state dict in place."""
        # fuse feedforward weights
        state_dict = concat_weights_and_biases(
            state_dict,
            ".*(proj_up_gate|proj_up).weight",
            ".proj_up_gate.weight",
            ".proj_up_gate_z.weight",
        )

        # fuse feedforward biases
        state_dict = concat_weights_and_biases(
            state_dict,
            ".*(proj_up_gate|proj_up).bias",
            ".proj_up_gate.bias",
            ".proj_up_gate_z.bias",
        )

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