Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1,89 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
+ 
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text
+*.ogg filter=lfs diff=lfs merge=lfs -text
+*.m4a filter=lfs diff=lfs merge=lfs -text
+*.aac filter=lfs diff=lfs merge=lfs -text
+*.opus filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.mov filter=lfs diff=lfs merge=lfs -text
+*.avi filter=lfs diff=lfs merge=lfs -text
+*.mkv filter=lfs diff=lfs merge=lfs -text
+*.webm filter=lfs diff=lfs merge=lfs -text
+*.obj filter=lfs diff=lfs merge=lfs -text
+*.glb filter=lfs diff=lfs merge=lfs -text
+*.gltf filter=lfs diff=lfs merge=lfs -text
+*.ply filter=lfs diff=lfs merge=lfs -text
+*.stl filter=lfs diff=lfs merge=lfs -text
+*.fbx filter=lfs diff=lfs merge=lfs -text
+
+data/inverse_folding/5YH2.cif filter=lfs diff=lfs merge=lfs -text
+data/inverse_folding/5YH2.pdb filter=lfs diff=lfs merge=lfs -text
+data/variant_prediction/BLAT_ECOLX_Ranganathan2015.csv filter=lfs diff=lfs merge=lfs -text
+
+ 
+
+ 
+weight/esm-main/examples/data/P62593.fasta filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/inverse_folding/data/4uv3.cif filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/inverse_folding/data/4uv3.pdb filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/inverse_folding/data/5YH2.cif filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/inverse_folding/data/5YH2.pdb filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/inverse_folding/notebook.ipynb filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/inverse_folding/notebook_multichain.ipynb filter=lfs diff=lfs merge=lfs -text
+
+weight/esm-main/examples/variant-prediction/data/BLAT_ECOLX_1_b0.5.a3m filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/variant-prediction/data/BLAT_ECOLX_Ranganathan2015.csv filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/variant-prediction/data/BLAT_ECOLX_Ranganathan2015_labeled.csv filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/variant-prediction/data/BLAT_ECOLX_Ranganathan2015_labeled_masked.csv filter=lfs diff=lfs merge=lfs -text
+
+ 
+
+*.pt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,283 @@
+---
+license: mit
+tasks:
+  - protein-structure-prediction
+frameworks:
+  - pytorch
+language:
+  - en
+  - zh
+tags:
+  - OneScience
+  - Life Sciences
+  - Protein Language Model
+  - Protein Structure Prediction
+  - Variant Effect Prediction
+  - ESM
+huggingface:
+  model: OneScience-Sugon/ESM
+  data_path: data/
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">ESM</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+ESM (Evolutionary Scale Modeling) is a family of protein language models released by Meta AI / FAIR. It can be used for protein representation extraction, structure prediction, variant effect scoring, and fixed-backbone sequence design.
+
+Paper: Evolutionary-scale prediction of atomic-level protein structure with a language model  
+https://www.science.org/doi/10.1126/science.ade2574
+
+# Model Description
+
+This model package provides PyTorch inference support for ESM-1, ESM-2, MSA Transformer, ESMFold, ESM-1v, and ESM-IF1, together with adaptations for running on DCUs. Sample data is distributed with the Hugging Face model repository `OneScience-Sugon/ESM`.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Protein representation extraction | Takes a FASTA file as input and outputs per-token, mean-pooled, BOS, or contact representations |
+| Protein structure prediction | Takes one or more amino acid sequences as input and outputs corresponding PDB structure files |
+| Variant effect scoring | Takes a wild-type sequence and a DMS mutation table as input and outputs mutation effect scores |
+| Fixed-backbone sequence design | Takes a PDB / CIF structure and chain ID as input and samples candidate sequences that satisfy the backbone constraints |
+| Structure-conditioned sequence scoring | Takes a structure and candidate sequences as input and computes their conditional log-likelihoods |
+| Hugging Face / OneCode execution | After downloading the model project, quickly verifies that the scripts run correctly in a life-sciences runtime environment |
+
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Try one-click AI4S development in the OneCode online environment:
+
+[Try one-click AI4S development](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- GPU or DCU is recommended.
+- A CPU can be used for import checks and lightweight configuration tests; full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or later is recommended, or a OneScience-recommended version matching the current cluster.
+
+
+
+
+
+**Environment Check**
+
+- NVIDIA GPU:
+
+```bash
+nvidia-smi
+```
+
+- Hygon DCU:
+
+```bash
+hy-smi
+```
+
+### Download the Model Package
+
+```bash
+hf download --model OneScience-Sugon/ESM --local-dir ./ESM
+cd ESM
+```
+
+This model package includes a small set of sample data that can be used directly to validate the default workflow.
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation supported
+pip install onescience[bio-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+```bash
+# If required libraries cannot be found, activate the CUDA compatibility environment as follows:
+source ${ROCM_PATH}/cuda/env.sh
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"
+```
+
+After installation, return to the model package directory:
+
+```bash
+cd ./ESM
+```
+
+### Training and Inference Data Overview
+
+The FASTA, PDB / CIF, and DMS files used in the ESM examples are distributed with the [Hugging Face model repository OneScience-Sugon/ESM](https://huggingface.co/OneScience-Sugon/ESM). After downloading the complete model package, the files are available under `data/`. This model package does not include a training entry point; the data is intended for example inference and workflow validation. You can also download only the data directory:
+
+```bash
+hf download --model OneScience-Sugon/ESM ESM/data --local-dir ./data
+```
+### Model Weights
+
+The repository includes multiple ESM model checkpoints under `weight/`; select the appropriate checkpoint for inference.
+
+### Preparing Weights
+
+Place the required ESM weights in the following directory:
+
+```text
+weight/
+  checkpoints/
+    esm2_t6_8M_UR50D.pt
+    esmfold_3B_v1.pt
+    esm1v_t33_650M_UR90S_1.pt
+    esm_if1_gvp4_t16_142M_UR50.pt
+    ...
+```
+
+When using a shared runtime environment, you can specify the weights location through an environment variable:
+
+```bash
+export ESM_WEIGHT_DIR=/path/to/esm/weight
+```
+
+The default example uses:
+
+- `weight/checkpoints/esm2_t6_8M_UR50D.pt`
+
+The ESMFold, ESM-1v, and ESM-IF1 examples require their respective weights to be available.
+
+### Default Example
+
+```bash
+bash scripts/infer.sh
+```
+
+The default example reads `data/fasta/few_proteins.fasta`, extracts protein representations using `esm2_t6_8M_UR50D.pt`, and saves the results to `outputs/embeddings/`.
+
+### Sequence Representation Extraction
+
+```bash
+python scripts/extract.py \
+  weight/checkpoints/esm2_t6_8M_UR50D.pt \
+  data/fasta/few_proteins.fasta \
+  outputs/embeddings \
+  --include mean per_tok \
+  --repr_layers 6
+```
+
+### ESMFold Structure Prediction
+
+```bash
+python scripts/fold.py \
+  -i data/fasta/few_proteins.fasta \
+  -o outputs/pdb \
+  --model-dir weight \
+  --cpu-only
+```
+
+The output directory will contain one or more `.pdb` files. For production GPU / DCU inference, remove `--cpu-only` and configure `--chunk-size` or `--max-tokens-per-batch` according to the available accelerator memory.
+
+You can also explicitly enable ESMFold via the default script:
+
+```bash
+RUN_ESMFOLD=1 bash scripts/infer.sh
+```
+
+### Inverse Folding — Sequence Sampling
+
+```bash
+python scripts/inverse_folding/sample_sequences.py \
+  data/inverse_folding/5YH2.pdb \
+  --chain A \
+  --outpath outputs/sampled_seqs.fasta \
+  --num-samples 1 \
+  --nogpu
+```
+
+### Inverse Folding — Sequence Scoring
+
+```bash
+python scripts/inverse_folding/score_log_likelihoods.py \
+  data/inverse_folding/5YH2.pdb \
+  data/inverse_folding/5YH2_mutated_seqs.fasta \
+  --chain A \
+  --outpath outputs/sequence_scores.csv \
+  --nogpu
+```
+
+### Variant Effect Prediction
+
+Variant effect prediction requires a wild-type sequence that is consistent with the mutation annotations in the DMS table:
+
+```bash
+python scripts/variant_prediction/predict.py \
+  --model-location esm1v_t33_650M_UR90S_1 \
+  --sequence "${ESM_VARIANT_SEQUENCE}" \
+  --dms-input data/variant_prediction/BLAT_ECOLX_Ranganathan2015.csv \
+  --mutation-col mutant \
+  --dms-output outputs/variant_prediction.csv \
+  --offset-idx 24 \
+  --scoring-strategy wt-marginals
+```
+
+# Data Format
+
+Sample data is stored under `data/` by default:
+
+```text
+data/
+  fasta/
+    few_proteins.fasta
+    some_proteins.fasta
+  inverse_folding/
+    5YH2.pdb
+    5YH2.cif
+    5YH2_mutated_seqs.fasta
+    example.json
+  variant_prediction/
+    BLAT_ECOLX_Ranganathan2015.csv
+    rho_pp.csv
+    aggregated_rho.csv
+    aggregated_rho_round3.csv
+```
+
+In this structure:
+
+- FASTA files are used for sequence representation extraction and structure prediction.
+- PDB / CIF files are used for inverse folding sampling and structure-conditioned sequence scoring.
+- Variant effect prediction CSV files must include a mutation column. The default column name is `mutant`, and mutations use notation such as `A123B`.
+- For custom DMS data, the wild-type amino acid at each mutated position in the sequence provided via `--sequence` must match the corresponding mutation annotation.
+
+# Verification
+
+Static import check:
+
+```bash
+python scripts/check_import_boundaries.py
+```
+
+Syntax check:
+
+```bash
+python -B -c "import ast, pathlib; [ast.parse(p.read_text(encoding='utf-8'), filename=str(p)) for root in ['model', 'scripts', 'tests'] for p in pathlib.Path(root).rglob('*.py')]"
+```
+
+# OneScience Official Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation & License
+
+- This repository is adapted from the open-source ESM model to support DCUs.
+- The ESM source code is licensed under the MIT License; see `LICENSE`. For the usage terms governing model weights and data, refer to the documentation provided by the respective publishers.
+- For scientific use, please cite the corresponding original ESM paper for each submodel used. For ESM-2 / ESMFold, cite: [Evolutionary-scale prediction of atomic-level protein structure with a language model](https://www.science.org/doi/10.1126/science.ade2574).

configuration.json

@@ -0,0 +1,2 @@
+{"framework":"Pytorch","task":"other"}
+

model/README.md

@@ -0,0 +1,10 @@
+# Local Model Source
+
+This directory contains OneScience `models` sources that ESM imports directly.
+
+- `esm/`: local copy of `onescience.models.esm`
+- `openfold/`: local copy of `onescience.models.openfold`, required by ESMFold
+- `protenix/layer_norm/`: optional fused layer norm dependency used by OpenFold when `LAYERNORM_TYPE=fast_layernorm`
+
+Other OneScience namespaces such as `onescience.datapipes`, `onescience.modules`, and `onescience.utils` are intentionally imported from the installed OneScience environment.
+

model/__init__.py

@@ -0,0 +1,2 @@
+"""Local model packages copied from OneScience for the standalone ESM project."""
+

model/esm/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from onescience.datapipes import esm as data  # noqa
+from onescience.datapipes.esm import Alphabet, BatchConverter, FastaBatchedDataset  # noqa
+
+from .version import version as __version__  # noqa
+from .esm1 import ProteinBertModel  # noqa
+from .esm2 import ESM2  # noqa
+from .msa_transformer import MSATransformer  #noqa
+from . import pretrained  # noqa

model/esm/esm1.py

@@ -0,0 +1,200 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from onescience.modules.esm import (
+    TransformerLayer,
+    LearnedPositionalEmbedding,
+    SinusoidalPositionalEmbedding,
+    RobertaLMHead,
+    ESM1bLayerNorm,
+    ContactPredictionHead,
+)
+
+
+class ProteinBertModel(nn.Module):
+    @classmethod
+    def add_args(cls, parser):
+        parser.add_argument(
+            "--num_layers", default=36, type=int, metavar="N", help="number of layers"
+        )
+        parser.add_argument(
+            "--embed_dim", default=1280, type=int, metavar="N", help="embedding dimension"
+        )
+        parser.add_argument(
+            "--logit_bias", action="store_true", help="whether to apply bias to logits"
+        )
+        parser.add_argument(
+            "--ffn_embed_dim",
+            default=5120,
+            type=int,
+            metavar="N",
+            help="embedding dimension for FFN",
+        )
+        parser.add_argument(
+            "--attention_heads",
+            default=20,
+            type=int,
+            metavar="N",
+            help="number of attention heads",
+        )
+
+    def __init__(self, args, alphabet):
+        super().__init__()
+        self.args = args
+        self.alphabet_size = len(alphabet)
+        self.padding_idx = alphabet.padding_idx
+        self.mask_idx = alphabet.mask_idx
+        self.cls_idx = alphabet.cls_idx
+        self.eos_idx = alphabet.eos_idx
+        self.prepend_bos = alphabet.prepend_bos
+        self.append_eos = alphabet.append_eos
+        self.emb_layer_norm_before = getattr(self.args, "emb_layer_norm_before", False)
+        if self.args.arch == "roberta_large":
+            self.model_version = "ESM-1b"
+            self._init_submodules_esm1b()
+        else:
+            self.model_version = "ESM-1"
+            self._init_submodules_esm1()
+
+    def _init_submodules_common(self):
+        self.embed_tokens = nn.Embedding(
+            self.alphabet_size, self.args.embed_dim, padding_idx=self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                TransformerLayer(
+                    self.args.embed_dim,
+                    self.args.ffn_embed_dim,
+                    self.args.attention_heads,
+                    add_bias_kv=(self.model_version != "ESM-1b"),
+                    use_esm1b_layer_norm=(self.model_version == "ESM-1b"),
+                )
+                for _ in range(self.args.layers)
+            ]
+        )
+
+        self.contact_head = ContactPredictionHead(
+            self.args.layers * self.args.attention_heads,
+            self.prepend_bos,
+            self.append_eos,
+            eos_idx=self.eos_idx,
+        )
+
+    def _init_submodules_esm1b(self):
+        self._init_submodules_common()
+        self.embed_scale = 1
+        self.embed_positions = LearnedPositionalEmbedding(
+            self.args.max_positions, self.args.embed_dim, self.padding_idx
+        )
+        self.emb_layer_norm_before = (
+            ESM1bLayerNorm(self.args.embed_dim) if self.emb_layer_norm_before else None
+        )
+        self.emb_layer_norm_after = ESM1bLayerNorm(self.args.embed_dim)
+        self.lm_head = RobertaLMHead(
+            embed_dim=self.args.embed_dim,
+            output_dim=self.alphabet_size,
+            weight=self.embed_tokens.weight,
+        )
+
+    def _init_submodules_esm1(self):
+        self._init_submodules_common()
+        self.embed_scale = math.sqrt(self.args.embed_dim)
+        self.embed_positions = SinusoidalPositionalEmbedding(self.args.embed_dim, self.padding_idx)
+        self.embed_out = nn.Parameter(torch.zeros((self.alphabet_size, self.args.embed_dim)))
+        self.embed_out_bias = None
+        if self.args.final_bias:
+            self.embed_out_bias = nn.Parameter(torch.zeros(self.alphabet_size))
+
+    def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False):
+        if return_contacts:
+            need_head_weights = True
+
+        assert tokens.ndim == 2
+        padding_mask = tokens.eq(self.padding_idx)  # B, T
+
+        x = self.embed_scale * self.embed_tokens(tokens)
+
+        if getattr(self.args, "token_dropout", False):
+            x.masked_fill_((tokens == self.mask_idx).unsqueeze(-1), 0.0)
+            # x: B x T x C
+            mask_ratio_train = 0.15 * 0.8
+            src_lengths = (~padding_mask).sum(-1)
+            mask_ratio_observed = (tokens == self.mask_idx).sum(-1).float() / src_lengths
+            x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]
+
+        x = x + self.embed_positions(tokens)
+
+        if self.model_version == "ESM-1b":
+            if self.emb_layer_norm_before:
+                x = self.emb_layer_norm_before(x)
+            if padding_mask is not None:
+                x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))
+
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+        if 0 in repr_layers:
+            hidden_representations[0] = x
+
+        if need_head_weights:
+            attn_weights = []
+
+        # (B, T, E) => (T, B, E)
+        x = x.transpose(0, 1)
+
+        if not padding_mask.any():
+            padding_mask = None
+
+        for layer_idx, layer in enumerate(self.layers):
+            x, attn = layer(
+                x, self_attn_padding_mask=padding_mask, need_head_weights=need_head_weights
+            )
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x.transpose(0, 1)
+            if need_head_weights:
+                # (H, B, T, T) => (B, H, T, T)
+                attn_weights.append(attn.transpose(1, 0))
+
+        if self.model_version == "ESM-1b":
+            x = self.emb_layer_norm_after(x)
+            x = x.transpose(0, 1)  # (T, B, E) => (B, T, E)
+
+            # last hidden representation should have layer norm applied
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x
+            x = self.lm_head(x)
+        else:
+            x = F.linear(x, self.embed_out, bias=self.embed_out_bias)
+            x = x.transpose(0, 1)  # (T, B, E) => (B, T, E)
+
+        result = {"logits": x, "representations": hidden_representations}
+        if need_head_weights:
+            # attentions: B x L x H x T x T
+            attentions = torch.stack(attn_weights, 1)
+            if self.model_version == "ESM-1":
+                # ESM-1 models have an additional null-token for attention, which we remove
+                attentions = attentions[..., :-1]
+            if padding_mask is not None:
+                attention_mask = 1 - padding_mask.type_as(attentions)
+                attention_mask = attention_mask.unsqueeze(1) * attention_mask.unsqueeze(2)
+                attentions = attentions * attention_mask[:, None, None, :, :]
+            result["attentions"] = attentions
+            if return_contacts:
+                contacts = self.contact_head(tokens, attentions)
+                result["contacts"] = contacts
+
+        return result
+
+    def predict_contacts(self, tokens):
+        return self(tokens, return_contacts=True)["contacts"]
+
+    @property
+    def num_layers(self):
+        return self.args.layers

model/esm/esm2.py

@@ -0,0 +1,152 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Union
+import torch
+import torch.nn as nn
+
+from onescience.datapipes.esm import Alphabet
+from onescience.modules.esm import (
+    ContactPredictionHead,
+    ESM1bLayerNorm,
+    RobertaLMHead,
+    TransformerLayer,
+)
+
+
+class ESM2(nn.Module):
+    def __init__(
+        self,
+        num_layers: int = 33,
+        embed_dim: int = 1280,
+        attention_heads: int = 20,
+        alphabet: Union[Alphabet, str] = "ESM-1b",
+        token_dropout: bool = True,
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        self.embed_dim = embed_dim
+        self.attention_heads = attention_heads
+        if not isinstance(alphabet, Alphabet):
+            alphabet = Alphabet.from_architecture(alphabet)
+        self.alphabet = alphabet
+        self.alphabet_size = len(alphabet)
+        self.padding_idx = alphabet.padding_idx
+        self.mask_idx = alphabet.mask_idx
+        self.cls_idx = alphabet.cls_idx
+        self.eos_idx = alphabet.eos_idx
+        self.prepend_bos = alphabet.prepend_bos
+        self.append_eos = alphabet.append_eos
+        self.token_dropout = token_dropout
+
+        self._init_submodules()
+
+    def _init_submodules(self):
+        self.embed_scale = 1
+        self.embed_tokens = nn.Embedding(
+            self.alphabet_size,
+            self.embed_dim,
+            padding_idx=self.padding_idx,
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                TransformerLayer(
+                    self.embed_dim,
+                    4 * self.embed_dim,
+                    self.attention_heads,
+                    add_bias_kv=False,
+                    use_esm1b_layer_norm=True,
+                    use_rotary_embeddings=True,
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+
+        self.contact_head = ContactPredictionHead(
+            self.num_layers * self.attention_heads,
+            self.prepend_bos,
+            self.append_eos,
+            eos_idx=self.eos_idx,
+        )
+        self.emb_layer_norm_after = ESM1bLayerNorm(self.embed_dim)
+
+        self.lm_head = RobertaLMHead(
+            embed_dim=self.embed_dim,
+            output_dim=self.alphabet_size,
+            weight=self.embed_tokens.weight,
+        )
+
+    def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False):
+        if return_contacts:
+            need_head_weights = True
+
+        assert tokens.ndim == 2
+        padding_mask = tokens.eq(self.padding_idx)  # B, T
+
+        x = self.embed_scale * self.embed_tokens(tokens)
+
+        if self.token_dropout:
+            x.masked_fill_((tokens == self.mask_idx).unsqueeze(-1), 0.0)
+            # x: B x T x C
+            mask_ratio_train = 0.15 * 0.8
+            src_lengths = (~padding_mask).sum(-1)
+            mask_ratio_observed = (tokens == self.mask_idx).sum(-1).to(x.dtype) / src_lengths
+            x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]
+
+        if padding_mask is not None:
+            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))
+
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+        if 0 in repr_layers:
+            hidden_representations[0] = x
+
+        if need_head_weights:
+            attn_weights = []
+
+        # (B, T, E) => (T, B, E)
+        x = x.transpose(0, 1)
+
+        if not padding_mask.any():
+            padding_mask = None
+
+        for layer_idx, layer in enumerate(self.layers):
+            x, attn = layer(
+                x,
+                self_attn_padding_mask=padding_mask,
+                need_head_weights=need_head_weights,
+            )
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x.transpose(0, 1)
+            if need_head_weights:
+                # (H, B, T, T) => (B, H, T, T)
+                attn_weights.append(attn.transpose(1, 0))
+
+        x = self.emb_layer_norm_after(x)
+        x = x.transpose(0, 1)  # (T, B, E) => (B, T, E)
+
+        # last hidden representation should have layer norm applied
+        if (layer_idx + 1) in repr_layers:
+            hidden_representations[layer_idx + 1] = x
+        x = self.lm_head(x)
+
+        result = {"logits": x, "representations": hidden_representations}
+        if need_head_weights:
+            # attentions: B x L x H x T x T
+            attentions = torch.stack(attn_weights, 1)
+            if padding_mask is not None:
+                attention_mask = 1 - padding_mask.type_as(attentions)
+                attention_mask = attention_mask.unsqueeze(1) * attention_mask.unsqueeze(2)
+                attentions = attentions * attention_mask[:, None, None, :, :]
+            result["attentions"] = attentions
+            if return_contacts:
+                contacts = self.contact_head(tokens, attentions)
+                result["contacts"] = contacts
+
+        return result
+
+    def predict_contacts(self, tokens):
+        return self(tokens, return_contacts=True)["contacts"]

model/esm/esmfold/v1/categorical_mixture.py

@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+
+
+class CategoricalMixture:
+    def __init__(self, param, bins=50, start=0, end=1):
+        # All tensors are of shape ..., bins.
+        self.logits = param
+        bins = torch.linspace(
+            start, end, bins + 1, device=self.logits.device, dtype=self.logits.dtype
+        )
+        self.v_bins = (bins[:-1] + bins[1:]) / 2
+
+    def log_prob(self, true):
+        # Shapes are:
+        #     self.probs: ... x bins
+        #     true      : ...
+        true_index = (
+            (
+                true.unsqueeze(-1)
+                - self.v_bins[
+                    [
+                        None,
+                    ]
+                    * true.ndim
+                ]
+            )
+            .abs()
+            .argmin(-1)
+        )
+        nll = self.logits.log_softmax(-1)
+        return torch.take_along_dim(nll, true_index.unsqueeze(-1), dim=-1).squeeze(-1)
+
+    def mean(self):
+        return (self.logits.softmax(-1) @ self.v_bins.unsqueeze(1)).squeeze(-1)
+
+
+def categorical_lddt(logits, bins=50):
+    # Logits are ..., 37, bins.
+    return CategoricalMixture(logits, bins=bins).mean()

model/esm/esmfold/v1/misc.py

@@ -0,0 +1,309 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import typing as T
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import nn
+from onescience.utils.openfold.np import residue_constants
+from onescience.utils.openfold.np.protein import Protein as OFProtein
+from onescience.utils.openfold.np.protein import to_pdb
+from onescience.utils.openfold.feats import atom14_to_atom37
+
+
+def encode_sequence(
+    seq: str,
+    residue_index_offset: T.Optional[int] = 512,
+    chain_linker: T.Optional[str] = "G" * 25,
+) -> T.Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if chain_linker is None:
+        chain_linker = ""
+    if residue_index_offset is None:
+        residue_index_offset = 0
+
+    chains = seq.split(":")
+    seq = chain_linker.join(chains)
+
+    unk_idx = residue_constants.restype_order_with_x["X"]
+    encoded = torch.tensor(
+        [residue_constants.restype_order_with_x.get(aa, unk_idx) for aa in seq]
+    )
+    residx = torch.arange(len(encoded))
+
+    if residue_index_offset > 0:
+        start = 0
+        for i, chain in enumerate(chains):
+            residx[start : start + len(chain) + len(chain_linker)] += (
+                i * residue_index_offset
+            )
+            start += len(chain) + len(chain_linker)
+
+    linker_mask = torch.ones_like(encoded, dtype=torch.float32)
+    chain_index = []
+    offset = 0
+    for i, chain in enumerate(chains):
+        if i > 0:
+            chain_index.extend([i - 1] * len(chain_linker))
+        chain_index.extend([i] * len(chain))
+        offset += len(chain)
+        linker_mask[offset : offset + len(chain_linker)] = 0
+        offset += len(chain_linker)
+
+    chain_index = torch.tensor(chain_index, dtype=torch.int64)
+
+    return encoded, residx, linker_mask, chain_index
+
+
+def batch_encode_sequences(
+    sequences: T.Sequence[str],
+    residue_index_offset: T.Optional[int] = 512,
+    chain_linker: T.Optional[str] = "G" * 25,
+) -> T.Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    aatype_list = []
+    residx_list = []
+    linker_mask_list = []
+    chain_index_list = []
+    for seq in sequences:
+        aatype_seq, residx_seq, linker_mask_seq, chain_index_seq = encode_sequence(
+            seq,
+            residue_index_offset=residue_index_offset,
+            chain_linker=chain_linker,
+        )
+        aatype_list.append(aatype_seq)
+        residx_list.append(residx_seq)
+        linker_mask_list.append(linker_mask_seq)
+        chain_index_list.append(chain_index_seq)
+
+    aatype = collate_dense_tensors(aatype_list)
+    mask = collate_dense_tensors(
+        [aatype.new_ones(len(aatype_seq)) for aatype_seq in aatype_list]
+    )
+    residx = collate_dense_tensors(residx_list)
+    linker_mask = collate_dense_tensors(linker_mask_list)
+    chain_index_list = collate_dense_tensors(chain_index_list, -1)
+
+    return aatype, mask, residx, linker_mask, chain_index_list
+
+
+def output_to_pdb(output: T.Dict) -> T.List[str]:
+    """Returns the pbd (file) string from the model given the model output."""
+    # atom14_to_atom37 must be called first, as it fails on latest numpy if the
+    # input is a numpy array. It will work if the input is a torch tensor.
+    final_atom_positions = atom14_to_atom37(output["positions"][-1], output)
+    output = {k: v.to("cpu").numpy() for k, v in output.items()}
+    final_atom_positions = final_atom_positions.cpu().numpy()
+    final_atom_mask = output["atom37_atom_exists"]
+    pdbs = []
+    for i in range(output["aatype"].shape[0]):
+        aa = output["aatype"][i]
+        pred_pos = final_atom_positions[i]
+        mask = final_atom_mask[i]
+        resid = output["residue_index"][i] + 1
+        pred = OFProtein(
+            aatype=aa,
+            atom_positions=pred_pos,
+            atom_mask=mask,
+            residue_index=resid,
+            b_factors=output["plddt"][i],
+            chain_index=output["chain_index"][i] if "chain_index" in output else None,
+        )
+        pdbs.append(to_pdb(pred))
+    return pdbs
+
+
+def collate_dense_tensors(
+    samples: T.List[torch.Tensor], pad_v: float = 0
+) -> torch.Tensor:
+    """
+    Takes a list of tensors with the following dimensions:
+        [(d_11,       ...,           d_1K),
+         (d_21,       ...,           d_2K),
+         ...,
+         (d_N1,       ...,           d_NK)]
+    and stack + pads them into a single tensor of:
+    (N, max_i=1,N { d_i1 }, ..., max_i=1,N {diK})
+    """
+    if len(samples) == 0:
+        return torch.Tensor()
+    if len(set(x.dim() for x in samples)) != 1:
+        raise RuntimeError(
+            f"Samples has varying dimensions: {[x.dim() for x in samples]}"
+        )
+    (device,) = tuple(set(x.device for x in samples))  # assumes all on same device
+    max_shape = [max(lst) for lst in zip(*[x.shape for x in samples])]
+    result = torch.empty(
+        len(samples), *max_shape, dtype=samples[0].dtype, device=device
+    )
+    result.fill_(pad_v)
+    for i in range(len(samples)):
+        result_i = result[i]
+        t = samples[i]
+        result_i[tuple(slice(0, k) for k in t.shape)] = t
+    return result
+
+
+class Attention(nn.Module):
+    def __init__(self, embed_dim, num_heads, head_width, gated=False):
+        super().__init__()
+        assert embed_dim == num_heads * head_width
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_width = head_width
+
+        self.proj = nn.Linear(embed_dim, embed_dim * 3, bias=False)
+        self.o_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.gated = gated
+        if gated:
+            self.g_proj = nn.Linear(embed_dim, embed_dim)
+            torch.nn.init.zeros_(self.g_proj.weight)
+            torch.nn.init.ones_(self.g_proj.bias)
+
+        self.rescale_factor = self.head_width**-0.5
+
+        torch.nn.init.zeros_(self.o_proj.bias)
+
+    def forward(self, x, mask=None, bias=None, indices=None):
+        """
+        Basic self attention with optional mask and external pairwise bias.
+        To handle sequences of different lengths, use mask.
+
+        Inputs:
+          x: batch of input sequneces (.. x L x C)
+          mask: batch of boolean masks where 1=valid, 0=padding position (.. x L_k). optional.
+          bias: batch of scalar pairwise attention biases (.. x Lq x Lk x num_heads). optional.
+
+        Outputs:
+          sequence projection (B x L x embed_dim), attention maps (B x L x L x num_heads)
+        """
+
+        t = rearrange(self.proj(x), "... l (h c) -> ... h l c", h=self.num_heads)
+        q, k, v = t.chunk(3, dim=-1)
+
+        q = self.rescale_factor * q
+        a = torch.einsum("...qc,...kc->...qk", q, k)
+
+        # Add external attention bias.
+        if bias is not None:
+            a = a + rearrange(bias, "... lq lk h -> ... h lq lk")
+
+        # Do not attend to padding tokens.
+        if mask is not None:
+            mask = repeat(
+                mask, "... lk -> ... h lq lk", h=self.num_heads, lq=q.shape[-2]
+            )
+            a = a.masked_fill(mask == False, -np.inf)
+
+        a = F.softmax(a, dim=-1)
+
+        y = torch.einsum("...hqk,...hkc->...qhc", a, v)
+        y = rearrange(y, "... h c -> ... (h c)", h=self.num_heads)
+
+        if self.gated:
+            y = self.g_proj(x).sigmoid() * y
+        y = self.o_proj(y)
+
+        return y, rearrange(a, "... lq lk h -> ... h lq lk")
+
+
+class Dropout(nn.Module):
+    """
+    Implementation of dropout with the ability to share the dropout mask
+    along a particular dimension.
+    """
+
+    def __init__(self, r: float, batch_dim: T.Union[int, T.List[int]]):
+        super(Dropout, self).__init__()
+
+        self.r = r
+        if type(batch_dim) == int:
+            batch_dim = [batch_dim]
+        self.batch_dim = batch_dim
+        self.dropout = nn.Dropout(self.r)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = list(x.shape)
+        if self.batch_dim is not None:
+            for bd in self.batch_dim:
+                shape[bd] = 1
+        return x * self.dropout(x.new_ones(shape))
+
+
+class SequenceToPair(nn.Module):
+    def __init__(self, sequence_state_dim, inner_dim, pairwise_state_dim):
+        super().__init__()
+
+        self.layernorm = nn.LayerNorm(sequence_state_dim)
+        self.proj = nn.Linear(sequence_state_dim, inner_dim * 2, bias=True)
+        self.o_proj = nn.Linear(2 * inner_dim, pairwise_state_dim, bias=True)
+
+        torch.nn.init.zeros_(self.proj.bias)
+        torch.nn.init.zeros_(self.o_proj.bias)
+
+    def forward(self, sequence_state):
+        """
+        Inputs:
+          sequence_state: B x L x sequence_state_dim
+
+        Output:
+          pairwise_state: B x L x L x pairwise_state_dim
+
+        Intermediate state:
+          B x L x L x 2*inner_dim
+        """
+
+        assert len(sequence_state.shape) == 3
+
+        s = self.layernorm(sequence_state)
+        s = self.proj(s)
+        q, k = s.chunk(2, dim=-1)
+
+        prod = q[:, None, :, :] * k[:, :, None, :]
+        diff = q[:, None, :, :] - k[:, :, None, :]
+
+        x = torch.cat([prod, diff], dim=-1)
+        x = self.o_proj(x)
+
+        return x
+
+
+class PairToSequence(nn.Module):
+    def __init__(self, pairwise_state_dim, num_heads):
+        super().__init__()
+
+        self.layernorm = nn.LayerNorm(pairwise_state_dim)
+        self.linear = nn.Linear(pairwise_state_dim, num_heads, bias=False)
+
+    def forward(self, pairwise_state):
+        """
+        Inputs:
+          pairwise_state: B x L x L x pairwise_state_dim
+
+        Output:
+          pairwise_bias: B x L x L x num_heads
+        """
+        assert len(pairwise_state.shape) == 4
+        z = self.layernorm(pairwise_state)
+        pairwise_bias = self.linear(z)
+        return pairwise_bias
+
+
+class ResidueMLP(nn.Module):
+    def __init__(self, embed_dim, inner_dim, norm=nn.LayerNorm, dropout=0):
+        super().__init__()
+
+        self.mlp = nn.Sequential(
+            norm(embed_dim),
+            nn.Linear(embed_dim, inner_dim),
+            nn.ReLU(),
+            nn.Linear(inner_dim, embed_dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return x + self.mlp(x)

model/esm/esmfold/v1/pretrained.py

@@ -0,0 +1,197 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from pathlib import Path
+
+import torch
+
+from model.esm.esmfold.v1.esmfold import ESMFold
+
+def _upgrade_state_dict(model_state):
+    """Map legacy ESMFold checkpoint keys to the current OpenFold modules."""
+    point_projection_prefixes = (
+        "trunk.structure_module.ipa.linear_q_points",
+        "trunk.structure_module.ipa.linear_kv_points",
+        "trunk.structure_module.ipa.linear_k_points",
+        "trunk.structure_module.ipa.linear_v_points",
+    )
+    for prefix in point_projection_prefixes:
+        for suffix in ("weight", "bias"):
+            old_key = f"{prefix}.{suffix}"
+            new_key = f"{prefix}.linear.{suffix}"
+            if old_key in model_state and new_key not in model_state:
+                model_state[new_key] = model_state.pop(old_key)
+    return model_state
+
+def _load_model(model_name):
+    if model_name.endswith(".pt"):  # local, treat as filepath
+        model_path = Path(model_name)
+        model_data = torch.load(str(model_path), map_location="cpu")
+    else:  # load from hub
+        url = f"https://dl.fbaipublicfiles.com/fair-esm/models/{model_name}.pt"
+        model_data = torch.hub.load_state_dict_from_url(url, progress=False, map_location="cpu")
+
+    cfg = model_data["cfg"]["model"]
+    model_state = _upgrade_state_dict(model_data["model"])
+    # model_state = model_data["model"]
+    model = ESMFold(esmfold_config=cfg)
+
+    expected_keys = set(model.state_dict().keys())
+    found_keys = set(model_state.keys())
+
+    missing_essential_keys = []
+    for missing_key in expected_keys - found_keys:
+        if not missing_key.startswith("esm."):
+            missing_essential_keys.append(missing_key)
+
+    if missing_essential_keys:
+        raise RuntimeError(f"Keys '{', '.join(missing_essential_keys)}' are missing.")
+
+    model.load_state_dict(model_state, strict=False)
+
+    return model
+
+
+def esmfold_v0():
+    """
+    ESMFold v0 model with 3B ESM-2, 48 folding blocks.
+    This version was used for the paper (Lin et al, 2022). It was trained
+    on all PDB chains until 2020-05, to ensure temporal holdout with CASP14
+    and the CAMEO validation and test set reported there.
+    """
+    return _load_model("esmfold_3B_v0")
+
+
+def esmfold_v1():
+    """
+    ESMFold v1 model using 3B ESM-2, 48 folding blocks.
+    ESMFold provides fast high accuracy atomic level structure prediction
+    directly from the individual sequence of a protein. ESMFold uses the ESM2
+    protein language model to extract meaningful representations from the
+    protein sequence.
+    """
+    return _load_model("esmfold_3B_v1")
+
+
+def esmfold_structure_module_only_8M():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_8M")
+
+
+def esmfold_structure_module_only_8M_270K():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_8M_270K")
+
+
+def esmfold_structure_module_only_35M():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_35M")
+
+
+def esmfold_structure_module_only_35M_270K():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_35M_270K")
+
+
+def esmfold_structure_module_only_150M():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_150M")
+
+
+def esmfold_structure_module_only_150M_270K():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_150M_270K")
+
+
+def esmfold_structure_module_only_650M():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_650M")
+
+
+def esmfold_structure_module_only_650M_270K():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_650M_270K")
+
+
+def esmfold_structure_module_only_3B():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_3B")
+
+
+def esmfold_structure_module_only_3B_270K():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_3B_270K")
+
+
+def esmfold_structure_module_only_15B():
+    """
+    ESMFold baseline model using 15B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    The 15B parameter ESM-2 was not trained out to 500K updates
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_15B")

model/esm/inverse_folding/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from . import gvp_transformer
+from . import util
+from . import multichain_util

model/esm/inverse_folding/features.py

@@ -0,0 +1,352 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# 
+# Portions of this file were adapted from the open source code for the following
+# two papers:
+#
+#   Ingraham, J., Garg, V., Barzilay, R., & Jaakkola, T. (2019). Generative
+#   models for graph-based protein design. Advances in Neural Information
+#   Processing Systems, 32.
+#
+#   Jing, B., Eismann, S., Suriana, P., Townshend, R. J. L., & Dror, R. (2020).
+#   Learning from Protein Structure with Geometric Vector Perceptrons. In
+#   International Conference on Learning Representations.
+#
+# MIT License
+# 
+# Copyright (c) 2020 Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael Townshend, Ron Dror
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# 
+# ================================================================
+# The below license applies to the portions of the code (parts of 
+# src/datasets.py and src/models.py) adapted from Ingraham, et al.
+# ================================================================
+# 
+# MIT License
+# 
+# Copyright (c) 2019 John Ingraham, Vikas Garg, Regina Barzilay, Tommi Jaakkola
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .gvp_utils import flatten_graph
+from .gvp_modules import GVP, LayerNorm
+from .util import normalize, norm, nan_to_num, rbf
+
+
+class GVPInputFeaturizer(nn.Module):
+
+    @staticmethod
+    def get_node_features(coords, coord_mask, with_coord_mask=True):
+        # scalar features
+        node_scalar_features = GVPInputFeaturizer._dihedrals(coords)
+        if with_coord_mask:
+            node_scalar_features = torch.cat([
+                node_scalar_features,
+                coord_mask.float().unsqueeze(-1)
+            ], dim=-1) 
+        # vector features
+        X_ca = coords[:, :, 1]
+        orientations = GVPInputFeaturizer._orientations(X_ca)
+        sidechains = GVPInputFeaturizer._sidechains(coords)
+        node_vector_features = torch.cat([orientations, sidechains.unsqueeze(-2)], dim=-2)
+        return node_scalar_features, node_vector_features
+
+    @staticmethod
+    def _orientations(X):
+        forward = normalize(X[:, 1:] - X[:, :-1])
+        backward = normalize(X[:, :-1] - X[:, 1:])
+        forward = F.pad(forward, [0, 0, 0, 1])
+        backward = F.pad(backward, [0, 0, 1, 0])
+        return torch.cat([forward.unsqueeze(-2), backward.unsqueeze(-2)], -2)
+    
+    @staticmethod
+    def _sidechains(X):
+        n, origin, c = X[:, :, 0], X[:, :, 1], X[:, :, 2]
+        c, n = normalize(c - origin), normalize(n - origin)
+        bisector = normalize(c + n)
+        perp = normalize(torch.cross(c, n, dim=-1))
+        vec = -bisector * math.sqrt(1 / 3) - perp * math.sqrt(2 / 3)
+        return vec 
+
+    @staticmethod
+    def _dihedrals(X, eps=1e-7):
+        X = torch.flatten(X[:, :, :3], 1, 2)
+        bsz = X.shape[0]
+        dX = X[:, 1:] - X[:, :-1]
+        U = normalize(dX, dim=-1)
+        u_2 = U[:, :-2]
+        u_1 = U[:, 1:-1]
+        u_0 = U[:, 2:]
+    
+        # Backbone normals
+        n_2 = normalize(torch.cross(u_2, u_1, dim=-1), dim=-1)
+        n_1 = normalize(torch.cross(u_1, u_0, dim=-1), dim=-1)
+    
+        # Angle between normals
+        cosD = torch.sum(n_2 * n_1, -1)
+        cosD = torch.clamp(cosD, -1 + eps, 1 - eps)
+        D = torch.sign(torch.sum(u_2 * n_1, -1)) * torch.acos(cosD)
+    
+        # This scheme will remove phi[0], psi[-1], omega[-1]
+        D = F.pad(D, [1, 2]) 
+        D = torch.reshape(D, [bsz, -1, 3])
+        # Lift angle representations to the circle
+        D_features = torch.cat([torch.cos(D), torch.sin(D)], -1)
+        return D_features
+
+    @staticmethod
+    def _positional_embeddings(edge_index, 
+                               num_embeddings=None,
+                               num_positional_embeddings=16,
+                               period_range=[2, 1000]):
+        # From https://github.com/jingraham/neurips19-graph-protein-design
+        num_embeddings = num_embeddings or num_positional_embeddings
+        d = edge_index[0] - edge_index[1]
+     
+        frequency = torch.exp(
+            torch.arange(0, num_embeddings, 2, dtype=torch.float32,
+                device=edge_index.device)
+            * -(np.log(10000.0) / num_embeddings)
+        )
+        angles = d.unsqueeze(-1) * frequency
+        E = torch.cat((torch.cos(angles), torch.sin(angles)), -1)
+        return E
+
+    @staticmethod
+    def _dist(X, coord_mask, padding_mask, top_k_neighbors, eps=1e-8):
+        """ Pairwise euclidean distances """
+        bsz, maxlen = X.size(0), X.size(1)
+        coord_mask_2D = torch.unsqueeze(coord_mask,1) * torch.unsqueeze(coord_mask,2)
+        residue_mask = ~padding_mask
+        residue_mask_2D = torch.unsqueeze(residue_mask,1) * torch.unsqueeze(residue_mask,2)
+        dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)
+        D = coord_mask_2D * norm(dX, dim=-1)
+    
+        # sorting preference: first those with coords, then among the residues that
+        # exist but are masked use distance in sequence as tie breaker, and then the
+        # residues that came from padding are last
+        seqpos = torch.arange(maxlen, device=X.device)
+        Dseq = torch.abs(seqpos.unsqueeze(1) - seqpos.unsqueeze(0)).repeat(bsz, 1, 1)
+        D_adjust = nan_to_num(D) + (~coord_mask_2D) * (1e8 + Dseq*1e6) + (
+            ~residue_mask_2D) * (1e10)
+    
+        if top_k_neighbors == -1:
+            D_neighbors = D_adjust
+            E_idx = seqpos.repeat(
+                    *D_neighbors.shape[:-1], 1)
+        else:
+            # Identify k nearest neighbors (including self)
+            k = min(top_k_neighbors, X.size(1))
+            D_neighbors, E_idx = torch.topk(D_adjust, k, dim=-1, largest=False)
+    
+        coord_mask_neighbors = (D_neighbors < 5e7)
+        residue_mask_neighbors = (D_neighbors < 5e9)
+        return D_neighbors, E_idx, coord_mask_neighbors, residue_mask_neighbors
+
+
+class Normalize(nn.Module):
+    def __init__(self, features, epsilon=1e-6):
+        super(Normalize, self).__init__()
+        self.gain = nn.Parameter(torch.ones(features))
+        self.bias = nn.Parameter(torch.zeros(features))
+        self.epsilon = epsilon
+
+    def forward(self, x, dim=-1):
+        mu = x.mean(dim, keepdim=True)
+        sigma = torch.sqrt(x.var(dim, keepdim=True) + self.epsilon)
+        gain = self.gain
+        bias = self.bias
+        # Reshape
+        if dim != -1:
+            shape = [1] * len(mu.size())
+            shape[dim] = self.gain.size()[0]
+            gain = gain.view(shape)
+            bias = bias.view(shape)
+        return gain * (x - mu) / (sigma + self.epsilon) + bias
+
+
+class DihedralFeatures(nn.Module):
+    def __init__(self, node_embed_dim):
+        """ Embed dihedral angle features. """
+        super(DihedralFeatures, self).__init__()
+        # 3 dihedral angles; sin and cos of each angle
+        node_in = 6
+        # Normalization and embedding
+        self.node_embedding = nn.Linear(node_in,  node_embed_dim, bias=True)
+        self.norm_nodes = Normalize(node_embed_dim)
+
+    def forward(self, X):
+        """ Featurize coordinates as an attributed graph """
+        V = self._dihedrals(X)
+        V = self.node_embedding(V)
+        V = self.norm_nodes(V)
+        return V
+
+    @staticmethod
+    def _dihedrals(X, eps=1e-7, return_angles=False):
+        # First 3 coordinates are N, CA, C
+        X = X[:,:,:3,:].reshape(X.shape[0], 3*X.shape[1], 3)
+
+        # Shifted slices of unit vectors
+        dX = X[:,1:,:] - X[:,:-1,:]
+        U = F.normalize(dX, dim=-1)
+        u_2 = U[:,:-2,:]
+        u_1 = U[:,1:-1,:]
+        u_0 = U[:,2:,:]
+        # Backbone normals
+        n_2 = F.normalize(torch.cross(u_2, u_1, dim=-1), dim=-1)
+        n_1 = F.normalize(torch.cross(u_1, u_0, dim=-1), dim=-1)
+
+        # Angle between normals
+        cosD = (n_2 * n_1).sum(-1)
+        cosD = torch.clamp(cosD, -1+eps, 1-eps)
+        D = torch.sign((u_2 * n_1).sum(-1)) * torch.acos(cosD)
+
+        # This scheme will remove phi[0], psi[-1], omega[-1]
+        D = F.pad(D, (1,2), 'constant', 0)
+        D = D.view((D.size(0), int(D.size(1)/3), 3))
+        phi, psi, omega = torch.unbind(D,-1)
+
+        if return_angles:
+            return phi, psi, omega
+
+        # Lift angle representations to the circle
+        D_features = torch.cat((torch.cos(D), torch.sin(D)), 2)
+        return D_features
+
+
+class GVPGraphEmbedding(GVPInputFeaturizer):
+
+    def __init__(self, args):
+        super().__init__()
+        self.top_k_neighbors = args.top_k_neighbors
+        self.num_positional_embeddings = 16
+        self.remove_edges_without_coords = True
+        node_input_dim = (7, 3)
+        edge_input_dim = (34, 1)
+        node_hidden_dim = (args.node_hidden_dim_scalar,
+                args.node_hidden_dim_vector)
+        edge_hidden_dim = (args.edge_hidden_dim_scalar,
+                args.edge_hidden_dim_vector)
+        self.embed_node = nn.Sequential(
+            GVP(node_input_dim, node_hidden_dim, activations=(None, None)),
+            LayerNorm(node_hidden_dim, eps=1e-4)
+        )
+        self.embed_edge = nn.Sequential(
+            GVP(edge_input_dim, edge_hidden_dim, activations=(None, None)),
+            LayerNorm(edge_hidden_dim, eps=1e-4)
+        )
+        self.embed_confidence = nn.Linear(16, args.node_hidden_dim_scalar)
+
+    def forward(self, coords, coord_mask, padding_mask, confidence):
+        with torch.no_grad():
+            node_features = self.get_node_features(coords, coord_mask)
+            edge_features, edge_index = self.get_edge_features(
+                coords, coord_mask, padding_mask)
+        node_embeddings_scalar, node_embeddings_vector = self.embed_node(node_features)
+        edge_embeddings = self.embed_edge(edge_features)
+
+        rbf_rep = rbf(confidence, 0., 1.)
+        node_embeddings = (
+            node_embeddings_scalar + self.embed_confidence(rbf_rep),
+            node_embeddings_vector
+        )
+
+        node_embeddings, edge_embeddings, edge_index = flatten_graph(
+            node_embeddings, edge_embeddings, edge_index)
+        return node_embeddings, edge_embeddings, edge_index
+
+    def get_edge_features(self, coords, coord_mask, padding_mask):
+        X_ca = coords[:, :, 1]
+        # Get distances to the top k neighbors
+        E_dist, E_idx, E_coord_mask, E_residue_mask = GVPInputFeaturizer._dist(
+                X_ca, coord_mask, padding_mask, self.top_k_neighbors)
+        # Flatten the graph to be batch size 1 for torch_geometric package 
+        dest = E_idx
+        B, L, k = E_idx.shape[:3]
+        src = torch.arange(L, device=E_idx.device).view([1, L, 1]).expand(B, L, k)
+        # After flattening, [2, B, E]
+        edge_index = torch.stack([src, dest], dim=0).flatten(2, 3)
+        # After flattening, [B, E]
+        E_dist = E_dist.flatten(1, 2)
+        E_coord_mask = E_coord_mask.flatten(1, 2).unsqueeze(-1)
+        E_residue_mask = E_residue_mask.flatten(1, 2)
+        # Calculate relative positional embeddings and distance RBF 
+        pos_embeddings = GVPInputFeaturizer._positional_embeddings(
+            edge_index,
+            num_positional_embeddings=self.num_positional_embeddings,
+        )
+        D_rbf = rbf(E_dist, 0., 20.)
+        # Calculate relative orientation 
+        X_src = X_ca.unsqueeze(2).expand(-1, -1, k, -1).flatten(1, 2)
+        X_dest = torch.gather(
+            X_ca,
+            1,
+            edge_index[1, :, :].unsqueeze(-1).expand([B, L*k, 3])
+        )
+        coord_mask_src = coord_mask.unsqueeze(2).expand(-1, -1, k).flatten(1, 2)
+        coord_mask_dest = torch.gather(
+            coord_mask,
+            1,
+            edge_index[1, :, :].expand([B, L*k])
+        )
+        E_vectors = X_src - X_dest
+        # For the ones without coordinates, substitute in the average vector
+        E_vector_mean = torch.sum(E_vectors * E_coord_mask, dim=1,
+                keepdims=True) / torch.sum(E_coord_mask, dim=1, keepdims=True)
+        E_vectors = E_vectors * E_coord_mask + E_vector_mean * ~(E_coord_mask)
+        # Normalize and remove nans 
+        edge_s = torch.cat([D_rbf, pos_embeddings], dim=-1)
+        edge_v = normalize(E_vectors).unsqueeze(-2)
+        edge_s, edge_v = map(nan_to_num, (edge_s, edge_v))
+        # Also add indications of whether the coordinates are present 
+        edge_s = torch.cat([
+            edge_s,
+            (~coord_mask_src).float().unsqueeze(-1),
+            (~coord_mask_dest).float().unsqueeze(-1),
+        ], dim=-1)
+        edge_index[:, ~E_residue_mask] = -1
+        if self.remove_edges_without_coords:
+            edge_index[:, ~E_coord_mask.squeeze(-1)] = -1
+        return (edge_s, edge_v), edge_index.transpose(0, 1) 

model/esm/inverse_folding/gvp_encoder.py

@@ -0,0 +1,56 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from argparse import Namespace
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .features import GVPGraphEmbedding
+from .gvp_modules import GVPConvLayer, LayerNorm 
+from .gvp_utils import unflatten_graph
+
+
+
+class GVPEncoder(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.embed_graph = GVPGraphEmbedding(args)
+
+        node_hidden_dim = (args.node_hidden_dim_scalar,
+                args.node_hidden_dim_vector)
+        edge_hidden_dim = (args.edge_hidden_dim_scalar,
+                args.edge_hidden_dim_vector)
+        
+        conv_activations = (F.relu, torch.sigmoid)
+        self.encoder_layers = nn.ModuleList(
+                GVPConvLayer(
+                    node_hidden_dim,
+                    edge_hidden_dim,
+                    drop_rate=args.dropout,
+                    vector_gate=True,
+                    attention_heads=0,
+                    n_message=3,
+                    conv_activations=conv_activations,
+                    n_edge_gvps=0,
+                    eps=1e-4,
+                    layernorm=True,
+                ) 
+            for i in range(args.num_encoder_layers)
+        )
+
+    def forward(self, coords, coord_mask, padding_mask, confidence):
+        node_embeddings, edge_embeddings, edge_index = self.embed_graph(
+                coords, coord_mask, padding_mask, confidence)
+        
+        for i, layer in enumerate(self.encoder_layers):
+            node_embeddings, edge_embeddings = layer(node_embeddings,
+                    edge_index, edge_embeddings)
+
+        node_embeddings = unflatten_graph(node_embeddings, coords.shape[0])
+        return node_embeddings

model/esm/inverse_folding/gvp_transformer_encoder.py

@@ -0,0 +1,184 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Contents of this file were adapted from the open source fairseq repository.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import math
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from onescience.modules.esm import SinusoidalPositionalEmbedding
+from .features import GVPInputFeaturizer, DihedralFeatures
+from .gvp_encoder import GVPEncoder
+from .transformer_layer import TransformerEncoderLayer
+from .util import nan_to_num, get_rotation_frames, rotate, rbf
+
+
+class GVPTransformerEncoder(nn.Module):
+    """
+    Transformer encoder consisting of *args.encoder.layers* layers. Each layer
+    is a :class:`TransformerEncoderLayer`.
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+        dictionary (~fairseq.data.Dictionary): encoding dictionary
+        embed_tokens (torch.nn.Embedding): input embedding
+    """
+
+    def __init__(self, args, dictionary, embed_tokens):
+        super().__init__()
+        self.args = args
+        self.dictionary = dictionary
+
+        self.dropout_module = nn.Dropout(args.dropout)
+
+        embed_dim = embed_tokens.embedding_dim
+        self.padding_idx = embed_tokens.padding_idx
+
+        self.embed_tokens = embed_tokens
+        self.embed_scale = math.sqrt(embed_dim)
+        self.embed_positions = SinusoidalPositionalEmbedding(
+            embed_dim,
+            self.padding_idx,
+        )
+        self.embed_gvp_input_features = nn.Linear(15, embed_dim)
+        self.embed_confidence = nn.Linear(16, embed_dim)
+        self.embed_dihedrals = DihedralFeatures(embed_dim)
+
+        gvp_args = argparse.Namespace()
+        for k, v in vars(args).items():
+            if k.startswith("gvp_"):
+                setattr(gvp_args, k[4:], v)
+        self.gvp_encoder = GVPEncoder(gvp_args)
+        gvp_out_dim = gvp_args.node_hidden_dim_scalar + (3 *
+                gvp_args.node_hidden_dim_vector)
+        self.embed_gvp_output = nn.Linear(gvp_out_dim, embed_dim)
+
+        self.layers = nn.ModuleList([])
+        self.layers.extend(
+            [self.build_encoder_layer(args) for i in range(args.encoder_layers)]
+        )
+        self.num_layers = len(self.layers)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+    def build_encoder_layer(self, args):
+        return TransformerEncoderLayer(args)
+
+    def forward_embedding(self, coords, padding_mask, confidence):
+        """
+        Args:
+            coords: N, CA, C backbone coordinates in shape length x 3 (atoms) x 3 
+            padding_mask: boolean Tensor (true for padding) of shape length
+            confidence: confidence scores between 0 and 1 of shape length
+        """
+        components = dict()
+        coord_mask = torch.all(torch.all(torch.isfinite(coords), dim=-1), dim=-1)
+        coords = nan_to_num(coords)
+        mask_tokens = (
+            padding_mask * self.dictionary.padding_idx + 
+            ~padding_mask * self.dictionary.get_idx("<mask>")
+        )
+        components["tokens"] = self.embed_tokens(mask_tokens) * self.embed_scale
+        components["diherals"] = self.embed_dihedrals(coords)
+
+        # GVP encoder
+        gvp_out_scalars, gvp_out_vectors = self.gvp_encoder(coords,
+                coord_mask, padding_mask, confidence)
+        R = get_rotation_frames(coords)
+        # Rotate to local rotation frame for rotation-invariance
+        gvp_out_features = torch.cat([
+            gvp_out_scalars,
+            rotate(gvp_out_vectors, R.transpose(-2, -1)).flatten(-2, -1),
+        ], dim=-1)
+        components["gvp_out"] = self.embed_gvp_output(gvp_out_features)
+
+        components["confidence"] = self.embed_confidence(
+             rbf(confidence, 0., 1.))
+
+        # In addition to GVP encoder outputs, also directly embed GVP input node
+        # features to the Transformer
+        scalar_features, vector_features = GVPInputFeaturizer.get_node_features(
+            coords, coord_mask, with_coord_mask=False)
+        features = torch.cat([
+            scalar_features,
+            rotate(vector_features, R.transpose(-2, -1)).flatten(-2, -1),
+        ], dim=-1)
+        components["gvp_input_features"] = self.embed_gvp_input_features(features)
+
+        embed = sum(components.values())
+        # for k, v in components.items():
+        #     print(k, torch.mean(v, dim=(0,1)), torch.std(v, dim=(0,1)))
+
+        x = embed
+        x = x + self.embed_positions(mask_tokens)
+        x = self.dropout_module(x)
+        return x, components 
+
+    def forward(
+        self,
+        coords,
+        encoder_padding_mask,
+        confidence,
+        return_all_hiddens: bool = False,
+    ):
+        """
+        Args:
+            coords (Tensor): backbone coordinates
+                shape batch_size x num_residues x num_atoms (3 for N, CA, C) x 3
+            encoder_padding_mask (ByteTensor): the positions of
+                  padding elements of shape `(batch_size x num_residues)`
+            confidence (Tensor): the confidence score of shape (batch_size x
+                num_residues). The value is between 0. and 1. for each residue
+                coordinate, or -1. if no coordinate is given
+            return_all_hiddens (bool, optional): also return all of the
+                intermediate hidden states (default: False).
+
+        Returns:
+            dict:
+                - **encoder_out** (Tensor): the last encoder layer's output of
+                  shape `(num_residues, batch_size, embed_dim)`
+                - **encoder_padding_mask** (ByteTensor): the positions of
+                  padding elements of shape `(batch_size, num_residues)`
+                - **encoder_embedding** (Tensor): the (scaled) embedding lookup
+                  of shape `(batch_size, num_residues, embed_dim)`
+                - **encoder_states** (List[Tensor]): all intermediate
+                  hidden states of shape `(num_residues, batch_size, embed_dim)`.
+                  Only populated if *return_all_hiddens* is True.
+        """
+        x, encoder_embedding = self.forward_embedding(coords,
+                encoder_padding_mask, confidence)
+        # account for padding while computing the representation
+        x = x * (1 - encoder_padding_mask.unsqueeze(-1).type_as(x))
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        encoder_states = []
+
+        if return_all_hiddens:
+            encoder_states.append(x)
+
+        # encoder layers
+        for layer in self.layers:
+            x = layer(
+                x, encoder_padding_mask=encoder_padding_mask
+            )
+            if return_all_hiddens:
+                assert encoder_states is not None
+                encoder_states.append(x)
+
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        return {
+            "encoder_out": [x],  # T x B x C
+            "encoder_padding_mask": [encoder_padding_mask],  # B x T
+            "encoder_embedding": [encoder_embedding],  # dictionary
+            "encoder_states": encoder_states,  # List[T x B x C]
+        }

model/esm/inverse_folding/gvp_utils.py

@@ -0,0 +1,68 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+
+def flatten_graph(node_embeddings, edge_embeddings, edge_index):
+    """
+    Flattens the graph into a batch size one (with disconnected subgraphs for
+    each example) to be compatible with pytorch-geometric package.
+    Args:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch size x nodes x node_embed_dim
+                - vector: shape batch size x nodes x node_embed_dim x 3
+        edge_embeddings: edge embeddings of in tuple form (scalar, vector)
+                - scalar: shape batch size x edges x edge_embed_dim
+                - vector: shape batch size x edges x edge_embed_dim x 3
+        edge_index: shape batch_size x 2 (source node and target node) x edges
+    Returns:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch total_nodes x node_embed_dim
+                - vector: shape batch total_nodes x node_embed_dim x 3
+        edge_embeddings: edge embeddings of in tuple form (scalar, vector)
+                - scalar: shape batch total_edges x edge_embed_dim
+                - vector: shape batch total_edges x edge_embed_dim x 3
+        edge_index: shape 2 x total_edges
+    """
+    x_s, x_v = node_embeddings
+    e_s, e_v = edge_embeddings
+    batch_size, N = x_s.shape[0], x_s.shape[1]
+    node_embeddings = (torch.flatten(x_s, 0, 1), torch.flatten(x_v, 0, 1))
+    edge_embeddings = (torch.flatten(e_s, 0, 1), torch.flatten(e_v, 0, 1))
+
+    edge_mask = torch.any(edge_index != -1, dim=1)
+    # Re-number the nodes by adding batch_idx * N to each batch
+    edge_index = edge_index + (torch.arange(batch_size, device=edge_index.device) *
+            N).unsqueeze(-1).unsqueeze(-1)
+    edge_index = edge_index.permute(1, 0, 2).flatten(1, 2)
+    edge_mask = edge_mask.flatten()
+    edge_index = edge_index[:, edge_mask] 
+    edge_embeddings = (
+        edge_embeddings[0][edge_mask, :],
+        edge_embeddings[1][edge_mask, :]
+    )
+    return node_embeddings, edge_embeddings, edge_index 
+
+
+def unflatten_graph(node_embeddings, batch_size):
+    """
+    Unflattens node embeddings.
+    Args:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch total_nodes x node_embed_dim
+                - vector: shape batch total_nodes x node_embed_dim x 3
+        batch_size: int
+    Returns:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch size x nodes x node_embed_dim
+                - vector: shape batch size x nodes x node_embed_dim x 3
+    """
+    x_s, x_v = node_embeddings
+    x_s = x_s.reshape(batch_size, -1, x_s.shape[1])
+    x_v = x_v.reshape(batch_size, -1, x_v.shape[1], x_v.shape[2])
+    return (x_s, x_v)
+
+

model/esm/inverse_folding/multichain_util.py

@@ -0,0 +1,152 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import biotite.structure
+import numpy as np
+import torch
+from typing import Sequence, Tuple, List
+
+from model.esm.inverse_folding.util import (
+    load_structure,
+    extract_coords_from_structure,
+    load_coords,
+    get_sequence_loss,
+    get_encoder_output,
+)
+
+
+def extract_coords_from_complex(structure: biotite.structure.AtomArray):
+    """
+    Args:
+        structure: biotite AtomArray
+    Returns:
+        Tuple (coords_list, seq_list)
+        - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+          coordinates representing the backbone of each chain
+        - seqs: Dictionary mapping chain ids to native sequences of each chain
+    """
+    coords = {}
+    seqs = {}
+    all_chains = biotite.structure.get_chains(structure)
+    for chain_id in all_chains:
+        chain = structure[structure.chain_id == chain_id]
+        coords[chain_id], seqs[chain_id] = extract_coords_from_structure(chain)
+    return coords, seqs
+
+
+def load_complex_coords(fpath, chains):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chains: the chain ids (the order matters for autoregressive model)
+    Returns:
+        Tuple (coords_list, seq_list)
+        - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+          coordinates representing the backbone of each chain
+        - seqs: Dictionary mapping chain ids to native sequences of each chain
+    """
+    structure = load_structure(fpath, chains)
+    return extract_coords_from_complex(structure)
+
+
+def _concatenate_coords(coords, target_chain_id, padding_length=10):
+    """
+    Args:
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        padding_length: Length of padding between concatenated chains
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates, a
+              concatenation of the chains with padding in between
+            - seq is the extracted sequence, with padding tokens inserted
+              between the concatenated chains
+    """
+    pad_coords = np.full((padding_length, 3, 3), np.nan, dtype=np.float32)
+    # For best performance, put the target chain first in concatenation.
+    coords_list = [coords[target_chain_id]]
+    for chain_id in coords:
+        if chain_id == target_chain_id:
+            continue
+        coords_list.append(pad_coords)
+        coords_list.append(coords[chain_id])
+    coords_concatenated = np.concatenate(coords_list, axis=0)
+    return coords_concatenated
+
+
+def sample_sequence_in_complex(model, coords, target_chain_id, temperature=1.,
+        padding_length=10):
+    """
+    Samples sequence for one chain in a complex.
+    Args:
+        model: An instance of the GVPTransformer model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        padding_length: padding length in between chains
+    Returns:
+        Sampled sequence for the target chain
+    """
+    target_chain_len = coords[target_chain_id].shape[0]
+    all_coords = _concatenate_coords(coords, target_chain_id)
+    device = next(model.parameters()).device
+
+    # Supply padding tokens for other chains to avoid unused sampling for speed
+    padding_pattern = ['<pad>'] * all_coords.shape[0]
+    for i in range(target_chain_len):
+        padding_pattern[i] = '<mask>'
+    sampled = model.sample(all_coords, partial_seq=padding_pattern,
+            temperature=temperature, device=device)
+    sampled = sampled[:target_chain_len]
+    return sampled
+
+
+def score_sequence_in_complex(model, alphabet, coords, target_chain_id,
+        target_seq, padding_length=10):
+    """
+    Scores sequence for one chain in a complex.
+    Args:
+        model: An instance of the GVPTransformer model
+        alphabet: Alphabet for the model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        target_seq: Target sequence for the target chain for scoring.
+        padding_length: padding length in between chains
+    Returns:
+        Tuple (ll_fullseq, ll_withcoord)
+        - ll_fullseq: Average log-likelihood over the full target chain
+        - ll_withcoord: Average log-likelihood in target chain excluding those
+            residues without coordinates
+    """
+    all_coords = _concatenate_coords(coords, target_chain_id)
+
+    loss, target_padding_mask = get_sequence_loss(model, alphabet, all_coords,
+            target_seq)
+    ll_fullseq = -np.sum(loss * ~target_padding_mask) / np.sum(
+            ~target_padding_mask)
+
+    # Also calculate average when excluding masked portions
+    coord_mask = np.all(np.isfinite(coords[target_chain_id]), axis=(-1, -2))
+    ll_withcoord = -np.sum(loss * coord_mask) / np.sum(coord_mask)
+    return ll_fullseq, ll_withcoord
+
+
+def get_encoder_output_for_complex(model, alphabet, coords, target_chain_id):
+    """
+    Args:
+        model: An instance of the GVPTransformer model
+        alphabet: Alphabet for the model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+    Returns:
+        Dictionary mapping chain id to encoder output for each chain
+    """
+    all_coords = _concatenate_coords(coords, target_chain_id)
+    all_rep = get_encoder_output(model, alphabet, all_coords)
+    target_chain_len = coords[target_chain_id].shape[0]
+    return all_rep[:target_chain_len]

model/esm/inverse_folding/transformer_decoder.py

@@ -0,0 +1,228 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Contents of this file were adapted from the open source fairseq repository.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Any, Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from onescience.modules.esm import SinusoidalPositionalEmbedding
+from .transformer_layer import TransformerDecoderLayer
+
+
+def fill_with_neg_inf(t):
+    """FP16-compatible function that fills a tensor with -inf."""
+    return t.float().fill_(float("-inf")).type_as(t)
+
+
+class TransformerDecoder(nn.Module):
+    """
+    Transformer decoder consisting of *args.decoder.layers* layers. Each layer
+    is a :class:`TransformerDecoderLayer`.
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+        dictionary (~fairseq.data.Dictionary): decoding dictionary
+        embed_tokens (torch.nn.Embedding): output embedding
+        no_encoder_attn (bool, optional): whether to attend to encoder outputs
+            (default: False).
+    """
+
+    def __init__(
+        self,
+        args,
+        dictionary,
+        embed_tokens,
+    ):
+        super().__init__()
+        self.args = args
+        self.dictionary = dictionary
+        self._future_mask = torch.empty(0)
+
+        self.dropout_module = nn.Dropout(args.dropout)
+
+        input_embed_dim = embed_tokens.embedding_dim
+        embed_dim = args.decoder_embed_dim
+        self.embed_dim = embed_dim
+
+        self.padding_idx = embed_tokens.padding_idx
+
+        self.embed_tokens = embed_tokens
+        self.embed_scale = math.sqrt(embed_dim)
+
+        self.project_in_dim = (
+            nn.Linear(input_embed_dim, embed_dim, bias=False)
+            if embed_dim != input_embed_dim
+            else None
+        )
+        self.embed_positions = SinusoidalPositionalEmbedding(
+            embed_dim,
+            self.padding_idx,
+        )
+
+        self.layers = nn.ModuleList([])
+        self.layers.extend(
+            [
+                self.build_decoder_layer(args)
+                for _ in range(args.decoder_layers)
+            ]
+        )
+        self.num_layers = len(self.layers)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+        self.build_output_projection(args, dictionary)
+
+    def build_output_projection(self, args, dictionary):
+        self.output_projection = nn.Linear(
+            args.decoder_embed_dim, len(dictionary), bias=False
+        )
+        nn.init.normal_(
+            self.output_projection.weight, mean=0, std=args.decoder_embed_dim ** -0.5
+        )
+
+    def build_decoder_layer(self, args):
+        return TransformerDecoderLayer(args)
+
+    def forward(
+        self,
+        prev_output_tokens,
+        encoder_out: Optional[Dict[str, List[Tensor]]] = None,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+        features_only: bool = False,
+        return_all_hiddens: bool = False,
+    ):
+        """
+        Args:
+            prev_output_tokens (LongTensor): previous decoder outputs of shape
+                `(batch, tgt_len)`, for teacher forcing
+            encoder_out (optional): output from the encoder, used for
+                encoder-side attention, should be of size T x B x C
+            incremental_state (dict): dictionary used for storing state during
+                :ref:`Incremental decoding`
+            features_only (bool, optional): only return features without
+                applying output layer (default: False).
+
+        Returns:
+            tuple:
+                - the decoder's output of shape `(batch, tgt_len, vocab)`
+                - a dictionary with any model-specific outputs
+        """
+
+        x, extra = self.extract_features(
+            prev_output_tokens,
+            encoder_out=encoder_out,
+            incremental_state=incremental_state,
+        )
+
+        if not features_only:
+            x = self.output_layer(x)
+        x = x.transpose(1, 2) # B x T x C -> B x C x T
+        return x, extra
+
+    def extract_features(
+        self,
+        prev_output_tokens,
+        encoder_out: Optional[Dict[str, List[Tensor]]],
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+    ):
+        """
+        Similar to *forward* but only return features.
+
+        Includes several features from "Jointly Learning to Align and
+        Translate with Transformer Models" (Garg et al., EMNLP 2019).
+
+        Returns:
+            tuple:
+                - the decoder's features of shape `(batch, tgt_len, embed_dim)`
+                - a dictionary with any model-specific outputs
+        """
+        bs, slen = prev_output_tokens.size()
+
+        enc: Optional[Tensor] = None
+        padding_mask: Optional[Tensor] = None
+        if encoder_out is not None and len(encoder_out["encoder_out"]) > 0:
+            enc = encoder_out["encoder_out"][0]
+            assert (
+                enc.size()[1] == bs
+            ), f"Expected enc.shape == (t, {bs}, c) got {enc.shape}"
+        if encoder_out is not None and len(encoder_out["encoder_padding_mask"]) > 0:
+            padding_mask = encoder_out["encoder_padding_mask"][0]
+
+        # embed positions
+        positions = self.embed_positions(
+            prev_output_tokens
+        )
+
+        if incremental_state is not None:
+            prev_output_tokens = prev_output_tokens[:, -1:]
+            positions = positions[:, -1:]
+
+        # embed tokens and positions
+        x = self.embed_scale * self.embed_tokens(prev_output_tokens)
+
+        if self.project_in_dim is not None:
+            x = self.project_in_dim(x)
+
+        x += positions
+
+        x = self.dropout_module(x)
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        self_attn_padding_mask: Optional[Tensor] = None
+        if prev_output_tokens.eq(self.padding_idx).any():
+            self_attn_padding_mask = prev_output_tokens.eq(self.padding_idx)
+
+        # decoder layers
+        attn: Optional[Tensor] = None
+        inner_states: List[Optional[Tensor]] = [x]
+        for idx, layer in enumerate(self.layers):
+            if incremental_state is None:
+                self_attn_mask = self.buffered_future_mask(x)
+            else:
+                self_attn_mask = None
+
+            x, layer_attn, _ = layer(
+                x,
+                enc,
+                padding_mask,
+                incremental_state,
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask,
+                need_attn=False,
+                need_head_weights=False,
+            )
+            inner_states.append(x)
+
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        # T x B x C -> B x C x T
+        x = x.transpose(0, 1)
+
+        return x, {"inner_states": inner_states}
+
+    def output_layer(self, features):
+        """Project features to the vocabulary size."""
+        return self.output_projection(features)
+
+    def buffered_future_mask(self, tensor):
+        dim = tensor.size(0)
+        # self._future_mask.device != tensor.device is not working in TorchScript. This is a workaround.
+        if (
+            self._future_mask.size(0) == 0
+            or (not self._future_mask.device == tensor.device)
+            or self._future_mask.size(0) < dim
+        ):
+            self._future_mask = torch.triu(
+                fill_with_neg_inf(torch.zeros([dim, dim])), 1
+            )
+        self._future_mask = self._future_mask.to(tensor)
+        return self._future_mask[:dim, :dim]

model/esm/inverse_folding/transformer_layer.py

@@ -0,0 +1,304 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Contents of this file were adapted from the open source fairseq repository.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from onescience.modules.attention import MultiheadAttention
+from torch import Tensor
+
+
+class TransformerEncoderLayer(nn.Module):
+    """Encoder layer block.
+    `layernorm -> dropout -> add residual`
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+    """
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.embed_dim = args.encoder_embed_dim
+        self.self_attn = self.build_self_attention(self.embed_dim, args)
+        self.self_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
+        self.dropout_module = nn.Dropout(args.dropout)
+        self.activation_fn = F.relu
+        self.fc1 = self.build_fc1(
+            self.embed_dim,
+            args.encoder_ffn_embed_dim,
+        )
+        self.fc2 = self.build_fc2(
+            args.encoder_ffn_embed_dim,
+            self.embed_dim,
+        )
+
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+    def build_fc1(self, input_dim, output_dim):
+        return nn.Linear(input_dim, output_dim)
+
+    def build_fc2(self, input_dim, output_dim):
+        return nn.Linear(input_dim, output_dim)
+
+    def build_self_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            embed_dim,
+            args.encoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=True,
+        )
+
+    def residual_connection(self, x, residual):
+        return residual + x
+
+    def forward(
+        self,
+        x,
+        encoder_padding_mask: Optional[Tensor],
+        attn_mask: Optional[Tensor] = None,
+    ):
+        """
+        Args:
+            x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
+            encoder_padding_mask (ByteTensor): binary ByteTensor of shape
+                `(batch, seq_len)` where padding elements are indicated by ``1``.
+            attn_mask (ByteTensor): binary tensor of shape `(tgt_len, src_len)`,
+                where `tgt_len` is the length of output and `src_len` is the
+                length of input, though here both are equal to `seq_len`.
+                `attn_mask[tgt_i, src_j] = 1` means that when calculating the
+                embedding for `tgt_i`, we exclude (mask out) `src_j`. This is
+                useful for strided self-attention.
+
+        Returns:
+            encoded output of shape `(seq_len, batch, embed_dim)`
+        """
+        # anything in original attn_mask = 1, becomes -1e8
+        # anything in original attn_mask = 0, becomes 0
+        # Note that we cannot use -inf here, because at some edge cases,
+        # the attention weight (before softmax) for some padded element in query
+        # will become -inf, which results in NaN in model parameters
+        if attn_mask is not None:
+            attn_mask = attn_mask.masked_fill(
+                attn_mask.to(torch.bool), -1e8 if x.dtype == torch.float32 else -1e4
+            )
+
+        residual = x
+        x = self.self_attn_layer_norm(x)
+        x, _ = self.self_attn(
+            query=x,
+            key=x,
+            value=x,
+            key_padding_mask=encoder_padding_mask,
+            need_weights=False,
+            attn_mask=attn_mask,
+        )
+        x = self.dropout_module(x)
+        x = self.residual_connection(x, residual)
+
+        residual = x
+        x = self.final_layer_norm(x)
+        x = self.activation_fn(self.fc1(x))
+        x = self.fc2(x)
+        x = self.dropout_module(x)
+        x = self.residual_connection(x, residual)
+        return x
+
+
+class TransformerDecoderLayer(nn.Module):
+    """Decoder layer block.
+    `layernorm -> dropout -> add residual`
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+        no_encoder_attn (bool, optional): whether to attend to encoder outputs
+            (default: False).
+    """
+
+    def __init__(
+        self, args, no_encoder_attn=False, add_bias_kv=False, add_zero_attn=False
+    ):
+        super().__init__()
+        self.embed_dim = args.decoder_embed_dim
+        self.dropout_module = nn.Dropout(args.dropout)
+
+        self.self_attn = self.build_self_attention(
+            self.embed_dim,
+            args,
+            add_bias_kv=add_bias_kv,
+            add_zero_attn=add_zero_attn,
+        )
+        self.nh = self.self_attn.num_heads
+        self.head_dim = self.self_attn.head_dim
+
+        self.activation_fn = F.relu
+
+        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+
+        if no_encoder_attn:
+            self.encoder_attn = None
+            self.encoder_attn_layer_norm = None
+        else:
+            self.encoder_attn = self.build_encoder_attention(self.embed_dim, args)
+            self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+
+        self.ffn_layernorm = (
+            LayerNorm(args.decoder_ffn_embed_dim)
+            if getattr(args, "scale_fc", False)
+            else None
+        )
+        self.w_resid = (
+            nn.Parameter(
+                torch.ones(
+                    self.embed_dim,
+                ),
+                requires_grad=True,
+            )
+            if getattr(args, "scale_resids", False)
+            else None
+        )
+
+        self.fc1 = self.build_fc1(
+            self.embed_dim,
+            args.decoder_ffn_embed_dim,
+        )
+        self.fc2 = self.build_fc2(
+            args.decoder_ffn_embed_dim,
+            self.embed_dim,
+        )
+
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.need_attn = True
+
+    def build_fc1(self, input_dim, output_dim):
+        return nn.Linear(input_dim, output_dim)
+
+    def build_fc2(self, input_dim, output_dim):
+        return nn.Linear(input_dim, output_dim)
+
+    def build_self_attention(
+        self, embed_dim, args, add_bias_kv=False, add_zero_attn=False
+    ):
+        return MultiheadAttention(
+            embed_dim,
+            args.decoder_attention_heads,
+            dropout=args.attention_dropout,
+            add_bias_kv=add_bias_kv,
+            add_zero_attn=add_zero_attn,
+            self_attention=True,
+        )
+
+    def build_encoder_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            embed_dim,
+            args.decoder_attention_heads,
+            kdim=args.encoder_embed_dim,
+            vdim=args.encoder_embed_dim,
+            dropout=args.attention_dropout,
+            encoder_decoder_attention=True,
+        )
+
+    def residual_connection(self, x, residual):
+        return residual + x
+
+    def forward(
+        self,
+        x,
+        encoder_out: Optional[torch.Tensor] = None,
+        encoder_padding_mask: Optional[torch.Tensor] = None,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+        prev_self_attn_state: Optional[List[torch.Tensor]] = None,
+        prev_attn_state: Optional[List[torch.Tensor]] = None,
+        self_attn_mask: Optional[torch.Tensor] = None,
+        self_attn_padding_mask: Optional[torch.Tensor] = None,
+        need_attn: bool = False,
+        need_head_weights: bool = False,
+    ):
+        """
+        Args:
+            x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
+            encoder_padding_mask (ByteTensor, optional): binary
+                ByteTensor of shape `(batch, src_len)` where padding
+                elements are indicated by ``1``.
+            need_attn (bool, optional): return attention weights
+            need_head_weights (bool, optional): return attention weights
+                for each head (default: return average over heads).
+
+        Returns:
+            encoded output of shape `(seq_len, batch, embed_dim)`
+        """
+        if need_head_weights:
+            need_attn = True
+
+        residual = x
+        x = self.self_attn_layer_norm(x)
+        if prev_self_attn_state is not None:
+            prev_key, prev_value = prev_self_attn_state[:2]
+            saved_state: Dict[str, Optional[Tensor]] = {
+                "prev_key": prev_key,
+                "prev_value": prev_value,
+            }
+            if len(prev_self_attn_state) >= 3:
+                saved_state["prev_key_padding_mask"] = prev_self_attn_state[2]
+            assert incremental_state is not None
+            self.self_attn._set_input_buffer(incremental_state, saved_state)
+        _self_attn_input_buffer = self.self_attn._get_input_buffer(incremental_state)
+        y = x
+
+        x, attn = self.self_attn(
+            query=x,
+            key=y,
+            value=y,
+            key_padding_mask=self_attn_padding_mask,
+            incremental_state=incremental_state,
+            need_weights=False,
+            attn_mask=self_attn_mask,
+        )
+        x = self.dropout_module(x)
+        x = self.residual_connection(x, residual)
+
+        if self.encoder_attn is not None and encoder_out is not None:
+            residual = x
+            x = self.encoder_attn_layer_norm(x)
+            if prev_attn_state is not None:
+                prev_key, prev_value = prev_attn_state[:2]
+                saved_state: Dict[str, Optional[Tensor]] = {
+                    "prev_key": prev_key,
+                    "prev_value": prev_value,
+                }
+                if len(prev_attn_state) >= 3:
+                    saved_state["prev_key_padding_mask"] = prev_attn_state[2]
+                assert incremental_state is not None
+                self.encoder_attn._set_input_buffer(incremental_state, saved_state)
+
+            x, attn = self.encoder_attn(
+                query=x,
+                key=encoder_out,
+                value=encoder_out,
+                key_padding_mask=encoder_padding_mask,
+                incremental_state=incremental_state,
+                static_kv=True,
+                need_weights=need_attn or (not self.training and self.need_attn),
+                need_head_weights=need_head_weights,
+            )
+            x = self.dropout_module(x)
+            x = self.residual_connection(x, residual)
+
+        residual = x
+        x = self.final_layer_norm(x)
+
+        x = self.activation_fn(self.fc1(x))
+        if self.ffn_layernorm is not None:
+            x = self.ffn_layernorm(x)
+        x = self.fc2(x)
+        x = self.dropout_module(x)
+        if self.w_resid is not None:
+            residual = torch.mul(self.w_resid, residual)
+        x = self.residual_connection(x, residual)
+        return x, attn, None

model/esm/inverse_folding/util.py

@@ -0,0 +1,323 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import math
+
+import biotite.structure
+from biotite.structure.io import pdbx, pdb
+from biotite.structure.residues import get_residues
+from biotite.structure import filter_backbone
+from biotite.structure import get_chains
+from biotite.sequence import ProteinSequence
+import numpy as np
+from scipy.spatial import transform
+from scipy.stats import special_ortho_group
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data as data
+from typing import Sequence, Tuple, List
+
+from onescience.datapipes.esm import BatchConverter
+
+
+def load_structure(fpath, chain=None):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chain: the chain id or list of chain ids to load
+    Returns:
+        biotite.structure.AtomArray
+    """
+    if fpath.endswith('cif'):
+        with open(fpath) as fin:
+            pdbxf = pdbx.PDBxFile.read(fin)
+        structure = pdbx.get_structure(pdbxf, model=1)
+    elif fpath.endswith('pdb'):
+        with open(fpath) as fin:
+            pdbf = pdb.PDBFile.read(fin)
+        structure = pdb.get_structure(pdbf, model=1)
+    bbmask = filter_backbone(structure)
+    structure = structure[bbmask]
+    all_chains = get_chains(structure)
+    if len(all_chains) == 0:
+        raise ValueError('No chains found in the input file.')
+    if chain is None:
+        chain_ids = all_chains
+    elif isinstance(chain, list):
+        chain_ids = chain
+    else:
+        chain_ids = [chain] 
+    for chain in chain_ids:
+        if chain not in all_chains:
+            raise ValueError(f'Chain {chain} not found in input file')
+    chain_filter = [a.chain_id in chain_ids for a in structure]
+    structure = structure[chain_filter]
+    return structure
+
+
+def extract_coords_from_structure(structure: biotite.structure.AtomArray):
+    """
+    Args:
+        structure: An instance of biotite AtomArray
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates
+            - seq is the extracted sequence
+    """
+    coords = get_atom_coords_residuewise(["N", "CA", "C"], structure)
+    residue_identities = get_residues(structure)[1]
+    seq = ''.join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+    return coords, seq
+
+
+def load_coords(fpath, chain):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chain: the chain id
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates
+            - seq is the extracted sequence
+    """
+    structure = load_structure(fpath, chain)
+    return extract_coords_from_structure(structure)
+
+
+def get_atom_coords_residuewise(atoms: List[str], struct: biotite.structure.AtomArray):
+    """
+    Example for atoms argument: ["N", "CA", "C"]
+    """
+    def filterfn(s, axis=None):
+        filters = np.stack([s.atom_name == name for name in atoms], axis=1)
+        sum = filters.sum(0)
+        if not np.all(sum <= np.ones(filters.shape[1])):
+            raise RuntimeError("structure has multiple atoms with same name")
+        index = filters.argmax(0)
+        coords = s[index].coord
+        coords[sum == 0] = float("nan")
+        return coords
+
+    return biotite.structure.apply_residue_wise(struct, struct, filterfn)
+
+
+def get_sequence_loss(model, alphabet, coords, seq):
+    device = next(model.parameters()).device
+    batch_converter = CoordBatchConverter(alphabet)
+    batch = [(coords, None, seq)]
+    coords, confidence, strs, tokens, padding_mask = batch_converter(
+        batch, device=device)
+
+    prev_output_tokens = tokens[:, :-1].to(device)
+    target = tokens[:, 1:]
+    target_padding_mask = (target == alphabet.padding_idx)
+    logits, _ = model.forward(coords, padding_mask, confidence, prev_output_tokens)
+    loss = F.cross_entropy(logits, target, reduction='none')
+    loss = loss[0].cpu().detach().numpy()
+    target_padding_mask = target_padding_mask[0].cpu().numpy()
+    return loss, target_padding_mask
+
+
+def score_sequence(model, alphabet, coords, seq):
+    loss, target_padding_mask = get_sequence_loss(model, alphabet, coords, seq)
+    ll_fullseq = -np.sum(loss * ~target_padding_mask) / np.sum(~target_padding_mask)
+    # Also calculate average when excluding masked portions
+    coord_mask = np.all(np.isfinite(coords), axis=(-1, -2))
+    ll_withcoord = -np.sum(loss * coord_mask) / np.sum(coord_mask)
+    return ll_fullseq, ll_withcoord
+
+
+def get_encoder_output(model, alphabet, coords):
+    device = next(model.parameters()).device
+    batch_converter = CoordBatchConverter(alphabet)
+    batch = [(coords, None, None)]
+    coords, confidence, strs, tokens, padding_mask = batch_converter(
+        batch, device=device)
+    encoder_out = model.encoder.forward(coords, padding_mask, confidence,
+            return_all_hiddens=False)
+    # remove beginning and end (bos and eos tokens)
+    return encoder_out['encoder_out'][0][1:-1, 0]
+
+
+def rotate(v, R):
+    """
+    Rotates a vector by a rotation matrix.
+    
+    Args:
+        v: 3D vector, tensor of shape (length x batch_size x channels x 3)
+        R: rotation matrix, tensor of shape (length x batch_size x 3 x 3)
+
+    Returns:
+        Rotated version of v by rotation matrix R.
+    """
+    R = R.unsqueeze(-3)
+    v = v.unsqueeze(-1)
+    return torch.sum(v * R, dim=-2)
+
+
+def get_rotation_frames(coords):
+    """
+    Returns a local rotation frame defined by N, CA, C positions.
+
+    Args:
+        coords: coordinates, tensor of shape (batch_size x length x 3 x 3)
+        where the third dimension is in order of N, CA, C
+
+    Returns:
+        Local relative rotation frames in shape (batch_size x length x 3 x 3)
+    """
+    v1 = coords[:, :, 2] - coords[:, :, 1]
+    v2 = coords[:, :, 0] - coords[:, :, 1]
+    e1 = normalize(v1, dim=-1)
+    u2 = v2 - e1 * torch.sum(e1 * v2, dim=-1, keepdim=True)
+    e2 = normalize(u2, dim=-1)
+    e3 = torch.cross(e1, e2, dim=-1)
+    R = torch.stack([e1, e2, e3], dim=-2)
+    return R
+
+
+def nan_to_num(ts, val=0.0):
+    """
+    Replaces nans in tensor with a fixed value.    
+    """
+    val = torch.tensor(val, dtype=ts.dtype, device=ts.device)
+    return torch.where(~torch.isfinite(ts), val, ts)
+
+
+def rbf(values, v_min, v_max, n_bins=16):
+    """
+    Returns RBF encodings in a new dimension at the end.
+    """
+    rbf_centers = torch.linspace(v_min, v_max, n_bins, device=values.device)
+    rbf_centers = rbf_centers.view([1] * len(values.shape) + [-1])
+    rbf_std = (v_max - v_min) / n_bins
+    v_expand = torch.unsqueeze(values, -1)
+    z = (values.unsqueeze(-1) - rbf_centers) / rbf_std
+    return torch.exp(-z ** 2)
+
+
+def norm(tensor, dim, eps=1e-8, keepdim=False):
+    """
+    Returns L2 norm along a dimension.
+    """
+    return torch.sqrt(
+            torch.sum(torch.square(tensor), dim=dim, keepdim=keepdim) + eps)
+
+
+def normalize(tensor, dim=-1):
+    """
+    Normalizes a tensor along a dimension after removing nans.
+    """
+    return nan_to_num(
+        torch.div(tensor, norm(tensor, dim=dim, keepdim=True))
+    )
+
+
+class CoordBatchConverter(BatchConverter):
+    def __call__(self, raw_batch: Sequence[Tuple[Sequence, str]], device=None):
+        """
+        Args:
+            raw_batch: List of tuples (coords, confidence, seq)
+            In each tuple,
+                coords: list of floats, shape L x 3 x 3
+                confidence: list of floats, shape L; or scalar float; or None
+                seq: string of length L
+        Returns:
+            coords: Tensor of shape batch_size x L x 3 x 3
+            confidence: Tensor of shape batch_size x L
+            strs: list of strings
+            tokens: LongTensor of shape batch_size x L
+            padding_mask: ByteTensor of shape batch_size x L
+        """
+        self.alphabet.cls_idx = self.alphabet.get_idx("<cath>") 
+        batch = []
+        for coords, confidence, seq in raw_batch:
+            if confidence is None:
+                confidence = 1.
+            if isinstance(confidence, float) or isinstance(confidence, int):
+                confidence = [float(confidence)] * len(coords)
+            if seq is None:
+                seq = 'X' * len(coords)
+            batch.append(((coords, confidence), seq))
+
+        coords_and_confidence, strs, tokens = super().__call__(batch)
+
+        # pad beginning and end of each protein due to legacy reasons
+        coords = [
+            F.pad(torch.tensor(cd), (0, 0, 0, 0, 1, 1), value=np.inf)
+            for cd, _ in coords_and_confidence
+        ]
+        confidence = [
+            F.pad(torch.tensor(cf), (1, 1), value=-1.)
+            for _, cf in coords_and_confidence
+        ]
+        coords = self.collate_dense_tensors(coords, pad_v=np.nan)
+        confidence = self.collate_dense_tensors(confidence, pad_v=-1.)
+        if device is not None:
+            coords = coords.to(device)
+            confidence = confidence.to(device)
+            tokens = tokens.to(device)
+        padding_mask = torch.isnan(coords[:,:,0,0])
+        coord_mask = torch.isfinite(coords.sum(-2).sum(-1))
+        confidence = confidence * coord_mask + (-1.) * padding_mask
+        return coords, confidence, strs, tokens, padding_mask
+
+    def from_lists(self, coords_list, confidence_list=None, seq_list=None, device=None):
+        """
+        Args:
+            coords_list: list of length batch_size, each item is a list of
+            floats in shape L x 3 x 3 to describe a backbone
+            confidence_list: one of
+                - None, default to highest confidence
+                - list of length batch_size, each item is a scalar
+                - list of length batch_size, each item is a list of floats of
+                    length L to describe the confidence scores for the backbone
+                    with values between 0. and 1.
+            seq_list: either None or a list of strings
+        Returns:
+            coords: Tensor of shape batch_size x L x 3 x 3
+            confidence: Tensor of shape batch_size x L
+            strs: list of strings
+            tokens: LongTensor of shape batch_size x L
+            padding_mask: ByteTensor of shape batch_size x L
+        """
+        batch_size = len(coords_list)
+        if confidence_list is None:
+            confidence_list = [None] * batch_size
+        if seq_list is None:
+            seq_list = [None] * batch_size
+        raw_batch = zip(coords_list, confidence_list, seq_list)
+        return self.__call__(raw_batch, device)
+
+    @staticmethod
+    def collate_dense_tensors(samples, pad_v):
+        """
+        Takes a list of tensors with the following dimensions:
+            [(d_11,       ...,           d_1K),
+             (d_21,       ...,           d_2K),
+             ...,
+             (d_N1,       ...,           d_NK)]
+        and stack + pads them into a single tensor of:
+        (N, max_i=1,N { d_i1 }, ..., max_i=1,N {diK})
+        """
+        if len(samples) == 0:
+            return torch.Tensor()
+        if len(set(x.dim() for x in samples)) != 1:
+            raise RuntimeError(
+                f"Samples has varying dimensions: {[x.dim() for x in samples]}"
+            )
+        (device,) = tuple(set(x.device for x in samples))  # assumes all on same device
+        max_shape = [max(lst) for lst in zip(*[x.shape for x in samples])]
+        result = torch.empty(
+            len(samples), *max_shape, dtype=samples[0].dtype, device=device
+        )
+        result.fill_(pad_v)
+        for i in range(len(samples)):
+            result_i = result[i]
+            t = samples[i]
+            result_i[tuple(slice(0, k) for k in t.shape)] = t
+        return result

model/esm/msa_transformer.py

@@ -0,0 +1,238 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+
+from onescience.modules.esm import (
+    AxialTransformerLayer,
+    LearnedPositionalEmbedding,
+    RobertaLMHead,
+    ESM1bLayerNorm,
+    ContactPredictionHead,
+)
+
+from onescience.modules.attention import RowSelfAttention, ColumnSelfAttention
+
+
+
+class MSATransformer(nn.Module):
+    @classmethod
+    def add_args(cls, parser):
+        # fmt: off
+        parser.add_argument(
+            "--num_layers",
+            default=12,
+            type=int,
+            metavar="N",
+            help="number of layers"
+        )
+        parser.add_argument(
+            "--embed_dim",
+            default=768,
+            type=int,
+            metavar="N",
+            help="embedding dimension"
+        )
+        parser.add_argument(
+            "--logit_bias",
+            action="store_true",
+            help="whether to apply bias to logits"
+        )
+        parser.add_argument(
+            "--ffn_embed_dim",
+            default=3072,
+            type=int,
+            metavar="N",
+            help="embedding dimension for FFN",
+        )
+        parser.add_argument(
+            "--attention_heads",
+            default=12,
+            type=int,
+            metavar="N",
+            help="number of attention heads",
+        )
+        parser.add_argument(
+            "--dropout",
+            default=0.1,
+            type=float,
+            help="Dropout to apply."
+        )
+        parser.add_argument(
+            "--attention_dropout",
+            default=0.1,
+            type=float,
+            help="Dropout to apply."
+        )
+        parser.add_argument(
+            "--activation_dropout",
+            default=0.1,
+            type=float,
+            help="Dropout to apply."
+        )
+        parser.add_argument(
+            "--max_tokens_per_msa",
+            default=2 ** 14,
+            type=int,
+            help=(
+                "Used during inference to batch attention computations in a single "
+                "forward pass. This allows increased input sizes with less memory."
+            ),
+        )
+        # fmt: on
+
+    def __init__(self, args, alphabet):
+        super().__init__()
+        self.args = args
+        self.alphabet_size = len(alphabet)
+        self.padding_idx = alphabet.padding_idx
+        self.mask_idx = alphabet.mask_idx
+        self.cls_idx = alphabet.cls_idx
+        self.eos_idx = alphabet.eos_idx
+        self.prepend_bos = alphabet.prepend_bos
+        self.append_eos = alphabet.append_eos
+
+        self.embed_tokens = nn.Embedding(
+            self.alphabet_size, self.args.embed_dim, padding_idx=self.padding_idx
+        )
+
+        if getattr(self.args, "embed_positions_msa", False):
+            emb_dim = getattr(self.args, "embed_positions_msa_dim", self.args.embed_dim)
+            self.msa_position_embedding = nn.Parameter(
+                0.01 * torch.randn(1, 1024, 1, emb_dim),
+                requires_grad=True,
+            )
+        else:
+            self.register_parameter("msa_position_embedding", None)
+
+        self.dropout_module = nn.Dropout(self.args.dropout)
+        self.layers = nn.ModuleList(
+            [
+                AxialTransformerLayer(
+                    self.args.embed_dim,
+                    self.args.ffn_embed_dim,
+                    self.args.attention_heads,
+                    self.args.dropout,
+                    self.args.attention_dropout,
+                    self.args.activation_dropout,
+                    getattr(self.args, "max_tokens_per_msa", self.args.max_tokens),
+                )
+                for _ in range(self.args.layers)
+            ]
+        )
+
+        self.contact_head = ContactPredictionHead(
+            self.args.layers * self.args.attention_heads,
+            self.prepend_bos,
+            self.append_eos,
+            eos_idx=self.eos_idx,
+        )
+        self.embed_positions = LearnedPositionalEmbedding(
+            self.args.max_positions,
+            self.args.embed_dim,
+            self.padding_idx,
+        )
+        self.emb_layer_norm_before = ESM1bLayerNorm(self.args.embed_dim)
+        self.emb_layer_norm_after = ESM1bLayerNorm(self.args.embed_dim)
+        self.lm_head = RobertaLMHead(
+            embed_dim=self.args.embed_dim,
+            output_dim=self.alphabet_size,
+            weight=self.embed_tokens.weight,
+        )
+
+    def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False):
+        if return_contacts:
+            need_head_weights = True
+
+        assert tokens.ndim == 3
+        batch_size, num_alignments, seqlen = tokens.size()
+        padding_mask = tokens.eq(self.padding_idx)  # B, R, C
+        if not padding_mask.any():
+            padding_mask = None
+
+        x = self.embed_tokens(tokens)
+        x += self.embed_positions(tokens.view(batch_size * num_alignments, seqlen)).view(x.size())
+        if self.msa_position_embedding is not None:
+            if x.size(1) > 1024:
+                raise RuntimeError(
+                    "Using model with MSA position embedding trained on maximum MSA "
+                    f"depth of 1024, but received {x.size(1)} alignments."
+                )
+            x += self.msa_position_embedding[:, :num_alignments]
+
+        x = self.emb_layer_norm_before(x)
+
+        x = self.dropout_module(x)
+
+        if padding_mask is not None:
+            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))
+
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+        if 0 in repr_layers:
+            hidden_representations[0] = x
+
+        if need_head_weights:
+            row_attn_weights = []
+            col_attn_weights = []
+
+        # B x R x C x D -> R x C x B x D
+        x = x.permute(1, 2, 0, 3)
+
+        for layer_idx, layer in enumerate(self.layers):
+            x = layer(
+                x,
+                self_attn_padding_mask=padding_mask,
+                need_head_weights=need_head_weights,
+            )
+            if need_head_weights:
+                x, col_attn, row_attn = x
+                # H x C x B x R x R -> B x H x C x R x R
+                col_attn_weights.append(col_attn.permute(2, 0, 1, 3, 4))
+                # H x B x C x C -> B x H x C x C
+                row_attn_weights.append(row_attn.permute(1, 0, 2, 3))
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x.permute(2, 0, 1, 3)
+
+        x = self.emb_layer_norm_after(x)
+        x = x.permute(2, 0, 1, 3)  # R x C x B x D -> B x R x C x D
+
+        # last hidden representation should have layer norm applied
+        if (layer_idx + 1) in repr_layers:
+            hidden_representations[layer_idx + 1] = x
+        x = self.lm_head(x)
+
+        result = {"logits": x, "representations": hidden_representations}
+        if need_head_weights:
+            # col_attentions: B x L x H x C x R x R
+            col_attentions = torch.stack(col_attn_weights, 1)
+            # row_attentions: B x L x H x C x C
+            row_attentions = torch.stack(row_attn_weights, 1)
+            result["col_attentions"] = col_attentions
+            result["row_attentions"] = row_attentions
+            if return_contacts:
+                contacts = self.contact_head(tokens, row_attentions)
+                result["contacts"] = contacts
+
+        return result
+
+    def predict_contacts(self, tokens):
+        return self(tokens, return_contacts=True)["contacts"]
+
+    @property
+    def num_layers(self):
+        return self.args.layers
+
+    def max_tokens_per_msa_(self, value: int) -> None:
+        """The MSA Transformer automatically batches attention computations when
+        gradients are disabled to allow you to pass in larger MSAs at test time than
+        you can fit in GPU memory. By default this occurs when more than 2^14 tokens
+        are passed in the input MSA. You can set this value to infinity to disable
+        this behavior.
+        """
+        for module in self.modules():
+            if isinstance(module, (RowSelfAttention, ColumnSelfAttention)):
+                module.max_tokens_per_msa = value

model/esm/pretrained.py

@@ -0,0 +1,553 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import re
+import urllib
+import warnings
+from argparse import Namespace
+from pathlib import Path
+
+import torch
+
+import model.esm as esm
+from onescience.datapipes.esm import Alphabet
+from model.esm.esm2 import ESM2
+
+
+def _has_regression_weights(model_name):
+    """Return whether we expect / require regression weights;
+    Right now that is all models except ESM-1v, ESM-IF, and partially trained ESM2 models"""
+    return not ("esm1v" in model_name or "esm_if" in model_name or "270K" in model_name or "500K" in model_name)
+
+
+def load_model_and_alphabet(model_name):
+    if model_name.endswith(".pt"):  # treat as filepath
+        return load_model_and_alphabet_local(model_name)
+    else:
+        return load_model_and_alphabet_hub(model_name)
+
+
+def load_hub_workaround(url):
+    try:
+        data = torch.hub.load_state_dict_from_url(url, progress=False, map_location="cpu")
+    except RuntimeError:
+        # Pytorch version issue - see https://github.com/pytorch/pytorch/issues/43106
+        fn = Path(url).name
+        data = torch.load(
+            f"{torch.hub.get_dir()}/checkpoints/{fn}",
+            map_location="cpu",
+        )
+    except urllib.error.HTTPError as e:
+        raise Exception(f"Could not load {url}, check if you specified a correct model name?")
+    return data
+
+
+def load_regression_hub(model_name):
+    url = f"https://dl.fbaipublicfiles.com/fair-esm/regression/{model_name}-contact-regression.pt"
+    regression_data = load_hub_workaround(url)
+    return regression_data
+
+
+def _download_model_and_regression_data(model_name):
+    url = f"https://dl.fbaipublicfiles.com/fair-esm/models/{model_name}.pt"
+    model_data = load_hub_workaround(url)
+    if _has_regression_weights(model_name):
+        regression_data = load_regression_hub(model_name)
+    else:
+        regression_data = None
+    return model_data, regression_data
+
+
+def load_model_and_alphabet_hub(model_name):
+    model_data, regression_data = _download_model_and_regression_data(model_name)
+    return load_model_and_alphabet_core(model_name, model_data, regression_data)
+
+
+def load_model_and_alphabet_local(model_location):
+    """Load from local path. The regression weights need to be co-located"""
+    model_location = Path(model_location)
+    model_data = torch.load(str(model_location), map_location="cpu")
+    model_name = model_location.stem
+    if _has_regression_weights(model_name):
+        regression_location = str(model_location.with_suffix("")) + "-contact-regression.pt"
+        regression_data = torch.load(regression_location, map_location="cpu")
+    else:
+        regression_data = None
+    return load_model_and_alphabet_core(model_name, model_data, regression_data)
+
+
+def has_emb_layer_norm_before(model_state):
+    """Determine whether layer norm needs to be applied before the encoder"""
+    return any(k.startswith("emb_layer_norm_before") for k, param in model_state.items())
+
+
+def _load_model_and_alphabet_core_v1(model_data):
+    import model.esm as esm  # since esm.inverse_folding is imported below, you actually have to re-import esm here
+
+    alphabet = esm.Alphabet.from_architecture(model_data["args"].arch)
+
+    if model_data["args"].arch == "roberta_large":
+        # upgrade state dict
+        pra = lambda s: "".join(s.split("encoder_")[1:] if "encoder" in s else s)
+        prs1 = lambda s: "".join(s.split("encoder.")[1:] if "encoder" in s else s)
+        prs2 = lambda s: "".join(
+            s.split("sentence_encoder.")[1:] if "sentence_encoder" in s else s
+        )
+        model_args = {pra(arg[0]): arg[1] for arg in vars(model_data["args"]).items()}
+        model_state = {prs1(prs2(arg[0])): arg[1] for arg in model_data["model"].items()}
+        model_state["embed_tokens.weight"][alphabet.mask_idx].zero_()  # For token drop
+        model_args["emb_layer_norm_before"] = has_emb_layer_norm_before(model_state)
+        model_type = esm.ProteinBertModel
+
+    elif model_data["args"].arch == "protein_bert_base":
+
+        # upgrade state dict
+        pra = lambda s: "".join(s.split("decoder_")[1:] if "decoder" in s else s)
+        prs = lambda s: "".join(s.split("decoder.")[1:] if "decoder" in s else s)
+        model_args = {pra(arg[0]): arg[1] for arg in vars(model_data["args"]).items()}
+        model_state = {prs(arg[0]): arg[1] for arg in model_data["model"].items()}
+        model_type = esm.ProteinBertModel
+    elif model_data["args"].arch == "msa_transformer":
+
+        # upgrade state dict
+        pra = lambda s: "".join(s.split("encoder_")[1:] if "encoder" in s else s)
+        prs1 = lambda s: "".join(s.split("encoder.")[1:] if "encoder" in s else s)
+        prs2 = lambda s: "".join(
+            s.split("sentence_encoder.")[1:] if "sentence_encoder" in s else s
+        )
+        prs3 = lambda s: s.replace("row", "column") if "row" in s else s.replace("column", "row")
+        model_args = {pra(arg[0]): arg[1] for arg in vars(model_data["args"]).items()}
+        model_state = {prs1(prs2(prs3(arg[0]))): arg[1] for arg in model_data["model"].items()}
+        if model_args.get("embed_positions_msa", False):
+            emb_dim = model_state["msa_position_embedding"].size(-1)
+            model_args["embed_positions_msa_dim"] = emb_dim  # initial release, bug: emb_dim==1
+
+        model_type = esm.MSATransformer
+
+    elif "invariant_gvp" in model_data["args"].arch:
+        import model.esm.inverse_folding
+
+        model_type = esm.inverse_folding.gvp_transformer.GVPTransformerModel
+        model_args = vars(model_data["args"])  # convert Namespace -> dict
+
+        def update_name(s):
+            # Map the module names in checkpoints trained with internal code to
+            # the updated module names in open source code
+            s = s.replace("W_v", "embed_graph.embed_node")
+            s = s.replace("W_e", "embed_graph.embed_edge")
+            s = s.replace("embed_scores.0", "embed_confidence")
+            s = s.replace("embed_score.", "embed_graph.embed_confidence.")
+            s = s.replace("seq_logits_projection.", "")
+            s = s.replace("embed_ingraham_features", "embed_dihedrals")
+            s = s.replace("embed_gvp_in_local_frame.0", "embed_gvp_output")
+            s = s.replace("embed_features_in_local_frame.0", "embed_gvp_input_features")
+            return s
+
+        model_state = {
+            update_name(sname): svalue
+            for sname, svalue in model_data["model"].items()
+            if "version" not in sname
+        }
+
+    else:
+        raise ValueError("Unknown architecture selected")
+
+    model = model_type(
+        Namespace(**model_args),
+        alphabet,
+    )
+
+    return model, alphabet, model_state
+
+
+def _load_model_and_alphabet_core_v2(model_data):
+    def upgrade_state_dict(state_dict):
+        """Removes prefixes 'model.encoder.sentence_encoder.' and 'model.encoder.'."""
+        prefixes = ["encoder.sentence_encoder.", "encoder."]
+        pattern = re.compile("^" + "|".join(prefixes))
+        state_dict = {pattern.sub("", name): param for name, param in state_dict.items()}
+        return state_dict
+
+    cfg = model_data["cfg"]["model"]
+    state_dict = model_data["model"]
+    state_dict = upgrade_state_dict(state_dict)
+    alphabet = Alphabet.from_architecture("ESM-1b")
+    model = ESM2(
+        num_layers=cfg.encoder_layers,
+        embed_dim=cfg.encoder_embed_dim,
+        attention_heads=cfg.encoder_attention_heads,
+        alphabet=alphabet,
+        token_dropout=cfg.token_dropout,
+    )
+    return model, alphabet, state_dict
+
+
+def load_model_and_alphabet_core(model_name, model_data, regression_data=None):
+    if regression_data is not None:
+        model_data["model"].update(regression_data["model"])
+
+    if model_name.startswith("esm2"):
+        model, alphabet, model_state = _load_model_and_alphabet_core_v2(model_data)
+    else:
+        model, alphabet, model_state = _load_model_and_alphabet_core_v1(model_data)
+
+    expected_keys = set(model.state_dict().keys())
+    found_keys = set(model_state.keys())
+
+    if regression_data is None:
+        expected_missing = {"contact_head.regression.weight", "contact_head.regression.bias"}
+        error_msgs = []
+        missing = (expected_keys - found_keys) - expected_missing
+        if missing:
+            error_msgs.append(f"Missing key(s) in state_dict: {missing}.")
+        unexpected = found_keys - expected_keys
+        if unexpected:
+            error_msgs.append(f"Unexpected key(s) in state_dict: {unexpected}.")
+
+        if error_msgs:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for {}:\n\t{}".format(
+                    model.__class__.__name__, "\n\t".join(error_msgs)
+                )
+            )
+        if expected_missing - found_keys:
+            warnings.warn(
+                "Regression weights not found, predicting contacts will not produce correct results."
+            )
+
+    model.load_state_dict(model_state, strict=regression_data is not None)
+
+    return model, alphabet
+
+
+def esm1_t34_670M_UR50S():
+    """34 layer transformer model with 670M params, trained on Uniref50 Sparse.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t34_670M_UR50S")
+
+
+def esm1_t34_670M_UR50D():
+    """34 layer transformer model with 670M params, trained on Uniref50 Dense.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t34_670M_UR50D")
+
+
+def esm1_t34_670M_UR100():
+    """34 layer transformer model with 670M params, trained on Uniref100.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t34_670M_UR100")
+
+
+def esm1_t12_85M_UR50S():
+    """12 layer transformer model with 85M params, trained on Uniref50 Sparse.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t12_85M_UR50S")
+
+
+def esm1_t6_43M_UR50S():
+    """6 layer transformer model with 43M params, trained on Uniref50 Sparse.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t6_43M_UR50S")
+
+
+def esm1b_t33_650M_UR50S():
+    """33 layer transformer model with 650M params, trained on Uniref50 Sparse.
+    This is our best performing model, which will be described in a future publication.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1b_t33_650M_UR50S")
+
+
+def esm_msa1_t12_100M_UR50S():
+    warnings.warn(
+        "This model had a minor bug in the positional embeddings, "
+        "please use ESM-MSA-1b: esm.pretrained.esm_msa1b_t12_100M_UR50S()",
+    )
+    return load_model_and_alphabet_hub("esm_msa1_t12_100M_UR50S")
+
+
+def esm_msa1b_t12_100M_UR50S():
+    return load_model_and_alphabet_hub("esm_msa1b_t12_100M_UR50S")
+
+
+def esm1v_t33_650M_UR90S():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 1 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_1")
+
+
+def esm1v_t33_650M_UR90S_1():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 1 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_1")
+
+
+def esm1v_t33_650M_UR90S_2():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 2 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_2")
+
+
+def esm1v_t33_650M_UR90S_3():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 3 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_3")
+
+
+def esm1v_t33_650M_UR90S_4():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 4 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_4")
+
+
+def esm1v_t33_650M_UR90S_5():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 5 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_5")
+
+
+def esm_if1_gvp4_t16_142M_UR50():
+    """Inverse folding model with 142M params, with 4 GVP-GNN layers, 8
+    Transformer encoder layers, and 8 Transformer decoder layers, trained on
+    CATH structures and 12 million alphafold2 predicted structures from UniRef50
+    sequences.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm_if1_gvp4_t16_142M_UR50")
+
+
+def esm2_t6_8M_UR50D():
+    """6 layer ESM-2 model with 8M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t6_8M_UR50D")
+
+
+def esm2_t12_35M_UR50D():
+    """12 layer ESM-2 model with 35M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t12_35M_UR50D")
+
+
+def esm2_t30_150M_UR50D():
+    """30 layer ESM-2 model with 150M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t30_150M_UR50D")
+
+
+def esm2_t33_650M_UR50D():
+    """33 layer ESM-2 model with 650M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t33_650M_UR50D")
+
+
+def esm2_t36_3B_UR50D():
+    """36 layer ESM-2 model with 3B params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t36_3B_UR50D")
+
+
+def esm2_t48_15B_UR50D():
+    """48 layer ESM-2 model with 15B params, trained on UniRef50.
+    If you have OOM while loading this model, please refer to README
+    on how to employ FSDP and ZeRO CPU offloading
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t48_15B_UR50D")
+
+
+def esmfold_v0():
+    """
+    ESMFold v0 model with 3B ESM-2, 48 folding blocks.
+    This version was used for the paper (Lin et al, 2022). It was trained 
+    on all PDB chains until 2020-05, to ensure temporal holdout with CASP14
+    and the CAMEO validation and test set reported there.
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_v0()
+
+
+def esmfold_v1():
+    """
+    ESMFold v1 model using 3B ESM-2, 48 folding blocks.
+    ESMFold provides fast high accuracy atomic level structure prediction
+    directly from the individual sequence of a protein. ESMFold uses the ESM2
+    protein language model to extract meaningful representations from the
+    protein sequence.
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_v1()
+
+def esmfold_structure_module_only_8M():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_8M()
+
+
+def esmfold_structure_module_only_8M_270K():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_8M_270K()
+
+
+def esmfold_structure_module_only_35M():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_35M()
+
+
+def esmfold_structure_module_only_35M_270K():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_35M_270K()
+
+
+def esmfold_structure_module_only_150M():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_150M()
+
+
+def esmfold_structure_module_only_150M_270K():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_150M_270K()
+
+
+def esmfold_structure_module_only_650M():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_650M()
+
+
+def esmfold_structure_module_only_650M_270K():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_650M_270K()
+
+
+def esmfold_structure_module_only_3B():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_3B()
+
+
+def esmfold_structure_module_only_3B_270K():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_3B_270K()
+
+
+def esmfold_structure_module_only_15B():
+    """
+    ESMFold baseline model using 15B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    The 15B parameter ESM-2 was not trained out to 500K updates
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_15B()

model/esm/version.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+version = "2.0.1"

model/openfold/dropout.py

@@ -0,0 +1,78 @@
+# Copyright 2021 AlQuraishi Laboratory
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+import torch.nn as nn
+from functools import partialmethod
+from typing import Union, List
+
+
+class Dropout(nn.Module):
+    """
+    Implementation of dropout with the ability to share the dropout mask
+    along a particular dimension.
+
+    If not in training mode, this module computes the identity function.
+    """
+
+    def __init__(self, r: float, batch_dim: Union[int, List[int]]):
+        """
+        Args:
+            r:
+                Dropout rate
+            batch_dim:
+                Dimension(s) along which the dropout mask is shared
+        """
+        super(Dropout, self).__init__()
+
+        self.r = r
+        if type(batch_dim) == int:
+            batch_dim = [batch_dim]
+        self.batch_dim = batch_dim
+        self.dropout = nn.Dropout(self.r)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x:
+                Tensor to which dropout is applied. Can have any shape
+                compatible with self.batch_dim
+        """
+        shape = list(x.shape)
+        if self.batch_dim is not None:
+            for bd in self.batch_dim:
+                shape[bd] = 1
+        mask = x.new_ones(shape)
+        mask = self.dropout(mask)
+        x *= mask
+        return x
+
+
+class DropoutRowwise(Dropout):
+    """
+    Convenience class for rowwise dropout as described in subsection
+    1.11.6.
+    """
+
+    __init__ = partialmethod(Dropout.__init__, batch_dim=-3)
+
+
+class DropoutColumnwise(Dropout):
+    """
+    Convenience class for columnwise dropout as described in subsection
+    1.11.6.
+    """
+
+    __init__ = partialmethod(Dropout.__init__, batch_dim=-2)

model/openfold/embedders.py

@@ -0,0 +1,984 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+from typing import Tuple, Optional
+
+from onescience.utils.openfold import all_atom_multimer
+from onescience.utils.openfold.feats import (  
+    pseudo_beta_fn,
+    dgram_from_positions,
+    build_template_angle_feat,
+    build_template_pair_feat,
+)
+from model.openfold.primitives import Linear, LayerNorm
+from model.openfold.template import (
+    TemplatePairStack,
+    TemplatePointwiseAttention,
+)
+from onescience.utils.openfold import geometry
+from onescience.utils.openfold.tensor_utils import add, one_hot, tensor_tree_map, dict_multimap
+
+
+class InputEmbedder(nn.Module):
+    """
+    Embeds a subset of the input features.
+
+    Implements Algorithms 3 (InputEmbedder) and 4 (relpos).
+    """
+
+    def __init__(
+        self,
+        tf_dim: int,
+        msa_dim: int,
+        c_z: int,
+        c_m: int,
+        relpos_k: int,
+        **kwargs,
+    ):
+        """
+        Args:
+            tf_dim:
+                Final dimension of the target features
+            msa_dim:
+                Final dimension of the MSA features
+            c_z:
+                Pair embedding dimension
+            c_m:
+                MSA embedding dimension
+            relpos_k:
+                Window size used in relative positional encoding
+        """
+        super(InputEmbedder, self).__init__()
+
+        self.tf_dim = tf_dim
+        self.msa_dim = msa_dim
+
+        self.c_z = c_z
+        self.c_m = c_m
+
+        self.linear_tf_z_i = Linear(tf_dim, c_z)
+        self.linear_tf_z_j = Linear(tf_dim, c_z)
+        self.linear_tf_m = Linear(tf_dim, c_m)
+        self.linear_msa_m = Linear(msa_dim, c_m)
+
+        # RPE stuff
+        self.relpos_k = relpos_k
+        self.no_bins = 2 * relpos_k + 1
+        self.linear_relpos = Linear(self.no_bins, c_z)
+
+    def relpos(self, ri: torch.Tensor):
+        """
+        Computes relative positional encodings
+
+        Implements Algorithm 4.
+
+        Args:
+            ri:
+                "residue_index" features of shape [*, N]
+        """
+        d = ri[..., None] - ri[..., None, :]
+        boundaries = torch.arange(
+            start=-self.relpos_k, end=self.relpos_k + 1, device=d.device
+        ) 
+        reshaped_bins = boundaries.view(((1,) * len(d.shape)) + (len(boundaries),))
+        d = d[..., None] - reshaped_bins
+        d = torch.abs(d)
+        d = torch.argmin(d, dim=-1)
+        d = nn.functional.one_hot(d, num_classes=len(boundaries)).float()
+        d = d.to(ri.dtype)
+        return self.linear_relpos(d)
+
+    def forward(
+        self,
+        tf: torch.Tensor,
+        ri: torch.Tensor,
+        msa: torch.Tensor,
+        inplace_safe: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            batch: Dict containing
+                "target_feat":
+                    Features of shape [*, N_res, tf_dim]
+                "residue_index":
+                    Features of shape [*, N_res]
+                "msa_feat":
+                    Features of shape [*, N_clust, N_res, msa_dim]
+        Returns:
+            msa_emb:
+                [*, N_clust, N_res, C_m] MSA embedding
+            pair_emb:
+                [*, N_res, N_res, C_z] pair embedding
+
+        """
+        # [*, N_res, c_z]
+        tf_emb_i = self.linear_tf_z_i(tf)
+        tf_emb_j = self.linear_tf_z_j(tf)
+
+        # [*, N_res, N_res, c_z]
+        pair_emb = self.relpos(ri.type(tf_emb_i.dtype))
+        pair_emb = add(pair_emb, 
+            tf_emb_i[..., None, :], 
+            inplace=inplace_safe
+        )
+        pair_emb = add(pair_emb, 
+            tf_emb_j[..., None, :, :], 
+            inplace=inplace_safe
+        )
+
+        # [*, N_clust, N_res, c_m]
+        n_clust = msa.shape[-3]
+        tf_m = (
+            self.linear_tf_m(tf)
+            .unsqueeze(-3)
+            .expand(((-1,) * len(tf.shape[:-2]) + (n_clust, -1, -1)))
+        )
+        msa_emb = self.linear_msa_m(msa) + tf_m
+
+        return msa_emb, pair_emb
+
+
+class InputEmbedderMultimer(nn.Module):
+    """
+    Embeds a subset of the input features.
+
+    Implements Algorithms 3 (InputEmbedder) and 4 (relpos).
+    """
+
+    def __init__(
+        self,
+        tf_dim: int,
+        msa_dim: int,
+        c_z: int,
+        c_m: int,
+        max_relative_idx: int,
+        use_chain_relative: bool,
+        max_relative_chain: int,
+        **kwargs,
+    ):
+        """
+        Args:
+            tf_dim:
+                Final dimension of the target features
+            msa_dim:
+                Final dimension of the MSA features
+            c_z:
+                Pair embedding dimension
+            c_m:
+                MSA embedding dimension
+            relpos_k:
+                Window size used in relative positional encoding
+        """
+        super(InputEmbedderMultimer, self).__init__()
+
+        self.tf_dim = tf_dim
+        self.msa_dim = msa_dim
+
+        self.c_z = c_z
+        self.c_m = c_m
+
+        self.linear_tf_z_i = Linear(tf_dim, c_z)
+        self.linear_tf_z_j = Linear(tf_dim, c_z)
+        self.linear_tf_m = Linear(tf_dim, c_m)
+        self.linear_msa_m = Linear(msa_dim, c_m)
+
+        # RPE stuff
+        self.max_relative_idx = max_relative_idx
+        self.use_chain_relative = use_chain_relative
+        self.max_relative_chain = max_relative_chain
+        if(self.use_chain_relative):
+            self.no_bins = (
+                2 * max_relative_idx + 2 +
+                1 +
+                2 * max_relative_chain + 2
+            )
+        else:
+            self.no_bins = 2 * max_relative_idx + 1
+        self.linear_relpos = Linear(self.no_bins, c_z)
+
+    def relpos(self, batch):
+        pos = batch["residue_index"]
+        asym_id = batch["asym_id"]
+        asym_id_same = (asym_id[..., None] == asym_id[..., None, :])
+        offset = pos[..., None] - pos[..., None, :]
+
+        clipped_offset = torch.clamp(
+            offset + self.max_relative_idx, 0, 2 * self.max_relative_idx
+        )
+
+        rel_feats = []
+        if(self.use_chain_relative):
+            final_offset = torch.where(
+                asym_id_same, 
+                clipped_offset,
+                (2 * self.max_relative_idx + 1) * 
+                torch.ones_like(clipped_offset)
+            )
+            boundaries = torch.arange(
+                start=0, end=2 * self.max_relative_idx + 2, device=final_offset.device
+            )
+            rel_pos = one_hot(
+                final_offset,
+                boundaries,
+            )
+
+            rel_feats.append(rel_pos)
+
+            entity_id = batch["entity_id"]
+            entity_id_same = (entity_id[..., None] == entity_id[..., None, :])
+            rel_feats.append(entity_id_same[..., None].to(dtype=rel_pos.dtype))
+
+            sym_id = batch["sym_id"]
+            rel_sym_id = sym_id[..., None] - sym_id[..., None, :]
+
+            max_rel_chain = self.max_relative_chain
+            clipped_rel_chain = torch.clamp(
+                rel_sym_id + max_rel_chain,
+                0,
+                2 * max_rel_chain,
+            )
+
+            final_rel_chain = torch.where(
+                entity_id_same,
+                clipped_rel_chain,
+                (2 * max_rel_chain + 1) *
+                torch.ones_like(clipped_rel_chain)
+            )
+
+            boundaries = torch.arange(
+                start=0, end=2 * max_rel_chain + 2, device=final_rel_chain.device
+            )
+            rel_chain = one_hot(
+                final_rel_chain,
+                boundaries,
+            )
+
+            rel_feats.append(rel_chain)
+        else:
+            boundaries = torch.arange(
+                start=0, end=2 * self.max_relative_idx + 1, device=clipped_offset.device
+            )
+            rel_pos = one_hot(
+                clipped_offset, boundaries,
+            )
+            rel_feats.append(rel_pos)
+
+        rel_feat = torch.cat(rel_feats, dim=-1).to(
+            self.linear_relpos.weight.dtype
+        )
+
+        return self.linear_relpos(rel_feat)
+
+    def forward(self, batch) -> Tuple[torch.Tensor, torch.Tensor]:
+        tf = batch["target_feat"] 
+        msa = batch["msa_feat"]
+
+        # [*, N_res, c_z]
+        tf_emb_i = self.linear_tf_z_i(tf)
+        tf_emb_j = self.linear_tf_z_j(tf)
+
+        # [*, N_res, N_res, c_z]
+        pair_emb = tf_emb_i[..., None, :] + tf_emb_j[..., None, :, :]
+        pair_emb = pair_emb + self.relpos(batch)
+
+        # [*, N_clust, N_res, c_m]
+        n_clust = msa.shape[-3]
+        tf_m = (
+            self.linear_tf_m(tf)
+            .unsqueeze(-3)
+            .expand(((-1,) * len(tf.shape[:-2]) + (n_clust, -1, -1)))
+        )
+        msa_emb = self.linear_msa_m(msa) + tf_m
+
+        return msa_emb, pair_emb
+
+
+class PreembeddingEmbedder(nn.Module):
+    """
+    Embeds the sequence pre-embedding passed to the model and the target_feat features.
+    """
+
+    def __init__(
+        self,
+        tf_dim: int,
+        preembedding_dim: int,
+        c_z: int,
+        c_m: int,
+        relpos_k: int,
+        **kwargs,
+    ):
+        """
+        Args:
+            tf_dim:
+                End channel dimension of the incoming target features
+            preembedding_dim:
+                End channel dimension of the incoming embeddings
+            c_z:
+                Pair embedding dimension
+            c_m:
+                Single-Seq embedding dimension
+            relpos_k:
+                Window size used in relative position encoding
+        """
+        super(PreembeddingEmbedder, self).__init__()
+
+        self.tf_dim = tf_dim
+        self.preembedding_dim = preembedding_dim
+
+        self.c_z = c_z
+        self.c_m = c_m
+
+        self.linear_tf_m = Linear(tf_dim, c_m)
+        self.linear_preemb_m = Linear(self.preembedding_dim, c_m)
+        self.linear_preemb_z_i = Linear(self.preembedding_dim, c_z)
+        self.linear_preemb_z_j = Linear(self.preembedding_dim, c_z)
+
+        # Relative Positional Encoding
+        self.relpos_k = relpos_k
+        self.no_bins = 2 * relpos_k + 1
+        self.linear_relpos = Linear(self.no_bins, c_z)
+
+    def relpos(self, ri: torch.Tensor):
+        """
+        Computes relative positional encodings
+        Args:
+            ri:
+                "residue_index" feature of shape [*, N]
+        Returns:
+                Relative positional encoding of protein using the
+                residue_index feature
+        """
+        d = ri[..., None] - ri[..., None, :]
+        boundaries = torch.arange(
+            start=-self.relpos_k, end=self.relpos_k + 1, device=d.device
+        )
+        reshaped_bins = boundaries.view(((1,) * len(d.shape)) + (len(boundaries),))
+        d = d[..., None] - reshaped_bins
+        d = torch.abs(d)
+        d = torch.argmin(d, dim=-1)
+        d = nn.functional.one_hot(d, num_classes=len(boundaries)).float()
+        d = d.to(ri.dtype)
+        return self.linear_relpos(d)
+
+    def forward(
+        self,
+        tf: torch.Tensor,
+        ri: torch.Tensor,
+        preemb: torch.Tensor,
+        inplace_safe: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        tf_m = (
+            self.linear_tf_m(tf)
+            .unsqueeze(-3)
+        )
+        preemb_emb = self.linear_preemb_m(preemb[..., None, :, :]) + tf_m
+        preemb_emb_i = self.linear_preemb_z_i(preemb)
+        preemb_emb_j = self.linear_preemb_z_j(preemb)
+
+        pair_emb = self.relpos(ri.type(preemb_emb_i.dtype))
+        pair_emb = add(pair_emb,
+                       preemb_emb_i[..., None, :],
+                       inplace=inplace_safe)
+        pair_emb = add(pair_emb,
+                       preemb_emb_j[..., None, :, :],
+                       inplace=inplace_safe)
+
+        return preemb_emb, pair_emb
+
+
+class RecyclingEmbedder(nn.Module):
+    """
+    Embeds the output of an iteration of the model for recycling.
+
+    Implements Algorithm 32.
+    """
+    def __init__(
+        self,
+        c_m: int,
+        c_z: int,
+        min_bin: float,
+        max_bin: float,
+        no_bins: int,
+        inf: float = 1e8,
+        **kwargs,
+    ):
+        """
+        Args:
+            c_m:
+                MSA channel dimension
+            c_z:
+                Pair embedding channel dimension
+            min_bin:
+                Smallest distogram bin (Angstroms)
+            max_bin:
+                Largest distogram bin (Angstroms)
+            no_bins:
+                Number of distogram bins
+        """
+        super(RecyclingEmbedder, self).__init__()
+
+        self.c_m = c_m
+        self.c_z = c_z
+        self.min_bin = min_bin
+        self.max_bin = max_bin
+        self.no_bins = no_bins
+        self.inf = inf
+
+        self.linear = Linear(self.no_bins, self.c_z)
+        self.layer_norm_m = LayerNorm(self.c_m)
+        self.layer_norm_z = LayerNorm(self.c_z)
+
+    def forward(
+        self,
+        m: torch.Tensor,
+        z: torch.Tensor,
+        x: torch.Tensor,
+        inplace_safe: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            m:
+                First row of the MSA embedding. [*, N_res, C_m]
+            z:
+                [*, N_res, N_res, C_z] pair embedding
+            x:
+                [*, N_res, 3] predicted C_beta coordinates
+        Returns:
+            m:
+                [*, N_res, C_m] MSA embedding update
+            z:
+                [*, N_res, N_res, C_z] pair embedding update
+        """
+        # [*, N, C_m]
+        m_update = self.layer_norm_m(m)
+        if(inplace_safe):
+            m.copy_(m_update)
+            m_update = m
+
+        # [*, N, N, C_z]
+        z_update = self.layer_norm_z(z)
+        if(inplace_safe):
+            z.copy_(z_update)
+            z_update = z
+
+        # This squared method might become problematic in FP16 mode.
+        bins = torch.linspace(
+            self.min_bin,
+            self.max_bin,
+            self.no_bins,
+            dtype=x.dtype,
+            device=x.device,
+            requires_grad=False,
+        )
+        squared_bins = bins ** 2
+        upper = torch.cat(
+            [squared_bins[1:], squared_bins.new_tensor([self.inf])], dim=-1
+        )
+        d = torch.sum(
+            (x[..., None, :] - x[..., None, :, :]) ** 2, dim=-1, keepdims=True
+        )
+
+        # [*, N, N, no_bins]
+        d = ((d > squared_bins) * (d < upper)).type(x.dtype)
+
+        # [*, N, N, C_z]
+        d = self.linear(d)
+        z_update = add(z_update, d, inplace_safe)
+
+        return m_update, z_update
+
+
+class TemplateSingleEmbedder(nn.Module):
+    """
+    Embeds the "template_angle_feat" feature.
+
+    Implements Algorithm 2, line 7.
+    """
+
+    def __init__(
+        self,
+        c_in: int,
+        c_out: int,
+        **kwargs,
+    ):
+        """
+        Args:
+            c_in:
+                Final dimension of "template_angle_feat"
+            c_out:
+                Output channel dimension
+        """
+        super(TemplateSingleEmbedder, self).__init__()
+
+        self.c_out = c_out
+        self.c_in = c_in
+
+        self.linear_1 = Linear(self.c_in, self.c_out, init="relu")
+        self.relu = nn.ReLU()
+        self.linear_2 = Linear(self.c_out, self.c_out, init="relu")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [*, N_templ, N_res, c_in] "template_angle_feat" features
+        Returns:
+            x: [*, N_templ, N_res, C_out] embedding
+        """
+        x = self.linear_1(x)
+        x = self.relu(x)
+        x = self.linear_2(x)
+
+        return x
+
+
+class TemplatePairEmbedder(nn.Module):
+    """
+    Embeds "template_pair_feat" features.
+
+    Implements Algorithm 2, line 9.
+    """
+
+    def __init__(
+        self,
+        c_in: int,
+        c_out: int,
+        **kwargs,
+    ):
+        """
+        Args:
+            c_in:
+
+            c_out:
+                Output channel dimension
+        """
+        super(TemplatePairEmbedder, self).__init__()
+
+        self.c_in = c_in
+        self.c_out = c_out
+
+        # Despite there being no relu nearby, the source uses that initializer
+        self.linear = Linear(self.c_in, self.c_out, init="relu")
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x:
+                [*, C_in] input tensor
+        Returns:
+            [*, C_out] output tensor
+        """
+        x = self.linear(x)
+
+        return x
+
+
+class ExtraMSAEmbedder(nn.Module):
+    """
+    Embeds unclustered MSA sequences.
+
+    Implements Algorithm 2, line 15
+    """
+    def __init__(
+        self,
+        c_in: int,
+        c_out: int,
+        **kwargs,
+    ):
+        """
+        Args:
+            c_in:
+                Input channel dimension
+            c_out:
+                Output channel dimension
+        """
+        super(ExtraMSAEmbedder, self).__init__()
+
+        self.c_in = c_in
+        self.c_out = c_out
+
+        self.linear = Linear(self.c_in, self.c_out)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x:
+                [*, N_extra_seq, N_res, C_in] "extra_msa_feat" features
+        Returns:
+            [*, N_extra_seq, N_res, C_out] embedding
+        """
+        x = self.linear(x)
+
+        return x
+
+
+class TemplateEmbedder(nn.Module):
+    def __init__(self, config):
+        super(TemplateEmbedder, self).__init__()
+        
+        self.config = config
+        self.template_single_embedder = TemplateSingleEmbedder(
+            **config["template_single_embedder"],
+        )
+        self.template_pair_embedder = TemplatePairEmbedder(
+            **config["template_pair_embedder"],
+        )
+        self.template_pair_stack = TemplatePairStack(
+            **config["template_pair_stack"],
+        )
+        self.template_pointwise_att = TemplatePointwiseAttention(
+            **config["template_pointwise_attention"],
+        )
+    
+    def forward(
+        self,
+        batch,
+        z, 
+        pair_mask,
+        templ_dim,
+        chunk_size,
+        _mask_trans=True,
+        use_deepspeed_evo_attention=False,
+        use_lma=False,
+        inplace_safe=False
+    ):
+        # Embed the templates one at a time (with a poor man's vmap)
+        pair_embeds = []
+        n = z.shape[-2]
+        n_templ = batch["template_aatype"].shape[templ_dim]
+
+        if (inplace_safe):
+            # We'll preallocate the full pair tensor now to avoid manifesting
+            # a second copy during the stack later on
+            t_pair = z.new_zeros(
+                z.shape[:-3] +
+                (n_templ, n, n, self.config.template_pair_embedder.c_out)
+            )
+
+        for i in range(n_templ):
+            idx = batch["template_aatype"].new_tensor(i)
+            single_template_feats = tensor_tree_map(
+                lambda t: torch.index_select(t, templ_dim, idx).squeeze(templ_dim),
+                batch,
+            )
+
+            # [*, N, N, C_t]
+            t = build_template_pair_feat(
+                single_template_feats,
+                use_unit_vector=self.config.use_unit_vector,
+                inf=self.config.inf,
+                eps=self.config.eps,
+                **self.config.distogram,
+            ).to(z.dtype)
+            t = self.template_pair_embedder(t)
+
+            if (inplace_safe):
+                t_pair[..., i, :, :, :] = t
+            else:
+                pair_embeds.append(t)
+
+            del t
+
+        if (not inplace_safe):
+            t_pair = torch.stack(pair_embeds, dim=templ_dim)
+
+        del pair_embeds
+
+        # [*, S_t, N, N, C_z]
+        t = self.template_pair_stack(
+            t_pair,
+            pair_mask.unsqueeze(-3).to(dtype=z.dtype),
+            chunk_size=chunk_size,
+            use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+            use_lma=use_lma,
+            inplace_safe=inplace_safe,
+            _mask_trans=_mask_trans,
+        )
+        del t_pair
+
+        # [*, N, N, C_z]
+        t = self.template_pointwise_att(
+            t,
+            z,
+            template_mask=batch["template_mask"].to(dtype=z.dtype),
+            use_lma=use_lma,
+        )
+
+        t_mask = torch.sum(batch["template_mask"], dim=-1) > 0
+        # Append singletons
+        t_mask = t_mask.reshape(
+            *t_mask.shape, *([1] * (len(t.shape) - len(t_mask.shape)))
+        )
+
+        if (inplace_safe):
+            t *= t_mask
+        else:
+            t = t * t_mask
+
+        ret = {}
+
+        ret.update({"template_pair_embedding": t})
+
+        del t
+
+        if self.config.embed_angles:
+            template_angle_feat = build_template_angle_feat(
+                batch
+            )
+
+            # [*, S_t, N, C_m]
+            a = self.template_single_embedder(template_angle_feat)
+
+            ret["template_single_embedding"] = a
+
+        return ret
+
+
+class TemplatePairEmbedderMultimer(nn.Module):
+    def __init__(self,
+        c_in: int,
+        c_out: int,
+        c_dgram: int,
+        c_aatype: int,
+    ):
+        super(TemplatePairEmbedderMultimer, self).__init__()
+
+        self.dgram_linear = Linear(c_dgram, c_out, init='relu')
+        self.aatype_linear_1 = Linear(c_aatype, c_out, init='relu')
+        self.aatype_linear_2 = Linear(c_aatype, c_out, init='relu')
+        self.query_embedding_layer_norm = LayerNorm(c_in)
+        self.query_embedding_linear = Linear(c_in, c_out, init='relu')
+        
+        self.pseudo_beta_mask_linear = Linear(1, c_out, init='relu')
+        self.x_linear = Linear(1, c_out, init='relu')
+        self.y_linear = Linear(1, c_out, init='relu')
+        self.z_linear = Linear(1, c_out, init='relu')
+        self.backbone_mask_linear = Linear(1, c_out, init='relu')
+
+    def forward(self,
+        template_dgram: torch.Tensor,
+        aatype_one_hot: torch.Tensor,
+        query_embedding: torch.Tensor,
+        pseudo_beta_mask: torch.Tensor,
+        backbone_mask: torch.Tensor,
+        multichain_mask_2d: torch.Tensor,
+        unit_vector: geometry.Vec3Array,
+    ) -> torch.Tensor:
+        act = 0.
+
+        pseudo_beta_mask_2d = (
+            pseudo_beta_mask[..., None] * pseudo_beta_mask[..., None, :]
+        )
+        pseudo_beta_mask_2d *= multichain_mask_2d
+        template_dgram *= pseudo_beta_mask_2d[..., None]
+        act += self.dgram_linear(template_dgram)
+        act += self.pseudo_beta_mask_linear(pseudo_beta_mask_2d[..., None])
+       
+        aatype_one_hot = aatype_one_hot.to(template_dgram.dtype)
+        act += self.aatype_linear_1(aatype_one_hot[..., None, :, :])
+        act += self.aatype_linear_2(aatype_one_hot[..., None, :])
+
+        backbone_mask_2d = (
+            backbone_mask[..., None] * backbone_mask[..., None, :]
+        )
+        backbone_mask_2d *= multichain_mask_2d
+        x, y, z = [(coord * backbone_mask_2d).to(dtype=query_embedding.dtype) for coord in unit_vector]
+        act += self.x_linear(x[..., None])
+        act += self.y_linear(y[..., None])
+        act += self.z_linear(z[..., None])
+       
+        act += self.backbone_mask_linear(backbone_mask_2d[..., None].to(dtype=query_embedding.dtype))
+
+        query_embedding = self.query_embedding_layer_norm(query_embedding)
+        act += self.query_embedding_linear(query_embedding)
+
+        return act
+
+
+class TemplateSingleEmbedderMultimer(nn.Module):
+    def __init__(self,
+        c_in: int,
+        c_out: int,
+    ):
+        super(TemplateSingleEmbedderMultimer, self).__init__()
+        self.template_single_embedder = Linear(c_in, c_out)
+        self.template_projector = Linear(c_out, c_out)
+    
+    def forward(self,
+        batch,
+        atom_pos,
+        aatype_one_hot,
+    ):
+        out = {}
+
+        dtype = batch["template_all_atom_positions"].dtype
+
+        template_chi_angles, template_chi_mask = (
+            all_atom_multimer.compute_chi_angles(
+                atom_pos,
+                batch["template_all_atom_mask"],
+                batch["template_aatype"],
+            )
+        )
+
+        template_features = torch.cat(
+            [
+                aatype_one_hot,
+                torch.sin(template_chi_angles) * template_chi_mask,
+                torch.cos(template_chi_angles) * template_chi_mask,
+                template_chi_mask,
+            ],
+            dim=-1,
+        ).to(dtype=dtype)
+
+        template_mask = template_chi_mask[..., 0].to(dtype=dtype)
+
+        template_activations = self.template_single_embedder(
+            template_features
+        )
+        template_activations = torch.nn.functional.relu(
+            template_activations
+        )
+        template_activations = self.template_projector(
+            template_activations,
+        )
+
+        out["template_single_embedding"] = (
+            template_activations
+        )
+        out["template_mask"] = template_mask
+
+        return out 
+
+ 
+class TemplateEmbedderMultimer(nn.Module):
+    def __init__(self, config):
+        super(TemplateEmbedderMultimer, self).__init__()
+        
+        self.config = config
+        self.template_pair_embedder = TemplatePairEmbedderMultimer(
+            **config["template_pair_embedder"],
+        )
+        self.template_single_embedder = TemplateSingleEmbedderMultimer(
+            **config["template_single_embedder"],
+        )
+        self.template_pair_stack = TemplatePairStack(
+            **config["template_pair_stack"],
+        )
+
+        self.linear_t = Linear(config.c_t, config.c_z)
+    
+    def forward(self, 
+        batch, 
+        z, 
+        padding_mask_2d, 
+        templ_dim,
+        chunk_size,
+        multichain_mask_2d,
+        _mask_trans=True,
+        use_deepspeed_evo_attention=False,
+        use_lma=False,
+        inplace_safe=False
+    ):
+        template_embeds = []
+        n_templ = batch["template_aatype"].shape[templ_dim]
+        for i in range(n_templ):
+            idx = batch["template_aatype"].new_tensor(i)
+            single_template_feats = tensor_tree_map(
+                lambda t: torch.index_select(t, templ_dim, idx),
+                batch,
+            )
+
+            single_template_embeds = {}
+            act = 0.
+
+            template_positions, pseudo_beta_mask = pseudo_beta_fn(
+                single_template_feats["template_aatype"],
+                single_template_feats["template_all_atom_positions"],
+                single_template_feats["template_all_atom_mask"])
+
+            template_dgram = dgram_from_positions(
+                template_positions,
+                inf=self.config.inf,
+                **self.config.distogram,
+            )
+
+            aatype_one_hot = torch.nn.functional.one_hot(
+                single_template_feats["template_aatype"], 22,
+            )
+            
+            raw_atom_pos = single_template_feats["template_all_atom_positions"]
+
+            # Vec3Arrays are required to be float32
+            atom_pos = geometry.Vec3Array.from_array(raw_atom_pos.to(dtype=torch.float32))
+
+            rigid, backbone_mask = all_atom_multimer.make_backbone_affine(
+                atom_pos,
+                single_template_feats["template_all_atom_mask"],
+                single_template_feats["template_aatype"],
+            )
+            points = rigid.translation
+            rigid_vec = rigid[..., None].inverse().apply_to_point(points)
+            unit_vector = rigid_vec.normalized()
+
+            pair_act = self.template_pair_embedder(
+                template_dgram,
+                aatype_one_hot,
+                z,
+                pseudo_beta_mask,
+                backbone_mask,
+                multichain_mask_2d,
+                unit_vector,
+            )
+
+            single_template_embeds["template_pair_embedding"] = pair_act
+            single_template_embeds.update(
+                self.template_single_embedder(
+                    single_template_feats,
+                    atom_pos,
+                    aatype_one_hot,
+                )
+            )
+            template_embeds.append(single_template_embeds)
+
+        template_embeds = dict_multimap(
+            partial(torch.cat, dim=templ_dim),
+            template_embeds,
+        )
+
+        # [*, S_t, N, N, C_z]
+        t = self.template_pair_stack(
+            template_embeds["template_pair_embedding"], 
+            padding_mask_2d.unsqueeze(-3).to(dtype=z.dtype), 
+            chunk_size=chunk_size,
+            use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+            use_lma=use_lma,
+            inplace_safe=inplace_safe,
+            _mask_trans=_mask_trans,
+        )
+        # [*, N, N, C_z]
+        t = torch.sum(t, dim=-4) / n_templ
+        t = torch.nn.functional.relu(t)
+        t = self.linear_t(t)
+        template_embeds["template_pair_embedding"] = t
+
+        return template_embeds

model/openfold/evoformer.py

@@ -0,0 +1,1219 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+import sys
+import torch
+import torch.nn as nn
+from typing import Tuple, Sequence, Optional
+from functools import partial
+from abc import ABC, abstractmethod
+
+from model.openfold.primitives import Linear, LayerNorm
+from model.openfold.dropout import DropoutRowwise, DropoutColumnwise
+from model.openfold.msa import (
+    MSARowAttentionWithPairBias,
+    MSAColumnAttention,
+    MSAColumnGlobalAttention,
+)
+from model.openfold.outer_product_mean import OuterProductMean
+from model.openfold.pair_transition import PairTransition
+from model.openfold.triangular_attention import (
+    TriangleAttention,
+    TriangleAttentionStartingNode,
+    TriangleAttentionEndingNode,
+)
+from model.openfold.triangular_multiplicative_update import (
+    TriangleMultiplicationOutgoing,
+    TriangleMultiplicationIncoming,
+    FusedTriangleMultiplicationIncoming,
+    FusedTriangleMultiplicationOutgoing
+)
+from onescience.utils.openfold.checkpointing import checkpoint_blocks, get_checkpoint_fn
+from onescience.utils.openfold.chunk_utils import chunk_layer, ChunkSizeTuner
+from onescience.utils.openfold.tensor_utils import add
+
+
+class MSATransition(nn.Module):
+    """
+    Feed-forward network applied to MSA activations after attention.
+
+    Implements Algorithm 9
+    """
+    def __init__(self, c_m, n):
+        """
+        Args:
+            c_m:
+                MSA channel dimension
+            n:
+                Factor multiplied to c_m to obtain the hidden channel
+                dimension
+        """
+        super(MSATransition, self).__init__()
+
+        self.c_m = c_m
+        self.n = n
+
+        self.layer_norm = LayerNorm(self.c_m)
+        self.linear_1 = Linear(self.c_m, self.n * self.c_m, init="relu")
+        self.relu = nn.ReLU()
+        self.linear_2 = Linear(self.n * self.c_m, self.c_m, init="final")
+
+    def _transition(self, m, mask):
+        m = self.layer_norm(m)
+        m = self.linear_1(m)
+        m = self.relu(m)
+        m = self.linear_2(m) * mask
+        return m
+
+    @torch.jit.ignore
+    def _chunk(self,
+        m: torch.Tensor,
+        mask: torch.Tensor,
+        chunk_size: int,
+    ) -> torch.Tensor:
+         return chunk_layer(
+             self._transition,
+             {"m": m, "mask": mask},
+             chunk_size=chunk_size,
+             no_batch_dims=len(m.shape[:-2]),
+         )
+
+    def forward(
+        self,
+        m: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        chunk_size: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            m:
+                [*, N_seq, N_res, C_m] MSA activation
+            mask:
+                [*, N_seq, N_res, C_m] MSA mask
+        Returns:
+            m:
+                [*, N_seq, N_res, C_m] MSA activation update
+        """
+        # DISCREPANCY: DeepMind forgets to apply the MSA mask here.
+        if mask is None:
+            mask = m.new_ones(m.shape[:-1])
+
+        mask = mask.unsqueeze(-1)
+
+        if chunk_size is not None:
+            m = self._chunk(m, mask, chunk_size)
+        else:
+            m = self._transition(m, mask)
+
+        return m
+
+
+class PairStack(nn.Module):
+    def __init__(
+        self,
+        c_z: int,
+        c_hidden_mul: int,
+        c_hidden_pair_att: int,
+        no_heads_pair: int,
+        transition_n: int,
+        pair_dropout: float,
+        fuse_projection_weights: bool,
+        inf: float,
+        eps: float
+    ):
+        super(PairStack, self).__init__()
+
+        if fuse_projection_weights:
+            self.tri_mul_out = FusedTriangleMultiplicationOutgoing(
+                c_z,
+                c_hidden_mul,
+            )
+            self.tri_mul_in = FusedTriangleMultiplicationIncoming(
+                c_z,
+                c_hidden_mul,
+            )
+        else:
+            self.tri_mul_out = TriangleMultiplicationOutgoing(
+                c_z,
+                c_hidden_mul,
+            )
+            self.tri_mul_in = TriangleMultiplicationIncoming(
+                c_z,
+                c_hidden_mul,
+            )
+
+        self.tri_att_start = TriangleAttention(
+            c_z,
+            c_hidden_pair_att,
+            no_heads_pair,
+            inf=inf,
+        )
+        self.tri_att_end = TriangleAttention(
+            c_z,
+            c_hidden_pair_att,
+            no_heads_pair,
+            inf=inf,
+        )
+
+        self.pair_transition = PairTransition(
+            c_z,
+            transition_n,
+        )
+
+        self.ps_dropout_row_layer = DropoutRowwise(pair_dropout)
+
+    def forward(self,
+        z: torch.Tensor,
+        pair_mask: torch.Tensor,
+        chunk_size: Optional[int] = None,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        inplace_safe: bool = False,
+        _mask_trans: bool = True,
+        _attn_chunk_size: Optional[int] = None
+    ) -> torch.Tensor:
+        # DeepMind doesn't mask these transitions in the source, so _mask_trans
+        # should be disabled to better approximate the exact activations of
+        # the original.
+        pair_trans_mask = pair_mask if _mask_trans else None
+
+        if (_attn_chunk_size is None):
+            _attn_chunk_size = chunk_size
+
+        tmu_update = self.tri_mul_out(
+            z,
+            mask=pair_mask,
+            inplace_safe=inplace_safe,
+            _add_with_inplace=True,
+        )
+        if (not inplace_safe):
+            z = z + self.ps_dropout_row_layer(tmu_update)
+        else:
+            z = tmu_update
+
+        del tmu_update
+
+        tmu_update = self.tri_mul_in(
+            z,
+            mask=pair_mask,
+            inplace_safe=inplace_safe,
+            _add_with_inplace=True,
+        )
+        if (not inplace_safe):
+            z = z + self.ps_dropout_row_layer(tmu_update)
+        else:
+            z = tmu_update
+
+        del tmu_update
+
+        z = add(z,
+                self.ps_dropout_row_layer(
+                    self.tri_att_start(
+                        z,
+                        mask=pair_mask,
+                        chunk_size=_attn_chunk_size,
+                        use_memory_efficient_kernel=False,
+                        use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                        use_lma=use_lma,
+                        inplace_safe=inplace_safe,
+                    )
+                ),
+                inplace=inplace_safe,
+                )
+
+        z = z.transpose(-2, -3)
+        if (inplace_safe):
+            z = z.contiguous()
+
+        z = add(z,
+                self.ps_dropout_row_layer(
+                    self.tri_att_end(
+                        z,
+                        mask=pair_mask.transpose(-1, -2),
+                        chunk_size=_attn_chunk_size,
+                        use_memory_efficient_kernel=False,
+                        use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                        use_lma=use_lma,
+                        inplace_safe=inplace_safe,
+                    )
+                ),
+                inplace=inplace_safe,
+                )
+
+        z = z.transpose(-2, -3)
+        if (inplace_safe):
+            z = z.contiguous()
+
+        z = add(z,
+                self.pair_transition(
+                    z, mask=pair_trans_mask, chunk_size=chunk_size,
+                ),
+                inplace=inplace_safe,
+        )
+
+        return z
+
+
+class MSABlock(nn.Module, ABC):
+    @abstractmethod
+    def __init__(self,
+        c_m: int,
+        c_z: int,
+        c_hidden_msa_att: int,
+        c_hidden_opm: int,
+        c_hidden_mul: int,
+        c_hidden_pair_att: int,
+        no_heads_msa: int,
+        no_heads_pair: int,
+        transition_n: int,
+        msa_dropout: float,
+        pair_dropout: float,
+        opm_first: bool,
+        fuse_projection_weights: bool,
+        inf: float,
+        eps: float,
+    ):
+        super(MSABlock, self).__init__()
+
+        self.opm_first = opm_first
+
+        self.msa_att_row = MSARowAttentionWithPairBias(
+            c_m=c_m,
+            c_z=c_z,
+            c_hidden=c_hidden_msa_att,
+            no_heads=no_heads_msa,
+            inf=inf,
+        )
+
+        self.msa_dropout_layer = DropoutRowwise(msa_dropout)
+
+        self.msa_transition = MSATransition(
+            c_m=c_m,
+            n=transition_n,
+        )
+
+        self.outer_product_mean = OuterProductMean(
+            c_m,
+            c_z,
+            c_hidden_opm,
+        )
+
+        self.pair_stack = PairStack(
+            c_z=c_z,
+            c_hidden_mul=c_hidden_mul,
+            c_hidden_pair_att=c_hidden_pair_att,
+            no_heads_pair=no_heads_pair,
+            transition_n=transition_n,
+            pair_dropout=pair_dropout,
+            fuse_projection_weights=fuse_projection_weights,
+            inf=inf,
+            eps=eps
+        )
+
+    def _compute_opm(self,
+        input_tensors: Sequence[torch.Tensor],
+        msa_mask: torch.Tensor,
+        chunk_size: Optional[int] = None,
+        inplace_safe: bool = False,
+        _offload_inference: bool = False
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        m, z = input_tensors
+
+        if (_offload_inference and inplace_safe):
+            # m: GPU, z: CPU
+            del m, z
+            assert (sys.getrefcount(input_tensors[1]) == 2)
+            input_tensors[1] = input_tensors[1].cpu()
+            m, z = input_tensors
+
+        opm = self.outer_product_mean(
+            m, mask=msa_mask, chunk_size=chunk_size, inplace_safe=inplace_safe
+        )
+
+        if (_offload_inference and inplace_safe):
+            # m: GPU, z: GPU
+            del m, z
+            assert (sys.getrefcount(input_tensors[0]) == 2)
+            input_tensors[1] = input_tensors[1].to(opm.device)
+            m, z = input_tensors
+
+        z = add(z, opm, inplace=inplace_safe)
+        del opm
+
+        return m, z
+
+    @abstractmethod
+    def forward(self,
+        m: Optional[torch.Tensor],
+        z: Optional[torch.Tensor],
+        msa_mask: torch.Tensor,
+        pair_mask: torch.Tensor,
+        chunk_size: Optional[int] = None,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        use_flash: bool = False,
+        inplace_safe: bool = False,
+        _mask_trans: bool = True,
+        _attn_chunk_size: Optional[int] = None,
+        _offload_inference: bool = False,
+        _offloadable_inputs: Optional[Sequence[torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        pass
+
+
+class EvoformerBlock(MSABlock):
+    def __init__(self,
+        c_m: int,
+        c_z: int,
+        c_hidden_msa_att: int,
+        c_hidden_opm: int,
+        c_hidden_mul: int,
+        c_hidden_pair_att: int,
+        no_heads_msa: int,
+        no_heads_pair: int,
+        transition_n: int,
+        msa_dropout: float,
+        pair_dropout: float,
+        no_column_attention: bool,
+        opm_first: bool,
+        fuse_projection_weights: bool,
+        inf: float,
+        eps: float,
+    ):
+        super(EvoformerBlock, self).__init__(c_m=c_m,
+                                             c_z=c_z,
+                                             c_hidden_msa_att=c_hidden_msa_att,
+                                             c_hidden_opm=c_hidden_opm,
+                                             c_hidden_mul=c_hidden_mul,
+                                             c_hidden_pair_att=c_hidden_pair_att,
+                                             no_heads_msa=no_heads_msa,
+                                             no_heads_pair=no_heads_pair,
+                                             transition_n=transition_n,
+                                             msa_dropout=msa_dropout,
+                                             pair_dropout=pair_dropout,
+                                             opm_first=opm_first,
+                                             fuse_projection_weights=fuse_projection_weights,
+                                             inf=inf,
+                                             eps=eps)
+
+        # Specifically, seqemb mode does not use column attention
+        self.no_column_attention = no_column_attention
+
+        if not self.no_column_attention:
+            self.msa_att_col = MSAColumnAttention(
+                c_m,
+                c_hidden_msa_att,
+                no_heads_msa,
+                inf=inf,
+            )
+
+    def forward(self,
+        m: Optional[torch.Tensor],
+        z: Optional[torch.Tensor],
+        msa_mask: torch.Tensor,
+        pair_mask: torch.Tensor,
+        chunk_size: Optional[int] = None,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        use_flash: bool = False,
+        inplace_safe: bool = False,
+        _mask_trans: bool = True,
+        _attn_chunk_size: Optional[int] = None,
+        _offload_inference: bool = False,
+        _offloadable_inputs: Optional[Sequence[torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        msa_trans_mask = msa_mask if _mask_trans else None
+
+        if(_attn_chunk_size is None):
+            _attn_chunk_size = chunk_size
+
+        if(_offload_inference and inplace_safe):
+            input_tensors = _offloadable_inputs
+            del _offloadable_inputs
+        else:
+            input_tensors = [m, z]
+
+        m, z = input_tensors
+
+        if self.opm_first:
+            del m, z
+
+            m, z = self._compute_opm(input_tensors=input_tensors,
+                                     msa_mask=msa_mask,
+                                     chunk_size=chunk_size,
+                                     inplace_safe=inplace_safe,
+                                     _offload_inference=_offload_inference)
+
+        m = add(m,
+                self.msa_dropout_layer(
+                    self.msa_att_row(
+                        m,
+                        z=z,
+                        mask=msa_mask,
+                        chunk_size=_attn_chunk_size,
+                        use_memory_efficient_kernel=False,
+                        use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                        use_lma=use_lma,
+                    )
+                ),
+                inplace=inplace_safe,
+                )
+
+        if (_offload_inference and inplace_safe):
+            # m: GPU, z: CPU
+            del m, z
+            assert (sys.getrefcount(input_tensors[1]) == 2)
+            input_tensors[1] = input_tensors[1].cpu()
+            torch.cuda.empty_cache()
+            m, z = input_tensors
+
+        # Specifically, column attention is not used in seqemb mode.
+        if not self.no_column_attention:
+            m = add(m,
+                    self.msa_att_col(
+                        m,
+                        mask=msa_mask,
+                        chunk_size=chunk_size,
+                        use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                        use_lma=use_lma,
+                        use_flash=use_flash,
+                    ),
+                    inplace=inplace_safe,
+                    )
+
+        m = add(
+            m,
+            self.msa_transition(
+                m, mask=msa_trans_mask, chunk_size=chunk_size,
+            ),
+            inplace=inplace_safe,
+        )
+
+        if not self.opm_first:
+            if (not inplace_safe):
+                input_tensors = [m, z]
+
+            del m, z
+
+            m, z = self._compute_opm(input_tensors=input_tensors,
+                                     msa_mask=msa_mask,
+                                     chunk_size=chunk_size,
+                                     inplace_safe=inplace_safe,
+                                     _offload_inference=_offload_inference)
+
+        if (_offload_inference and inplace_safe):
+            # m: CPU, z: GPU
+            del m, z
+            assert (sys.getrefcount(input_tensors[0]) == 2)
+            device = input_tensors[0].device
+            input_tensors[0] = input_tensors[0].cpu()
+            input_tensors[1] = input_tensors[1].to(device)
+            m, z = input_tensors
+
+        if (not inplace_safe):
+            input_tensors = [m, z]
+
+        del m, z
+
+        z = self.pair_stack(
+            z=input_tensors[1],
+            pair_mask=pair_mask,
+            chunk_size=chunk_size,
+            use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+            use_lma=use_lma,
+            inplace_safe=inplace_safe,
+            _mask_trans=_mask_trans,
+            _attn_chunk_size=_attn_chunk_size
+        )
+
+        if (_offload_inference and inplace_safe):
+            # m: GPU, z: GPU
+            device = z.device
+            assert (sys.getrefcount(input_tensors[0]) == 2)
+            input_tensors[0] = input_tensors[0].to(device)
+            m, _ = input_tensors
+        else:
+            m = input_tensors[0]
+
+        return m, z
+
+
+class ExtraMSABlock(MSABlock):
+    """ 
+        Almost identical to the standard EvoformerBlock, except in that the
+        ExtraMSABlock uses GlobalAttention for MSA column attention and
+        requires more fine-grained control over checkpointing. Separated from
+        its twin to preserve the TorchScript-ability of the latter.
+    """
+    def __init__(self,
+        c_m: int,
+        c_z: int,
+        c_hidden_msa_att: int,
+        c_hidden_opm: int,
+        c_hidden_mul: int,
+        c_hidden_pair_att: int,
+        no_heads_msa: int,
+        no_heads_pair: int,
+        transition_n: int,
+        msa_dropout: float,
+        pair_dropout: float,
+        opm_first: bool,
+        fuse_projection_weights: bool,
+        inf: float,
+        eps: float,
+        ckpt: bool,
+    ):
+        super(ExtraMSABlock, self).__init__(c_m=c_m,
+                                            c_z=c_z,
+                                            c_hidden_msa_att=c_hidden_msa_att,
+                                            c_hidden_opm=c_hidden_opm,
+                                            c_hidden_mul=c_hidden_mul,
+                                            c_hidden_pair_att=c_hidden_pair_att,
+                                            no_heads_msa=no_heads_msa,
+                                            no_heads_pair=no_heads_pair,
+                                            transition_n=transition_n,
+                                            msa_dropout=msa_dropout,
+                                            pair_dropout=pair_dropout,
+                                            opm_first=opm_first,
+                                            fuse_projection_weights=fuse_projection_weights,
+                                            inf=inf,
+                                            eps=eps)
+
+        self.ckpt = ckpt
+
+        self.msa_att_col = MSAColumnGlobalAttention(
+            c_in=c_m,
+            c_hidden=c_hidden_msa_att,
+            no_heads=no_heads_msa,
+            inf=inf,
+            eps=eps,
+        )
+
+    def forward(self,
+        m: Optional[torch.Tensor],
+        z: Optional[torch.Tensor],
+        msa_mask: torch.Tensor,
+        pair_mask: torch.Tensor,
+        chunk_size: Optional[int] = None,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        inplace_safe: bool = False,
+        _mask_trans: bool = True,
+        _attn_chunk_size: Optional[int] = None,
+        _offload_inference: bool = False,
+        _offloadable_inputs: Optional[Sequence[torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if(_attn_chunk_size is None):
+            _attn_chunk_size = chunk_size
+       
+        if(_offload_inference and inplace_safe):
+            input_tensors = _offloadable_inputs
+            del _offloadable_inputs
+        else:
+            input_tensors = [m, z]
+
+        m, z = input_tensors
+
+        if self.opm_first:
+            del m, z
+
+            m, z = self._compute_opm(input_tensors=input_tensors,
+                                     msa_mask=msa_mask,
+                                     chunk_size=chunk_size,
+                                     inplace_safe=inplace_safe,
+                                     _offload_inference=_offload_inference)
+
+        m = add(m,
+            self.msa_dropout_layer(
+                self.msa_att_row(
+                    m.clone() if torch.is_grad_enabled() else m, 
+                    z=z.clone() if torch.is_grad_enabled() else z, 
+                    mask=msa_mask, 
+                    chunk_size=_attn_chunk_size,
+                    use_lma=use_lma,
+                    use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                    use_memory_efficient_kernel=not (use_lma or use_deepspeed_evo_attention),
+                    _checkpoint_chunks=
+                        self.ckpt if torch.is_grad_enabled() else False,
+                )
+            ),
+            inplace=inplace_safe,
+        )
+
+        if (not inplace_safe):
+            input_tensors = [m, z]
+
+        del m, z
+
+        def fn(input_tensors):
+            m, z = input_tensors
+
+            if (_offload_inference and inplace_safe):
+                # m: GPU, z: CPU
+                del m, z
+                assert (sys.getrefcount(input_tensors[1]) == 2)
+                input_tensors[1] = input_tensors[1].cpu()
+                torch.cuda.empty_cache()
+                m, z = input_tensors
+
+            m = add(m,
+                    self.msa_att_col(
+                        m,
+                        mask=msa_mask,
+                        chunk_size=chunk_size,
+                        use_lma=use_lma,
+                    ),
+                    inplace=inplace_safe,
+                    )
+
+            m = add(
+                m,
+                self.msa_transition(
+                    m, mask=msa_mask, chunk_size=chunk_size,
+                ),
+                inplace=inplace_safe,
+            )
+
+            if not self.opm_first:
+                if (not inplace_safe):
+                    input_tensors = [m, z]
+
+                del m, z
+
+                m, z = self._compute_opm(input_tensors=input_tensors,
+                                         msa_mask=msa_mask,
+                                         chunk_size=chunk_size,
+                                         inplace_safe=inplace_safe,
+                                         _offload_inference=_offload_inference)
+
+            if (_offload_inference and inplace_safe):
+                # m: CPU, z: GPU
+                del m, z
+                assert (sys.getrefcount(input_tensors[0]) == 2)
+                device = input_tensors[0].device
+                input_tensors[0] = input_tensors[0].cpu()
+                input_tensors[1] = input_tensors[1].to(device)
+                m, z = input_tensors
+
+            if (not inplace_safe):
+                input_tensors = [m, z]
+
+            del m, z
+
+            z = self.pair_stack(
+                input_tensors[1],
+                pair_mask=pair_mask,
+                chunk_size=chunk_size,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma,
+                inplace_safe=inplace_safe,
+                _mask_trans=_mask_trans,
+                _attn_chunk_size=_attn_chunk_size
+            )
+
+            m = input_tensors[0]
+            if (_offload_inference and inplace_safe):
+                # m: GPU, z: GPU
+                device = z.device
+                del m
+                assert (sys.getrefcount(input_tensors[0]) == 2)
+                input_tensors[0] = input_tensors[0].to(device)
+                m, _ = input_tensors
+
+            return m, z
+
+        if (torch.is_grad_enabled() and self.ckpt):
+            checkpoint_fn = get_checkpoint_fn()
+            m, z = checkpoint_fn(fn, input_tensors)
+        else:
+            m, z = fn(input_tensors)
+
+        return m, z
+
+
+class EvoformerStack(nn.Module):
+    """
+    Main Evoformer trunk.
+
+    Implements Algorithm 6.
+    """
+
+    def __init__(
+        self,
+        c_m: int,
+        c_z: int,
+        c_hidden_msa_att: int,
+        c_hidden_opm: int,
+        c_hidden_mul: int,
+        c_hidden_pair_att: int,
+        c_s: int,
+        no_heads_msa: int,
+        no_heads_pair: int,
+        no_blocks: int,
+        transition_n: int,
+        msa_dropout: float,
+        pair_dropout: float,
+        no_column_attention: bool,
+        opm_first: bool,
+        fuse_projection_weights: bool,
+        blocks_per_ckpt: int,
+        inf: float,
+        eps: float,
+        clear_cache_between_blocks: bool = False, 
+        tune_chunk_size: bool = False,
+        **kwargs,
+    ):
+        """
+        Args:
+            c_m:
+                MSA channel dimension
+            c_z:
+                Pair channel dimension
+            c_hidden_msa_att:
+                Hidden dimension in MSA attention
+            c_hidden_opm:
+                Hidden dimension in outer product mean module
+            c_hidden_mul:
+                Hidden dimension in multiplicative updates
+            c_hidden_pair_att:
+                Hidden dimension in triangular attention
+            c_s:
+                Channel dimension of the output "single" embedding
+            no_heads_msa:
+                Number of heads used for MSA attention
+            no_heads_pair:
+                Number of heads used for pair attention
+            no_blocks:
+                Number of Evoformer blocks in the stack
+            transition_n:
+                Factor by which to multiply c_m to obtain the MSATransition
+                hidden dimension
+            msa_dropout:
+                Dropout rate for MSA activations
+            pair_dropout:
+                Dropout used for pair activations
+            no_column_attention:
+                When True, doesn't use column attention. Required for running
+                sequence embedding mode
+            opm_first:
+                When True, Outer Product Mean is performed at the beginning of
+                the Evoformer block instead of after the MSA Stack.
+                Used in Multimer pipeline.
+            fuse_projection_weights:
+                When True, uses FusedTriangleMultiplicativeUpdate variant in
+                the Pair Stack. Used in Multimer pipeline.
+            blocks_per_ckpt:
+                Number of Evoformer blocks in each activation checkpoint
+            clear_cache_between_blocks:
+                Whether to clear CUDA's GPU memory cache between blocks of the
+                stack. Slows down each block but can reduce fragmentation
+            tune_chunk_size:
+                Whether to dynamically tune the module's chunk size
+        """
+        super(EvoformerStack, self).__init__()
+
+        self.blocks_per_ckpt = blocks_per_ckpt
+        self.clear_cache_between_blocks = clear_cache_between_blocks
+
+        self.blocks = nn.ModuleList()
+
+        for _ in range(no_blocks):
+            block = EvoformerBlock(
+                c_m=c_m,
+                c_z=c_z,
+                c_hidden_msa_att=c_hidden_msa_att,
+                c_hidden_opm=c_hidden_opm,
+                c_hidden_mul=c_hidden_mul,
+                c_hidden_pair_att=c_hidden_pair_att,
+                no_heads_msa=no_heads_msa,
+                no_heads_pair=no_heads_pair,
+                transition_n=transition_n,
+                msa_dropout=msa_dropout,
+                pair_dropout=pair_dropout,
+                no_column_attention=no_column_attention,
+                opm_first=opm_first,
+                fuse_projection_weights=fuse_projection_weights,
+                inf=inf,
+                eps=eps,
+            )
+            self.blocks.append(block)
+
+        self.linear = Linear(c_m, c_s)
+
+        self.tune_chunk_size = tune_chunk_size
+        self.chunk_size_tuner = None
+        if(tune_chunk_size):
+            self.chunk_size_tuner = ChunkSizeTuner()
+
+    def _prep_blocks(self, 
+        m: torch.Tensor, 
+        z: torch.Tensor, 
+        chunk_size: int,
+        use_deepspeed_evo_attention: bool,
+        use_lma: bool,
+        use_flash: bool,
+        msa_mask: Optional[torch.Tensor],
+        pair_mask: Optional[torch.Tensor],
+        inplace_safe: bool,
+        _mask_trans: bool,
+    ):
+        blocks = [
+            partial(
+                b,
+                msa_mask=msa_mask,
+                pair_mask=pair_mask,
+                chunk_size=chunk_size,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma,
+                use_flash=use_flash,
+                inplace_safe=inplace_safe,
+                _mask_trans=_mask_trans,
+            )
+            for b in self.blocks
+        ]
+
+        if(self.clear_cache_between_blocks):
+            def block_with_cache_clear(block, *args, **kwargs):
+                torch.cuda.empty_cache()
+                return block(*args, **kwargs)
+
+            blocks = [partial(block_with_cache_clear, b) for b in blocks]
+
+        if(chunk_size is not None and self.chunk_size_tuner is not None):
+            assert(not self.training)
+            tuned_chunk_size = self.chunk_size_tuner.tune_chunk_size(
+                representative_fn=blocks[0],
+                # We don't want to write in-place during chunk tuning runs
+                args=(m.clone(), z.clone(),),
+                min_chunk_size=chunk_size,
+            )
+            blocks = [
+                partial(b, 
+                    chunk_size=tuned_chunk_size,
+                    # A temporary measure to address torch's occasional
+                    # inability to allocate large tensors
+                    _attn_chunk_size=max(chunk_size, tuned_chunk_size // 4),
+                ) for b in blocks
+            ]
+
+        return blocks
+
+    def _forward_offload(self,
+        input_tensors: Sequence[torch.Tensor],
+        msa_mask: torch.Tensor,
+        pair_mask: torch.Tensor,
+        chunk_size: int,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        use_flash: bool = False,
+        _mask_trans: bool = True,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        assert(not (self.training or torch.is_grad_enabled()))
+        blocks = self._prep_blocks(
+            # We are very careful not to create references to these tensors in
+            # this function
+            m=input_tensors[0],
+            z=input_tensors[1],
+            chunk_size=chunk_size,
+            use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+            use_lma=use_lma,
+            use_flash=use_flash,
+            msa_mask=msa_mask,
+            pair_mask=pair_mask,
+            inplace_safe=True,
+            _mask_trans=_mask_trans,
+        )
+
+        for b in blocks:
+            m, z = b(
+                None, 
+                None, 
+                _offload_inference=True,
+                _offloadable_inputs=input_tensors,
+            )
+            input_tensors[0] = m
+            input_tensors[1] = z
+            del m, z
+        
+        m, z = input_tensors
+        
+        s = self.linear(m[..., 0, :, :])
+        
+        return m, z, s
+
+    def forward(self,
+        m: torch.Tensor,
+        z: torch.Tensor,
+        msa_mask: torch.Tensor,
+        pair_mask: torch.Tensor,
+        chunk_size: int,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        use_flash: bool = False,
+        inplace_safe: bool = False,
+        _mask_trans: bool = True,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            m:
+                [*, N_seq, N_res, C_m] MSA embedding
+            z:
+                [*, N_res, N_res, C_z] pair embedding
+            msa_mask:
+                [*, N_seq, N_res] MSA mask
+            pair_mask:
+                [*, N_res, N_res] pair mask
+            chunk_size: 
+                Inference-time subbatch size. Acts as a minimum if 
+                self.tune_chunk_size is True
+            use_deepspeed_evo_attention:
+                Whether to use DeepSpeed memory efficient kernel.
+                Mutually exclusive with use_lma and use_flash.
+            use_lma:
+                Whether to use low-memory attention during inference.
+                Mutually exclusive with use_flash and use_deepspeed_evo_attention.
+            use_flash: 
+                Whether to use FlashAttention where possible. Mutually 
+                exclusive with use_lma and use_deepspeed_evo_attention.
+        Returns:
+            m:
+                [*, N_seq, N_res, C_m] MSA embedding
+            z:
+                [*, N_res, N_res, C_z] pair embedding
+            s:
+                [*, N_res, C_s] single embedding (or None if extra MSA stack)
+        """ 
+        blocks = self._prep_blocks(
+            m=m,
+            z=z,
+            chunk_size=chunk_size,
+            use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+            use_lma=use_lma,
+            use_flash=use_flash,
+            msa_mask=msa_mask,
+            pair_mask=pair_mask,
+            inplace_safe=inplace_safe,
+            _mask_trans=_mask_trans,
+        )
+
+        blocks_per_ckpt = self.blocks_per_ckpt
+        if(not torch.is_grad_enabled()):
+            blocks_per_ckpt = None
+        
+        m, z = checkpoint_blocks(
+            blocks,
+            args=(m, z),
+            blocks_per_ckpt=blocks_per_ckpt,
+        )
+
+        s = self.linear(m[..., 0, :, :])
+
+        return m, z, s
+
+
+class ExtraMSAStack(nn.Module):
+    """
+    Implements Algorithm 18.
+    """
+    def __init__(self,
+        c_m: int,
+        c_z: int,
+        c_hidden_msa_att: int,
+        c_hidden_opm: int,
+        c_hidden_mul: int,
+        c_hidden_pair_att: int,
+        no_heads_msa: int,
+        no_heads_pair: int,
+        no_blocks: int,
+        transition_n: int,
+        msa_dropout: float,
+        pair_dropout: float,
+        opm_first: bool,
+        fuse_projection_weights: bool,
+        inf: float,
+        eps: float,
+        ckpt: bool,
+        clear_cache_between_blocks: bool = False,
+        tune_chunk_size: bool = False,
+        **kwargs,
+    ):
+        super(ExtraMSAStack, self).__init__()
+ 
+        self.ckpt = ckpt
+        self.clear_cache_between_blocks = clear_cache_between_blocks
+        self.blocks = nn.ModuleList()
+        for _ in range(no_blocks):
+            block = ExtraMSABlock(
+                c_m=c_m,
+                c_z=c_z,
+                c_hidden_msa_att=c_hidden_msa_att,
+                c_hidden_opm=c_hidden_opm,
+                c_hidden_mul=c_hidden_mul,
+                c_hidden_pair_att=c_hidden_pair_att,
+                no_heads_msa=no_heads_msa,
+                no_heads_pair=no_heads_pair,
+                transition_n=transition_n,
+                msa_dropout=msa_dropout,
+                pair_dropout=pair_dropout,
+                opm_first=opm_first,
+                fuse_projection_weights=fuse_projection_weights,
+                inf=inf,
+                eps=eps,
+                ckpt=False,
+            )
+            self.blocks.append(block)
+            
+        self.tune_chunk_size = tune_chunk_size
+        self.chunk_size_tuner = None
+        if(tune_chunk_size):
+            self.chunk_size_tuner = ChunkSizeTuner()
+
+    def _prep_blocks(self, 
+        m: torch.Tensor, 
+        z: torch.Tensor, 
+        chunk_size: int,
+        use_deepspeed_evo_attention: bool,
+        use_lma: bool,
+        msa_mask: Optional[torch.Tensor],
+        pair_mask: Optional[torch.Tensor],
+        inplace_safe: bool,
+        _mask_trans: bool,
+    ):
+        blocks = [
+            partial(
+                b, 
+                msa_mask=msa_mask, 
+                pair_mask=pair_mask, 
+                chunk_size=chunk_size,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma,
+                inplace_safe=inplace_safe,
+                _mask_trans=_mask_trans,
+            ) for b in self.blocks
+        ]
+
+        def clear_cache(b, *args, **kwargs):
+            torch.cuda.empty_cache()
+            return b(*args, **kwargs)
+
+        if(self.clear_cache_between_blocks):
+            blocks = [partial(clear_cache, b) for b in blocks]
+
+        if(chunk_size is not None and self.chunk_size_tuner is not None):
+            tuned_chunk_size = self.chunk_size_tuner.tune_chunk_size(
+                representative_fn=blocks[0],
+                # Tensors cloned to avoid getting written to in-place
+                # A corollary is that chunk size tuning should be disabled for
+                # large N, when z gets really big
+                args=(m.clone(), z.clone(),),
+                min_chunk_size=chunk_size,
+            )
+            blocks = [
+                partial(b, 
+                    chunk_size=tuned_chunk_size,
+                    # A temporary measure to address torch's occasional
+                    # inability to allocate large tensors
+                    _attn_chunk_size=max(chunk_size, tuned_chunk_size // 4),
+                ) for b in blocks
+            ]
+
+        return blocks
+
+    def _forward_offload(self,
+        input_tensors: Sequence[torch.Tensor],
+        chunk_size: int,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        msa_mask: Optional[torch.Tensor] = None,
+        pair_mask: Optional[torch.Tensor] = None,
+        _mask_trans: bool = True,
+    ) -> torch.Tensor:
+        assert(not (self.training or torch.is_grad_enabled()))
+        blocks = self._prep_blocks(
+            # We are very careful not to create references to these tensors in
+            # this function
+            m=input_tensors[0],
+            z=input_tensors[1],
+            chunk_size=chunk_size,
+            use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+            use_lma=use_lma,
+            msa_mask=msa_mask,
+            pair_mask=pair_mask,
+            inplace_safe=True,
+            _mask_trans=_mask_trans,
+        )
+
+        for b in blocks:
+            m, z = b(
+                None, 
+                None, 
+                _offload_inference=True,
+                _offloadable_inputs=input_tensors,
+            )
+            input_tensors[0] = m
+            input_tensors[1] = z
+            del m, z
+
+        return input_tensors[1]
+
+    def forward(self,
+        m: torch.Tensor,
+        z: torch.Tensor,
+        msa_mask: Optional[torch.Tensor],
+        pair_mask: Optional[torch.Tensor],
+        chunk_size: int,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        inplace_safe: bool = False,
+        _mask_trans: bool = True,
+    ) -> torch.Tensor:
+        """
+        Args:
+            m:
+                [*, N_extra, N_res, C_m] extra MSA embedding
+            z:
+                [*, N_res, N_res, C_z] pair embedding
+            chunk_size: Inference-time subbatch size for Evoformer modules
+            use_deepspeed_evo_attention: Whether to use DeepSpeed memory-efficient kernel
+            use_lma: Whether to use low-memory attention during inference
+            msa_mask:
+                Optional [*, N_extra, N_res] MSA mask
+            pair_mask:
+                Optional [*, N_res, N_res] pair mask
+        Returns:
+            [*, N_res, N_res, C_z] pair update
+        """
+        checkpoint_fn = get_checkpoint_fn()
+        blocks = self._prep_blocks(
+            m=m,
+            z=z,
+            chunk_size=chunk_size,
+            use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+            use_lma=use_lma,
+            msa_mask=msa_mask,
+            pair_mask=pair_mask,
+            inplace_safe=inplace_safe,
+            _mask_trans=_mask_trans,
+        )
+
+        for b in blocks:
+            if(self.ckpt and torch.is_grad_enabled()):
+                m, z = checkpoint_fn(b, m, z)
+            else:
+                m, z = b(m, z)
+
+        return z

model/openfold/heads.py

@@ -0,0 +1,267 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+
+from model.openfold.primitives import Linear, LayerNorm
+from onescience.utils.openfold.loss import (
+    compute_plddt,
+    compute_tm,
+    compute_predicted_aligned_error,
+)
+from onescience.utils.openfold.precision_utils import is_fp16_enabled
+
+
+class AuxiliaryHeads(nn.Module):
+    def __init__(self, config):
+        super(AuxiliaryHeads, self).__init__()
+
+        self.plddt = PerResidueLDDTCaPredictor(
+            **config["lddt"],
+        )
+
+        self.distogram = DistogramHead(
+            **config["distogram"],
+        )
+
+        self.masked_msa = MaskedMSAHead(
+            **config["masked_msa"],
+        )
+
+        self.experimentally_resolved = ExperimentallyResolvedHead(
+            **config["experimentally_resolved"],
+        )
+
+        if config.tm.enabled:
+            self.tm = TMScoreHead(
+                **config.tm,
+            )
+
+        self.config = config
+
+    def forward(self, outputs):
+        aux_out = {}
+        lddt_logits = self.plddt(outputs["sm"]["single"])
+        aux_out["lddt_logits"] = lddt_logits
+
+        # Required for relaxation later on
+        aux_out["plddt"] = compute_plddt(lddt_logits)
+
+        distogram_logits = self.distogram(outputs["pair"])
+        aux_out["distogram_logits"] = distogram_logits
+
+        masked_msa_logits = self.masked_msa(outputs["msa"])
+        aux_out["masked_msa_logits"] = masked_msa_logits
+
+        experimentally_resolved_logits = self.experimentally_resolved(
+            outputs["single"]
+        )
+        aux_out[
+            "experimentally_resolved_logits"
+        ] = experimentally_resolved_logits
+
+        if self.config.tm.enabled:
+            tm_logits = self.tm(outputs["pair"])
+            aux_out["tm_logits"] = tm_logits
+            aux_out["ptm_score"] = compute_tm(
+                tm_logits, **self.config.tm
+            )
+            asym_id = outputs.get("asym_id")
+            if asym_id is not None:
+                aux_out["iptm_score"] = compute_tm(
+                    tm_logits, asym_id=asym_id, interface=True, **self.config.tm
+                )
+                aux_out["weighted_ptm_score"] = (self.config.tm["iptm_weight"] * aux_out["iptm_score"]
+                                                 + self.config.tm["ptm_weight"] * aux_out["ptm_score"])
+
+            aux_out.update(
+                compute_predicted_aligned_error(
+                    tm_logits,
+                    **self.config.tm,
+                )
+            )
+
+        return aux_out
+
+
+class PerResidueLDDTCaPredictor(nn.Module):
+    def __init__(self, no_bins, c_in, c_hidden):
+        super(PerResidueLDDTCaPredictor, self).__init__()
+
+        self.no_bins = no_bins
+        self.c_in = c_in
+        self.c_hidden = c_hidden
+
+        self.layer_norm = LayerNorm(self.c_in)
+
+        self.linear_1 = Linear(self.c_in, self.c_hidden, init="relu")
+        self.linear_2 = Linear(self.c_hidden, self.c_hidden, init="relu")
+        self.linear_3 = Linear(self.c_hidden, self.no_bins, init="final")
+
+        self.relu = nn.ReLU()
+
+    def forward(self, s):
+        s = self.layer_norm(s)
+        s = self.linear_1(s)
+        s = self.relu(s)
+        s = self.linear_2(s)
+        s = self.relu(s)
+        s = self.linear_3(s)
+
+        return s
+
+
+class DistogramHead(nn.Module):
+    """
+    Computes a distogram probability distribution.
+
+    For use in computation of distogram loss, subsection 1.9.8
+    """
+
+    def __init__(self, c_z, no_bins, **kwargs):
+        """
+        Args:
+            c_z:
+                Input channel dimension
+            no_bins:
+                Number of distogram bins
+        """
+        super(DistogramHead, self).__init__()
+
+        self.c_z = c_z
+        self.no_bins = no_bins
+
+        self.linear = Linear(self.c_z, self.no_bins, init="final")
+
+    def _forward(self, z):  # [*, N, N, C_z]
+        """
+        Args:
+            z:
+                [*, N_res, N_res, C_z] pair embedding
+        Returns:
+            [*, N, N, no_bins] distogram probability distribution
+        """
+        # [*, N, N, no_bins]
+        logits = self.linear(z)
+        logits = logits + logits.transpose(-2, -3)
+        return logits
+    
+    def forward(self, z): 
+        if(is_fp16_enabled()):
+            with torch.cuda.amp.autocast(enabled=False):
+                return self._forward(z.float())
+        else:
+            return self._forward(z)
+
+
+class TMScoreHead(nn.Module):
+    """
+    For use in computation of TM-score, subsection 1.9.7
+    """
+
+    def __init__(self, c_z, no_bins, **kwargs):
+        """
+        Args:
+            c_z:
+                Input channel dimension
+            no_bins:
+                Number of bins
+        """
+        super(TMScoreHead, self).__init__()
+
+        self.c_z = c_z
+        self.no_bins = no_bins
+
+        self.linear = Linear(self.c_z, self.no_bins, init="final")
+
+    def forward(self, z):
+        """
+        Args:
+            z:
+                [*, N_res, N_res, C_z] pairwise embedding
+        Returns:
+            [*, N_res, N_res, no_bins] prediction
+        """
+        # [*, N, N, no_bins]
+        logits = self.linear(z)
+        return logits
+
+
+class MaskedMSAHead(nn.Module):
+    """
+    For use in computation of masked MSA loss, subsection 1.9.9
+    """
+
+    def __init__(self, c_m, c_out, **kwargs):
+        """
+        Args:
+            c_m:
+                MSA channel dimension
+            c_out:
+                Output channel dimension
+        """
+        super(MaskedMSAHead, self).__init__()
+
+        self.c_m = c_m
+        self.c_out = c_out
+
+        self.linear = Linear(self.c_m, self.c_out, init="final")
+
+    def forward(self, m):
+        """
+        Args:
+            m:
+                [*, N_seq, N_res, C_m] MSA embedding
+        Returns:
+            [*, N_seq, N_res, C_out] reconstruction
+        """
+        # [*, N_seq, N_res, C_out]
+        logits = self.linear(m)
+        return logits
+
+
+class ExperimentallyResolvedHead(nn.Module):
+    """
+    For use in computation of "experimentally resolved" loss, subsection
+    1.9.10
+    """
+
+    def __init__(self, c_s, c_out, **kwargs):
+        """
+        Args:
+            c_s:
+                Input channel dimension
+            c_out:
+                Number of distogram bins
+        """
+        super(ExperimentallyResolvedHead, self).__init__()
+
+        self.c_s = c_s
+        self.c_out = c_out
+
+        self.linear = Linear(self.c_s, self.c_out, init="final")
+
+    def forward(self, s):
+        """
+        Args:
+            s:
+                [*, N_res, C_s] single embedding
+        Returns:
+            [*, N, C_out] logits
+        """
+        # [*, N, C_out]
+        logits = self.linear(s)
+        return logits

model/openfold/model.py

@@ -0,0 +1,591 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from functools import partial
+import weakref
+
+import torch
+import torch.nn as nn
+
+from onescience.datapipes.openfold import data_transforms_multimer
+from onescience.utils.openfold.feats import (
+    pseudo_beta_fn,
+    build_extra_msa_feat,
+    dgram_from_positions,
+    atom14_to_atom37,
+)
+from onescience.utils.openfold.tensor_utils import masked_mean
+from model.openfold.embedders import (
+    InputEmbedder,
+    InputEmbedderMultimer,
+    RecyclingEmbedder,
+    TemplateEmbedder,
+    TemplateEmbedderMultimer,
+    ExtraMSAEmbedder,
+    PreembeddingEmbedder,
+)
+from model.openfold.evoformer import EvoformerStack, ExtraMSAStack
+from model.openfold.heads import AuxiliaryHeads
+from model.openfold.structure_module import StructureModule
+from model.openfold.template import (
+    TemplatePairStack,
+    TemplatePointwiseAttention,
+    embed_templates_average,
+    embed_templates_offload,
+)
+import onescience.utils.openfold.np.residue_constants as residue_constants
+from onescience.utils.openfold.feats import (
+    pseudo_beta_fn,
+    build_extra_msa_feat,
+    build_template_angle_feat,
+    build_template_pair_feat,
+    atom14_to_atom37,
+)
+from onescience.utils.openfold.loss import (
+    compute_plddt,
+)
+from onescience.utils.openfold.tensor_utils import (
+    add,
+    dict_multimap,
+    tensor_tree_map,
+)
+
+
+class AlphaFold(nn.Module):
+    """
+    Alphafold 2.
+
+    Implements Algorithm 2 (but with training).
+    """
+
+    def __init__(self, config):
+        """
+        Args:
+            config:
+                A dict-like config object (like the one in config.py)
+        """
+        super(AlphaFold, self).__init__()
+
+        self.globals = config.globals
+        self.config = config.model
+        self.template_config = self.config.template
+        self.extra_msa_config = self.config.extra_msa
+        self.seqemb_mode = config.globals.seqemb_mode_enabled
+
+        # Main trunk + structure module
+        if self.globals.is_multimer:
+            self.input_embedder = InputEmbedderMultimer(
+                **self.config["input_embedder"]
+            )
+        elif self.seqemb_mode:
+            # If using seqemb mode, embed the sequence embeddings passed
+            # to the model ("preembeddings") instead of embedding the sequence
+            self.input_embedder = PreembeddingEmbedder(
+                **self.config["preembedding_embedder"],
+            )
+        else:
+            self.input_embedder = InputEmbedder(
+                **self.config["input_embedder"],
+            )
+
+        self.recycling_embedder = RecyclingEmbedder(
+            **self.config["recycling_embedder"],
+        )
+
+        if self.template_config.enabled:
+            if self.globals.is_multimer:
+                self.template_embedder = TemplateEmbedderMultimer(
+                    self.template_config,
+                )
+            else:
+                self.template_embedder = TemplateEmbedder(
+                    self.template_config,
+                )
+
+        if self.extra_msa_config.enabled:
+            self.extra_msa_embedder = ExtraMSAEmbedder(
+                **self.extra_msa_config["extra_msa_embedder"],
+            )
+            self.extra_msa_stack = ExtraMSAStack(
+                **self.extra_msa_config["extra_msa_stack"],
+            )
+
+        self.evoformer = EvoformerStack(
+            **self.config["evoformer_stack"],
+        )
+
+        self.structure_module = StructureModule(
+            is_multimer=self.globals.is_multimer,
+            **self.config["structure_module"],
+        )
+        self.aux_heads = AuxiliaryHeads(
+            self.config["heads"],
+        )
+
+    def embed_templates(self, batch, feats, z, pair_mask, templ_dim, inplace_safe):
+        if self.globals.is_multimer:
+            asym_id = feats["asym_id"]
+            multichain_mask_2d = (
+                asym_id[..., None] == asym_id[..., None, :]
+            )
+            template_embeds = self.template_embedder(
+                batch,
+                z,
+                pair_mask.to(dtype=z.dtype),
+                templ_dim,
+                chunk_size=self.globals.chunk_size,
+                multichain_mask_2d=multichain_mask_2d,
+                use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
+                use_lma=self.globals.use_lma,
+                inplace_safe=inplace_safe,
+                _mask_trans=self.config._mask_trans
+            )
+            feats["template_torsion_angles_mask"] = (
+                template_embeds["template_mask"]
+            )
+        else:
+            if self.template_config.offload_templates:
+                return embed_templates_offload(self,
+                                               batch, z, pair_mask, templ_dim, inplace_safe=inplace_safe,
+                                               )
+            elif self.template_config.average_templates:
+                return embed_templates_average(self,
+                                               batch, z, pair_mask, templ_dim, inplace_safe=inplace_safe,
+                                               )
+
+            template_embeds = self.template_embedder(
+                batch,
+                z,
+                pair_mask.to(dtype=z.dtype),
+                templ_dim,
+                chunk_size=self.globals.chunk_size,
+                use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
+                use_lma=self.globals.use_lma,
+                inplace_safe=inplace_safe,
+                _mask_trans=self.config._mask_trans
+            )
+
+        return template_embeds
+
+    def tolerance_reached(self, prev_pos, next_pos, mask, eps=1e-8) -> bool:
+        """
+        Early stopping criteria based on criteria used in
+        AF2Complex: https://www.nature.com/articles/s41467-022-29394-2
+        Args:
+          prev_pos: Previous atom positions in atom37/14 representation
+          next_pos: Current atom positions in atom37/14 representation
+          mask: 1-D sequence mask
+          eps: Epsilon used in square root calculation
+        Returns:
+          Whether to stop recycling early based on the desired tolerance.
+        """
+
+        def distances(points):
+            """Compute all pairwise distances for a set of points."""
+            d = points[..., None, :] - points[..., None, :, :]
+            return torch.sqrt(torch.sum(d ** 2, dim=-1))
+
+        if self.config.recycle_early_stop_tolerance < 0:
+            return False
+
+        ca_idx = residue_constants.atom_order['CA']
+        sq_diff = (distances(prev_pos[..., ca_idx, :]) - distances(next_pos[..., ca_idx, :])) ** 2
+        mask = mask[..., None] * mask[..., None, :]
+        sq_diff = masked_mean(mask=mask, value=sq_diff, dim=list(range(len(mask.shape))))
+        diff = torch.sqrt(sq_diff + eps).item()
+        return diff <= self.config.recycle_early_stop_tolerance
+
+    def iteration(self, feats, prevs, _recycle=True):
+        # Primary output dictionary
+        outputs = {}
+
+        # This needs to be done manually for DeepSpeed's sake
+        dtype = next(self.parameters()).dtype
+        for k in feats:
+            if feats[k].dtype == torch.float32:
+                feats[k] = feats[k].to(dtype=dtype)
+
+        # Grab some data about the input
+        batch_dims = feats["target_feat"].shape[:-2]
+        no_batch_dims = len(batch_dims)
+        n = feats["target_feat"].shape[-2]
+        n_seq = feats["msa_feat"].shape[-3]
+        device = feats["target_feat"].device
+
+        # Controls whether the model uses in-place operations throughout
+        # The dual condition accounts for activation checkpoints
+        inplace_safe = not (self.training or torch.is_grad_enabled())
+
+        # Prep some features
+        seq_mask = feats["seq_mask"]
+        pair_mask = seq_mask[..., None] * seq_mask[..., None, :]
+        msa_mask = feats["msa_mask"]
+
+        if self.globals.is_multimer:
+            # Initialize the MSA and pair representations
+            # m: [*, S_c, N, C_m]
+            # z: [*, N, N, C_z]
+            m, z = self.input_embedder(feats)
+        elif self.seqemb_mode:
+            # Initialize the SingleSeq and pair representations
+            # m: [*, 1, N, C_m]
+            # z: [*, N, N, C_z]
+            m, z = self.input_embedder(
+                feats["target_feat"],
+                feats["residue_index"],
+                feats["seq_embedding"]
+            )
+        else:
+            # Initialize the MSA and pair representations
+            # m: [*, S_c, N, C_m]
+            # z: [*, N, N, C_z]
+            m, z = self.input_embedder(
+                feats["target_feat"],
+                feats["residue_index"],
+                feats["msa_feat"],
+                inplace_safe=inplace_safe,
+            )
+
+        # Unpack the recycling embeddings. Removing them from the list allows 
+        # them to be freed further down in this function, saving memory
+        m_1_prev, z_prev, x_prev = reversed([prevs.pop() for _ in range(3)])
+
+        # Initialize the recycling embeddings, if needs be 
+        if None in [m_1_prev, z_prev, x_prev]:
+            # [*, N, C_m]
+            m_1_prev = m.new_zeros(
+                (*batch_dims, n, self.config.input_embedder.c_m),
+                requires_grad=False,
+            )
+
+            # [*, N, N, C_z]
+            z_prev = z.new_zeros(
+                (*batch_dims, n, n, self.config.input_embedder.c_z),
+                requires_grad=False,
+            )
+
+            # [*, N, 3]
+            x_prev = z.new_zeros(
+                (*batch_dims, n, residue_constants.atom_type_num, 3),
+                requires_grad=False,
+            )
+
+        pseudo_beta_x_prev = pseudo_beta_fn(
+            feats["aatype"], x_prev, None
+        ).to(dtype=z.dtype)
+
+        # The recycling embedder is memory-intensive, so we offload first
+        if self.globals.offload_inference and inplace_safe:
+            m = m.cpu()
+            z = z.cpu()
+
+        # m_1_prev_emb: [*, N, C_m]
+        # z_prev_emb: [*, N, N, C_z]
+        m_1_prev_emb, z_prev_emb = self.recycling_embedder(
+            m_1_prev,
+            z_prev,
+            pseudo_beta_x_prev,
+            inplace_safe=inplace_safe,
+        )
+
+        del pseudo_beta_x_prev
+
+        if self.globals.offload_inference and inplace_safe:
+            m = m.to(m_1_prev_emb.device)
+            z = z.to(z_prev.device)
+
+        # [*, S_c, N, C_m]
+        m[..., 0, :, :] += m_1_prev_emb
+
+        # [*, N, N, C_z]
+        z = add(z, z_prev_emb, inplace=inplace_safe)
+
+        # Deletions like these become significant for inference with large N,
+        # where they free unused tensors and remove references to others such
+        # that they can be offloaded later
+        del m_1_prev, z_prev, m_1_prev_emb, z_prev_emb
+
+        # Embed the templates + merge with MSA/pair embeddings
+        if self.config.template.enabled:
+            template_feats = {
+                k: v for k, v in feats.items() if k.startswith("template_")
+            }
+
+            template_embeds = self.embed_templates(
+                template_feats,
+                feats,
+                z,
+                pair_mask.to(dtype=z.dtype),
+                no_batch_dims,
+                inplace_safe=inplace_safe,
+            )
+
+            # [*, N, N, C_z]
+            z = add(z,
+                    template_embeds.pop("template_pair_embedding"),
+                    inplace_safe,
+                    )
+
+            if (
+                "template_single_embedding" in template_embeds
+            ):
+                # [*, S = S_c + S_t, N, C_m]
+                m = torch.cat(
+                    [m, template_embeds["template_single_embedding"]],
+                    dim=-3
+                )
+
+                # [*, S, N]
+                if not self.globals.is_multimer:
+                    torsion_angles_mask = feats["template_torsion_angles_mask"]
+                    msa_mask = torch.cat(
+                        [feats["msa_mask"], torsion_angles_mask[..., 2]],
+                        dim=-2
+                    )
+                else:
+                    msa_mask = torch.cat(
+                        [feats["msa_mask"], template_embeds["template_mask"]],
+                        dim=-2,
+                    )
+
+        # Embed extra MSA features + merge with pairwise embeddings
+        if self.config.extra_msa.enabled:
+            if self.globals.is_multimer:
+                extra_msa_fn = data_transforms_multimer.build_extra_msa_feat
+            else:
+                extra_msa_fn = build_extra_msa_feat
+
+            # [*, S_e, N, C_e]
+            extra_msa_feat = extra_msa_fn(feats).to(dtype=z.dtype)
+            a = self.extra_msa_embedder(extra_msa_feat)
+
+            if self.globals.offload_inference:
+                # To allow the extra MSA stack (and later the evoformer) to
+                # offload its inputs, we remove all references to them here
+                input_tensors = [a, z]
+                del a, z
+
+                # [*, N, N, C_z]
+                z = self.extra_msa_stack._forward_offload(
+                    input_tensors,
+                    msa_mask=feats["extra_msa_mask"].to(dtype=m.dtype),
+                    chunk_size=self.globals.chunk_size,
+                    use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
+                    use_lma=self.globals.use_lma,
+                    pair_mask=pair_mask.to(dtype=m.dtype),
+                    _mask_trans=self.config._mask_trans,
+                )
+
+                del input_tensors
+            else:
+                # [*, N, N, C_z]
+                z = self.extra_msa_stack(
+                    a, z,
+                    msa_mask=feats["extra_msa_mask"].to(dtype=m.dtype),
+                    chunk_size=self.globals.chunk_size,
+                    use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
+                    use_lma=self.globals.use_lma,
+                    pair_mask=pair_mask.to(dtype=m.dtype),
+                    inplace_safe=inplace_safe,
+                    _mask_trans=self.config._mask_trans,
+                )
+
+        # Run MSA + pair embeddings through the trunk of the network
+        # m: [*, S, N, C_m]
+        # z: [*, N, N, C_z]
+        # s: [*, N, C_s]          
+        if self.globals.offload_inference:
+            input_tensors = [m, z]
+            del m, z
+            m, z, s = self.evoformer._forward_offload(
+                input_tensors,
+                msa_mask=msa_mask.to(dtype=input_tensors[0].dtype),
+                pair_mask=pair_mask.to(dtype=input_tensors[1].dtype),
+                chunk_size=self.globals.chunk_size,
+                use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
+                use_lma=self.globals.use_lma,
+                _mask_trans=self.config._mask_trans,
+            )
+
+            del input_tensors
+        else:
+            m, z, s = self.evoformer(
+                m,
+                z,
+                msa_mask=msa_mask.to(dtype=m.dtype),
+                pair_mask=pair_mask.to(dtype=z.dtype),
+                chunk_size=self.globals.chunk_size,
+                use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
+                use_lma=self.globals.use_lma,
+                use_flash=self.globals.use_flash,
+                inplace_safe=inplace_safe,
+                _mask_trans=self.config._mask_trans,
+            )
+
+        outputs["msa"] = m[..., :n_seq, :, :]
+        outputs["pair"] = z
+        outputs["single"] = s
+
+        del z
+
+        # Predict 3D structure
+        outputs["sm"] = self.structure_module(
+            outputs,
+            feats["aatype"],
+            mask=feats["seq_mask"].to(dtype=s.dtype),
+            inplace_safe=inplace_safe,
+            _offload_inference=self.globals.offload_inference,
+        )
+        outputs["final_atom_positions"] = atom14_to_atom37(
+            outputs["sm"]["positions"][-1], feats
+        )
+        outputs["final_atom_mask"] = feats["atom37_atom_exists"]
+        outputs["final_affine_tensor"] = outputs["sm"]["frames"][-1]
+
+        # Save embeddings for use during the next recycling iteration
+
+        # [*, N, C_m]
+        m_1_prev = m[..., 0, :, :]
+
+        # [*, N, N, C_z]
+        z_prev = outputs["pair"]
+
+        early_stop = False
+        if self.globals.is_multimer:
+            early_stop = self.tolerance_reached(x_prev, outputs["final_atom_positions"], seq_mask)
+
+        del x_prev
+
+        # [*, N, 3]
+        x_prev = outputs["final_atom_positions"]
+
+        return outputs, m_1_prev, z_prev, x_prev, early_stop
+
+    def _disable_activation_checkpointing(self):
+        self.template_embedder.template_pair_stack.blocks_per_ckpt = None
+        self.evoformer.blocks_per_ckpt = None
+
+        for b in self.extra_msa_stack.blocks:
+            b.ckpt = False
+
+    def _enable_activation_checkpointing(self):
+        self.template_embedder.template_pair_stack.blocks_per_ckpt = (
+            self.config.template.template_pair_stack.blocks_per_ckpt
+        )
+        self.evoformer.blocks_per_ckpt = (
+            self.config.evoformer_stack.blocks_per_ckpt
+        )
+
+        for b in self.extra_msa_stack.blocks:
+            b.ckpt = self.config.extra_msa.extra_msa_stack.ckpt
+
+    def forward(self, batch):
+        """
+        Args:
+            batch:
+                Dictionary of arguments outlined in Algorithm 2. Keys must
+                include the official names of the features in the
+                supplement subsection 1.2.9.
+
+                The final dimension of each input must have length equal to
+                the number of recycling iterations.
+
+                Features (without the recycling dimension):
+
+                    "aatype" ([*, N_res]):
+                        Contrary to the supplement, this tensor of residue
+                        indices is not one-hot.
+                    "target_feat" ([*, N_res, C_tf])
+                        One-hot encoding of the target sequence. C_tf is
+                        config.model.input_embedder.tf_dim.
+                    "residue_index" ([*, N_res])
+                        Tensor whose final dimension consists of
+                        consecutive indices from 0 to N_res.
+                    "msa_feat" ([*, N_seq, N_res, C_msa])
+                        MSA features, constructed as in the supplement.
+                        C_msa is config.model.input_embedder.msa_dim.
+                    "seq_mask" ([*, N_res])
+                        1-D sequence mask
+                    "msa_mask" ([*, N_seq, N_res])
+                        MSA mask
+                    "pair_mask" ([*, N_res, N_res])
+                        2-D pair mask
+                    "extra_msa_mask" ([*, N_extra, N_res])
+                        Extra MSA mask
+                    "template_mask" ([*, N_templ])
+                        Template mask (on the level of templates, not
+                        residues)
+                    "template_aatype" ([*, N_templ, N_res])
+                        Tensor of template residue indices (indices greater
+                        than 19 are clamped to 20 (Unknown))
+                    "template_all_atom_positions"
+                        ([*, N_templ, N_res, 37, 3])
+                        Template atom coordinates in atom37 format
+                    "template_all_atom_mask" ([*, N_templ, N_res, 37])
+                        Template atom coordinate mask
+                    "template_pseudo_beta" ([*, N_templ, N_res, 3])
+                        Positions of template carbon "pseudo-beta" atoms
+                        (i.e. C_beta for all residues but glycine, for
+                        for which C_alpha is used instead)
+                    "template_pseudo_beta_mask" ([*, N_templ, N_res])
+                        Pseudo-beta mask
+        """
+        # Initialize recycling embeddings
+        m_1_prev, z_prev, x_prev = None, None, None
+        prevs = [m_1_prev, z_prev, x_prev]
+
+        is_grad_enabled = torch.is_grad_enabled()
+
+        # Main recycling loop
+        num_iters = batch["aatype"].shape[-1]
+        early_stop = False
+        num_recycles = 0
+        for cycle_no in range(num_iters):
+            # Select the features for the current recycling cycle
+            fetch_cur_batch = lambda t: t[..., cycle_no]
+            feats = tensor_tree_map(fetch_cur_batch, batch)
+
+            # Enable grad iff we're training and it's the final recycling layer
+            is_final_iter = cycle_no == (num_iters - 1) or early_stop
+            with torch.set_grad_enabled(is_grad_enabled and is_final_iter):
+                if is_final_iter:
+                    # Sidestep AMP bug (PyTorch issue #65766)
+                    if torch.is_autocast_enabled():
+                        torch.clear_autocast_cache()
+
+                # Run the next iteration of the model
+                outputs, m_1_prev, z_prev, x_prev, early_stop = self.iteration(
+                    feats,
+                    prevs,
+                    _recycle=(num_iters > 1)
+                )
+
+                num_recycles += 1
+
+                if not is_final_iter:
+                    del outputs
+                    prevs = [m_1_prev, z_prev, x_prev]
+                    del m_1_prev, z_prev, x_prev
+                else:
+                    break
+
+        outputs["num_recycles"] = torch.tensor(num_recycles, device=feats["aatype"].device)
+
+        if "asym_id" in batch:
+            outputs["asym_id"] = feats["asym_id"]
+
+        # Run auxiliary heads
+        outputs.update(self.aux_heads(outputs))
+
+        return outputs

model/openfold/msa.py

@@ -0,0 +1,476 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from functools import partial
+import math
+import torch
+import torch.nn as nn
+from typing import Optional, List, Tuple
+
+from model.openfold.primitives import (
+    Linear, 
+    LayerNorm,
+    Attention, 
+    GlobalAttention, 
+    _attention_chunked_trainable,
+)
+from onescience.utils.openfold.checkpointing import get_checkpoint_fn
+from onescience.utils.openfold.chunk_utils import chunk_layer
+from onescience.utils.openfold.tensor_utils import (
+    permute_final_dims,
+    flatten_final_dims,
+)
+
+
+class MSAAttention(nn.Module):
+    def __init__(
+        self,
+        c_in,
+        c_hidden,
+        no_heads,
+        pair_bias=False,
+        c_z=None,
+        inf=1e9,
+    ):
+        """
+        Args:
+            c_in:
+                Input channel dimension
+            c_hidden:
+                Per-head hidden channel dimension
+            no_heads:
+                Number of attention heads
+            pair_bias:
+                Whether to use pair embedding bias
+            c_z:
+                Pair embedding channel dimension. Ignored unless pair_bias
+                is true
+            inf:
+                A large number to be used in computing the attention mask
+        """
+        super(MSAAttention, self).__init__()
+
+        self.c_in = c_in
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.pair_bias = pair_bias
+        self.c_z = c_z
+        self.inf = inf
+
+        self.layer_norm_m = LayerNorm(self.c_in)
+
+        self.layer_norm_z = None
+        self.linear_z = None
+        if self.pair_bias:
+            self.layer_norm_z = LayerNorm(self.c_z)
+            self.linear_z = Linear(
+                self.c_z, self.no_heads, bias=False, init="normal"
+            )
+        
+        self.mha = Attention(
+            self.c_in, 
+            self.c_in, 
+            self.c_in, 
+            self.c_hidden, 
+            self.no_heads,
+        )
+
+    @torch.jit.ignore
+    def _chunk(self, 
+        m: torch.Tensor,
+        biases: Optional[List[torch.Tensor]],
+        chunk_size: int,
+        use_memory_efficient_kernel: bool,
+        use_deepspeed_evo_attention: bool,
+        use_lma: bool,
+        use_flash: bool,
+        flash_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        def fn(m, biases, flash_mask):
+            m = self.layer_norm_m(m)
+            return self.mha(
+                q_x=m, 
+                kv_x=m, 
+                biases=biases,
+                use_memory_efficient_kernel=use_memory_efficient_kernel,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma,
+                use_flash=use_flash,
+                flash_mask=flash_mask,
+            )
+
+        inputs = {"m": m}
+        if(biases is not None):
+            inputs["biases"] = biases
+        else:
+            fn = partial(fn, biases=None)
+        if(use_flash and flash_mask is not None):
+            inputs["flash_mask"] = flash_mask
+        else:
+            fn = partial(fn, flash_mask=None)
+
+        return chunk_layer(
+            fn,
+            inputs,
+            chunk_size=chunk_size,
+            no_batch_dims=len(m.shape[:-2])
+        )
+
+    def _prep_inputs(self,
+        m: torch.Tensor,
+        z: Optional[torch.Tensor],
+        mask: Optional[torch.Tensor],
+        inplace_safe: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 
+        n_seq, n_res = m.shape[-3:-1]
+        if mask is None:
+            # [*, N_seq, N_res]
+            mask = m.new_ones(
+                m.shape[:-3] + (n_seq, n_res),
+            )
+
+        # [*, N_seq, 1, 1, N_res]
+        mask_bias = (self.inf * (mask - 1))[..., :, None, None, :]
+
+        if (self.pair_bias and 
+            z is not None and                       # For the 
+            self.layer_norm_z is not None and       # benefit of
+            self.linear_z is not None               # TorchScript
+        ):
+            chunks = []
+
+            for i in range(0, z.shape[-3], 256):
+                z_chunk = z[..., i: i + 256, :, :]
+
+                # [*, N_res, N_res, C_z]
+                z_chunk = self.layer_norm_z(z_chunk)
+            
+                # [*, N_res, N_res, no_heads]
+                z_chunk = self.linear_z(z_chunk)
+
+                chunks.append(z_chunk)
+            
+            z = torch.cat(chunks, dim=-3)
+            
+            # [*, 1, no_heads, N_res, N_res]
+            z = permute_final_dims(z, (2, 0, 1)).unsqueeze(-4)
+
+        return m, mask_bias, z
+
+    @torch.jit.ignore
+    def _chunked_msa_attn(self,
+        m: torch.Tensor,
+        z: Optional[torch.Tensor],
+        mask: Optional[torch.Tensor],
+        chunk_logits: int,
+        checkpoint: bool,
+        inplace_safe: bool = False
+    ) -> torch.Tensor:
+        """ 
+        MSA attention with training-time chunking of the softmax computation.
+        Saves memory in the extra MSA stack. Probably obviated by our fused 
+        attention kernel, which is now used by default.
+        """
+        MSA_DIM = -4
+
+        def _get_qkv(m, z):
+            m, mask_bias, z = self._prep_inputs(
+                m, z, mask, inplace_safe=inplace_safe
+            )
+            m = self.layer_norm_m(m)
+            q, k, v = self.mha._prep_qkv(m, m)
+            return m, q, k, v, mask_bias, z
+
+        checkpoint_fn = get_checkpoint_fn()
+
+        if(torch.is_grad_enabled() and checkpoint):
+            m, q, k, v, mask_bias, z = checkpoint_fn(_get_qkv, m, z)
+        else:
+            m, q, k, v, mask_bias, z = _get_qkv(m, z)
+       
+        o = _attention_chunked_trainable(
+            query=q, 
+            key=k, 
+            value=v, 
+            biases=[mask_bias, z], 
+            chunk_size=chunk_logits, 
+            chunk_dim=MSA_DIM,
+            checkpoint=checkpoint,
+        )
+
+        if(torch.is_grad_enabled() and checkpoint):
+            # Storing an additional m here is far from ideal
+            m = checkpoint_fn(self.mha._wrap_up, o, m)
+        else:
+            m = self.mha._wrap_up(o, m)
+
+        return m
+
+    def forward(self, 
+        m: torch.Tensor, 
+        z: Optional[torch.Tensor] = None, 
+        mask: Optional[torch.Tensor] = None, 
+        chunk_size: Optional[int] = None,
+        use_memory_efficient_kernel: bool = False,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        use_flash: bool = False,
+        inplace_safe: bool = False,
+        _chunk_logits: Optional[int] = None,
+        _checkpoint_chunks: Optional[bool] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            m:
+                [*, N_seq, N_res, C_m] MSA embedding
+            z:
+                [*, N_res, N_res, C_z] pair embedding. Required only if
+                pair_bias is True
+            mask:
+                [*, N_seq, N_res] MSA mask
+            chunk_size:
+                Size of chunks into which the inputs are split along their
+                batch dimensions. A low value decreases memory overhead at the 
+                cost of slower execution. Chunking is not performed by default.
+                
+        """
+        if(_chunk_logits is not None):
+            return self._chunked_msa_attn(
+                m=m, z=z, mask=mask, 
+                chunk_logits=_chunk_logits, 
+                checkpoint=_checkpoint_chunks,
+                inplace_safe=inplace_safe,
+            )
+       
+        if(use_flash):
+            assert z is None
+            biases = None
+        else:    
+            m, mask_bias, z = self._prep_inputs(
+                m, z, mask, inplace_safe=inplace_safe
+            )
+    
+            biases = [mask_bias]
+            if(z is not None):
+                biases.append(z)
+
+        if chunk_size is not None:
+            m = self._chunk(
+                m, 
+                biases, 
+                chunk_size,
+                use_memory_efficient_kernel=use_memory_efficient_kernel,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma,
+                use_flash=use_flash,
+                flash_mask=mask,
+            )
+        else:
+            m = self.layer_norm_m(m)
+            m = self.mha(
+                q_x=m, 
+                kv_x=m, 
+                biases=biases,
+                use_memory_efficient_kernel=use_memory_efficient_kernel,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma,
+                use_flash=use_flash,
+                flash_mask=mask,
+            )
+
+        return m
+
+
+class MSARowAttentionWithPairBias(MSAAttention):
+    """
+    Implements Algorithm 7.
+    """
+
+    def __init__(self, c_m, c_z, c_hidden, no_heads, inf=1e9):
+        """
+        Args:
+            c_m:
+                Input channel dimension
+            c_z:
+                Pair embedding channel dimension
+            c_hidden:
+                Per-head hidden channel dimension
+            no_heads:
+                Number of attention heads
+            inf:
+                Large number used to construct attention masks
+        """
+        super(MSARowAttentionWithPairBias, self).__init__(
+            c_m,
+            c_hidden,
+            no_heads,
+            pair_bias=True,
+            c_z=c_z,
+            inf=inf,
+        )
+
+
+class MSAColumnAttention(nn.Module):
+    """
+    Implements Algorithm 8.
+
+    By rights, this should also be a subclass of MSAAttention. Alas,
+    most inheritance isn't supported by TorchScript.
+    """
+
+    def __init__(self, c_m, c_hidden, no_heads, inf=1e9):
+        """
+        Args:
+            c_m:
+                MSA channel dimension
+            c_hidden:
+                Per-head hidden channel dimension
+            no_heads:
+                Number of attention heads
+            inf:
+                Large number used to construct attention masks
+        """
+        super(MSAColumnAttention, self).__init__()
+        
+        self.c_m = c_m
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.inf = inf
+
+        self._msa_att = MSAAttention(
+            c_in=c_m,
+            c_hidden=c_hidden,
+            no_heads=no_heads,
+            pair_bias=False,
+            c_z=None,
+            inf=inf,
+        )
+
+    def forward(self, 
+        m: torch.Tensor, 
+        mask: Optional[torch.Tensor] = None, 
+        chunk_size: Optional[int] = None,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        use_flash: bool = False,
+    ) -> torch.Tensor:
+        """
+        Args:
+            m:
+                [*, N_seq, N_res, C_m] MSA embedding
+            mask:
+                [*, N_seq, N_res] MSA mask
+            chunk_size:
+                Size of chunks into which the inputs are split along their
+                batch dimensions. A low value decreases memory overhead at the 
+                cost of slower execution. Chunking is not performed by default.
+        """ 
+        # [*, N_res, N_seq, C_in]
+        m = m.transpose(-2, -3)
+        if mask is not None:
+            mask = mask.transpose(-1, -2)
+
+        m = self._msa_att(
+            m, 
+            mask=mask, 
+            chunk_size=chunk_size,
+            use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+            use_lma=use_lma,
+            use_flash=use_flash,
+        )
+
+        # [*, N_seq, N_res, C_in]
+        m = m.transpose(-2, -3)
+        if mask is not None:
+            mask = mask.transpose(-1, -2)
+
+        return m
+
+
+class MSAColumnGlobalAttention(nn.Module):
+    def __init__(
+        self, c_in, c_hidden, no_heads, inf=1e9, eps=1e-10,
+    ):
+        super(MSAColumnGlobalAttention, self).__init__()
+
+        self.c_in = c_in
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.inf = inf
+        self.eps = eps
+
+        self.layer_norm_m = nn.LayerNorm(c_in)
+
+        self.global_attention = GlobalAttention(
+            c_in=c_in,
+            c_hidden=c_hidden,
+            no_heads=no_heads,
+            inf=inf,
+            eps=eps,
+        )
+
+    @torch.jit.ignore
+    def _chunk(self,
+        m: torch.Tensor,
+        mask: torch.Tensor,
+        chunk_size: int,
+        use_lma: bool = False,
+    ) -> torch.Tensor:
+        mha_input = {
+            "m": m,
+            "mask": mask,
+        }
+
+        def fn(m, mask):
+            m = self.layer_norm_m(m)
+            return self.global_attention(m, mask, use_lma=use_lma)
+
+        return chunk_layer(
+            fn,
+            mha_input,
+            chunk_size=chunk_size,
+            no_batch_dims=len(m.shape[:-2]),
+        )
+
+    def forward(
+        self, 
+        m: torch.Tensor, 
+        mask: Optional[torch.Tensor] = None, 
+        chunk_size: Optional[int] = None,
+        use_lma: bool = False,
+    ) -> torch.Tensor:
+        n_seq, n_res, c_in = m.shape[-3:]
+
+        if mask is None:
+            # [*, N_seq, N_res]
+            mask = torch.ones(
+                m.shape[:-1],
+                dtype=m.dtype,
+                device=m.device,
+            ).detach()
+
+        # [*, N_res, N_seq, C_in]
+        m = m.transpose(-2, -3)
+        mask = mask.transpose(-1, -2)
+
+        if chunk_size is not None:
+            m = self._chunk(m, mask, chunk_size, use_lma=use_lma) 
+        else:
+            m = self.layer_norm_m(m)
+            m = self.global_attention(m=m, mask=mask, use_lma=use_lma)
+
+        # [*, N_seq, N_res, C_in]
+        m = m.transpose(-2, -3)
+
+        return m

model/openfold/outer_product_mean.py

@@ -0,0 +1,159 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import partial
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from model.openfold.primitives import Linear
+from onescience.utils.openfold.chunk_utils import chunk_layer
+from onescience.utils.openfold.precision_utils import is_fp16_enabled
+
+
+class OuterProductMean(nn.Module):
+    """
+    Implements Algorithm 10.
+    """
+
+    def __init__(self, c_m, c_z, c_hidden, eps=1e-3, bias: bool=True):
+        """
+        Args:
+            c_m:
+                MSA embedding channel dimension
+            c_z:
+                Pair embedding channel dimension
+            c_hidden:
+                Hidden channel dimension
+        """
+        super(OuterProductMean, self).__init__()
+
+        self.c_m = c_m
+        self.c_z = c_z
+        self.c_hidden = c_hidden
+        self.eps = eps
+
+        self.layer_norm = nn.LayerNorm(c_m)
+        self.linear_1 = Linear(c_m, c_hidden, bias=bias)
+        self.linear_2 = Linear(c_m, c_hidden, bias=bias)
+        self.linear_out = Linear(c_hidden ** 2, c_z, init="final")
+
+    def _opm(self, a, b):
+        # [*, N_res, N_res, C, C]
+        outer = torch.einsum("...bac,...dae->...bdce", a, b)
+
+        # [*, N_res, N_res, C * C]
+        outer = outer.reshape(outer.shape[:-2] + (-1,))
+
+        # [*, N_res, N_res, C_z]
+        outer = self.linear_out(outer)
+
+        return outer
+
+    @torch.jit.ignore
+    def _chunk(self, 
+        a: torch.Tensor, 
+        b: torch.Tensor, 
+        chunk_size: int
+    ) -> torch.Tensor:
+        # Since the "batch dim" in this case is not a true batch dimension
+        # (in that the shape of the output depends on it), we need to
+        # iterate over it ourselves
+        a_reshape = a.reshape((-1,) + a.shape[-3:])
+        b_reshape = b.reshape((-1,) + b.shape[-3:])
+        out = []
+        for a_prime, b_prime in zip(a_reshape, b_reshape):
+            outer = chunk_layer(
+                partial(self._opm, b=b_prime),
+                {"a": a_prime},
+                chunk_size=chunk_size,
+                no_batch_dims=1,
+            )
+            out.append(outer)
+
+        # For some cursed reason making this distinction saves memory
+        if(len(out) == 1):
+            outer = out[0].unsqueeze(0)
+        else:
+            outer = torch.stack(out, dim=0)
+
+        outer = outer.reshape(a.shape[:-3] + outer.shape[1:])
+
+        return outer
+
+    def _forward(self, 
+        m: torch.Tensor, 
+        mask: Optional[torch.Tensor] = None,
+        chunk_size: Optional[int] = None,
+        inplace_safe: bool = False,
+    ) -> torch.Tensor:
+        """
+        Args:
+            m:
+                [*, N_seq, N_res, C_m] MSA embedding
+            mask:
+                [*, N_seq, N_res] MSA mask
+        Returns:
+            [*, N_res, N_res, C_z] pair embedding update
+        """
+        if mask is None:
+            mask = m.new_ones(m.shape[:-1])
+
+        # [*, N_seq, N_res, C_m]
+        ln = self.layer_norm(m)
+
+        # [*, N_seq, N_res, C]
+        mask = mask.unsqueeze(-1)
+        a = self.linear_1(ln) 
+        a = a * mask
+        
+        b = self.linear_2(ln) 
+        b = b * mask
+
+        del ln
+
+        a = a.transpose(-2, -3)
+        b = b.transpose(-2, -3)
+
+        if chunk_size is not None:
+            outer = self._chunk(a, b, chunk_size)
+        else:
+            outer = self._opm(a, b)
+
+        # [*, N_res, N_res, 1]
+        norm = torch.einsum("...abc,...adc->...bdc", mask, mask)
+        norm = norm + self.eps
+
+        # [*, N_res, N_res, C_z]
+        if(inplace_safe):
+            outer /= norm
+        else:
+            outer = outer / norm
+
+        return outer
+
+    def forward(self,
+                m: torch.Tensor,
+                mask: Optional[torch.Tensor] = None,
+                chunk_size: Optional[int] = None,
+                inplace_safe: bool = False,
+    ) -> torch.Tensor:
+        if(is_fp16_enabled()):
+            with torch.cuda.amp.autocast(enabled=False):
+                return self._forward(m.float(), mask, chunk_size, inplace_safe)
+        else:
+            return self._forward(m, mask, chunk_size, inplace_safe)
+        

model/openfold/pair_transition.py

@@ -0,0 +1,99 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from model.openfold.primitives import Linear, LayerNorm
+from onescience.utils.openfold.chunk_utils import chunk_layer
+
+
+class PairTransition(nn.Module):
+    """
+    Implements Algorithm 15.
+    """
+
+    def __init__(self, c_z, n):
+        """
+        Args:
+            c_z:
+                Pair transition channel dimension
+            n:
+                Factor by which c_z is multiplied to obtain hidden channel
+                dimension
+        """
+        super(PairTransition, self).__init__()
+
+        self.c_z = c_z
+        self.n = n
+
+        self.layer_norm = LayerNorm(self.c_z)
+        self.linear_1 = Linear(self.c_z, self.n * self.c_z, init="relu")
+        self.relu = nn.ReLU()
+        self.linear_2 = Linear(self.n * self.c_z, c_z, init="final")
+
+    def _transition(self, z, mask):
+        # [*, N_res, N_res, C_z]
+        z = self.layer_norm(z)
+        
+        # [*, N_res, N_res, C_hidden]
+        z = self.linear_1(z)
+        z = self.relu(z)
+
+        # [*, N_res, N_res, C_z]
+        z = self.linear_2(z) 
+        z = z * mask
+
+        return z
+
+    @torch.jit.ignore
+    def _chunk(self,
+        z: torch.Tensor,
+        mask: torch.Tensor,
+        chunk_size: int,
+    ) -> torch.Tensor:
+        return chunk_layer(
+            self._transition,
+            {"z": z, "mask": mask},
+            chunk_size=chunk_size,
+            no_batch_dims=len(z.shape[:-2]),
+        )
+
+    def forward(self, 
+        z: torch.Tensor, 
+        mask: Optional[torch.Tensor] = None,
+        chunk_size: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            z:
+                [*, N_res, N_res, C_z] pair embedding
+        Returns:
+            [*, N_res, N_res, C_z] pair embedding update
+        """
+        # DISCREPANCY: DeepMind forgets to apply the mask in this module.
+        if mask is None:
+            mask = z.new_ones(z.shape[:-1])
+
+        # [*, N_res, N_res, 1]
+        mask = mask.unsqueeze(-1)
+
+        if chunk_size is not None:
+            z = self._chunk(z, mask, chunk_size)
+        else:
+            z = self._transition(z=z, mask=mask)
+
+        return z

model/openfold/primitives.py

@@ -0,0 +1,902 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import importlib
+import math
+from typing import Optional, Callable, List, Tuple
+import numpy as np
+
+deepspeed_is_installed = importlib.util.find_spec("deepspeed") is not None
+ds4s_is_installed = deepspeed_is_installed and importlib.util.find_spec("deepspeed.ops.deepspeed4science") is not None
+if deepspeed_is_installed:
+    import deepspeed
+
+if ds4s_is_installed:
+    from deepspeed.ops.deepspeed4science import DS4Sci_EvoformerAttention
+
+fa_is_installed = importlib.util.find_spec("flash_attn") is not None
+if fa_is_installed:
+    from flash_attn.bert_padding import unpad_input
+    from flash_attn.flash_attn_interface import flash_attn_varlen_kvpacked_func
+
+import torch
+import torch.nn as nn
+from scipy.stats import truncnorm
+
+from onescience.utils.openfold.checkpointing import get_checkpoint_fn
+from onescience.utils.openfold.kernel.attention_core import attention_core
+from onescience.utils.openfold.precision_utils import is_fp16_enabled
+from onescience.utils.openfold.tensor_utils import (
+    permute_final_dims,
+    flatten_final_dims,
+)
+
+
+DEFAULT_LMA_Q_CHUNK_SIZE = 1024
+DEFAULT_LMA_KV_CHUNK_SIZE = 4096
+
+
+def _prod(nums):
+    out = 1
+    for n in nums:
+        out = out * n
+    return out
+
+
+def _calculate_fan(linear_weight_shape, fan="fan_in"):
+    fan_out, fan_in = linear_weight_shape
+
+    if fan == "fan_in":
+        f = fan_in
+    elif fan == "fan_out":
+        f = fan_out
+    elif fan == "fan_avg":
+        f = (fan_in + fan_out) / 2
+    else:
+        raise ValueError("Invalid fan option")
+
+    return f
+
+
+def trunc_normal_init_(weights, scale=1.0, fan="fan_in"):
+    shape = weights.shape
+    f = _calculate_fan(shape, fan)
+    scale = scale / max(1, f)
+    a = -2
+    b = 2
+    std = math.sqrt(scale) / truncnorm.std(a=a, b=b, loc=0, scale=1)
+    size = _prod(shape)
+    samples = truncnorm.rvs(a=a, b=b, loc=0, scale=std, size=size)
+    samples = np.reshape(samples, shape)
+    with torch.no_grad():
+        weights.copy_(torch.tensor(samples, device=weights.device))
+
+
+def lecun_normal_init_(weights):
+    trunc_normal_init_(weights, scale=1.0)
+
+
+def he_normal_init_(weights):
+    trunc_normal_init_(weights, scale=2.0)
+
+
+def glorot_uniform_init_(weights):
+    nn.init.xavier_uniform_(weights, gain=1)
+
+
+def final_init_(weights):
+    with torch.no_grad():
+        weights.fill_(0.0)
+
+
+def gating_init_(weights):
+    with torch.no_grad():
+        weights.fill_(0.0)
+
+
+def normal_init_(weights):
+    torch.nn.init.kaiming_normal_(weights, nonlinearity="linear")
+
+
+def ipa_point_weights_init_(weights):
+    with torch.no_grad():
+        softplus_inverse_1 = 0.541324854612918
+        weights.fill_(softplus_inverse_1)
+
+
+class Linear(nn.Linear):
+    """
+    A Linear layer with built-in nonstandard initializations. Called just
+    like torch.nn.Linear.
+
+    Implements the initializers in 1.11.4, plus some additional ones found
+    in the code.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        bias: bool = True,
+        init: str = "default",
+        init_fn: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None,
+        precision=None
+    ):
+        """
+        Args:
+            in_dim:
+                The final dimension of inputs to the layer
+            out_dim:
+                The final dimension of layer outputs
+            bias:
+                Whether to learn an additive bias. True by default
+            init:
+                The initializer to use. Choose from:
+
+                "default": LeCun fan-in truncated normal initialization
+                "relu": He initialization w/ truncated normal distribution
+                "glorot": Fan-average Glorot uniform initialization
+                "gating": Weights=0, Bias=1
+                "normal": Normal initialization with std=1/sqrt(fan_in)
+                "final": Weights=0, Bias=0
+
+                Overridden by init_fn if the latter is not None.
+            init_fn:
+                A custom initializer taking weight and bias as inputs.
+                Overrides init if not None.
+        """
+        super(Linear, self).__init__(in_dim, out_dim, bias=bias)
+
+        if bias:
+            with torch.no_grad():
+                self.bias.fill_(0)
+
+        with torch.no_grad():
+            if init_fn is not None:
+                init_fn(self.weight, self.bias)
+            else:
+                if init == "default":
+                    lecun_normal_init_(self.weight)
+                elif init == "relu":
+                    he_normal_init_(self.weight)
+                elif init == "glorot":
+                    glorot_uniform_init_(self.weight)
+                elif init == "gating":
+                    gating_init_(self.weight)
+                    if bias:
+                        self.bias.fill_(1.0)
+                elif init == "normal":
+                    normal_init_(self.weight)
+                elif init == "final":
+                    final_init_(self.weight)
+                else:
+                    raise ValueError("Invalid init string.")
+
+        self.precision = precision
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        d = input.dtype
+        deepspeed_is_initialized = (
+                deepspeed_is_installed and
+                deepspeed.comm.comm.is_initialized()
+        )
+        if self.precision is not None:
+            with torch.cuda.amp.autocast(enabled=False):
+                bias = self.bias.to(dtype=self.precision) if self.bias is not None else None
+                return nn.functional.linear(input.to(dtype=self.precision),
+                                            self.weight.to(dtype=self.precision),
+                                            bias).to(dtype=d)
+
+        if d is torch.bfloat16 and not deepspeed_is_initialized:
+            with torch.cuda.amp.autocast(enabled=False):
+                bias = self.bias.to(dtype=d) if self.bias is not None else None
+                return nn.functional.linear(input, self.weight.to(dtype=d), bias)
+
+        return nn.functional.linear(input, self.weight, self.bias)
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, c_in, eps=1e-5):
+        super(LayerNorm, self).__init__()
+        
+        self.c_in = (c_in,)
+        self.eps = eps
+
+        self.weight = nn.Parameter(torch.ones(c_in))
+        self.bias = nn.Parameter(torch.zeros(c_in))
+
+    def forward(self, x): 
+        d = x.dtype
+        deepspeed_is_initialized = (
+            deepspeed_is_installed and 
+            deepspeed.comm.comm.is_initialized()
+        )
+        if d is torch.bfloat16 and not deepspeed_is_initialized:
+            with torch.cuda.amp.autocast(enabled=False):
+                out = nn.functional.layer_norm(
+                    x, 
+                    self.c_in, 
+                    self.weight.to(dtype=d), 
+                    self.bias.to(dtype=d), 
+                    self.eps
+                )
+        else:
+            out = nn.functional.layer_norm(
+                x,
+                self.c_in,
+                self.weight,
+                self.bias,
+                self.eps,
+            )
+
+        return out
+import os
+fastln_is_installed = os.getenv("LAYERNORM_TYPE", None) == "fast_layernorm"
+if fastln_is_installed:
+    # LayerNorm is a time bottomneck, so we use a custom implementation.
+    from model.protenix.layer_norm.layer_norm import FusedLayerNorm
+
+class OpenFoldLayerNorm(nn.Module):
+    def __init__(
+        self,
+        c_in,
+        create_scale: bool = True,
+        create_offset: bool = True,
+        eps=1e-5,
+    ):
+        super(OpenFoldLayerNorm, self).__init__()
+
+        self.c_in = (c_in,)
+        self.create_scale = create_scale
+        self.create_offset = create_offset
+        self.eps = eps
+
+        if self.create_scale:
+            self.weight = nn.Parameter(torch.ones(c_in))
+        else:
+            self.weight = None
+        if self.create_offset:
+            self.bias = nn.Parameter(torch.zeros(c_in))
+        else:
+            self.bias = None
+
+    def forward(self, x):
+        d = x.dtype
+        deepspeed_is_initialized = (
+            deepspeed_is_installed and deepspeed.comm.comm.is_initialized()
+        )
+        if d is torch.bfloat16 and not deepspeed_is_initialized:
+            with torch.cuda.amp.autocast(enabled=False):
+                out = nn.functional.layer_norm(
+                    x,
+                    self.c_in,
+                    self.weight.to(dtype=d) if self.weight is not None else None,
+                    self.bias.to(dtype=d) if self.bias is not None else None,
+                    self.eps,
+                )
+        else:
+            out = nn.functional.layer_norm(
+                x,
+                self.c_in,
+                self.weight,
+                self.bias,
+                self.eps,
+            )
+        return out
+
+
+# Keep the function name for code simplicity
+def ProtenixLayerNorm(
+    c_in,
+    create_scale: bool = True,
+    create_offset: bool = True,
+    eps: float = 1e-5,
+):
+    # if specify "fast_layernorm" and fastln_is_installed, use the FusedLayerNorm,
+    # Otherwise, OpenFoldLayerNorm is used!
+    if fastln_is_installed:
+        # print("use fast layernorm")
+        return FusedLayerNorm(
+            c_in, create_scale=create_scale, create_offset=create_offset, eps=eps
+        )
+    # print("use openfold layernorm")
+    return OpenFoldLayerNorm(c_in, create_scale, create_offset, eps)
+  
+
+@torch.jit.ignore
+def softmax_no_cast(t: torch.Tensor, dim: int = -1) -> torch.Tensor:
+    """
+        Softmax, but without automatic casting to fp32 when the input is of
+        type bfloat16
+    """
+    d = t.dtype
+    deepspeed_is_initialized = (
+        deepspeed_is_installed and 
+        deepspeed.comm.comm.is_initialized()
+    )
+    if d is torch.bfloat16 and not deepspeed_is_initialized:
+        with torch.cuda.amp.autocast(enabled=False):
+            s = torch.nn.functional.softmax(t, dim=dim)
+    else:
+        s = torch.nn.functional.softmax(t, dim=dim)
+
+    return s
+
+
+#@torch.jit.script
+def _attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, biases: List[torch.Tensor]) -> torch.Tensor:
+    # [*, H, C_hidden, K]
+    key = permute_final_dims(key, (1, 0))
+
+    # [*, H, Q, K]
+    a = torch.matmul(query, key)
+
+    for b in biases:
+        a += b
+
+    a = softmax_no_cast(a, -1)
+
+    # [*, H, Q, C_hidden]
+    a = torch.matmul(a, value)
+
+    return a
+
+
+@torch.jit.ignore
+def _attention_chunked_trainable(
+    query, key, value, biases, chunk_size, chunk_dim, checkpoint, 
+):
+    if checkpoint and len(biases) > 2:
+        raise ValueError(
+            "Checkpointed version permits only permits two bias terms"
+        )
+
+    def _checkpointable_attention(q, k, v, b1, b2):
+        bs = [b for b in [b1, b2] if b is not None]
+        a = _attention(q, k, v, bs)
+        return a
+
+    o_chunks = []
+    checkpoint_fn = get_checkpoint_fn()
+    count = query.shape[chunk_dim]
+    for start in range(0, count, chunk_size):
+        end = start + chunk_size
+        idx = [slice(None)] * len(query.shape)
+        idx[chunk_dim] = slice(start, end)
+        idx_tup = tuple(idx)
+        q_chunk = query[idx_tup]
+        k_chunk = key[idx_tup]
+        v_chunk = value[idx_tup]
+
+        def _slice_bias(b):
+            idx[chunk_dim] = (
+                slice(start, end) if b.shape[chunk_dim] != 1 else slice(None)
+            )
+            return b[tuple(idx)]
+
+        if checkpoint:
+            bias_1_chunk, bias_2_chunk = [
+                _slice_bias(b) if b is not None else None
+                for b in (biases + [None, None])[:2]
+            ]
+
+            o_chunk = checkpoint_fn(_checkpointable_attention,
+                q_chunk, k_chunk, v_chunk, bias_1_chunk, bias_2_chunk
+            )
+        else:
+            bias_chunks = [
+                _slice_bias(b) for b in biases
+            ]
+
+            o_chunk = _attention(q_chunk, k_chunk, v_chunk, bias_chunks)
+            
+        o_chunk = o_chunk.transpose(-2, -3)
+        o_chunks.append(o_chunk)
+
+    o = torch.cat(o_chunks, dim=chunk_dim)
+    return o
+
+
+class Attention(nn.Module):
+    """
+    Standard multi-head attention using AlphaFold's default layer
+    initialization. Allows multiple bias vectors.
+    """
+    def __init__(
+        self,
+        c_q: int,
+        c_k: int,
+        c_v: int,
+        c_hidden: int,
+        no_heads: int,
+        gating: bool = True,
+        bias: bool = True
+    ):
+        """
+        Args:
+            c_q:
+                Input dimension of query data
+            c_k:
+                Input dimension of key data
+            c_v:
+                Input dimension of value data
+            c_hidden:
+                Per-head hidden dimension
+            no_heads:
+                Number of attention heads
+            gating:
+                Whether the output should be gated using query data
+        """
+        super(Attention, self).__init__()
+
+        self.c_q = c_q
+        self.c_k = c_k
+        self.c_v = c_v
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.gating = gating
+
+        # DISCREPANCY: c_hidden is not the per-head channel dimension, as
+        # stated in the supplement, but the overall channel dimension.
+
+        self.linear_q = Linear(
+            self.c_q, self.c_hidden * self.no_heads, bias=False, init="glorot"
+        )
+        self.linear_k = Linear(
+            self.c_k, self.c_hidden * self.no_heads, bias=False, init="glorot"
+        )
+        self.linear_v = Linear(
+            self.c_v, self.c_hidden * self.no_heads, bias=False, init="glorot"
+        )
+        self.linear_o = Linear(
+            self.c_hidden * self.no_heads, self.c_q, bias=bias, init="final"
+        )
+
+        self.linear_g = None
+        if self.gating:
+            self.linear_g = Linear(
+                self.c_q, self.c_hidden * self.no_heads, bias=bias, init="gating"
+            )
+
+        self.sigmoid = nn.Sigmoid()
+
+    def _prep_qkv(self,
+        q_x: torch.Tensor, 
+        kv_x: torch.Tensor,
+        apply_scale: bool = True
+    ) -> Tuple[
+        torch.Tensor, torch.Tensor, torch.Tensor
+    ]:
+        # [*, Q/K/V, H * C_hidden]
+        q = self.linear_q(q_x)
+        k = self.linear_k(kv_x)
+        v = self.linear_v(kv_x)
+
+        # [*, Q/K, H, C_hidden]
+        q = q.view(q.shape[:-1] + (self.no_heads, -1))
+        k = k.view(k.shape[:-1] + (self.no_heads, -1))
+        v = v.view(v.shape[:-1] + (self.no_heads, -1))
+
+        # [*, H, Q/K, C_hidden]
+        q = q.transpose(-2, -3)
+        k = k.transpose(-2, -3)
+        v = v.transpose(-2, -3)
+
+        if apply_scale:
+            q /= math.sqrt(self.c_hidden)
+
+        return q, k, v
+
+    def _wrap_up(self,
+        o: torch.Tensor, 
+        q_x: torch.Tensor
+    ) -> torch.Tensor:
+        if self.linear_g is not None:
+            g = self.sigmoid(self.linear_g(q_x))
+        
+            # [*, Q, H, C_hidden]
+            g = g.view(g.shape[:-1] + (self.no_heads, -1))
+            o = o * g
+
+        # [*, Q, H * C_hidden]
+        o = flatten_final_dims(o, 2)
+
+        # [*, Q, C_q]
+        o = self.linear_o(o)
+
+        return o
+
+    def forward(
+        self,
+        q_x: torch.Tensor,
+        kv_x: torch.Tensor,
+        biases: Optional[List[torch.Tensor]] = None,
+        use_memory_efficient_kernel: bool = False,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        lma_q_chunk_size: int = DEFAULT_LMA_Q_CHUNK_SIZE,
+        lma_kv_chunk_size: int = DEFAULT_LMA_KV_CHUNK_SIZE,
+        use_flash: bool = False,
+        flash_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            q_x:
+                [*, Q, C_q] query data
+            kv_x:
+                [*, K, C_k] key data
+            biases:
+                List of biases that broadcast to [*, H, Q, K]
+            use_memory_efficient_kernel:
+                Whether to use a custom memory-efficient attention kernel.
+                This should be the default choice for most. If none of the
+                "use_<...>" flags are True, a stock PyTorch implementation
+                is used instead
+            use_deepspeed_evo_attention:
+                Whether to use DeepSpeed memory-efficient attention kernel.
+                If none of the "use_<...>" flags are True, a stock PyTorch
+                implementation is used instead
+            use_lma:
+                Whether to use low-memory attention (Staats & Rabe 2021). If
+                none of the "use_<...>" flags are True, a stock PyTorch 
+                implementation is used instead
+            lma_q_chunk_size:
+                Query chunk size (for LMA)
+            lma_kv_chunk_size:
+                Key/Value chunk size (for LMA)
+        Returns
+            [*, Q, C_q] attention update
+        """
+        if use_lma and (lma_q_chunk_size is None or lma_kv_chunk_size is None):
+            raise ValueError(
+                "If use_lma is specified, lma_q_chunk_size and "
+                "lma_kv_chunk_size must be provided"
+            )
+
+        if use_flash and biases is not None:
+            raise ValueError(
+                "use_flash is incompatible with the bias option. For masking, "
+                "use flash_mask instead"
+            )
+
+        attn_options = [use_memory_efficient_kernel, use_deepspeed_evo_attention, use_lma, use_flash]
+        if sum(attn_options) > 1:
+            raise ValueError(
+                "Choose at most one alternative attention algorithm"
+            )
+
+        if biases is None:
+            biases = []
+        
+        # DeepSpeed attention kernel applies scaling internally
+        q, k, v = self._prep_qkv(q_x, kv_x,
+                                 apply_scale=not use_deepspeed_evo_attention)
+
+        if is_fp16_enabled():
+            use_memory_efficient_kernel = False
+        
+        if use_memory_efficient_kernel:
+            if len(biases) > 2:
+                raise ValueError(
+                    "If use_memory_efficient_kernel is True, you may only "
+                    "provide up to two bias terms"
+                )
+            o = attention_core(q, k, v, *((biases + [None] * 2)[:2]))
+            o = o.transpose(-2, -3)
+        elif use_deepspeed_evo_attention:
+            if len(biases) > 2:
+                raise ValueError(
+                    "If use_deepspeed_evo_attention is True, you may only "
+                    "provide up to two bias terms"
+                )
+            o = _deepspeed_evo_attn(q, k, v, biases)
+        elif use_lma:
+            biases = [
+                b.expand(b.shape[:-2] + (q_x.shape[-2],) + (kv_x.shape[-2],)) 
+                for b in biases
+            ]
+            o = _lma(q, k, v, biases, lma_q_chunk_size, lma_kv_chunk_size)
+            o = o.transpose(-2, -3)
+        elif use_flash:
+            o = _flash_attn(q, k, v, flash_mask)
+        else:
+            o = _attention(q, k, v, biases)
+            o = o.transpose(-2, -3)
+
+        o = self._wrap_up(o, q_x)
+
+        return o
+
+
+class GlobalAttention(nn.Module):
+    def __init__(self, c_in, c_hidden, no_heads, inf, eps):
+        super(GlobalAttention, self).__init__()
+
+        self.c_in = c_in
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.inf = inf
+        self.eps = eps
+
+        self.linear_q = Linear(
+            c_in, c_hidden * no_heads, bias=False, init="glorot"
+        )
+
+        self.linear_k = Linear(
+            c_in, c_hidden, bias=False, init="glorot",
+        )
+        self.linear_v = Linear(
+            c_in, c_hidden, bias=False, init="glorot",
+        )
+        self.linear_g = Linear(c_in, c_hidden * no_heads, init="gating")
+        self.linear_o = Linear(c_hidden * no_heads, c_in, init="final")
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, 
+        m: torch.Tensor, 
+        mask: torch.Tensor,
+        use_lma: bool = False,
+    ) -> torch.Tensor:
+        # [*, N_res, C_in]
+        q = torch.sum(m * mask.unsqueeze(-1), dim=-2) / (
+            torch.sum(mask, dim=-1)[..., None] + self.eps
+        )
+
+        # [*, N_res, H * C_hidden]
+        q = self.linear_q(q)
+        q *= (self.c_hidden ** (-0.5))
+
+        # [*, N_res, H, C_hidden]
+        q = q.view(q.shape[:-1] + (self.no_heads, -1))
+
+        # [*, N_res, N_seq, C_hidden]
+        k = self.linear_k(m)
+        v = self.linear_v(m)
+
+        bias = (self.inf * (mask - 1))[..., :, None, :]
+        if not use_lma:
+            # [*, N_res, H, N_seq]
+            a = torch.matmul(
+                q,
+                k.transpose(-1, -2),  # [*, N_res, C_hidden, N_seq]
+            )
+            a += bias
+            a = softmax_no_cast(a)
+
+            # [*, N_res, H, C_hidden]
+            o = torch.matmul(
+                a,
+                v,
+            )
+        else:
+            o = _lma(
+                q, 
+                k, 
+                v, 
+                [bias], 
+                DEFAULT_LMA_Q_CHUNK_SIZE, 
+                DEFAULT_LMA_KV_CHUNK_SIZE
+            )
+
+        # [*, N_res, N_seq, C_hidden]
+        g = self.sigmoid(self.linear_g(m))
+
+        # [*, N_res, N_seq, H, C_hidden]
+        g = g.view(g.shape[:-1] + (self.no_heads, -1))
+
+        # [*, N_res, N_seq, H, C_hidden]
+        o = o.unsqueeze(-3) * g
+
+        # [*, N_res, N_seq, H * C_hidden]
+        o = o.reshape(o.shape[:-2] + (-1,))
+
+        # [*, N_res, N_seq, C_in]
+        m = self.linear_o(o)
+
+        return m
+
+
+@torch.jit.ignore
+def _deepspeed_evo_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    biases: List[torch.Tensor],
+):
+    """""
+    Compute attention using the DeepSpeed DS4Sci_EvoformerAttention kernel.
+
+    Args:
+        q:
+            [*, H, Q, C_hidden] query data
+        k:
+            [*, H, K, C_hidden] key data
+        v:
+            [*, H, V, C_hidden] value data
+        biases:
+            List of biases that broadcast to [*, H, Q, K]
+    """
+
+    if not ds4s_is_installed:
+        raise ValueError(
+            "_deepspeed_evo_attn requires that DeepSpeed be installed "
+            "and that the deepspeed.ops.deepspeed4science package exists"
+        )
+
+    def reshape_dims(x):
+        no_batch_dims = len(x.shape[:-3])
+        if no_batch_dims < 2:
+            return x.reshape(*((1,) * (2 - no_batch_dims) + x.shape))
+        if no_batch_dims > 2:
+            return x.reshape(*((x.shape[0], -1) + x.shape[-3:]))
+        return x
+
+    # [*, Q/K, H, C_hidden]
+    q = q.transpose(-2, -3)
+    k = k.transpose(-2, -3)
+    v = v.transpose(-2, -3)
+
+    # Reshape tensors to match expected input shape [B, N, Q/K, H, C_hidden]
+    # for DS4Sci_EvoformerAttention() by adding or flattening batch dims as needed.
+    orig_shape = q.shape
+    if len(orig_shape[:-3]) != 2:
+        q = reshape_dims(q)
+        k = reshape_dims(k)
+        v = reshape_dims(v)
+        biases = [reshape_dims(b) for b in biases]
+
+    # DeepSpeed attn. kernel requires inputs to be type bf16 or fp16
+    # Cast to bf16 so kernel can be used during inference
+    orig_dtype = q.dtype
+    if orig_dtype not in [torch.bfloat16, torch.float16]:
+        o = DS4Sci_EvoformerAttention(q.to(dtype=torch.bfloat16),
+                                      k.to(dtype=torch.bfloat16),
+                                      v.to(dtype=torch.bfloat16),
+                                      [b.to(dtype=torch.bfloat16) for b in biases])
+
+        o = o.to(dtype=orig_dtype)
+    else:
+        o = DS4Sci_EvoformerAttention(q, k, v, biases)
+
+    o = o.reshape(orig_shape)
+    return o
+
+
+def _lma(
+    q: torch.Tensor, 
+    k: torch.Tensor, 
+    v: torch.Tensor, 
+    biases: List[torch.Tensor], 
+    q_chunk_size: int, 
+    kv_chunk_size: int,
+):
+    no_q, no_kv = q.shape[-2], k.shape[-2]
+
+    # [*, H, Q, C_hidden]
+    o = q.new_zeros(q.shape)
+    for q_s in range(0, no_q, q_chunk_size):
+        q_chunk = q[..., q_s: q_s + q_chunk_size, :]
+        large_bias_chunks = [
+            b[..., q_s: q_s + q_chunk_size, :] for b in biases
+        ]
+
+        maxes = []
+        weights = []
+        values = []
+        for kv_s in range(0, no_kv, kv_chunk_size):
+            k_chunk = k[..., kv_s: kv_s + kv_chunk_size, :]
+            v_chunk = v[..., kv_s: kv_s + kv_chunk_size, :]
+            small_bias_chunks = [
+                b[..., kv_s: kv_s + kv_chunk_size] for b in large_bias_chunks
+            ]
+
+            a = torch.einsum(
+                "...hqd,...hkd->...hqk", q_chunk, k_chunk,
+            )
+       
+            for b in small_bias_chunks:
+                a += b
+        
+            max_a = torch.max(a, dim=-1, keepdim=True)[0]
+            exp_a = torch.exp(a - max_a)
+            exp_v = torch.einsum("...hvf,...hqv->...hqf", v_chunk, exp_a)
+ 
+            maxes.append(max_a.detach().squeeze(-1))
+            weights.append(torch.sum(exp_a, dim=-1))
+            values.append(exp_v)
+
+        chunk_max = torch.stack(maxes, dim=-3)
+        chunk_weights = torch.stack(weights, dim=-3)
+        chunk_values = torch.stack(values, dim=-4)
+
+        global_max = torch.max(chunk_max, dim=-3, keepdim=True)[0]
+        max_diffs = torch.exp(chunk_max - global_max)
+        chunk_values = chunk_values * max_diffs.unsqueeze(-1)
+        chunk_weights = chunk_weights * max_diffs
+
+        all_values = torch.sum(chunk_values, dim=-4)
+        all_weights = torch.sum(chunk_weights.unsqueeze(-1), dim=-4)
+
+        q_chunk_out = all_values / all_weights
+
+        o[..., q_s: q_s + q_chunk_size, :] = q_chunk_out
+
+    return o
+
+
+@torch.jit.ignore
+def _flash_attn(q, k, v, kv_mask):
+    if not fa_is_installed:
+        raise ValueError(
+            "_flash_attn requires that FlashAttention be installed"
+        )
+   
+    batch_dims = q.shape[:-3]
+    no_heads, n, c = q.shape[-3:]
+    dtype = q.dtype
+
+    q = q.half()
+    k = k.half()
+    v = v.half()
+    kv_mask = kv_mask.half()
+
+    # [*, B, N, H, C]
+    q = q.transpose(-2, -3)
+    k = k.transpose(-2, -3)
+    v = v.transpose(-2, -3)
+
+    # [B_flat, N, H, C]
+    q = q.reshape(-1, *q.shape[-3:])
+    k = k.reshape(-1, *k.shape[-3:])
+    v = v.reshape(-1, *v.shape[-3:])
+
+    # Flattened batch size
+    batch_size = q.shape[0]
+    
+    # [B_flat * N, H, C]
+    q = q.reshape(-1, *q.shape[-2:])
+    
+    q_max_s = n
+    q_cu_seqlens = torch.arange(
+        0, (batch_size + 1) * n, step=n, dtype=torch.int32, device=q.device
+    )
+
+    # [B_flat, N, 2, H, C]
+    kv = torch.stack([k, v], dim=-3) 
+    kv_shape = kv.shape
+    
+    # [B_flat, N, 2 * H * C]
+    kv = kv.reshape(*kv.shape[:-3], -1) 
+    
+    kv_unpad, _, kv_cu_seqlens, kv_max_s, _ = unpad_input(kv, kv_mask)
+    kv_unpad = kv_unpad.reshape(-1, *kv_shape[-3:])
+   
+    out = flash_attn_varlen_kvpacked_func(
+        q,
+        kv_unpad,
+        q_cu_seqlens,
+        kv_cu_seqlens,
+        q_max_s,
+        kv_max_s,
+        dropout_p=0.,
+        softmax_scale=1.,  # q has been scaled already
+    )
+  
+    # [*, B, N, H, C]
+    out = out.reshape(*batch_dims, n, no_heads, c) 
+
+    out = out.to(dtype=dtype)
+
+    return out

model/openfold/structure_module.py

@@ -0,0 +1,1252 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from functools import reduce
+import importlib
+import math
+import sys
+from operator import mul
+
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple, Sequence, Union
+
+from model.openfold.primitives import Linear, LayerNorm, ipa_point_weights_init_
+from onescience.utils.openfold.np.residue_constants import (
+    restype_rigid_group_default_frame,
+    restype_atom14_to_rigid_group,
+    restype_atom14_mask,
+    restype_atom14_rigid_group_positions,
+)
+from onescience.utils.openfold.geometry.quat_rigid import QuatRigid
+from onescience.utils.openfold.geometry.rigid_matrix_vector import Rigid3Array
+from onescience.utils.openfold.geometry.vector import Vec3Array, square_euclidean_distance
+from onescience.utils.openfold.feats import (
+    frames_and_literature_positions_to_atom14_pos,
+    torsion_angles_to_frames,
+)
+from onescience.utils.openfold.precision_utils import is_fp16_enabled
+from onescience.utils.openfold.rigid_utils import Rotation, Rigid
+from onescience.utils.openfold.tensor_utils import (
+    dict_multimap,
+    permute_final_dims,
+    flatten_final_dims,
+)
+
+attn_core_inplace_cuda = importlib.import_module("attn_core_inplace_cuda")
+
+
+class AngleResnetBlock(nn.Module):
+    def __init__(self, c_hidden):
+        """
+        Args:
+            c_hidden:
+                Hidden channel dimension
+        """
+        super(AngleResnetBlock, self).__init__()
+
+        self.c_hidden = c_hidden
+
+        self.linear_1 = Linear(self.c_hidden, self.c_hidden, init="relu")
+        self.linear_2 = Linear(self.c_hidden, self.c_hidden, init="final")
+
+        self.relu = nn.ReLU()
+
+    def forward(self, a: torch.Tensor) -> torch.Tensor:
+
+        s_initial = a
+
+        a = self.relu(a)
+        a = self.linear_1(a)
+        a = self.relu(a)
+        a = self.linear_2(a)
+
+        return a + s_initial
+
+
+class AngleResnet(nn.Module):
+    """
+    Implements Algorithm 20, lines 11-14
+    """
+
+    def __init__(self, c_in, c_hidden, no_blocks, no_angles, epsilon):
+        """
+        Args:
+            c_in:
+                Input channel dimension
+            c_hidden:
+                Hidden channel dimension
+            no_blocks:
+                Number of resnet blocks
+            no_angles:
+                Number of torsion angles to generate
+            epsilon:
+                Small constant for normalization
+        """
+        super(AngleResnet, self).__init__()
+
+        self.c_in = c_in
+        self.c_hidden = c_hidden
+        self.no_blocks = no_blocks
+        self.no_angles = no_angles
+        self.eps = epsilon
+
+        self.linear_in = Linear(self.c_in, self.c_hidden)
+        self.linear_initial = Linear(self.c_in, self.c_hidden)
+
+        self.layers = nn.ModuleList()
+        for _ in range(self.no_blocks):
+            layer = AngleResnetBlock(c_hidden=self.c_hidden)
+            self.layers.append(layer)
+
+        self.linear_out = Linear(self.c_hidden, self.no_angles * 2)
+
+        self.relu = nn.ReLU()
+
+    def forward(
+        self, s: torch.Tensor, s_initial: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            s:
+                [*, C_hidden] single embedding
+            s_initial:
+                [*, C_hidden] single embedding as of the start of the
+                StructureModule
+        Returns:
+            [*, no_angles, 2] predicted angles
+        """
+        # NOTE: The ReLU's applied to the inputs are absent from the supplement
+        # pseudocode but present in the source. For maximal compatibility with
+        # the pretrained weights, I'm going with the source.
+
+        # [*, C_hidden]
+        s_initial = self.relu(s_initial)
+        s_initial = self.linear_initial(s_initial)
+        s = self.relu(s)
+        s = self.linear_in(s)
+        s = s + s_initial
+
+        for l in self.layers:
+            s = l(s)
+
+        s = self.relu(s)
+
+        # [*, no_angles * 2]
+        s = self.linear_out(s)
+
+        # [*, no_angles, 2]
+        s = s.view(s.shape[:-1] + (-1, 2))
+
+        unnormalized_s = s
+        norm_denom = torch.sqrt(
+            torch.clamp(
+                torch.sum(s ** 2, dim=-1, keepdim=True),
+                min=self.eps,
+            )
+        )
+        s = s / norm_denom
+
+        return unnormalized_s, s
+
+
+class PointProjection(nn.Module):
+    def __init__(self,
+        c_hidden: int,
+        num_points: int,
+        no_heads: int,
+        is_multimer: bool,
+        return_local_points: bool = False,
+    ):
+        super().__init__()
+        self.return_local_points = return_local_points
+        self.no_heads = no_heads
+        self.num_points = num_points
+        self.is_multimer = is_multimer
+
+        # Multimer requires this to be run with fp32 precision during training
+        precision = torch.float32 if self.is_multimer else None
+        self.linear = Linear(c_hidden, no_heads * 3 * num_points, precision=precision)
+
+    def forward(self, 
+        activations: torch.Tensor, 
+        rigids: Union[Rigid, Rigid3Array],
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        # TODO: Needs to run in high precision during training
+        points_local = self.linear(activations)
+        out_shape = points_local.shape[:-1] + (self.no_heads, self.num_points, 3)
+
+        if self.is_multimer:
+            points_local = points_local.view(
+                points_local.shape[:-1] + (self.no_heads, -1)
+            )
+
+        points_local = torch.split(
+            points_local, points_local.shape[-1] // 3, dim=-1
+        )
+
+        points_local = torch.stack(points_local, dim=-1).view(out_shape)
+
+        points_global = rigids[..., None, None].apply(points_local)
+
+        if(self.return_local_points):
+            return points_global, points_local
+
+        return points_global
+
+
+class InvariantPointAttention(nn.Module):
+    """
+    Implements Algorithm 22.
+    """
+    def __init__(
+        self,
+        c_s: int,
+        c_z: int,
+        c_hidden: int,
+        no_heads: int,
+        no_qk_points: int,
+        no_v_points: int,
+        inf: float = 1e5,
+        eps: float = 1e-8,
+        is_multimer: bool = False,
+    ):
+        """
+        Args:
+            c_s:
+                Single representation channel dimension
+            c_z:
+                Pair representation channel dimension
+            c_hidden:
+                Hidden channel dimension
+            no_heads:
+                Number of attention heads
+            no_qk_points:
+                Number of query/key points to generate
+            no_v_points:
+                Number of value points to generate
+        """
+        super(InvariantPointAttention, self).__init__()
+
+        self.c_s = c_s
+        self.c_z = c_z
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.no_qk_points = no_qk_points
+        self.no_v_points = no_v_points
+        self.inf = inf
+        self.eps = eps
+        self.is_multimer = is_multimer
+
+        # These linear layers differ from their specifications in the
+        # supplement. There, they lack bias and use Glorot initialization.
+        # Here as in the official source, they have bias and use the default
+        # Lecun initialization.
+        hc = self.c_hidden * self.no_heads
+        self.linear_q = Linear(self.c_s, hc, bias=(not is_multimer))
+
+        self.linear_q_points = PointProjection(
+            self.c_s,
+            self.no_qk_points,
+            self.no_heads,
+            self.is_multimer
+        )
+
+        if(is_multimer):
+            self.linear_k = Linear(self.c_s, hc, bias=False)
+            self.linear_v = Linear(self.c_s, hc, bias=False)
+            self.linear_k_points = PointProjection(
+                self.c_s,
+                self.no_qk_points,
+                self.no_heads,
+                self.is_multimer
+            )
+
+            self.linear_v_points = PointProjection(
+                self.c_s,
+                self.no_v_points,
+                self.no_heads,
+                self.is_multimer
+            )
+        else:
+            self.linear_kv = Linear(self.c_s, 2 * hc)
+            self.linear_kv_points = PointProjection(
+                self.c_s,
+                self.no_qk_points + self.no_v_points,
+                self.no_heads,
+                self.is_multimer
+            )
+
+        self.linear_b = Linear(self.c_z, self.no_heads)
+
+        self.head_weights = nn.Parameter(torch.zeros((no_heads)))
+        ipa_point_weights_init_(self.head_weights)
+
+        concat_out_dim = self.no_heads * (
+            self.c_z + self.c_hidden + self.no_v_points * 4
+        )
+        self.linear_out = Linear(concat_out_dim, self.c_s, init="final")
+
+        self.softmax = nn.Softmax(dim=-1)
+        self.softplus = nn.Softplus()
+
+    def forward(
+        self,
+        s: torch.Tensor,
+        z: torch.Tensor,
+        r: Union[Rigid, Rigid3Array],
+        mask: torch.Tensor,
+        inplace_safe: bool = False,
+        _offload_inference: bool = False,
+        _z_reference_list: Optional[Sequence[torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            s:
+                [*, N_res, C_s] single representation
+            z:
+                [*, N_res, N_res, C_z] pair representation
+            r:
+                [*, N_res] transformation object
+            mask:
+                [*, N_res] mask
+        Returns:
+            [*, N_res, C_s] single representation update
+        """
+        if (_offload_inference and inplace_safe):
+            z = _z_reference_list
+        else:
+            z = [z]
+
+        #######################################
+        # Generate scalar and point activations
+        #######################################
+        # [*, N_res, H * C_hidden]
+        q = self.linear_q(s)
+
+        # [*, N_res, H, C_hidden]
+        q = q.view(q.shape[:-1] + (self.no_heads, -1))
+
+        # [*, N_res, H, P_qk]
+        q_pts = self.linear_q_points(s, r)
+
+        # The following two blocks are equivalent
+        # They're separated only to preserve compatibility with old AF weights
+        if(self.is_multimer):
+            # [*, N_res, H * C_hidden]
+            k = self.linear_k(s)
+            v = self.linear_v(s)
+
+            # [*, N_res, H, C_hidden]
+            k = k.view(k.shape[:-1] + (self.no_heads, -1))
+            v = v.view(v.shape[:-1] + (self.no_heads, -1))
+
+            # [*, N_res, H, P_qk, 3]
+            k_pts = self.linear_k_points(s, r)
+
+            # [*, N_res, H, P_v, 3]
+            v_pts = self.linear_v_points(s, r)
+        else:
+            # [*, N_res, H * 2 * C_hidden]
+            kv = self.linear_kv(s)
+
+            # [*, N_res, H, 2 * C_hidden]
+            kv = kv.view(kv.shape[:-1] + (self.no_heads, -1))
+
+            # [*, N_res, H, C_hidden]
+            k, v = torch.split(kv, self.c_hidden, dim=-1)
+
+            kv_pts = self.linear_kv_points(s, r)
+
+            # [*, N_res, H, P_q/P_v, 3]
+            k_pts, v_pts = torch.split(
+                kv_pts, [self.no_qk_points, self.no_v_points], dim=-2
+            )
+
+        ##########################
+        # Compute attention scores
+        ##########################
+        # [*, N_res, N_res, H]
+        b = self.linear_b(z[0])
+
+        if (_offload_inference):
+            assert (sys.getrefcount(z[0]) == 2)
+            z[0] = z[0].cpu()
+
+        # [*, H, N_res, N_res]
+        if (is_fp16_enabled()):
+            with torch.cuda.amp.autocast(enabled=False):
+                a = torch.matmul(
+                    permute_final_dims(q.float(), (1, 0, 2)),  # [*, H, N_res, C_hidden]
+                    permute_final_dims(k.float(), (1, 2, 0)),  # [*, H, C_hidden, N_res]
+                )
+        else:
+            a = torch.matmul(
+                permute_final_dims(q, (1, 0, 2)),  # [*, H, N_res, C_hidden]
+                permute_final_dims(k, (1, 2, 0)),  # [*, H, C_hidden, N_res]
+            )
+
+        a *= math.sqrt(1.0 / (3 * self.c_hidden))
+        a += (math.sqrt(1.0 / 3) * permute_final_dims(b, (2, 0, 1)))
+
+        # [*, N_res, N_res, H, P_q, 3]
+        pt_att = q_pts.unsqueeze(-4) - k_pts.unsqueeze(-5)
+
+        if (inplace_safe):
+            pt_att *= pt_att
+        else:
+            pt_att = pt_att ** 2
+
+        pt_att = sum(torch.unbind(pt_att, dim=-1))
+
+        head_weights = self.softplus(self.head_weights).view(
+            *((1,) * len(pt_att.shape[:-2]) + (-1, 1))
+        )
+        head_weights = head_weights * math.sqrt(
+            1.0 / (3 * (self.no_qk_points * 9.0 / 2))
+        )
+
+        if (inplace_safe):
+            pt_att *= head_weights
+        else:
+            pt_att = pt_att * head_weights
+
+        # [*, N_res, N_res, H]
+        pt_att = torch.sum(pt_att, dim=-1) * (-0.5)
+
+        # [*, N_res, N_res]
+        square_mask = mask.unsqueeze(-1) * mask.unsqueeze(-2)
+        square_mask = self.inf * (square_mask - 1)
+
+        # [*, H, N_res, N_res]
+        pt_att = permute_final_dims(pt_att, (2, 0, 1))
+
+        if (inplace_safe):
+            a += pt_att
+            del pt_att
+            a += square_mask.unsqueeze(-3)
+            # in-place softmax
+            attn_core_inplace_cuda.forward_(
+                a,
+                reduce(mul, a.shape[:-1]),
+                a.shape[-1],
+            )
+        else:
+            a = a + pt_att
+            a = a + square_mask.unsqueeze(-3)
+            a = self.softmax(a)
+
+        ################
+        # Compute output
+        ################
+        # [*, N_res, H, C_hidden]
+        o = torch.matmul(
+            a, v.transpose(-2, -3).to(dtype=a.dtype)
+        ).transpose(-2, -3)
+
+        # [*, N_res, H * C_hidden]
+        o = flatten_final_dims(o, 2)
+
+        # [*, H, 3, N_res, P_v]
+        if (inplace_safe):
+            v_pts = permute_final_dims(v_pts, (1, 3, 0, 2))
+            o_pt = [
+                torch.matmul(a, v.to(a.dtype))
+                for v in torch.unbind(v_pts, dim=-3)
+            ]
+            o_pt = torch.stack(o_pt, dim=-3)
+        else:
+            o_pt = torch.sum(
+                (
+                        a[..., None, :, :, None]
+                        * permute_final_dims(v_pts, (1, 3, 0, 2))[..., None, :, :]
+                ),
+                dim=-2,
+            )
+
+        # [*, N_res, H, P_v, 3]
+        o_pt = permute_final_dims(o_pt, (2, 0, 3, 1))
+        o_pt = r[..., None, None].invert_apply(o_pt)
+
+        # [*, N_res, H * P_v]
+        o_pt_norm = flatten_final_dims(
+            torch.sqrt(torch.sum(o_pt ** 2, dim=-1) + self.eps), 2
+        )
+
+        # [*, N_res, H * P_v, 3]
+        o_pt = o_pt.reshape(*o_pt.shape[:-3], -1, 3)
+        o_pt = torch.unbind(o_pt, dim=-1)
+
+        if (_offload_inference):
+            z[0] = z[0].to(o_pt.device)
+
+        # [*, N_res, H, C_z]
+        o_pair = torch.matmul(a.transpose(-2, -3), z[0].to(dtype=a.dtype))
+
+        # [*, N_res, H * C_z]
+        o_pair = flatten_final_dims(o_pair, 2)
+
+        # [*, N_res, C_s]
+        s = self.linear_out(
+            torch.cat(
+                (o, *o_pt, o_pt_norm, o_pair), dim=-1
+            ).to(dtype=z[0].dtype)
+        )
+
+        return s
+
+
+#TODO: This module follows the refactoring done in IPA for multimer. Running the regular IPA above
+# in multimer mode should be equivalent, but tests do not pass unless using this version. Determine
+# whether or not the increase in test error matters in practice.
+class InvariantPointAttentionMultimer(nn.Module):
+    """
+    Implements Algorithm 22.
+    """
+    def __init__(
+        self,
+        c_s: int,
+        c_z: int,
+        c_hidden: int,
+        no_heads: int,
+        no_qk_points: int,
+        no_v_points: int,
+        inf: float = 1e5,
+        eps: float = 1e-8,
+        is_multimer: bool = True,
+    ):
+        """
+        Args:
+            c_s:
+                Single representation channel dimension
+            c_z:
+                Pair representation channel dimension
+            c_hidden:
+                Hidden channel dimension
+            no_heads:
+                Number of attention heads
+            no_qk_points:
+                Number of query/key points to generate
+            no_v_points:
+                Number of value points to generate
+        """
+        super(InvariantPointAttentionMultimer, self).__init__()
+
+        self.c_s = c_s
+        self.c_z = c_z
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.no_qk_points = no_qk_points
+        self.no_v_points = no_v_points
+        self.inf = inf
+        self.eps = eps
+
+        # These linear layers differ from their specifications in the
+        # supplement. There, they lack bias and use Glorot initialization.
+        # Here as in the official source, they have bias and use the default
+        # Lecun initialization.
+        hc = self.c_hidden * self.no_heads
+        self.linear_q = Linear(self.c_s, hc, bias=False)
+
+        self.linear_q_points = PointProjection(
+            self.c_s,
+            self.no_qk_points,
+            self.no_heads,
+            is_multimer=True
+        )
+
+        self.linear_k = Linear(self.c_s, hc, bias=False)
+        self.linear_v = Linear(self.c_s, hc, bias=False)
+        self.linear_k_points = PointProjection(
+            self.c_s,
+            self.no_qk_points,
+            self.no_heads,
+            is_multimer=True
+        )
+
+        self.linear_v_points = PointProjection(
+            self.c_s,
+            self.no_v_points,
+            self.no_heads,
+            is_multimer=True
+        )
+
+        self.linear_b = Linear(self.c_z, self.no_heads)
+
+        self.head_weights = nn.Parameter(torch.zeros((no_heads)))
+        ipa_point_weights_init_(self.head_weights)
+
+        concat_out_dim = self.no_heads * (
+            self.c_z + self.c_hidden + self.no_v_points * 4
+        )
+        self.linear_out = Linear(concat_out_dim, self.c_s, init="final")
+
+        self.softmax = nn.Softmax(dim=-2)
+
+    def forward(
+        self,
+        s: torch.Tensor,
+        z: Optional[torch.Tensor],
+        r: Union[Rigid, Rigid3Array],
+        mask: torch.Tensor,
+        inplace_safe: bool = False,
+        _offload_inference: bool = False,
+        _z_reference_list: Optional[Sequence[torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            s:
+                [*, N_res, C_s] single representation
+            z:
+                [*, N_res, N_res, C_z] pair representation
+            r:
+                [*, N_res] transformation object
+            mask:
+                [*, N_res] mask
+        Returns:
+            [*, N_res, C_s] single representation update
+        """
+        if(_offload_inference and inplace_safe):
+            z = _z_reference_list
+        else:
+            z = [z]
+
+        a = 0.
+
+        point_variance = (max(self.no_qk_points, 1) * 9.0 / 2)
+        point_weights = math.sqrt(1.0 / point_variance)
+
+        softplus = lambda x: torch.logaddexp(x, torch.zeros_like(x))
+
+        head_weights = softplus(self.head_weights)
+        point_weights = point_weights * head_weights
+
+        #######################################
+        # Generate scalar and point activations
+        #######################################
+
+        # [*, N_res, H, P_qk]
+        q_pts = Vec3Array.from_array(self.linear_q_points(s, r))
+
+        # [*, N_res, H, P_qk, 3]
+        k_pts = Vec3Array.from_array(self.linear_k_points(s, r))
+
+        pt_att = square_euclidean_distance(q_pts.unsqueeze(-3), k_pts.unsqueeze(-4), epsilon=0.)
+        pt_att = torch.sum(pt_att * point_weights[..., None], dim=-1) * (-0.5)
+        pt_att = pt_att.to(dtype=s.dtype)
+        a = a + pt_att
+
+        scalar_variance = max(self.c_hidden, 1) * 1.
+        scalar_weights = math.sqrt(1.0 / scalar_variance)
+
+        # [*, N_res, H * C_hidden]
+        q = self.linear_q(s)
+        k = self.linear_k(s)
+
+        # [*, N_res, H, C_hidden]
+        q = q.view(q.shape[:-1] + (self.no_heads, -1))
+        k = k.view(k.shape[:-1] + (self.no_heads, -1))
+
+        q = q * scalar_weights
+        a = a + torch.einsum('...qhc,...khc->...qkh', q, k)
+
+        ##########################
+        # Compute attention scores
+        ##########################
+        # [*, N_res, N_res, H]
+        b = self.linear_b(z[0])
+
+        if (_offload_inference):
+            assert (sys.getrefcount(z[0]) == 2)
+            z[0] = z[0].cpu()
+
+        a = a + b
+
+        # [*, N_res, N_res]
+        square_mask = mask.unsqueeze(-1) * mask.unsqueeze(-2)
+        square_mask = self.inf * (square_mask - 1)
+
+        a = a + square_mask.unsqueeze(-1)
+        a = a * math.sqrt(1. / 3)  # Normalize by number of logit terms (3)
+        a = self.softmax(a)
+
+        # [*, N_res, H * C_hidden]
+        v = self.linear_v(s)
+
+        # [*, N_res, H, C_hidden]
+        v = v.view(v.shape[:-1] + (self.no_heads, -1))
+
+        o = torch.einsum('...qkh, ...khc->...qhc', a, v)
+
+        # [*, N_res, H * C_hidden]
+        o = flatten_final_dims(o, 2)
+
+        # [*, N_res, H, P_v, 3]
+        v_pts = Vec3Array.from_array(self.linear_v_points(s, r))
+
+        # [*, N_res, H, P_v]
+        o_pt = v_pts[..., None, :, :, :] * a.unsqueeze(-1)
+        o_pt = o_pt.sum(dim=-3)
+        # o_pt = Vec3Array(
+        #     torch.sum(a.unsqueeze(-1) * v_pts[..., None, :, :, :].x, dim=-3),
+        #     torch.sum(a.unsqueeze(-1) * v_pts[..., None, :, :, :].y, dim=-3),
+        #     torch.sum(a.unsqueeze(-1) * v_pts[..., None, :, :, :].z, dim=-3),
+        # )
+
+        # [*, N_res, H * P_v, 3]
+        o_pt = o_pt.reshape(o_pt.shape[:-2] + (-1,))
+
+        # [*, N_res, H, P_v]
+        o_pt = r[..., None].apply_inverse_to_point(o_pt)
+        o_pt_flat = [o_pt.x, o_pt.y, o_pt.z]
+        o_pt_flat = [x.to(dtype=a.dtype) for x in o_pt_flat]
+
+        # [*, N_res, H * P_v]
+        o_pt_norm = o_pt.norm(epsilon=1e-8)
+
+        if (_offload_inference):
+            z[0] = z[0].to(o_pt.x.device)
+
+        o_pair = torch.einsum('...ijh, ...ijc->...ihc', a, z[0].to(dtype=a.dtype))
+
+        # [*, N_res, H * C_z]
+        o_pair = flatten_final_dims(o_pair, 2)
+
+        # [*, N_res, C_s]
+        s = self.linear_out(
+            torch.cat(
+                (o, *o_pt_flat, o_pt_norm, o_pair), dim=-1
+            ).to(dtype=z[0].dtype)
+        )
+
+        return s
+
+
+class BackboneUpdate(nn.Module):
+    """
+    Implements part of Algorithm 23.
+    """
+
+    def __init__(self, c_s):
+        """
+        Args:
+            c_s:
+                Single representation channel dimension
+        """
+        super(BackboneUpdate, self).__init__()
+
+        self.c_s = c_s
+
+        self.linear = Linear(self.c_s, 6, init="final")
+
+    def forward(self, s: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            [*, N_res, C_s] single representation
+        Returns:
+            [*, N_res, 6] update vector 
+        """
+        # [*, 6]
+        update = self.linear(s)
+
+        return update 
+
+
+class StructureModuleTransitionLayer(nn.Module):
+    def __init__(self, c):
+        super(StructureModuleTransitionLayer, self).__init__()
+
+        self.c = c
+
+        self.linear_1 = Linear(self.c, self.c, init="relu")
+        self.linear_2 = Linear(self.c, self.c, init="relu")
+        self.linear_3 = Linear(self.c, self.c, init="final")
+
+        self.relu = nn.ReLU()
+
+    def forward(self, s):
+        s_initial = s
+        s = self.linear_1(s)
+        s = self.relu(s)
+        s = self.linear_2(s)
+        s = self.relu(s)
+        s = self.linear_3(s)
+
+        s = s + s_initial
+
+        return s
+
+
+class StructureModuleTransition(nn.Module):
+    def __init__(self, c, num_layers, dropout_rate):
+        super(StructureModuleTransition, self).__init__()
+
+        self.c = c
+        self.num_layers = num_layers
+        self.dropout_rate = dropout_rate
+
+        self.layers = nn.ModuleList()
+        for _ in range(self.num_layers):
+            l = StructureModuleTransitionLayer(self.c)
+            self.layers.append(l)
+
+        self.dropout = nn.Dropout(self.dropout_rate)
+        self.layer_norm = LayerNorm(self.c)
+
+    def forward(self, s):
+        for l in self.layers:
+            s = l(s)
+
+        s = self.dropout(s)
+        s = self.layer_norm(s)
+
+        return s
+
+
+class StructureModule(nn.Module):
+    def __init__(
+        self,
+        c_s,
+        c_z,
+        c_ipa,
+        c_resnet,
+        no_heads_ipa,
+        no_qk_points,
+        no_v_points,
+        dropout_rate,
+        no_blocks,
+        no_transition_layers,
+        no_resnet_blocks,
+        no_angles,
+        trans_scale_factor,
+        epsilon,
+        inf,
+        is_multimer=False,
+        **kwargs,
+    ):
+        """
+        Args:
+            c_s:
+                Single representation channel dimension
+            c_z:
+                Pair representation channel dimension
+            c_ipa:
+                IPA hidden channel dimension
+            c_resnet:
+                Angle resnet (Alg. 23 lines 11-14) hidden channel dimension
+            no_heads_ipa:
+                Number of IPA heads
+            no_qk_points:
+                Number of query/key points to generate during IPA
+            no_v_points:
+                Number of value points to generate during IPA
+            dropout_rate:
+                Dropout rate used throughout the layer
+            no_blocks:
+                Number of structure module blocks
+            no_transition_layers:
+                Number of layers in the single representation transition
+                (Alg. 23 lines 8-9)
+            no_resnet_blocks:
+                Number of blocks in the angle resnet
+            no_angles:
+                Number of angles to generate in the angle resnet
+            trans_scale_factor:
+                Scale of single representation transition hidden dimension
+            epsilon:
+                Small number used in angle resnet normalization
+            inf:
+                Large number used for attention masking
+        """
+        super(StructureModule, self).__init__()
+
+        self.c_s = c_s
+        self.c_z = c_z
+        self.c_ipa = c_ipa
+        self.c_resnet = c_resnet
+        self.no_heads_ipa = no_heads_ipa
+        self.no_qk_points = no_qk_points
+        self.no_v_points = no_v_points
+        self.dropout_rate = dropout_rate
+        self.no_blocks = no_blocks
+        self.no_transition_layers = no_transition_layers
+        self.no_resnet_blocks = no_resnet_blocks
+        self.no_angles = no_angles
+        self.trans_scale_factor = trans_scale_factor
+        self.epsilon = epsilon
+        self.inf = inf
+        self.is_multimer = is_multimer
+
+        # Buffers to be lazily initialized later
+        # self.default_frames
+        # self.group_idx
+        # self.atom_mask
+        # self.lit_positions
+
+        self.layer_norm_s = LayerNorm(self.c_s)
+        self.layer_norm_z = LayerNorm(self.c_z)
+
+        self.linear_in = Linear(self.c_s, self.c_s)
+
+        ipa = InvariantPointAttention if not self.is_multimer else InvariantPointAttentionMultimer
+        self.ipa = ipa(
+            self.c_s,
+            self.c_z,
+            self.c_ipa,
+            self.no_heads_ipa,
+            self.no_qk_points,
+            self.no_v_points,
+            inf=self.inf,
+            eps=self.epsilon,
+            is_multimer=self.is_multimer,
+        )
+
+        self.ipa_dropout = nn.Dropout(self.dropout_rate)
+        self.layer_norm_ipa = LayerNorm(self.c_s)
+
+        self.transition = StructureModuleTransition(
+            self.c_s,
+            self.no_transition_layers,
+            self.dropout_rate,
+        )
+
+        if self.is_multimer:
+            self.bb_update = QuatRigid(self.c_s, full_quat=False)
+        else:
+            self.bb_update = BackboneUpdate(self.c_s)
+
+        self.angle_resnet = AngleResnet(
+            self.c_s,
+            self.c_resnet,
+            self.no_resnet_blocks,
+            self.no_angles,
+            self.epsilon,
+        )
+
+    def _forward_monomer(
+        self,
+        evoformer_output_dict,
+        aatype,
+        mask=None,
+        inplace_safe=False,
+        _offload_inference=False,
+    ):
+        """
+        Args:
+            evoformer_output_dict:
+                Dictionary containing:
+                    "single":
+                        [*, N_res, C_s] single representation
+                    "pair":
+                        [*, N_res, N_res, C_z] pair representation
+            aatype:
+                [*, N_res] amino acid indices
+            mask:
+                Optional [*, N_res] sequence mask
+        Returns:
+            A dictionary of outputs
+        """
+        s = evoformer_output_dict["single"]
+
+        if mask is None:
+            # [*, N]
+            mask = s.new_ones(s.shape[:-1])
+
+        # [*, N, C_s]
+        s = self.layer_norm_s(s)
+
+        # [*, N, N, C_z]
+        z = self.layer_norm_z(evoformer_output_dict["pair"])
+
+        z_reference_list = None
+        if (_offload_inference):
+            assert (sys.getrefcount(evoformer_output_dict["pair"]) == 2)
+            evoformer_output_dict["pair"] = evoformer_output_dict["pair"].cpu()
+            z_reference_list = [z]
+            z = None
+
+        # [*, N, C_s]
+        s_initial = s
+        s = self.linear_in(s)
+
+        # [*, N]
+        rigids = Rigid.identity(
+            s.shape[:-1], 
+            s.dtype, 
+            s.device, 
+            self.training,
+            fmt="quat",
+        )
+        outputs = []
+        for i in range(self.no_blocks):
+            # [*, N, C_s]
+            s = s + self.ipa(
+                s, 
+                z, 
+                rigids, 
+                mask, 
+                inplace_safe=inplace_safe,
+                _offload_inference=_offload_inference, 
+                _z_reference_list=z_reference_list
+            )
+            s = self.ipa_dropout(s)
+            s = self.layer_norm_ipa(s)
+            s = self.transition(s)
+           
+            # [*, N]
+            rigids = rigids.compose_q_update_vec(self.bb_update(s))
+
+            # To hew as closely as possible to AlphaFold, we convert our
+            # quaternion-based transformations to rotation-matrix ones
+            # here
+            backb_to_global = Rigid(
+                Rotation(
+                    rot_mats=rigids.get_rots().get_rot_mats(), 
+                    quats=None
+                ),
+                rigids.get_trans(),
+            )
+
+            backb_to_global = backb_to_global.scale_translation(
+                self.trans_scale_factor
+            )
+
+            # [*, N, 7, 2]
+            unnormalized_angles, angles = self.angle_resnet(s, s_initial)
+
+            all_frames_to_global = self.torsion_angles_to_frames(
+                backb_to_global,
+                angles,
+                aatype,
+            )
+
+            pred_xyz = self.frames_and_literature_positions_to_atom14_pos(
+                all_frames_to_global,
+                aatype,
+            )
+
+            scaled_rigids = rigids.scale_translation(self.trans_scale_factor)
+            
+            preds = {
+                "frames": scaled_rigids.to_tensor_7(),
+                "sidechain_frames": all_frames_to_global.to_tensor_4x4(),
+                "unnormalized_angles": unnormalized_angles,
+                "angles": angles,
+                "positions": pred_xyz,
+                "states": s,
+            }
+
+            outputs.append(preds)
+
+            rigids = rigids.stop_rot_gradient()
+
+        del z, z_reference_list
+
+        if (_offload_inference):
+            evoformer_output_dict["pair"] = (
+                evoformer_output_dict["pair"].to(s.device)
+            )
+
+        outputs = dict_multimap(torch.stack, outputs)
+        outputs["single"] = s
+
+        return outputs
+
+    def _forward_multimer(
+            self,
+            evoformer_output_dict,
+            aatype,
+            mask=None,
+            inplace_safe=False,
+            _offload_inference=False,
+    ):
+        s = evoformer_output_dict["single"]
+
+        if mask is None:
+            # [*, N]
+            mask = s.new_ones(s.shape[:-1])
+
+        # [*, N, C_s]
+        s = self.layer_norm_s(s)
+
+        # [*, N, N, C_z]
+        z = self.layer_norm_z(evoformer_output_dict["pair"])
+
+        z_reference_list = None
+        if (_offload_inference):
+            assert (sys.getrefcount(evoformer_output_dict["pair"]) == 2)
+            evoformer_output_dict["pair"] = evoformer_output_dict["pair"].cpu()
+            z_reference_list = [z]
+            z = None
+
+        # [*, N, C_s]
+        s_initial = s
+        s = self.linear_in(s)
+
+        # [*, N]
+        rigids = Rigid3Array.identity(
+            s.shape[:-1], 
+            s.device, 
+        )
+        outputs = []
+        for i in range(self.no_blocks):
+            # [*, N, C_s]
+            s = s + self.ipa(
+                s,
+                z,
+                rigids,
+                mask,
+                inplace_safe=inplace_safe,
+                _offload_inference=_offload_inference,
+                _z_reference_list=z_reference_list
+            )
+            s = self.ipa_dropout(s)
+            s = self.layer_norm_ipa(s)
+            s = self.transition(s)
+
+            # [*, N]
+            rigids = rigids @ self.bb_update(s)
+
+            # [*, N, 7, 2]
+            unnormalized_angles, angles = self.angle_resnet(s, s_initial)
+
+            all_frames_to_global = self.torsion_angles_to_frames(
+                rigids.scale_translation(self.trans_scale_factor),
+                angles,
+                aatype,
+            )
+
+            pred_xyz = self.frames_and_literature_positions_to_atom14_pos(
+                all_frames_to_global,
+                aatype,
+            )
+            
+            preds = {
+                "frames": rigids.scale_translation(self.trans_scale_factor).to_tensor(),
+                "sidechain_frames": all_frames_to_global.to_tensor_4x4(),
+                "unnormalized_angles": unnormalized_angles,
+                "angles": angles,
+                "positions": pred_xyz,
+            }
+
+            preds = {k: v.to(dtype=s.dtype) for k, v in preds.items()}
+
+            outputs.append(preds)
+
+            rigids = rigids.stop_rot_gradient()
+
+        del z, z_reference_list
+
+        if (_offload_inference):
+            evoformer_output_dict["pair"] = (
+                evoformer_output_dict["pair"].to(s.device)
+            )
+
+        outputs = dict_multimap(torch.stack, outputs)
+        outputs["single"] = s
+
+        return outputs
+
+    def forward(
+        self,
+        evoformer_output_dict,
+        aatype,
+        mask=None,
+        inplace_safe=False,
+        _offload_inference=False,
+    ):
+        """
+        Args:
+            s:
+                [*, N_res, C_s] single representation
+            z:
+                [*, N_res, N_res, C_z] pair representation
+            aatype:
+                [*, N_res] amino acid indices
+            mask:
+                Optional [*, N_res] sequence mask
+        Returns:
+            A dictionary of outputs
+        """
+        if(self.is_multimer):
+            outputs = self._forward_multimer(evoformer_output_dict, aatype, mask, inplace_safe, _offload_inference)
+        else:
+            outputs = self._forward_monomer(evoformer_output_dict, aatype, mask, inplace_safe, _offload_inference)
+
+        return outputs
+
+    def _init_residue_constants(self, float_dtype, device):
+        if not hasattr(self, "default_frames"):
+            self.register_buffer(
+                "default_frames",
+                torch.tensor(
+                    restype_rigid_group_default_frame,
+                    dtype=float_dtype,
+                    device=device,
+                    requires_grad=False,
+                ),
+                persistent=False,
+            )
+        if not hasattr(self, "group_idx"):
+            self.register_buffer(
+                "group_idx",
+                torch.tensor(
+                    restype_atom14_to_rigid_group,
+                    device=device,
+                    requires_grad=False,
+                ),
+                persistent=False,
+            )
+        if not hasattr(self, "atom_mask"):
+            self.register_buffer(
+                "atom_mask",
+                torch.tensor(
+                    restype_atom14_mask,
+                    dtype=float_dtype,
+                    device=device,
+                    requires_grad=False,
+                ),
+                persistent=False,
+            )
+        if not hasattr(self, "lit_positions"):
+            self.register_buffer(
+                "lit_positions",
+                torch.tensor(
+                    restype_atom14_rigid_group_positions,
+                    dtype=float_dtype,
+                    device=device,
+                    requires_grad=False,
+                ),
+                persistent=False,
+            )
+
+    def torsion_angles_to_frames(self, r, alpha, f):
+        # Lazily initialize the residue constants on the correct device
+        self._init_residue_constants(alpha.dtype, alpha.device)
+        # Separated purely to make testing less annoying
+        return torsion_angles_to_frames(r, alpha, f, self.default_frames)
+
+    def frames_and_literature_positions_to_atom14_pos(
+            self, r, f  # [*, N, 8]  # [*, N]
+    ):
+        # Lazily initialize the residue constants on the correct device
+        self._init_residue_constants(r.dtype, r.device)
+        return frames_and_literature_positions_to_atom14_pos(
+            r,
+            f,
+            self.default_frames,
+            self.group_idx,
+            self.atom_mask,
+            self.lit_positions,
+        )

model/openfold/template.py

@@ -0,0 +1,693 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from functools import partial
+import math
+import sys
+from typing import Optional, List
+
+import torch
+import torch.nn as nn
+
+from model.openfold.primitives import LayerNorm, Attention
+from model.openfold.dropout import (
+    DropoutRowwise,
+    DropoutColumnwise,
+)
+from model.openfold.pair_transition import PairTransition
+from model.openfold.triangular_attention import (
+    TriangleAttentionStartingNode,
+    TriangleAttentionEndingNode,
+)
+from model.openfold.triangular_multiplicative_update import (
+    TriangleMultiplicationOutgoing,
+    TriangleMultiplicationIncoming,
+    FusedTriangleMultiplicationOutgoing,
+    FusedTriangleMultiplicationIncoming
+)
+from onescience.utils.openfold.checkpointing import checkpoint_blocks
+from onescience.utils.openfold.chunk_utils import (
+    chunk_layer,
+    ChunkSizeTuner,
+)
+from onescience.utils.openfold.feats import (
+    build_template_angle_feat,
+    build_template_pair_feat,
+)
+from onescience.utils.openfold.tensor_utils import (
+    add,
+    permute_final_dims,
+    tensor_tree_map,
+)
+
+
+class TemplatePointwiseAttention(nn.Module):
+    """
+    Implements Algorithm 17.
+    """
+
+    def __init__(self, c_t, c_z, c_hidden, no_heads, inf, **kwargs):
+        """
+        Args:
+            c_t:
+                Template embedding channel dimension
+            c_z:
+                Pair embedding channel dimension
+            c_hidden:
+                Hidden channel dimension
+        """
+        super(TemplatePointwiseAttention, self).__init__()
+
+        self.c_t = c_t
+        self.c_z = c_z
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.inf = inf
+
+        self.mha = Attention(
+            self.c_z,
+            self.c_t,
+            self.c_t,
+            self.c_hidden,
+            self.no_heads,
+            gating=False,
+        )
+
+    def _chunk(self,
+               z: torch.Tensor,
+               t: torch.Tensor,
+               biases: List[torch.Tensor],
+               chunk_size: int,
+               use_lma: bool = False,
+               ) -> torch.Tensor:
+        mha_inputs = {
+            "q_x": z,
+            "kv_x": t,
+            "biases": biases,
+        }
+        return chunk_layer(
+            partial(self.mha, use_lma=use_lma),
+            mha_inputs,
+            chunk_size=chunk_size,
+            no_batch_dims=len(z.shape[:-2]),
+        )
+
+    def forward(self,
+                t: torch.Tensor,
+                z: torch.Tensor,
+                template_mask: Optional[torch.Tensor] = None,
+                # This module suffers greatly from a small chunk size
+                chunk_size: Optional[int] = 256,
+                use_lma: bool = False,
+                ) -> torch.Tensor:
+        """
+        Args:
+            t:
+                [*, N_templ, N_res, N_res, C_t] template embedding
+            z:
+                [*, N_res, N_res, C_t] pair embedding
+            template_mask:
+                [*, N_templ] template mask
+        Returns:
+            [*, N_res, N_res, C_z] pair embedding update
+        """
+        if template_mask is None:
+            template_mask = t.new_ones(t.shape[:-3])
+
+        bias = self.inf * (template_mask[..., None, None, None, None, :] - 1)
+
+        # [*, N_res, N_res, 1, C_z]
+        z = z.unsqueeze(-2)
+
+        # [*, N_res, N_res, N_temp, C_t]
+        t = permute_final_dims(t, (1, 2, 0, 3))
+
+        # [*, N_res, N_res, 1, C_z]
+        biases = [bias]
+        if chunk_size is not None and not self.training:
+            z = self._chunk(z, t, biases, chunk_size, use_lma=use_lma)
+        else:
+            z = self.mha(q_x=z, kv_x=t, biases=biases, use_lma=use_lma)
+
+        # [*, N_res, N_res, C_z]
+        z = z.squeeze(-2)
+
+        return z
+
+
+class TemplatePairStackBlock(nn.Module):
+    def __init__(
+        self,
+        c_t: int,
+        c_hidden_tri_att: int,
+        c_hidden_tri_mul: int,
+        no_heads: int,
+        pair_transition_n: int,
+        dropout_rate: float,
+        tri_mul_first: bool,
+        fuse_projection_weights: bool,
+        inf: float,
+        **kwargs,
+    ):
+        super(TemplatePairStackBlock, self).__init__()
+
+        self.c_t = c_t
+        self.c_hidden_tri_att = c_hidden_tri_att
+        self.c_hidden_tri_mul = c_hidden_tri_mul
+        self.no_heads = no_heads
+        self.pair_transition_n = pair_transition_n
+        self.dropout_rate = dropout_rate
+        self.inf = inf
+        self.tri_mul_first = tri_mul_first
+
+        self.dropout_row = DropoutRowwise(self.dropout_rate)
+        self.dropout_col = DropoutColumnwise(self.dropout_rate)
+
+        self.tri_att_start = TriangleAttentionStartingNode(
+            self.c_t,
+            self.c_hidden_tri_att,
+            self.no_heads,
+            inf=inf,
+        )
+        self.tri_att_end = TriangleAttentionEndingNode(
+            self.c_t,
+            self.c_hidden_tri_att,
+            self.no_heads,
+            inf=inf,
+        )
+
+        if fuse_projection_weights:
+            self.tri_mul_out = FusedTriangleMultiplicationOutgoing(
+                self.c_t,
+                self.c_hidden_tri_mul,
+            )
+            self.tri_mul_in = FusedTriangleMultiplicationIncoming(
+                self.c_t,
+                self.c_hidden_tri_mul,
+            )
+        else:
+            self.tri_mul_out = TriangleMultiplicationOutgoing(
+                self.c_t,
+                self.c_hidden_tri_mul,
+            )
+            self.tri_mul_in = TriangleMultiplicationIncoming(
+                self.c_t,
+                self.c_hidden_tri_mul,
+            )
+
+        self.pair_transition = PairTransition(
+            self.c_t,
+            self.pair_transition_n,
+        )
+
+    def tri_att_start_end(self,
+                          single: torch.Tensor,
+                          _attn_chunk_size: Optional[int],
+                          single_mask: torch.Tensor,
+                          use_deepspeed_evo_attention: bool,
+                          use_lma: bool,
+                          inplace_safe: bool):
+        single = add(single,
+                     self.dropout_row(
+                         self.tri_att_start(
+                             single,
+                             chunk_size=_attn_chunk_size,
+                             mask=single_mask,
+                             use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                             use_lma=use_lma,
+                             inplace_safe=inplace_safe,
+                         )
+                     ),
+                     inplace_safe,
+                     )
+
+        single = add(single,
+                     self.dropout_col(
+                         self.tri_att_end(
+                             single,
+                             chunk_size=_attn_chunk_size,
+                             mask=single_mask,
+                             use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                             use_lma=use_lma,
+                             inplace_safe=inplace_safe,
+                         )
+                     ),
+                     inplace_safe,
+                     )
+
+        return single
+
+    def tri_mul_out_in(self,
+                       single: torch.Tensor,
+                       single_mask: torch.Tensor,
+                       inplace_safe: bool):
+        tmu_update = self.tri_mul_out(
+            single,
+            mask=single_mask,
+            inplace_safe=inplace_safe,
+            _add_with_inplace=True,
+        )
+        if not inplace_safe:
+            single = single + self.dropout_row(tmu_update)
+        else:
+            single = tmu_update
+
+        del tmu_update
+
+        tmu_update = self.tri_mul_in(
+            single,
+            mask=single_mask,
+            inplace_safe=inplace_safe,
+            _add_with_inplace=True,
+        )
+        if not inplace_safe:
+            single = single + self.dropout_row(tmu_update)
+        else:
+            single = tmu_update
+
+        del tmu_update
+
+        return single
+
+    def forward(self,
+                z: torch.Tensor,
+                mask: torch.Tensor,
+                chunk_size: Optional[int] = None,
+                use_deepspeed_evo_attention: bool = False,
+                use_lma: bool = False,
+                inplace_safe: bool = False,
+                _mask_trans: bool = True,
+                _attn_chunk_size: Optional[int] = None,
+                ):
+        if _attn_chunk_size is None:
+            _attn_chunk_size = chunk_size
+
+        single_templates = [
+            t.unsqueeze(-4) for t in torch.unbind(z, dim=-4)
+        ]
+        single_templates_masks = [
+            m.unsqueeze(-3) for m in torch.unbind(mask, dim=-3)
+        ]
+
+        for i in range(len(single_templates)):
+            single = single_templates[i]
+            single_mask = single_templates_masks[i]
+
+            if self.tri_mul_first:
+                single = self.tri_att_start_end(single=self.tri_mul_out_in(single=single,
+                                                                           single_mask=single_mask,
+                                                                           inplace_safe=inplace_safe),
+                                                _attn_chunk_size=_attn_chunk_size,
+                                                single_mask=single_mask,
+                                                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                                                use_lma=use_lma,
+                                                inplace_safe=inplace_safe)
+            else:
+                single = self.tri_mul_out_in(
+                    single=self.tri_att_start_end(single=single,
+                                                  _attn_chunk_size=_attn_chunk_size,
+                                                  single_mask=single_mask,
+                                                  use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                                                  use_lma=use_lma,
+                                                  inplace_safe=inplace_safe),
+                    single_mask=single_mask,
+                    inplace_safe=inplace_safe)
+
+            single = add(single,
+                         self.pair_transition(
+                             single,
+                             mask=single_mask if _mask_trans else None,
+                             chunk_size=chunk_size,
+                         ),
+                         inplace_safe,
+                         )
+
+            if not inplace_safe:
+                single_templates[i] = single
+
+        if not inplace_safe:
+            z = torch.cat(single_templates, dim=-4)
+
+        return z
+
+
+class TemplatePairStack(nn.Module):
+    """
+    Implements Algorithm 16.
+    """
+
+    def __init__(
+        self,
+        c_t,
+        c_hidden_tri_att,
+        c_hidden_tri_mul,
+        no_blocks,
+        no_heads,
+        pair_transition_n,
+        dropout_rate,
+        tri_mul_first,
+        fuse_projection_weights,
+        blocks_per_ckpt,
+        tune_chunk_size: bool = False,
+        inf=1e9,
+        **kwargs,
+    ):
+        """
+        Args:
+            c_t:
+                Template embedding channel dimension
+            c_hidden_tri_att:
+                Per-head hidden dimension for triangular attention
+            c_hidden_tri_att:
+                Hidden dimension for triangular multiplication
+            no_blocks:
+                Number of blocks in the stack
+            pair_transition_n:
+                Scale of pair transition (Alg. 15) hidden dimension
+            dropout_rate:
+                Dropout rate used throughout the stack
+            blocks_per_ckpt:
+                Number of blocks per activation checkpoint. None disables
+                activation checkpointing
+        """
+        super(TemplatePairStack, self).__init__()
+
+        self.blocks_per_ckpt = blocks_per_ckpt
+
+        self.blocks = nn.ModuleList()
+        for _ in range(no_blocks):
+            block = TemplatePairStackBlock(
+                c_t=c_t,
+                c_hidden_tri_att=c_hidden_tri_att,
+                c_hidden_tri_mul=c_hidden_tri_mul,
+                no_heads=no_heads,
+                pair_transition_n=pair_transition_n,
+                dropout_rate=dropout_rate,
+                tri_mul_first=tri_mul_first,
+                fuse_projection_weights=fuse_projection_weights,
+                inf=inf,
+            )
+            self.blocks.append(block)
+
+        self.layer_norm = LayerNorm(c_t)
+
+        self.tune_chunk_size = tune_chunk_size
+        self.chunk_size_tuner = None
+        if tune_chunk_size:
+            self.chunk_size_tuner = ChunkSizeTuner()
+
+    def forward(
+        self,
+        t: torch.tensor,
+        mask: torch.tensor,
+        chunk_size: int,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        inplace_safe: bool = False,
+        _mask_trans: bool = True,
+    ):
+        """
+        Args:
+            t:
+                [*, N_templ, N_res, N_res, C_t] template embedding
+            mask:
+                [*, N_templ, N_res, N_res] mask
+        Returns:
+            [*, N_templ, N_res, N_res, C_t] template embedding update
+        """
+        if mask.shape[-3] == 1:
+            expand_idx = list(mask.shape)
+            expand_idx[-3] = t.shape[-4]
+            mask = mask.expand(*expand_idx)
+
+        blocks = [
+            partial(
+                b,
+                mask=mask,
+                chunk_size=chunk_size,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma,
+                inplace_safe=inplace_safe,
+                _mask_trans=_mask_trans,
+            )
+            for b in self.blocks
+        ]
+
+        if chunk_size is not None and self.chunk_size_tuner is not None:
+            assert (not self.training)
+            tuned_chunk_size = self.chunk_size_tuner.tune_chunk_size(
+                representative_fn=blocks[0],
+                args=(t.clone(),),
+                min_chunk_size=chunk_size,
+            )
+            blocks = [
+                partial(b,
+                        chunk_size=tuned_chunk_size,
+                        _attn_chunk_size=max(chunk_size, tuned_chunk_size // 4),
+                        ) for b in blocks
+            ]
+
+        t, = checkpoint_blocks(
+            blocks=blocks,
+            args=(t,),
+            blocks_per_ckpt=self.blocks_per_ckpt if self.training else None,
+        )
+
+        t = self.layer_norm(t)
+
+        return t
+
+
+def embed_templates_offload(
+    model,
+    batch,
+    z,
+    pair_mask,
+    templ_dim,
+    template_chunk_size=256,
+    inplace_safe=False,
+):
+    """
+    Args:
+        model: 
+            An AlphaFold model object
+        batch: 
+            An AlphaFold input batch. See documentation of AlphaFold.
+        z: 
+            A [*, N, N, C_z] pair embedding
+        pair_mask: 
+            A [*, N, N] pair mask
+        templ_dim: 
+            The template dimension of the template tensors in batch
+        template_chunk_size: 
+            Integer value controlling how quickly the offloaded pair embedding
+            tensor is brought back into GPU memory. In dire straits, can be
+            lowered to reduce memory consumption of this function even more.
+    Returns:
+        A dictionary of template pair and angle embeddings.
+    
+    A version of the "embed_templates" method of the AlphaFold class that
+    offloads the large template pair tensor to CPU. Slower but more frugal 
+    with GPU memory than the original. Useful for long-sequence inference.
+    """
+    # Embed the templates one at a time (with a poor man's vmap)
+    pair_embeds_cpu = []
+    n = z.shape[-2]
+    n_templ = batch["template_aatype"].shape[templ_dim]
+    for i in range(n_templ):
+        idx = batch["template_aatype"].new_tensor(i)
+        single_template_feats = tensor_tree_map(
+            lambda t: torch.index_select(t, templ_dim, idx).squeeze(templ_dim),
+            batch,
+        )
+
+        # [*, N, N, C_t]
+        t = build_template_pair_feat(
+            single_template_feats,
+            use_unit_vector=model.config.template.use_unit_vector,
+            inf=model.config.template.inf,
+            eps=model.config.template.eps,
+            **model.config.template.distogram,
+        ).to(z.dtype)
+        t = model.template_pair_embedder(t)
+
+        # [*, 1, N, N, C_z]
+        t = model.template_pair_stack(
+            t.unsqueeze(templ_dim),
+            pair_mask.unsqueeze(-3).to(dtype=z.dtype),
+            chunk_size=model.globals.chunk_size,
+            use_deepspeed_evo_attention=model.globals.use_deepspeed_evo_attention,
+            use_lma=model.globals.use_lma,
+            inplace_safe=inplace_safe,
+            _mask_trans=model.config._mask_trans,
+        )
+
+        assert (sys.getrefcount(t) == 2)
+
+        pair_embeds_cpu.append(t.cpu())
+
+        del t
+
+    # Preallocate the output tensor
+    t = z.new_zeros(z.shape)
+
+    for i in range(0, n, template_chunk_size):
+        pair_chunks = [
+            p[..., i: i + template_chunk_size, :, :] for p in pair_embeds_cpu
+        ]
+        pair_chunk = torch.cat(pair_chunks, dim=templ_dim).to(device=z.device)
+        z_chunk = z[..., i: i + template_chunk_size, :, :]
+        att_chunk = model.template_pointwise_att(
+            pair_chunk,
+            z_chunk,
+            template_mask=batch["template_mask"].to(dtype=z.dtype),
+            use_lma=model.globals.use_lma,
+        )
+
+        t[..., i: i + template_chunk_size, :, :] = att_chunk
+
+    del pair_chunks
+
+    if inplace_safe:
+        t = t * (torch.sum(batch["template_mask"], dim=-1) > 0)
+    else:
+        t *= (torch.sum(batch["template_mask"], dim=-1) > 0)
+
+    ret = {}
+    if model.config.template.embed_angles:
+        template_angle_feat = build_template_angle_feat(
+            batch,
+        )
+
+        # [*, N, C_m]
+        a = model.template_single_embedder(template_angle_feat)
+
+        ret["template_single_embedding"] = a
+
+    ret.update({"template_pair_embedding": t})
+
+    return ret
+
+
+def embed_templates_average(
+    model,
+    batch,
+    z,
+    pair_mask,
+    templ_dim,
+    templ_group_size=2,
+    inplace_safe=False,
+):
+    """
+    Args:
+        model: 
+            An AlphaFold model object
+        batch: 
+            An AlphaFold input batch. See documentation of AlphaFold.
+        z: 
+            A [*, N, N, C_z] pair embedding
+        pair_mask: 
+            A [*, N, N] pair mask
+        templ_dim: 
+            The template dimension of the template tensors in batch
+        templ_group_size: 
+            Granularity of the approximation. Larger values trade memory for 
+            greater proximity to the original function
+    Returns:
+        A dictionary of template pair and angle embeddings.
+
+    A memory-efficient approximation of the "embed_templates" method of the 
+    AlphaFold class. Instead of running pointwise attention over pair 
+    embeddings for all of the templates at the same time, it splits templates 
+    into groups of size templ_group_size, computes embeddings for each group 
+    normally, and then averages the group embeddings. In our experiments, this 
+    approximation has a minimal effect on the quality of the resulting 
+    embedding, while its low memory footprint allows the number of templates 
+    to scale almost indefinitely.
+    """
+    # Embed the templates one at a time (with a poor man's vmap)
+    n = z.shape[-2]
+    n_templ = batch["template_aatype"].shape[templ_dim]
+    out_tensor = z.new_zeros(z.shape)
+    for i in range(0, n_templ, templ_group_size):
+        def slice_template_tensor(t):
+            s = [slice(None) for _ in t.shape]
+            s[templ_dim] = slice(i, i + templ_group_size)
+            return t[s]
+
+        template_feats = tensor_tree_map(
+            slice_template_tensor,
+            batch,
+        )
+
+        # [*, N, N, C_t]
+        t = build_template_pair_feat(
+            template_feats,
+            use_unit_vector=model.config.template.use_unit_vector,
+            inf=model.config.template.inf,
+            eps=model.config.template.eps,
+            **model.config.template.distogram,
+        ).to(z.dtype)
+
+        # [*, S_t, N, N, C_z]
+        t = model.template_pair_embedder(t)
+        t = model.template_pair_stack(
+            t,
+            pair_mask.unsqueeze(-3).to(dtype=z.dtype),
+            chunk_size=model.globals.chunk_size,
+            use_deepspeed_evo_attention=model.globals.use_deepspeed_evo_attention,
+            use_lma=model.globals.use_lma,
+            inplace_safe=inplace_safe,
+            _mask_trans=model.config._mask_trans,
+        )
+
+        t = model.template_pointwise_att(
+            t,
+            z,
+            template_mask=template_feats["template_mask"].to(dtype=z.dtype),
+            use_lma=model.globals.use_lma,
+        )
+
+        denom = math.ceil(n_templ / templ_group_size)
+        if inplace_safe:
+            t /= denom
+        else:
+            t = t / denom
+
+        if inplace_safe:
+            out_tensor += t
+        else:
+            out_tensor = out_tensor + t
+
+        del t
+
+    if inplace_safe:
+        out_tensor *= (torch.sum(batch["template_mask"], dim=-1) > 0)
+    else:
+        out_tensor = out_tensor * (torch.sum(batch["template_mask"], dim=-1) > 0)
+
+    ret = {}
+    if model.config.template.embed_angles:
+        template_angle_feat = build_template_angle_feat(
+            batch,
+        )
+
+        # [*, N, C_m]
+        a = model.template_single_embedder(template_angle_feat)
+
+        ret["template_single_embedding"] = a
+
+    ret.update({"template_pair_embedding": out_tensor})
+
+    return ret

model/openfold/torchscript.py

@@ -0,0 +1,215 @@
+# Copyright 2021 AlQuraishi Laboratory
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+
+from model.openfold.dropout import (
+    DropoutRowwise,
+    DropoutColumnwise,
+)
+from model.openfold.evoformer import (
+    EvoformerBlock,
+    EvoformerStack,
+)
+from model.openfold.outer_product_mean import OuterProductMean
+from model.openfold.msa import (
+    MSARowAttentionWithPairBias, 
+    MSAColumnAttention,
+    MSAColumnGlobalAttention,
+)
+from model.openfold.pair_transition import PairTransition
+from model.openfold.primitives import Attention, GlobalAttention
+from model.openfold.structure_module import (
+    InvariantPointAttention,
+    BackboneUpdate,
+)
+from model.openfold.template import TemplatePairStackBlock
+from model.openfold.triangular_attention import (
+    TriangleAttentionStartingNode,
+    TriangleAttentionEndingNode,
+)
+from model.openfold.triangular_multiplicative_update import (
+    TriangleMultiplicationOutgoing,
+    TriangleMultiplicationIncoming,
+)
+
+
+def script_preset_(model: torch.nn.Module):
+    """
+    TorchScript a handful of low-level but frequently used submodule types 
+    that are known to be scriptable.
+
+    Args:
+        model: 
+            A torch.nn.Module. It should contain at least some modules from 
+            this repository, or this function won't do anything.
+    """
+    script_submodules_(
+        model, 
+        [
+            nn.Dropout,
+            Attention,
+            GlobalAttention,
+            EvoformerBlock,
+            #TemplatePairStackBlock,
+        ], 
+        attempt_trace=False,
+        batch_dims=None,
+    ) 
+
+    
+def _get_module_device(module: torch.nn.Module) -> torch.device:
+    """
+    Fetches the device of a module, assuming that all of the module's
+    parameters reside on a single device
+
+    Args:
+        module: A torch.nn.Module
+    Returns:
+        The module's device
+    """
+    return next(module.parameters()).device
+
+
+def _trace_module(module, batch_dims=None):
+    if(batch_dims is None):
+        batch_dims = ()
+
+    # Stand-in values
+    n_seq = 10
+    n_res = 10
+
+    device = _get_module_device(module)
+
+    def msa(channel_dim):
+        return torch.rand(
+            (*batch_dims, n_seq, n_res, channel_dim),
+            device=device,
+        )
+
+    def pair(channel_dim):
+        return torch.rand(
+            (*batch_dims, n_res, n_res, channel_dim),
+            device=device,
+        )
+
+    if(isinstance(module, MSARowAttentionWithPairBias)):
+        inputs = {
+            "forward": (
+                msa(module.c_in), # m
+                pair(module.c_z), # z
+                torch.randint(
+                    0, 2, 
+                    (*batch_dims, n_seq, n_res)
+                ), # mask
+            ),
+        }
+    elif(isinstance(module, MSAColumnAttention)):
+        inputs = {
+            "forward": (
+                msa(module.c_in), # m
+                torch.randint(
+                    0, 2, 
+                    (*batch_dims, n_seq, n_res)
+                ), # mask
+            ),
+        }
+    elif(isinstance(module, OuterProductMean)):
+        inputs = {
+            "forward": (
+                msa(module.c_m),
+                torch.randint(
+                    0, 2,
+                    (*batch_dims, n_seq, n_res)
+                )
+            )
+        }
+    else:
+        raise TypeError(
+            f"tracing is not supported for modules of type {type(module)}"
+        )
+
+    return torch.jit.trace_module(module, inputs)
+
+
+def _script_submodules_helper_(
+    model,
+    types,
+    attempt_trace,
+    to_trace,
+):
+    for name, child in model.named_children():
+        if(types is None or any(isinstance(child, t) for t in types)):
+            try:
+                scripted = torch.jit.script(child)
+                setattr(model, name, scripted)
+                continue
+            except (RuntimeError, torch.jit.frontend.NotSupportedError) as e:
+                if(attempt_trace):
+                    to_trace.add(type(child))
+                else:
+                    raise e
+        
+        _script_submodules_helper_(child, types, attempt_trace, to_trace)
+
+
+def _trace_submodules_(
+    model,
+    types,
+    batch_dims=None,
+):
+    for name, child in model.named_children():
+        if(any(isinstance(child, t) for t in types)):
+            traced = _trace_module(child, batch_dims=batch_dims)
+            setattr(model, name, traced)
+        else:
+            _trace_submodules_(child, types, batch_dims=batch_dims)
+
+
+def script_submodules_(
+    model: nn.Module,
+    types: Optional[Sequence[type]] = None,
+    attempt_trace: Optional[bool] = True,
+    batch_dims: Optional[Tuple[int]] = None,
+):
+    """
+    Convert all submodules whose types match one of those in the input 
+    list to recursively scripted equivalents in place. To script the entire
+    model, just call torch.jit.script on it directly.
+
+    When types is None, all submodules are scripted.
+
+    Args:
+        model: 
+            A torch.nn.Module
+        types: 
+            A list of types of submodules to script
+        attempt_trace: 
+            Whether to attempt to trace specified modules if scripting 
+            fails. Recall that tracing eliminates all conditional 
+            logic---with great tracing comes the mild responsibility of 
+            having to remember to ensure that the modules in question 
+            perform the same computations no matter what.
+    """
+    to_trace = set()
+
+    # Aggressively script as much as possible first...
+    _script_submodules_helper_(model, types, attempt_trace, to_trace)
+  
+    # ... and then trace stragglers.
+    if(attempt_trace and len(to_trace) > 0):
+        _trace_submodules_(model, to_trace, batch_dims=batch_dims)

model/openfold/triangular_attention.py

@@ -0,0 +1,168 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import partialmethod, partial
+import math
+from typing import Optional, List
+
+import torch
+import torch.nn as nn
+
+from model.openfold.primitives import Linear, LayerNorm, Attention
+from onescience.utils.openfold.chunk_utils import chunk_layer
+from onescience.utils.openfold.tensor_utils import (
+    permute_final_dims,
+    flatten_final_dims,
+)
+
+
+class TriangleAttention(nn.Module):
+    def __init__(
+        self, c_in, c_hidden, no_heads, starting=True, inf=1e9, bias: bool=True
+    ):
+        """
+        Args:
+            c_in:
+                Input channel dimension
+            c_hidden:
+                Overall hidden channel dimension (not per-head)
+            no_heads:
+                Number of attention heads
+        """
+        super(TriangleAttention, self).__init__()
+
+        self.c_in = c_in
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.starting = starting
+        self.inf = inf
+
+        self.layer_norm = LayerNorm(self.c_in)
+
+        self.linear = Linear(c_in, self.no_heads, bias=False, init="normal")
+        if bias==False:
+            self.mha = Attention(
+                self.c_in, self.c_in, self.c_in, self.c_hidden, self.no_heads, bias=False
+            )
+        else:
+            self.mha = Attention(
+                self.c_in, self.c_in, self.c_in, self.c_hidden, self.no_heads
+            )
+
+    @torch.jit.ignore
+    def _chunk(self,
+        x: torch.Tensor,
+        biases: List[torch.Tensor],
+        chunk_size: int,
+        use_memory_efficient_kernel: bool = False,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        inplace_safe: bool = False,
+    ) -> torch.Tensor:
+        "triangle! triangle!"
+        mha_inputs = {
+            "q_x": x,
+            "kv_x": x,
+            "biases": biases,
+        }
+
+        return chunk_layer(
+            partial(
+                self.mha, 
+                use_memory_efficient_kernel=use_memory_efficient_kernel,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma
+            ),
+            mha_inputs,
+            chunk_size=chunk_size,
+            no_batch_dims=len(x.shape[:-2]),
+            _out=x if inplace_safe else None,
+        )
+
+    def forward(self, 
+        x: torch.Tensor, 
+        mask: Optional[torch.Tensor] = None,
+        chunk_size: Optional[int] = None,
+        use_memory_efficient_kernel: bool = False,
+        use_deepspeed_evo_attention: bool = False,
+        use_lma: bool = False,
+        inplace_safe: bool = False,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x:
+                [*, I, J, C_in] input tensor (e.g. the pair representation)
+        Returns:
+            [*, I, J, C_in] output tensor
+        """ 
+        if mask is None:
+            # [*, I, J]
+            mask = x.new_ones(
+                x.shape[:-1],
+            )
+
+        if(not self.starting):
+            x = x.transpose(-2, -3)
+            mask = mask.transpose(-1, -2)
+
+        # [*, I, J, C_in]
+        x = self.layer_norm(x)
+
+        # [*, I, 1, 1, J]
+        mask_bias = (self.inf * (mask - 1))[..., :, None, None, :]
+
+        # [*, H, I, J]
+        triangle_bias = permute_final_dims(self.linear(x), (2, 0, 1))
+
+        # [*, 1, H, I, J]
+        triangle_bias = triangle_bias.unsqueeze(-4)
+
+        biases = [mask_bias, triangle_bias]
+
+        if chunk_size is not None:
+            x = self._chunk(
+                x, 
+                biases, 
+                chunk_size, 
+                use_memory_efficient_kernel=use_memory_efficient_kernel,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma,
+                inplace_safe=inplace_safe,
+            )
+        else:
+            x = self.mha(
+                q_x=x, 
+                kv_x=x, 
+                biases=biases, 
+                use_memory_efficient_kernel=use_memory_efficient_kernel,
+                use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+                use_lma=use_lma
+            )
+
+        if(not self.starting):
+            x = x.transpose(-2, -3)
+
+        return x
+
+
+# Implements Algorithm 13
+TriangleAttentionStartingNode = TriangleAttention
+
+
+class TriangleAttentionEndingNode(TriangleAttention):
+    """
+    Implements Algorithm 14.
+    """
+    __init__ = partialmethod(TriangleAttention.__init__, starting=False)

model/openfold/triangular_multiplicative_update.py

@@ -0,0 +1,647 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import partialmethod
+from typing import Optional
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+
+from model.openfold.primitives import Linear, LayerNorm
+from onescience.utils.openfold.chunk_utils import chunk_layer
+from onescience.utils.openfold.precision_utils import is_fp16_enabled
+from onescience.utils.openfold.tensor_utils import add, permute_final_dims
+
+
+class BaseTriangleMultiplicativeUpdate(nn.Module, ABC):
+    """
+    Implements Algorithms 11 and 12.
+    """
+    @abstractmethod
+    def __init__(self, c_z, c_hidden, _outgoing, bias=True):
+        """
+        Args:
+            c_z:
+                Input channel dimension
+            c:
+                Hidden channel dimension
+        """
+        super(BaseTriangleMultiplicativeUpdate, self).__init__()
+        self.c_z = c_z
+        self.c_hidden = c_hidden
+        self._outgoing = _outgoing
+        self.bias = bias
+
+        self.linear_g = Linear(self.c_z, self.c_z, bias=bias, init="gating")
+        self.linear_z = Linear(self.c_hidden, self.c_z, bias=bias, init="final")
+
+        self.layer_norm_in = LayerNorm(self.c_z)
+        self.layer_norm_out = LayerNorm(self.c_hidden)
+
+        self.sigmoid = nn.Sigmoid()
+
+    def _combine_projections(self,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        _inplace_chunk_size: Optional[int] = None
+    ) -> torch.Tensor:
+        if(self._outgoing):
+            a = permute_final_dims(a, (2, 0, 1))
+            b = permute_final_dims(b, (2, 1, 0))
+        else:
+            a = permute_final_dims(a, (2, 1, 0))
+            b = permute_final_dims(b,  (2, 0, 1))
+
+        if(_inplace_chunk_size is not None):
+            # To be replaced by torch vmap
+            for i in range(0, a.shape[-3], _inplace_chunk_size):
+                a_chunk = a[..., i: i + _inplace_chunk_size, :, :]
+                b_chunk = b[..., i: i + _inplace_chunk_size, :, :]
+                a[..., i: i + _inplace_chunk_size, :, :] = (
+                    torch.matmul(
+                        a_chunk,
+                        b_chunk,
+                    )
+                )
+
+            p = a
+        else:
+            p = torch.matmul(a, b)
+
+        return permute_final_dims(p, (1, 2, 0))
+
+    @abstractmethod
+    def forward(self,
+        z: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        inplace_safe: bool = False,
+        _add_with_inplace: bool = False
+    ) -> torch.Tensor:
+        """
+        Args:
+            x:
+                [*, N_res, N_res, C_z] input tensor
+            mask:
+                [*, N_res, N_res] input mask
+        Returns:
+            [*, N_res, N_res, C_z] output tensor
+        """
+        pass
+
+
+class TriangleMultiplicativeUpdate(BaseTriangleMultiplicativeUpdate):
+    """
+    Implements Algorithms 11 and 12.
+    """
+    def __init__(self, c_z, c_hidden, _outgoing=True, bias: bool=True):
+        """
+        Args:
+            c_z:
+                Input channel dimension
+            c:
+                Hidden channel dimension
+        """
+        super(TriangleMultiplicativeUpdate, self).__init__(c_z=c_z,
+                                                           c_hidden=c_hidden,
+                                                           _outgoing=_outgoing,
+                                                           bias=bias
+                                                          )
+
+        self.linear_a_p = Linear(self.c_z, self.c_hidden, bias=bias)
+        self.linear_a_g = Linear(self.c_z, self.c_hidden, bias=bias, init="gating")
+        self.linear_b_p = Linear(self.c_z, self.c_hidden, bias=bias,)
+        self.linear_b_g = Linear(self.c_z, self.c_hidden, bias=bias, init="gating")
+
+    def _inference_forward(self,
+        z: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        inplace_chunk_size: Optional[int] = None,
+        with_add: bool = True,
+    ):
+        """
+        Args:
+            z:
+                A [*, N, N, C_z] pair representation
+            mask:
+                A [*, N, N] pair mask
+            inplace_chunk_size:
+                Size of chunks used in the main computation. Increase to trade
+                memory for speed.
+            with_add:
+                If True, z is overwritten with (z + update). Otherwise, it is
+                overwritten with (update).
+        Returns:
+            A reference to the overwritten z
+
+        More memory-efficient, inference-only version of the forward function.
+        Uses in-place operations, fusion of the addition that happens after
+        this module in the Evoformer, a smidge of recomputation, and 
+        a cache of overwritten values to lower peak memory consumption of this
+        module from 5x the size of the input tensor z to 2.5x its size. Useful
+        for inference on extremely long sequences. 
+        
+        It works as follows. We will make reference to variables used in the
+        default forward implementation below. Naively, triangle multiplication
+        attention requires the manifestation of 5 tensors the size of z:
+        1) z, the "square" input tensor, 2) a, the first projection of z, 
+        3) b, the second projection of b, 4) g, a z-sized mask, and 5) a 
+        z-sized tensor for intermediate computations. For large N, this is 
+        prohibitively expensive; for N=4000, for example, z is more than 8GB 
+        alone. To avoid this problem, we compute b, g, and all intermediate
+        tensors in small chunks, noting that the chunks required to compute a
+        chunk of the output depend only on the tensor a and corresponding 
+        vertical and horizontal chunks of z. This suggests an algorithm that 
+        loops over pairs of chunks of z: hereafter "columns" and "rows" of
+        z, even though each "column" and "row" in fact contains
+        inplace_chunk_size contiguous true columns and rows of z. Writing 
+        output chunks to a new tensor would bring total memory consumption
+        down to 3x the size of z. However, more memory can be saved by writing
+        output chunks directly to z in-place. WLOG, we choose to write output
+        chunks vertically, overwriting the ith "column" of z at the end of
+        the ith iteration of the main loop. Despite this overwriting, the 
+        ith column is always one column ahead of previously overwritten columns 
+        and can be recovered directly from z. After the first iteration,
+        however, the ith row of z is always at least partially overwritten. For
+        this reason, we introduce the z-cache, a tensor one-half the size of 
+        z. The z-cache initially contains the left half (2nd and 3rd quadrants)
+        of z. For 0 < i < N/2, the missing left part of the ith row of z is
+        recovered from this cache at the beginning of the ith iteration. Once i 
+        exceeds n/2, the cache is "reoriented" to encompass the 3rd and 4th 
+        quadrants of z instead. Though the 3rd quadrant of the original z is 
+        entirely overwritten at this point, it can be recovered from the z-cache 
+        itself. Thereafter, the ith row of z can be recovered in its entirety 
+        from the reoriented z-cache. After the final iteration, z has been 
+        completely overwritten and contains the triangular multiplicative 
+        update. If with_add is True, it instead contains the sum of z and the
+        triangular multiplicative update. In either case, peak memory 
+        consumption is just 2.5x the size of z, disregarding memory used for 
+        chunks and other small variables.
+        """
+        if mask is None:
+            mask = z.new_ones(z.shape[:-1])
+
+        mask = mask.unsqueeze(-1)
+       
+        def compute_projection_helper(pair, mask, a=True):
+            if(a):
+                linear_g = self.linear_a_g
+                linear_p = self.linear_a_p
+            else:
+                linear_g = self.linear_b_g
+                linear_p = self.linear_b_p
+            
+            pair = self.layer_norm_in(pair)
+            p = linear_g(pair)
+            p.sigmoid_()
+            p *= linear_p(pair)
+            p *= mask
+            p = permute_final_dims(p, (2, 0, 1))
+            return p
+
+        def compute_projection(pair, mask, a=True, chunked=True): 
+            need_transpose = self._outgoing ^ a
+            if(not chunked):
+                p = compute_projection_helper(pair, mask, a)
+                if(need_transpose):
+                    p = p.transpose(-1, -2)
+            else:
+                # This computation is chunked so as not to exceed our 2.5x 
+                # budget with a large intermediate tensor
+                linear_g = self.linear_a_g if a else self.linear_b_g
+                #c = linear_g.bias.shape[-1]
+                if self.bias:
+                    c = linear_g.bias.shape[-1]
+                else:
+                    c = linear_g.weight.shape[0]
+                out_shape = pair.shape[:-3] + (c,) + pair.shape[-3:-1]
+                p = pair.new_zeros(out_shape)
+                for i in range(0, pair.shape[-3], inplace_chunk_size):
+                    pair_chunk = pair[..., i: i + inplace_chunk_size, :, :]
+                    mask_chunk = mask[..., i: i + inplace_chunk_size, :, :]
+                    pair_chunk = compute_projection_helper(
+                        pair[..., i: i + inplace_chunk_size, :, :],
+                        mask[..., i: i + inplace_chunk_size, :, :], 
+                        a,
+                    )
+                    if(need_transpose):
+                        pair_chunk = pair_chunk.transpose(-1, -2)
+                        p[..., i: i + inplace_chunk_size] = pair_chunk
+                    else:
+                        p[..., i: i + inplace_chunk_size, :] = pair_chunk
+                    
+                    del pair_chunk
+
+            return p
+
+        # We start by fully manifesting a. In addition to the input, this
+        # brings total memory consumption to 2x z (disregarding size of chunks)
+        # [*, N, N, c]
+        a = compute_projection(z, mask, True, chunked=True)
+        #if bias==True:
+        #    a = compute_projection(z, mask, True, True, chunked=True)
+        #else:
+        #    a = compute_projection(z, mask, False, True, chunked=True)
+
+        if(inplace_chunk_size is not None):
+            n = a.shape[-1]
+            half_n = n // 2 + n % 2
+            row_dim = -3
+            col_dim = -2
+            b_chunk_dim = row_dim if self._outgoing else col_dim
+            
+            def empty_slicer(t):
+                return [slice(None) for _ in t.shape]
+            
+            def slice_tensor(t, start, end, dim):
+                # Slices start:end from the dim dimension of t
+                s = empty_slicer(t)
+                s[dim] = slice(start, end)
+                return t[s]
+
+            def flip_z_cache_(z_cache, z):
+                # "Reorient" the z_cache (see below), filling it with quadrants
+                # 3---recovered from the z_cache---and 4---recovered from z---
+                # of the input tensor z. 
+                quadrant_3 = slice_tensor(
+                    z_cache, half_n, None, row_dim
+                )
+                z_cache = z_cache.transpose(row_dim, col_dim)
+                
+                # If n is odd, we need to shrink the z_cache by one row
+                z_cache = z_cache[..., :(n // 2), :, :]
+
+                # Move the 3rd quadrant of z into the 
+                first_half_slicer = empty_slicer(z_cache)
+                first_half_slicer[col_dim] = slice(0, half_n)
+                z_cache[first_half_slicer] = quadrant_3
+               
+                # Get the fourth quadrant of z
+                quadrant_4 = slice_tensor(z, half_n, None, row_dim)
+                quadrant_4 = slice_tensor(
+                    quadrant_4, half_n, None, col_dim
+                )
+
+                # Insert said quadrant into the rotated z-cache
+                quadrant_3_slicer = empty_slicer(z_cache)
+                quadrant_3_slicer[col_dim] = slice(half_n, None)
+
+                z_cache[quadrant_3_slicer] = quadrant_4
+
+                return z_cache
+
+            # Initialize the z cache to the left half of z.
+            z_cache_shape = list(z.shape)
+            z_cache_shape[col_dim] = half_n
+            z_cache = z.new_zeros(z_cache_shape)
+            z_cache_slicer = empty_slicer(z_cache)
+            z_cache_slicer[col_dim] = slice(0, half_n)
+            z_cache.copy_(z[z_cache_slicer])
+            z_cache_rotated = False
+
+            # We need to reorient the z-cache at the halfway point, and we 
+            # don't want a single chunk to straddle that point. We contract one
+            # of the chunks in the middle to address that problem.
+            i_range = list(range(0, half_n, inplace_chunk_size))
+            initial_offsets = [
+                i_2 - i_1 for i_1, i_2 in zip(i_range, i_range[1:] + [half_n])
+            ]
+            after_half = list(range(half_n, n, inplace_chunk_size))
+            after_half_offsets = [inplace_chunk_size for _ in after_half]
+            combined_range_with_offsets = zip(
+                i_range + after_half, initial_offsets + after_half_offsets
+            )
+            for i, offset in combined_range_with_offsets:
+                if(not z_cache_rotated and i >= half_n):
+                    z_cache = flip_z_cache_(z_cache, z)
+                    z_cache_rotated = True
+
+                z_chunk_b = slice_tensor(
+                    z, i, i + offset, b_chunk_dim,
+                )
+                mask_chunk = slice_tensor(
+                    mask, i, i + offset, b_chunk_dim,
+                )
+
+                z_chunk_b = z_chunk_b.clone()
+                if(b_chunk_dim == col_dim):
+                    z_chunk_b = slice_tensor(
+                        z, i, i + offset, col_dim
+                    )
+                else: # b_chunk_dim == row_dim
+                    # In this case, the b-dimension (b_chunk_dim) is partially 
+                    # overwritten at the end of each iteration. We need to 
+                    # restore the missing component from the z-cache.
+                    if(not z_cache_rotated):
+                        z_chunk_slicer = empty_slicer(z_chunk_b)
+                        z_chunk_slicer[col_dim] = slice(0, half_n)
+                        z_chunk_b[z_chunk_slicer] = slice_tensor(
+                            z_cache, i, i + offset, row_dim,
+                        )
+                    else:
+                        z_cache_offset = i - half_n
+                        z_chunk_b = slice_tensor(
+                            z_cache, 
+                            z_cache_offset, z_cache_offset + offset, 
+                            row_dim
+                        )
+                b_chunk = compute_projection(z_chunk_b, mask_chunk, a=False, chunked=False)
+                #if bias==True:
+                #    b_chunk = compute_projection(
+                #        z_chunk_b, mask_chunk, True, a=False, chunked=False
+                #    )
+                #else:
+                #    b_chunk = compute_projection(z_chunk_b, mask_chunk, False, a=False, chunked=False)
+                del z_chunk_b
+
+                x_chunk = torch.matmul(
+                     a,
+                     b_chunk,
+                )
+                x_chunk = permute_final_dims(x_chunk, (1, 2, 0))
+                x_chunk = self.layer_norm_out(x_chunk)
+                x_chunk = self.linear_z(x_chunk)
+
+                # The g dimension (col_dim) is parallel to and ahead of the 
+                # overwrites in z. We can extract the g chunk normally.
+                z_chunk_g = slice_tensor(
+                    z, i, i + offset, col_dim
+                )
+                g_chunk = self.linear_g(self.layer_norm_in(z_chunk_g)) 
+                g_chunk.sigmoid_()
+                del z_chunk_g
+                
+                x_chunk *= g_chunk
+
+                # Write the columns into z in-place
+                z_slicer = empty_slicer(z)
+                z_slicer[col_dim] = slice(i, i + offset)
+                if(with_add):
+                    z[z_slicer] += x_chunk
+                else:
+                    z[z_slicer] = x_chunk
+        else:
+            b = compute_projection(z, mask, False, False)
+            #if bias==True:
+            #    b = compute_projection(z, mask, True, False, False)
+            #else:
+            #    b = compute_projection(z, mask, False, False, False)   
+            x = torch.matmul(a, b)
+            x = self.layer_norm_out(x)
+            x = self.linear_z(x)
+            g = self.linear_g(z)
+            g.sigmoid_()
+            x *= g
+            if(with_add):
+                z += x
+            else:
+                z = x
+
+        return z
+
+    def forward(self, 
+        z: torch.Tensor, 
+        mask: Optional[torch.Tensor] = None,
+        inplace_safe: bool = False,
+        _add_with_inplace: bool = False,
+        _inplace_chunk_size: Optional[int] = 256,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x:
+                [*, N_res, N_res, C_z] input tensor
+            mask:
+                [*, N_res, N_res] input mask
+        Returns:
+            [*, N_res, N_res, C_z] output tensor
+        """
+        if(inplace_safe):
+            x = self._inference_forward(
+                z, 
+                mask, 
+                inplace_chunk_size=_inplace_chunk_size,
+                with_add=_add_with_inplace,
+            )
+            return x
+
+        if mask is None:
+            mask = z.new_ones(z.shape[:-1])
+
+        mask = mask.unsqueeze(-1)
+        
+        z = self.layer_norm_in(z)
+        a = mask
+        a = a * self.sigmoid(self.linear_a_g(z)) 
+        a = a * self.linear_a_p(z)
+        b = mask
+        b = b * self.sigmoid(self.linear_b_g(z))
+        b = b * self.linear_b_p(z)
+
+        # Prevents overflow of torch.matmul in combine projections in
+        # reduced-precision modes
+        a_std = a.std()
+        b_std = b.std()
+        if(is_fp16_enabled() and a_std != 0. and b_std != 0.):
+            a = a / a.std()
+            b = b / b.std()
+
+        if(is_fp16_enabled()):
+            with torch.cuda.amp.autocast(enabled=False):
+                x = self._combine_projections(a.float(), b.float())
+        else:
+            x = self._combine_projections(a, b)
+        
+        del a, b
+        x = self.layer_norm_out(x)
+        x = self.linear_z(x)
+        g = self.sigmoid(self.linear_g(z))
+        x = x * g
+
+        return x
+
+
+class TriangleMultiplicationOutgoing(TriangleMultiplicativeUpdate):
+    """
+    Implements Algorithm 11.
+    """
+    __init__ = partialmethod(TriangleMultiplicativeUpdate.__init__, _outgoing=True)
+
+
+class TriangleMultiplicationIncoming(TriangleMultiplicativeUpdate):
+    """
+    Implements Algorithm 12.
+    """
+    __init__ = partialmethod(TriangleMultiplicativeUpdate.__init__, _outgoing=False)
+
+class ProtenixTriangleMultiplicationOutgoing(TriangleMultiplicativeUpdate):
+    """
+    Implements Algorithm 11.
+    """
+    __init__ = partialmethod(TriangleMultiplicativeUpdate.__init__, _outgoing=True, bias=False)
+
+class ProtenixTriangleMultiplicationIncoming(TriangleMultiplicativeUpdate):
+    """
+    Implements Algorithm 12.
+    """
+    __init__ = partialmethod(TriangleMultiplicativeUpdate.__init__, _outgoing=False, bias=False)
+
+
+class FusedTriangleMultiplicativeUpdate(BaseTriangleMultiplicativeUpdate):
+    """
+    Implements Algorithms 11 and 12.
+    """
+
+    def __init__(self, c_z, c_hidden, _outgoing=True):
+        """
+        Args:
+            c_z:
+                Input channel dimension
+            c:
+                Hidden channel dimension
+        """
+        super(FusedTriangleMultiplicativeUpdate, self).__init__(c_z=c_z,
+                                                                c_hidden=c_hidden,
+                                                                _outgoing=_outgoing)
+
+        self.linear_ab_p = Linear(self.c_z, self.c_hidden * 2)
+        self.linear_ab_g = Linear(self.c_z, self.c_hidden * 2, init="gating")
+
+    def _inference_forward(self,
+                           z: torch.Tensor,
+                           mask: Optional[torch.Tensor] = None,
+                           _inplace_chunk_size: Optional[int] = None,
+                           with_add: bool = True,
+                           ):
+        """
+        Args:
+            z:
+                A [*, N, N, C_z] pair representation
+            mask:
+                A [*, N, N] pair mask
+            with_add:
+                If True, z is overwritten with (z + update). Otherwise, it is
+                overwritten with (update).
+        Returns:
+            A reference to the overwritten z
+        """
+        if mask is None:
+            mask = z.new_ones(z.shape[:-1])
+
+        mask = mask.unsqueeze(-1)
+
+        def compute_projection_helper(pair, mask):
+            p = self.linear_ab_g(pair)
+            p.sigmoid_()
+            p *= self.linear_ab_p(pair)
+            p *= mask
+
+            return p
+
+        def compute_projection(pair, mask):
+            p = compute_projection_helper(pair, mask)
+            left = p[..., :self.c_hidden]
+            right = p[..., self.c_hidden:]
+
+            return left, right
+
+        z_norm_in = self.layer_norm_in(z)
+        a, b = compute_projection(z_norm_in, mask)
+        x = self._combine_projections(a, b, _inplace_chunk_size=_inplace_chunk_size)
+        x = self.layer_norm_out(x)
+        x = self.linear_z(x)
+        g = self.linear_g(z_norm_in)
+        g.sigmoid_()
+        x *= g
+        if (with_add):
+            z += x
+        else:
+            z = x
+
+        return z
+
+    def forward(self,
+                z: torch.Tensor,
+                mask: Optional[torch.Tensor] = None,
+                inplace_safe: bool = False,
+                _add_with_inplace: bool = False,
+                _inplace_chunk_size: Optional[int] = 256
+                ) -> torch.Tensor:
+        """
+        Args:
+            x:
+                [*, N_res, N_res, C_z] input tensor
+            mask:
+                [*, N_res, N_res] input mask
+        Returns:
+            [*, N_res, N_res, C_z] output tensor
+        """
+        if (inplace_safe):
+            x = self._inference_forward(
+                z,
+                mask,
+                _inplace_chunk_size=_inplace_chunk_size,
+                with_add=_add_with_inplace,
+            )
+            return x
+
+        if mask is None:
+            mask = z.new_ones(z.shape[:-1])
+
+        mask = mask.unsqueeze(-1)
+
+        z = self.layer_norm_in(z)
+        ab = mask
+        ab = ab * self.sigmoid(self.linear_ab_g(z))
+        ab = ab * self.linear_ab_p(z)
+
+        a = ab[..., :self.c_hidden]
+        b = ab[..., self.c_hidden:]
+
+        # Prevents overflow of torch.matmul in combine projections in
+        # reduced-precision modes
+        a_std = a.std()
+        b_std = b.std()
+        if (is_fp16_enabled() and a_std != 0. and b_std != 0.):
+            a = a / a.std()
+            b = b / b.std()
+
+        if (is_fp16_enabled()):
+            with torch.cuda.amp.autocast(enabled=False):
+                x = self._combine_projections(a.float(), b.float())
+        else:
+            x = self._combine_projections(a, b)
+
+        del a, b
+        x = self.layer_norm_out(x)
+        x = self.linear_z(x)
+        g = self.sigmoid(self.linear_g(z))
+        x = x * g
+
+        return x
+
+
+class FusedTriangleMultiplicationOutgoing(FusedTriangleMultiplicativeUpdate):
+    """
+    Implements Algorithm 11.
+    """
+    __init__ = partialmethod(FusedTriangleMultiplicativeUpdate.__init__, _outgoing=True)
+
+
+class FusedTriangleMultiplicationIncoming(FusedTriangleMultiplicativeUpdate):
+    """
+    Implements Algorithm 12.
+    """
+    __init__ = partialmethod(FusedTriangleMultiplicativeUpdate.__init__, _outgoing=False)
+

weight/esm-main/scripts/atlas/v0/full/tarballs/tm_.90_1_plddt_.70_.80.txt

@@ -0,0 +1 @@
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_.90_1_plddt_.70_.80_00.tar.gz

weight/esm-main/scripts/atlas/v0/full/tarballs/tm_0_.40_plddt_0_.40.txt

@@ -0,0 +1,46 @@
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_00.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_01.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_02.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_03.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_04.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_05.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_06.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_07.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_08.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_09.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_10.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_11.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_12.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_13.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_14.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_15.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_16.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_17.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_18.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_19.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_20.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_21.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_22.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_23.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_24.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_25.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_26.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_27.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_28.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_29.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_30.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_31.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_32.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_33.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_34.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_35.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_36.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_37.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_38.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_39.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_40.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_41.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_42.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_43.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_44.tar.gz
+https://dl.fbaipublicfiles.com/esmatlas/v0/full/tarballs/tm_0_.40_plddt_0_.40_45.tar.gz

weight/esm-main/scripts/atlas/v2023_02/full/esm2_embeddings/tm_.70_.80_plddt_.70_.80.txt

@@ -0,0 +1,13 @@
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_00.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_01.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_02.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_03.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_04.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_05.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_06.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_07.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_08.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_09.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_10.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_11.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.70_.80_12.npz

weight/esm-main/scripts/atlas/v2023_02/full/esm2_embeddings/tm_.70_.80_plddt_.80_.90.txt

@@ -0,0 +1,2 @@
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.80_.90_00.npz
+https://dl.fbaipublicfiles.com/esmatlas/v2023_02/full/lm_reps/tm_.70_.80_plddt_.80_.90_01.npz