Migrate application to hugginface

.gitignore

@@ -0,0 +1,4 @@
+*/*.pyc
+*.pcy
+*/__pycache__
+__pycache__

README.md

@@ -1,14 +1,93 @@
----
-title: Prot Xlstm Variant Fitness
-emoji: ⚡
-colorFrom: indigo
-colorTo: blue
-sdk: streamlit
-sdk_version: 1.43.2
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: 'This application enables to inspect mutational effects on a '
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# prot-xLSTM-app
+
+
+
+## Getting started
+
+To make it easy for you to get started with GitLab, here's a list of recommended next steps.
+
+Already a pro? Just edit this README.md and make it your own. Want to make it easy? [Use the template at the bottom](#editing-this-readme)!
+
+## Add your files
+
+- [ ] [Create](https://docs.gitlab.com/ee/user/project/repository/web_editor.html#create-a-file) or [upload](https://docs.gitlab.com/ee/user/project/repository/web_editor.html#upload-a-file) files
+- [ ] [Add files using the command line](https://docs.gitlab.com/ee/gitlab-basics/add-file.html#add-a-file-using-the-command-line) or push an existing Git repository with the following command:
+
+```
+cd existing_repo
+git remote add origin https://git.bioinf.jku.at/chemoinformatics/prot-xlstm-app.git
+git branch -M main
+git push -uf origin main
+```
+
+## Integrate with your tools
+
+- [ ] [Set up project integrations](https://git.bioinf.jku.at/chemoinformatics/prot-xlstm-app/-/settings/integrations)
+
+## Collaborate with your team
+
+- [ ] [Invite team members and collaborators](https://docs.gitlab.com/ee/user/project/members/)
+- [ ] [Create a new merge request](https://docs.gitlab.com/ee/user/project/merge_requests/creating_merge_requests.html)
+- [ ] [Automatically close issues from merge requests](https://docs.gitlab.com/ee/user/project/issues/managing_issues.html#closing-issues-automatically)
+- [ ] [Enable merge request approvals](https://docs.gitlab.com/ee/user/project/merge_requests/approvals/)
+- [ ] [Set auto-merge](https://docs.gitlab.com/ee/user/project/merge_requests/merge_when_pipeline_succeeds.html)
+
+## Test and Deploy
+
+Use the built-in continuous integration in GitLab.
+
+- [ ] [Get started with GitLab CI/CD](https://docs.gitlab.com/ee/ci/quick_start/)
+- [ ] [Analyze your code for known vulnerabilities with Static Application Security Testing (SAST)](https://docs.gitlab.com/ee/user/application_security/sast/)
+- [ ] [Deploy to Kubernetes, Amazon EC2, or Amazon ECS using Auto Deploy](https://docs.gitlab.com/ee/topics/autodevops/requirements.html)
+- [ ] [Use pull-based deployments for improved Kubernetes management](https://docs.gitlab.com/ee/user/clusters/agent/)
+- [ ] [Set up protected environments](https://docs.gitlab.com/ee/ci/environments/protected_environments.html)
+
+***
+
+# Editing this README
+
+When you're ready to make this README your own, just edit this file and use the handy template below (or feel free to structure it however you want - this is just a starting point!). Thanks to [makeareadme.com](https://www.makeareadme.com/) for this template.
+
+## Suggestions for a good README
+
+Every project is different, so consider which of these sections apply to yours. The sections used in the template are suggestions for most open source projects. Also keep in mind that while a README can be too long and detailed, too long is better than too short. If you think your README is too long, consider utilizing another form of documentation rather than cutting out information.
+
+## Name
+Choose a self-explaining name for your project.
+
+## Description
+Let people know what your project can do specifically. Provide context and add a link to any reference visitors might be unfamiliar with. A list of Features or a Background subsection can also be added here. If there are alternatives to your project, this is a good place to list differentiating factors.
+
+## Badges
+On some READMEs, you may see small images that convey metadata, such as whether or not all the tests are passing for the project. You can use Shields to add some to your README. Many services also have instructions for adding a badge.
+
+## Visuals
+Depending on what you are making, it can be a good idea to include screenshots or even a video (you'll frequently see GIFs rather than actual videos). Tools like ttygif can help, but check out Asciinema for a more sophisticated method.
+
+## Installation
+Within a particular ecosystem, there may be a common way of installing things, such as using Yarn, NuGet, or Homebrew. However, consider the possibility that whoever is reading your README is a novice and would like more guidance. Listing specific steps helps remove ambiguity and gets people to using your project as quickly as possible. If it only runs in a specific context like a particular programming language version or operating system or has dependencies that have to be installed manually, also add a Requirements subsection.
+
+## Usage
+Use examples liberally, and show the expected output if you can. It's helpful to have inline the smallest example of usage that you can demonstrate, while providing links to more sophisticated examples if they are too long to reasonably include in the README.
+
+## Support
+Tell people where they can go to for help. It can be any combination of an issue tracker, a chat room, an email address, etc.
+
+## Roadmap
+If you have ideas for releases in the future, it is a good idea to list them in the README.
+
+## Contributing
+State if you are open to contributions and what your requirements are for accepting them.
+
+For people who want to make changes to your project, it's helpful to have some documentation on how to get started. Perhaps there is a script that they should run or some environment variables that they need to set. Make these steps explicit. These instructions could also be useful to your future self.
+
+You can also document commands to lint the code or run tests. These steps help to ensure high code quality and reduce the likelihood that the changes inadvertently break something. Having instructions for running tests is especially helpful if it requires external setup, such as starting a Selenium server for testing in a browser.
+
+## Authors and acknowledgment
+Show your appreciation to those who have contributed to the project.
+
+## License
+For open source projects, say how it is licensed.
+
+## Project status
+If you have run out of energy or time for your project, put a note at the top of the README saying that development has slowed down or stopped completely. Someone may choose to fork your project or volunteer to step in as a maintainer or owner, allowing your project to keep going. You can also make an explicit request for maintainers.

app.py

@@ -0,0 +1,151 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "" #disable cuda
+
+import streamlit as st
+import numpy as np
+import torch
+import time
+from Bio import SeqIO
+
+from protxlstm.applications.fitness_prediction import single_mutation_landscape_xlstm, create_mutation_df
+from protxlstm.applications.msa_sampler import sample_msa
+from protxlstm.models.xlstm import xLSTMLMHeadModel
+from protxlstm.utils import load_model
+import io
+
+from frontend.constants import info_text, citation_text
+
+DEFAULT_SEQUENCE = "MTARGLALGLLLLLLCPAQVFSQSCVWYGECGIAYGDKRYNCEYSGPPKPLPKDGYDLVQELCPGFFFGNVSLCCDVRQLQTLKDNLQLPLQFLSRCPSCFYNLLNLFCELTCSPRQSQFLNVTATEDYVDPVTNQTKTNVKELQYYVGQSFANAMYNACRDVEAPSSNDKALGLLCGKDADACNATNWIEYMFNKDNGQAPFTITPVFSDFPVHGMEPMNNATKGCDESVDEVTAPCSCQDCSIVCGPKPQPPPPPAPWTILGLDAMYVIMWITYMAFLLVFFGAFFAVWCYRKRYFVSEYTPIDSNIAFSVNASDKGEASCCDPVSAAFEGCLRRLFTRWGSFCVRNPGCVIFFSLVFITACSSGLVFVRVTTNPVDLWSAPSSQARLEKEYFDQHFGPFFRTEQLIIRAPLTDKHIYQPYPSGADVPFGPPLDIQILHQVLDLQIAIENITASYDNETVTLQDICLAPLSPYNTNCTILSVLNYFQNSHSVLDHKKGDDFFVYADYHTHFLYCVRAPASLNDTSLLHDPCLGTFGGPVFPWLVLGGYDDQNYNNATALVITFPVNNYYNDTEKLQRAQAWEKEFINFVKNYKNPNLTISFTAERSIEDELNRESDSDVFTVVISYAIMFLYISLALGHMKSCRRLLVDSKVSLGIAGILIVLSSVACSLGVFSYIGLPLTLIVIEVIPFLVLAVGVDNIFILVQAYQRDERLQGETLDQQLGRVLGEVAPSMFLSSFSETVAFFLGALSVMPAVHTFSLFAGLAVFIDFLLQITCFVSLLGLDIKRQEKNRLDIFCCVRGAEDGTSVQASESCLFRFFKNSYSPLLLKDWMRPIVIAIFVGVLSFSIAVLNKVDIGLDQSLSMPDDSYMVDYFKSISQYLHAGPPVYFVLEEGHDYTSSKGQNMVCGGMGCNNDSLVQQIFNAAQLDNYTRIGFAPSSWIDDYFDWVKPQSSCCRVDNITDQFCNASVVDPACVRCRPLTPEGKQRPQGGDFMRFLPMFLSDNPNPKCGKGGHAAYSSAVNILLGHGTRVGATYFMTYHTVLQTSADFIDALKKARLIASNVTETMGINGSAYRVFPYSVFYVFYEQYLTIIDDTIFNLGVSLGAIFLVTMVLLGCELWSAVIMCATIAMVLVNMFGVMWLWGISLNAVSLVNLVMSCGISVEFCSHITRAFTVSMKGSRVERAEEALAHMGSSVFSGITLTKFGGIVVLAFAKSQIFQIFYFRMYLAMVLLGATHGLIFLPVLLSYIGPSVNKAKSCATEERYKGTERERLLNF"
+
+mutation_positions = []
+msa_file = None
+
+if 'fitness_done' not in st.session_state:
+    st.session_state.fitness_done = False
+    st.session_state.mutations = None
+    st.session_state.fitness_duration = None
+    st.session_state.target_sequence = ""
+    st.session_state.context_sequences = []
+    st.session_state.num_context_sequences = 25
+
+def run_model():
+    try:
+        st.session_state.fitness_duration = time.time()
+        checkpoint = "protxlstm/checkpoints/small"
+        num_context_tokens = 2**15
+        df_mutations = create_mutation_df(st.session_state.target_sequence, mutation_positions)
+        if msa_file != None and st.session_state.num_context_sequences != 0:
+            def load_sequences_from_msa_file(file_obj):
+                text_io = io.TextIOWrapper(file_obj, encoding="utf-8")
+                sequences = [str(record.seq) for record in SeqIO.parse(text_io, "fasta")]
+                return sequences
+            msa_sequences = [msa.upper() for msa in load_sequences_from_msa_file(msa_file)]
+            st.session_state.context_sequences = sample_msa(msa_sequences, max_context_sequences=st.session_state.num_context_sequences, context_length=num_context_tokens)
+            st.session_state.context_sequences += [st.session_state.target_sequence]
+
+        config_update_kwargs = {
+            "mlstm_backend": "chunkwise_variable",
+            "mlstm_chunksize": 1024,
+            "mlstm_return_last_state": True}
+
+        model = load_model(
+            checkpoint,
+            model_class=xLSTMLMHeadModel,
+            device='cpu',
+            dtype=torch.bfloat16,
+            **config_update_kwargs,
+            )
+        model = model.eval()
+        st.session_state.mutations, _ = single_mutation_landscape_xlstm(model, df_mutations, st.session_state.context_sequences, chunk_chunk_size=2**15)
+        print("fitness_done")
+        st.session_state.fitness_done = True
+        st.session_state.fitness_duration = time.time() - st.session_state.fitness_duration
+    except Exception as e:
+        print(e)
+
+# PAGE STYLE (mainly for custom aa selection)
+st.set_page_config(layout="wide")
+st.markdown(
+    """
+    <style>
+    .stButtonGroup button {
+        padding: 0px 1px 0px 1px !important;
+        border: 0 solid transparent !important;
+        min-height: 0px !important;
+        line-height: 120% !important;
+        height: auto !important;
+    }
+    .stSidebar {
+        width: 600px !important;
+    }
+    </style>
+    """,
+    unsafe_allow_html=True
+)
+
+
+with st.sidebar:
+    st.title("Prot-xLSTM Variant Fitness")
+
+    # LOAD SEQUENCE
+    st.session_state.target_sequence = st.text_area(
+        "Target protein sequence",
+        placeholder=DEFAULT_SEQUENCE,
+        value=st.session_state.target_sequence
+    )
+    if st.button("Load sequence"):
+        if st.session_state.target_sequence == "":
+            st.session_state.target_sequence = DEFAULT_SEQUENCE
+
+    # MANAGE CONTEXT SEQUENCES
+    context_type = st.selectbox(
+        "Choose how to enter context",
+        ("Enter manually", "Use MSA file"),
+        index=None,
+        placeholder="Choose context",
+    )
+    if context_type == 'Enter manually':
+        context_sequence_str = st.text_area(
+            "Enter context protein sequences (seperated by comma)",
+            placeholder=DEFAULT_SEQUENCE,
+        )
+        st.session_state.context_sequences = context_sequence_str.split(",") + [st.session_state.target_sequence]
+    elif context_type == 'Use MSA file':
+        msa_file = st.file_uploader("Choose MSA file")
+        st.session_state.num_context_sequences = st.number_input("How many of these sequences should be used?", min_value=0, step=1, value=25)
+    else:
+        st.session_state.context_sequences = [st.session_state.target_sequence]
+
+if st.session_state.target_sequence != "":
+    with st.container():
+
+        # MUTATION POSITION SELECTION
+        aas = list(st.session_state.target_sequence)
+        mutation_indices = np.arange(1, len(aas)+1)
+        mutation_positions = st.segmented_control(
+            "Choose mutation positions (click to select)", mutation_indices, selection_mode="multi", format_func=lambda i: aas[i-1],
+        )
+        st.button("Check Fitness", on_click=run_model)
+
+        # DISPLAY RESULTS
+        if st.session_state.fitness_done:
+            st.metric(label="Running time", value=f"{st.session_state.fitness_duration:.2f} sec.")
+            selected_pos = st.selectbox(
+                "Visualized mutation position",
+                st.session_state.mutations['position'].unique()
+            )
+            selected_data = st.session_state.mutations.where(st.session_state.mutations['position'] == selected_pos)
+            st.bar_chart(selected_data, x='mutation', y='effect', horizontal=True)
+            st.dataframe(st.session_state.mutations, use_container_width=True)
+
+# TUTORIAL
+with st.expander("Info & Tutorial", expanded=True):
+    st.subheader("Tutorial")
+    st.markdown("**1.** Choose a target protein sequence (leave empty to use a sample sequence) and press 'Load Sequence'")
+    st.markdown("**2.** Enter or upload you context sequences. (leave empty to use no context)")
+    st.markdown("**3.** Choose which amino acids to mutate (click on the AA's to select them) and press 'Check Fitness'")
+    st.subheader("General Information")
+    st.markdown(info_text, unsafe_allow_html=True)
+    st.markdown("")
+    st.subheader("Cite us / BibTex")
+    st.code(citation_text, language=None)

frontend/constants.py

@@ -0,0 +1,41 @@
+info_text = ("""
+             <div style="text-align: justify"> 
+             This application enables to inspect mutational effects on a 
+             predefined protein sequence.<br>
+             <br>
+             </div>
+             
+             <div style="text-align: justify"> 
+             It is built on the Prot-xLSTM backbone model, an xLSTM model specifically 
+             trained on protein sequences. Prot-xLSTM was trained using the 
+             Fill-In-the-Middle (FIM) objective, which allows it to perform sequence 
+             inpainting. Additionally, the model can be provided with a potentially 
+             large set of homologous sequences to enhance its predictions.<br>
+             <br>
+             </div>
+             
+             <div style="text-align: justify"> 
+             For further information please refer, to: <a href="https://openreview.net/forum?id=IjbXZdugdj" target="_blank">https://openreview.net/forum?id=IjbXZdugdj</a>. <br>
+             <br>
+             
+             This Hugging Face application is based on the following GitHub repository: 
+             <a href="https://github.com/ml-jku/Prot-xLSTM?tab=readme-ov-file" target="_blank">https://github.com/ml-jku/Prot-xLSTM?tab=readme-ov-file</a>. <br>
+             The streamlit application was developed by Elias Bürger.
+             </div>
+             
+             <div style="text-align: justify"> 
+             Please cite us as follows: <br> 
+             </div> 
+             """)
+citation_text = """
+                 @misc{
+                    schmidinger2024bioxlstmgenerativemodelingrepresentation,
+                    title={Bio-xLSTM: Generative modeling, representation and in-context learning of biological and chemical sequences},
+                    author={Niklas Schmidinger and Lisa Schneckenreiter and Philipp Seidl and Johannes Schimunek and Pieter-Jan Hoedt and Johannes Brandstetter and Andreas Mayr and Sohvi Luukkonen and Sepp Hochreiter and Günter Klambauer},
+                    year={2024},
+                    eprint={2411.04165},
+                    archivePrefix={arXiv},
+                    primaryClass={q-bio.BM},
+                    url={https://arxiv.org/abs/2411.04165},
+                  }  
+                """

prot_xlstm_env.yml

@@ -0,0 +1,39 @@
+name: prot_xlstm_app
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+  - defaults
+dependencies:
+  - cuda=12.1
+  - cuda-nvcc=12.1
+  - gxx_linux-64=11.2.0
+  - python=3.11
+  - pip
+  - pytorch=2.2.0
+  - pytorch-cuda=12.1
+  - cmake
+  - ninja
+  - pip:
+    - accelerate>=0.26.0
+    - biopython #==1.83
+    - bottleneck #==1.4.2
+    - dacite #==1.8.1
+    - ipykernel #==6.29.3  
+    - mamba_ssm==1.2.0
+    - matplotlib #==3.8.4
+    - numpy<2.0 #==1.26.4
+    - omegaconf #==2.3.0
+    - pandas #==2.2.2
+    - pyhmmer #==0.10.15
+    - rich #==13.7.1
+    - scipy #==1.13.0
+    - seaborn #==0.13.2
+    - torchmetrics #==1.2.1
+    - tqdm #==4.66.4
+    - transformers==4.44.2
+    - tueplots #==0.0.17
+    - wandb #==0.17.0
+    - streamlit #==1.43.2
+
+

protxlstm/__init__.py

@@ -0,0 +1 @@
+# __version__ = "0.0.1"

protxlstm/applications/fitness_prediction.py

@@ -0,0 +1,214 @@
+import os
+import numpy as np
+import pandas as pd
+
+import torch
+from tqdm.auto import tqdm
+
+from protxlstm.applications.msa_sampler import MSASampler
+from protxlstm.generation import generate_sequence
+from protxlstm.utils import AA_TO_ID, tokenizer, ID_TO_AA
+
+
+def precompute_context_state(model, sequences, chunk_chunk_size=2**15):
+    """
+    Precompute the output states for a fixed context that remains the same across generations.
+    Returns the hidden states to continue generation later.
+    """
+    device = next(model.parameters()).device
+
+    input_ids, pos_ids = prepare_context(sequences)
+    state = None
+    
+    for chunk in range(input_ids.shape[1]//chunk_chunk_size+1):
+
+        start_idx = chunk*chunk_chunk_size
+        end_idx = min((chunk+1)*chunk_chunk_size, input_ids.shape[1])
+
+        if start_idx == end_idx:
+            pass
+        
+        else:
+            input_ids_chunk = input_ids[:, start_idx:end_idx].to(device)
+            pos_ids_chunk = pos_ids[:, start_idx:end_idx].to(device)
+
+            with torch.no_grad():
+                outputs = model(input_ids=input_ids_chunk,
+                                position_ids=pos_ids_chunk,
+                                state=state,
+                                output_hidden_states=True,
+                                return_dict=True)
+                state = outputs.state
+
+    # Return the hidden states for reuse
+    return state
+
+def prepare_context(sequences):
+    tokenized_sequences = tokenizer(sequences, concatenate=False)
+    pos_ids = torch.cat([torch.arange(0, len(seq), dtype=torch.int64) for seq in tokenized_sequences], 0)[None, :]
+    input_ids = torch.cat(tokenized_sequences, 0)[None, :].to(torch.int64)
+    return input_ids, pos_ids
+
+def prepare_single_mutation_target(target, mut_pos):
+
+    pos_ids = torch.arange(target.shape[1], dtype=torch.int64)[None,:] # default position ids
+    t = torch.ones((target.shape[0], 1), dtype=torch.int64)
+    new_target = torch.cat([
+        target[:,:mut_pos], # WT sequence until mutated position
+        AA_TO_ID["<mask-1>"] * t, # Mask token at the muated position
+        target[:,mut_pos+1:],  # WT sequence after mutated position
+        AA_TO_ID["<eos>"] * t, # End of sequence token
+        AA_TO_ID["<mask-1>"] * t, # Mask token
+        ], dim=1)
+    new_pos_ids = torch.cat([
+        pos_ids,
+        0 * t, # end of sequence
+        mut_pos * t, # mutation position
+        ], dim=1)
+
+    is_fim_dict = { AA_TO_ID["<mask-1>"] : pos_ids[:,mut_pos].squeeze().item()}
+
+    return new_target, new_pos_ids, is_fim_dict
+
+def single_mutation_landscape_xlstm(model, single_mutations, context_sequences, chunk_chunk_size=2**15):
+
+    device = next(model.parameters()).device
+
+    # Tokenize WT target sequence    
+    wt_tokens = tokenizer([context_sequences[-1]], concatenate=True)
+    
+    # Precompute hidden state of context
+    context_state = precompute_context_state(model, context_sequences, chunk_chunk_size=chunk_chunk_size)
+
+    mutation_positions = sorted(single_mutations.position.unique())    
+    all_logits = np.zeros((len(mutation_positions), 20))
+
+    # Iterate over all mutated positions
+    for i, pos in tqdm(enumerate(mutation_positions), total=len(mutation_positions), desc="Generating mutational landscape"): # This loop can be parallelized
+        
+        # Prepare target
+        wt_aa_id = wt_tokens[0, pos+1].int().item() # wild type AA index
+        target_tokens, target_pos_ids, _ = prepare_single_mutation_target(wt_tokens, pos+1)
+
+        with torch.no_grad():
+            outputs = model(input_ids=target_tokens.to(device),
+                            position_ids=target_pos_ids.to(device),
+                            state=context_state,
+            )
+
+        # Extact logits and compute mutational effect
+        logits = outputs.logits.clone().detach()  # Raw logits
+        logits_mut = logits[0, -1, 4:24].log_softmax(-1)  # Log-softmax for mutation prediction: (4-24) correspond to natural NNs
+        mut_effects = logits_mut - logits_mut[wt_aa_id - 4]  # Subtract log probability of ground truth
+        all_logits[i,:] = logits_mut.cpu()
+        single_mutations.loc[single_mutations.position == pos, 'effect'] = single_mutations.loc[single_mutations.position == pos, 'mutation_idx'].apply(lambda x : mut_effects[x-4].item())
+    
+    return single_mutations, all_logits
+
+def single_mutation_landscape_mamba(model, single_mutations, context_sequences):
+
+    # Prepare context sequences
+    context_tokens, context_pos_ids = prepare_context(context_sequences)
+
+    # Tokenize WT target sequence
+    wt_tokens = tokenizer([context_sequences[-1]], concatenate=True)
+
+    mutation_positions = sorted(single_mutations.position.unique())    
+    all_logits = np.zeros((len(mutation_positions), 20))
+    
+    # Iterate over all mutated positions
+    for i, pos in tqdm(enumerate(mutation_positions), total=len(mutation_positions), desc="Generating mutational landscape"): # This loop can be parallelized
+        
+        # Prepare target
+        wt_aa_id = wt_tokens[0, pos+1].int().item() # wild type AA index
+        target_tokens, target_pos_ids, is_fim_dict = prepare_single_mutation_target(wt_tokens, pos+1)
+       
+        # Merge context and target
+        device = next(model.parameters()).device
+        context_tokens = torch.cat([context_tokens, target_tokens], dim=1).to(device)
+        context_pos_ids = torch.cat([context_pos_ids, target_pos_ids], dim=1).to(device)
+
+        # Generate fim-token prediction
+        output = generate_sequence(
+                model,
+                context_tokens,
+                position_ids=context_pos_ids,
+                is_fim=is_fim_dict,
+                max_length=1,
+                temperature=1.0,
+                top_k=0,
+                top_p=0.0,
+                return_dict_in_generate=True,
+                output_scores=True,
+                eos_token_id=AA_TO_ID["<cls>"],
+                device=device
+            )
+        
+        # Extact logits and compute mutational effect
+        logits = torch.tensor(output["scores"])  # Raw logits
+        logits_mut = logits[0, 0, 4:24].log_softmax(-1)  # Log-softmax for mutation prediction: (4-24) correspond to natural NNs
+        mut_effects = logits_mut - logits_mut[wt_aa_id - 4]  # Subtract log probability of ground truth
+        all_logits[i,:] = logits_mut.cpu()
+
+        single_mutations.loc[single_mutations.position == pos, 'effect'] = single_mutations.loc[single_mutations.position == pos, 'mutation_idx'].apply(lambda x : mut_effects[x-4].item())
+    
+    return single_mutations, all_logits
+
+def single_mutation_landscape_retrieval(single_mutations, msa_sequences, msa_weights_path):
+
+    # One-hot encode MSA sequences
+    msa_tokens = np.array([[AA_TO_ID[aa.upper()] for aa in seq] for seq in msa_sequences])
+    one_hot_tokens = np.zeros((len(msa_tokens), len(msa_tokens[0]), 40))
+    one_hot_tokens[np.arange(len(msa_tokens))[:, None], np.arange(len(msa_tokens[0])), msa_tokens] = 1
+
+    #Load/compute weights
+    if os.path.exists(msa_weights_path):
+        weights = np.load(msa_weights_path)
+    else:
+        sampler = MSASampler(0.98, 0.7)
+        weights = sampler.get_weights(msa_tokens)[1]
+        np.save(msa_weights_path, weights)
+    assert one_hot_tokens.shape[0] == weights.shape[0]
+
+    # Apply sequence weights, normalize amino acid probabilities per position, and convert to a PyTorch tensor.
+    one_hot_tokens = one_hot_tokens * weights[:, None, None]
+    one_hot_tokens = one_hot_tokens.sum(0)
+    one_hot_tokens = one_hot_tokens[:, 4:24] + 1 / len(msa_sequences)
+    one_hot_tokens_sum = one_hot_tokens.sum(-1)
+    one_hot_tokens = one_hot_tokens / one_hot_tokens_sum[:, None]
+    one_hot_tokens = torch.tensor(one_hot_tokens).float()    
+
+    # Compute mutational effects
+    wild_type = msa_tokens[0]
+    logits = one_hot_tokens.log()
+    logits = logits - logits[torch.arange(len(logits)), wild_type - 4][:, None]
+
+    single_mutations['retrieval_effect'] = single_mutations.apply(
+        lambda row: logits[row['position'], row['mutation_idx'] - 4].item(), axis=1)
+    
+    return single_mutations
+
+
+def create_mutation_df(sequence, mutation_positions):
+    """
+    Generate a DataFrame containing all possible mutations at specified positions in a sequence.
+
+    Args:
+        sequence (str): The original sequence to mutate.
+        mutation_positions (list of int): List of positions to mutate (1-based index).
+
+    Returns:
+        pd.DataFrame:
+            - 'mutation': formatted mutation string (e.g., 'A10G' for Ala at position 10 to Gly).
+            - 'position': 0-based position in the sequence.
+            - 'mutation_idx': numeric index for the mutation.
+    """
+    
+    AAs = {k: v for k, v in ID_TO_AA.items() if 4 <= k <= 23}
+    mutation_data = []
+    for position in mutation_positions:
+        wt = sequence[position - 1] 
+        for idx, aa in AAs.items():
+            mutation = f"{wt}{position}{aa}"
+            mutation_data.append({'mutation': mutation, 'position': position - 1, 'mutation_idx': idx})
+    return pd.DataFrame(mutation_data)

protxlstm/applications/generation_utils/create_sequence_df.py

@@ -0,0 +1,85 @@
+import numpy as np
+import pickle
+import pandas as pd
+
+from protxlstm.dataloaders import ProteinMemmapDataset
+from protxlstm.utils import decode_sequence, reorder_masked_sequence
+
+
+def create_sequence_df(model_name, family_idx, parameters_list=None, num_sequences = 100, data_dir="./data/"):
+
+    #load dataset
+    dataset = ProteinMemmapDataset(
+            msa_memmap_path=f"{data_dir}open_protein_set_memmap.dat",
+            msa_memmap_meta_path=f"{data_dir}open_protein_set_memmap_indices.csv",
+            subset_path=f"{data_dir}/cluster_testing_set.txt",
+            sample=False,
+            max_msa_len=-1,
+            reverse=False,
+            seed=0,
+            troubleshoot=False,
+            fim_strategy="multiple_span",
+            always_mask=False,
+            max_position_embeddings=2048,
+            max_seq_position_embeddings=512,
+            add_position_ids="1d",
+            mask_fraction=0.2,
+            max_patches=5,
+        )
+    
+    family_id = list(dataset.dataset_meta["msa_id"])[family_idx]
+
+    if model_name == "natural":
+
+        data = dataset[family_idx]
+        sequence_df = pd.DataFrame(columns=["family", "family_id", "sequence", "sequence_length"])
+        tokens = data["input_ids"][None,:]
+        all_context = decode_sequence(tokens[0].cpu().numpy())
+        list_sequences_msa = [reorder_masked_sequence(elem+"<cls>") for elem in all_context.split("<cls>")[1:-1]]
+
+        rd_idxs = np.random.choice(len(list_sequences_msa), num_sequences, replace=False)
+        natural_sequences = [seq for i, seq in enumerate(list_sequences_msa) if i in rd_idxs]
+
+        df_dict = {"family": [family_idx]*len(natural_sequences), 
+                    "family_id": [family_id]*len(natural_sequences),
+                    "sequence": natural_sequences,
+                    "sequence_length": [len(seq) for seq in natural_sequences]}
+
+        sequence_df = pd.concat([sequence_df, pd.DataFrame(df_dict)], ignore_index = True)
+
+    else:
+        
+        sequence_df = pd.DataFrame(columns=["family", "family_id",  "n_seqs_ctx", "temperature", "top_k", "top_p", "original_sequence", "sequence", "sequence_length", "perplexity"])
+        
+        if parameters_list is None:
+            parameters_list = [(10,1.,10,1.), (10,1.,15,1.), (10,1.,10,0.95), (10,0.9,10,0.95), (10,0.8,10,0.9),
+                    (100,1.,10,1.), (100,1.,15,1.), (100,1.,10,0.95), (100,0.9,10,0.95), (100,0.8,10,0.9),
+                    (500,1.,10,1.), (500,1.,15,1.), (500,1.,10,0.95), (500,0.9,10,0.95), (500,0.8,10,0.9),
+                    (1000,1.,10,1.), (1000,1.,15,1.), (1000,1.,10,0.95), (1000,0.9,10,0.95), (1000,0.8,10,0.9),
+                    (-1,1.,10,1.), (-1,1.,15,1.), (-1,1.,10,0.95), (-1,0.9,10,0.95), (-1,0.8,10,0.9)]
+        
+        for param in parameters_list:
+            n_seqs_ctx, temperature, top_k, top_p = param
+
+            with open(f"evaluation/generation/generated_sequences/{model_name}/{family_idx}_{param}_{num_sequences}", "rb") as f:
+                gen_seqs = pickle.load(f)
+
+            original_sequences =  list(gen_seqs[family_idx][param].keys())
+            reordered_sequences = [reorder_masked_sequence(seq) for seq in original_sequences]
+            perplexities = [gen_seqs[family_idx][param][seq]["perplexity"] for seq in original_sequences]
+            df_dict = {"family": [family_idx]*len(original_sequences), 
+                        "family_id": [family_id]*len(original_sequences),
+                        "n_seqs_ctx": [n_seqs_ctx]*len(original_sequences),
+                        "temperature": [temperature]*len(original_sequences),
+                        "top_k": [top_k]*len(original_sequences),
+                        "top_p": [top_p]*len(original_sequences),
+                        "original_sequence": original_sequences,
+                        "sequence": reordered_sequences,
+                        "sequence_length": [len(seq) for seq in reordered_sequences],
+                        "perplexity": perplexities
+                        }
+
+            sequence_df = pd.concat([sequence_df, pd.DataFrame(df_dict)], ignore_index = True)
+
+    return sequence_df
+

protxlstm/applications/generation_utils/score_hamming.py

@@ -0,0 +1,80 @@
+import numpy as np
+from tqdm import tqdm
+import pandas as pd
+from Bio import Align
+
+from protxlstm.dataloaders import ProteinMemmapDataset
+from protxlstm.utils import decode_sequence, reorder_masked_sequence
+
+
+aligner = Align.PairwiseAligner()
+aligner.mode = 'global'
+aligner.match_score = 1
+aligner.mismatch_score = -1
+aligner.open_gap_score = -1
+aligner.extend_gap_score = -1
+
+def align_sequences(ref_seq, query_seq, print_alignments=False):
+    def hamming_str(s1,s2):
+        assert len(s1) == len(s2)
+        return sum(np.array(list(s1)) != np.array(list(s2)))/len(s1)
+    alignments = aligner.align(ref_seq, query_seq)
+    if print_alignments:
+        print("Score = %.1f:" % alignments[0].score)
+        print(alignments[0])
+    return hamming_str(alignments[0][0], alignments[0][1]), alignments[0][0], alignments[0][1]
+
+
+def score_hamming(sequence_df, family_idx, data_dir = f"./data/"):
+
+    assert len(set(list(sequence_df["family"]))) == 1 and sequence_df["family"].iloc[0] == family_idx
+
+    #load dataset
+    dataset = ProteinMemmapDataset(
+            msa_memmap_path=f"{data_dir}open_protein_set_memmap.dat",
+            msa_memmap_meta_path=f"{data_dir}open_protein_set_memmap_indices.csv",
+            subset_path=f"{data_dir}/cluster_testing_set.txt",
+            sample=False,
+            max_msa_len=-1,
+            reverse=False,
+            seed=0,
+            troubleshoot=False,
+            fim_strategy="multiple_span",
+            always_mask=False,
+            max_position_embeddings=2048,
+            max_seq_position_embeddings=512,
+            add_position_ids="1d",
+            mask_fraction=0.2,
+            max_patches=5,
+        )
+    
+    # Select a sample of the dataset to be the input
+    data = dataset[family_idx]
+    tokens = data["input_ids"][None,:]
+    all_context = decode_sequence(tokens[0].cpu().numpy())
+    list_sequences_msa = [reorder_masked_sequence(elem+"<cls>") for elem in all_context.split("<cls>")[1:-1]]
+
+    # sequence_df["hamming"] = pd.Series(dtype=object)
+    sequence_df["min_hamming"] = pd.Series()
+    sequence_df["median_hamming"] = pd.Series()
+    sequence_df["mean_hamming"] = pd.Series()
+    sequence_df["std_hamming"] = pd.Series()
+
+    for seq in tqdm(list(sequence_df["sequence"])):
+
+        all_hamming = []
+        for ctx_seq in list_sequences_msa:
+            if ctx_seq == seq:
+                continue
+            else:
+                hamming, _, _ = align_sequences(ctx_seq, seq , print_alignments=False)
+                all_hamming.append(hamming)
+
+        # sequence_df.loc[sequence_df["sequence"] == seq, "hamming"] = [all_hamming]
+        sequence_df.loc[sequence_df["sequence"] == seq, "min_hamming"] = np.min(all_hamming)
+        sequence_df.loc[sequence_df["sequence"] == seq, "median_hamming"] = np.median(all_hamming)
+        sequence_df.loc[sequence_df["sequence"] == seq, "mean_hamming"] = np.mean(all_hamming)
+        sequence_df.loc[sequence_df["sequence"] == seq, "std_hamming"] = np.std(all_hamming)
+
+    return sequence_df
+        

protxlstm/applications/generation_utils/score_hmmer.py

@@ -0,0 +1,102 @@
+import string
+from Bio import SeqIO
+import pyhmmer
+from tqdm import tqdm
+
+alphabet = pyhmmer.easel.Alphabet.amino()
+
+# This is an efficient way to delete lowercase characters and insertion characters from a string
+deletekeys = dict.fromkeys(string.ascii_lowercase)
+deletekeys["."] = None
+deletekeys["*"] = None
+translation = str.maketrans(deletekeys)
+
+def remove_insertions(sequence: str) -> str:
+    """ Removes any insertions into the sequence. Needed to load aligned sequences in an MSA. """
+    return sequence.translate(translation)
+
+def read_msa(filename: str):
+    """ Reads the sequences from an MSA file, automatically removes insertions."""
+    return [(record.description, remove_insertions(str(record.seq))) for record in SeqIO.parse(filename, "fasta")]
+
+def read_msa_unaligned(filename: str):
+    """ Reads the sequences from an MSA file, removes only . - and * characters."""
+    return [(record.description, str(record.seq).replace(".","").replace("-","").replace("*","").upper()) for record in SeqIO.parse(filename, "fasta")]
+
+def check_msa(msa):
+    """ Checks if there are any repeated sequences in the MSA"""
+    seqs = set()
+    for el in msa:
+        seqs.add(el[1])
+    assert len(seqs) == len(msa), "There are repeated sequences in the MSA"
+    
+def make_hmm_from_a3m_msa(msa_filepath, hmm_filename=None):
+    # Load MSA from a3m
+    msa_tup = read_msa(msa_filepath)
+    # check_msa(msa_tup)
+    # Create digitized MSA block
+    all_seqs = [pyhmmer.easel.TextSequence(name=str(i).encode("utf-8"), sequence=seq) for i, (idz, seq) in enumerate(msa_tup)]
+    msa  = pyhmmer.easel.TextMSA(name=b"msa", sequences=all_seqs)
+    msa = msa.digitize(alphabet)
+    # Fit HMM
+    builder = pyhmmer.plan7.Builder(alphabet)
+    background = pyhmmer.plan7.Background(alphabet)
+    hmm, _, _ = builder.build_msa(msa, background)
+    if hmm_filename is not None:
+        with open(f"{hmm_filename}.hmm", "wb") as output_file:
+            hmm.write(output_file)
+    return hmm
+
+def align_and_score_sequences_in_a3m_with_hmm(hmm, sequences_path=None, sequences_list=None):
+    if sequences_list is not None:
+        msa = sequences_list
+        all_seqs = [pyhmmer.easel.TextSequence(name=str(i).encode("utf-8"), sequence=seq) for i, seq in enumerate(sequences_list)]
+    elif sequences_path is not None:
+        # Load sequences from a3m
+        msa = read_msa_unaligned(sequences_path)
+        all_seqs = [pyhmmer.easel.TextSequence(name=str(i).encode("utf-8"), sequence=seq) for i, (idz, seq) in enumerate(msa)]
+    else:
+        raise NotImplementedError("Missing sequences to align/score")
+    # Create digitized Sequence block
+    seq_block = pyhmmer.easel.TextSequenceBlock(all_seqs)
+    seq_block = seq_block.digitize(alphabet)
+    # Get all hits from the hmm
+    background = pyhmmer.plan7.Background(alphabet)
+    pipeline = pyhmmer.plan7.Pipeline(alphabet, background=background, bias_filter=False, F1=1.0, F2=1.0, F3=1.0)
+    hits = pipeline.search_hmm(hmm, seq_block)
+    if len(hits) != len(msa):
+        print(f"Number of hits: {len(hits)} is different from the number of sequences in the MSA: {len(msa)}")
+    # Extract hits
+    all_hits = {}
+    for hit in hits:
+        idz, score, evalue = hit.name, hit.score, hit.evalue
+        i = int(idz.decode("utf-8"))
+        seq = msa[i][1] if sequences_path is not None else sequences_list[i]
+        all_hits[seq] = {"score": score, "evalue": evalue}
+    return all_hits
+
+
+def score_hmmer(sequence_df, family_idx, data_dir = f"./data/"):
+
+    assert len(set(list(sequence_df["family"]))) == 1 and sequence_df["family"].iloc[0] == family_idx
+    
+    family_id = sequence_df["family_id"].iloc[0]
+    msa_filepath = f"{data_dir}/a3m_files/{family_id}/a3m/uniclust30.a3m"
+    try:
+        hmm = make_hmm_from_a3m_msa(msa_filepath)
+    except:
+        raise Exception(f"Missing MSA of family {family_id}")
+    
+    # align sequences
+    sequences = list(sequence_df["sequence"])
+    scores = align_and_score_sequences_in_a3m_with_hmm(hmm, sequences_list=sequences)
+
+    # save the scores associated to each sequence in the main df in the columns "score" and "evalue"
+    for seq in tqdm(sequences):
+        sequence_df.loc[sequence_df["sequence"] == seq, "score_gen"] = scores[seq]["score"] if seq in scores.keys() else 0
+        sequence_df.loc[sequence_df["sequence"] == seq, "evalue_gen"] = scores[seq]["evalue"] if seq in scores.keys() else 1
+
+    return sequence_df
+
+
+            

protxlstm/applications/generation_utils/score_structure.py

@@ -0,0 +1,55 @@
+from Bio.PDB import PDBParser
+import torch
+from tqdm import tqdm
+from transformers import EsmForProteinFolding
+
+from protxlstm.utils import MASK_TO_ID
+
+
+pdb_parser = PDBParser()
+
+
+def compute_structure(seq, model):
+    def keep_sequence(seq, l):
+        if len(seq) > l:
+            return False
+        for mm in list(MASK_TO_ID.keys())+["<eos>", "<pad>", "<unk>", "<mask>", "<cls>", "<null_1>", "." , "-"]:
+            if mm in seq:
+                return False
+        return True
+    keep = keep_sequence(seq, l=750)
+    if keep:
+        with torch.no_grad():
+            output = model.infer([seq])
+        # pdb = model.output_to_pdb(output)
+        ptm = output["ptm"].item()
+        pae = output["predicted_aligned_error"].cpu().numpy()
+        mean_plddt = ((output["plddt"] * output["atom37_atom_exists"]).sum(dim=(1, 2)) / output["atom37_atom_exists"].sum(dim=(1, 2))).item()
+        pos_plddt = ((output["plddt"] * output["atom37_atom_exists"]).sum(dim=(2,)) / output["atom37_atom_exists"].sum(dim=(2,))).cpu().numpy()
+    else:
+        print(f"Sequence is invalid.")
+        ptm, pae, mean_plddt, pos_plddt = 0, 0 ,0 , 0
+    return ptm, pae, mean_plddt, pos_plddt
+
+
+def score_structure(sequence_df, family_idx):
+
+    assert len(set(list(sequence_df["family"]))) == 1 and sequence_df["family"].iloc[0] == family_idx
+
+    device="cuda:0"
+
+    # Import the folding model
+    model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True)
+
+    model = model.cuda(device)
+    model.esm = model.esm.half()
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    sequences = list(sequence_df["sequence"])
+    for seq in tqdm(sequences):
+
+        ptm, pae, mean_plddt, pos_plddt = compute_structure(seq, model)
+        sequence_df.loc[sequence_df["sequence"] == seq, "ptm"] = ptm
+        sequence_df.loc[sequence_df["sequence"] == seq, "mean_plddt"] = mean_plddt
+        
+    return sequence_df

protxlstm/applications/msa_sampler.py

@@ -0,0 +1,196 @@
+# Original code from ProtMamba under Apache License 2.0.
+#
+# Modifications made by Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Modify handling of weights in `MSASampler`
+
+import math
+import os
+from typing import Optional, Callable
+
+import numpy as np
+import torch
+
+from protxlstm.utils import AA_TO_ID
+
+
+def compute_hamming_csim_torch(
+        seqs: torch.Tensor,
+        ungapped_msa: torch.Tensor,
+        gap_token: int,
+        gap_token_mask: int,
+) -> torch.Tensor:
+    return (seqs.unsqueeze(1) == ungapped_msa).sum(dim=2)
+
+def _compute_homology_weights(
+        ungapped_msa: np.ndarray,
+        gap_token: int,
+        gap_token_mask: int,
+        theta: float,
+        hamming_csim_func: Callable,
+        max_memory: int = 20,
+        can_use_torch: bool = True,
+) -> np.ndarray:
+    use_torch = can_use_torch
+    if use_torch:
+        hamming_csim_func = compute_hamming_csim_torch
+    batch_size = math.floor(
+        2
+        * 1024
+        * 1024
+        * 1024
+        / (ungapped_msa.shape[0] * ungapped_msa.shape[1])
+        * max_memory
+        / 40
+    )
+
+    batch_size = 1 if batch_size == 0 else batch_size
+
+    neighbors = []
+    if not use_torch:
+        masked_ungapped_msa = ungapped_msa.copy()
+    else:
+        ungapped_msa = torch.from_numpy(ungapped_msa).byte()
+        masked_ungapped_msa = ungapped_msa.clone()
+    masked_ungapped_msa[masked_ungapped_msa == gap_token] = gap_token_mask
+    for b_start in range(0, len(ungapped_msa), batch_size):
+        b_end = b_start + batch_size
+        seqs = ungapped_msa[b_start:b_end]
+
+        sim = hamming_csim_func(
+            seqs=seqs,
+            ungapped_msa=masked_ungapped_msa,
+            gap_token=gap_token,
+            gap_token_mask=gap_token_mask,
+        )
+        if not use_torch:
+            sim = sim / (seqs != gap_token).sum(axis=1, keepdims=True)
+            d = 1 - sim
+            d = d.clamp(0, 1)
+            this_neighbors = (d <= theta).sum(axis=1)
+        else:
+            sim = sim / (seqs != gap_token).sum(dim=1, keepdim=True)
+            d = 1 - sim
+            # fillna
+            d[torch.isnan(d)] = 0
+            d = d.clamp(0, 1)
+            this_neighbors = (d <= theta).sum(dim=1).cpu()
+        neighbors.append(this_neighbors)
+    return np.concatenate(neighbors)
+
+def compute_homology_weights(
+        ungapped_msa: np.ndarray,
+        theta: float = 0.2,
+        gap_token: int = AA_TO_ID["-"],
+        gap_token_mask: int = 255,
+        hamming_csim_func: Callable = compute_hamming_csim_torch,
+) -> tuple[int, np.ndarray]:
+    """
+    Calculate the effective number of sequences and sampling probability for the NEIGHBORS and NEIGHBORS_NO_LIMIT sampling methods using numpy.
+
+    Parameters:
+
+        ungapped_msa (np.ndarray): The MSA (from .fa).
+        theta (float, optional): A parameter used to determine the similarity between sequences. Default is 0.2.
+        gap_token (int, optional): The token representing gaps in the (Uniprot21 encoded) MSA. Default is 20.
+        gap_token_mask (int): token for masking gaps. should be a token not representing any other value.
+
+    Returns:
+
+        tuple[int, np.ndarray]: A tuple containing the effective number of sequences and the sampling probability for each sequence in the MSA.
+    """
+    neighbors = _compute_homology_weights(
+        ungapped_msa=ungapped_msa,
+        gap_token=gap_token,
+        gap_token_mask=gap_token_mask,
+        theta=theta,
+        hamming_csim_func=hamming_csim_func,
+    )
+    n_eff = np.sum(1 / neighbors)
+
+    p = 1 / neighbors
+    p /= np.sum(p)
+    return n_eff, p
+
+class MSASampler:
+
+    def __init__(self, max_similarity, max_dissimilarity, force_include_first=True):
+        self.max_similarity = max_similarity
+        self.max_dissimilarity = max_dissimilarity
+        self.force_include_first = force_include_first
+        self.theta = 0.2
+
+    def _get_sim_filtered_idxs(self, msa: np.ndarray) -> np.ndarray:
+        nonnormalized_sim = (msa == msa[[0]]).sum(axis=1)
+        normfactor = msa.shape[1]
+        norm_sim = nonnormalized_sim / normfactor
+
+        assert (norm_sim.min() >= 0) and (norm_sim.max() <= 1)
+        dsim = 1 - norm_sim
+
+        max_sim_filter = norm_sim <= self.max_similarity
+        max_dissim_filter = dsim <= self.max_dissimilarity
+        return np.where(max_sim_filter & max_dissim_filter)[0]
+
+    def get_weights(
+            self, msa: np.ndarray,
+    ) -> tuple[Optional[float], Optional[np.ndarray]]:
+        return compute_homology_weights(
+            ungapped_msa=msa,
+            theta=self.theta,
+            gap_token_mask=255,
+
+        )
+
+    def get_sample_idxs(
+            self,
+            msa: np.ndarray,
+            size: int = 1,
+            random = False,
+            msa_weights_path = None,
+            seed = 0,
+    ) -> np.ndarray:
+        
+        np.random.seed(seed)
+        
+        if random:
+            return np.random.choice(len(msa), replace=False, size=size) if len(msa) >= size else np.arange(len(msa))
+        
+        msa = np.array([[AA_TO_ID[aa] for aa in seq.upper()][:len(msa[0])] for seq in msa], dtype=np.uint8)
+
+        if msa_weights_path and os.path.exists(msa_weights_path):
+            weights = np.load(msa_weights_path)
+        elif msa_weights_path:
+            os.makedirs(os.path.dirname(msa_weights_path), exist_ok=True)
+            _, weights = self.get_weights(
+                msa=msa,
+            )
+            np.save(msa_weights_path, weights)
+        else:
+            _, weights = self.get_weights(
+                msa=msa,
+             )
+
+
+        original_msa_sample_idxs = np.arange(len(msa))
+        sample_idxs = self._get_sim_filtered_idxs(msa)
+        original_msa_sample_idxs = original_msa_sample_idxs[sample_idxs]
+
+        if self.force_include_first:
+            original_msa_sample_idxs = np.concatenate(
+                [[0], original_msa_sample_idxs[original_msa_sample_idxs != 0]]
+            )
+        return np.random.choice(len(msa), replace=False, size=size, p=weights / weights.sum()) if len(msa) >= size else original_msa_sample_idxs
+    
+def sample_msa(msa_sequences, msa_weights_path=None, context_length=200_000, max_context_sequences=200, seed=0, sort=True):
+    """Sample MSA sequences for the context"""    
+    n_sequences = min( context_length // len(msa_sequences[0]), len(msa_sequences) if max_context_sequences == 0 else max_context_sequences ) - 1 
+    sampler = MSASampler(0.98, 0.7, force_include_first=False)
+    sample_idx = sampler.get_sample_idxs(
+        msa_sequences, size=n_sequences, msa_weights_path=msa_weights_path, seed=seed
+        )
+    
+    # Sort sequences from least similar to most similar and add wild type target sequence
+    if sort:
+        context_sequences = [msa_sequences[i] for i in sample_idx][::-1]
+    
+    return context_sequences

protxlstm/applications/sample_sequences.py

@@ -0,0 +1,200 @@
+import torch
+from tqdm import tqdm
+import pickle
+import os
+import argparse
+import json
+
+from protxlstm.dataloaders import ProteinMemmapDataset
+from protxlstm.generation import generate_sequence
+from protxlstm.utils import (
+    AA_TO_ID,
+    load_model,
+) 
+from protxlstm.models.xlstm import xLSTMLMHeadModel
+from protxlstm.models.mamba import MambaLMHeadModelwithPosids
+
+
+def sample_sequences(dataset,
+                     model,
+                     family_idx,
+                     params,
+                     n_samples_per_family,
+                     max_length=1000,
+                     chunk_chunk_size=2**15,
+                     save_path=None,
+                     device="cuda:0"):
+    """
+    Function to sample sequences from the model. Given a dataset, a list of families (their indexes in the dataset)
+    and a set of generating parameters, it generates `n_samples_per_family` sequences for each family and each parameter set.
+    The function returns a dictionary with the following structure:
+    gen_seqs = {family_idx: {parameters: {sequence: perplexity}}}
+    The parameters are in a list of tuples with the following structure:    
+    parameters_list = [(nr_seqs_ctx, temperature, top_k, top_p)]
+    """        
+    gen_seqs = {}
+    gen_seqs[family_idx] = {}
+    gen_seqs[family_idx][params] = {}
+    print(f"Sampling sequences for family {family_idx} and parameters {params}.")
+
+    n_seqs_ctx , temperature, top_k, top_p = params
+    for _ in tqdm(range(n_samples_per_family)):
+        # Sample the dataset to get the input
+        data = dataset[family_idx]
+        tokens = data["input_ids"][None,:].to(device)
+        pos_ids = data["position_ids"][None,:].to(device)
+
+        start_seqs = torch.argwhere(tokens[0]==0)[:,0].cpu().numpy()
+
+        n_seqs_ctx = len(start_seqs) if len(start_seqs) < n_seqs_ctx else n_seqs_ctx
+        L = start_seqs[n_seqs_ctx]+1
+        context_tokens = tokens[:,:L]
+        context_pos_ids = pos_ids[:,:L]
+        is_fim={}
+
+        # Generate the new sequence               
+        output = generate_sequence(model,
+                                context_tokens,
+                                position_ids=context_pos_ids,
+                                is_fim=is_fim,
+                                max_length=(L+max_length),
+                                temperature=temperature,
+                                top_k=top_k,
+                                top_p=top_p,
+                                return_dict_in_generate=True,
+                                output_scores=True,
+                                eos_token_id=torch.tensor([AA_TO_ID["<cls>"]]).to(device),
+                                chunk_chunk_size=chunk_chunk_size,
+                                device=device)
+        
+        # Get the perplexity of the generated sequence
+        output_seq = output["generated"] 
+        loss = torch.nn.functional.cross_entropy(torch.from_numpy(output["scores"]).permute(0, 2, 1),
+                                                torch.from_numpy(output["generated_tokens"][0][None,:]))
+        
+        # save only sequences with length < max_length
+        if len(output_seq[0]) < max_length:
+
+            gen_seqs[family_idx][params][output_seq[0]] = {"perplexity": torch.exp(loss).item()}
+
+    if save_path is not None:
+        if not os.path.exists("evaluation/generation/generated_sequences"):
+            os.mkdir("evaluation/generation/generated_sequences")
+        if not os.path.exists(save_path):
+            os.mkdir(save_path)
+        with open(f'{save_path}/{family_idx}_{params}_{n_samples_per_family}', "wb") as f:
+            pickle.dump(gen_seqs, f)
+        print(f"Sequences saved for family {family_idx} and parameters {params}")
+
+    return gen_seqs
+
+def generate_sequences(model_name,
+                    checkpoint,
+                    family_idxs=[],
+                    parameters_list=[],
+                    n_samples_per_family = 100,
+                    chunk_size=1024,
+                    chunk_chunk_size=2**15,
+                    data_dir="data/",
+                    device="cuda:0"
+                    ):
+    
+    # Load the test dataset
+    fim_strategy = "multiple_span"
+    mask_fraction = 0.2
+
+    dataset = ProteinMemmapDataset(
+            msa_memmap_path=f"{data_dir}open_protein_set_memmap.dat",
+            msa_memmap_meta_path=f"{data_dir}open_protein_set_memmap_indices.csv",
+            subset_path=f"{data_dir}cluster_testing_set.txt",
+            sample=False,
+            max_msa_len=-1,
+            reverse=False,
+            seed=0,
+            troubleshoot=False,
+            fim_strategy=fim_strategy,
+            always_mask=False,
+            max_position_embeddings=2048,
+            max_seq_position_embeddings=512,
+            add_position_ids="1d",
+            mask_fraction=mask_fraction
+        )
+
+    if model_name == "xlstm":
+        model_class = xLSTMLMHeadModel
+    elif model_name == "mamba":
+        model_class = MambaLMHeadModelwithPosids
+
+    save_path = f"evaluation/generation/generated_sequences/{checkpoint.split('/')[-1]}"
+
+    if model_name == "xlstm":
+        config_update_kwargs = {
+                "mlstm_backend": "chunkwise_variable",
+                "mlstm_chunksize": chunk_size,
+                "mlstm_return_last_state": True
+            }
+    else:
+        config_update_kwargs = {}
+
+  
+    #load the model
+    model = load_model(checkpoint,
+                    model_class=model_class,
+                    device=device,
+                    dtype=torch.bfloat16,
+                    **config_update_kwargs,
+                    )
+    model = model.eval()
+    print("Model loaded.")
+
+    for family_idx in family_idxs:
+        for params in parameters_list:
+            params = tuple(params)
+            if not os.path.exists(f'{save_path}/{family_idx}_{params}_{n_samples_per_family}'):
+                gen_seqs = sample_sequences(
+                        dataset=dataset,
+                        model=model,
+                        family_idx=family_idx,
+                        params=params,
+                        n_samples_per_family=n_samples_per_family,
+                        chunk_chunk_size=chunk_chunk_size,
+                        save_path=save_path,
+                        device=device)
+                
+                print(f"Sampled {len(gen_seqs[family_idx][params])} valid sequences.")
+            else:
+                print(f"Sequences for family {family_idx} and parameters {params} already exist.")
+                
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(
+        description="Generate sequences."
+    )
+    parser.add_argument("--model_name", type=str, help="Either 'xlstm' or 'mamba'.")
+    parser.add_argument("--checkpoint", type=str, help="Path to model checkpoint.")
+    parser.add_argument("--family_idxs", type=str, help="List of family indices.")
+    parser.add_argument("--parameters_list", type=str, help="List of sampling parameters.")
+    parser.add_argument("--n_samples_per_family", type=int, default=100, help="Number of sequences to sample per family and parameter set.")
+    parser.add_argument("--chunk_size", type=int, default=1024, help="Chunk size for xLSTM context encoding.")
+    parser.add_argument("--chunk_chunk_size", type=int, default=2*15, help="Length of context sequence part processed at once.")
+    parser.add_argument("--data_dir", type=str, default="data/", help="Path to dataset.")
+    parser.add_argument("--device", type=str, default="cuda:0", help="Device.")
+
+    args = parser.parse_args()
+
+    family_idxs = json.loads(args.family_idxs)
+    parameters_list = json.loads(args.parameters_list)
+
+    # Run sequence generation
+    generate_sequences(
+        model_name=args.model_name,
+        checkpoint=args.checkpoint,
+        family_idxs=family_idxs,
+        parameters_list=parameters_list,
+        n_samples_per_family=args.n_samples_per_family,
+        chunk_size=args.chunk_size,
+        chunk_chunk_size=args.chunk_chunk_size,
+        data_dir=args.data_dir,
+        device=args.device,
+        )

protxlstm/applications/score_sequences.py

@@ -0,0 +1,58 @@
+import argparse
+import os
+import pickle
+
+from generation_utils.create_sequence_df import create_sequence_df
+from generation_utils.score_hamming import score_hamming
+from generation_utils.score_hmmer import score_hmmer
+from generation_utils.score_structure import score_structure
+
+
+def score_sequences(model_name,
+                    family_idx,
+                    num_sequences = 100,
+                    data_dir = "data/"):
+    
+    if os.path.isfile(f"evaluation/generation/evaluations/{model_name}/sequence_df_{family_idx}"):
+        with open(f"evaluation/generation/evaluations/{model_name}/sequence_df_{family_idx}", "rb") as f:
+            sequence_df = pickle.load(f)
+    else:
+        sequence_df = create_sequence_df(model_name, family_idx, data_dir = data_dir, num_sequences = num_sequences)
+        if not os.path.exists("evaluation/generation/evaluations/"):
+            os.mkdir("evaluation/generation/evaluations/")
+        if not os.path.exists(f"evaluation/generation/evaluations/{model_name}/"):
+            os.mkdir(f"evaluation/generation/evaluations/{model_name}/")
+        with open(f"evaluation/generation/evaluations/{model_name}/sequence_df_{family_idx}", "wb") as f:
+            pickle.dump(sequence_df, f)
+
+    if not "min_hamming" in sequence_df.columns:
+        sequence_df = score_hamming(sequence_df, family_idx, data_dir)
+        with open(f"evaluation/generation/evaluations/{model_name}/sequence_df_{family_idx}", "wb") as f:
+            pickle.dump(sequence_df, f)
+
+    if not "score_gen" in sequence_df.columns:
+        sequence_df = score_hmmer(sequence_df, family_idx, data_dir)
+        with open(f"evaluation/generation/evaluations/{model_name}/sequence_df_{family_idx}", "wb") as f:
+            pickle.dump(sequence_df, f)
+
+    if not "ptm" in sequence_df.columns:
+        sequence_df = score_structure(sequence_df, family_idx)
+        with open(f"evaluation/generation/evaluations/{model_name}/sequence_df_{family_idx}", "wb") as f:
+            pickle.dump(sequence_df, f)
+
+    return sequence_df
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(
+        description="Generate sequences."
+    )
+    parser.add_argument("--model_name", type=str, help="Either 'xlstm' or 'mamba'.")
+    parser.add_argument("--family_idx", type=int, help="Family index.")
+    parser.add_argument("--num_sequences", type=int, default=100, help="Number of sequences.")
+    parser.add_argument("--data_dir", type=str, default="./data/", help="Path to dataset.")
+
+    args = parser.parse_args()
+
+    sequence_df = score_sequences(args.model_name, args.family_idx, args.num_sequences, args.data_dir)

protxlstm/checkpoints/small/config.json

@@ -0,0 +1 @@
+{"mlstm_block": {"mlstm": {"proj_factor": 2.0, "round_proj_up_dim_up": true, "round_proj_up_to_multiple_of": 64, "_proj_up_dim": 1024, "conv1d_kernel_size": 4, "qkv_proj_blocksize": 4, "num_heads": 4, "embedding_dim": 512, "bias": false, "dropout": 0.0, "context_length": 2048, "backend": "chunkwise", "chunk_size": 1024, "return_last_state": false, "_num_blocks": 16, "_inner_embedding_dim": 1024}}, "slstm_block": {"slstm": {"hidden_size": 512, "num_heads": 4, "num_states": 4, "backend": "cuda", "function": "slstm", "bias_init": "powerlaw_blockdependent", "recurrent_weight_init": "zeros", "_block_idx": 0, "_num_blocks": 16, "num_gates": 4, "gradient_recurrent_cut": false, "gradient_recurrent_clipval": null, "forward_clipval": null, "batch_size": 8, "input_shape": "BSGNH", "internal_input_shape": "SBNGH", "output_shape": "BNSH", "constants": {}, "dtype": "bfloat16", "dtype_b": "float32", "dtype_r": "bfloat16", "dtype_w": "bfloat16", "dtype_g": "bfloat16", "dtype_s": "bfloat16", "dtype_a": "float32", "enable_automatic_mixed_precision": true, "initial_val": 0.0, "embedding_dim": 512, "conv1d_kernel_size": 4, "dropout": 0.0}, "feedforward": {"proj_factor": 1.3, "round_proj_up_dim_up": true, "round_proj_up_to_multiple_of": 64, "_proj_up_dim": 0, "act_fn": "gelu", "embedding_dim": -1, "dropout": 0.0, "bias": false, "ff_type": "ffn_gated", "_num_blocks": 1}, "_num_blocks": 16, "_block_idx": 0}, "context_length": 2048, "num_blocks": 16, "embedding_dim": 512, "add_post_blocks_norm": true, "bias": false, "dropout": 0.0, "checkpoint_blocks": true, "slstm_at": [], "_block_map": "0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0", "vocab_size": 38, "tie_weights": false, "weight_decay_on_embedding": false, "add_embedding_dropout": false, "position_embeddings": "rot_1d", "max_position_embeddings": 2048, "max_seq_position_embeddings": 512, "rope_base_frequency": 500000}

protxlstm/checkpoints/small/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bcae4c2a893afed859ca9b5926d24e8f7d0b22eca198e4aa950b80909be8e50
+size 207533690

protxlstm/checkpoints/small/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ba59f1fd544a5d9f6c4adb40730cf90b8f69a772df838246f724586cb1d602a
+size 103773526

protxlstm/checkpoints/small/rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5480500274058efdf9caa959d35a42a948ff3dc8536e082b9bc22f2ecd423108
+size 14244

protxlstm/checkpoints/small/scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25a26dbf75285e97210697d5608b6d76ef35aa0d2879be319ef2785f881153b9
+size 1000

protxlstm/data.py

@@ -0,0 +1,60 @@
+import csv
+import os
+
+import numpy as np
+from tqdm import tqdm
+
+from protxlstm.utils import load_sequences_from_msa_file, tokenizer
+
+def process_msa(msa_item):
+    msa_name, msa_path = msa_item
+    # Load an a3m file with all the context sequences
+    msa = load_sequences_from_msa_file(msa_path)
+    # Tokenize the sequences and concatenate them into a single array
+    tokens = tokenizer(msa, concatenate=True)
+    tokens = tokens.numpy()[0]
+    return msa_name, tokens
+
+def main(data_dir, output_dir):
+    msa_paths = {k: os.path.join(data_dir, k, 'a3m/uniclust30.a3m') for k in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, k))}
+    msa_items = list(msa_paths.items())
+
+    dataset_dictionary = {}
+    total_length = 0
+
+    # First pass: calculate total length of all concatenated arrays
+    for item in tqdm(msa_items):
+        try:
+            k, v = process_msa(item)
+            dataset_dictionary[k] = v
+            total_length += len(v)
+        except:
+            print(f"Error processing {item}")
+
+    # Initialize the memmap array with the calculated total length
+    memmap_path = os.path.join(output_dir, 'open_protein_set_memmap.dat')
+    concatenated_array = np.memmap(memmap_path, dtype='int8', mode='w+', shape=(total_length,))
+
+    with open(f'{output_dir}/open_protein_set_memmap_indices.csv', 'w', newline='') as csvfile:
+        csvwriter = csv.writer(csvfile)
+        
+        csvwriter.writerow(['msa_id', 'Start', 'End'])
+        
+        start_index = 0
+        for key, array in dataset_dictionary.items():
+            end_index = start_index + len(array) - 1
+            concatenated_array[start_index:end_index + 1] = array  # Write to memmap
+            csvwriter.writerow([key, start_index, end_index])
+            start_index = end_index + 1
+
+    # Ensure the data is written to disk
+    concatenated_array.flush()
+
+
+if __name__ == "__main__":
+    data_dir = 'data/a3m_files'
+    output_dir = 'data/'
+    main(data_dir, output_dir)
+
+
+

protxlstm/dataloaders.py

@@ -0,0 +1,249 @@
+# Original code from ProtMamba under Apache License 2.0.
+#
+# Modifications made by Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Uniclust30_Dataset renamed to ProteinMemmapDataset
+#       - Dataset input file format changed for more efficient dataloading
+#       - Option to use only a subset
+#   - DataCollatorForUniclust30Dataset renamed to ProteinDataCollator
+#       - Add sequence padding
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader, Dataset
+from typing import Dict, Optional, Sequence
+
+from protxlstm.fim import MultipleSpanFIM, NoFIM, SingleSpanFIM
+from protxlstm.utils import AA_TO_ID
+
+
+# Make dataset
+class ProteinMemmapDataset(Dataset):
+    """
+    ProteinMemmapDataset is a PyTorch Dataset class for handling memory-mapped datasets of protein multiple sequence alignments (MSAs).
+    
+    This class imports MSA data stored in memmap format and associated metadata CSVs. It supports flexible
+    data sampling strategies and inpainting methods for sequence manipulation and training purposes.
+    
+    Args:
+        msa_memmap_path (str): Path to the memory-mapped file containing the MSA clusters.
+        msa_memmap_meta_path (str): Path to the CSV file with metadata linking MSA Cluster IDs and indices in the memmap array.
+        subset_path (str, optional): Path to a CSV file specifying a subset of cluster IDs to use.
+        sample (bool, optional): If True, randomly samples sequences from each cluster; otherwise, loads all sequences and shuffles them.
+        max_msa_len (int, optional): Maximum length of the MSA sequences to include. Defaults to -1 (no limit).
+        reverse (bool, optional): If True, reverses sequences with a probability of 0.5 and moves the last token to the front.
+        seed (int, optional): Random seed for reproducibility. Defaults to 42.
+        troubleshoot (bool, optional): If True, prints debugging information. Defaults to False.
+        fim_strategy (str, optional): Strategy for inpainting ("no-scramble", "one_span", or "multiple_span").
+        max_patches (int, optional): Number of patches for inpainting. Used when fim_strategy is "multiple_span".
+        mask_fraction (float, optional): Fraction of the patches to mask. Used when fim_strategy is "multiple_span".
+        always_mask (bool, optional): If True, ensures masking is applied in the inpainting process.
+        max_position_embeddings (int, optional): Maximum position embeddings. Defaults to 2048.
+        max_seq_position_embeddings (int, optional): Maximum sequence position embeddings for 2D positional IDs. Defaults to 512.
+        add_position_ids (str, optional): Type of position IDs to add ("none", "1d", or "2d"). Defaults to "1d".
+    """
+    
+    _FIM = {"no-scramble": NoFIM, "one_span": SingleSpanFIM, "multiple_span": MultipleSpanFIM}
+    _POSIDS = {"none", "1d", "2d"}
+
+    def __init__(self, 
+                 msa_memmap_path=None,
+                 msa_memmap_meta_path=None,
+                 subset_path=None,
+                 sample=False,
+                 max_msa_len=-1,
+                 reverse=False,
+                 seed=42,
+                 troubleshoot=False,
+                 fim_strategy="no-scramble",
+                 max_patches=5,
+                 mask_fraction=0.2,
+                 always_mask=False,
+                 max_position_embeddings=2048,
+                 max_seq_position_embeddings=512,
+                 add_position_ids="1d", ):
+        
+        np.random.seed(seed)
+
+        if msa_memmap_path:
+            self.dataset = np.memmap(msa_memmap_path, dtype=np.int8, mode='r')
+            self.dataset_meta = pd.read_csv(msa_memmap_meta_path)
+            if subset_path:
+                subset_ids = pd.read_csv(subset_path, header=None, names=['ID'])['ID'].tolist()
+                self.dataset_meta = self.dataset_meta[self.dataset_meta['msa_id'].isin(subset_ids)]
+        else:
+            self.dataset = None
+
+        self.sample = sample
+        self.max_msa_len = max_msa_len
+        self.reverse = reverse
+        self.fim_strategy = fim_strategy
+        if fim_strategy in ProteinMemmapDataset._FIM:
+            self.fim = ProteinMemmapDataset._FIM[fim_strategy](max_patches=max_patches,
+                                                             mask_fraction=mask_fraction,
+                                                             always_mask=always_mask,
+                                                             add_position_ids=add_position_ids != "none",
+                                                             troubleshoot=troubleshoot)
+        else:
+            raise ValueError(f'Fill in the middle stragy "{fim_strategy}" not recognized.')
+        
+        self.max_position_embeddings = max_position_embeddings
+        self.max_seq_position_embeddings = max_seq_position_embeddings
+        self.add_position_ids = add_position_ids
+
+        self.troubleshoot = troubleshoot
+
+    def __len__(self):
+        # meta dataframe has one row for each MSA cluster
+        return len(self.dataset_meta)
+
+    def __getitem__(self, idx):
+        # get all the sequences in the cluster
+        sequences = self.get_sequences(idx)
+        # get total number of sequences in the cluster and choose how many to sample
+        orig_num_sequences = len(self.get_index_start_of_sequences(sequences))
+        num_sequences = np.random.randint(1, orig_num_sequences + 1) if self.sample else orig_num_sequences
+        # sample the sequences
+        sequences, position_ids = self.sample_sequences(sequences, num_sequences)
+        # with probability 0.5, reverse the sequences and move the last token to the front
+        sequences, position_ids = self.reverse_sequences(sequences, position_ids) if (
+                self.reverse and np.random.rand() > 0.5) else sequences, position_ids
+        # limit the length of the MSA
+        sequences = sequences[:self.max_msa_len] if self.max_msa_len > 0 else sequences
+        if self.add_position_ids != "none":
+            position_ids = position_ids[:self.max_msa_len] if self.max_msa_len > 0 else position_ids
+        # convert to tensor
+        sequences = torch.asarray(sequences, dtype=torch.int64)
+        position_ids = torch.asarray(position_ids, dtype=torch.int64).clamp(0,
+                                                                            self.max_position_embeddings - 1) if self.add_position_ids!="none" else None
+
+        if self.troubleshoot:
+            print(
+                f"Cluster {idx} has {orig_num_sequences} sequences, of which {num_sequences} sampled now. Total MSA length: {len(sequences)}")
+        if self.add_position_ids == "1d":
+            return dict(input_ids=sequences, position_ids=position_ids, labels=sequences)
+        if self.add_position_ids == "2d":
+            seq_position_ids = (sequences == AA_TO_ID["<cls>"]).int().cumsum(-1).clamp(0,
+                                                                                       self.max_seq_position_embeddings - 1).contiguous()
+            return dict(input_ids=sequences, position_ids=position_ids, seq_position_ids=seq_position_ids,
+                        labels=sequences)
+        return dict(input_ids=sequences, labels=sequences)
+    
+    def get_msa_id(self, idx):
+        """Get the MSA ID in the cluster with index `idx`."""
+        cluster_meta = self.dataset_meta.iloc[idx]
+        return cluster_meta.msa_id
+    
+    def get_idx_from_msa_id(self, msa_id):
+        """Get `idx` with the MSA ID"""
+        return self.dataset_meta[self.dataset_meta.msa_id == msa_id].index[0]
+
+    def get_sequences(self, idx):
+        """Get the sequences in the cluster with index `idx`."""
+        cluster_meta = self.dataset_meta.iloc[idx]
+        sequences = self.dataset[cluster_meta.Start : cluster_meta.End]
+        return sequences
+
+    def get_index_start_of_sequences(self, sequences):
+        """Get the positions of the start of each sequence in the cluster."""
+        return np.where(sequences == 0)[0]
+
+    def reverse_sequences(self, sequence, position_ids=None):
+        """Reverse the sequences and move the last token to the front."""
+        sequence = sequence[::-1]
+        if position_ids is not None:
+            position_ids = position_ids[::-1]
+        return np.concatenate([sequence[-1:], sequence[:-1]]), np.concatenate(
+            [position_ids[-1:], position_ids[:-1]]) if position_ids is not None else None
+
+    def sample_sequences(self, sequences, num_sequences, shuffle=True):
+        """Sample `num_sequences` from the sequences in the cluster."""
+        L = len(sequences)
+        # get the indexes of the start of each sequence
+        inds = self.get_index_start_of_sequences(sequences)
+        # check that there are sequences in the cluster and that there are enough of them
+        assert len(inds) > 0, "No sequences found in cluster."
+        assert len(inds) >= num_sequences, "Not enough sequences in cluster."
+        # sample n_sequences randomly from the sequences
+        if shuffle:
+            which_seqs = np.random.choice(np.arange(len(inds)), num_sequences, replace=False)
+        else:
+            which_seqs = np.arange(len(inds))[-num_sequences:]
+        # get the tuples of start and end indexes of the sequences
+        tuples = [(inds[i], inds[i + 1]) if i < len(inds) - 1 else (inds[i], L) for i in which_seqs]
+        if self.troubleshoot:
+            print(f"Sampled sequences: {tuples}")
+        # concatenate the sequences
+        sequences, position_ids = self.fim.apply(sequences, tuples)
+        return sequences, position_ids
+
+
+
+def make_dataloader(dataset):
+    """Basic function to make a dataloader.
+    """
+    dataloader = DataLoader(dataset)
+    return dataloader
+
+
+class ProteinDataCollator(object):
+    """
+    Collate examples into a batch, and pad batch to a specified maximum sequence length,
+    or to the longest sequence in the batch if max_sequence_length is None.
+    """
+    def __init__(self, max_sequence_length: Optional[int] = None):
+        """
+        Initialize the collator with an optional max_sequence_length.
+        
+        Args:
+            max_sequence_length (Optional[int]): The maximum sequence length to pad/truncate to.
+                                                 If None, pad to the longest sequence in the batch.
+        """
+        self.max_sequence_length = max_sequence_length
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "input_ids"))
+        
+        longest_seq = max(len(seq) for seq in input_ids)
+        if self.max_sequence_length is None:
+            max_len = longest_seq
+        else:
+            max_len = self.max_sequence_length
+
+        input_ids = self.pad_sequences(input_ids, max_len, padding_value=AA_TO_ID["<pad>"])
+        
+        labels = self.pad_sequences(labels, longest_seq, padding_value=AA_TO_ID["<pad>"])
+        labels = self.pad_sequences(labels, max_len, padding_value=-100)
+
+        return_dict = dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(AA_TO_ID["<pad>"])
+        )
+
+        if "position_ids" in instances[0]:
+
+            position_ids = [instance["position_ids"] for instance in instances]
+            position_ids = self.pad_sequences(position_ids, max_len, padding_value=0)
+            return_dict["position_ids"] = position_ids
+            
+            if "seq_position_ids" in instances[0]:
+                seq_position_ids = [instance["seq_position_ids"] for instance in instances]
+                seq_position_ids = self.pad_sequences(seq_position_ids, max_len, padding_value=0)
+                return_dict["seq_position_ids"] = seq_position_ids
+
+        return return_dict
+        
+    def pad_sequences(self, seqs, max_length, padding_value):
+        # truncate long sequences (redundant, already done in __getitem__, maybe safe to remove)
+        seqs = [seq[:max_length] for seq in seqs]
+
+        # pad to same length
+        seqs = torch.nn.utils.rnn.pad_sequence(seqs, batch_first=True, padding_value=padding_value)
+
+        # pad to max length
+        padding = max_length - seqs.size(1)
+        seqs = torch.nn.functional.pad(seqs, (0, padding), value=padding_value)
+
+        return seqs

protxlstm/fim.py

@@ -0,0 +1,203 @@
+
+# Original code from ProtMamba under Apache License 2.0.
+
+from protxlstm.utils import MASK_TO_ID, AA_TO_ID
+import numpy as np
+
+class AbstractFIM(object):
+    def __init__(self,
+                 max_patches=5,
+                 mask_fraction=0.2,
+                 always_mask=False,
+                 mask_tokens=MASK_TO_ID,
+                 eos_token=AA_TO_ID["<eos>"],
+                 add_position_ids=False,
+                 troubleshoot=False):
+        """
+        This class is designed to concatenate sequences based on different scrambling strategies.
+        It takes a list of sequences, tuples indicating the start and end indices of each sequence,
+        an optional number of patches to sample, and a scrambling strategy as inputs.
+        """
+        self.troubleshoot = troubleshoot
+        self.max_patches = max_patches
+        self.mask_fraction = mask_fraction
+        self.mask_tokens = mask_tokens
+        assert len(
+            self.mask_tokens) >= self.max_patches, "Number of mask tokens must be bigger than max number of patches."
+        self.eos_token = eos_token
+        self.add_position_ids = add_position_ids
+        self.always_mask = always_mask
+
+    def apply(self, sequences, tuples):
+        """
+        This function concatenates the sequences scrambling each one according to the scrambling strategy.
+        """
+        input_ids, position_ids = [], []
+        for t in tuples:
+            seq, pos = self.fim(sequences, t)
+            input_ids.extend(seq)
+            if self.add_position_ids:
+                position_ids.extend(pos)
+        if self.add_position_ids:
+            return input_ids, position_ids
+        return input_ids, None
+
+    def fim(self, sequences, t):
+        """
+        This function concatenates the sequence's parts based on the scrambling strategy.
+        """
+        raise NotImplementedError
+
+
+class NoFIM(AbstractFIM):
+    def __init__(self,
+                 max_patches=5,
+                 mask_fraction=0.2,
+                 always_mask=False,
+                 mask_tokens=MASK_TO_ID,
+                 eos_token=AA_TO_ID["<eos>"],
+                 add_position_ids=False,
+                 troubleshoot=False):
+        super().__init__(max_patches, mask_fraction, always_mask, mask_tokens, eos_token, add_position_ids, troubleshoot)
+
+    def fim(self, sequences, t):
+        """
+        This function keeps the sequence identical without any scrambling.
+        """
+        if self.add_position_ids:
+            position_ids = np.arange(t[0], t[1]) - t[0]
+            return sequences[t[0]:t[1]], position_ids
+        return sequences[t[0]:t[1]], None
+
+
+class SingleSpanFIM(AbstractFIM):
+
+    def __init__(self,
+                 max_patches=5,
+                 mask_fraction=0.2,
+                 always_mask=False,
+                 mask_tokens=MASK_TO_ID,
+                 eos_token=AA_TO_ID["<eos>"],
+                 add_position_ids=False,
+                 troubleshoot=False):
+        super().__init__(max_patches, mask_fraction, always_mask, mask_tokens, eos_token, add_position_ids, troubleshoot)
+
+    def fim(self, sequences, t):
+        """
+        This function creates and concatenates parts of the sequences based on the OpenAI scrambling strategy.
+        It randomly selects two indices within the range of the given tuple,
+        splits the sequence into three parts based on these indices, and then concatenates them with the
+        masked patch at the end
+        """
+        new_tuple = tuple(np.sort(np.random.choice(np.arange(t[0] + 1, t[1]), 2, replace=False)))
+        part1 = sequences[t[0]:new_tuple[0]]
+        part2 = sequences[new_tuple[0]:new_tuple[1]]
+        part3 = sequences[new_tuple[1]:t[1]]
+        sequence = np.concatenate([part1, [self.mask_tokens["<mask-1>"]], part3, [self.mask_tokens["<mask-1>"]], part2])
+        position_ids_sequence = None
+        if self.add_position_ids:
+            position_ids = np.arange(t[0], t[1]) - t[0]
+            position_ids_part1 = position_ids[t[0]:new_tuple[0]]
+            position_ids_part2 = position_ids[new_tuple[0]:new_tuple[1]]
+            position_ids_part3 = position_ids[new_tuple[1]:t[1]]
+            position_ids_sequence = np.concatenate(
+                [position_ids_part1, [position_ids_part2[0]], position_ids_part3, [position_ids_part2[0]],
+                 position_ids_part2])
+
+        return sequence, position_ids_sequence
+
+
+class MultipleSpanFIM(AbstractFIM):
+    def __init__(self,
+                 max_patches=5,
+                 mask_fraction=0.2,
+                 always_mask=False,
+                 mask_tokens=MASK_TO_ID,
+                 eos_token=AA_TO_ID["<eos>"],
+                 add_position_ids=False,
+                 troubleshoot=False):
+        super().__init__(max_patches, mask_fraction, always_mask, mask_tokens, eos_token, add_position_ids, troubleshoot)
+
+    def fim(self, sequences, t):
+        """
+        This function creates and concatenates parts of the sequences based on the inpaint scrambling strategy.
+        It randomly selects `2*num_patches` indices within the range of the given tuple,
+        splits the sequence into unmasked and masked parts based on these indices, and then concatenates them.
+        The number of patches is sampled from a poisson distribution with upper limit `self.max_patches` and average 1.
+        The concatenation is done by joining all unmaksed parts (interleaved with mask tokens) and afterwards
+        all masked parts (interleaved with mask tokens). At the end of the unmasked parts, a special token is added
+        to indicate the end of the unmasked parts, and at the end of the masked parts, a special token is added
+        to indicate the end of the masked parts.
+        """
+        # sample num_patches from a discrete poisson distribution with upper limit L
+        def sample_lengths(start, end):
+            """
+            Sample a length uniformly from 1 to max_L*self.mask_fraction (must be bigger than 1).
+            If the length is larger than max_L, return max_L.
+            """
+            max_L = end - start
+            length = np.random.randint(1, max(int(max_L * self.mask_fraction), 2))
+            return min(length, max_L)
+
+        # sample num_patches from a discrete poisson distribution with upper limit max_patches
+        num_patches = 1000
+        while num_patches > self.max_patches:
+            num_patches = np.random.poisson(1)
+        if self.always_mask:
+            num_patches = max(num_patches, 1)
+        # sample num_patches starting points for the masked positions (+ final position)
+        start_patches = list(np.sort(np.random.choice(np.arange(t[0] + 1, t[1]),
+                                                      num_patches,
+                                                      replace=False))) + [t[1]]
+        # sample num_patches lengths of the patches
+        len_patches = [sample_lengths(start_patches[i], start_patches[i + 1])
+                       for i in range(len(start_patches) - 1)]
+        # create masked tuples with start and end indices of the patches
+        masked_tuples = [(start_patches[i], start_patches[i] + len_patches[i]) for i in range(len(start_patches) - 1)]
+        # split the sequences into unmasked and masked parts
+        unmasked_sequence, masked_sequence, unmasked_position_ids, masked_position_ids = self.split_sequences(sequences,
+                                                                                                              t,
+                                                                                                              masked_tuples)
+
+        if self.troubleshoot:
+            print(f"For sequence in {t}: sampled {num_patches=}, {start_patches=}, {len_patches=}, {masked_tuples=}")
+        # concatenate the unmasked and masked parts
+        return unmasked_sequence + masked_sequence, unmasked_position_ids + masked_position_ids if self.add_position_ids else None
+
+    def split_sequences(self, sequences, t, masked_tuples):
+        """
+        This function splits the sequences into unmasked and masked parts based on the given tuples.
+        Args:
+            t (tuple): The start and end index of each sequence.
+            masked_tuples (list): A list of tuples specifying the indices for masked regions.
+        Returns:
+            unmasked_parts (list): The unmasked parts of the sequences interleaved with mask_tokens.
+            masked_parts (list): The masked parts of the sequences interleaved with mask_tokens.
+        """
+        unmasked_parts, masked_parts = [], []
+        unmasked_positions, masked_positions = [], []
+        position_ids = None
+        start, end = t
+        if self.add_position_ids:
+            position_ids = np.arange(start, end) - start
+        for i, region in enumerate(masked_tuples):
+            mask_token = self.mask_tokens[f"<mask-{i + 1}>"]
+            unmasked_parts.extend(sequences[start:region[0]])
+            unmasked_parts.append(mask_token)
+            masked_parts.append(mask_token)
+            masked_parts.extend(sequences[region[0]:region[1]])
+            if self.add_position_ids:
+                unmasked_positions.extend(position_ids[start-t[0]:region[0]-t[0]])
+                unmasked_positions.append(position_ids[region[0]-t[0]])
+                masked_positions.append(position_ids[region[0]-t[0]])
+                masked_positions.extend(position_ids[region[0]-t[0]:region[1]-t[0]])
+
+            start = region[1]
+        unmasked_parts.extend(sequences[start:end])
+        if self.add_position_ids:
+            unmasked_positions.extend(position_ids[start-t[0]:end-t[0]])
+        if len(masked_tuples) > 0:
+            unmasked_parts.append(self.eos_token)
+            if self.add_position_ids:
+                unmasked_positions.append(0)
+        return unmasked_parts, masked_parts, unmasked_positions, masked_positions

protxlstm/generation.py

@@ -0,0 +1,384 @@
+# Original code from ProtMamba under Apache License 2.0.
+#
+# Modifications made by Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Add option to pass input state for generation
+#   - Add functions to generate sequences with xlstm
+
+import numpy as np
+import torch
+from protxlstm.mamba_utils_generation  import (
+    InferenceParams,
+    GenerationMixin,
+    GreedySearchDecoderOnlyOutput,
+    modify_logits_for_top_p_filtering, 
+    modify_logits_for_min_p_filtering, 
+    modify_logit_for_repetition_penalty,
+    SampleDecoderOnlyOutput,
+    update_graph_cache
+)
+
+from protxlstm.utils import AA_TO_ID, decode_sequence
+
+def sample_safe(logits, top_k=1, top_p=0.0, min_p=0.0, temperature=1.0):
+    """Sample from top-k logits.
+    Arguments:
+        logits: Tensor of shape (batch_size, vocab_size)
+    """
+    if top_k == 1:  # Short-circuit for greedy decoding
+        return logits.argmax(dim=-1)
+    else:
+        if top_p > 0.0:
+            assert top_p <= 1.0, "top-p should be in (0, 1]."
+        if top_k > 0:
+            top_k = min(top_k, logits.size(-1))  # Safety check
+            logits_top, indices = torch.topk(logits, top_k, dim=-1)
+            if temperature != 1.0:
+                logits_top /= temperature
+            modify_logits_for_top_p_filtering(logits_top, top_p)
+
+            return indices[
+                torch.arange(indices.shape[0], device=indices.device),
+                torch.multinomial(
+                    torch.softmax(logits_top, dim=-1), num_samples=1
+                ).squeeze(dim=-1),
+            ]
+        else:
+            if min_p > 0.0:
+                logits_top = logits.clone()
+                max_prob = logits_top[..., 0].item()
+                min_prob = max_prob * min_p
+                modify_logits_for_min_p_filtering(logits_top, min_p)
+                if temperature != 1.0:
+                    logits_top /= temperature
+                return torch.multinomial(
+                    torch.softmax(logits_top, dim=-1), num_samples=1
+                ).squeeze(dim=-1)
+            # Clone so that when we modify for top_p we don't change the original logits
+            logits_top = logits / temperature if temperature != 1.0 else logits.clone()
+            modify_logits_for_top_p_filtering(logits_top, top_p)
+            return torch.multinomial(
+                torch.softmax(logits_top, dim=-1), num_samples=1
+            ).squeeze(dim=-1)
+
+
+@torch.inference_mode()
+def decode_safe(
+    input_ids,
+    position_ids,
+    seq_position_ids,
+    is_fim,
+    model,
+    max_length,
+    state=None,
+    top_k=1,
+    top_p=0.0,
+    min_p=0.0,
+    temperature=1.0,
+    repetition_penalty=1.0,
+    eos_token_id=None,
+    teacher_outputs=None,
+    vocab_size=None,
+    cg=False,
+    enable_timing=False,
+    streamer = None,
+    chunk_chunk_size = 2**15,
+):
+    """Decoding, either greedy or with top-k or top-p sampling.
+    If top-k = 0, don't limit the number of candidates (pure sampling).
+    Top-k and top-p can be used together. If top_k > 0 and top_p > 0, then top-k is applied first,
+    then top-p.
+    We assume that all sequences in the same batch have the same length.
+
+    Arguments:
+        input_ids: (batch, seq_len)
+        max_length: int
+        is_fim: dictionary with mask indices and associated position indices
+        teacher_outputs (optional): (batch, seq_len). If provided, instead of sampling from the
+            logits, the next token is taken from the teacher_outputs. Useful for testing.
+    Returns: GreedySearchDecoderOnlyOutput or SampleDecoderOnlyOutput, with the following fields:
+        sequences: (batch, max_length)
+        scores: tuples of (batch, vocab_size)
+    """
+    if streamer is not None:
+        streamer.put(input_ids.cpu())
+
+    batch_size, seqlen_og = input_ids.shape
+    teacher_output_len = teacher_outputs.shape[1] if teacher_outputs is not None else 0
+    if cg:
+        if not hasattr(model, "_decoding_cache"):
+            model._decoding_cache = None
+        model._decoding_cache = update_graph_cache(
+            model,
+            model._decoding_cache,
+            batch_size,
+            seqlen_og,
+            max_length,
+        )
+        inference_params = model._decoding_cache.inference_params
+        inference_params.reset(max_length, batch_size)
+    else:
+        inference_params = InferenceParams(
+            max_seqlen=max_length, max_batch_size=batch_size
+        )
+
+    def get_logits(input_ids, position_ids, seq_position_ids, inference_params):
+        decoding = inference_params.seqlen_offset > 0
+        if not cg or not decoding:
+            logits = model(
+                input_ids,
+                position_ids=position_ids,
+                seq_position_ids=seq_position_ids,
+                inference_params=inference_params,
+                num_last_tokens=1,
+            ).logits.squeeze(dim=1)
+        else:
+            logits = model._decoding_cache.run(
+                input_ids,
+                position_ids,
+                inference_params.seqlen_offset,
+                seq_position_ids=seq_position_ids,
+            ).squeeze(dim=1)
+        return logits[..., :vocab_size] if vocab_size is not None else logits
+    
+    def get_xlstm_logits_step(input_ids, position_ids, seq_position_ids, state):
+
+        if not input_ids.shape[1] == 1:
+
+            for i in range(input_ids.shape[1]):
+                if position_ids != None:
+                    token_position_ids = position_ids[:,i:(i+1)]
+                else:
+                    token_position_ids = None
+                if seq_position_ids != None:
+                    token_seq_position_ids = seq_position_ids[:,i:(i+1)]
+                else:
+                    token_seq_position_ids = None
+                logits, state = model.step(input_ids[:,i:(i+1)], state, position_ids=token_position_ids, seq_position_ids=token_seq_position_ids)
+            
+        else:
+            logits, state = model.step(input_ids, state, position_ids=position_ids, seq_position_ids=seq_position_ids)
+
+        logits = logits.squeeze(dim=1)
+        if vocab_size is not None:
+            logits = logits[..., :vocab_size]
+
+        return logits, state
+    
+    def get_xlstm_logits_chunkwise(input_ids, position_ids, seq_position_ids, chunk_chunk_size=2**15, state=None):
+
+        assert model.config.config_dataclass.mlstm_block.mlstm.backend == "chunkwise_variable"
+
+        for chunk in range(input_ids.shape[1]//chunk_chunk_size+1):
+
+            start_idx = chunk*chunk_chunk_size
+            end_idx = min((chunk+1)*chunk_chunk_size, input_ids.shape[1])
+
+            if start_idx == end_idx:
+                pass
+            
+            else:
+                input_ids_chunk = input_ids[:, start_idx:end_idx]
+
+                if not position_ids == None:
+                    position_ids_chunk = position_ids[:, start_idx:end_idx]
+                else:
+                    position_ids_chunk = None
+
+                if not seq_position_ids == None:
+                    seq_position_ids_chunk = seq_position_ids[:, start_idx:end_idx]
+                else:
+                    seq_position_ids_chunk = None   
+
+                outputs = model(input_ids_chunk, position_ids=position_ids_chunk, seq_position_ids=seq_position_ids_chunk, state=state)
+                logits, state = outputs.logits, outputs.state
+
+        logits = logits[:,-1,:]
+        logits = logits.squeeze(dim=1)
+        if vocab_size is not None:
+            logits = logits[..., :vocab_size]
+
+        return logits, state    
+           
+    def sample_tokens(logits, inference_params):
+        if (
+            teacher_outputs is None
+            or teacher_output_len <= inference_params.seqlen_offset
+        ):
+            token = sample_safe(
+                logits, top_k=top_k, top_p=top_p, min_p=min_p, temperature=temperature
+            )
+        else:
+            token = teacher_outputs[:, inference_params.seqlen_offset]
+        # return rearrange(token, "b -> b 1")
+        return token.unsqueeze(1)
+
+    def get_fim_position_id(
+        last_position_ids, sampled_tokens, is_fim, repeat_next=False
+    ):
+        if type(is_fim) is dict:
+            val = int(last_position_ids) + 1
+            should_repeat_next = False
+            if is_fim and int(sampled_tokens) in is_fim:
+                val = is_fim[int(sampled_tokens)]
+                should_repeat_next = True
+            elif repeat_next:
+                val = int(last_position_ids)
+            return torch.full_like(last_position_ids, fill_value=val), should_repeat_next
+        else:
+            t = [get_fim_position_id(last_position_ids_, sampled_tokens_, is_fim_dict, repeat_next) for
+                 (last_position_ids_, sampled_tokens_, is_fim_dict) in
+                 zip(last_position_ids, sampled_tokens, is_fim)]
+            return torch.stack([t_[0] for t_ in t], dim=0), t[0][1]
+
+    def should_stop(current_token, inference_params):
+        if inference_params.seqlen_offset == 0:
+            return False
+        if eos_token_id is not None and (current_token == eos_token_id).any():
+            if current_token.shape[1] > 1:
+                raise NotImplementedError("Batched eos_token_id not supported")
+            return True
+        if inference_params.seqlen_offset >= max_length - 1:
+            return True
+        return False
+
+    start = torch.cuda.Event(enable_timing=enable_timing)
+    end = torch.cuda.Event(enable_timing=enable_timing)
+
+    if enable_timing:
+        start.record()
+    scores, sequences = [], [input_ids]
+    new_position_ids, new_seq_position_ids = [position_ids], [seq_position_ids]
+    sequences_cat = input_ids
+    repeat_next = False
+    if position_ids.shape[0] > 1:
+        raise NotImplementedError("Batched generation with position_ids not supported")
+    
+    encode_context=True
+    while not should_stop(sequences[-1], inference_params):
+
+        from protxlstm.models.xlstm import xLSTMLMHeadModel
+        if isinstance(model, xLSTMLMHeadModel):
+            if encode_context:
+                with torch.no_grad():
+                    logits, state = get_xlstm_logits_chunkwise(sequences[-1], new_position_ids[-1], new_seq_position_ids[-1], state=state, chunk_chunk_size=chunk_chunk_size)
+                encode_context = False
+            else:
+                logits, state = get_xlstm_logits_step(sequences[-1], new_position_ids[-1], new_seq_position_ids[-1], state=state)
+        else:
+            logits = get_logits(sequences[-1], new_position_ids[-1], new_seq_position_ids[-1], inference_params)
+
+        scores.append(logits)
+    
+        inference_params.seqlen_offset += sequences[-1].shape[1]
+        if repetition_penalty == 1.0:
+            sampled_tokens = sample_tokens(scores[-1], inference_params)
+        else:
+            logits = modify_logit_for_repetition_penalty(
+                scores[-1].clone(), sequences_cat, repetition_penalty
+            )
+            sampled_tokens = sample_tokens(logits, inference_params)
+            sequences_cat = torch.cat([sequences_cat, sampled_tokens], dim=1)
+        sequences.append(sampled_tokens)
+        # Update position_ids
+        if position_ids is not None:
+            last_position_ids, repeat_next = get_fim_position_id(
+                new_position_ids[-1][:, -1:], sampled_tokens, is_fim, repeat_next
+            )
+            new_position_ids.append(last_position_ids)
+        # Update seq_position_ids
+        if seq_position_ids is not None:
+            new_seq_position_ids.append(new_seq_position_ids[-1][:, -1:])
+
+        if streamer is not None:
+            streamer.put(sampled_tokens.cpu())
+    if streamer is not None:
+        streamer.end()
+    if enable_timing:
+        end.record()
+        torch.cuda.synchronize()
+        print(f"Prompt processing + decoding time: {(start.elapsed_time(end)):.0f}ms")
+    output_cls = (
+        GreedySearchDecoderOnlyOutput if top_k == 1 else SampleDecoderOnlyOutput
+    )
+    return output_cls(sequences=torch.cat(sequences, dim=1), scores=tuple(scores))
+
+
+class GenerationMixinSafe(GenerationMixin):
+
+    def generate(
+        self,
+        input_ids,
+        position_ids,
+        seq_position_ids,
+        is_fim=None,
+        state=None,
+        max_length=1,
+        top_k=1,
+        top_p=0.0,
+        min_p=0.0,
+        temperature=1.0,
+        return_dict_in_generate=False,
+        output_scores=False,
+        chunk_chunk_size=2**15,
+        **kwargs,
+    ):
+
+        output = decode_safe(
+            input_ids,
+            position_ids,
+            seq_position_ids,
+            is_fim,
+            self,
+            max_length,
+            state=state,
+            top_k=top_k,
+            top_p=top_p,
+            min_p=min_p,
+            temperature=temperature,
+            chunk_chunk_size=chunk_chunk_size,
+            **kwargs,
+        )
+        if not output_scores:
+            output.scores = None
+        return output if return_dict_in_generate else output.sequences
+
+
+def generate_sequence(model, tokens, position_ids=None, seq_position_ids=None, state=None, is_fim=False, max_length=2000, temperature=1., top_p=0.0, top_k=1,
+                      return_dict_in_generate=False, output_scores=False, eos_token_id=AA_TO_ID["<cls>"], device="cuda", chunk_chunk_size=2**15):
+    """Generating, either greedy or with top-k or top-p sampling.
+    If top-k = 0, don't limit the number of candidates (pure sampling).
+    Top-k and top-p can be used together. If top_k > 0 and top_p > 0, then top-k is applied first,
+    then top-p. We assume that all sequences in the same batch have the same length.
+    """
+    input_ids = tokens.to(device)
+    position_ids = position_ids.to(device) if position_ids is not None else None
+    seq_position_ids = seq_position_ids.to(device) if seq_position_ids is not None else None
+    # generate sequence
+    out = model.generate(input_ids=input_ids,
+                         position_ids=position_ids,
+                         seq_position_ids=seq_position_ids,
+                         is_fim=is_fim,
+                         state=state,
+                         max_length=max_length,
+                         temperature=temperature,
+                         top_p=top_p,
+                         top_k=top_k,
+                         return_dict_in_generate=return_dict_in_generate,
+                         output_scores=output_scores,
+                         eos_token_id=eos_token_id,
+                         chunk_chunk_size=chunk_chunk_size,
+                         )
+    sequences = out.sequences
+    dic = {"input": [decode_sequence(seq) for seq in sequences[:, :input_ids.shape[-1]].cpu().numpy()],
+            "generated": [decode_sequence(seq) for seq in sequences[:, input_ids.shape[-1]:].cpu().numpy()],
+            "input_tokens": [seq for seq in sequences[:, :input_ids.shape[-1]].cpu().numpy()],
+            "generated_tokens": [seq for seq in sequences[:, input_ids.shape[-1]:].cpu().numpy()]}
+    if output_scores:
+        dic["scores"] = np.array([el.to(torch.float32).cpu().numpy() for el in out.scores]).transpose(1, 0, 2)
+    return dic
+
+
+
+
+
+

protxlstm/index.html

@@ -0,0 +1,16 @@
+<!DOCTYPE HTML PUBLIC "-//W3C//DTD HTML 3.2 Final//EN">
+<html>
+ <head>
+  <title>Index of /research/Bio-xLSTM/downloads/Prot-xLSTM/checkpoints/protxlstm_26M_30B</title>
+ </head>
+ <body>
+<h1>Index of /research/Bio-xLSTM/downloads/Prot-xLSTM/checkpoints/protxlstm_26M_30B</h1>
+<pre><img src="/icons/blank.gif" alt="Icon "> <a href="?C=N;O=D">Name</a>                                                  <a href="?C=M;O=A">Last modified</a>      <a href="?C=S;O=A">Size</a>  <a href="?C=D;O=A">Description</a><hr><img src="/icons/back.gif" alt="[PARENTDIR]"> <a href="/research/Bio-xLSTM/downloads/Prot-xLSTM/checkpoints/">Parent Directory</a>                                                           -   
+<img src="/icons/unknown.gif" alt="[   ]"> <a href="config.json">config.json</a>                                           2024-11-04 14:36  1.8K  
+<img src="/icons/unknown.gif" alt="[   ]"> <a href="optimizer.pt">optimizer.pt</a>                                          2024-11-04 14:36  198M  
+<img src="/icons/binary.gif" alt="[   ]"> <a href="pytorch_model.bin">pytorch_model.bin</a>                                     2024-11-04 14:36   99M  
+<img src="/icons/unknown.gif" alt="[   ]"> <a href="rng_state.pth">rng_state.pth</a>                                         2024-11-04 14:36   14K  
+<img src="/icons/unknown.gif" alt="[   ]"> <a href="scheduler.pt">scheduler.pt</a>                                          2024-11-04 14:36  1.0K  
+<img src="/icons/unknown.gif" alt="[   ]"> <a href="trainer_state.json">trainer_state.json</a>                                    2024-11-04 14:36  2.4M  
+<hr></pre>
+</body></html>

protxlstm/mamba_utils_generation.py

@@ -0,0 +1,382 @@
+# From: https://github.com/state-spaces/mamba/blob/main/mamba_ssm/utils/generation.py
+# Copyright (c) 2023, Albert Gu, Tri Dao.
+import gc
+from dataclasses import dataclass, field
+from typing import Callable, Optional
+
+import torch
+from torch import Tensor
+from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer
+
+
+@dataclass
+class InferenceParams:
+    """Inference parameters that are passed to the main model in order
+    to efficienly calculate and store the context during inference."""
+
+    max_seqlen: int
+    max_batch_size: int
+    seqlen_offset: int = 0
+    batch_size_offset: int = 0
+    key_value_memory_dict: dict = field(default_factory=dict)
+    lengths_per_sample: Optional[Tensor] = None
+
+    def reset(self, max_seqlen, max_batch_size):
+        self.max_seqlen = max_seqlen
+        self.max_batch_size = max_batch_size
+        self.seqlen_offset = 0
+        if self.lengths_per_sample is not None:
+            self.lengths_per_sample.zero_()
+
+
+def modify_logits_for_min_p_filtering(logits, min_p):
+    """Set the logits for none min_p values to -inf. Done in-place."""
+    if min_p <= 0.0 or min_p >= 1.0:
+        return
+    indices_to_remove = logits < min_p
+    logits.masked_fill_(indices_to_remove, float("-Inf"))
+# https://github.com/NVIDIA/Megatron-LM/blob/0bb597b42c53355a567aba2a1357cc34b9d99ddd/megatron/text_generation/sampling.py
+# https://github.com/huggingface/transformers/blob/a44985b41cfa2de48a5e1de7f1f93b7483da25d1/src/transformers/generation/logits_process.py#L231
+def modify_logits_for_top_k_filtering(logits, top_k):
+    """Set the logits for none top-k values to -inf. Done in-place."""
+    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+    logits.masked_fill_(indices_to_remove, float("-Inf"))
+
+
+# https://github.com/NVIDIA/Megatron-LM/blob/0bb597b42c53355a567aba2a1357cc34b9d99ddd/megatron/text_generation/sampling.py
+# https://github.com/huggingface/transformers/blob/a44985b41cfa2de48a5e1de7f1f93b7483da25d1/src/transformers/generation/logits_process.py#L170
+def modify_logits_for_top_p_filtering(logits, top_p):
+    """Set the logits for none top-p values to -inf. Done in-place."""
+    if top_p <= 0.0 or top_p >= 1.0:
+        return
+    # First sort and calculate cumulative sum of probabilities.
+    sorted_logits, sorted_indices = torch.sort(logits, descending=False)
+    cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+    # Remove tokens with cumulative top_p above the threshold (token with 0 are kept)
+    sorted_indices_to_remove = cumulative_probs <= (1 - top_p)
+    # scatter sorted tensors to original indexing
+    indices_to_remove = sorted_indices_to_remove.scatter(
+        1, sorted_indices, sorted_indices_to_remove
+    )
+    logits.masked_fill_(indices_to_remove, float("-inf"))
+
+
+def modify_logit_for_repetition_penalty(logits, prev_output_tokens, repetition_penalty=1.0):
+    """Apply repetition penalty. See https://arxiv.org/abs/1909.05858
+    logits: (batch_size, vocab_size)
+    prev_output_tokens: (batch_size, seq_len)
+    """
+    if repetition_penalty == 1.0:
+        return logits
+    score = torch.gather(logits, 1, prev_output_tokens)
+    # if score < 0 then repetition penalty has to be multiplied to reduce the previous token probability
+    score = torch.where(score < 0, score * repetition_penalty, score / repetition_penalty)
+    logits.scatter_(1, prev_output_tokens, score)
+    return logits
+
+
+def sample(logits, top_k=1, top_p=0.0, min_p=0.0, temperature=1.0):
+    """Sample from top-k logits.
+    Arguments:
+        logits: Tensor of shape (batch_size, vocab_size)
+    """
+    if top_k == 1:  # Short-circuit for greedy decoding
+        return logits.argmax(dim=-1)
+    else:
+        if top_p > 0.0:
+            assert top_p <= 1.0, "top-p should be in (0, 1]."
+        if top_k > 0:
+            top_k = min(top_k, logits.size(-1))  # Safety check
+            logits_top, indices = torch.topk(logits, top_k, dim=-1)
+            if temperature != 1.0:
+                logits_top /= temperature
+            modify_logits_for_top_p_filtering(logits_top, top_p)
+            return indices[
+                torch.arange(indices.shape[0], device=indices.device),
+                torch.multinomial(torch.softmax(logits_top, dim=-1), num_samples=1).squeeze(dim=-1),
+            ]
+        else:
+            if min_p > 0.0:
+                logits_top = logits.clone()
+                max_prob = logits_top[..., 0].item()
+                min_prob = max_prob * min_p
+                modify_logits_for_min_p_filtering(logits_top, min_prob)
+                if temperature != 1.0:
+                    logits_top /= temperature
+                return torch.multinomial(torch.softmax(logits_top, dim=-1), num_samples=1).squeeze(dim=-1)
+            # Clone so that when we modify for top_p we don't change the original logits
+            logits_top = logits / temperature if temperature != 1.0 else logits.clone()
+            modify_logits_for_top_p_filtering(logits_top, top_p)
+            return torch.multinomial(torch.softmax(logits_top, dim=-1), num_samples=1).squeeze(
+                dim=-1
+            )
+
+
+@torch.inference_mode()
+def decode(
+    input_ids,
+    model,
+    max_length,
+    top_k=1,
+    top_p=0.0,
+    min_p=0.0,
+    temperature=1.0,
+    repetition_penalty=1.0,
+    eos_token_id=None,
+    teacher_outputs=None,
+    vocab_size=None,
+    cg=False,
+    enable_timing=False,
+    streamer: Optional[TextStreamer] = None
+):
+    """Decoding, either greedy or with top-k or top-p sampling.
+    If top-k = 0, don't limit the number of candidates (pure sampling).
+    Top-k and top-p can be used together. If top_k > 0 and top_p > 0, then top-k is applied first,
+    then top-p.
+    We assume that all sequences in the same batch have the same length.
+
+    Arguments:
+        input_ids: (batch, seq_len)
+        max_length: int
+        teacher_outputs (optional): (batch, seq_len). If provided, instead of sampling from the
+            logits, the next token is taken from the teacher_outputs. Useful for testing.
+    Returns: GreedySearchDecoderOnlyOutput or SampleDecoderOnlyOutput, with the following fields:
+        sequences: (batch, max_length)
+        scores: tuples of (batch, vocab_size)
+    """
+    if streamer is not None:
+        streamer.put(input_ids.cpu())
+
+    batch_size, seqlen_og = input_ids.shape
+    teacher_output_len = teacher_outputs.shape[1] if teacher_outputs is not None else 0
+    if cg:
+        if not hasattr(model, "_decoding_cache"):
+            model._decoding_cache = None
+        model._decoding_cache = update_graph_cache(
+            model,
+            model._decoding_cache,
+            batch_size,
+            seqlen_og,
+            max_length,
+        )
+        inference_params = model._decoding_cache.inference_params
+        inference_params.reset(max_length, batch_size)
+    else:
+        inference_params = InferenceParams(max_seqlen=max_length, max_batch_size=batch_size)
+
+    def get_logits(input_ids, inference_params):
+        decoding = inference_params.seqlen_offset > 0
+        if decoding:
+            position_ids = torch.full(
+                (batch_size, 1),
+                inference_params.seqlen_offset,
+                dtype=torch.long,
+                device=input_ids.device,
+            )
+        else:
+            position_ids = None
+        if not cg or not decoding:
+            logits = model(
+                input_ids,
+                position_ids=position_ids,
+                inference_params=inference_params,
+                num_last_tokens=1,
+            ).logits.squeeze(dim=1)
+        else:
+            logits = model._decoding_cache.run(
+                input_ids, position_ids, inference_params.seqlen_offset
+            ).squeeze(dim=1)
+        return logits[..., :vocab_size] if vocab_size is not None else logits
+
+    def sample_tokens(logits, inference_params):
+        if teacher_outputs is None or teacher_output_len <= inference_params.seqlen_offset:
+            token = sample(logits, top_k=top_k, top_p=top_p, min_p=min_p, temperature=temperature)
+        else:
+            token = teacher_outputs[:, inference_params.seqlen_offset]
+        # return rearrange(token, "b -> b 1")
+        return token.unsqueeze(1)
+
+    def should_stop(current_token, inference_params):
+        if inference_params.seqlen_offset == 0:
+            return False
+        if eos_token_id is not None and (current_token == eos_token_id).all():
+            return True
+        if inference_params.seqlen_offset >= max_length - 1:
+            return True
+        return False
+
+    start = torch.cuda.Event(enable_timing=enable_timing)
+    end = torch.cuda.Event(enable_timing=enable_timing)
+
+    if enable_timing:
+        start.record()
+    scores, sequences = [], [input_ids]
+    sequences_cat = input_ids
+    while not should_stop(sequences[-1], inference_params):
+        scores.append(get_logits(sequences[-1], inference_params))
+        inference_params.seqlen_offset += sequences[-1].shape[1]
+        if repetition_penalty == 1.0:
+            sampled_tokens = sample_tokens(scores[-1], inference_params)
+        else:
+            logits = modify_logit_for_repetition_penalty(
+                scores[-1].clone(), sequences_cat, repetition_penalty
+            )
+            sampled_tokens = sample_tokens(logits, inference_params)
+            sequences_cat = torch.cat([sequences_cat, sampled_tokens], dim=1)
+        sequences.append(sampled_tokens)
+        if streamer is not None:
+            streamer.put(sampled_tokens.cpu())
+    if streamer is not None:
+        streamer.end()
+    if enable_timing:
+        end.record()
+        torch.cuda.synchronize()
+        print(f"Prompt processing + decoding time: {(start.elapsed_time(end)):.0f}ms")
+    output_cls = GreedySearchDecoderOnlyOutput if top_k == 1 else SampleDecoderOnlyOutput
+    return output_cls(sequences=torch.cat(sequences, dim=1), scores=tuple(scores))
+
+
+class GenerationMixin:
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        raise NotImplementedError
+
+    def generate(
+        self,
+        input_ids,
+        max_length,
+        top_k=1,
+        top_p=0.0,
+        min_p=0.0,
+        temperature=1.0,
+        return_dict_in_generate=False,
+        output_scores=False,
+        **kwargs,
+    ):
+        output = decode(
+            input_ids, self, max_length, top_k=top_k, top_p=top_p, min_p = min_p, temperature=temperature, **kwargs
+        )
+        if not output_scores:
+            output.scores = None
+        return output if return_dict_in_generate else output.sequences
+
+
+@dataclass
+class DecodingCGCache:
+    max_batch_size: int = 0
+    max_seqlen: int = 0
+    device = None
+    dtype = None
+    callables: dict = field(default_factory=dict)
+    mempool = None
+    inference_params: Optional[InferenceParams] = None
+    run: Optional[Callable] = None
+
+
+@torch.inference_mode()
+def update_graph_cache(
+    model,
+    cache,
+    batch_size,
+    seqlen_og,
+    max_seqlen,
+    decoding_seqlens=(1,),
+    dtype=None,
+    n_warmups=2,
+):
+    if cache is None:
+        cache = DecodingCGCache()
+    param_example = next(iter(model.parameters()))
+    device = param_example.device
+    if dtype is None:
+        dtype = param_example.dtype
+    if (
+        (device, dtype) != (cache.device, cache.dtype)
+        or batch_size > cache.max_batch_size
+        or max_seqlen > cache.max_seqlen
+    ):  # Invalidate the cache
+        cache.callables = {}
+        cache.mempool = None
+        cache.inference_params = None
+        gc.collect()
+        cache.device, cache.dtype = device, dtype
+        cache.max_batch_size, cache.max_seqlen = batch_size, max_seqlen
+        assert hasattr(model, "allocate_inference_cache"), "CUDA graph decoding requires that the model has a method allocate_inference_cache"
+        inf_cache = model.allocate_inference_cache(batch_size, max_seqlen, dtype)
+        lengths_per_sample = torch.full((batch_size,), seqlen_og, dtype=torch.int32, device=device)
+        cache.inference_params = InferenceParams(
+            max_seqlen=max_seqlen,
+            max_batch_size=batch_size,
+            seqlen_offset=seqlen_og,
+            key_value_memory_dict=inf_cache,
+            lengths_per_sample=lengths_per_sample,
+        )
+        cache.mempool = torch.cuda.graphs.graph_pool_handle()
+    for decoding_seqlen in decoding_seqlens:
+        if (batch_size, decoding_seqlen) not in cache.callables:
+            cache.callables[batch_size, decoding_seqlen] = capture_graph(
+                model,
+                cache.inference_params,
+                batch_size,
+                max_seqlen,
+                decoding_seqlen=decoding_seqlen,
+                mempool=cache.mempool,
+                n_warmups=n_warmups,
+            )
+
+    def dispatch(input_ids, position_ids, seqlen):
+        batch_size, decoding_seqlen = input_ids.shape[:2]
+        return cache.callables[batch_size, decoding_seqlen](input_ids, position_ids, seqlen)
+
+    cache.run = dispatch
+    cache.inference_params.seqlen_offset = 0  # Reset so it's not confusing
+    return cache
+
+
+def capture_graph(
+    model, inference_params, batch_size, max_seqlen, decoding_seqlen=1, mempool=None, n_warmups=2
+):
+    device = next(iter(model.parameters())).device
+    input_ids = torch.full((batch_size, decoding_seqlen), 0, dtype=torch.long, device=device)
+    position_ids = torch.full((batch_size, decoding_seqlen), 0, dtype=torch.long, device=device)
+    seqlen_offset_og = inference_params.seqlen_offset
+    inference_params.seqlen_offset = max_seqlen - decoding_seqlen
+    inference_params.lengths_per_sample[:] = inference_params.seqlen_offset
+
+    # Warmup before capture
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(n_warmups):
+            logits = model(
+                input_ids,
+                position_ids=position_ids,
+                inference_params=inference_params,
+                num_last_tokens=decoding_seqlen,
+            ).logits
+        s.synchronize()
+        # This might be needed for correctness if we run with NCCL_GRAPH_MIXING_SUPPORT=0,
+        # which requires that graph launch and non-captured launch to not overlap (I think,
+        # that's how I interpret the documentation). I'm not sure if this is required.
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+    torch.cuda.current_stream().wait_stream(s)
+    # Captures the graph
+    # To allow capture, automatically sets a side stream as the current stream in the context
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, pool=mempool):
+        logits = model(
+            input_ids,
+            position_ids=position_ids,
+            inference_params=inference_params,
+            num_last_tokens=decoding_seqlen,
+        ).logits
+
+    def run(new_input_ids, new_position_ids, seqlen):
+        inference_params.lengths_per_sample[:] = seqlen
+        input_ids.copy_(new_input_ids)
+        position_ids.copy_(new_position_ids)
+        graph.replay()
+        return logits.clone()
+
+    inference_params.seqlen_offset = seqlen_offset_og
+    return run

protxlstm/models/llama.py

@@ -0,0 +1,342 @@
+import json
+import math
+import os
+from collections import namedtuple
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig
+
+from protxlstm.xlstm.components.rotary_position import compute_freqs_cis
+
+# Note: generation capabilities are not implemented for the transformer
+
+class TransformerConfig(PretrainedConfig):
+
+    model_type = "llama"
+
+    def __init__(
+        self,
+        d_model,
+        n_layer,
+        n_heads,
+        n_kv_heads,
+        bidirectional,
+        vocab_size,
+        hidden_dim,
+        multiple_of,  # MLP hidden layer size will be multiple of
+        norm_eps,
+        max_length,
+        dropout,
+        max_position_embeddings,
+        rope_base_frequency,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        # default hyperparameters for the Llama 7B model
+        self.dim = d_model
+        self.n_layers = n_layer
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.causal_attention = not bidirectional
+        self.vocab_size = vocab_size
+        self.hidden_dim = hidden_dim
+        self.multiple_of = multiple_of
+        self.norm_eps = norm_eps
+        self.max_seq_len = max_length
+        self.dropout = dropout
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_base_frequency = rope_base_frequency
+
+class RMSNorm_transformer(torch.nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t, freqs).float()  # type: ignore
+    freqs_cos = torch.cos(freqs)  # real part
+    freqs_sin = torch.sin(freqs)  # imaginary part
+    return freqs_cos, freqs_sin
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(shape)
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    # reshape xq and xk to match the complex representation
+    xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
+    xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)
+
+    # reshape freqs_cos and freqs_sin for broadcasting
+    freqs_cos = reshape_for_broadcast(freqs_cos, xq_r)
+    freqs_sin = reshape_for_broadcast(freqs_sin, xq_r)
+
+    # apply rotation using real numbers
+    xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
+    xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
+    xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin
+    xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos
+
+    # flatten last two dimensions
+    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3)
+    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3)
+
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
+    bs, slen, n_kv_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        x[:, :, :, None, :]
+        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
+        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
+    )
+
+class Attention(nn.Module):
+    def __init__(self, args: TransformerConfig):
+        super().__init__()
+        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
+        assert args.n_heads % self.n_kv_heads == 0
+        model_parallel_size = 1
+        self.n_local_heads = args.n_heads // model_parallel_size
+        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
+        self.n_rep = self.n_local_heads // self.n_local_kv_heads
+        self.head_dim = args.dim // args.n_heads
+        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+        self.attn_dropout = nn.Dropout(args.dropout)
+        self.resid_dropout = nn.Dropout(args.dropout)
+        self.dropout = args.dropout
+        self.causal_attention = args.causal_attention
+
+        # use flash attention or a manual implementation?
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash and self.causal_attention:
+            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+            mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
+            mask = torch.triu(mask, diagonal=1)
+            self.register_buffer("mask", mask)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs_cos: torch.Tensor,
+        freqs_sin: torch.Tensor,
+    ):
+        bsz, seqlen, _ = x.shape
+
+        # QKV
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
+
+        # RoPE relative positional embeddings
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)
+
+        # grouped multiquery attention: expand out keys and values
+        xk = repeat_kv(xk, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
+        xv = repeat_kv(xv, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
+
+        # make heads into a batch dimension
+        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
+        xk = xk.transpose(1, 2)
+        xv = xv.transpose(1, 2)
+
+        # flash implementation
+        if self.flash:
+            output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None, dropout_p=self.dropout if self.training else 0.0, is_causal=self.causal_attention)
+        else:
+            # manual implementation
+            scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
+            if self.causal_attention:
+                scores = scores + self.mask[:, :, :seqlen, :seqlen]   # (bs, n_local_heads, seqlen, cache_len + seqlen)
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            scores = self.attn_dropout(scores)
+            output = torch.matmul(scores, xv)  # (bs, n_local_heads, seqlen, head_dim)
+
+        # restore time as batch dimension and concat heads
+        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+
+        # final projection into the residual stream
+        output = self.wo(output)
+        output = self.resid_dropout(output)
+        return output
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
+        super().__init__()
+        if hidden_dim is None:
+            hidden_dim = 4 * dim
+            hidden_dim = int(2 * hidden_dim / 3)
+            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, layer_id: int, args: TransformerConfig):
+        super().__init__()
+        self.n_heads = args.n_heads
+        self.dim = args.dim
+        self.head_dim = args.dim // args.n_heads
+        self.attention = Attention(args)
+        self.feed_forward = FeedForward(
+            dim=args.dim,
+            hidden_dim=args.hidden_dim,
+            multiple_of=args.multiple_of,
+            dropout=args.dropout,
+        )
+        self.layer_id = layer_id
+        self.attention_norm = RMSNorm_transformer(args.dim, eps=args.norm_eps)
+        self.ffn_norm = RMSNorm_transformer(args.dim, eps=args.norm_eps)
+
+    def forward(self, x, freqs_cos, freqs_sin):
+        h = x + self.attention.forward(self.attention_norm(x), freqs_cos, freqs_sin)
+        out = h + self.feed_forward.forward(self.ffn_norm(h))
+        return out
+
+class Transformer(nn.Module):
+    
+    last_loss: Optional[torch.Tensor]
+
+    def __init__(self, params: TransformerConfig):
+        super().__init__()
+        self.params = params
+        self.vocab_size = params.vocab_size
+        self.n_layers = params.n_layers
+
+        self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
+        self.dropout = nn.Dropout(params.dropout)
+        self.layers = torch.nn.ModuleList()
+        for layer_id in range(params.n_layers):
+            self.layers.append(TransformerBlock(layer_id, params))
+        self.layer_head_dim = self.layers[0].head_dim
+
+        self.norm = RMSNorm_transformer(params.dim, eps=params.norm_eps)
+        self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
+
+        # share the unembedding parameters with the embedding parameters
+        self.tok_embeddings.weight = self.output.weight # https://paperswithcode.com/method/weight-tying
+
+        # some useful precompute for the RoPE relative positional embeddings
+        # freqs_cos, freqs_sin = precompute_freqs_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len)
+        # self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        # self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+        # init all weights
+        self.apply(self._init_weights)
+        # apply special scaled init to the residual projections, per GPT-2 paper
+        for pn, p in self.named_parameters():
+            if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * params.n_layers))
+
+        # Initialize attribute for the loss of the last forward call. This will be set if the forward is called with a targets tensor.
+        self.last_loss = None
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+        _bsz, seqlen = tokens.shape
+        h = self.tok_embeddings(tokens)
+        h = self.dropout(h)
+        # freqs_cos = self.freqs_cos[:seqlen]
+        # freqs_sin = self.freqs_sin[:seqlen]
+
+        if 'position_ids' in kwargs:
+            freqs_cos, freqs_sin = compute_freqs_cis(kwargs.pop("position_ids"), self.layer_head_dim, theta=self.params.rope_base_frequency)
+        else: 
+            raise ValueError('Llama model only implemented with RoPEs')
+
+        freqs_cos = freqs_cos.squeeze()
+        freqs_sin = freqs_sin.squeeze()
+
+        for layer in self.layers:
+            h = layer(h, freqs_cos, freqs_sin)
+        h = self.norm(h)
+
+        if targets is not None:
+            logits = self.output(h)
+            self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            logits = self.output(h)
+            self.last_loss = None
+
+        return logits
+    
+class TransformerLMHeadModel(nn.Module):
+
+    def __init__(
+        self,
+        config: TransformerConfig,  
+    ) -> None:
+        
+        super().__init__()
+
+        self.config = config
+        
+        self.backbone = Transformer(config)
+
+    def forward(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0):
+        """
+        num_last_tokens: if > 0, only return the logits for the last n tokens
+        """
+
+        lm_logits = self.backbone(input_ids, position_ids=position_ids)
+    
+        CausalLMOutput = namedtuple("CausalLMOutput", ["loss", "logits"])
+        return CausalLMOutput(loss=None, logits=lm_logits)
+
+    def save_pretrained(self, save_directory):
+        """
+        Save the model and its configuration file to a directory.
+        """
+
+        # Ensure save_directory exists
+        os.makedirs(save_directory, exist_ok=True)
+
+        # Save the model's state_dict
+        model_path = os.path.join(save_directory, "pytorch_model.bin")
+        torch.save(self.state_dict(), model_path)
+
+        # Save the configuration of the model
+        config_path = os.path.join(save_directory, "config.json")
+        with open(config_path, "w") as f:
+            json.dump(self.config.to_dict(), f)

protxlstm/models/mamba.py

@@ -0,0 +1,833 @@
+# Original code from ProtMamba under Apache License 2.0.
+
+import json
+import os
+from collections import namedtuple
+from dataclasses import dataclass, field
+from functools import partial
+
+from mamba_ssm.models.config_mamba import MambaConfig
+from mamba_ssm.modules.mamba_simple import Block, Mamba
+from mamba_ssm.models.mixer_seq_simple import MixerModel, _init_weights
+from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
+from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+from transformers import PretrainedConfig
+
+from protxlstm.generation import GenerationMixinSafe
+
+@dataclass
+class MambaConfig(PretrainedConfig):
+    d_model: int = 2560
+    n_layer: int = 64
+    vocab_size: int = 50277
+    ssm_cfg: dict = field(default_factory=dict)
+    rms_norm: bool = True
+    residual_in_fp32: bool = True
+    fused_add_norm: bool = True
+    pad_vocab_size_multiple: int = 8
+    max_position_embeddings: int = 2048
+
+def create_block(
+    d_model,
+    ssm_cfg=None,
+    norm_epsilon=1e-5,
+    rms_norm=False,
+    residual_in_fp32=False,
+    fused_add_norm=False,
+    layer_idx=None,
+    device=None,
+    dtype=None,
+    checkpoint_mixer=False,
+):
+    if ssm_cfg is None:
+        ssm_cfg = {}
+    factory_kwargs = {"device": device, "dtype": dtype}
+    mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)
+    norm_cls = partial(
+        nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs
+    )
+    block = Block(
+        d_model,
+        mixer_cls,
+        norm_cls=norm_cls,
+        fused_add_norm=fused_add_norm,
+        residual_in_fp32=residual_in_fp32,
+    )
+    block.layer_idx = layer_idx
+    if checkpoint_mixer:
+        block.mixer = CheckpointedModule(block.mixer)
+    return block
+
+class CheckpointedModule(torch.nn.Module):
+    def __init__(self, layer):
+        super().__init__()
+        self.ckpt_layer = layer
+
+    def forward(self, x, *args, **kwargs):
+        return checkpoint(self.ckpt_layer, x, use_reentrant=False)
+
+    # def state_dict(self, **kwargs):
+    #     # Get the state dict of the underlying layer
+    #     layer_state_dict = self.ckpt_layer.state_dict(**kwargs)
+    #     # Create a new state dict with the original keys
+    #     state_dict = {k.replace('ckpt_layer.', ''): v for k, v in layer_state_dict.items()}
+    #     return state_dict
+
+class MixerModelSafe(MixerModel):
+    """
+    Overwrite the forward method to allow saving intermediate layers.
+    """
+
+    def forward(self, input_ids, inference_params=None, save_layer=[]):
+        hidden_states = self.embedding(input_ids)
+        residual = None
+        if len(save_layer) > 0:
+            hidden_states_dict = {}
+        for i, layer in enumerate(self.layers):
+            hidden_states, residual = layer(
+                hidden_states, residual, inference_params=inference_params
+            )
+            if i + 1 in save_layer:
+                hidden_states_dict[i + 1] = (
+                    hidden_states.detach().cpu().to(torch.float).numpy()
+                )
+        if len(save_layer) > 0:
+            return hidden_states_dict
+
+        if not self.fused_add_norm:
+            residual = (
+                (hidden_states + residual) if residual is not None else hidden_states
+            )
+            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
+        else:
+            # Set prenorm=False here since we don't need the residual
+            fused_add_norm_fn = (
+                rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
+            )
+            hidden_states = fused_add_norm_fn(
+                hidden_states,
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+            )
+        return hidden_states
+
+class MixerModelWithPosids(nn.Module):
+    r"""Mixer model for Mamba but we add positional encodings to the input embeddings."""
+
+    def __init__(
+        self,
+        d_model: int,
+        n_layer: int,
+        vocab_size: int,
+        max_position_embeddings: int,
+        ssm_cfg=None,
+        norm_epsilon: float = 1e-5,
+        rms_norm: bool = False,
+        initializer_cfg=None,
+        fused_add_norm=False,
+        residual_in_fp32=False,
+        device=None,
+        dtype=None,
+        checkpoint_mixer=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+
+        self.embedding = nn.Embedding(vocab_size, d_model // 2, **factory_kwargs)
+        self.position_embedding = nn.Embedding(
+            max_position_embeddings, d_model - d_model // 2, **factory_kwargs
+        )
+
+        # We change the order of residual and layer norm:
+        # Instead of LN -> Attn / MLP -> Add, we do:
+        # Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and
+        # the main branch (output of MLP / Mixer). The model definition is unchanged.
+        # This is for performance reason: we can fuse add + layer_norm.
+        self.fused_add_norm = fused_add_norm
+        if self.fused_add_norm:
+            if layer_norm_fn is None or rms_norm_fn is None:
+                raise ImportError("Failed to import Triton LayerNorm / RMSNorm kernels")
+
+        self.layers = nn.ModuleList(
+            [
+                create_block(
+                    d_model,
+                    ssm_cfg=ssm_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=i,
+                    checkpoint_mixer=checkpoint_mixer,
+                    **factory_kwargs,
+                )
+                for i in range(n_layer)
+            ]
+        )
+
+        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(
+            d_model, eps=norm_epsilon, **factory_kwargs
+        )
+
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_layer,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return {
+            i: layer.allocate_inference_cache(
+                batch_size, max_seqlen, dtype=dtype, **kwargs
+            )
+            for i, layer in enumerate(self.layers)
+        }
+
+    def forward(self, input_ids, position_ids, inference_params=None, save_layer=[]):
+        hidden_states = torch.cat(
+            [
+                self.embedding(input_ids),
+                self.position_embedding(position_ids),
+            ],
+            -1,
+        )
+        residual = None
+        if len(save_layer) > 0:
+            hidden_states_dict = {}
+        for i, layer in enumerate(self.layers):
+            hidden_states, residual = layer(
+                hidden_states, residual, inference_params=inference_params
+            )
+            if i + 1 in save_layer:
+                hidden_states_dict[i + 1] = (
+                    hidden_states.detach().cpu().to(torch.float).numpy()
+                )
+        if len(save_layer) > 0:
+            return hidden_states_dict
+
+        if not self.fused_add_norm:
+            residual = (
+                (hidden_states + residual) if residual is not None else hidden_states
+            )
+            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
+        else:
+            fused_add_norm_fn = (
+                rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
+            )
+            hidden_states = fused_add_norm_fn(
+                hidden_states,
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+            )
+        return hidden_states
+
+class MixerModelWith2DPosids(nn.Module):
+    r"""Mixer model for Mamba but we add positional encodings to the input embeddings."""
+
+    def __init__(
+        self,
+        d_model: int,
+        n_layer: int,
+        vocab_size: int,
+        max_position_embeddings: int,
+        max_sequence_position_embeddings: int = 512,
+        ssm_cfg=None,
+        norm_epsilon: float = 1e-5,
+        rms_norm: bool = False,
+        initializer_cfg=None,
+        fused_add_norm=False,
+        residual_in_fp32=False,
+        device=None,
+        dtype=None,
+        checkpoint_mixer=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+
+        self.embedding = nn.Embedding(
+            vocab_size, d_model - 2 * d_model // 4, **factory_kwargs
+        )
+        self.position_embedding = nn.Embedding(
+            max_position_embeddings, d_model // 4, **factory_kwargs
+        )
+        self.seq_position_embedding = nn.Embedding(
+            max_sequence_position_embeddings, d_model // 4, **factory_kwargs
+        )
+        self.d_embeddings = d_model - 2 * d_model // 4
+
+        # We change the order of residual and layer norm:
+        # Instead of LN -> Attn / MLP -> Add, we do:
+        # Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and
+        # the main branch (output of MLP / Mixer). The model definition is unchanged.
+        # This is for performance reason: we can fuse add + layer_norm.
+        self.fused_add_norm = fused_add_norm
+        if self.fused_add_norm:
+            if layer_norm_fn is None or rms_norm_fn is None:
+                raise ImportError("Failed to import Triton LayerNorm / RMSNorm kernels")
+
+        self.layers = nn.ModuleList(
+            [
+                create_block(
+                    d_model,
+                    ssm_cfg=ssm_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=i,
+                    checkpoint_mixer=checkpoint_mixer,
+                    **factory_kwargs,
+                )
+                for i in range(n_layer)
+            ]
+        )
+
+        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(
+            d_model, eps=norm_epsilon, **factory_kwargs
+        )
+
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_layer,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return {
+            i: layer.allocate_inference_cache(
+                batch_size, max_seqlen, dtype=dtype, **kwargs
+            )
+            for i, layer in enumerate(self.layers)
+        }
+
+    def forward(
+        self,
+        input_ids,
+        position_ids,
+        seq_position_ids,
+        inference_params=None,
+        save_layer=[],
+    ):
+        hidden_states = torch.cat(
+            [
+                self.embedding(input_ids),
+                self.position_embedding(position_ids),
+                self.seq_position_embedding(seq_position_ids),
+            ],
+            -1,
+        )
+        residual = None
+        if len(save_layer) > 0:
+            hidden_states_dict = {}
+        for i, layer in enumerate(self.layers):
+            hidden_states, residual = layer(
+                hidden_states, residual, inference_params=inference_params
+            )
+            if i + 1 in save_layer:
+                hidden_states_dict[i + 1] = (
+                    hidden_states.detach().cpu().to(torch.float).numpy()
+                )
+        if len(save_layer) > 0:
+            return hidden_states_dict
+
+        if not self.fused_add_norm:
+            residual = (
+                (hidden_states + residual) if residual is not None else hidden_states
+            )
+            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
+        else:
+            fused_add_norm_fn = (
+                rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
+            )
+            hidden_states = fused_add_norm_fn(
+                hidden_states,
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+            )
+        return hidden_states
+
+class MambaLMHeadModelSafe(nn.Module, GenerationMixinSafe):
+
+    def __init__(
+        self,
+        config: MambaConfig,
+        initializer_cfg=None,
+        device=None,
+        dtype=None,
+        checkpoint_mixer=False,
+    ) -> None:
+        self.config = config
+        d_model = config.d_model
+        n_layer = config.n_layer
+        vocab_size = config.vocab_size
+        ssm_cfg = config.ssm_cfg
+        rms_norm = config.rms_norm
+        residual_in_fp32 = config.residual_in_fp32
+        fused_add_norm = config.fused_add_norm
+        pad_vocab_size_multiple = config.pad_vocab_size_multiple
+        factory_kwargs = {"device": device, "dtype": dtype}
+        if checkpoint_mixer:
+            raise NotImplementedError(
+                "Checkpointing is not yet supported for MambaLMHeadModelSafe"
+            )
+
+        super().__init__()
+        if vocab_size % pad_vocab_size_multiple != 0:
+            vocab_size += pad_vocab_size_multiple - (
+                vocab_size % pad_vocab_size_multiple
+            )
+        self.backbone = MixerModelSafe(
+            d_model=d_model,
+            n_layer=n_layer,
+            vocab_size=vocab_size,
+            ssm_cfg=ssm_cfg,
+            rms_norm=rms_norm,
+            initializer_cfg=initializer_cfg,
+            fused_add_norm=fused_add_norm,
+            residual_in_fp32=residual_in_fp32,
+            **factory_kwargs,
+        )
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)
+
+        # Initialize weights and apply final processing
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_layer,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+        self.tie_weights()
+
+    def tie_weights(self):
+        self.lm_head.weight = self.backbone.embedding.weight
+
+    def clip_grad_norm_(self, max_norm, norm_type=2.0):
+        r"""Clip the norm of the gradients for the model.
+        Args:
+            max_norm (float or int): The maximum norm of the gradients.
+                The gradients are modified in-place.
+            norm_type (float or int): The type of the used p-norm. Can be 'inf' for infinity norm.
+        Returns:
+            Total norm of the parameters (viewed as a single vector).
+        """
+        return torch.nn.utils.clip_grad_value_(self.parameters(), max_norm)
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return self.backbone.allocate_inference_cache(
+            batch_size, max_seqlen, dtype=dtype, **kwargs
+        )
+
+    def forward(
+        self,
+        input_ids,
+        position_ids=None,
+        inference_params=None,
+        num_last_tokens=0,
+        save_layer=[],
+        *args,
+        **kwargs,
+    ):
+        """
+        "position_ids" is just to be compatible with Transformer generation. We don't use it.
+        num_last_tokens: if > 0, only return the logits for the last n tokens
+        """
+        return self.protected_forward(
+            input_ids, position_ids, inference_params, num_last_tokens, save_layer
+        )
+
+    def protected_forward(
+        self,
+        input_ids,
+        position_ids=None,
+        inference_params=None,
+        num_last_tokens=0,
+        save_layer=[],
+    ):
+        hidden_states = self.backbone(
+            input_ids, inference_params=inference_params, save_layer=save_layer
+        )
+        if len(save_layer) > 0:
+            return hidden_states
+        if num_last_tokens > 0:
+            hidden_states = hidden_states[:, -num_last_tokens:]
+        lm_logits = self.lm_head(hidden_states)
+        CausalLMOutput = namedtuple("CausalLMOutput", ["loss", "logits"])
+        return CausalLMOutput(loss=None, logits=lm_logits)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name, device=None, dtype=None, **kwargs):
+        config_data = load_config_hf(pretrained_model_name)
+        config = MambaConfig(**config_data)
+        model = cls(config, device=device, dtype=dtype, **kwargs)
+        model.load_state_dict(
+            load_state_dict_hf(pretrained_model_name, device=device, dtype=dtype),
+            strict=False,
+        )
+        return model
+
+    def save_pretrained(self, save_directory):
+        """
+        Minimal implementation of save_pretrained for MambaLMHeadModel.
+        Save the model and its configuration file to a directory.
+        """
+        # Ensure save_directory exists
+        os.makedirs(save_directory, exist_ok=True)
+
+        # Save the model's state_dict
+        model_path = os.path.join(save_directory, "pytorch_model.bin")
+        torch.save(self.state_dict(), model_path)
+
+        # Save the configuration of the model
+        config_path = os.path.join(save_directory, "config.json")
+        with open(config_path, "w") as f:
+            json.dump(self.config.__dict__, f)
+
+class MambaLMHeadModelwithPosids(nn.Module, GenerationMixinSafe):
+
+    def __init__(
+        self,
+        config: MambaConfig,
+        initializer_cfg=None,
+        device=None,
+        dtype=None,
+        checkpoint_mixer=False,
+    ) -> None:
+        self.config = config
+        d_model = config.d_model
+        n_layer = config.n_layer
+        vocab_size = config.vocab_size
+        max_position_embeddings = config.max_position_embeddings
+        ssm_cfg = config.ssm_cfg
+        rms_norm = config.rms_norm
+        residual_in_fp32 = config.residual_in_fp32
+        fused_add_norm = config.fused_add_norm
+        pad_vocab_size_multiple = config.pad_vocab_size_multiple
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        super().__init__()
+        if vocab_size % pad_vocab_size_multiple != 0:
+            vocab_size += pad_vocab_size_multiple - (
+                vocab_size % pad_vocab_size_multiple
+            )
+        self.backbone = MixerModelWithPosids(
+            d_model=d_model,
+            n_layer=n_layer,
+            vocab_size=vocab_size,
+            max_position_embeddings=max_position_embeddings,
+            ssm_cfg=ssm_cfg,
+            rms_norm=rms_norm,
+            initializer_cfg=initializer_cfg,
+            fused_add_norm=fused_add_norm,
+            residual_in_fp32=residual_in_fp32,
+            checkpoint_mixer=checkpoint_mixer,
+            **factory_kwargs,
+        )
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)
+
+        # Initialize weights and apply final processing
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_layer,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+        self.tie_weights()
+
+    def tie_weights(self):
+        self.lm_head.weight = self.backbone.embedding.weight
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return self.backbone.allocate_inference_cache(
+            batch_size, max_seqlen, dtype=dtype, **kwargs
+        )
+
+    def forward(
+        self,
+        input_ids,
+        position_ids=None,
+        inference_params=None,
+        num_last_tokens=0,
+        save_layer=[],
+        *args,
+        **kwargs,
+    ):
+        """
+        "position_ids" is just to be compatible with Transformer generation. We don't use it.
+        num_last_tokens: if > 0, only return the logits for the last n tokens
+        """
+        return self.protected_forward(
+            input_ids, position_ids, inference_params, num_last_tokens, save_layer
+        )
+
+    def protected_forward(
+        self,
+        input_ids,
+        position_ids=None,
+        inference_params=None,
+        num_last_tokens=0,
+        save_layer=[],
+    ):
+        hidden_states = self.backbone(
+            input_ids,
+            position_ids=position_ids,
+            inference_params=inference_params,
+            save_layer=save_layer,
+        )
+        if len(save_layer) > 0:
+            return hidden_states
+        hidden_states = hidden_states[:, :, : self.config.d_model // 2]
+        if num_last_tokens > 0:
+            hidden_states = hidden_states[:, -num_last_tokens:]
+        lm_logits = self.lm_head(hidden_states)
+        CausalLMOutput = namedtuple("CausalLMOutput", ["loss", "logits"])
+        return CausalLMOutput(loss=None, logits=lm_logits)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name,
+        device=None,
+        dtype=None,
+        checkpoint_mixer=False,
+        **kwargs,
+    ):
+        config_data = load_config_hf(pretrained_model_name)
+        config = MambaConfig(**config_data)
+        model = cls(
+            config,
+            device=device,
+            dtype=dtype,
+            checkpoint_mixer=checkpoint_mixer,
+            **kwargs,
+        )
+        state_dict = load_state_dict_hf(
+            pretrained_model_name, device=device, dtype=dtype
+        )
+        if state_dict.keys() != model.state_dict().keys():
+            if checkpoint_mixer:
+                for key in model.state_dict().keys():
+                    if "ckpt_layer" in key:
+                        state_dict[key] = state_dict.pop(key.replace("ckpt_layer.", ""))
+                print(
+                    "Using a model that was pretrained without gradient checkpointing and now want to use it. Changed the keys of the state_dict to match the model's keys."
+                )
+            else:
+                for key in list(state_dict.keys()):
+                    if "ckpt_layer" in key:
+                        state_dict[key.replace("ckpt_layer.", "")] = state_dict.pop(key)
+                print(
+                    "Using a model that was pretrained with gradient checkpointing but now do not want to use it. Changed the keys of the state_dict to match the model's keys."
+                )
+            assert (
+                state_dict.keys() == model.state_dict().keys()
+            ), "The keys of the state_dict do not match the model's keys."
+        model.load_state_dict(state_dict)
+        return model
+
+    def save_pretrained(self, save_directory):
+        """
+        Minimal implementation of save_pretrained for MambaLMHeadModel.
+        Save the model and its configuration file to a directory.
+        """
+        # Ensure save_directory exists
+        os.makedirs(save_directory, exist_ok=True)
+
+        # Save the model's state_dict
+        model_path = os.path.join(save_directory, "pytorch_model.bin")
+        torch.save(self.state_dict(), model_path)
+
+        # Save the configuration of the model
+        config_path = os.path.join(save_directory, "config.json")
+        with open(config_path, "w") as f:
+            json.dump(self.config.__dict__, f)
+
+class MambaLMHeadModelwith2DPosids(nn.Module, GenerationMixinSafe):
+
+    def __init__(
+        self,
+        config: MambaConfig,
+        initializer_cfg=None,
+        device=None,
+        dtype=None,
+        checkpoint_mixer=False,
+    ) -> None:
+        self.config = config
+        d_model = config.d_model
+        n_layer = config.n_layer
+        vocab_size = config.vocab_size
+        max_position_embeddings = config.max_position_embeddings
+        ssm_cfg = config.ssm_cfg
+        rms_norm = config.rms_norm
+        residual_in_fp32 = config.residual_in_fp32
+        fused_add_norm = config.fused_add_norm
+        pad_vocab_size_multiple = config.pad_vocab_size_multiple
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        super().__init__()
+        if vocab_size % pad_vocab_size_multiple != 0:
+            vocab_size += pad_vocab_size_multiple - (
+                vocab_size % pad_vocab_size_multiple
+            )
+        self.backbone = MixerModelWith2DPosids(
+            d_model=d_model,
+            n_layer=n_layer,
+            vocab_size=vocab_size,
+            max_position_embeddings=max_position_embeddings,
+            ssm_cfg=ssm_cfg,
+            rms_norm=rms_norm,
+            initializer_cfg=initializer_cfg,
+            fused_add_norm=fused_add_norm,
+            residual_in_fp32=residual_in_fp32,
+            checkpoint_mixer=checkpoint_mixer,
+            **factory_kwargs,
+        )
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)
+
+        # Initialize weights and apply final processing
+        self.apply(
+            partial(
+                _init_weights,
+                n_layer=n_layer,
+                **(initializer_cfg if initializer_cfg is not None else {}),
+            )
+        )
+        self.tie_weights()
+
+    def tie_weights(self):
+        self.lm_head.weight = self.backbone.embedding.weight
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        return self.backbone.allocate_inference_cache(
+            batch_size, max_seqlen, dtype=dtype, **kwargs
+        )
+
+    def forward(
+        self,
+        input_ids,
+        position_ids=None,
+        seq_position_ids=None,
+        inference_params=None,
+        num_last_tokens=0,
+        save_layer=[],
+        *args,
+        **kwargs,
+    ):
+        """
+        "position_ids" is just to be compatible with Transformer generation. We don't use it.
+        num_last_tokens: if > 0, only return the logits for the last n tokens
+        """
+        return self.protected_forward(
+            input_ids,
+            position_ids,
+            seq_position_ids,
+            inference_params,
+            num_last_tokens,
+            save_layer,
+        )
+
+    def protected_forward(
+        self,
+        input_ids,
+        position_ids=None,
+        seq_position_ids=None,
+        inference_params=None,
+        num_last_tokens=0,
+        save_layer=[],
+    ):
+        hidden_states = self.backbone(
+            input_ids,
+            position_ids=position_ids,
+            seq_position_ids=seq_position_ids,
+            inference_params=inference_params,
+            save_layer=save_layer,
+        )
+        if len(save_layer) > 0:
+            return hidden_states
+        hidden_states = hidden_states[:, :, : self.backbone.d_embeddings]
+        if num_last_tokens > 0:
+            hidden_states = hidden_states[:, -num_last_tokens:]
+        lm_logits = self.lm_head(hidden_states)
+        CausalLMOutput = namedtuple("CausalLMOutput", ["loss", "logits"])
+        return CausalLMOutput(loss=None, logits=lm_logits)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name,
+        device=None,
+        dtype=None,
+        checkpoint_mixer=False,
+        **kwargs,
+    ):
+        config_data = load_config_hf(pretrained_model_name)
+        config = MambaConfig(**config_data)
+        model = cls(
+            config,
+            device=device,
+            dtype=dtype,
+            checkpoint_mixer=checkpoint_mixer,
+            **kwargs,
+        )
+        state_dict = load_state_dict_hf(
+            pretrained_model_name, device=device, dtype=dtype
+        )
+        if state_dict.keys() != model.state_dict().keys():
+            if checkpoint_mixer:
+                for key in model.state_dict().keys():
+                    if "ckpt_layer" in key:
+                        state_dict[key] = state_dict.pop(key.replace("ckpt_layer.", ""))
+                print(
+                    "Using a model that was pretrained without gradient checkpointing and now want to use it. Changed the keys of the state_dict to match the model's keys."
+                )
+            else:
+                for key in list(state_dict.keys()):
+                    if "ckpt_layer" in key:
+                        state_dict[key.replace("ckpt_layer.", "")] = state_dict.pop(key)
+                print(
+                    "Using a model that was pretrained with gradient checkpointing but now do not want to use it. Changed the keys of the state_dict to match the model's keys."
+                )
+            assert (
+                state_dict.keys() == model.state_dict().keys()
+            ), "The keys of the state_dict do not match the model's keys."
+        model.load_state_dict(state_dict)
+        return model
+
+    def save_pretrained(self, save_directory):
+        """
+        Minimal implementation of save_pretrained for MambaLMHeadModel.
+        Save the model and its configuration file to a directory.
+        """
+        # Ensure save_directory exists
+        os.makedirs(save_directory, exist_ok=True)
+
+        # Save the model's state_dict
+        model_path = os.path.join(save_directory, "pytorch_model.bin")
+        torch.save(self.state_dict(), model_path)
+
+        # Save the configuration of the model
+        config_path = os.path.join(save_directory, "config.json")
+        with open(config_path, "w") as f:
+            json.dump(self.config.__dict__, f)

protxlstm/models/xlstm.py

@@ -0,0 +1,180 @@
+__all__ = [
+    "xLSTMConfig", 
+    "xLSTMLMHeadModel",
+]
+
+import json
+import os
+from collections import namedtuple
+from dataclasses import asdict
+
+import torch
+import torch.nn as nn
+from dacite import Config as DaciteConfig, from_dict
+from omegaconf import OmegaConf
+from transformers import PretrainedConfig
+
+from protxlstm.generation import GenerationMixinSafe
+from protxlstm.utils import load_config_hf, load_state_dict_hf
+from protxlstm.xlstm.xlstm_lm_model import xLSTMLMModel, xLSTMLMModelConfig
+
+
+class xLSTMConfig(PretrainedConfig):
+
+    def __init__(self):
+        self.config_dataclass = xLSTMLMModelConfig()
+
+    def init_from_dict(self, config: dict):
+        config = OmegaConf.create(config)
+        self.config_dataclass = from_dict(
+            data_class=xLSTMLMModelConfig,
+            data=OmegaConf.to_container(config),
+            config=DaciteConfig(strict=True),
+        )
+        return self
+
+    def to_dict(self):
+        return asdict(self.config_dataclass)
+    
+
+class xLSTMLMHeadModel(nn.Module, GenerationMixinSafe):
+
+    def __init__(self, config: xLSTMConfig) -> None:
+        super().__init__()
+
+        self.config = config
+        self.backbone = xLSTMLMModel(self.config.config_dataclass)
+        self.backbone.reset_parameters()
+
+        self.setup()
+
+
+    def setup(self):
+        
+        if 'LOCAL_RANK' in os.environ:
+            current_device = int(os.environ['LOCAL_RANK'])
+        else:
+            if 'SLURM_LOCALID' in os.environ:
+                current_device = int(os.environ['SLURM_LOCALID'])
+            else:
+                current_device = 0
+
+        #torch.cuda.set_device(f'cuda:{current_device}')
+
+        #self.backbone = self.backbone.to("cuda")
+
+
+    def forward(
+        self,
+        input_ids,
+        state=None,
+        position_ids=None,
+        seq_position_ids=None,
+        inference_params=None,
+        num_last_tokens=0,
+        save_layer=[],
+        **kwargs,
+    ):
+
+        if self.config.config_dataclass.mlstm_block.mlstm.return_last_state:
+            lm_logits, state = self.backbone(input_ids, position_ids=position_ids, seq_position_ids=seq_position_ids, state=state)
+            CausalLMOutput = namedtuple("CausalLMOutput", ["loss", "logits", "state"])
+            return CausalLMOutput(loss=None, logits=lm_logits, state=state)
+        else: 
+            lm_logits = self.backbone(input_ids, position_ids=position_ids, seq_position_ids=seq_position_ids, state=state)
+            CausalLMOutput = namedtuple("CausalLMOutput", ["loss", "logits"])
+            return CausalLMOutput(loss=None, logits=lm_logits)
+    
+    def step(
+        self,
+        input_ids,
+        state=None,
+        position_ids=None,
+        seq_position_ids=None,
+        inference_params=None,
+        num_last_tokens=0,
+        save_layer=[],
+        **kwargs,
+    ):
+
+        lm_logits, state = self.backbone.step(
+            input_ids, state=state, position_ids=position_ids, seq_position_ids=seq_position_ids
+        )
+
+        return lm_logits, state
+    
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name,
+        device=None,
+        dtype=None,
+        mlstm_backend=None,
+        mlstm_chunksize=None,
+        checkpoint_blocks=None,
+        rope_base_frequency=None,
+        mlstm_return_last_state=None,
+    ):
+        # Load the checkpoint config
+        config_dict = load_config_hf(pretrained_model_name)
+
+        # update rope base frequency
+        if rope_base_frequency is not None and config_dict.get("rope_base_frequency", None) != rope_base_frequency:
+            config_dict["rope_base_frequency"] = rope_base_frequency
+        
+        # update mlstm backend
+        if mlstm_backend is not None and config_dict["mlstm_block"]["mlstm"].get("backend", None) != mlstm_backend:
+            assert mlstm_backend in ["chunkwise", "chunkwise_variable", "parallel"], "invalid mlstm backend."
+            config_dict["mlstm_block"]["mlstm"]["backend"] = mlstm_backend
+
+        # update mlstm chunksize
+        if mlstm_chunksize is not None and config_dict["mlstm_block"]["mlstm"].get("chunk_size", None) != mlstm_chunksize:
+            config_dict["mlstm_block"]["mlstm"]["chunk_size"] = mlstm_chunksize
+
+        # update activation checkpointing
+        if checkpoint_blocks is not None:
+            config_dict["checkpoint_blocks"] = checkpoint_blocks
+
+        if mlstm_return_last_state is not None:
+            config_dict["mlstm_block"]["mlstm"]["return_last_state"] = mlstm_return_last_state
+            
+        if "slstm_block" in config_dict:
+            config_dict.pop("slstm_block")
+
+        if "slstm_at" in config_dict:
+            config_dict.pop("slstm_at")
+
+        config = xLSTMConfig().init_from_dict(config_dict)
+        
+        model = cls(config)
+
+        state_dict = load_state_dict_hf(
+            pretrained_model_name, device=device, dtype=dtype
+        )
+        assert (
+            state_dict.keys() == model.state_dict().keys()
+        ), "The keys of the state_dict do not match the model's keys."
+
+        model.load_state_dict(state_dict)
+        
+        return model
+
+    def save_pretrained(self, save_directory):
+        """
+        Save the model and its configuration file to a directory.
+        """
+
+        # Ensure save_directory exists
+        os.makedirs(save_directory, exist_ok=True)
+
+        # Save the model's state_dict
+        model_path = os.path.join(save_directory, "pytorch_model.bin")
+        torch.save(self.state_dict(), model_path)
+
+        # Save the configuration of the model
+        config_path = os.path.join(save_directory, "config.json")
+        with open(config_path, "w") as f:
+            json.dump(self.config.to_dict(), f)
+
+

protxlstm/plot_utils.py

@@ -0,0 +1,26 @@
+
+cd = { # use dependent on model-type!!
+    "xLSTM": "#3073AD",
+    "Transformers": "#4B9D7A",
+    "Mamba": "#DF8953",
+    "S4": "#D275AB",
+    "Hyena": "#E86A61",
+}
+
+def setup_matplotlib():
+  import matplotlib.pyplot as plt
+  from tueplots import bundles, axes
+  bundles.icml2022()
+  plt.rcParams.update(bundles.icml2022())
+  plt.rcParams.update(axes.lines(base_width=0.5))
+  plt.rcParams["text.usetex"] = False
+  plt.rcParams['font.family'] = "sans-serif"
+  plt.rcParams['font.serif'] = 'Arial'
+  plt.rcParams['legend.edgecolor'] = 'grey'
+  plt.rcParams['legend.framealpha'] = 0.7
+  plt.rcParams['lines.linewidth'] = 1.2
+  plt.rcParams['axes.grid'] = True
+  plt.rcParams['axes.grid.axis'] = 'both'
+  plt.rcParams['grid.alpha'] = 0.2
+  plt.rcParams['axes.grid'] = True
+  plt.rcParams['axes.prop_cycle'] = plt.cycler(color=cd.values())

protxlstm/train.py

@@ -0,0 +1,338 @@
+# Original code from ProtMamba under Apache License 2.0.
+#
+# Modifications made by Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Extended to training of xlstm and transformer-based models
+#   - Predefined splits instead of on-the-fly creation
+#   - Option to overwrite config parameters from the command line
+#   - wandb logging
+
+import argparse
+import os
+
+import torch
+from omegaconf import OmegaConf
+from transformers import TrainingArguments
+
+from protxlstm.dataloaders import ProteinMemmapDataset, ProteinDataCollator
+from protxlstm.models.xlstm import xLSTMConfig, xLSTMLMHeadModel
+from protxlstm.models.llama import TransformerConfig, TransformerLMHeadModel
+from protxlstm.trainer import ProtTrainer, EarlyStoppingCallback, get_last_checkpoint
+from protxlstm.utils import (
+    AA_TO_ID,
+    compute_metrics,
+    is_zero_rank,
+    parse_override_args,
+    print_number_of_parameters,
+    print_zero_rank,
+    set_optimizer_and_scheduler,
+    setup_wandb,
+    load_model,
+)
+
+def run(config):
+    """
+    Run training loop.
+   
+    Args:
+       config (dict): dictionary with the configuration parameters.     
+    """ 
+
+    if config.model_type == 'llama':
+        pe_kwargs = {
+            'max_position_embeddings' : config["model"]["max_position_embeddings"],
+            'add_position_ids' : '1d',
+        }
+    elif config.model_type == 'mamba':
+        from protxlstm.models.mamba import MambaConfig, MambaLMHeadModelSafe, MambaLMHeadModelwithPosids, MambaLMHeadModelwith2DPosids
+        pe_kwargs = {
+            'max_position_embeddings' : config["model"]["max_position_embeddings"],
+            'max_seq_position_embeddings' : config["model"]["max_seq_position_embeddings"],
+            'add_position_ids' : config["model"]["add_position_ids"]
+        }
+    else:
+        position_embeddings = config["model"]["position_embeddings"]
+        assert position_embeddings in ["none", "abs_1d", "abs_2d", "rot_1d", "rot_2d"]
+        if position_embeddings != "none":
+            position_embeddings = position_embeddings.split("_")[-1]
+        pe_kwargs = {
+            'max_position_embeddings' : config["model"]["max_position_embeddings"],
+            'max_seq_position_embeddings' : config["model"]["max_seq_position_embeddings"],
+            'add_position_ids' : position_embeddings
+        }
+
+    # Setup WandB
+    wandb_run_name = setup_wandb(config)
+
+    # Load datasets
+    dataset_params = {
+        "msa_memmap_path": config["msa_memmap_path"],
+        "msa_memmap_meta_path": config["msa_memmap_meta_path"],
+        "sample": config["sample_sequences"],
+        "max_msa_len": config["max_msa_len"],
+        "reverse": False,
+        "seed": config["seed_sequence_sampling"],
+        "troubleshoot": False,
+        "fim_strategy": config["fim_strategy"],
+        "always_mask": config["always_mask"],
+        **pe_kwargs,
+    }
+    train_dataset = ProteinMemmapDataset(subset_path=config["train_set"], **dataset_params)
+    valid_dataset = ProteinMemmapDataset(subset_path=config["valid_set"], **dataset_params)
+    train_eval_dataset = ProteinMemmapDataset(subset_path=config["train_eval_set"], **dataset_params)
+
+    print(f'Train set size: {len(train_dataset)} Train eval set size: {len(train_eval_dataset)} Valid set size: {len(valid_dataset)}')
+    
+    assert (
+        len(AA_TO_ID) == config["model"]["vocab_size"]
+    ), f"Vocab size in the config file does not match the one in the code. I should be {len(AA_TO_ID)}"
+    
+    # Create data collator for batched training
+    data_collator = ProteinDataCollator(max_sequence_length=config["max_msa_len"])
+ 
+    # Check datatypes
+    if config["dtype"] == "float32":
+        dtype = torch.float32
+    elif config["dtype"] == "bfloat16":
+        dtype = torch.bfloat16
+    else:
+        raise ValueError("dtype must be either float32 or bfloat16")
+    
+    # Initialize model
+    if config.model_type == 'xlstm':
+        
+        # Load model for finetuning
+        if config.finetune_model_path:
+            # These fields are updated in the config loaded from the checkpoint
+            config_update_kwargs = {
+                "mlstm_backend": config["model"]["mlstm_block"]["mlstm"]["backend"],
+                "mlstm_chunksize": config["model"]["mlstm_block"]["mlstm"]["chunk_size"],
+                "checkpoint_blocks": config["model"]["checkpoint_blocks"],
+                "rope_base_frequency": config["model"]["rope_base_frequency"]
+            }      
+            model = load_model(
+                config.finetune_model_path,
+                model_class=xLSTMLMHeadModel,
+                device="cuda",
+                dtype=dtype,
+                **config_update_kwargs
+            )     
+        else:
+            # Create new mode
+            xlstm_config = xLSTMConfig().init_from_dict(config["model"])
+            model = xLSTMLMHeadModel(xlstm_config)                  
+
+    elif config.model_type == 'mamba':
+
+        _mamba_model = {
+            "none": MambaLMHeadModelSafe, 
+            "1d": MambaLMHeadModelwithPosids, 
+            "2d": MambaLMHeadModelwith2DPosids, 
+            }
+        Mamba = _mamba_model[config['model']["add_position_ids"]]
+
+        # Load model for finetuning
+        if config.finetune_model_path:     
+            model = load_model(
+                config.finetune_model_path,
+                model_class=Mamba,
+                device="cuda",
+                dtype=dtype,
+                checkpoint_mixer=config["checkpoint_mixer"],
+            )     
+        else:
+            # Create new mode
+            mamba_config = MambaConfig(d_model=config['model']["d_model"],
+                                    n_layer=config['model']["n_layer"],
+                                    vocab_size=config['model']["vocab_size"],
+                                    residual_in_fp32=config['model']["residual_in_fp32"])
+            model = Mamba(mamba_config, dtype=dtype, checkpoint_mixer=config['model']["checkpoint_mixer"])
+        
+    elif config.model_type == 'llama':
+        
+        llama_config = TransformerConfig(
+            d_model=config["model"]["d_model"],
+            n_layer=config["model"]["n_layer"],
+            n_heads=config["model"]["n_heads"],
+            n_kv_heads=config["model"]["n_kv_heads"],
+            bidirectional=config["model"]["bidirectional"],
+            hidden_dim=config["model"]["hidden_dim"],
+            multiple_of=config["model"]["multiple_of"],
+            norm_eps=config["model"]["norm_eps"],
+            max_length=config["model"]["max_length"],
+            vocab_size=config["model"]["vocab_size"],
+            dropout=config["model"]["dropout"],
+            max_position_embeddings=config["model"]["max_position_embeddings"],
+            rope_base_frequency=config["model"]["rope_base_frequency"],
+
+        )
+
+        model = TransformerLMHeadModel(llama_config)
+
+    else:
+        raise ValueError(f"Unsupported model_type: {config.model_type}. Expected 'xlstm', 'mamba', or 'llama'.")
+
+    
+    # TODO: Improve what we want print
+    if is_zero_rank():
+        print_number_of_parameters(model)
+    print_zero_rank(f"dtype: {config['dtype']}")
+    print_zero_rank(f"Epochs: {config['num_epochs']}")
+    print_zero_rank(f"Batch size per GPU: {config['batch_size']}")
+    print_zero_rank(f"Gradient accumulation steps: {config['gradient_accumulation_steps']}")
+    eff_batch_size = config["batch_size"] * config["gradient_accumulation_steps"]
+    nr_gpus = torch.cuda.device_count()
+    print_zero_rank(f"GPUS: {nr_gpus}")
+    eff_batch_size *= nr_gpus
+    print_zero_rank(f"Effective batch size: {eff_batch_size}")
+    print_zero_rank(
+        f"Steps per training epoch: {len(train_dataset) // config['batch_size']}, eff. steps: {len(train_dataset) // eff_batch_size}"
+    )
+    print_zero_rank(f"Steps per evaluation epoch: {len(valid_dataset) // config['batch_size']}")
+    print_zero_rank(f"Max MSA length: {config['max_msa_len']}")
+    ev_epochs = round(
+        config["eval_steps"] * config["batch_size"] / len(train_dataset), 3
+    )
+    print_zero_rank(
+        f"Evaluation every {config['eval_steps']} steps, i.e. {ev_epochs} epochs. Effectively every {config['eval_steps']*config['gradient_accumulation_steps']} steps, i.e. {ev_epochs*config['gradient_accumulation_steps']} epochs."
+    )
+    if config.model_type == 'xlstm' and config["model"]["checkpoint_blocks"]:
+        print_zero_rank("Using gradient checkpointing")
+    if config["compute_only_fim_loss"]:
+        print_zero_rank("Computing only FIM loss for training")
+
+    # Training callbacks
+    es_callback = EarlyStoppingCallback(
+        train_path=config["output_dir"] + '/' + wandb_run_name, config=config
+    )
+    callbacks = [es_callback]    
+
+    # Optimizer and Schedulers
+    optimizer, scheduler = set_optimizer_and_scheduler(
+        config, 
+        len(train_dataset), 
+        model.parameters()
+    )
+
+    # Find checkpoint if available
+    last_checkpoint = None
+    if config.finetune_model_path is None:
+        path = os.path.join(config["output_dir"], wandb_run_name)
+        if os.path.exists(path):
+            last_checkpoint = get_last_checkpoint(path)
+            if last_checkpoint is None:
+                print_zero_rank("No checkpoint found, starting training from scratch.")
+            else:
+                print_zero_rank(f"Resuming training from the last checkpoint: {last_checkpoint}")
+
+    # Create trainer
+    trainer = ProtTrainer(
+        model=model,
+        train_dataset=train_dataset,
+        eval_dataset={"valid": valid_dataset, "train": train_eval_dataset},
+        optimizers=(optimizer, scheduler),
+        args=TrainingArguments(
+            run_name=wandb_run_name,
+            local_rank=int(os.getenv('LOCAL_RANK', '0')),
+            learning_rate=config["learning_rate"],
+            num_train_epochs=config["num_epochs"],
+            per_device_train_batch_size=config["batch_size"],
+            per_device_eval_batch_size=config["batch_size"],
+            gradient_accumulation_steps=config["gradient_accumulation_steps"],
+            eval_accumulation_steps=config["eval_accumulation_steps"],
+            eval_strategy="steps",
+            max_grad_norm=config["max_grad_norm"],
+            bf16=config["dtype"] == "bfloat16",
+            dataloader_num_workers=32,
+            logging_steps=config["logging_steps"],
+            eval_steps=config["eval_steps"],
+            save_steps=config["save_steps"],
+            output_dir=config["output_dir"] + '/' + wandb_run_name,
+            logging_dir=config["output_dir"] + '/' + wandb_run_name,
+            report_to="wandb" if is_zero_rank() else None,
+            log_on_each_node=False,
+            overwrite_output_dir=False,
+            push_to_hub=False,
+            label_names=["labels"],
+        ),
+        compute_only_fim_loss=config["compute_only_fim_loss"],
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+        callbacks=callbacks,
+    )    
+
+    # Train model
+    while True:
+        if last_checkpoint is None and trainer.state.global_step == 0:
+            eval_results = trainer.evaluate()
+            print_zero_rank(
+                f">>> Initial validation perplexity: {eval_results['eval_valid_perplexity/batch']:.2f}"
+            )
+        else:
+            print_zero_rank(f"Resuming training from the last checkpoint: {last_checkpoint}")
+        # Train
+        trainer.train(resume_from_checkpoint=last_checkpoint)
+        
+        # Break training when the number of epochs is reached
+        if (
+            not es_callback.should_restart
+            or trainer.state.epoch >= config["num_epochs"]
+        ):
+            eval_results = trainer.evaluate()
+            print_zero_rank(
+                f">>> Final Perplexity: {eval_results['eval_valid_perplexity/batch']:.2f}"
+            )
+            break
+        # If the training was interrupted because of a loss spike, restart from the last checkpoint
+        last_checkpoint = es_callback.checkpoint_path
+
+    return trainer
+
+if __name__ == "__main__":
+
+    # Default configuration file paths
+    default_model_config = "configs/xlstm_default_config.yaml"
+    default_train_config = "configs/train_default_config.yaml"
+
+    parser = argparse.ArgumentParser(
+        description="Train or finetune a model with the provided configuration."
+    )
+    parser.add_argument(
+        "--model_config_path",
+        type=str,
+        default=default_model_config,
+        help=f"Path to the model configuration file (default: {default_model_config})"
+    )
+    parser.add_argument(
+        "--train_config_path",
+        type=str,
+        default=default_train_config,
+        help=f"Path to the training and dataset configuration file (default: {default_train_config})"
+    )
+    parser.add_argument(
+        "overrides",
+        nargs=argparse.REMAINDER,
+        help="Override configuration values using key=value format.",
+    )
+
+    args = parser.parse_args()  
+
+    # Check if the default config files exist, or raise an error
+    if not os.path.exists(args.model_config_path):
+        raise FileNotFoundError(f"Model config file not found: {args.model_config_path}")
+    if not os.path.exists(args.train_config_path):
+        raise FileNotFoundError(f"Train config file not found: {args.train_config_path}")
+
+    # Load the model and training configurations
+    model_config = OmegaConf.load(args.model_config_path)
+    train_config = OmegaConf.load(args.train_config_path)
+
+    # Merge the model and training configurations
+    config = OmegaConf.merge(model_config, train_config)
+
+    # Parse overrides
+    if args.overrides:
+        overrides = parse_override_args(args.overrides)
+        config.merge_with(OmegaConf.create(overrides))
+
+    # Run the training/finetuning process
+    run(config)

protxlstm/trainer.py

@@ -0,0 +1,123 @@
+# Original code from ProtMamba under Apache License 2.0.
+#
+# Modifications made by Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - MambaTrainer renamed to ProtTrainer
+
+import os
+import re
+
+import torch
+from transformers import Trainer, TrainerCallback
+
+from protxlstm.utils import AA_TO_ID, find_fim_indices
+
+class ProtTrainer(Trainer):
+    """
+    Base HuggingFace Trainer used for training.
+    
+    from https://github.com/havenhq/mamba-chat/blob/main/trainer/mamba_trainer.py"""
+    def __init__(self, compute_only_fim_loss, **kwargs,):
+        super().__init__(**kwargs)
+        self.compute_only_fim_loss = compute_only_fim_loss
+        
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        input_ids = inputs.pop("input_ids")
+        labels = inputs.pop("labels")
+        if "seq_position_ids" in inputs and "position_ids" in inputs:
+            position_ids = inputs.pop("position_ids")
+            seq_position_ids = inputs.pop("seq_position_ids")
+            output = model(input_ids, position_ids=position_ids, seq_position_ids=seq_position_ids)
+        elif "position_ids" in inputs:
+            position_ids = inputs.pop("position_ids")
+            output = model(input_ids, position_ids=position_ids)
+        else:
+            output = model(input_ids)
+        lm_logits = output.logits
+
+        labels = labels.to(lm_logits.device)
+        shift_logits = lm_logits[:, :-1, :].contiguous()
+        labels = labels[:, 1:].contiguous()
+
+        loss_fct = torch.nn.CrossEntropyLoss()
+        if self.compute_only_fim_loss:
+            # start and end tokens
+            is_cls_tokens = (labels == AA_TO_ID["<cls>"])
+            is_eos_tokens = (labels == AA_TO_ID["<eos>"])
+            bool_fim = find_fim_indices(is_cls_tokens, is_eos_tokens)
+            # include also the cls token
+            bool_fim = bool_fim | is_cls_tokens
+            inds = torch.where(bool_fim)
+            lm_loss = loss_fct(shift_logits[inds[0], inds[1], :], labels[bool_fim])
+        else:
+            lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
+
+        return (lm_loss, output) if return_outputs else lm_loss
+    
+    def save_model(self, output_dir, _internal_call): 
+        if int(os.getenv('LOCAL_RANK', '0')) == 0:
+            self.model.save_pretrained(output_dir)
+            
+PREFIX_CHECKPOINT_DIR = "checkpoint"
+_re_checkpoint = re.compile(r"^" + PREFIX_CHECKPOINT_DIR + r"\-(\d+)$")
+
+def get_last_checkpoint(folder, max_steps=None):
+    content = os.listdir(folder)
+    checkpoints = [
+        path
+        for path in content
+        if _re_checkpoint.search(path) is not None and os.path.isdir(os.path.join(folder, path))
+    ]
+    if len(checkpoints) == 0:
+        return
+    
+    max_steps = max_steps if max_steps is not None else float("inf")
+    # func = lambda x: int(_re_checkpoint.search(x).groups()[0])
+    def func(x):
+        num = int(_re_checkpoint.search(x).groups()[0])
+        return num if num < max_steps else -1
+    return os.path.join(folder, max(checkpoints, key=func))
+
+class EarlyStoppingCallback(TrainerCallback):
+    def __init__(self, train_path, config=None):
+        self.step_counter_reset = 0
+        self.step_counter_stop = 0
+        self.best_loss = None
+        self.train_path = train_path
+        self.patience = config["patience"]
+        self.metric_name = config["early_stopping_metric"]
+        self.checkpoint_path = None
+        self.should_restart = False
+        self.eval_steps = config["eval_steps"]
+        self.loss_increase_factor = config["loss_increase_factor"]
+    
+    def get_checkpoint_path(self, max_steps):
+        last_checkpoint = None
+        if os.path.exists(self.train_path):
+            last_checkpoint = get_last_checkpoint(self.train_path, max_steps)
+            if last_checkpoint is None:
+                print("No checkpoint found, starting training from scratch.")
+            else:
+                print(f"Max checkpoint allowed: {max_steps}, restarting from {last_checkpoint}.")
+        return last_checkpoint
+
+    def on_evaluate(self, args, state, control, model, metrics, **kwargs):
+        if self.metric_name in metrics:
+            if self.best_loss is None:
+                self.best_loss = metrics[self.metric_name]
+            elif self.best_loss*self.loss_increase_factor < metrics[self.metric_name]:
+                self.step_counter += 1
+                if self.step_counter >= self.patience:
+                    checkpoint_path = self.get_checkpoint_path(max_steps=(state.global_step-self.patience*self.eval_steps))
+                    control.should_training_stop = True
+                    self.checkpoint_path = checkpoint_path
+                    self.should_restart = True 
+            else:
+                self.step_counter = 0
+                self.best_loss = min(self.best_loss, metrics[self.metric_name])
+                self.should_restart = False
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        self.step_counter = 0
+        self.best_loss = None
+        self.should_restart = False

protxlstm/utils.py

@@ -0,0 +1,482 @@
+# Some of the objects in this file come from ProtMamba and mamba both under Apache License 2.0.
+
+import json
+import os
+
+import numpy as np
+import rich
+import torch
+from Bio import SeqIO
+from omegaconf import DictConfig, OmegaConf
+from torch.optim import AdamW
+import wandb
+from transformers import (
+    get_constant_schedule_with_warmup,
+    get_cosine_schedule_with_warmup,
+    get_cosine_with_hard_restarts_schedule_with_warmup,
+)
+from transformers.utils import WEIGHTS_NAME, CONFIG_NAME
+from transformers.utils.hub import cached_file
+
+__all__ = ['AA_TO_ID', 'MASK_TO_ID', 'ID_TO_AA', 'load_model', 'encode_sequence', 'decode_sequence', 'clean_sequence', 'tokenizer',
+           'reorder_masked_sequence', 'load_sequences_from_msa_file', 'prepare_dataset_for_fim_generation',
+           'prepare_tokens', 'prepare_target', 'print_number_of_parameters', 'find_fim_indices',
+           'compute_metrics', 'compute_metrics_with_std', 'print_config', 'print_zero_rank', 'is_zero_rank']
+
+# Constants
+AA_TO_ID = {'<cls>': 0,
+            '<pad>': 1,
+            '<eos>': 2,
+            '<unk>': 3,
+            'L': 4,
+            'A': 5,
+            'G': 6,
+            'V': 7,
+            'S': 8,
+            'E': 9,
+            'R': 10,
+            'T': 11,
+            'I': 12,
+            'D': 13,
+            'P': 14,
+            'K': 15,
+            'Q': 16,
+            'N': 17,
+            'F': 18,
+            'Y': 19,
+            'M': 20,
+            'H': 21,
+            'W': 22,
+            'C': 23,
+            'X': 24,
+            'B': 25,
+            'U': 26,
+            'Z': 27,
+            'O': 28,
+            '.': 29,
+            '-': 30,
+            '<null_1>': 31,
+            '<mask>': 32}
+
+MASK_TO_ID = {"<mask-1>": 33,
+              "<mask-2>": 34,
+              "<mask-3>": 35,
+              "<mask-4>": 36,
+              "<mask-5>": 37,}
+
+AA_TO_ID.update(MASK_TO_ID)
+
+ID_TO_AA = {v: k for k, v in AA_TO_ID.items()}
+
+# Logging & prints
+def setup_wandb(config):
+
+    # WandB setup
+    os.environ["WANDB_PROJECT"] = config["wandb_project"]
+    os.environ["WANDB_ENTITY"] = config["wandb_entity"]
+    os.environ["WANDB_MODE"] = config["wandb_mode"]
+
+    if config['model_type'] == 'xlstm':
+        pe = config['model']['add_position_ids']
+        pe  = 'None' if pe == 'none' else 'AbsPE' if pe == 'abs_1d' else 'AbsPE2' if pe == 'abs_2d' else 'RoPE' if pe == 'rot_1d' else pe == 'rot_2d'
+        wandb_run_name = f"{config['model_type']}_l{config['model']['num_blocks']}_d{config['model']['embedding_dim']}_{pe}_s{config['max_msa_len']}_lr{config['learning_rate']}"
+    elif config['model_type'] == 'mamba':
+        pe = config['model']['add_position_ids']
+        pe = 'None' if pe == 'none' else 'AbsPE' if pe == '1d' else pe == '2d'
+        wandb_run_name = f"{config['model_type']}_l{config['model']['n_layer']}_d{config['model']['d_model']}_{pe}_s{config['max_msa_len']}_lr{config['learning_rate']}"
+    elif config['model_type'] == 'llama':
+        pe = 'RoPE'
+        wandb_run_name = f"{config['model_type']}_l{config['model']['n_layer']}_d{config['model']['d_model']}_dh{config['model']['hidden_dim']}_{prepare_dataset_for_fim_generation}_s{config['max_msa_len']}_lr{config['learning_rate']}_sched-{config['scheduler']}"
+
+    if config['name_prefix']:
+        wandb_run_name = str(config['name_prefix']) + '_' + wandb_run_name
+    if config['name_suffix']:
+        wandb_run_name = wandb_run_name + '_' + str(config['name_suffix'])
+
+    if is_zero_rank():
+        wandb.init(
+            project=config["wandb_project"],
+            entity=config["wandb_entity"],
+            mode=config["wandb_mode"],
+            name=wandb_run_name)
+        config_dict = OmegaConf.to_container(config, resolve=True)
+        wandb.config.update(config_dict)
+    return wandb_run_name
+
+def is_zero_rank():
+    return int(os.getenv('LOCAL_RANK', '0')) == 0
+
+def print_zero_rank(var):
+    if is_zero_rank():
+        print(var)
+
+def print_number_of_parameters(model):
+    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    formatted_num_params = f"{num_params:_}"
+    print("Number of trainable parameters: ", formatted_num_params)
+
+# Sequence tools
+def encode_sequence(sequence):
+    """Tokenize a sequence of amino acids and add a cls token at the beginning."""
+    tokenized_sequence = [AA_TO_ID[aa] if aa in AA_TO_ID else AA_TO_ID['<unk>'] for aa in sequence]
+    return [AA_TO_ID['<cls>']] + tokenized_sequence
+
+def decode_sequence(sequence):
+    """Decode a sequence of tokens."""
+    return "".join([ID_TO_AA[token] if token in ID_TO_AA else "<unk>" for token in sequence])
+
+def clean_sequence(sequence):
+    """Remove gaps and convert all residues to upper case."""
+    return sequence.replace("-", "").upper()
+
+def tokenizer(sequence_list, concatenate=True):
+    """Tokenize a collection of sequences. If the sequences are aligned, the gaps will be removed
+    and the insertions (lower case) will be promoted to upper case."""
+    # clean and encode all sequences
+    sequence_list = [encode_sequence(clean_sequence(sequence)) for sequence in sequence_list]
+    if concatenate:
+        # concatenate all sequences
+        sequences = np.concatenate(sequence_list)
+        # convert to tensor and add batch dimension
+        return torch.asarray(sequences, dtype=torch.int8)[None,:]
+    else:
+        return [torch.asarray(sequence, dtype=torch.int8) for sequence in sequence_list]
+
+def reorder_masked_sequence(mask_seq, return_ids=False):
+    """
+    Reorder a masked sequence to fill the masked positions with the tokens
+    that should be there but are positioned after the <eos> token.
+    """
+    mask_seq = mask_seq.split("<cls>")[0]
+    try:
+        # Split the sequence and masks
+        seq, masks = mask_seq.split("<eos>")
+    except:
+        return mask_seq
+    full_seq = ""
+    ids_mask = []
+    # Iterate over each mask tag
+    for mm in ["<mask-1>", "<mask-2>", "<mask-3>", "<mask-4>", "<mask-5>","<mask-?>"]:
+        try:
+            # Split the sequence in before and after the mask tag
+            seq1, seq2 = seq.split(mm)
+            if mm=="<mask-1>":
+                # If the mask is the first one, add the sequence before the mask and update the masks
+                masks = masks.split("<mask-1>")[1]
+                full_seq += seq1
+            else:
+                # If the mask is not the first one, insert the mask between the two sequence parts
+                masks1, masks2 = masks.split(mm)
+                ids_mask += [(len(full_seq), len(full_seq)+len(masks1))]
+                full_seq += masks1 + seq1
+                # Update the masks
+                masks = masks2 
+            # Update the sequence with the part after the mask
+            seq = seq2
+        except:
+            # If the mask is not found, add the remaining sequence
+            ids_mask += [(len(full_seq), len(full_seq)+len(masks))]
+            full_seq += masks + seq
+            break
+    if return_ids:
+        return full_seq, ids_mask
+    return full_seq
+
+def load_sequences_from_msa_file(file_path):
+    """Load a collection of sequences from an a3m file."""
+    with open(file_path, "r") as f:
+        sequences = [str(record.seq) for record in SeqIO.parse(f, "fasta")]
+    return sequences
+
+def prepare_dataset_for_fim_generation(tokens, pos_ids):
+    """
+    Function to transform the tokenized training dataset into a format that can be used for FIM generation.
+    Splits the input tokens and pos_ids into the FIM part (of the last sequence) and the context part (all
+    the previous sequences and the masked part of the last sequence).
+    Also returns a dictionary with the positions of the mask tokens in the FIM part.
+    """
+    def find_mask_positions(tokens_fim):
+        """
+        Function to find the positions of the mask tokens in the FIM part of the last sequence.
+        """
+        bool_mask = None
+        inds_masks = []
+        for ind in MASK_TO_ID.values():
+            tmp_bool = tokens_fim[0].cpu().numpy() == ind
+            bool_mask = tmp_bool if bool_mask is None else bool_mask | tmp_bool
+            inds_masks += [ind]
+        return bool_mask, inds_masks
+    # find where the FIM part of the last sequence starts
+    start_last_fim = np.where(tokens[0].cpu().numpy() == AA_TO_ID["<eos>"])[0][-1]
+    start_next_seqs = np.where(tokens[0,start_last_fim+1:].cpu().numpy() == AA_TO_ID["<cls>"])[0]
+    end_last_fim = start_last_fim+ 1 +start_next_seqs[0] if len(start_next_seqs) > 0 else tokens.shape[1]
+    # split tokens and pos_ids into FIM part and context part
+    tokens_to_fim = tokens[:,:start_last_fim+1]
+    pos_ids_to_fim = pos_ids[:,:start_last_fim+1]
+    tokens_fim = tokens[:,start_last_fim+1:end_last_fim]
+    pos_ids_fim = pos_ids[:,start_last_fim+1:end_last_fim]
+    # find positions of mask tokens
+    bool_mask, inds_masks = find_mask_positions(tokens_fim)
+    masked_positions = pos_ids_fim[0,bool_mask]
+    mask_dict = {ind: int(pos) for ind, pos in zip(inds_masks, masked_positions)}
+    return tokens_to_fim, pos_ids_to_fim, tokens_fim, pos_ids_fim, mask_dict
+
+# Metrics
+def find_fim_indices(is_cls_tokens, is_eos_tokens):
+    """Function to find the indices of the FIM tokens in the sequences.
+    """
+    # add a cls token at the beginning
+    is_cls_tokens = torch.cat([torch.ones_like(is_cls_tokens[:, :1]), is_cls_tokens], dim=1)
+    is_eos_tokens = torch.cat([torch.zeros_like(is_eos_tokens[:, :1]), is_eos_tokens], dim=1)
+    # both eos and cls tokens
+    bol = is_cls_tokens | is_eos_tokens
+    tmp = torch.zeros_like(is_cls_tokens, dtype=torch.int)
+    tmp[torch.nonzero(is_cls_tokens, as_tuple=True)] = 1
+    tmp[torch.nonzero(is_eos_tokens, as_tuple=True)] = -1
+    bol1 = torch.clone(bol)
+    for batch_ind in range(tmp.size(0)):
+        tmp1 = tmp[batch_ind,bol[batch_ind]]
+        # find all positions where a 1 if preceeded by a -1
+        tmp1 = tmp1[:-1]*tmp1[1:]
+        # add the first element to make the sequence start with a 1
+        tmp1 = torch.cat([torch.ones_like(tmp1[:1]).to(tmp1.device), tmp1])
+        new_bol = tmp1<0
+        # bool array True only in the positions where a 1 is preceeded by a -1
+        bol1[batch_ind,bol[batch_ind]] = False if new_bol.size(0) == 0 else new_bol
+    cumulative_sum = torch.cumsum(bol1, dim=1)
+    # Use modulo operation to get the desired tensor
+    bol2 = cumulative_sum % 2 == 1
+    bol2[is_eos_tokens]= False
+    return bol2[:,1:]
+
+def compute_metrics(eval_pred):
+    predictions, labels = eval_pred
+    predictions = torch.tensor(predictions).permute(0, 2, 1)
+    labels = torch.tensor(labels)
+    # shift labels to align them with predictions and remove last prediction to match the length
+    predictions = predictions[:, :, :-1].contiguous()
+    labels = labels[:, 1:].contiguous()
+    # compute unreduced elementwise loss
+    unreduced_loss = torch.nn.functional.cross_entropy(predictions, labels, reduction="none")
+    # compute reconstruction accuracy
+    reconstruction = (predictions.argmax(1) == labels)
+
+    # start and end tokens
+    is_cls_tokens = (labels == AA_TO_ID["<cls>"])
+    is_eos_tokens = (labels == AA_TO_ID["<eos>"])
+    # fill in the middle tokens
+    if False:
+        fim_tokens = torch.zeros(is_cls_tokens.size(0), is_cls_tokens.size(1), dtype=torch.bool)
+        in_mask_vector = torch.zeros(is_cls_tokens.size(0), dtype=torch.bool)
+        for j in range(is_cls_tokens.size(1)):
+            in_mask_vector = in_mask_vector & ~is_cls_tokens[:, j]
+            fim_tokens[:, j] = in_mask_vector
+            in_mask_vector = in_mask_vector | is_eos_tokens[:, j]
+    fim_tokens = find_fim_indices(is_cls_tokens, is_eos_tokens)
+        
+    number_sequences = torch.cumsum(torch.cat([torch.zeros(is_cls_tokens.size(0),1, dtype=torch.int32), is_cls_tokens[:,:-1]],1), -1)
+    # fist, second and last sequence tokens
+    first_sequence_tokens = ((~fim_tokens & (labels < 33)) | fim_tokens) & (number_sequences == 0)
+    second_sequence_tokens = ((~fim_tokens & (labels < 33)) | fim_tokens) & (number_sequences == 1)
+    last_sequence_tokens = ((~fim_tokens & (labels < 33)) | fim_tokens) & (number_sequences == (number_sequences.max(1).values[:, None] - 1))
+    # end of mask tokens
+    end_of_masks = (fim_tokens & (labels > 33)) | is_cls_tokens | is_eos_tokens
+
+    return {
+        "loss/all": torch.mean(unreduced_loss).item(),
+        "loss/end_span": torch.mean(unreduced_loss[end_of_masks]).item(),
+        "perplexity/seq": torch.mean(torch.exp(torch.mean(unreduced_loss, dim=1))).item(),
+        "perplexity/end_span": torch.exp(torch.mean(unreduced_loss[end_of_masks])).item(),
+        "perplexity/batch": torch.exp(torch.mean(unreduced_loss)).item(),
+        "perplexity/first_seq": torch.exp(torch.mean(unreduced_loss[first_sequence_tokens])).item(),
+        "perplexity/second_seq": torch.exp(torch.mean(unreduced_loss[second_sequence_tokens])).item(),
+        "perplexity/last_seq": torch.exp(torch.mean(unreduced_loss[last_sequence_tokens])).item(),
+        "perplexity/fim": torch.exp(torch.mean(unreduced_loss[fim_tokens])).item(),
+        "reconstruction/all": torch.mean(reconstruction.float()).item(),
+        "reconstruction/end_span": torch.mean(reconstruction[end_of_masks].float()).item(),
+        "reconstruction/first_seq": torch.mean(reconstruction[first_sequence_tokens].float()).item(),
+        "reconstruction/second_seq": torch.mean(reconstruction[second_sequence_tokens].float()).item(),
+        "reconstruction/last_seq": torch.mean(reconstruction[last_sequence_tokens].float()).item(),
+        "reconstruction/fim": torch.mean(reconstruction[fim_tokens].float()).item(),
+        }
+
+def compute_metrics_with_std(eval_pred):
+    predictions, labels = eval_pred
+    predictions = torch.tensor(predictions).permute(0, 2, 1)
+    labels = torch.tensor(labels)
+    # shift labels to align them with predictions and remove last prediction to match the length
+    predictions = predictions[:, :, :-1].contiguous()
+    labels = labels[:, 1:].contiguous()
+    # compute unreduced elementwise loss
+    unreduced_loss = torch.nn.functional.cross_entropy(predictions, labels, reduction="none")
+    # compute reconstruction accuracy
+    reconstruction = (predictions.argmax(1) == labels)
+
+    # start and end tokens
+    is_cls_tokens = (labels == AA_TO_ID["<cls>"])
+    is_eos_tokens = (labels == AA_TO_ID["<eos>"])
+    # fill in the middle tokens
+    if False:
+        fim_tokens = torch.zeros(is_cls_tokens.size(0), is_cls_tokens.size(1), dtype=torch.bool)
+        in_mask_vector = torch.zeros(is_cls_tokens.size(0), dtype=torch.bool)
+        for j in range(is_cls_tokens.size(1)):
+            in_mask_vector = in_mask_vector & ~is_cls_tokens[:, j]
+            fim_tokens[:, j] = in_mask_vector
+            in_mask_vector = in_mask_vector | is_eos_tokens[:, j]
+    fim_tokens = find_fim_indices(is_cls_tokens, is_eos_tokens)
+        
+    number_sequences = torch.cumsum(torch.cat([torch.zeros(is_cls_tokens.size(0),1, dtype=torch.int32), is_cls_tokens[:,:-1]],1), -1)
+    # fist, second and last sequence tokens
+    first_sequence_tokens = ((~fim_tokens & (labels < 33)) | fim_tokens) & (number_sequences == 0)
+    second_sequence_tokens = ((~fim_tokens & (labels < 33)) | fim_tokens) & (number_sequences == 1)
+    last_sequence_tokens = ((~fim_tokens & (labels < 33)) | fim_tokens) & (number_sequences == (number_sequences.max(1).values[:, None] - 1))
+    # end of mask tokens
+    end_of_masks = (fim_tokens & (labels > 33)) | is_cls_tokens | is_eos_tokens
+    
+    def perplexities_per_seq_for_subset(unreduced_loss, subset):
+        return torch.exp(torch.nanmean(torch.where(subset, unreduced_loss, torch.tensor(float('nan'))), dim=1))
+    
+    return{
+        # Loss
+        "loss/all": torch.mean(unreduced_loss).item(),
+        "loss/std": torch.std(unreduced_loss).item(), 
+        "loss/end_span": torch.mean(unreduced_loss[end_of_masks]).item(),
+        "loss/end_span_std": torch.std(unreduced_loss[end_of_masks]).item(),
+
+        # Perplexity of all tokens
+        "perplexity/batch": torch.exp(torch.mean(unreduced_loss)).item(),
+        "perplexity/batch_std": torch.exp(torch.std(unreduced_loss)).item(), # Fix
+        
+        # Perplexity per sequence
+        "perplexity/seq": torch.mean(torch.exp(torch.mean(unreduced_loss, dim=1))).item(),
+        "perplexity/seq_std": torch.std(torch.exp(torch.mean(unreduced_loss, dim=1))).item(),
+        "perplexity/end_span": torch.exp(torch.mean(unreduced_loss[end_of_masks])).item(),
+        "perplexity/end_span_std": torch.std(torch.exp(unreduced_loss[end_of_masks])).item(),   
+        
+        "perplexity/first_seq": torch.mean(perplexities_per_seq_for_subset(unreduced_loss, first_sequence_tokens)).item(),
+        "perplexity/first_seq_std": torch.std(perplexities_per_seq_for_subset(unreduced_loss, first_sequence_tokens)).item(),
+        "perplexity/second_seq": torch.mean(perplexities_per_seq_for_subset(unreduced_loss, second_sequence_tokens)).item(),
+        "perplexity/second_seq_std": torch.std(perplexities_per_seq_for_subset(unreduced_loss, second_sequence_tokens)).item(),
+        "perplexity/last_seq": torch.mean(perplexities_per_seq_for_subset(unreduced_loss, last_sequence_tokens)).item(),
+        "perplexity/last_seq_std": torch.std(perplexities_per_seq_for_subset(unreduced_loss, last_sequence_tokens)).item(),
+        "perplexity/fim": torch.mean(perplexities_per_seq_for_subset(unreduced_loss, fim_tokens)).item(),
+        "perplexity/fim_std": torch.std(perplexities_per_seq_for_subset(unreduced_loss, fim_tokens)).item(),
+        "reconstruction/all": torch.mean(reconstruction.float()).item(),
+        "reconstruction/std": torch.std(reconstruction.float()).item(),
+        "reconstruction/end_span": torch.mean(reconstruction[end_of_masks].float()).item(),
+        "reconstruction/end_span_std": torch.std(reconstruction[end_of_masks].float()).item(),
+        "reconstruction/first_seq": torch.mean(reconstruction[first_sequence_tokens].float()).item(),
+        "reconstruction/first_seq_std": torch.std(reconstruction[first_sequence_tokens].float()).item(),
+        "reconstruction/second_seq": torch.mean(reconstruction[second_sequence_tokens].float()).item(),
+        "reconstruction/second_seq_std": torch.std(reconstruction[second_sequence_tokens].float()).item(),
+        "reconstruction/last_seq": torch.mean(reconstruction[last_sequence_tokens].float()).item(),
+        "reconstruction/last_seq_std": torch.std(reconstruction[last_sequence_tokens].float()).item(),
+        "reconstruction/fim": torch.mean(reconstruction[fim_tokens].float()).item(),
+        "reconstruction/fim_std": torch.std(reconstruction[fim_tokens].float()).item(),
+    }
+
+# Others
+def set_optimizer_and_scheduler(config, ntrain, parameters):
+
+    # Set optimizer
+    optimizer = AdamW(
+        parameters,
+        lr=config["learning_rate"],
+        betas=(config["beta1"], config["beta2"]),
+        weight_decay=config["weight_decay"],
+    )
+
+    eff_batch_size = config["batch_size"] * config["gradient_accumulation_steps"] * torch.cuda.device_count()
+
+    # Set scheduler
+    if config["scheduler"] == "cosine":
+        print_zero_rank("Using cosine scheduler")
+        scheduler = get_cosine_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=config["warmup_steps"],
+            num_training_steps=config["num_epochs"] * ntrain // eff_batch_size,
+        )
+    if config["scheduler"] == "cosine-restarts":
+        scheduler = get_cosine_with_hard_restarts_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=config["warmup_steps"],
+            num_training_steps=config["num_epochs"] * ntrain // eff_batch_size,
+            num_cycles=config["num_cycles"],
+        )
+    elif config["scheduler"] == "constant":
+        print_zero_rank("Using constant scheduler with warmup")
+        scheduler = get_constant_schedule_with_warmup(
+            optimizer, num_warmup_steps=config["warmup_steps"]
+        )
+    else:
+        raise ValueError("Scheduler must be either cosine or constant") 
+    
+    # Finetuning and no optimizer/scheduler reset
+    if config.finetune_model_path and not config.restart_optimizer_and_scheduler:
+        optimizer.load_state_dict(torch.load(config.finetune_model_path + "/optimizer.pt"))
+        for param_group in optimizer.param_groups:
+            param_group['initial_lr'] = config['learning_rate']
+            param_group['lr'] = config['learning_rate']
+
+        scheduler.load_state_dict(torch.load(config.finetune_model_path + "/scheduler.pt"))
+        scheduler.base_lrs = [config['learning_rate']]
+        scheduler._last_lr = [config['learning_rate']]  
+
+    return optimizer, scheduler 
+
+def parse_override_args(override_args):
+    overrides = {}
+    for arg in override_args:
+        key, value = arg.split("=")
+        keys = key.split(".")
+        sub_dict = overrides
+        for sub_key in keys[:-1]:
+            if sub_key not in sub_dict:
+                sub_dict[sub_key] = {}
+            sub_dict = sub_dict[sub_key]
+        # Convert value to appropriate type
+        if value == 'True':
+            value = True
+        elif value == 'False':
+            value = False
+        elif value == 'None':
+            value = None
+        else:
+            try:
+                value = int(value)
+            except ValueError:
+                try:
+                    value = float(value)
+                except ValueError:
+                    pass
+        sub_dict[keys[-1]] = value
+    return overrides
+
+def load_model(
+    model_path,
+    device,
+    model_class,
+    dtype=torch.bfloat16,
+    **kwargs
+):
+    model = model_class.from_pretrained(
+        model_path, device=device, dtype=dtype, **kwargs
+    )
+    return model
+
+# https://github.com/state-spaces/mamba/blob/main/mamba_ssm/utils/hf.py
+def load_config_hf(model_name):
+    resolved_archive_file = cached_file(model_name, CONFIG_NAME, _raise_exceptions_for_missing_entries=False)
+    return json.load(open(resolved_archive_file))
+
+# https://github.com/state-spaces/mamba/blob/main/mamba_ssm/utils/hf.py
+def load_state_dict_hf(model_name, device=None, dtype=None):
+    # If not fp32, then we don't want to load directly to the GPU
+    mapped_device = "cpu" if dtype not in [torch.float32, None] else device
+    resolved_archive_file = cached_file(model_name, WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False)
+    return torch.load(resolved_archive_file, map_location=mapped_device)
+    # Convert dtype before moving to GPU to save memory
+    if dtype is not None:
+        state_dict = {k: v.to(dtype=dtype) for k, v in state_dict.items()}
+    state_dict = {k: v.to(device=device) for k, v in state_dict.items()}
+    return state_dict

protxlstm/xlstm/__init__.py

@@ -0,0 +1,6 @@
+from .blocks.mlstm.block import mLSTMBlock, mLSTMBlockConfig
+from .blocks.mlstm.layer import mLSTMLayer, mLSTMLayerConfig
+from .components.feedforward import FeedForwardConfig, GatedFeedForward
+from .components.rotary_position import compute_freqs_cis, apply_rotary_emb
+from .xlstm_block_stack import xLSTMBlockStack, xLSTMBlockStackConfig
+from .xlstm_lm_model import xLSTMLMModel, xLSTMLMModelConfig

protxlstm/xlstm/blocks/mlstm/__init__.py

@@ -0,0 +1 @@
+

protxlstm/xlstm/blocks/mlstm/backends.py

@@ -0,0 +1,314 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck
+
+# Modified by Pieter-Jan Hoedt, Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Fix numerical issues between parallel and stepwise backends
+#   - Make chunkwise implementation compatible with stepwise backend and variable sequence lengths
+
+
+import math
+from typing import Union, Tuple, Optional
+import torch
+
+
+def parallel_stabilized_simple(
+    queries: torch.Tensor,
+    keys: torch.Tensor,
+    values: torch.Tensor,
+    igate_preact: torch.Tensor,
+    fgate_preact: torch.Tensor,
+    lower_triangular_matrix: torch.Tensor = None,
+    stabilize_rowwise: bool = True,
+    eps: float = 1e-6,
+    **kwargs,
+) -> torch.Tensor:
+    """This is the mLSTM cell in parallel form.
+    This version is stabilized. We control the range of exp() arguments by
+    ensuring that they are always smaller than 0.0 by subtracting the maximum.
+
+    Args:
+        queries (torch.Tensor): (B, NH, S, DH)
+        keys (torch.Tensor): (B, NH, S, DH)
+        values (torch.Tensor): (B, NH, S, DH)
+        igate_preact (torch.Tensor): (B, NH, S, 1)
+        fgate_preact (torch.Tensor): (B, NH, S, 1)
+        lower_triangular_matrix (torch.Tensor, optional): (S,S). Defaults to None.
+        stabilize_rowwise (bool, optional): Wether to stabilize the combination matrix C rowwise (take maximum per row).
+            Alternative: Subtract the maximum over all rows. Defaults to True.
+
+    Returns:
+        torch.Tensor: (B, NH, S, DH), h_tilde_state
+    """
+
+    B, NH, S, DH = queries.shape
+    _dtype, _device = queries.dtype, queries.device
+
+    # forget gate matrix
+    log_fgates = torch.nn.functional.logsigmoid(fgate_preact)  # (B, NH, S, 1)
+    if lower_triangular_matrix is None or S < lower_triangular_matrix.size(-1):
+        ltr = torch.tril(torch.ones((S, S), dtype=torch.bool, device=_device))
+    else:
+        ltr = lower_triangular_matrix
+    assert (
+        ltr.dtype == torch.bool
+    ), f"lower_triangular_matrix must be of dtype bool, got {ltr.dtype}"
+
+    log_f_mat = torch.tril(log_fgates.repeat(1, 1, 1, S), diagonal=-1)
+    log_prod_f_mat = torch.cumsum(log_f_mat, dim=-2)
+    # Causal masking & selection of the correct submatrix, such that forgetgate at timestep t is not applied
+    # to the input at timestep t
+    log_fg_matrix = torch.where(ltr, log_prod_f_mat, -float("inf"))  # (B, NH, S, S)
+
+    # gate decay matrix D (combination of forget gate and input gate)
+    log_D_matrix = log_fg_matrix + igate_preact.transpose(-2, -1)  # (B, NH, S, S)
+    # D matrix stabilization
+    if stabilize_rowwise:
+        max_log_D, _ = torch.max(log_D_matrix, dim=-1, keepdim=True)  # (B, NH, S, 1)
+    else:
+        max_log_D = torch.max(log_D_matrix.view(B, NH, -1), dim=-1, keepdim=True)[
+            0
+        ].unsqueeze(-1)
+        # (B, NH, 1, 1)
+    log_D_matrix_stabilized = log_D_matrix - max_log_D  # (B, NH, S, S)
+    D_matrix = torch.exp(log_D_matrix_stabilized)  # (B, NH, S, S)
+
+    keys_scaled = keys / math.sqrt(DH)
+
+    # combination matrix C
+    qk_matrix = queries @ keys_scaled.transpose(-2, -1)  # (B, NH, S, S)
+    C_matrix = qk_matrix * D_matrix  # (B, NH, S, S)
+    normalizer = torch.maximum(
+        C_matrix.sum(dim=-1, keepdim=True).abs(), torch.exp(-max_log_D)
+    )  # (B, NH, S, 1)
+    # (B, NH, S, S)
+    C_matrix_normalized = C_matrix / (normalizer + eps)
+
+    # retrieved values
+    h_tilde_state = C_matrix_normalized @ values  # (B, NH, S, DH)
+
+    return h_tilde_state
+
+
+def chunkwise_simple(
+    queries: torch.Tensor,
+    keys: torch.Tensor,  # B, NH, S, DH
+    values: torch.Tensor,  # B, NH, S, DH
+    igate_preact: torch.Tensor,  # B, NH, S
+    fgate_preact: torch.Tensor,  # B, NH, S
+    initial_C: Optional[torch.Tensor] = None,  # B, NH, DH, DH
+    initial_n: Optional[torch.Tensor] = None,  # B, NH, DH, 1
+    initial_m: Optional[torch.Tensor] = None,  # B, NH, 1, 1
+    chunk_size: int = 64,  # optimize this
+    return_last_state: bool = False,
+    eps: float = 1e-6,
+    **kwargs,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    B, NH, S, DH = queries.shape
+    NS, CS = S // chunk_size, chunk_size
+    _dtype, _device = queries.dtype, queries.device
+
+    # form chunks
+    q = queries.view(B, NH, NS, CS, DH)
+    k = keys.view(B, NH, NS, CS, DH) / math.sqrt(DH)
+    v = values.view(B, NH, NS, CS, DH)
+
+    # forget gates
+    log_fgates = torch.nn.functional.logsigmoid(fgate_preact).view(B, NH, NS, CS)
+    log_fgates_acc = log_fgates.cumsum(dim=3)
+    igate_preact = igate_preact.view(B, NH, NS, CS)
+
+    log_fgates_rep = log_fgates[:, :, :, :, None].repeat(1, 1, 1, 1, CS)
+    log_fg_matrix = torch.tril(log_fgates_rep, diagonal=-1)
+    log_prod_fg_matrix = torch.cumsum(log_fg_matrix, dim=3)
+    
+    loggates = (igate_preact + log_prod_fg_matrix[:, :, :, -1]).unsqueeze(-1)
+    m_loc, _ = torch.max(loggates, dim=3, keepdim=True)
+    loggates = loggates - m_loc
+
+    kv = k.transpose(-1, -2) @ (v * (loggates).exp())
+    ksum = (k * (loggates).exp()).sum(dim=-2)
+    C = torch.zeros((B, NH, NS + 1, DH, DH), device=kv.device, dtype=kv.dtype)
+    n = torch.zeros((B, NH, NS + 1, DH, 1), device=kv.device, dtype=kv.dtype)
+    if initial_C is not None:
+        C[:, :, 0] = initial_C
+    if initial_n is not None:
+        n[:, :, 0] = initial_n
+
+    m = torch.zeros((B, NH, NS + 1, 1, 1), device=kv.device, dtype=kv.dtype)
+    if initial_m is not None:
+        m[:, :, 0] = initial_m
+
+    for i in range(1, NS + 1):
+        m[:, :, i] = torch.maximum(
+            log_fgates_acc[:, :, i - 1, -1, None, None] + m[:, :, i - 1],
+            m_loc[:, :, i - 1],
+        )
+        C[:, :, i] = (
+            C[:, :, i - 1].clone()
+            * (
+                log_fgates_acc[:, :, i - 1, -1, None, None]
+                + m[:, :, i - 1]
+                - m[:, :, i]
+            ).exp()
+            + kv[:, :, i - 1] * (m_loc[:, :, i - 1] - m[:, :, i]).exp()
+        )
+        n[:, :, i] = (
+            n[:, :, i - 1].clone()
+            * (
+                log_fgates_acc[:, :, i - 1, None, -1:]
+                + m[:, :, i - 1]
+                - m[:, :, i]
+            ).exp()
+            + ksum[:, :, i - 1, :, None] * (m_loc[:, :, i - 1] - m[:, :, i]).exp()
+        )
+
+    log_fg_matrix = log_prod_fg_matrix - torch.triu(
+        torch.full([1, 1, 1, CS, CS], float("inf")).to(q), diagonal=1
+    )
+
+    # gate decay matrix D (combination of forget gate and input gate)
+    log_D_matrix = log_fg_matrix + igate_preact[:, :, :, :, None].transpose(
+        -2, -1
+    )  # (B, NH, NS, CS, CS)
+    D_max, _ = torch.max(log_D_matrix, dim=-1, keepdim=True)
+
+    stab = torch.maximum(D_max, m[:, :, :-1, :] + log_fgates_acc[:, :, :, :, None])
+    inter_C = (
+        q * (m[:, :, :-1, :] + log_fgates_acc[:, :, :, :, None] - stab).exp()
+    ) @ C[:, :, :-1]
+    inter_n = (
+        q * (m[:, :, :-1, :] + log_fgates_acc[:, :, :, :, None] - stab).exp()
+    ) @ n[:, :, :-1, :]
+
+    # D matrix stabilization
+    log_D_matrix_stabilized = log_D_matrix - stab  # (B, NH, NS, CS, CS)
+    D_matrix = torch.exp(log_D_matrix_stabilized)  # (B, NH, NS, CS, CS)
+
+    # combination matrix C
+    qk_matrix = q @ k.transpose(-2, -1)  # (B, NH, NS, CS, CS)
+    E_matrix = qk_matrix * D_matrix  # (B, NH, NS, CS, CS)
+
+    normalizer = torch.maximum(
+        (E_matrix.sum(dim=-1, keepdim=True) + inter_n).abs(),
+        torch.exp(-stab),
+    )  # (B, NH, NS, CS, 1)
+
+    E_matrix_normalized = E_matrix / (normalizer + eps)
+
+    # retrieved values
+    intra = E_matrix_normalized @ v  # (B, NH, S, DH)
+    inter = inter_C / (normalizer + eps)
+
+    if return_last_state:
+        return (intra + inter).view((B, NH, S, DH)), (C[:, :, -1], n[:, :, -1], m[:, :, -1])
+    else:
+        return (intra + inter).view((B, NH, S, DH))
+
+
+# chunkwise backend adapted to handle inputs which are not cleanly divisible by chunk_size
+def chunkwise_variable(
+    queries: torch.Tensor,
+    keys: torch.Tensor,
+    values: torch.Tensor,
+    igate_preact: torch.Tensor,
+    fgate_preact: torch.Tensor,
+    initial_C: Optional[torch.Tensor] = None,
+    initial_n: Optional[torch.Tensor] = None,
+    initial_m: Optional[torch.Tensor] = None,
+    chunk_size: int = 64,
+    return_last_state: bool = False,
+    eps: float = 1e-6,
+    **kwargs,
+) -> Union[
+    torch.Tensor, Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
+]:
+    """"
+    Wrapper around chunkwise_simple to allow sequences with arbitrary lengths
+    """
+    tail_size = queries.shape[-2] % chunk_size
+    if tail_size == 0 or queries.shape[-2] < chunk_size:
+        return chunkwise_simple(
+            queries, keys, values, igate_preact, fgate_preact,
+            initial_C, initial_n, initial_m, chunk_size if tail_size == 0 else tail_size, 
+            return_last_state, eps, **kwargs
+        )
+
+    sections = [queries.shape[-2] - tail_size, tail_size]
+    head_args, tail_args = zip(*(torch.split(x, sections, dim=-2) for x in [
+        queries, keys, values, igate_preact, fgate_preact
+    ]))
+    head_out, state = chunkwise_simple(
+        *head_args, initial_C, initial_n, initial_m,
+        chunk_size=chunk_size, return_last_state=True, eps=eps, **kwargs
+    )
+    tail_out = chunkwise_simple(
+        *tail_args, *state, chunk_size=tail_size, 
+        return_last_state=return_last_state, eps=eps, **kwargs
+    )
+
+    if return_last_state:
+        return torch.cat([head_out, tail_out[0]], dim=-2), tail_out[-1]
+    else:
+        return torch.cat([head_out, tail_out], dim=-2)
+
+
+def recurrent_step_stabilized_simple(
+    c_state: torch.Tensor,
+    n_state: torch.Tensor,
+    m_state: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    igate_preact: torch.Tensor,
+    fgate_preact: torch.Tensor,
+    eps: float = 1e-6,
+    **kwargs,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    """This is a single step of the mLSTM operation in recurrent form.
+
+    Args:
+        c_state (torch.Tensor): (B, NH, DH, DH)
+        n_state (torch.Tensor): (B, NH, DH, 1)
+        m_state (torch.Tensor): (B, NH, 1, 1)
+        q (torch.Tensor): (B, NH, 1, DH)
+        k (torch.Tensor): (B, NH, 1, DH)
+        v (torch.Tensor): (B, NH, 1, DH)
+        igate_preact (torch.Tensor): (B, NH, 1, 1)
+        fgate_preact (torch.Tensor): (B, NH, 1, 1)
+
+    Returns:
+        tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+            (hidden_state [B, NH, DH], (c_state_new [B, NH, DH, DH], n_state_new [B, NH, DH, 1]], m_state_new [B, NH, 1, 1]))
+    """
+    B, NH, S, DH = q.shape
+    # projections
+    q, k, v = (
+        q.squeeze_(2).unsqueeze(-1),
+        k.squeeze_(2).unsqueeze(-1),
+        v.squeeze_(2).unsqueeze(-1),
+    )  # (B, NH, DH, 1)
+
+    # gates
+    log_fg_act = torch.nn.functional.logsigmoid(fgate_preact)  # (B, NH, 1, 1)
+
+    # update rule
+    m_state_new = torch.max(log_fg_act + m_state, igate_preact)  # (B, NH, 1, 1)
+
+    fg_act = torch.exp(log_fg_act + m_state - m_state_new)  # (B, NH, 1, 1)
+    ig_act = torch.exp(igate_preact - m_state_new)  # (B, NH, 1, 1)
+
+    k_scaled = k / math.sqrt(DH)
+
+    c_state_new = fg_act * c_state + ig_act * (
+        k_scaled @ v.transpose(-1, -2)
+    )  # (B, NH, DH, DH)
+    n_state_new = fg_act * n_state + ig_act * k_scaled  # (B, NH, DH, 1)
+
+    h_num = q.transpose(-1, -2) @ c_state_new  # (B, NH, 1, DH)
+
+    qn_dotproduct = q.transpose(-1, -2) @ n_state_new  # (B, NH, 1, 1)
+    max_val = torch.exp(-m_state_new)  # (B, NH, 1, 1)
+    h_denom = torch.maximum(qn_dotproduct.abs(), max_val) + eps
+    h = h_num / h_denom  # (B, NH, 1, DH) / (B, NH, 1, 1) = (B, NH, 1, DH)
+
+    return h, (c_state_new, n_state_new, m_state_new)

protxlstm/xlstm/blocks/mlstm/block.py

@@ -0,0 +1,27 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck
+
+# Modified by Pieter-Jan Hoedt, Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Remove sLSTM
+
+
+from dataclasses import dataclass, field
+
+from ..xlstm_block import xLSTMBlock, xLSTMBlockConfig
+from .layer import mLSTMLayerConfig
+
+
+@dataclass
+class mLSTMBlockConfig:
+    mlstm: mLSTMLayerConfig = field(default_factory=mLSTMLayerConfig)
+
+    def __post_init__(self):
+        self.mlstm.__post_init__()
+
+
+class mLSTMBlock(xLSTMBlock):
+
+    config_class = mLSTMBlockConfig
+
+    def __init__(self, config: mLSTMBlockConfig) -> None:
+        super().__init__(config=xLSTMBlockConfig(mlstm=config.mlstm, feedforward=None))

protxlstm/xlstm/blocks/mlstm/cell.py

@@ -0,0 +1,212 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck
+
+# Modified by Pieter-Jan Hoedt, Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Add references to chunkwise backends
+#   - Modify forward to take and return state
+
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from functools import partial
+
+from ...components.init import bias_linspace_init_
+from ...components.ln import MultiHeadLayerNorm
+from .backends import parallel_stabilized_simple, chunkwise_simple, chunkwise_variable, recurrent_step_stabilized_simple
+
+
+@dataclass
+class mLSTMCellConfig:
+    context_length: int = -1
+    embedding_dim: int = -1
+    num_heads: int = -1
+    backend: str = "parallel" # "chunkwise"
+    chunk_size: int = 64
+    return_last_state: bool = False
+
+
+class mLSTMCell(nn.Module):
+    config_class = mLSTMCellConfig
+
+    def __init__(self, config: mLSTMCellConfig):
+        super().__init__()
+        self.config = config
+
+        if self.config.return_last_state == True:
+            assert config.backend != "parallel", "Parallel backend cannot return state - set return_last_state to False or use a chunkwise backend."
+
+        if config.backend == "parallel":
+            self.backend_fn = parallel_stabilized_simple
+        elif config.backend == "chunkwise":
+            chunkwise_backend = partial(chunkwise_simple, chunk_size=config.chunk_size, return_last_state=config.return_last_state)
+            self.backend_fn = chunkwise_backend
+        elif config.backend == "chunkwise_variable":
+            chunkwise_backend = partial(chunkwise_variable, chunk_size=config.chunk_size, return_last_state=config.return_last_state)
+            self.backend_fn = chunkwise_backend
+        else:
+            raise ValueError(f"Unknown mLSTM backend: {config.backend}")
+        self.backend_fn_step = recurrent_step_stabilized_simple
+
+        self.igate = nn.Linear(3 * config.embedding_dim, config.num_heads)
+        self.fgate = nn.Linear(3 * config.embedding_dim, config.num_heads)
+
+        self.outnorm = MultiHeadLayerNorm(ndim=config.embedding_dim, weight=True, bias=False)
+
+        if config.backend == "parallel":
+            self.register_buffer(
+                "causal_mask",
+                torch.tril(torch.ones(config.context_length, config.context_length, dtype=torch.bool)),
+                persistent=False,
+            )
+        else:
+            self.causal_mask = None
+
+        self.reset_parameters()
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, state = None, **kwargs) -> torch.Tensor:
+        B, S, _ = q.shape  # (B, S, H)
+
+        if_gate_input = torch.cat([q, k, v], dim=-1)
+        q = q.view(B, S, self.config.num_heads, -1)  # (B, S, NH, DH)
+        k = k.view(B, S, self.config.num_heads, -1)  # (B, S, NH, DH)
+        v = v.view(B, S, self.config.num_heads, -1)  # (B, S, NH, DH)
+
+        q = q.transpose(1, 2)  # (B, NH, S, DH)
+        k = k.transpose(1, 2)  # (B, NH, S, DH)
+        v = v.transpose(1, 2)  # (B, NH, S, DH)
+
+        # compute input and forget gate pre-activations
+        igate_preact = self.igate(if_gate_input)  # (B, S, NH)
+        igate_preact = igate_preact.transpose(-1, -2).unsqueeze(-1)  # (B, NH, S, 1)
+        fgate_preact = self.fgate(if_gate_input)  # (B, S, NH)
+        fgate_preact = fgate_preact.transpose(-1, -2).unsqueeze(-1)  # (B, NH, S, 1)#
+
+        if state != None and self.config.backend in ["chunkwise", "chunkwise_variable"]:
+
+            initial_C, initial_n, initial_m = state
+            
+            if self.config.return_last_state:
+
+                h_state, mlstm_state = self.backend_fn(
+                    queries=q,
+                    keys=k,
+                    values=v,
+                    igate_preact=igate_preact,
+                    fgate_preact=fgate_preact,
+                    initial_C=initial_C,
+                    initial_n=initial_n,
+                    initial_m=initial_m,
+                    lower_triangular_matrix=self.causal_mask,
+                )
+
+            else:
+                h_state = self.backend_fn(
+                    queries=q,
+                    keys=k,
+                    values=v,
+                    igate_preact=igate_preact,
+                    fgate_preact=fgate_preact,
+                    initial_C=initial_C,
+                    initial_n=initial_n,
+                    initial_m=initial_m,
+                    lower_triangular_matrix=self.causal_mask,
+                )  # (B, NH, S, DH)
+
+        else:
+            if self.config.return_last_state:
+                h_state, mlstm_state = self.backend_fn(
+                    queries=q,
+                    keys=k,
+                    values=v,
+                    igate_preact=igate_preact,
+                    fgate_preact=fgate_preact,
+                    lower_triangular_matrix=self.causal_mask,
+                )
+
+            else:
+                h_state = self.backend_fn(
+                    queries=q,
+                    keys=k,
+                    values=v,
+                    igate_preact=igate_preact,
+                    fgate_preact=fgate_preact,
+                    lower_triangular_matrix=self.causal_mask,
+                )  # (B, NH, S, DH)
+
+        h_state_norm = self.outnorm(h_state)  # (B, NH, S, DH)
+        h_state_norm = h_state_norm.transpose(1, 2).reshape(B, S, -1)  # (B, NH, S, DH) -> (B, S, NH, DH) -> (B, S, H)
+
+        if self.config.return_last_state:
+            return h_state_norm, mlstm_state
+        else:
+            return h_state_norm
+    
+    def step(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        mlstm_state: tuple[torch.Tensor, torch.Tensor, torch.Tensor] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+        B, S, _ = q.shape  # (B, S, H)
+        assert S == 1, f"mLSTMCell.step only supports sequence length S=1, but got S={S}."
+
+        if_gate_input = torch.cat([q, k, v], dim=-1)
+        q = q.view(B, S, self.config.num_heads, -1)  # (B, S, NH, DH)
+        k = k.view(B, S, self.config.num_heads, -1)  # (B, S, NH, DH)
+        v = v.view(B, S, self.config.num_heads, -1)  # (B, S, NH, DH)
+
+        _, _, NH, DH = q.shape
+
+        q = q.transpose(1, 2)  # (B, NH, S, DH)
+        k = k.transpose(1, 2)  # (B, NH, S, DH)
+        v = v.transpose(1, 2)  # (B, NH, S, DH)
+
+        # compute input and forget gate pre-activations
+        igate_preact = self.igate(if_gate_input)  # (B, S, NH)
+        igate_preact = igate_preact.transpose(-1, -2).unsqueeze(-1)  # (B, NH, S, 1)
+        fgate_preact = self.fgate(if_gate_input)  # (B, S, NH)
+        fgate_preact = fgate_preact.transpose(-1, -2).unsqueeze(-1)  # (B, NH, S, 1)
+
+        if mlstm_state is None:
+            c_state = torch.zeros(size=(B, NH, DH, DH), device=q.device, dtype=q.dtype)
+            n_state = torch.zeros(size=(B, NH, DH, 1), device=q.device, dtype=q.dtype)
+            m_state = torch.zeros(size=(B, NH, 1, 1), device=q.device, dtype=q.dtype)
+        else:
+            c_state, n_state, m_state = mlstm_state
+            c_state = c_state.to(device=q.device, dtype=q.dtype)
+            n_state = n_state.to(device=q.device, dtype=q.dtype)
+            m_state = m_state.to(device=q.device, dtype=q.dtype)
+
+        assert c_state.shape == (B, NH, DH, DH), f"Expected c_state shape {(B, NH, DH, DH)}, but got {c_state.shape}."
+        assert n_state.shape == (B, NH, DH, 1), f"Expected n_state shape {(B, NH, DH, 1)}, but got {n_state.shape}."
+        assert m_state.shape == (B, NH, 1, 1), f"Expected m_state shape {(B, NH, 1, 1)}, but got {m_state.shape}."
+
+        h_state, mlstm_state = self.backend_fn_step(
+            c_state=c_state,
+            n_state=n_state,
+            m_state=m_state,
+            q=q,
+            k=k,
+            v=v,
+            igate_preact=igate_preact,
+            fgate_preact=fgate_preact,
+        )  # (B, NH, 1 DH), ((B, NH, DH, DH), (B, NH, DH, 1), (B, NH, 1, 1))
+
+        h_state_norm = self.outnorm(h_state)  # (B, NH, S, DH)
+        h_state_norm = h_state_norm.transpose(1, 2).reshape(B, S, -1)  # (B, NH, S, DH) -> (B, S, NH, DH) -> (B, S, H)
+
+        return h_state_norm, mlstm_state
+    
+
+    def reset_parameters(self):
+        self.outnorm.reset_parameters()
+        # forget gate initialization
+        torch.nn.init.zeros_(self.fgate.weight)
+        bias_linspace_init_(self.fgate.bias, start=3.0, end=6.0)
+        # input gate initialization
+        torch.nn.init.zeros_(self.igate.weight)
+        torch.nn.init.normal_(self.igate.bias, mean=0.0, std=0.1)

protxlstm/xlstm/blocks/mlstm/layer.py

@@ -0,0 +1,217 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck
+
+# Modified by Pieter-Jan Hoedt, Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Modify forward to take and return state
+
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ...components.conv import CausalConv1d, CausalConv1dConfig
+from ...components.init import small_init_init_, wang_init_
+from ...components.linear_headwise import (
+    LinearHeadwiseExpand,
+    LinearHeadwiseExpandConfig,
+)
+from ...utils import UpProjConfigMixin
+from ...components.rotary_position import apply_rotary_emb
+from .cell import mLSTMCell, mLSTMCellConfig
+
+
+@dataclass
+class mLSTMLayerConfig(UpProjConfigMixin):
+    conv1d_kernel_size: int = 4
+    qkv_proj_blocksize: int = 4
+    num_heads: int = 4
+    proj_factor: float = 2.0
+
+    # will be set toplevel config
+    embedding_dim: int = -1
+    bias: bool = False
+    dropout: float = 0.0
+    context_length: int = -1
+    backend: str = "parallel"  # "chunkwise"
+    chunk_size: int = 64
+    return_last_state: bool = False
+
+    _num_blocks: int = 1
+    _inner_embedding_dim: int = None
+
+    def __post_init__(self):
+        self._set_proj_up_dim(embedding_dim=self.embedding_dim)
+        self._inner_embedding_dim = self._proj_up_dim
+
+
+class mLSTMLayer(nn.Module):
+    config_class = mLSTMLayerConfig
+
+    def __init__(self, config: mLSTMLayerConfig):
+        super().__init__()
+        self.config = config
+
+        self.proj_up = nn.Linear(
+            in_features=self.config.embedding_dim,
+            out_features=2 * self.config._inner_embedding_dim,
+            bias=self.config.bias,
+        )
+
+        num_proj_heads = round(self.config._inner_embedding_dim // self.config.qkv_proj_blocksize)
+        self.q_proj = LinearHeadwiseExpand(
+            config=LinearHeadwiseExpandConfig(
+                in_features=self.config._inner_embedding_dim,
+                num_heads=num_proj_heads,
+                bias=self.config.bias,
+            )
+        )
+        self.k_proj = LinearHeadwiseExpand(
+            config=LinearHeadwiseExpandConfig(
+                in_features=self.config._inner_embedding_dim,
+                num_heads=num_proj_heads,
+                bias=self.config.bias,
+            )
+        )
+        self.v_proj = LinearHeadwiseExpand(
+            config=LinearHeadwiseExpandConfig(
+                in_features=self.config._inner_embedding_dim,
+                num_heads=num_proj_heads,
+                bias=self.config.bias,
+            )
+        )
+
+        self.conv1d = CausalConv1d(
+            config=CausalConv1dConfig(
+                feature_dim=self.config._inner_embedding_dim,
+                kernel_size=self.config.conv1d_kernel_size,
+            )
+        )
+        self.conv_act_fn = nn.SiLU()
+        self.mlstm_cell = mLSTMCell(
+            config=mLSTMCellConfig(
+                context_length=self.config.context_length,
+                embedding_dim=self.config._inner_embedding_dim,
+                num_heads=self.config.num_heads,
+                backend=self.config.backend,
+                chunk_size=self.config.chunk_size,
+                return_last_state = self.config.return_last_state
+            )
+        )
+        self.ogate_act_fn = nn.SiLU()
+
+        self.learnable_skip = nn.Parameter(torch.ones(self.config._inner_embedding_dim, requires_grad=True))
+
+        self.proj_down = nn.Linear(
+            in_features=self.config._inner_embedding_dim,
+            out_features=self.config.embedding_dim,
+            bias=self.config.bias,
+        )
+        self.dropout = nn.Dropout(self.config.dropout)
+        self.reset_parameters()
+
+    def forward(self, x: torch.Tensor, freqs_cos=None, freqs_sin=None, state=None, **kwargs) -> torch.Tensor:
+        B, S, _ = x.shape
+
+        # up-projection
+        x_inner = self.proj_up(x)
+        x_mlstm, z = torch.split(x_inner, split_size_or_sections=self.config._inner_embedding_dim, dim=-1)
+
+        # mlstm branch
+        if state != None:
+            mlstm_state = state["mlstm_state"]
+            conv_state = state["conv_state"][0]
+        else:
+            mlstm_state, conv_state = None, None
+
+        if self.config.return_last_state:
+            x_mlstm_conv, conv_state = self.conv1d(x_mlstm, conv_state = conv_state, return_last_state = True)
+        else:
+            x_mlstm_conv = self.conv1d(x_mlstm, conv_state = conv_state)
+                
+        x_mlstm_conv_act = self.conv_act_fn(x_mlstm_conv)
+
+        q = self.q_proj(x_mlstm_conv_act)
+        k = self.k_proj(x_mlstm_conv_act)
+        v = self.v_proj(x_mlstm)
+
+        if freqs_cos is not None and freqs_sin is not None:
+            q, k = apply_rotary_emb(q, k, freqs_cos, freqs_sin)
+
+        if self.config.return_last_state:
+            h_tilde_state, mlstm_state = self.mlstm_cell(q=q, k=k, v=v, state=mlstm_state, **kwargs)
+        else:
+            h_tilde_state = self.mlstm_cell(q=q, k=k, v=v, state=mlstm_state, **kwargs)
+
+        h_tilde_state_skip = h_tilde_state + (self.learnable_skip * x_mlstm_conv_act)
+
+        # output / z branch
+        h_state = h_tilde_state_skip * self.ogate_act_fn(z)
+
+        # down-projection
+        y = self.dropout(self.proj_down(h_state))
+
+        if self.config.return_last_state:
+            return y, {"mlstm_state": mlstm_state, "conv_state": (conv_state,)}
+        else:
+            return y
+    
+    def step(
+        self,
+        x: torch.Tensor,
+        freqs_cos=None, 
+        freqs_sin=None,
+        mlstm_state: tuple[torch.Tensor, torch.Tensor, torch.Tensor] = None,
+        conv_state: tuple[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, dict[str, tuple[torch.Tensor, ...]]]:
+        B, S, _ = x.shape
+
+        # up-projection
+        x_inner = self.proj_up(x)
+        x_mlstm, z = torch.split(x_inner, split_size_or_sections=self.config._inner_embedding_dim, dim=-1)
+
+        # mlstm branch
+        x_mlstm_conv, conv_state = self.conv1d.step(x_mlstm, conv_state=conv_state)
+        x_mlstm_conv_act = self.conv_act_fn(x_mlstm_conv)
+
+        q = self.q_proj(x_mlstm_conv_act)
+        k = self.k_proj(x_mlstm_conv_act)
+        v = self.v_proj(x_mlstm)
+
+        if freqs_cos is not None and freqs_sin is not None:
+            q, k = apply_rotary_emb(q, k, freqs_cos, freqs_sin)
+
+        h_tilde_state, mlstm_state = self.mlstm_cell.step(q=q, k=k, v=v, mlstm_state=mlstm_state)
+
+        h_tilde_state_skip = h_tilde_state + (self.learnable_skip * x_mlstm_conv_act)
+
+        # output / z branch
+        h_state = h_tilde_state_skip * self.ogate_act_fn(z)
+
+        # down-projection
+        y = self.dropout(self.proj_down(h_state))
+        return y, {"mlstm_state": mlstm_state, "conv_state": conv_state}
+
+    def reset_parameters(self):
+        # init inproj
+        small_init_init_(self.proj_up.weight, dim=self.config.embedding_dim)
+        if self.proj_up.bias is not None:
+            nn.init.zeros_(self.proj_up.bias)
+        # init outproj
+        wang_init_(self.proj_down.weight, dim=self.config.embedding_dim, num_blocks=self.config._num_blocks)
+        if self.proj_down.bias is not None:
+            nn.init.zeros_(self.proj_down.bias)
+
+        nn.init.ones_(self.learnable_skip)
+
+        def _init_qkv_proj(qkv_proj: LinearHeadwiseExpand):
+            # use the embedding dim instead of the inner embedding dim
+            small_init_init_(qkv_proj.weight, dim=self.config.embedding_dim)
+            if qkv_proj.bias is not None:
+                nn.init.zeros_(qkv_proj.bias)
+
+        _init_qkv_proj(self.q_proj)
+        _init_qkv_proj(self.k_proj)
+        _init_qkv_proj(self.v_proj)
+
+        self.mlstm_cell.reset_parameters()

protxlstm/xlstm/blocks/xlstm_block.py

@@ -0,0 +1,111 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck
+
+# Modified by Pieter-Jan Hoedt, Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Remove sLSTM
+#   - Modify forward to take and return state
+
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import nn
+
+from ..components.feedforward import FeedForwardConfig, create_feedforward
+from ..components.ln import LayerNorm
+from .mlstm.layer import mLSTMLayer, mLSTMLayerConfig
+
+"""An xLSTM block can be either an sLSTM Block or an mLSTM Block.
+In this repository only mLSTM is implemented.
+
+It contains the pre-LayerNorms and the skip connections.
+"""
+
+
+@dataclass
+class xLSTMBlockConfig:
+    mlstm: Optional[mLSTMLayerConfig] = None
+
+    feedforward: Optional[FeedForwardConfig] = None
+
+    _num_blocks: int = 1
+    _block_idx: int = 0
+
+    def __post_init__(self):
+        assert (
+            self.mlstm is not None
+        ), "mlstm config must be provided"
+
+        embedding_dim = (
+            self.mlstm.embedding_dim
+        )
+
+        self.mlstm._num_blocks = self._num_blocks
+        self.mlstm._block_idx = self._block_idx
+
+        if self.feedforward:
+            self.feedforward.embedding_dim = embedding_dim
+            self.feedforward._num_blocks = self._num_blocks
+            self.feedforward.__post_init__()
+
+
+class xLSTMBlock(nn.Module):
+
+    config_class = xLSTMBlockConfig
+
+    def __init__(self, config: xLSTMBlockConfig) -> None:
+        super().__init__()
+        self.config = config
+        embedding_dim = (
+            self.config.mlstm.embedding_dim
+        )
+
+        self.xlstm_norm = LayerNorm(ndim=embedding_dim, weight=True, bias=False)
+
+        if self.config.mlstm is not None:
+            self.xlstm = mLSTMLayer(config=self.config.mlstm)
+        else:
+            raise ValueError("mlstm must be provided")
+
+        if self.config.feedforward is not None:
+            self.ffn_norm = LayerNorm(
+                ndim=self.config.feedforward.embedding_dim, weight=True, bias=False
+            )
+            self.ffn = create_feedforward(config=self.config.feedforward)
+        else:
+            self.ffn_norm = None
+            self.ffn = None
+
+        self.reset_parameters()
+
+    def forward(self, x: torch.Tensor, state=None, **kwargs) -> torch.Tensor:
+        if self.config.mlstm.return_last_state:
+            x_xlstm, xlstm_state = self.xlstm(self.xlstm_norm(x), state=state, **kwargs)
+            x = x + x_xlstm
+        else:
+            x = x + self.xlstm(self.xlstm_norm(x), state=state, **kwargs)
+
+        if self.ffn is not None:
+            x = x + self.ffn(self.ffn_norm(x), **kwargs)
+
+        if self.config.mlstm.return_last_state:
+            return x, xlstm_state
+        else:
+            return x
+    
+    def step(self, x: torch.Tensor, **kwargs) -> tuple[torch.Tensor, dict[str, tuple[torch.Tensor, ...]]]:
+        x_xlstm, xlstm_state = self.xlstm.step(self.xlstm_norm(x), **kwargs)
+        x = x + x_xlstm
+        if self.ffn is not None:
+            x = x + self.ffn(self.ffn_norm(x), **kwargs)
+        return x, xlstm_state
+    
+    def reset_parameters(self) -> None:
+        
+        self.xlstm.reset_parameters()
+        self.xlstm_norm.reset_parameters()
+    
+        if self.ffn is not None:
+            self.ffn.reset_parameters()
+            self.ffn_norm.reset_parameters()

protxlstm/xlstm/components/conv.py

@@ -0,0 +1,163 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck, Korbinian Pöppel
+
+# Modified by Pieter-Jan Hoedt, Niklas Schmidinger, Lisa Schneckenreiter and Sohvi Luukkonen 
+#   - Modify forward to take and return state
+
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+
+# from einops import rearrange
+from torch import nn
+
+
+@dataclass
+class CausalConv1dConfig:
+    feature_dim: int = None  # F
+    kernel_size: int = 4
+    causal_conv_bias: bool = True
+    channel_mixing: bool = False
+    conv1d_kwargs: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        assert self.kernel_size >= 0, "kernel_size must be >= 0"
+
+
+def conv1d_step(
+    x: torch.Tensor,
+    conv_state: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    B: batch size
+    S: sequence length
+    D: feature dimension
+    KS: kernel size
+    Args:
+        x (torch.Tensor): (B, S, D)
+        conv_state (torch.Tensor): (B, KS, D)
+        conv1d_weight (torch.Tensor): (KS, D)
+    """
+    assert (
+        x.shape[0] == conv_state.shape[0]
+    ), f"x has batch size {x.shape[0]} but conv_state has batch size {conv_state.shape[0]}"
+    assert (
+        x.shape[2] == conv_state.shape[2]
+    ), f"x has feature dimension {x.shape[2]} but conv_state has feature dimension {conv_state.shape[2]}"
+    assert x.shape[1] == 1, f"x has sequence length {x.shape[1]} but it should be 1"
+    conv_state.copy_(torch.roll(conv_state, shifts=-1, dims=1))
+    conv_state[:, -1:, :] = x
+    y = torch.sum(conv_state * conv1d_weight, dim=1, keepdim=True)
+    if conv1d_bias is not None:
+        y += conv1d_bias
+    return y, conv_state
+
+
+class CausalConv1d(nn.Module):
+    config_class = CausalConv1dConfig
+    """
+    Implements causal depthwise convolution of a time series tensor.
+    Input:  Tensor of shape (B,T,F), i.e. (batch, time, feature)
+    Output: Tensor of shape (B,T,F)
+
+    Args:
+        feature_dim: number of features in the input tensor
+        kernel_size: size of the kernel for the depthwise convolution
+        causal_conv_bias: whether to use bias in the depthwise convolution
+        channel_mixing: whether to use channel mixing (i.e. groups=1) or not (i.e. groups=feature_dim)
+                        If True, it mixes the convolved features across channels.
+                        If False, all the features are convolved independently.
+    """
+
+    def __init__(self, config: CausalConv1dConfig):
+        super().__init__()
+        self.config = config
+        self.groups = self.config.feature_dim
+        if self.config.channel_mixing:
+            self.groups = 1
+        if self.config.kernel_size == 0:
+            self.conv = None  # Noop
+        else:
+            self.pad = (
+                self.config.kernel_size - 1
+            )  # padding of this size assures temporal causality.
+            self.conv = nn.Conv1d(
+                in_channels=self.config.feature_dim,
+                out_channels=self.config.feature_dim,
+                kernel_size=self.config.kernel_size,
+                padding=self.pad,
+                groups=self.groups,
+                bias=self.config.causal_conv_bias,
+                **self.config.conv1d_kwargs,
+            )
+        # B, C, L
+        self.reset_parameters()
+
+    def reset_parameters(self, **kwargs):
+        self.conv.reset_parameters()
+
+    def _create_weight_decay_optim_groups(
+        self,
+    ) -> tuple[set[nn.Parameter], set[nn.Parameter]]:
+        if self.config.kernel_size == 0:
+            return (), ()
+        else:
+            weight_decay = (self.conv.weight,)
+            no_weight_decay = ()
+            if self.config.causal_conv_bias:
+                no_weight_decay += (self.conv.bias,)
+            return weight_decay, no_weight_decay
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        conv_state: Optional[torch.Tensor] = None,
+        return_last_state: bool = False,
+    ) -> torch.Tensor:
+        if conv_state is not None:
+            conv_state = conv_state[:,-self.pad:]
+            x = torch.cat([conv_state, x], dim=1)
+
+        if self.config.kernel_size == 0:
+            return x
+        y = x.transpose(2, 1)  # (B,F,T) tensor - now in the right shape for conv layer.
+        y = self.conv(y)  # (B,F,T+pad) tensor
+        if conv_state is not None:
+            y = y[:, :, conv_state.shape[1] :]
+
+        if return_last_state:
+            return y[:, :, : -self.pad].transpose(2, 1), x[:, -self.config.kernel_size:]   #[:, -self.pad :]
+        else:
+            return y[:, :, : -self.pad].transpose(2, 1)
+
+    def step(
+        self,
+        x: torch.Tensor,
+        conv_state: tuple[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor]]:
+
+        if self.config.kernel_size == 0:
+            return x, conv_state
+
+        B, S, D = x.shape
+
+        if conv_state is None:
+            conv_state = (
+                torch.zeros(
+                    size=(B, self.config.kernel_size, D),
+                    device=self.conv.weight.device,
+                    dtype=self.conv.weight.dtype,
+                ),
+            )
+
+        y, conv_state = conv1d_step(
+            x,
+            conv_state[0],
+            self.conv.weight[:, 0, :].transpose(0, 1),  # rearrange(, "D 1 KS -> KS D")
+            conv1d_bias=self.conv.bias if self.config.causal_conv_bias else None,
+        )
+        return y, (conv_state,)

protxlstm/xlstm/components/feedforward.py

@@ -0,0 +1,88 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck
+from dataclasses import dataclass
+from typing import Callable, Literal
+
+import torch
+from torch import nn
+
+from ..utils import UpProjConfigMixin
+from .init import small_init_init_, wang_init_
+
+_act_fn_registry = {
+    "gelu": nn.functional.gelu,
+    "relu": nn.functional.relu,
+    "relu^2": lambda x: torch.square(nn.functional.relu(x)),
+    "sigmoid": nn.functional.sigmoid,
+    "swish": nn.functional.silu,
+    "selu": nn.functional.selu,
+}
+
+
+def get_act_fn(act_fn_name: str) -> Callable[[torch.Tensor], torch.Tensor]:
+    if act_fn_name in _act_fn_registry:
+        return _act_fn_registry[act_fn_name]
+    else:
+        assert (
+            False
+        ), f'Unknown activation function name "{act_fn_name}". Available activation functions are: {str(_act_fn_cls_registry.keys())}'
+
+
+@dataclass
+class FeedForwardConfig(UpProjConfigMixin):
+    proj_factor: float = 1.3
+    act_fn: str = "gelu"
+    embedding_dim: int = -1
+    dropout: float = 0.0
+    bias: bool = False
+    ff_type: Literal["ffn_gated"] = "ffn_gated"
+
+    _num_blocks: int = 1
+
+    def __post_init__(self):
+        self._set_proj_up_dim(embedding_dim=self.embedding_dim)
+        assert self.act_fn in _act_fn_registry, f"Unknown activation function {self.act_fn}"
+
+
+class GatedFeedForward(nn.Module):
+    config_class = FeedForwardConfig
+
+    def __init__(self, config: FeedForwardConfig):
+        super().__init__()
+        self.config = config
+
+        self.proj_up = nn.Linear(
+            in_features=self.config.embedding_dim,
+            out_features=2 * self.config._proj_up_dim,
+            bias=self.config.bias,
+        )
+        self.proj_down = nn.Linear(
+            in_features=self.config._proj_up_dim,
+            out_features=self.config.embedding_dim,
+            bias=self.config.bias,
+        )
+
+        self.act_fn = get_act_fn(self.config.act_fn)
+
+        self.dropout = nn.Dropout(self.config.dropout)
+        self.reset_parameters()
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        gate_preact, up_proj = self.proj_up(x).split(self.config._proj_up_dim, dim=-1)
+        x = self.dropout(self.proj_down(self.act_fn(gate_preact) * up_proj))
+        return x
+
+    def reset_parameters(self):
+        small_init_init_(self.proj_up.weight, dim=self.config.embedding_dim)
+        if self.proj_up.bias is not None:
+            nn.init.zeros_(self.proj_up.bias)
+        wang_init_(self.proj_down.weight, dim=self.config.embedding_dim, num_blocks=self.config._num_blocks)
+        if self.proj_down.bias is not None:
+            nn.init.zeros_(self.proj_down.bias)
+
+
+def create_feedforward(config: FeedForwardConfig) -> nn.Module:
+    if config.ff_type == "ffn_gated":
+        return GatedFeedForward(config)
+    else:
+        raise ValueError(f"Unknown feedforward type {config.ff_type}")

protxlstm/xlstm/components/init.py

@@ -0,0 +1,32 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck
+import math
+
+import torch
+
+
+def bias_linspace_init_(param: torch.Tensor, start: float = 3.4, end: float = 6.0) -> torch.Tensor:
+    """Linearly spaced bias init across dimensions."""
+    assert param.dim() == 1, f"param must be 1-dimensional (typically a bias), got {param.dim()}"
+    n_dims = param.shape[0]
+    init_vals = torch.linspace(start, end, n_dims)
+    with torch.no_grad():
+        param.copy_(init_vals)
+    return param
+
+
+def small_init_init_(param: torch.Tensor, dim: int) -> torch.Tensor:
+    """Fills the input Tensor with values according to the method described in Transformers without Tears: Improving
+    the Normalization of Self-Attention - Nguyen, T. & Salazar, J. (2019), using a normal distribution.
+    Adopted from https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/init_functions.py.
+    """
+    std = math.sqrt(2 / (5 * dim))
+    torch.nn.init.normal_(param, mean=0.0, std=std)
+    return param
+
+
+def wang_init_(param: torch.Tensor, dim: int, num_blocks: int):
+    """Adopted from https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/init_functions.py."""
+    std = 2 / num_blocks / math.sqrt(dim)
+    torch.nn.init.normal_(param, mean=0.0, std=std)
+    return param

protxlstm/xlstm/components/linear_headwise.py

@@ -0,0 +1,92 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck, Korbininan Pöppel
+from dataclasses import dataclass
+
+from math import sqrt
+import torch
+
+# from einops import einsum, rearrange
+from torch import nn
+
+
+@dataclass
+class LinearHeadwiseExpandConfig:
+    in_features: int = 0
+    # this is the number of heads that the in_features are split into
+    # if num_heads=1, this is a normal linear layer
+    # if num_heads>1, the in_features are split into num_heads and each head is projected separately
+    # if num_heads=in_features, each feature is projected separately
+    num_heads: int = -1
+    expand_factor_up: float = 1
+
+    # this is internally computed
+    # but can be overwritten if you want to use a different output dimension
+    # if > 0 the expand factor is ignored
+    _out_features: int = -1
+
+    bias: bool = True
+    trainable_weight: bool = True
+    trainable_bias: bool = True
+
+    def __post_init__(self):
+        assert self.num_heads > 0, "num_heads must be set"
+        assert self.num_heads <= self.in_features, "num_heads must be <= in_features"
+        assert (
+            self.in_features % self.num_heads == 0
+        ), "in_features must be a multiple of num_heads"
+
+        if self._out_features < 0:
+            self._out_features = round(self.expand_factor_up * self.in_features)
+
+
+class LinearHeadwiseExpand(nn.Module):
+    """This is a structured projection layer that projects the input to a higher dimension.
+    It only allows integer up-projection factors, i.e. the output dimension is a multiple of the input dimension.
+    """
+
+    config_class = LinearHeadwiseExpandConfig
+
+    def __init__(self, config: LinearHeadwiseExpandConfig):
+        super().__init__()
+        self.config = config
+        in_features = self.config.in_features
+        num_heads = self.config.num_heads
+        out_features_per_head = config._out_features // num_heads
+        self.weight = nn.Parameter(
+            torch.empty(num_heads, out_features_per_head, in_features // num_heads),
+            requires_grad=config.trainable_weight,
+        )
+        if config.bias:
+            self.bias = nn.Parameter(
+                torch.empty(config._out_features), requires_grad=config.trainable_bias
+            )
+        else:
+            self.bias = None
+        self.reset_parameters()
+
+    def reset_parameters(self, **kwargs):
+        # small init
+        nn.init.normal_(
+            self.weight.data, mean=0.0, std=sqrt(2 / 5 / self.weight.shape[-1])
+        )
+        if self.bias is not None:
+            nn.init.zeros_(self.bias.data)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        x = x.view(*shape[:-1], self.config.num_heads, -1)
+        x = torch.einsum("...hd,hod->...ho", x, self.weight)
+        x = x.reshape(*shape[:-1], -1)
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+    def extra_repr(self):
+        return (
+            f"in_features={self.config.in_features}, "
+            f"num_heads={self.config.num_heads}, "
+            f"expand_factor_up={self.config.expand_factor_up}, "
+            f"bias={self.config.bias}, "
+            f"trainable_weight={self.config.trainable_weight}, "
+            f"trainable_bias={self.config.trainable_bias}, "
+        )

protxlstm/xlstm/components/ln.py

@@ -0,0 +1,68 @@
+# Copyright (c) NXAI GmbH and its affiliates 2024
+# Maximilian Beck, Korbinian Pöppel
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class LayerNorm(nn.Module):
+    """LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False."""
+
+    def __init__(
+        self,
+        ndim: int = -1,
+        weight: bool = True,
+        bias: bool = False,
+        eps: float = 1e-5,
+        residual_weight: bool = True,
+    ):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(ndim)) if weight else None
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+        self.eps = eps
+        self.residual_weight = residual_weight
+        self.ndim = ndim
+        self.reset_parameters()
+
+    @property
+    def weight_proxy(self) -> torch.Tensor:
+        if self.weight is None:
+            return None
+        if self.residual_weight:
+            return 1.0 + self.weight
+        else:
+            return self.weight
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.layer_norm(
+            input, normalized_shape=(self.ndim,), weight=self.weight_proxy, bias=self.bias, eps=self.eps
+        )
+
+    def reset_parameters(self):
+        if self.weight_proxy is not None:
+            if self.residual_weight:
+                nn.init.zeros_(self.weight)
+            else:
+                nn.init.ones_(self.weight)
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+
+class MultiHeadLayerNorm(LayerNorm):
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        assert input.dim() == 4, "Input must be 4D tensor (B, NH, S, DH)"
+        B, NH, S, DH = input.shape
+
+        gn_in_1 = input.transpose(1, 2)  # (B, S, NH, DH)
+        gn_in_2 = gn_in_1.reshape(B * S, NH * DH)  # (B * S, NH * DH)
+        out = F.group_norm(
+            gn_in_2,
+            num_groups=NH,
+            weight=self.weight_proxy,
+            bias=self.bias,
+            eps=self.eps,
+        )
+        # (B * S), (NH * DH) -> (B, S, NH, DH) -> (B, NH, S, DH)
+        out = out.view(B, S, NH, DH).transpose(1, 2)
+        return out

protxlstm/xlstm/components/rotary_position.py

@@ -0,0 +1,35 @@
+import torch
+from typing import Tuple
+import math
+
+
+def compute_freqs_cis(t: torch.Tensor, head_dim: int, theta: float = 10_000.0):
+    freqs = theta ** (-torch.arange(0, head_dim, 2).float() / head_dim)
+    freqs = t.unsqueeze(-1) * freqs.to(t.device)  # type: ignore
+    freqs_cos = torch.cos(freqs)  # real part
+    freqs_sin = torch.sin(freqs)  # imaginary part
+    return freqs_cos, freqs_sin
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    # reshape xq and xk to match the complex representation
+    xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
+    xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)
+
+    # apply rotation using real numbers
+    xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
+    xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
+    xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin
+    xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos
+
+    # flatten last two dimensions
+    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(-2)
+    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(-2)
+
+    return xq_out.type_as(xq), xk_out.type_as(xk)