Apply coding standards from 1cb5747 (files only)

README.md

@@ -16,11 +16,12 @@ Load the published model, fold two protein chains together, and write an mmCIF
 file. The example omits `num_sampling_steps` and uses the model default.
 
 ```python
-from pathlib import Path
-
 import torch
+
+from pathlib import Path
 from transformers import AutoModel
 
+
 model = AutoModel.from_pretrained(
     "Synthyra/ESMFold2-Fast",
     trust_remote_code=True,
@@ -109,11 +110,13 @@ state mixture and projection, followed by one trainable transformer probe.
 
 ```python
 import torch
+
 from transformers import (
     AutoModelForSequenceClassification,
     AutoModelForTokenClassification,
 )
 
+
 model_id = "Synthyra/ESMFold2-Fast"
 sequence_model = AutoModelForSequenceClassification.from_pretrained(
     model_id, num_labels=2, trust_remote_code=True
@@ -123,11 +126,11 @@ token_model = AutoModelForTokenClassification.from_pretrained(
 ).eval()
 sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"]
 batch = sequence_model.prepare_classifier_inputs(sequences)
-biological = batch["attention_mask"].bool()
+biological = batch["attention_mask"].bool()  # (b, l)
 
-sequence_labels = torch.zeros(len(sequences), dtype=torch.long)
-token_labels = torch.full_like(batch["input_ids"], -100)
-token_labels[biological] = 0
+sequence_labels = torch.zeros(len(sequences), dtype=torch.long)  # (b,)
+token_labels = torch.full_like(batch["input_ids"], -100)  # (b, l)
+token_labels[biological] = 0  # selected biological positions; labels stay (b, l)
 
 with torch.inference_mode():
     sequence_output = sequence_model(**batch, labels=sequence_labels)
@@ -147,6 +150,7 @@ python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
 ```python
 from peft import LoraConfig, TaskType, get_peft_model
 
+
 peft_model = get_peft_model(
     sequence_model,
     LoraConfig(

fastplms/attention/_core.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import warnings
 import torch
+
 from collections import OrderedDict
 from collections.abc import Callable
 from dataclasses import dataclass

fastplms/attention/_kernel_lock.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import json
 import os
+
 from pathlib import Path
 from typing import Any
 

fastplms/attention/interfaces.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import torch
+
 from collections.abc import Mapping
 from functools import partial
 from typing import Any

fastplms/embeddings/batches.py

@@ -1,378 +1,379 @@
-"""Execute model-specific batches and return ordered residue-aware CPU tensors."""
-
-from __future__ import annotations
-
-import torch
-from collections.abc import Callable, Iterator, Sequence
-from contextlib import contextmanager
-from dataclasses import dataclass, field
-from typing import Any
-from torch import Tensor
-
-from .identity import _model_device
-from .inputs import _planned_batches
-from .pooling import Pooler
-from .types import EmbeddingBatch, EmbeddingInput, EmbeddingRecord
-
-
-_MAX_PARTI_RESIDUES = 2_048
-
-
-def _validate_parti_length(M: Tensor) -> None:
-    """Reject an oversized attention graph before model inference."""
-
-    # M: (b, l)
-    n_residues = int(M.to(dtype=torch.int64).sum(dim=1).max().item())
-    if n_residues > _MAX_PARTI_RESIDUES:
-        raise ValueError(f"parti supports at most {_MAX_PARTI_RESIDUES:,} biological residues.")
-
-
-def select_hidden_state_embeddings(
-    last_hidden_state: Tensor,
-    hidden_states: tuple[Tensor, ...] | None,
-    *,
-    hidden_state_index: int = -1,
-    store_all_hidden_states: bool = False,
-) -> Tensor:
-    """Select one hidden state or stack every state without changing values."""
-    # last_hidden_state and each hidden_states entry: (b, l, d)
-    if store_all_hidden_states:
-        if not hidden_states:
-            raise ValueError("store_all_hidden_states requires model hidden states.")
-        # H has shape (b, n, l, d), where n follows the model's output order.
-        return torch.stack(hidden_states, dim=1)  # (b, n, l, d)
-    if hidden_state_index == -1:
-        return last_hidden_state  # (b, l, d)
-    if not hidden_states:
-        raise ValueError("hidden_state_index requires model hidden states.")
-    return hidden_states[hidden_state_index]  # (b, l, d)
-
-
-def _residue_embeddings(X: Tensor, M: Tensor) -> list[Tensor]:
-    """Copy every sample's biological residues to the host in one transfer.
-
-    Boolean indexing packs the selected rows in batch order, so splitting the
-    packed rows by residue count gives the values that indexing each sample
-    would. Each returned tensor owns its storage, as a per-sample copy does.
-    """
-    # X: (b, l, d); M: (b, l)
-    residue_counts = M.sum(dim=1).tolist()  # b counts r_i
-    packed = X[M].detach().cpu()  # (sum of r_i, d)
-    return [sample.clone() for sample in torch.split(packed, residue_counts)]  # each: (r_i, d)
-
-
-@contextmanager
-def _temporary_eval(model: Any) -> Iterator[None]:
-    was_training = getattr(model, "training", None)
-    eval_method = getattr(model, "eval", None)
-    train_method = getattr(model, "train", None)
-    if (
-        not isinstance(was_training, bool)
-        or not callable(eval_method)
-        or not callable(train_method)
-    ):
-        yield
-        return
-    eval_method()
-    try:
-        yield
-    finally:
-        train_method(was_training)
-
-
-def _biological_residue_mask(
-    input_ids: Tensor,
-    attention_mask: Tensor,
-    tokenizer: Any,
-) -> Tensor:
-    """Remove padding and tokenizer-declared special tokens from M."""
-
-    # input_ids, attention_mask: (b, l)
-    M = attention_mask.to(dtype=torch.bool)  # (b, l)
-    special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ()))
-    if special_ids:
-        specials = torch.tensor(  # (n_special,)
-            special_ids,
-            device=input_ids.device,
-            dtype=input_ids.dtype,
-        )
-        M = M & ~torch.isin(input_ids, specials)  # (b, l)
-    return M  # (b, l)
-
-
-def _generic_embedding_batch(
-    model: Any,
-    sequences: list[str],
-    *,
-    tokenizer: Any | None,
-    max_length: int | None,
-    truncate: bool,
-    need_attentions: bool,
-    model_kwargs: dict[str, Any],
-) -> EmbeddingBatch:
-    config = getattr(model, "config", None)
-    model_type = str(getattr(config, "model_type", "")).lower()
-    if tokenizer is None:
-        tokenizer = getattr(model, "tokenizer", None)
-
-    if tokenizer is None and model_type == "e1":
-        output = model._embed(sequences, return_attention_mask=True, **model_kwargs)
-        if not isinstance(output, tuple) or len(output) != 2:
-            raise TypeError("E1 _embed must return (X, residue_mask).")
-        X, M = output  # (b, l, d), (b, l)
-        preparer = getattr(model, "prep_tokens", None)
-        if preparer is not None and hasattr(preparer, "get_batch_kwargs"):
-            prepared = preparer.get_batch_kwargs(sequences, device=X.device)
-            input_ids = prepared["input_ids"]  # (b, l)
-            boundary_ids = preparer.boundary_token_ids.to(  # (n_boundary,)
-                device=input_ids.device, dtype=input_ids.dtype
-            )
-            # E1 wraps each raw sequence in BOS, context-label, terminal-label,
-            # and EOS tokens. Only amino-acid rows are biological residues.
-            M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids)  # (b, l)
-        if need_attentions:
-            raise ValueError("parti is not available for tokenizer-free E1 embedding.")
-        return EmbeddingBatch(  # X: (b, l, d); residue_mask: (b, l)
-            X=X,
-            residue_mask=M.to(dtype=torch.bool),
-        )
-    if tokenizer is None:
-        raise ValueError("A tokenizer is required for this model's embedding path.")
-
-    tokenize_kwargs: dict[str, Any] = {
-        "return_tensors": "pt",
-        "padding": True,
-        "truncation": truncate,
-    }
-    if max_length is not None and truncate:
-        # ``max_length`` is a biological-residue limit. Tokenizer limits include
-        # boundary tokens, so reserve their declared width instead of dropping
-        # residues at the exact boundary.
-        special_token_count = 0
-        num_special_tokens_to_add = getattr(tokenizer, "num_special_tokens_to_add", None)
-        if callable(num_special_tokens_to_add):
-            special_token_count = int(num_special_tokens_to_add(pair=False))
-        tokenize_kwargs["max_length"] = max_length + special_token_count
-    sequence_tokenizer = getattr(model, "_tokenize_sequence_batch", None)
-    if callable(sequence_tokenizer):
-        encoded = sequence_tokenizer(sequences, tokenizer=tokenizer, **tokenize_kwargs)
-    else:
-        encoded = tokenizer(sequences, **tokenize_kwargs)
-    device = _model_device(model)
-    input_ids = encoded["input_ids"].to(device)  # (b, l)
-    attention_mask = encoded.get(  # (b, l)
-        "attention_mask",
-        input_ids.new_ones(input_ids.shape),
-    ).to(device)
-    M = _biological_residue_mask(input_ids, attention_mask, tokenizer)  # (b, l)
-    if need_attentions:
-        # Validate l before either the backbone or its quadratic attention graph
-        # is materialized. M has shape (b, l).
-        _validate_parti_length(M)
-    X = model._embed(input_ids, attention_mask, **model_kwargs)  # (b, l, d)
-    attentions = None
-    if need_attentions:
-        output = model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            output_attentions=True,
-            return_dict=True,
-        )
-        attentions = getattr(output, "attentions", None)  # each: (b, h, l, l)
-        if attentions is None:
-            raise ValueError("The model did not return attentions required by parti.")
-    return EmbeddingBatch(  # X: (b, l, d); M: (b, l)
-        X=X,
-        residue_mask=M,
-        attentions=attentions,
-    )
-
-
-@dataclass(eq=False)
-class BatchExecutor:
-    """Model and batch policy for one bounded embedding window at a time."""
-
-    model: Any
-    batch_size: int
-    max_tokens_per_batch: int | None
-    max_length: int | None
-    truncate: bool
-    model_kwargs: dict[str, Any]
-    hidden_state_source: str
-    normalized_decoder_inputs: tuple[str, ...] | None
-    decoder_input_ids: Tensor | None
-    decoder_attention_mask: Tensor | None
-    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None
-    tokenizer: Any | None
-    store_all_hidden_states: bool
-    full_embeddings: bool
-    dtype: torch.dtype | None
-    pooler: Pooler | None
-    attention_backend: str | None
-    need_attentions: bool
-    model_type: str = field(init=False)
-    resolved_tokenizer: Any = field(init=False)
-
-    def __post_init__(self) -> None:
-        config = getattr(self.model, "config", None)
-        self.model_type = str(getattr(config, "model_type", "")).lower()
-        self.resolved_tokenizer = (
-            self.tokenizer if self.tokenizer is not None else getattr(self.model, "tokenizer", None)
-        )
-
-    def run_window(
-        self,
-        window_records: Sequence[EmbeddingInput],
-        *,
-        window_start: int,
-    ) -> tuple[list[EmbeddingRecord], dict[str, tuple[int, int]]]:
-        """Restore source order after length-bucketed inference and pooling."""
-
-        pool_slices: dict[str, tuple[int, int]] = {}
-        window_results: dict[int, EmbeddingRecord] = {}
-        for local_positions in _planned_batches(
-            window_records,
-            range(len(window_records)),
-            batch_size=self.batch_size,
-            max_tokens_per_batch=self.max_tokens_per_batch,
-            max_length=self.max_length,
-            truncate=self.truncate,
-        ):
-            batch_positions = [window_start + position for position in local_positions]
-            batch_records = [window_records[position] for position in local_positions]
-            sequences = [
-                record.sequence[: self.max_length]
-                if self.truncate and self.max_length is not None
-                else record.sequence
-                for record in batch_records
-            ]
-            batch_model_kwargs = dict(self.model_kwargs)
-            if self.model_type == "fast_ankh" or self.hidden_state_source == "decoder":
-                batch_model_kwargs["hidden_state_source"] = self.hidden_state_source
-                if self.normalized_decoder_inputs is not None:
-                    batch_model_kwargs["decoder_inputs"] = [
-                        self.normalized_decoder_inputs[position] for position in batch_positions
-                    ]
-                if self.decoder_input_ids is not None:
-                    # decoder_input_ids: (n_records, l_decoder)
-                    indices = torch.tensor(  # (b,)
-                        batch_positions,
-                        device=self.decoder_input_ids.device,
-                        dtype=torch.long,
-                    )
-                    batch_model_kwargs["decoder_input_ids"] = (  # (b, l_decoder)
-                        self.decoder_input_ids.index_select(0, indices)
-                    )
-                if self.decoder_attention_mask is not None:
-                    # decoder_attention_mask: (n_records, l_decoder)
-                    indices = torch.tensor(  # (b,)
-                        batch_positions,
-                        device=self.decoder_attention_mask.device,
-                        dtype=torch.long,
-                    )
-                    batch_model_kwargs["decoder_attention_mask"] = (
-                        self.decoder_attention_mask.index_select(0, indices)  # (b, l_decoder)
-                    )
-            custom_batch = self._embedding_batch_fn or getattr(self.model, "_embedding_batch", None)
-            if custom_batch is not None:
-                if self.model_type == "fast_ankh":
-                    batch = custom_batch(
-                        sequences,
-                        tokenizer=self.resolved_tokenizer,
-                        max_length=self.max_length,
-                        truncate=self.truncate,
-                        need_attentions=self.need_attentions,
-                        **batch_model_kwargs,
-                    )
-                else:
-                    batch = custom_batch(sequences, **batch_model_kwargs)
-                if not isinstance(batch, EmbeddingBatch):
-                    raise TypeError("_embedding_batch must return EmbeddingBatch.")
-            else:
-                batch = _generic_embedding_batch(
-                    self.model,
-                    sequences,
-                    tokenizer=self.tokenizer,
-                    max_length=self.max_length,
-                    truncate=self.truncate,
-                    need_attentions=self.need_attentions,
-                    model_kwargs=batch_model_kwargs,
-                )
-            X = batch.X  # (b, l, d) or (b, n_states, l, d)
-            raw_mask = batch.residue_mask  # (b, l)
-            if not isinstance(X, Tensor) or not isinstance(raw_mask, Tensor):
-                raise TypeError("Embedding batches must provide Tensor X and residue_mask.")
-            if X.is_meta or raw_mask.is_meta:
-                raise ValueError("Embedding batches cannot contain meta tensors.")
-            if not X.is_floating_point():
-                raise TypeError("Embedding batches must use a floating-point X dtype.")
-            if raw_mask.is_complex() or not bool(torch.isfinite(raw_mask).all()):
-                raise ValueError("Embedding residue_mask must contain finite binary values.")
-            if not bool(((raw_mask == 0) | (raw_mask == 1)).all()):
-                raise ValueError("Embedding residue_mask must contain finite binary values.")
-            M = raw_mask.to(device=X.device, dtype=torch.bool)  # (b, l)
-            valid_X_shape = (
-                X.ndim == 3
-                and X.shape[0] == len(batch_records)
-                and X.shape[-1] > 0
-                and M.shape == X.shape[:2]
-            )
-            valid_all_states_shape = (
-                X.ndim == 4
-                and self.store_all_hidden_states
-                and self.full_embeddings
-                and X.shape[0] == len(batch_records)
-                and X.shape[1] > 0
-                and X.shape[-1] > 0
-                and M.shape == (X.shape[0], X.shape[2])
-            )
-            if not (valid_X_shape or valid_all_states_shape):
-                raise ValueError(
-                    "Embedding batches must provide X with shape (b, l, d), or "
-                    "(b, states, l, d) when storing all hidden states, and "
-                    "residue_mask with shape (b, l)."
-                )
-            if not bool(M.any(dim=1).all()):
-                raise ValueError("Every embedding sample must contain a biological residue.")
-            finite_selected = (  # X.shape
-                torch.isfinite(X) | ~M.unsqueeze(-1)
-                if X.ndim == 3
-                else torch.isfinite(X) | ~M[:, None, :, None]
-            )
-            if not bool(finite_selected.all()):
-                raise ValueError("Biological residue embeddings produced non-finite output.")
-            if self.need_attentions:
-                # Validate the biological graph only after mask integrity is established.
-                _validate_parti_length(M)
-            if self.dtype is not None:
-                X = X.to(dtype=self.dtype)  # unchanged shape
-
-            if self.full_embeddings:
-                if X.ndim == 4:
-                    values = [
-                        X_i[:, M_i, :].detach().cpu()  # (n_states, r_i, d)
-                        for X_i, M_i in zip(X, M, strict=True)
-                    ]
-                else:
-                    values = _residue_embeddings(X, M)  # each: (r_i, d)
-            else:
-                if self.pooler is None:
-                    raise RuntimeError(
-                        "Pooled embedding output was requested without an initialized pooler."
-                    )
-                Y = self.pooler(  # (b, n_poolers * d)
-                    X,
-                    M,
-                    attentions=batch.attentions,
-                    attention_backend=self.attention_backend,
-                )
-                pool_slices = self.pooler.output_slices(X.shape[-1])
-                values = list(Y.detach().cpu().unbind(0))  # each: (n_poolers * d,)
-            for position, record, value in zip(batch_positions, batch_records, values, strict=True):
-                window_results[position] = EmbeddingRecord(record.id, record.sequence, value)
-
-        new_records = [
-            window_results[position]
-            for position in range(window_start, window_start + len(window_records))
-        ]
-        return new_records, pool_slices
+"""Execute model-specific batches and return ordered residue-aware CPU tensors."""
+
+from __future__ import annotations
+
+import torch
+
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any
+from torch import Tensor
+
+from .identity import _model_device
+from .inputs import _planned_batches
+from .pooling import Pooler
+from .types import EmbeddingBatch, EmbeddingInput, EmbeddingRecord
+
+
+_MAX_PARTI_RESIDUES = 2_048
+
+
+def _validate_parti_length(M: Tensor) -> None:
+    """Reject an oversized attention graph before model inference."""
+
+    # M: (b, l)
+    n_residues = int(M.to(dtype=torch.int64).sum(dim=1).max().item())
+    if n_residues > _MAX_PARTI_RESIDUES:
+        raise ValueError(f"parti supports at most {_MAX_PARTI_RESIDUES:,} biological residues.")
+
+
+def select_hidden_state_embeddings(
+    last_hidden_state: Tensor,
+    hidden_states: tuple[Tensor, ...] | None,
+    *,
+    hidden_state_index: int = -1,
+    store_all_hidden_states: bool = False,
+) -> Tensor:
+    """Select one hidden state or stack every state without changing values."""
+    # last_hidden_state and each hidden_states entry: (b, l, d)
+    if store_all_hidden_states:
+        if not hidden_states:
+            raise ValueError("store_all_hidden_states requires model hidden states.")
+        # H has shape (b, n, l, d), where n follows the model's output order.
+        return torch.stack(hidden_states, dim=1)  # (b, n, l, d)
+    if hidden_state_index == -1:
+        return last_hidden_state  # (b, l, d)
+    if not hidden_states:
+        raise ValueError("hidden_state_index requires model hidden states.")
+    return hidden_states[hidden_state_index]  # (b, l, d)
+
+
+def _residue_embeddings(X: Tensor, M: Tensor) -> list[Tensor]:
+    """Copy every sample's biological residues to the host in one transfer.
+
+    Boolean indexing packs the selected rows in batch order, so splitting the
+    packed rows by residue count gives the values that indexing each sample
+    would. Each returned tensor owns its storage, as a per-sample copy does.
+    """
+    # X: (b, l, d); M: (b, l)
+    residue_counts = M.sum(dim=1).tolist()  # b counts r_i
+    packed = X[M].detach().cpu()  # (sum of r_i, d)
+    return [sample.clone() for sample in torch.split(packed, residue_counts)]  # each: (r_i, d)
+
+
+@contextmanager
+def _temporary_eval(model: Any) -> Iterator[None]:
+    was_training = getattr(model, "training", None)
+    eval_method = getattr(model, "eval", None)
+    train_method = getattr(model, "train", None)
+    if (
+        not isinstance(was_training, bool)
+        or not callable(eval_method)
+        or not callable(train_method)
+    ):
+        yield
+        return
+    eval_method()
+    try:
+        yield
+    finally:
+        train_method(was_training)
+
+
+def _biological_residue_mask(
+    input_ids: Tensor,
+    attention_mask: Tensor,
+    tokenizer: Any,
+) -> Tensor:
+    """Remove padding and tokenizer-declared special tokens from M."""
+
+    # input_ids, attention_mask: (b, l)
+    M = attention_mask.to(dtype=torch.bool)  # (b, l)
+    special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ()))
+    if special_ids:
+        specials = torch.tensor(  # (n_special,)
+            special_ids,
+            device=input_ids.device,
+            dtype=input_ids.dtype,
+        )
+        M = M & ~torch.isin(input_ids, specials)  # (b, l)
+    return M  # (b, l)
+
+
+def _generic_embedding_batch(
+    model: Any,
+    sequences: list[str],
+    *,
+    tokenizer: Any | None,
+    max_length: int | None,
+    truncate: bool,
+    need_attentions: bool,
+    model_kwargs: dict[str, Any],
+) -> EmbeddingBatch:
+    config = getattr(model, "config", None)
+    model_type = str(getattr(config, "model_type", "")).lower()
+    if tokenizer is None:
+        tokenizer = getattr(model, "tokenizer", None)
+
+    if tokenizer is None and model_type == "e1":
+        output = model._embed(sequences, return_attention_mask=True, **model_kwargs)
+        if not isinstance(output, tuple) or len(output) != 2:
+            raise TypeError("E1 _embed must return (X, residue_mask).")
+        X, M = output  # (b, l, d), (b, l)
+        preparer = getattr(model, "prep_tokens", None)
+        if preparer is not None and hasattr(preparer, "get_batch_kwargs"):
+            prepared = preparer.get_batch_kwargs(sequences, device=X.device)
+            input_ids = prepared["input_ids"]  # (b, l)
+            boundary_ids = preparer.boundary_token_ids.to(  # (n_boundary,)
+                device=input_ids.device, dtype=input_ids.dtype
+            )
+            # E1 wraps each raw sequence in BOS, context-label, terminal-label,
+            # and EOS tokens. Only amino-acid rows are biological residues.
+            M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids)  # (b, l)
+        if need_attentions:
+            raise ValueError("parti is not available for tokenizer-free E1 embedding.")
+        return EmbeddingBatch(  # X: (b, l, d); residue_mask: (b, l)
+            X=X,
+            residue_mask=M.to(dtype=torch.bool),
+        )
+    if tokenizer is None:
+        raise ValueError("A tokenizer is required for this model's embedding path.")
+
+    tokenize_kwargs: dict[str, Any] = {
+        "return_tensors": "pt",
+        "padding": True,
+        "truncation": truncate,
+    }
+    if max_length is not None and truncate:
+        # ``max_length`` is a biological-residue limit. Tokenizer limits include
+        # boundary tokens, so reserve their declared width instead of dropping
+        # residues at the exact boundary.
+        special_token_count = 0
+        num_special_tokens_to_add = getattr(tokenizer, "num_special_tokens_to_add", None)
+        if callable(num_special_tokens_to_add):
+            special_token_count = int(num_special_tokens_to_add(pair=False))
+        tokenize_kwargs["max_length"] = max_length + special_token_count
+    sequence_tokenizer = getattr(model, "_tokenize_sequence_batch", None)
+    if callable(sequence_tokenizer):
+        encoded = sequence_tokenizer(sequences, tokenizer=tokenizer, **tokenize_kwargs)
+    else:
+        encoded = tokenizer(sequences, **tokenize_kwargs)
+    device = _model_device(model)
+    input_ids = encoded["input_ids"].to(device)  # (b, l)
+    attention_mask = encoded.get(  # (b, l)
+        "attention_mask",
+        input_ids.new_ones(input_ids.shape),
+    ).to(device)
+    M = _biological_residue_mask(input_ids, attention_mask, tokenizer)  # (b, l)
+    if need_attentions:
+        # Validate l before either the backbone or its quadratic attention graph
+        # is materialized. M has shape (b, l).
+        _validate_parti_length(M)
+    X = model._embed(input_ids, attention_mask, **model_kwargs)  # (b, l, d)
+    attentions = None
+    if need_attentions:
+        output = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_attentions=True,
+            return_dict=True,
+        )
+        attentions = getattr(output, "attentions", None)  # each: (b, h, l, l)
+        if attentions is None:
+            raise ValueError("The model did not return attentions required by parti.")
+    return EmbeddingBatch(  # X: (b, l, d); M: (b, l)
+        X=X,
+        residue_mask=M,
+        attentions=attentions,
+    )
+
+
+@dataclass(eq=False)
+class BatchExecutor:
+    """Model and batch policy for one bounded embedding window at a time."""
+
+    model: Any
+    batch_size: int
+    max_tokens_per_batch: int | None
+    max_length: int | None
+    truncate: bool
+    model_kwargs: dict[str, Any]
+    hidden_state_source: str
+    normalized_decoder_inputs: tuple[str, ...] | None
+    decoder_input_ids: Tensor | None
+    decoder_attention_mask: Tensor | None
+    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None
+    tokenizer: Any | None
+    store_all_hidden_states: bool
+    full_embeddings: bool
+    dtype: torch.dtype | None
+    pooler: Pooler | None
+    attention_backend: str | None
+    need_attentions: bool
+    model_type: str = field(init=False)
+    resolved_tokenizer: Any = field(init=False)
+
+    def __post_init__(self) -> None:
+        config = getattr(self.model, "config", None)
+        self.model_type = str(getattr(config, "model_type", "")).lower()
+        self.resolved_tokenizer = (
+            self.tokenizer if self.tokenizer is not None else getattr(self.model, "tokenizer", None)
+        )
+
+    def run_window(
+        self,
+        window_records: Sequence[EmbeddingInput],
+        *,
+        window_start: int,
+    ) -> tuple[list[EmbeddingRecord], dict[str, tuple[int, int]]]:
+        """Restore source order after length-bucketed inference and pooling."""
+
+        pool_slices: dict[str, tuple[int, int]] = {}
+        window_results: dict[int, EmbeddingRecord] = {}
+        for local_positions in _planned_batches(
+            window_records,
+            range(len(window_records)),
+            batch_size=self.batch_size,
+            max_tokens_per_batch=self.max_tokens_per_batch,
+            max_length=self.max_length,
+            truncate=self.truncate,
+        ):
+            batch_positions = [window_start + position for position in local_positions]
+            batch_records = [window_records[position] for position in local_positions]
+            sequences = [
+                record.sequence[: self.max_length]
+                if self.truncate and self.max_length is not None
+                else record.sequence
+                for record in batch_records
+            ]
+            batch_model_kwargs = dict(self.model_kwargs)
+            if self.model_type == "fast_ankh" or self.hidden_state_source == "decoder":
+                batch_model_kwargs["hidden_state_source"] = self.hidden_state_source
+                if self.normalized_decoder_inputs is not None:
+                    batch_model_kwargs["decoder_inputs"] = [
+                        self.normalized_decoder_inputs[position] for position in batch_positions
+                    ]
+                if self.decoder_input_ids is not None:
+                    # decoder_input_ids: (n_records, l_decoder)
+                    indices = torch.tensor(  # (b,)
+                        batch_positions,
+                        device=self.decoder_input_ids.device,
+                        dtype=torch.long,
+                    )
+                    batch_model_kwargs["decoder_input_ids"] = (  # (b, l_decoder)
+                        self.decoder_input_ids.index_select(0, indices)
+                    )
+                if self.decoder_attention_mask is not None:
+                    # decoder_attention_mask: (n_records, l_decoder)
+                    indices = torch.tensor(  # (b,)
+                        batch_positions,
+                        device=self.decoder_attention_mask.device,
+                        dtype=torch.long,
+                    )
+                    batch_model_kwargs["decoder_attention_mask"] = (
+                        self.decoder_attention_mask.index_select(0, indices)  # (b, l_decoder)
+                    )
+            custom_batch = self._embedding_batch_fn or getattr(self.model, "_embedding_batch", None)
+            if custom_batch is not None:
+                if self.model_type == "fast_ankh":
+                    batch = custom_batch(
+                        sequences,
+                        tokenizer=self.resolved_tokenizer,
+                        max_length=self.max_length,
+                        truncate=self.truncate,
+                        need_attentions=self.need_attentions,
+                        **batch_model_kwargs,
+                    )
+                else:
+                    batch = custom_batch(sequences, **batch_model_kwargs)
+                if not isinstance(batch, EmbeddingBatch):
+                    raise TypeError("_embedding_batch must return EmbeddingBatch.")
+            else:
+                batch = _generic_embedding_batch(
+                    self.model,
+                    sequences,
+                    tokenizer=self.tokenizer,
+                    max_length=self.max_length,
+                    truncate=self.truncate,
+                    need_attentions=self.need_attentions,
+                    model_kwargs=batch_model_kwargs,
+                )
+            X = batch.X  # (b, l, d) or (b, n_states, l, d)
+            raw_mask = batch.residue_mask  # (b, l)
+            if not isinstance(X, Tensor) or not isinstance(raw_mask, Tensor):
+                raise TypeError("Embedding batches must provide Tensor X and residue_mask.")
+            if X.is_meta or raw_mask.is_meta:
+                raise ValueError("Embedding batches cannot contain meta tensors.")
+            if not X.is_floating_point():
+                raise TypeError("Embedding batches must use a floating-point X dtype.")
+            if raw_mask.is_complex() or not bool(torch.isfinite(raw_mask).all()):
+                raise ValueError("Embedding residue_mask must contain finite binary values.")
+            if not bool(((raw_mask == 0) | (raw_mask == 1)).all()):
+                raise ValueError("Embedding residue_mask must contain finite binary values.")
+            M = raw_mask.to(device=X.device, dtype=torch.bool)  # (b, l)
+            valid_X_shape = (
+                X.ndim == 3
+                and X.shape[0] == len(batch_records)
+                and X.shape[-1] > 0
+                and M.shape == X.shape[:2]
+            )
+            valid_all_states_shape = (
+                X.ndim == 4
+                and self.store_all_hidden_states
+                and self.full_embeddings
+                and X.shape[0] == len(batch_records)
+                and X.shape[1] > 0
+                and X.shape[-1] > 0
+                and M.shape == (X.shape[0], X.shape[2])
+            )
+            if not (valid_X_shape or valid_all_states_shape):
+                raise ValueError(
+                    "Embedding batches must provide X with shape (b, l, d), or "
+                    "(b, states, l, d) when storing all hidden states, and "
+                    "residue_mask with shape (b, l)."
+                )
+            if not bool(M.any(dim=1).all()):
+                raise ValueError("Every embedding sample must contain a biological residue.")
+            finite_selected = (  # X.shape
+                torch.isfinite(X) | ~M.unsqueeze(-1)
+                if X.ndim == 3
+                else torch.isfinite(X) | ~M[:, None, :, None]
+            )
+            if not bool(finite_selected.all()):
+                raise ValueError("Biological residue embeddings produced non-finite output.")
+            if self.need_attentions:
+                # Validate the biological graph only after mask integrity is established.
+                _validate_parti_length(M)
+            if self.dtype is not None:
+                X = X.to(dtype=self.dtype)  # unchanged shape
+
+            if self.full_embeddings:
+                if X.ndim == 4:
+                    values = [
+                        X_i[:, M_i, :].detach().cpu()  # (n_states, r_i, d)
+                        for X_i, M_i in zip(X, M, strict=True)
+                    ]
+                else:
+                    values = _residue_embeddings(X, M)  # each: (r_i, d)
+            else:
+                if self.pooler is None:
+                    raise RuntimeError(
+                        "Pooled embedding output was requested without an initialized pooler."
+                    )
+                Y = self.pooler(  # (b, n_poolers * d)
+                    X,
+                    M,
+                    attentions=batch.attentions,
+                    attention_backend=self.attention_backend,
+                )
+                pool_slices = self.pooler.output_slices(X.shape[-1])
+                values = list(Y.detach().cpu().unbind(0))  # each: (n_poolers * d,)
+            for position, record, value in zip(batch_positions, batch_records, values, strict=True):
+                window_results[position] = EmbeddingRecord(record.id, record.sequence, value)
+
+        new_records = [
+            window_results[position]
+            for position in range(window_start, window_start + len(window_records))
+        ]
+        return new_records, pool_slices

fastplms/embeddings/identity.py

@@ -1,510 +1,511 @@
-"""Deterministic identity for embedding inputs, models, tokenizers, and execution."""
-
-from __future__ import annotations
-
-import hashlib
-import json
-import platform
-import torch
-from collections.abc import Iterable, Mapping, Sequence
-from pathlib import Path
-from typing import Any
-from torch import Tensor
-
+"""Deterministic identity for embedding inputs, models, tokenizers, and execution."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import platform
+import torch
+
+from collections.abc import Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
+from torch import Tensor
+
 from .inputs import _InputSpool
 from .storage import tensor_sha256
 from .types import EmbeddingInput
-
-
-_RUN_FINGERPRINT_SCHEMA_VERSION = 3
-_MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2
-
-
-def _model_device(model: Any) -> torch.device:
-    try:
-        return torch.device(next(model.parameters()).device)
-    except (AttributeError, StopIteration):
-        return torch.device("cpu")
-
-
-def _attention_backend(model: Any) -> str | None:
-    config = getattr(model, "config", None)
-    for name in ("_attn_implementation", "attn_implementation", "attn_backend"):
-        value = getattr(config, name, None)
-        if value:
-            return str(value)
-    return None
-
-
-def _attention_kernel_metadata(backend: str | None) -> dict[str, Any] | None:
-    if backend not in {"flash_attention_2", "flash_attention_3"}:
-        return None
-    from fastplms.registry import get_model_registry
-
-    spec = get_model_registry().attention_kernels[backend]
-    return {
-        "repository": spec.repository,
-        "revision": spec.revision,
-        "version": spec.version,
-        "expected_variant": spec.expected_variant,
-        "dtypes": list(spec.dtypes),
-    }
-
-
-def _fingerprint_jsonable(value: Any) -> Any:
-    if isinstance(value, Mapping):
-        return {str(key): _fingerprint_jsonable(item) for key, item in value.items()}
-    if isinstance(value, (list, tuple)):
-        return [_fingerprint_jsonable(item) for item in value]
-    if isinstance(value, (set, frozenset)):
-        return sorted((_fingerprint_jsonable(item) for item in value), key=repr)
-    if isinstance(value, Path):
-        return str(value)
-    if isinstance(value, Tensor):
-        return {
-            "dtype": str(value.dtype).removeprefix("torch."),
-            "shape": list(value.shape),
-            "sha256": tensor_sha256(value),
-        }
-    if isinstance(value, torch.dtype):
-        return str(value).removeprefix("torch.")
-    if isinstance(value, torch.device):
-        return str(value)
-    if value is None or isinstance(value, (str, int, float, bool)):
-        return value
-    return {
-        "class": f"{value.__class__.__module__}.{value.__class__.__qualname__}",
-        "value": str(value),
-    }
-
-
-def _tokenizer_content_sha256(tokenizer: Any) -> str:
-    content: dict[str, Any] = {
-        "init_kwargs": getattr(tokenizer, "init_kwargs", None),
-        "special_tokens_map": getattr(tokenizer, "special_tokens_map", None),
-        "model_max_length": getattr(tokenizer, "model_max_length", None),
-        "padding_side": getattr(tokenizer, "padding_side", None),
-        "truncation_side": getattr(tokenizer, "truncation_side", None),
-    }
-    get_vocab = getattr(tokenizer, "get_vocab", None)
-    if callable(get_vocab):
-        content["vocabulary"] = get_vocab()
-    get_added_vocab = getattr(tokenizer, "get_added_vocab", None)
-    if callable(get_added_vocab):
-        content["added_vocabulary"] = get_added_vocab()
-    backend = getattr(tokenizer, "backend_tokenizer", None)
-    backend_to_str = getattr(backend, "to_str", None)
-    if callable(backend_to_str):
-        content["backend"] = backend_to_str()
-    serialized = json.dumps(
-        _fingerprint_jsonable(content),
-        sort_keys=True,
-        separators=(",", ":"),
-        ensure_ascii=False,
-    ).encode()
-    return hashlib.sha256(serialized).hexdigest()
-
-
-def _tokenizer_metadata(model: Any, tokenizer: Any | None) -> dict[str, Any]:
-    resolved = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
-    if resolved is None:
-        # Raw-sequence families such as E1 retain their loader context on the
-        # model/encoder rather than exposing a Transformers tokenizer. Bind the
-        # non-secret source policy to resume identity without serializing a Hub
-        # token or forcing lazy tokenizer initialization.
-        for candidate in (model, getattr(model, "model", None)):
-            settings = getattr(candidate, "__dict__", {}).get("_fastplms_tokenizer_kwargs")
-            if isinstance(settings, Mapping):
-                token_value = settings.get("token")
-                return {
-                    "mode": "native-sequence",
-                    "source": (
-                        str(settings.get("tokenizer_source"))
-                        if settings.get("tokenizer_source") is not None
-                        else None
-                    ),
-                    "revision": settings.get("revision"),
-                    "cache_dir": (
-                        str(settings.get("cache_dir"))
-                        if settings.get("cache_dir") is not None
-                        else None
-                    ),
-                    "local_files_only": bool(settings.get("local_files_only", False)),
-                    "token_policy": (
-                        "disabled"
-                        if token_value is False
-                        else "provided"
-                        if token_value is not None
-                        else "default"
-                    ),
-                }
-        return {"mode": "native-sequence"}
-    return {
-        "mode": "tokenizer",
-        "class": f"{resolved.__class__.__module__}.{resolved.__class__.__qualname__}",
-        "name_or_path": getattr(resolved, "name_or_path", None),
-        "vocab_size": getattr(resolved, "vocab_size", None),
-        "special_token_ids": list(getattr(resolved, "all_special_ids", ())),
-        "content_sha256": _tokenizer_content_sha256(resolved),
-    }
-
-
-def _software_versions() -> dict[str, str | None]:
-    try:
-        import fastplms
-
-        fastplms_version = fastplms.__version__
-    except (AttributeError, ImportError):
-        fastplms_version = None
-    try:
-        import safetensors
-
-        safetensors_version = safetensors.__version__
-    except ImportError:
-        safetensors_version = None
-    try:
-        import transformers
-
-        transformers_version = transformers.__version__
-    except ImportError:
-        transformers_version = None
-    return {
-        "fastplms": fastplms_version,
-        "python": platform.python_version(),
-        "safetensors": safetensors_version,
-        "torch": torch.__version__,
-        "torch_cuda": torch.version.cuda,
-        "transformers": transformers_version,
-    }
-
-
-def _adapter_identity_metadata(model: Any) -> dict[str, Any] | None:
-    """Return deterministic PEFT/adapter identity without tensor payloads."""
-
-    peft_config = getattr(model, "peft_config", None)
-    if not isinstance(peft_config, Mapping) or not peft_config:
-        return None
-    configurations: dict[str, Any] = {}
-    for name, config in sorted(peft_config.items(), key=lambda item: str(item[0])):
-        to_dict = getattr(config, "to_dict", None)
-        if callable(to_dict):
-            value = to_dict()
-        else:
-            try:
-                value = vars(config)
-            except TypeError:
-                value = config
-        configurations[str(name)] = _fingerprint_jsonable(value)
-    active_adapters = getattr(model, "active_adapters", None)
-    if callable(active_adapters):
-        active_adapters = active_adapters()
-    return {
-        "active": _fingerprint_jsonable(active_adapters),
-        "configurations": configurations,
-    }
-
-
-def _execution_identity_metadata(model: Any) -> dict[str, Any]:
-    """Capture runtime policy that can change persisted numerical results."""
-
-    parameter_dtypes = sorted(
-        {
-            str(parameter.dtype).removeprefix("torch.")
-            for parameter in getattr(model, "parameters", lambda: ())()
-        }
-    )
-    return {
-        "device": _model_device(model).type,
-        "hf_device_map": _fingerprint_jsonable(getattr(model, "hf_device_map", None)),
-        "parameter_dtypes": parameter_dtypes,
-        "software": _software_versions(),
-    }
-
-
-def _first_metadata_value(*values: Any) -> Any:
-    for value in values:
-        if isinstance(value, str):
-            if value.strip():
-                return value
-        elif value is not None:
-            return value
-    return None
-
-
-def _model_identity_metadata(model: Any) -> dict[str, Any]:
-    """Resolve model and checkpoint identity, including local artifact fallbacks."""
-
-    config = getattr(model, "config", None)
-    checkpoint_revision = _first_metadata_value(
-        getattr(config, "fastplms_checkpoint_revision", None),
-        getattr(config, "_commit_hash", None),
-    )
-    return {
-        "model_id": _first_metadata_value(
-            getattr(config, "fastplms_model_id", None),
-            getattr(config, "_name_or_path", None),
-        ),
-        "model_revision": _first_metadata_value(
-            getattr(config, "_commit_hash", None),
-            checkpoint_revision,
-        ),
-        "checkpoint_repo_id": getattr(config, "fastplms_checkpoint_repo_id", None),
-        "checkpoint_revision": checkpoint_revision,
-        "checkpoint_hash": _first_metadata_value(
-            getattr(model, "checkpoint_hash", None),
-            getattr(config, "checkpoint_hash", None),
-            getattr(config, "fastplms_checkpoint_hash", None),
-        ),
-        "weights_revision": getattr(config, "fastplms_weights_revision", None),
-        "runtime_revision": getattr(config, "fastplms_runtime_revision", None),
-        "source_tree_sha256": getattr(config, "fastplms_source_tree_sha256", None),
-        "runtime_bundle_sha256": getattr(config, "fastplms_runtime_bundle_sha256", None),
-    }
-
-
-def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]:
-    """Yield X in logical row-major order without materializing a full copy."""
-
-    # X: (...)
-    if X.numel() == 0:
-        return
-    if X.ndim == 0:
-        yield X
-        return
-    trailing_elements = 1
-    for size in X.shape[1:]:
-        trailing_elements *= int(size)
-    if trailing_elements <= max_elements:
-        rows_per_chunk = max(1, max_elements // trailing_elements)
-        for start in range(0, X.shape[0], rows_per_chunk):
-            yield X[start : start + rows_per_chunk]  # (chunk_rows, ...)
-        return
-    for row in X:
-        yield from _bounded_tensor_chunks(row, max_elements)
-
-
-def _model_state_sha256(model: Any) -> str:
-    """Hash named parameters and persistent buffers using bounded CPU copies."""
-
-    # Never cache this digest from tensor identity or ``Tensor._version``.
-    # ``Parameter.data`` and independent tensor aliases can mutate shared storage
-    # without changing either signal, while persisted resume identity must bind
-    # the authoritative bytes visible at the start of this run.
-    state = model.state_dict(keep_vars=True)
-    digest = hashlib.sha256()
-    for name, value in sorted(state.items()):
-        if not isinstance(value, Tensor):
-            raise TypeError(f"Model state entry {name!r} is not a tensor.")
-        if value.is_meta:
-            raise ValueError(
-                f"Cannot fingerprint meta-device model state entry {name!r}; pass "
-                "model_state_fingerprint with a caller-owned state identity."
-            )
-        header = json.dumps(
-            {
-                "name": name,
-                "dtype": str(value.dtype).removeprefix("torch."),
-                "shape": list(value.shape),
-            },
-            sort_keys=True,
-            separators=(",", ":"),
-        ).encode()
-        digest.update(len(header).to_bytes(8, "big"))
-        digest.update(header)
-        max_elements = max(1, _MODEL_STATE_HASH_CHUNK_BYTES // value.element_size())
-        for chunk in _bounded_tensor_chunks(value.detach(), max_elements):
-            cpu_chunk = chunk.to(device="cpu").contiguous()  # chunk.shape
-            digest.update(cpu_chunk.reshape(-1).view(torch.uint8).numpy().tobytes())
-    return digest.hexdigest()
-
-
-def _input_sha256(records: Iterable[EmbeddingInput]) -> str:
-    """Hash an ordered input stream without constructing a duplicate JSON payload."""
-
-    precomputed = getattr(records, "input_fingerprint", None)
-    if isinstance(precomputed, str):
-        return precomputed
-    digest = hashlib.sha256()
-    count = 0
-    for record in records:
-        count += 1
-        for value in (record.id, record.sequence):
-            encoded = value.encode("utf-8")
-            digest.update(len(encoded).to_bytes(8, "big"))
-            digest.update(encoded)
-    digest.update(count.to_bytes(8, "big"))
-    return digest.hexdigest()
-
-
-def _run_fingerprint(
-    model: Any,
-    records: Sequence[EmbeddingInput],
-    *,
-    pooling: Sequence[str],
-    full_embeddings: bool,
-    max_length: int | None,
-    truncate: bool,
-    dtype: torch.dtype | None,
-    model_kwargs: dict[str, Any],
-    tokenizer_metadata: dict[str, Any],
-    model_state_fingerprint: str | None,
-    persist_output: bool,
-    embedding_context: Mapping[str, Any],
-    batch_size: int,
-    batch_window_size: int,
-    max_tokens_per_batch: int | None,
-) -> tuple[str, str, str | None, str]:
-    input_fingerprint = _input_sha256(records)
-    attention_backend = _attention_backend(model)
-    model_identity = _model_identity_metadata(model)
-    if model_state_fingerprint is None and persist_output:
-        resolved_model_state_fingerprint = _model_state_sha256(model)
-        model_state_fingerprint_source = "computed"
-    elif model_state_fingerprint is not None:
-        resolved_model_state_fingerprint = model_state_fingerprint.strip()
-        if not resolved_model_state_fingerprint:
-            raise ValueError("model_state_fingerprint must not be empty.")
-        model_state_fingerprint_source = "caller"
-    else:
-        resolved_model_state_fingerprint = None
-        model_state_fingerprint_source = "not-computed"
-    payload = {
-        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-        "input_fingerprint": input_fingerprint,
-        "model_state_fingerprint": resolved_model_state_fingerprint,
-        "model_state_fingerprint_source": model_state_fingerprint_source,
-        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
-        **model_identity,
-        "attention_backend": attention_backend,
-        "attention_kernel": _attention_kernel_metadata(attention_backend),
-        "layer": repr(
-            getattr(model, "embedding_layer", model_kwargs.get("hidden_state_index", -1))
-        ),
-        "projection": getattr(model, "embedding_projection", None),
-        "esmc_source": getattr(model, "_esmc_source", None),
-        "esmc_revision": getattr(model, "_esmc_source_revision", None),
-        "esmc_files": getattr(model, "_esmc_source_files", None),
-        "token_policy": getattr(model, "embedding_token_policy", None),
-        "tokenizer": tokenizer_metadata,
-        "adapter": _adapter_identity_metadata(model),
-        "execution": _execution_identity_metadata(model),
-        "embedding_context": _fingerprint_jsonable(embedding_context),
-        "pooling": list(pooling),
-        "full_embeddings": full_embeddings,
-        "max_length": max_length,
-        "truncate": truncate,
-        "dtype": str(dtype) if dtype is not None else None,
-        "batching": {
-            "batch_size": batch_size,
-            "batch_window_size": batch_window_size,
-            "max_tokens_per_batch": max_tokens_per_batch,
-            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
-        },
-        "model_kwargs": {
-            key: _fingerprint_jsonable(value) for key, value in sorted(model_kwargs.items())
-        },
-        "residue_mask_policy": "attention-mask-minus-special-tokens",
-    }
-    run_fingerprint = hashlib.sha256(
-        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
-    ).hexdigest()
-    return (
-        input_fingerprint,
-        run_fingerprint,
-        resolved_model_state_fingerprint,
-        model_state_fingerprint_source,
-    )
-
-
-def _ordered_string_sha256(values: Sequence[str]) -> str:
-    digest = hashlib.sha256()
-    for value in values:
-        encoded = value.encode("utf-8")
-        digest.update(len(encoded).to_bytes(8, "big"))
-        digest.update(encoded)
-    digest.update(len(values).to_bytes(8, "big"))
-    return digest.hexdigest()
-
-
-def _embedding_context(
-    model: Any,
-    records: Sequence[EmbeddingInput],
-    *,
-    hidden_state_source: str,
-    decoder_inputs: Sequence[str] | None,
-    decoder_input_ids: Tensor | None,
-    decoder_attention_mask: Tensor | None,
-    model_kwargs: Mapping[str, Any],
-) -> tuple[dict[str, Any], tuple[str, ...] | None]:
-    if hidden_state_source not in {"encoder", "decoder"}:
-        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
-    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
-    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
-        raise TypeError("hidden_state_index must be an integer.")
-    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
-    if not isinstance(store_all_hidden_states, bool):
-        raise TypeError("store_all_hidden_states must be a boolean.")
-    normalized_decoder_inputs: tuple[str, ...] | None = None
-    has_decoder_inputs = decoder_inputs is not None
-    has_decoder_ids = decoder_input_ids is not None
-    if hidden_state_source == "encoder":
-        if has_decoder_inputs or has_decoder_ids or decoder_attention_mask is not None:
-            raise ValueError("Decoder inputs are only valid when hidden_state_source='decoder'.")
-    else:
-        if has_decoder_inputs == has_decoder_ids:
-            raise ValueError(
-                "Decoder embedding requires exactly one of decoder_inputs or decoder_input_ids."
-            )
-    decoder_input_fingerprint: str | None = None
-    if decoder_inputs is not None:
-        if isinstance(decoder_inputs, (str, bytes)) or not isinstance(decoder_inputs, Sequence):
-            raise TypeError("decoder_inputs must be an aligned sequence of strings.")
-        normalized_decoder_inputs = tuple(decoder_inputs)
-        if not all(isinstance(value, str) and value for value in normalized_decoder_inputs):
-            raise ValueError("decoder_inputs must contain non-empty strings.")
-        if len(normalized_decoder_inputs) != len(records):
-            raise ValueError("decoder_inputs must align one-to-one with embedding inputs.")
-        decoder_input_fingerprint = _ordered_string_sha256(normalized_decoder_inputs)
-        if decoder_attention_mask is not None:
-            raise ValueError("decoder_attention_mask requires decoder_input_ids.")
-    if decoder_input_ids is not None:
-        if not isinstance(decoder_input_ids, Tensor) or decoder_input_ids.ndim != 2:
-            raise ValueError("decoder_input_ids must have shape (batch, sequence).")
-        if decoder_input_ids.shape[0] != len(records):
-            raise ValueError("decoder_input_ids must align one-to-one with embedding inputs.")
-        if decoder_input_ids.dtype == torch.bool or decoder_input_ids.is_floating_point():
-            raise TypeError("decoder_input_ids must use an integer token dtype.")
-        decoder_input_fingerprint = tensor_sha256(decoder_input_ids)
-    decoder_mask_fingerprint: str | None = None
-    if decoder_attention_mask is not None:
-        if not isinstance(decoder_attention_mask, Tensor):
-            raise TypeError("decoder_attention_mask must be a tensor.")
-        if decoder_input_ids is None or decoder_attention_mask.shape != decoder_input_ids.shape:
-            raise ValueError("decoder_attention_mask must match decoder_input_ids shape.")
-        decoder_mask_fingerprint = tensor_sha256(decoder_attention_mask)
-
-    context: dict[str, Any] = {
-        "hidden_state_source": hidden_state_source,
-        "hidden_state_index": hidden_state_index,
-        "store_all_hidden_states": store_all_hidden_states,
-        "decoder_input_fingerprint": decoder_input_fingerprint,
-        "decoder_attention_mask_fingerprint": decoder_mask_fingerprint,
-        "decoder_alignment": "input-position" if hidden_state_source == "decoder" else None,
-    }
-    metadata_hook = getattr(model, "_embedding_metadata", None)
-    model_metadata: Mapping[str, Any] | None = None
-    if callable(metadata_hook):
-        model_metadata = metadata_hook(**context)
-        if not isinstance(model_metadata, Mapping):
-            raise TypeError("_embedding_metadata must return a mapping.")
-        context["model_embedding"] = _fingerprint_jsonable(model_metadata)
-    if hidden_state_source == "decoder":
-        has_decoder_batch = callable(getattr(model, "_embedding_batch", None))
-        declares_decoder_stack = (
-            model_metadata is not None and model_metadata.get("hidden_state_stack") == "decoder"
-        )
-        if not has_decoder_batch or not declares_decoder_stack:
-            raise ValueError(
-                f"{model.__class__.__name__} does not declare decoder embedding support."
-            )
-    return context, normalized_decoder_inputs
+
+
+_RUN_FINGERPRINT_SCHEMA_VERSION = 3
+_MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2
+
+
+def _model_device(model: Any) -> torch.device:
+    try:
+        return torch.device(next(model.parameters()).device)
+    except (AttributeError, StopIteration):
+        return torch.device("cpu")
+
+
+def _attention_backend(model: Any) -> str | None:
+    config = getattr(model, "config", None)
+    for name in ("_attn_implementation", "attn_implementation", "attn_backend"):
+        value = getattr(config, name, None)
+        if value:
+            return str(value)
+    return None
+
+
+def _attention_kernel_metadata(backend: str | None) -> dict[str, Any] | None:
+    if backend not in {"flash_attention_2", "flash_attention_3"}:
+        return None
+    from fastplms.registry import get_model_registry
+
+    spec = get_model_registry().attention_kernels[backend]
+    return {
+        "repository": spec.repository,
+        "revision": spec.revision,
+        "version": spec.version,
+        "expected_variant": spec.expected_variant,
+        "dtypes": list(spec.dtypes),
+    }
+
+
+def _fingerprint_jsonable(value: Any) -> Any:
+    if isinstance(value, Mapping):
+        return {str(key): _fingerprint_jsonable(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_fingerprint_jsonable(item) for item in value]
+    if isinstance(value, (set, frozenset)):
+        return sorted((_fingerprint_jsonable(item) for item in value), key=repr)
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, Tensor):
+        return {
+            "dtype": str(value.dtype).removeprefix("torch."),
+            "shape": list(value.shape),
+            "sha256": tensor_sha256(value),
+        }
+    if isinstance(value, torch.dtype):
+        return str(value).removeprefix("torch.")
+    if isinstance(value, torch.device):
+        return str(value)
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    return {
+        "class": f"{value.__class__.__module__}.{value.__class__.__qualname__}",
+        "value": str(value),
+    }
+
+
+def _tokenizer_content_sha256(tokenizer: Any) -> str:
+    content: dict[str, Any] = {
+        "init_kwargs": getattr(tokenizer, "init_kwargs", None),
+        "special_tokens_map": getattr(tokenizer, "special_tokens_map", None),
+        "model_max_length": getattr(tokenizer, "model_max_length", None),
+        "padding_side": getattr(tokenizer, "padding_side", None),
+        "truncation_side": getattr(tokenizer, "truncation_side", None),
+    }
+    get_vocab = getattr(tokenizer, "get_vocab", None)
+    if callable(get_vocab):
+        content["vocabulary"] = get_vocab()
+    get_added_vocab = getattr(tokenizer, "get_added_vocab", None)
+    if callable(get_added_vocab):
+        content["added_vocabulary"] = get_added_vocab()
+    backend = getattr(tokenizer, "backend_tokenizer", None)
+    backend_to_str = getattr(backend, "to_str", None)
+    if callable(backend_to_str):
+        content["backend"] = backend_to_str()
+    serialized = json.dumps(
+        _fingerprint_jsonable(content),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+    ).encode()
+    return hashlib.sha256(serialized).hexdigest()
+
+
+def _tokenizer_metadata(model: Any, tokenizer: Any | None) -> dict[str, Any]:
+    resolved = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
+    if resolved is None:
+        # Raw-sequence families such as E1 retain their loader context on the
+        # model/encoder rather than exposing a Transformers tokenizer. Bind the
+        # non-secret source policy to resume identity without serializing a Hub
+        # token or forcing lazy tokenizer initialization.
+        for candidate in (model, getattr(model, "model", None)):
+            settings = getattr(candidate, "__dict__", {}).get("_fastplms_tokenizer_kwargs")
+            if isinstance(settings, Mapping):
+                token_value = settings.get("token")
+                return {
+                    "mode": "native-sequence",
+                    "source": (
+                        str(settings.get("tokenizer_source"))
+                        if settings.get("tokenizer_source") is not None
+                        else None
+                    ),
+                    "revision": settings.get("revision"),
+                    "cache_dir": (
+                        str(settings.get("cache_dir"))
+                        if settings.get("cache_dir") is not None
+                        else None
+                    ),
+                    "local_files_only": bool(settings.get("local_files_only", False)),
+                    "token_policy": (
+                        "disabled"
+                        if token_value is False
+                        else "provided"
+                        if token_value is not None
+                        else "default"
+                    ),
+                }
+        return {"mode": "native-sequence"}
+    return {
+        "mode": "tokenizer",
+        "class": f"{resolved.__class__.__module__}.{resolved.__class__.__qualname__}",
+        "name_or_path": getattr(resolved, "name_or_path", None),
+        "vocab_size": getattr(resolved, "vocab_size", None),
+        "special_token_ids": list(getattr(resolved, "all_special_ids", ())),
+        "content_sha256": _tokenizer_content_sha256(resolved),
+    }
+
+
+def _software_versions() -> dict[str, str | None]:
+    try:
+        import fastplms
+
+        fastplms_version = fastplms.__version__
+    except (AttributeError, ImportError):
+        fastplms_version = None
+    try:
+        import safetensors
+
+        safetensors_version = safetensors.__version__
+    except ImportError:
+        safetensors_version = None
+    try:
+        import transformers
+
+        transformers_version = transformers.__version__
+    except ImportError:
+        transformers_version = None
+    return {
+        "fastplms": fastplms_version,
+        "python": platform.python_version(),
+        "safetensors": safetensors_version,
+        "torch": torch.__version__,
+        "torch_cuda": torch.version.cuda,
+        "transformers": transformers_version,
+    }
+
+
+def _adapter_identity_metadata(model: Any) -> dict[str, Any] | None:
+    """Return deterministic PEFT/adapter identity without tensor payloads."""
+
+    peft_config = getattr(model, "peft_config", None)
+    if not isinstance(peft_config, Mapping) or not peft_config:
+        return None
+    configurations: dict[str, Any] = {}
+    for name, config in sorted(peft_config.items(), key=lambda item: str(item[0])):
+        to_dict = getattr(config, "to_dict", None)
+        if callable(to_dict):
+            value = to_dict()
+        else:
+            try:
+                value = vars(config)
+            except TypeError:
+                value = config
+        configurations[str(name)] = _fingerprint_jsonable(value)
+    active_adapters = getattr(model, "active_adapters", None)
+    if callable(active_adapters):
+        active_adapters = active_adapters()
+    return {
+        "active": _fingerprint_jsonable(active_adapters),
+        "configurations": configurations,
+    }
+
+
+def _execution_identity_metadata(model: Any) -> dict[str, Any]:
+    """Capture runtime policy that can change persisted numerical results."""
+
+    parameter_dtypes = sorted(
+        {
+            str(parameter.dtype).removeprefix("torch.")
+            for parameter in getattr(model, "parameters", lambda: ())()
+        }
+    )
+    return {
+        "device": _model_device(model).type,
+        "hf_device_map": _fingerprint_jsonable(getattr(model, "hf_device_map", None)),
+        "parameter_dtypes": parameter_dtypes,
+        "software": _software_versions(),
+    }
+
+
+def _first_metadata_value(*values: Any) -> Any:
+    for value in values:
+        if isinstance(value, str):
+            if value.strip():
+                return value
+        elif value is not None:
+            return value
+    return None
+
+
+def _model_identity_metadata(model: Any) -> dict[str, Any]:
+    """Resolve model and checkpoint identity, including local artifact fallbacks."""
+
+    config = getattr(model, "config", None)
+    checkpoint_revision = _first_metadata_value(
+        getattr(config, "fastplms_checkpoint_revision", None),
+        getattr(config, "_commit_hash", None),
+    )
+    return {
+        "model_id": _first_metadata_value(
+            getattr(config, "fastplms_model_id", None),
+            getattr(config, "_name_or_path", None),
+        ),
+        "model_revision": _first_metadata_value(
+            getattr(config, "_commit_hash", None),
+            checkpoint_revision,
+        ),
+        "checkpoint_repo_id": getattr(config, "fastplms_checkpoint_repo_id", None),
+        "checkpoint_revision": checkpoint_revision,
+        "checkpoint_hash": _first_metadata_value(
+            getattr(model, "checkpoint_hash", None),
+            getattr(config, "checkpoint_hash", None),
+            getattr(config, "fastplms_checkpoint_hash", None),
+        ),
+        "weights_revision": getattr(config, "fastplms_weights_revision", None),
+        "runtime_revision": getattr(config, "fastplms_runtime_revision", None),
+        "source_tree_sha256": getattr(config, "fastplms_source_tree_sha256", None),
+        "runtime_bundle_sha256": getattr(config, "fastplms_runtime_bundle_sha256", None),
+    }
+
+
+def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]:
+    """Yield X in logical row-major order without materializing a full copy."""
+
+    # X: (...)
+    if X.numel() == 0:
+        return
+    if X.ndim == 0:
+        yield X
+        return
+    trailing_elements = 1
+    for size in X.shape[1:]:
+        trailing_elements *= int(size)
+    if trailing_elements <= max_elements:
+        rows_per_chunk = max(1, max_elements // trailing_elements)
+        for start in range(0, X.shape[0], rows_per_chunk):
+            yield X[start : start + rows_per_chunk]  # (chunk_rows, ...)
+        return
+    for row in X:
+        yield from _bounded_tensor_chunks(row, max_elements)
+
+
+def _model_state_sha256(model: Any) -> str:
+    """Hash named parameters and persistent buffers using bounded CPU copies."""
+
+    # Never cache this digest from tensor identity or ``Tensor._version``.
+    # ``Parameter.data`` and independent tensor aliases can mutate shared storage
+    # without changing either signal, while persisted resume identity must bind
+    # the authoritative bytes visible at the start of this run.
+    state = model.state_dict(keep_vars=True)
+    digest = hashlib.sha256()
+    for name, value in sorted(state.items()):
+        if not isinstance(value, Tensor):
+            raise TypeError(f"Model state entry {name!r} is not a tensor.")
+        if value.is_meta:
+            raise ValueError(
+                f"Cannot fingerprint meta-device model state entry {name!r}; pass "
+                "model_state_fingerprint with a caller-owned state identity."
+            )
+        header = json.dumps(
+            {
+                "name": name,
+                "dtype": str(value.dtype).removeprefix("torch."),
+                "shape": list(value.shape),
+            },
+            sort_keys=True,
+            separators=(",", ":"),
+        ).encode()
+        digest.update(len(header).to_bytes(8, "big"))
+        digest.update(header)
+        max_elements = max(1, _MODEL_STATE_HASH_CHUNK_BYTES // value.element_size())
+        for chunk in _bounded_tensor_chunks(value.detach(), max_elements):
+            cpu_chunk = chunk.to(device="cpu").contiguous()  # chunk.shape
+            digest.update(cpu_chunk.reshape(-1).view(torch.uint8).numpy().tobytes())
+    return digest.hexdigest()
+
+
+def _input_sha256(records: Iterable[EmbeddingInput]) -> str:
+    """Hash an ordered input stream without constructing a duplicate JSON payload."""
+
+    precomputed = getattr(records, "input_fingerprint", None)
+    if isinstance(precomputed, str):
+        return precomputed
+    digest = hashlib.sha256()
+    count = 0
+    for record in records:
+        count += 1
+        for value in (record.id, record.sequence):
+            encoded = value.encode("utf-8")
+            digest.update(len(encoded).to_bytes(8, "big"))
+            digest.update(encoded)
+    digest.update(count.to_bytes(8, "big"))
+    return digest.hexdigest()
+
+
+def _run_fingerprint(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    *,
+    pooling: Sequence[str],
+    full_embeddings: bool,
+    max_length: int | None,
+    truncate: bool,
+    dtype: torch.dtype | None,
+    model_kwargs: dict[str, Any],
+    tokenizer_metadata: dict[str, Any],
+    model_state_fingerprint: str | None,
+    persist_output: bool,
+    embedding_context: Mapping[str, Any],
+    batch_size: int,
+    batch_window_size: int,
+    max_tokens_per_batch: int | None,
+) -> tuple[str, str, str | None, str]:
+    input_fingerprint = _input_sha256(records)
+    attention_backend = _attention_backend(model)
+    model_identity = _model_identity_metadata(model)
+    if model_state_fingerprint is None and persist_output:
+        resolved_model_state_fingerprint = _model_state_sha256(model)
+        model_state_fingerprint_source = "computed"
+    elif model_state_fingerprint is not None:
+        resolved_model_state_fingerprint = model_state_fingerprint.strip()
+        if not resolved_model_state_fingerprint:
+            raise ValueError("model_state_fingerprint must not be empty.")
+        model_state_fingerprint_source = "caller"
+    else:
+        resolved_model_state_fingerprint = None
+        model_state_fingerprint_source = "not-computed"
+    payload = {
+        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+        "input_fingerprint": input_fingerprint,
+        "model_state_fingerprint": resolved_model_state_fingerprint,
+        "model_state_fingerprint_source": model_state_fingerprint_source,
+        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
+        **model_identity,
+        "attention_backend": attention_backend,
+        "attention_kernel": _attention_kernel_metadata(attention_backend),
+        "layer": repr(
+            getattr(model, "embedding_layer", model_kwargs.get("hidden_state_index", -1))
+        ),
+        "projection": getattr(model, "embedding_projection", None),
+        "esmc_source": getattr(model, "_esmc_source", None),
+        "esmc_revision": getattr(model, "_esmc_source_revision", None),
+        "esmc_files": getattr(model, "_esmc_source_files", None),
+        "token_policy": getattr(model, "embedding_token_policy", None),
+        "tokenizer": tokenizer_metadata,
+        "adapter": _adapter_identity_metadata(model),
+        "execution": _execution_identity_metadata(model),
+        "embedding_context": _fingerprint_jsonable(embedding_context),
+        "pooling": list(pooling),
+        "full_embeddings": full_embeddings,
+        "max_length": max_length,
+        "truncate": truncate,
+        "dtype": str(dtype) if dtype is not None else None,
+        "batching": {
+            "batch_size": batch_size,
+            "batch_window_size": batch_window_size,
+            "max_tokens_per_batch": max_tokens_per_batch,
+            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
+        },
+        "model_kwargs": {
+            key: _fingerprint_jsonable(value) for key, value in sorted(model_kwargs.items())
+        },
+        "residue_mask_policy": "attention-mask-minus-special-tokens",
+    }
+    run_fingerprint = hashlib.sha256(
+        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
+    ).hexdigest()
+    return (
+        input_fingerprint,
+        run_fingerprint,
+        resolved_model_state_fingerprint,
+        model_state_fingerprint_source,
+    )
+
+
+def _ordered_string_sha256(values: Sequence[str]) -> str:
+    digest = hashlib.sha256()
+    for value in values:
+        encoded = value.encode("utf-8")
+        digest.update(len(encoded).to_bytes(8, "big"))
+        digest.update(encoded)
+    digest.update(len(values).to_bytes(8, "big"))
+    return digest.hexdigest()
+
+
+def _embedding_context(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    *,
+    hidden_state_source: str,
+    decoder_inputs: Sequence[str] | None,
+    decoder_input_ids: Tensor | None,
+    decoder_attention_mask: Tensor | None,
+    model_kwargs: Mapping[str, Any],
+) -> tuple[dict[str, Any], tuple[str, ...] | None]:
+    if hidden_state_source not in {"encoder", "decoder"}:
+        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
+    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
+    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
+        raise TypeError("hidden_state_index must be an integer.")
+    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
+    if not isinstance(store_all_hidden_states, bool):
+        raise TypeError("store_all_hidden_states must be a boolean.")
+    normalized_decoder_inputs: tuple[str, ...] | None = None
+    has_decoder_inputs = decoder_inputs is not None
+    has_decoder_ids = decoder_input_ids is not None
+    if hidden_state_source == "encoder":
+        if has_decoder_inputs or has_decoder_ids or decoder_attention_mask is not None:
+            raise ValueError("Decoder inputs are only valid when hidden_state_source='decoder'.")
+    else:
+        if has_decoder_inputs == has_decoder_ids:
+            raise ValueError(
+                "Decoder embedding requires exactly one of decoder_inputs or decoder_input_ids."
+            )
+    decoder_input_fingerprint: str | None = None
+    if decoder_inputs is not None:
+        if isinstance(decoder_inputs, (str, bytes)) or not isinstance(decoder_inputs, Sequence):
+            raise TypeError("decoder_inputs must be an aligned sequence of strings.")
+        normalized_decoder_inputs = tuple(decoder_inputs)
+        if not all(isinstance(value, str) and value for value in normalized_decoder_inputs):
+            raise ValueError("decoder_inputs must contain non-empty strings.")
+        if len(normalized_decoder_inputs) != len(records):
+            raise ValueError("decoder_inputs must align one-to-one with embedding inputs.")
+        decoder_input_fingerprint = _ordered_string_sha256(normalized_decoder_inputs)
+        if decoder_attention_mask is not None:
+            raise ValueError("decoder_attention_mask requires decoder_input_ids.")
+    if decoder_input_ids is not None:
+        if not isinstance(decoder_input_ids, Tensor) or decoder_input_ids.ndim != 2:
+            raise ValueError("decoder_input_ids must have shape (batch, sequence).")
+        if decoder_input_ids.shape[0] != len(records):
+            raise ValueError("decoder_input_ids must align one-to-one with embedding inputs.")
+        if decoder_input_ids.dtype == torch.bool or decoder_input_ids.is_floating_point():
+            raise TypeError("decoder_input_ids must use an integer token dtype.")
+        decoder_input_fingerprint = tensor_sha256(decoder_input_ids)
+    decoder_mask_fingerprint: str | None = None
+    if decoder_attention_mask is not None:
+        if not isinstance(decoder_attention_mask, Tensor):
+            raise TypeError("decoder_attention_mask must be a tensor.")
+        if decoder_input_ids is None or decoder_attention_mask.shape != decoder_input_ids.shape:
+            raise ValueError("decoder_attention_mask must match decoder_input_ids shape.")
+        decoder_mask_fingerprint = tensor_sha256(decoder_attention_mask)
+
+    context: dict[str, Any] = {
+        "hidden_state_source": hidden_state_source,
+        "hidden_state_index": hidden_state_index,
+        "store_all_hidden_states": store_all_hidden_states,
+        "decoder_input_fingerprint": decoder_input_fingerprint,
+        "decoder_attention_mask_fingerprint": decoder_mask_fingerprint,
+        "decoder_alignment": "input-position" if hidden_state_source == "decoder" else None,
+    }
+    metadata_hook = getattr(model, "_embedding_metadata", None)
+    model_metadata: Mapping[str, Any] | None = None
+    if callable(metadata_hook):
+        model_metadata = metadata_hook(**context)
+        if not isinstance(model_metadata, Mapping):
+            raise TypeError("_embedding_metadata must return a mapping.")
+        context["model_embedding"] = _fingerprint_jsonable(model_metadata)
+    if hidden_state_source == "decoder":
+        has_decoder_batch = callable(getattr(model, "_embedding_batch", None))
+        declares_decoder_stack = (
+            model_metadata is not None and model_metadata.get("hidden_state_stack") == "decoder"
+        )
+        if not has_decoder_batch or not declares_decoder_stack:
+            raise ValueError(
+                f"{model.__class__.__name__} does not declare decoder embedding support."
+            )
+    return context, normalized_decoder_inputs

fastplms/embeddings/inputs.py

@@ -1,263 +1,264 @@
-"""Normalize ordered inputs and plan bounded windows without retaining a full stream."""
-
-from __future__ import annotations
-
-import hashlib
-import sqlite3
-import tempfile
-from collections.abc import Iterable, Iterator, Mapping, Sequence
-from pathlib import Path
-from typing import overload
-
-from .types import EmbeddingInput
-
-
-def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]:
-    """Yield FASTA records in source order without reading the file into memory."""
-
-    identifier: str | None = None
-    sequence_parts: list[str] = []
-    found_record = False
-    with Path(path).open("r", encoding="utf-8") as handle:
-        for line_number, raw_line in enumerate(handle, start=1):
-            line = raw_line.strip()
-            if not line:
-                continue
-            if line.startswith(">"):
-                if identifier is not None:
-                    found_record = True
-                    yield EmbeddingInput(identifier, "".join(sequence_parts))
-                identifier = line[1:].strip().split(maxsplit=1)[0]
-                if not identifier:
-                    raise ValueError(f"Missing FASTA identifier on line {line_number}.")
-                sequence_parts = []
-            else:
-                if identifier is None:
-                    raise ValueError(
-                        f"Sequence data precedes the first FASTA header on line {line_number}."
-                    )
-                sequence_parts.append("".join(line.split()))
-    if identifier is not None:
-        found_record = True
-        yield EmbeddingInput(identifier, "".join(sequence_parts))
-    if not found_record:
-        raise ValueError(f"No FASTA records found in {path}.")
-
-
-def parse_fasta(path: str | Path) -> list[EmbeddingInput]:
-    """Parse FASTA records while preserving identifiers, order, and duplicates."""
-
-    return list(iter_fasta(path))
-
-
-def _normalize_input_item(
-    position: int,
-    item: str | EmbeddingInput | tuple[str, str],
-) -> EmbeddingInput:
-    if isinstance(item, EmbeddingInput):
-        return item
-    if isinstance(item, str):
-        return EmbeddingInput(str(position), item)
-    if isinstance(item, tuple) and len(item) == 2:
-        return EmbeddingInput(str(item[0]), str(item[1]))
-    raise TypeError(
-        "inputs must contain sequences, EmbeddingInput values, or (id, sequence) tuples."
-    )
-
-
-class _InputSpool(Sequence[EmbeddingInput]):
-    """Immutable disk-backed normalized inputs with an incremental digest."""
-
-    def __init__(
-        self,
-        values: Iterable[str | EmbeddingInput | tuple[str, str]],
-    ) -> None:
-        self._temporary: tempfile.TemporaryDirectory[str] | None = tempfile.TemporaryDirectory(
-            prefix="fastplms-inputs-"
-        )
-        self.path = Path(self._temporary.name) / "inputs.sqlite"
-        self._connection: sqlite3.Connection | None = sqlite3.connect(self.path)
-        self._connection.execute(
-            "CREATE TABLE inputs ("
-            "position INTEGER PRIMARY KEY, input_id TEXT NOT NULL, sequence TEXT NOT NULL)"
-        )
-        digest = hashlib.sha256()
-        count = 0
-        pending: list[tuple[int, str, str]] = []
-        try:
-            for position, item in enumerate(values):
-                record = _normalize_input_item(position, item)
-                for value in (record.id, record.sequence):
-                    encoded = value.encode("utf-8")
-                    digest.update(len(encoded).to_bytes(8, "big"))
-                    digest.update(encoded)
-                pending.append((position, record.id, record.sequence))
-                count += 1
-                if len(pending) == 1_024:
-                    self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
-                    pending.clear()
-            if pending:
-                self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
-            if count == 0:
-                raise ValueError("inputs must contain at least one sequence.")
-            self._connection.commit()
-            self._connection.close()
-            self._connection = sqlite3.connect(
-                f"{self.path.resolve().as_uri()}?mode=ro",
-                uri=True,
-            )
-        except BaseException:
-            self.close()
-            raise
-        digest.update(count.to_bytes(8, "big"))
-        self.input_fingerprint = digest.hexdigest()
-        self._count = count
-
-    def _require_connection(self) -> sqlite3.Connection:
-        if self._connection is None:
-            raise RuntimeError("Input spool is closed.")
-        return self._connection
-
-    def __len__(self) -> int:
-        return self._count
-
-    def __iter__(self) -> Iterator[EmbeddingInput]:
-        cursor = self._require_connection().execute(
-            "SELECT input_id, sequence FROM inputs ORDER BY position"
-        )
-        while rows := cursor.fetchmany(1_024):
-            for input_id, sequence in rows:
-                yield EmbeddingInput(input_id, sequence)
-
-    @overload
-    def __getitem__(self, index: int, /) -> EmbeddingInput: ...
-
-    @overload
-    def __getitem__(self, index: slice, /) -> list[EmbeddingInput]: ...
-
-    def __getitem__(self, index: int | slice) -> EmbeddingInput | list[EmbeddingInput]:
-        connection = self._require_connection()
-
-        if isinstance(index, slice):
-            start, stop, step = index.indices(self._count)
-            if step != 1:
-                return [self[position] for position in range(start, stop, step)]
-            rows = connection.execute(
-                "SELECT input_id, sequence FROM inputs "
-                "WHERE position >= ? AND position < ? ORDER BY position",
-                (start, stop),
-            ).fetchall()
-            return [EmbeddingInput(input_id, sequence) for input_id, sequence in rows]
-        position = index + self._count if index < 0 else index
-        if position < 0 or position >= self._count:
-            raise IndexError(index)
-        row = connection.execute(
-            "SELECT input_id, sequence FROM inputs WHERE position = ?", (position,)
-        ).fetchone()
-        if row is None:
-            raise IndexError(index)
-        return EmbeddingInput(row[0], row[1])
-
-    def close(self) -> None:
-        connection = getattr(self, "_connection", None)
-        if connection is not None:
-            connection.close()
-            self._connection = None
-        temporary = getattr(self, "_temporary", None)
-        if temporary is not None:
-            temporary.cleanup()
-            self._temporary = None
-
-    def __del__(self) -> None:
-        self.close()
-
-
-def _normalize_inputs(
-    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
-    *,
-    disk_backed: bool,
-) -> Sequence[EmbeddingInput]:
-    is_fasta_path = isinstance(inputs, Path)
-    if isinstance(inputs, str):
-        try:
-            is_fasta_path = Path(inputs).is_file()
-        except OSError:
-            is_fasta_path = False
-    should_spool = disk_backed or is_fasta_path or not isinstance(inputs, (str, Sequence, Mapping))
-    values: Iterable[str | EmbeddingInput | tuple[str, str]]
-    if isinstance(inputs, Path):
-        values = iter_fasta(inputs)
-    elif isinstance(inputs, str):
-        values = iter_fasta(inputs) if is_fasta_path else [inputs]
-    elif isinstance(inputs, Mapping):
-        values = inputs.items()
-    else:
-        values = inputs
-    if should_spool:
-        return _InputSpool(values)
-    records: list[EmbeddingInput] = []
-    for position, item in enumerate(values):
-        records.append(_normalize_input_item(position, item))
-    if not records:
-        raise ValueError("inputs must contain at least one sequence.")
-    return records
-
-
-def _validate_untruncated_lengths(
-    records: Sequence[EmbeddingInput],
-    *,
-    max_length: int | None,
-    truncate: bool,
-) -> None:
-    """Fail before inference when a biological-residue limit would be exceeded."""
-
-    if max_length is None or truncate:
-        return
-    for position, record in enumerate(records):
-        residue_count = len(record.sequence)
-        if residue_count > max_length:
-            raise ValueError(
-                f"Input at position {position} with id {record.id!r} has "
-                f"{residue_count} biological residues, exceeding max_length={max_length} "
-                "while truncate=False."
-            )
-
-
-def _planned_batches(
-    records: Sequence[EmbeddingInput],
-    positions: range,
-    *,
-    batch_size: int,
-    max_tokens_per_batch: int | None,
-    max_length: int | None,
-    truncate: bool,
-) -> Iterator[list[int]]:
-    """Length-bucket one bounded window while retaining stable output positions."""
-
-    def effective_length(position: int) -> int:
-        length = len(records[position].sequence)
-        return min(length, max_length) if truncate and max_length is not None else length
-
-    ordered = sorted(positions, key=lambda position: (-effective_length(position), position))
-    batch: list[int] = []
-    longest = 0
-    for position in ordered:
-        length = effective_length(position)
-        if max_tokens_per_batch is not None and length > max_tokens_per_batch:
-            raise ValueError(
-                f"Input at position {position} has {length} residues, exceeding "
-                f"max_tokens_per_batch={max_tokens_per_batch}."
-            )
-        candidate_longest = max(longest, length)
-        exceeds_tokens = (
-            max_tokens_per_batch is not None
-            and candidate_longest * (len(batch) + 1) > max_tokens_per_batch
-        )
-        if batch and (len(batch) >= batch_size or exceeds_tokens):
-            yield batch
-            batch = []
-            longest = 0
-        batch.append(position)
-        longest = max(longest, length)
-    if batch:
-        yield batch
+"""Normalize ordered inputs and plan bounded windows without retaining a full stream."""
+
+from __future__ import annotations
+
+import hashlib
+import sqlite3
+import tempfile
+
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from pathlib import Path
+from typing import overload
+
+from .types import EmbeddingInput
+
+
+def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]:
+    """Yield FASTA records in source order without reading the file into memory."""
+
+    identifier: str | None = None
+    sequence_parts: list[str] = []
+    found_record = False
+    with Path(path).open("r", encoding="utf-8") as handle:
+        for line_number, raw_line in enumerate(handle, start=1):
+            line = raw_line.strip()
+            if not line:
+                continue
+            if line.startswith(">"):
+                if identifier is not None:
+                    found_record = True
+                    yield EmbeddingInput(identifier, "".join(sequence_parts))
+                identifier = line[1:].strip().split(maxsplit=1)[0]
+                if not identifier:
+                    raise ValueError(f"Missing FASTA identifier on line {line_number}.")
+                sequence_parts = []
+            else:
+                if identifier is None:
+                    raise ValueError(
+                        f"Sequence data precedes the first FASTA header on line {line_number}."
+                    )
+                sequence_parts.append("".join(line.split()))
+    if identifier is not None:
+        found_record = True
+        yield EmbeddingInput(identifier, "".join(sequence_parts))
+    if not found_record:
+        raise ValueError(f"No FASTA records found in {path}.")
+
+
+def parse_fasta(path: str | Path) -> list[EmbeddingInput]:
+    """Parse FASTA records while preserving identifiers, order, and duplicates."""
+
+    return list(iter_fasta(path))
+
+
+def _normalize_input_item(
+    position: int,
+    item: str | EmbeddingInput | tuple[str, str],
+) -> EmbeddingInput:
+    if isinstance(item, EmbeddingInput):
+        return item
+    if isinstance(item, str):
+        return EmbeddingInput(str(position), item)
+    if isinstance(item, tuple) and len(item) == 2:
+        return EmbeddingInput(str(item[0]), str(item[1]))
+    raise TypeError(
+        "inputs must contain sequences, EmbeddingInput values, or (id, sequence) tuples."
+    )
+
+
+class _InputSpool(Sequence[EmbeddingInput]):
+    """Immutable disk-backed normalized inputs with an incremental digest."""
+
+    def __init__(
+        self,
+        values: Iterable[str | EmbeddingInput | tuple[str, str]],
+    ) -> None:
+        self._temporary: tempfile.TemporaryDirectory[str] | None = tempfile.TemporaryDirectory(
+            prefix="fastplms-inputs-"
+        )
+        self.path = Path(self._temporary.name) / "inputs.sqlite"
+        self._connection: sqlite3.Connection | None = sqlite3.connect(self.path)
+        self._connection.execute(
+            "CREATE TABLE inputs ("
+            "position INTEGER PRIMARY KEY, input_id TEXT NOT NULL, sequence TEXT NOT NULL)"
+        )
+        digest = hashlib.sha256()
+        count = 0
+        pending: list[tuple[int, str, str]] = []
+        try:
+            for position, item in enumerate(values):
+                record = _normalize_input_item(position, item)
+                for value in (record.id, record.sequence):
+                    encoded = value.encode("utf-8")
+                    digest.update(len(encoded).to_bytes(8, "big"))
+                    digest.update(encoded)
+                pending.append((position, record.id, record.sequence))
+                count += 1
+                if len(pending) == 1_024:
+                    self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+                    pending.clear()
+            if pending:
+                self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+            if count == 0:
+                raise ValueError("inputs must contain at least one sequence.")
+            self._connection.commit()
+            self._connection.close()
+            self._connection = sqlite3.connect(
+                f"{self.path.resolve().as_uri()}?mode=ro",
+                uri=True,
+            )
+        except BaseException:
+            self.close()
+            raise
+        digest.update(count.to_bytes(8, "big"))
+        self.input_fingerprint = digest.hexdigest()
+        self._count = count
+
+    def _require_connection(self) -> sqlite3.Connection:
+        if self._connection is None:
+            raise RuntimeError("Input spool is closed.")
+        return self._connection
+
+    def __len__(self) -> int:
+        return self._count
+
+    def __iter__(self) -> Iterator[EmbeddingInput]:
+        cursor = self._require_connection().execute(
+            "SELECT input_id, sequence FROM inputs ORDER BY position"
+        )
+        while rows := cursor.fetchmany(1_024):
+            for input_id, sequence in rows:
+                yield EmbeddingInput(input_id, sequence)
+
+    @overload
+    def __getitem__(self, index: int, /) -> EmbeddingInput: ...
+
+    @overload
+    def __getitem__(self, index: slice, /) -> list[EmbeddingInput]: ...
+
+    def __getitem__(self, index: int | slice) -> EmbeddingInput | list[EmbeddingInput]:
+        connection = self._require_connection()
+
+        if isinstance(index, slice):
+            start, stop, step = index.indices(self._count)
+            if step != 1:
+                return [self[position] for position in range(start, stop, step)]
+            rows = connection.execute(
+                "SELECT input_id, sequence FROM inputs "
+                "WHERE position >= ? AND position < ? ORDER BY position",
+                (start, stop),
+            ).fetchall()
+            return [EmbeddingInput(input_id, sequence) for input_id, sequence in rows]
+        position = index + self._count if index < 0 else index
+        if position < 0 or position >= self._count:
+            raise IndexError(index)
+        row = connection.execute(
+            "SELECT input_id, sequence FROM inputs WHERE position = ?", (position,)
+        ).fetchone()
+        if row is None:
+            raise IndexError(index)
+        return EmbeddingInput(row[0], row[1])
+
+    def close(self) -> None:
+        connection = getattr(self, "_connection", None)
+        if connection is not None:
+            connection.close()
+            self._connection = None
+        temporary = getattr(self, "_temporary", None)
+        if temporary is not None:
+            temporary.cleanup()
+            self._temporary = None
+
+    def __del__(self) -> None:
+        self.close()
+
+
+def _normalize_inputs(
+    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
+    *,
+    disk_backed: bool,
+) -> Sequence[EmbeddingInput]:
+    is_fasta_path = isinstance(inputs, Path)
+    if isinstance(inputs, str):
+        try:
+            is_fasta_path = Path(inputs).is_file()
+        except OSError:
+            is_fasta_path = False
+    should_spool = disk_backed or is_fasta_path or not isinstance(inputs, (str, Sequence, Mapping))
+    values: Iterable[str | EmbeddingInput | tuple[str, str]]
+    if isinstance(inputs, Path):
+        values = iter_fasta(inputs)
+    elif isinstance(inputs, str):
+        values = iter_fasta(inputs) if is_fasta_path else [inputs]
+    elif isinstance(inputs, Mapping):
+        values = inputs.items()
+    else:
+        values = inputs
+    if should_spool:
+        return _InputSpool(values)
+    records: list[EmbeddingInput] = []
+    for position, item in enumerate(values):
+        records.append(_normalize_input_item(position, item))
+    if not records:
+        raise ValueError("inputs must contain at least one sequence.")
+    return records
+
+
+def _validate_untruncated_lengths(
+    records: Sequence[EmbeddingInput],
+    *,
+    max_length: int | None,
+    truncate: bool,
+) -> None:
+    """Fail before inference when a biological-residue limit would be exceeded."""
+
+    if max_length is None or truncate:
+        return
+    for position, record in enumerate(records):
+        residue_count = len(record.sequence)
+        if residue_count > max_length:
+            raise ValueError(
+                f"Input at position {position} with id {record.id!r} has "
+                f"{residue_count} biological residues, exceeding max_length={max_length} "
+                "while truncate=False."
+            )
+
+
+def _planned_batches(
+    records: Sequence[EmbeddingInput],
+    positions: range,
+    *,
+    batch_size: int,
+    max_tokens_per_batch: int | None,
+    max_length: int | None,
+    truncate: bool,
+) -> Iterator[list[int]]:
+    """Length-bucket one bounded window while retaining stable output positions."""
+
+    def effective_length(position: int) -> int:
+        length = len(records[position].sequence)
+        return min(length, max_length) if truncate and max_length is not None else length
+
+    ordered = sorted(positions, key=lambda position: (-effective_length(position), position))
+    batch: list[int] = []
+    longest = 0
+    for position in ordered:
+        length = effective_length(position)
+        if max_tokens_per_batch is not None and length > max_tokens_per_batch:
+            raise ValueError(
+                f"Input at position {position} has {length} residues, exceeding "
+                f"max_tokens_per_batch={max_tokens_per_batch}."
+            )
+        candidate_longest = max(longest, length)
+        exceeds_tokens = (
+            max_tokens_per_batch is not None
+            and candidate_longest * (len(batch) + 1) > max_tokens_per_batch
+        )
+        if batch and (len(batch) >= batch_size or exceeds_tokens):
+            yield batch
+            batch = []
+            longest = 0
+        batch.append(position)
+        longest = max(longest, length)
+    if batch:
+        yield batch

fastplms/embeddings/output.py

@@ -1,215 +1,215 @@
-"""Resume validation and transactional publication of ordered embedding windows."""
-
-from __future__ import annotations
-
-from collections.abc import Sequence
-from pathlib import Path
-from typing import Any
-
-from .identity import _RUN_FINGERPRINT_SCHEMA_VERSION
-from .pooling import Pooler
-from .storage import (
-    SafetensorsStreamWriter,
-    append_sqlite_records,
-    initialize_sqlite_run,
-    load_result,
-    load_sqlite_result,
-    safetensors_result_exists,
-    save_result,
-    tensor_sha256,
-    update_sqlite_run_metadata,
-)
-from .types import EmbeddingInput, EmbeddingRecord, EmbeddingResult, LazyTensorReference
-
-
-def _output_exists(path: str | Path, format: str) -> bool:
-    path = Path(path)
-    if format == "sqlite":
-        return path.is_file()
-    return safetensors_result_exists(path)
-
-
-def _output_descriptor(position: int, record: EmbeddingRecord) -> dict[str, Any]:
-    tensor = record.tensor
-    if isinstance(tensor, LazyTensorReference):
-        dtype = tensor.dtype
-        shape = tensor.shape
-        digest = tensor.sha256
-    else:
-        dtype = str(tensor.dtype).removeprefix("torch.")
-        shape = tuple(tensor.shape)
-        digest = tensor_sha256(tensor)
-    return {
-        "position": position,
-        "id": record.id,
-        "dtype": dtype,
-        "shape": shape,
-        "sha256": digest,
-    }
-
-
-class EmbeddingOutput:
-    """Own the resumable prefix and the commit state of one output destination."""
-
-    def __init__(
-        self,
-        records: Sequence[EmbeddingInput],
-        *,
-        output: str | Path | None,
-        format: str,
-        resume: bool,
-        shard_size: int,
-        run_fingerprint: str,
-        input_fingerprint: str,
-        model_state_fingerprint: str | None,
-        model_state_fingerprint_source: str,
-        pooler: Pooler | None,
-        pooling_names: Sequence[str],
-    ) -> None:
-        self.output = output
-        self.format = format
-        self.shard_size = shard_size
-        self.completed: EmbeddingResult | None = None
-        output_already_exists = output is not None and _output_exists(output, format)
-        existing: EmbeddingResult | None = None
-        self.start_position = 0
-        if output is not None and resume and output_already_exists:
-            if format == "sqlite":
-                try:
-                    existing = load_sqlite_result(output, run_id=run_fingerprint)
-                except KeyError:
-                    existing = load_result(output, format=format)
-            else:
-                existing = load_result(output, format=format)
-            if existing.metadata.get("fingerprint_schema_version") != (
-                _RUN_FINGERPRINT_SCHEMA_VERSION
-            ):
-                raise ValueError(
-                    "Existing embeddings use an incompatible run fingerprint schema; "
-                    "choose another output or set resume=False."
-                )
-            if existing.metadata.get("run_fingerprint") != run_fingerprint:
-                raise ValueError(
-                    "Existing embeddings were produced by a different run fingerprint; "
-                    "choose another output or set resume=False."
-                )
-            if len(existing) > len(records):
-                raise ValueError(
-                    "Existing embeddings are not an ordered prefix of the requested inputs."
-                )
-            prefix_matches = all(
-                (observed.id, observed.sequence) == (expected.id, expected.sequence)
-                for expected, observed in zip(records, existing, strict=False)
-            )
-            if not prefix_matches:
-                raise ValueError(
-                    "Existing embeddings are not an ordered prefix of the requested inputs."
-                )
-            if len(existing) == len(records) and existing.metadata.get("complete", True):
-                self.completed = existing
-                return
-            self.start_position = len(existing)
-
-        self.sqlite_run_id: str | None = None
-        self.sqlite_replace_on_first_commit = False
-        self.sqlite_initial_metadata: dict[str, Any] | None = None
-        if output is not None and format == "sqlite":
-            self.sqlite_initial_metadata = {
-                "format_version": 1,
-                "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-                "run_fingerprint": run_fingerprint,
-                "input_fingerprint": input_fingerprint,
-                "model_state_fingerprint": model_state_fingerprint,
-                "model_state_fingerprint_source": model_state_fingerprint_source,
-                "complete": False,
-            }
-            self.sqlite_run_id = run_fingerprint
-            if not resume and output_already_exists:
-                try:
-                    load_sqlite_result(output, run_id=run_fingerprint)
-                except KeyError:
-                    pass
-                else:
-                    # Keep an exact prior run readable until replacement inference
-                    # has produced the first complete commit window.
-                    self.sqlite_replace_on_first_commit = True
-            if not self.sqlite_replace_on_first_commit:
-                initialize_sqlite_run(
-                    output,
-                    self.sqlite_initial_metadata,
-                    resume=resume,
-                )
-
-        stream_safetensors = output is not None and format == "safetensors"
-        self.output_records: list[EmbeddingRecord] = (
-            [] if self.sqlite_run_id is not None or stream_safetensors else list(existing or ())
-        )
-        self.output_descriptors: list[dict[str, Any]] | None = [] if output is None else None
-        self.pool_slices: dict[str, tuple[int, int]] = {}
-        if existing and pooler is not None:
-            pooled_width = existing[0].load_tensor().shape[-1]
-            if pooled_width % len(pooling_names) != 0:
-                raise ValueError("Stored pooled width is inconsistent with pooling metadata.")
-            self.pool_slices = pooler.output_slices(pooled_width // len(pooling_names))
-
-        self.safetensors_writer: SafetensorsStreamWriter | None = None
-        if stream_safetensors:
-            if output is None:
-                raise RuntimeError(
-                    "Safetensors streaming was enabled without an output destination."
-                )
-            transactional_overwrite = output_already_exists and not resume
-            self.safetensors_writer = SafetensorsStreamWriter(
-                output,
-                {
-                    "format_version": 1,
-                    "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-                    "run_fingerprint": run_fingerprint,
-                    "input_fingerprint": input_fingerprint,
-                    "model_state_fingerprint": model_state_fingerprint,
-                    "model_state_fingerprint_source": model_state_fingerprint_source,
-                    "complete": False,
-                },
-                shard_size=shard_size,
-                existing=existing or (),
-                reuse_existing=bool(resume and existing is not None),
-                publish_initial=not transactional_overwrite,
-                publish_incremental=not transactional_overwrite,
-            )
-
-    def append(self, window_start: int, new_records: list[EmbeddingRecord]) -> None:
-        """Commit a complete ordered window at the storage format's granularity."""
-
-        if self.output_descriptors is not None:
-            self.output_descriptors.extend(
-                _output_descriptor(window_start + offset, record)
-                for offset, record in enumerate(new_records)
-            )
-        if self.output is not None and self.sqlite_run_id is not None:
-            append_sqlite_records(
-                self.output,
-                self.sqlite_run_id,
-                window_start,
-                new_records,
-                replace_metadata=(
-                    self.sqlite_initial_metadata if self.sqlite_replace_on_first_commit else None
-                ),
-            )
-            self.sqlite_replace_on_first_commit = False
-        elif self.safetensors_writer is not None:
-            self.safetensors_writer.append(new_records)
-        else:
-            self.output_records.extend(new_records)
-
-    def finish(self, metadata: dict[str, Any]) -> EmbeddingResult:
-        """Publish completion only after every window has committed."""
-
-        if self.output is not None and self.sqlite_run_id is not None:
-            update_sqlite_run_metadata(self.output, self.sqlite_run_id, metadata)
-            return load_sqlite_result(self.output, run_id=self.sqlite_run_id)
-        if self.safetensors_writer is not None:
-            return self.safetensors_writer.publish(complete=True, metadata=metadata)
-        result = EmbeddingResult(self.output_records, metadata)
-        if self.output is not None:
-            return save_result(result, self.output, format=self.format, shard_size=self.shard_size)
-        return result
+"""Resume validation and transactional publication of ordered embedding windows."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from .identity import _RUN_FINGERPRINT_SCHEMA_VERSION
+from .pooling import Pooler
+from .storage import (
+    SafetensorsStreamWriter,
+    append_sqlite_records,
+    initialize_sqlite_run,
+    load_result,
+    load_sqlite_result,
+    safetensors_result_exists,
+    save_result,
+    tensor_sha256,
+    update_sqlite_run_metadata,
+)
+from .types import EmbeddingInput, EmbeddingRecord, EmbeddingResult, LazyTensorReference
+
+
+def _output_exists(path: str | Path, format: str) -> bool:
+    path = Path(path)
+    if format == "sqlite":
+        return path.is_file()
+    return safetensors_result_exists(path)
+
+
+def _output_descriptor(position: int, record: EmbeddingRecord) -> dict[str, Any]:
+    tensor = record.tensor
+    if isinstance(tensor, LazyTensorReference):
+        dtype = tensor.dtype
+        shape = tensor.shape
+        digest = tensor.sha256
+    else:
+        dtype = str(tensor.dtype).removeprefix("torch.")
+        shape = tuple(tensor.shape)
+        digest = tensor_sha256(tensor)
+    return {
+        "position": position,
+        "id": record.id,
+        "dtype": dtype,
+        "shape": shape,
+        "sha256": digest,
+    }
+
+
+class EmbeddingOutput:
+    """Own the resumable prefix and the commit state of one output destination."""
+
+    def __init__(
+        self,
+        records: Sequence[EmbeddingInput],
+        *,
+        output: str | Path | None,
+        format: str,
+        resume: bool,
+        shard_size: int,
+        run_fingerprint: str,
+        input_fingerprint: str,
+        model_state_fingerprint: str | None,
+        model_state_fingerprint_source: str,
+        pooler: Pooler | None,
+        pooling_names: Sequence[str],
+    ) -> None:
+        self.output = output
+        self.format = format
+        self.shard_size = shard_size
+        self.completed: EmbeddingResult | None = None
+        output_already_exists = output is not None and _output_exists(output, format)
+        existing: EmbeddingResult | None = None
+        self.start_position = 0
+        if output is not None and resume and output_already_exists:
+            if format == "sqlite":
+                try:
+                    existing = load_sqlite_result(output, run_id=run_fingerprint)
+                except KeyError:
+                    existing = load_result(output, format=format)
+            else:
+                existing = load_result(output, format=format)
+            if existing.metadata.get("fingerprint_schema_version") != (
+                _RUN_FINGERPRINT_SCHEMA_VERSION
+            ):
+                raise ValueError(
+                    "Existing embeddings use an incompatible run fingerprint schema; "
+                    "choose another output or set resume=False."
+                )
+            if existing.metadata.get("run_fingerprint") != run_fingerprint:
+                raise ValueError(
+                    "Existing embeddings were produced by a different run fingerprint; "
+                    "choose another output or set resume=False."
+                )
+            if len(existing) > len(records):
+                raise ValueError(
+                    "Existing embeddings are not an ordered prefix of the requested inputs."
+                )
+            prefix_matches = all(
+                (observed.id, observed.sequence) == (expected.id, expected.sequence)
+                for expected, observed in zip(records, existing, strict=False)
+            )
+            if not prefix_matches:
+                raise ValueError(
+                    "Existing embeddings are not an ordered prefix of the requested inputs."
+                )
+            if len(existing) == len(records) and existing.metadata.get("complete", True):
+                self.completed = existing
+                return
+            self.start_position = len(existing)
+
+        self.sqlite_run_id: str | None = None
+        self.sqlite_replace_on_first_commit = False
+        self.sqlite_initial_metadata: dict[str, Any] | None = None
+        if output is not None and format == "sqlite":
+            self.sqlite_initial_metadata = {
+                "format_version": 1,
+                "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+                "run_fingerprint": run_fingerprint,
+                "input_fingerprint": input_fingerprint,
+                "model_state_fingerprint": model_state_fingerprint,
+                "model_state_fingerprint_source": model_state_fingerprint_source,
+                "complete": False,
+            }
+            self.sqlite_run_id = run_fingerprint
+            if not resume and output_already_exists:
+                try:
+                    load_sqlite_result(output, run_id=run_fingerprint)
+                except KeyError:
+                    pass
+                else:
+                    # Keep an exact prior run readable until replacement inference
+                    # has produced the first complete commit window.
+                    self.sqlite_replace_on_first_commit = True
+            if not self.sqlite_replace_on_first_commit:
+                initialize_sqlite_run(
+                    output,
+                    self.sqlite_initial_metadata,
+                    resume=resume,
+                )
+
+        stream_safetensors = output is not None and format == "safetensors"
+        self.output_records: list[EmbeddingRecord] = (
+            [] if self.sqlite_run_id is not None or stream_safetensors else list(existing or ())
+        )
+        self.output_descriptors: list[dict[str, Any]] | None = [] if output is None else None
+        self.pool_slices: dict[str, tuple[int, int]] = {}
+        if existing and pooler is not None:
+            pooled_width = existing[0].load_tensor().shape[-1]
+            if pooled_width % len(pooling_names) != 0:
+                raise ValueError("Stored pooled width is inconsistent with pooling metadata.")
+            self.pool_slices = pooler.output_slices(pooled_width // len(pooling_names))
+
+        self.safetensors_writer: SafetensorsStreamWriter | None = None
+        if stream_safetensors:
+            if output is None:
+                raise RuntimeError(
+                    "Safetensors streaming was enabled without an output destination."
+                )
+            transactional_overwrite = output_already_exists and not resume
+            self.safetensors_writer = SafetensorsStreamWriter(
+                output,
+                {
+                    "format_version": 1,
+                    "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+                    "run_fingerprint": run_fingerprint,
+                    "input_fingerprint": input_fingerprint,
+                    "model_state_fingerprint": model_state_fingerprint,
+                    "model_state_fingerprint_source": model_state_fingerprint_source,
+                    "complete": False,
+                },
+                shard_size=shard_size,
+                existing=existing or (),
+                reuse_existing=bool(resume and existing is not None),
+                publish_initial=not transactional_overwrite,
+                publish_incremental=not transactional_overwrite,
+            )
+
+    def append(self, window_start: int, new_records: list[EmbeddingRecord]) -> None:
+        """Commit a complete ordered window at the storage format's granularity."""
+
+        if self.output_descriptors is not None:
+            self.output_descriptors.extend(
+                _output_descriptor(window_start + offset, record)
+                for offset, record in enumerate(new_records)
+            )
+        if self.output is not None and self.sqlite_run_id is not None:
+            append_sqlite_records(
+                self.output,
+                self.sqlite_run_id,
+                window_start,
+                new_records,
+                replace_metadata=(
+                    self.sqlite_initial_metadata if self.sqlite_replace_on_first_commit else None
+                ),
+            )
+            self.sqlite_replace_on_first_commit = False
+        elif self.safetensors_writer is not None:
+            self.safetensors_writer.append(new_records)
+        else:
+            self.output_records.extend(new_records)
+
+    def finish(self, metadata: dict[str, Any]) -> EmbeddingResult:
+        """Publish completion only after every window has committed."""
+
+        if self.output is not None and self.sqlite_run_id is not None:
+            update_sqlite_run_metadata(self.output, self.sqlite_run_id, metadata)
+            return load_sqlite_result(self.output, run_id=self.sqlite_run_id)
+        if self.safetensors_writer is not None:
+            return self.safetensors_writer.publish(complete=True, metadata=metadata)
+        result = EmbeddingResult(self.output_records, metadata)
+        if self.output is not None:
+            return save_result(result, self.output, format=self.format, shard_size=self.shard_size)
+        return result

fastplms/embeddings/pooling.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import math
 import torch
+
 from collections.abc import Sequence
 from torch import Tensor
 

fastplms/embeddings/runner.py

@@ -1,425 +1,426 @@
-"""Coordinate input preparation, run identity, batch execution, and publication."""
-
-from __future__ import annotations
-
-import torch
-from collections.abc import Callable, Iterable, Mapping, Sequence
-from pathlib import Path
-from typing import Any
+"""Coordinate input preparation, run identity, batch execution, and publication."""
+
+from __future__ import annotations
+
+import torch
+
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
 from torch import Tensor
 
 from . import identity
 from .batches import (
-    BatchExecutor,
-    _residue_embeddings as _residue_embeddings,
-    _temporary_eval,
-    select_hidden_state_embeddings as select_hidden_state_embeddings,
-)
-from .identity import (
-    _RUN_FINGERPRINT_SCHEMA_VERSION,
-    _adapter_identity_metadata,
-    _attention_backend,
-    _attention_kernel_metadata,
-    _embedding_context,
-    _execution_identity_metadata,
-    _fingerprint_jsonable,
-    _model_identity_metadata,
-    _run_fingerprint,
+    BatchExecutor,
+    _residue_embeddings as _residue_embeddings,
+    _temporary_eval,
+    select_hidden_state_embeddings as select_hidden_state_embeddings,
+)
+from .identity import (
+    _RUN_FINGERPRINT_SCHEMA_VERSION,
+    _adapter_identity_metadata,
+    _attention_backend,
+    _attention_kernel_metadata,
+    _embedding_context,
+    _execution_identity_metadata,
+    _fingerprint_jsonable,
+    _model_identity_metadata,
+    _run_fingerprint,
     _tokenizer_metadata,
-)
-from .inputs import (
-    _InputSpool,
-    _normalize_inputs,
-    _validate_untruncated_lengths,
-    iter_fasta as iter_fasta,
-    parse_fasta as parse_fasta,
-)
-from .output import EmbeddingOutput
-from .pooling import Pooler
-from .types import EmbeddingBatch, EmbeddingInput, EmbeddingResult
-
-
-_DEFAULT_BATCH_WINDOW_MULTIPLIER = 16
-_SUPPORTED_STORAGE_FORMATS = frozenset({"safetensors", "sqlite"})
-
-
-def embed_dataset(
-    model: Any,
-    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
-    *,
-    batch_size: int = 2,
-    pooling: str | Sequence[str] | None = None,
-    full_embeddings: bool = False,
-    output: str | Path | None = None,
-    format: str = "safetensors",
-    resume: bool = True,
-    tokenizer: Any | None = None,
-    max_length: int | None = None,
-    truncate: bool = True,
-    dtype: torch.dtype | None = torch.float32,
-    shard_size: int = 2 * 1024**3,
-    model_state_fingerprint: str | None = None,
-    batch_window_size: int | None = None,
-    max_tokens_per_batch: int | None = None,
-    hidden_state_source: str = "encoder",
-    decoder_inputs: Sequence[str] | None = None,
-    decoder_input_ids: Tensor | None = None,
-    decoder_attention_mask: Tensor | None = None,
-    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None = None,
-    _embedding_batch_identity: Mapping[str, Any] | None = None,
-    _allowed_unsupported_pooling: Sequence[str] = (),
-    **model_kwargs: Any,
-) -> EmbeddingResult:
-    """Embed protein sequences with stable ordering and residue-only pooling."""
-
-    for name, value in (
-        ("batch_size", batch_size),
-        ("shard_size", shard_size),
-    ):
-        if not isinstance(value, int) or isinstance(value, bool):
-            raise TypeError(f"{name} must be a positive integer.")
-        if value <= 0:
-            raise ValueError(f"{name} must be a positive integer.")
-    for optional_name, optional_value in (
-        ("max_length", max_length),
-        ("max_tokens_per_batch", max_tokens_per_batch),
-        ("batch_window_size", batch_window_size),
-    ):
-        if optional_value is not None and (
-            not isinstance(optional_value, int) or isinstance(optional_value, bool)
-        ):
-            raise TypeError(f"{optional_name} must be a positive integer when provided.")
-        if optional_value is not None and optional_value <= 0:
-            raise ValueError(f"{optional_name} must be a positive integer when provided.")
-    for name, value in (
-        ("full_embeddings", full_embeddings),
-        ("resume", resume),
-        ("truncate", truncate),
-    ):
-        if not isinstance(value, bool):
-            raise TypeError(f"{name} must be a boolean.")
-    if not isinstance(format, str):
-        raise TypeError("format must be a string.")
-    if output is not None and not isinstance(output, (str, Path)):
-        raise TypeError("output must be a path or None.")
-    if model_state_fingerprint is not None and (
-        not isinstance(model_state_fingerprint, str) or not model_state_fingerprint
-    ):
-        raise ValueError("model_state_fingerprint must be a non-empty string when provided.")
-    if hidden_state_source not in {"encoder", "decoder"}:
-        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
-    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
-    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
-        raise TypeError("hidden_state_index must be an integer.")
-    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
-    if not isinstance(store_all_hidden_states, bool):
-        raise TypeError("store_all_hidden_states must be a boolean.")
-    if decoder_input_ids is not None:
-        if not isinstance(decoder_input_ids, Tensor):
-            raise TypeError("decoder_input_ids must be a tensor.")
-        if decoder_input_ids.is_meta:
-            raise ValueError("decoder_input_ids cannot be a meta tensor.")
-        if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[1] == 0:
-            raise ValueError("decoder_input_ids must have non-empty shape (batch, sequence).")
-        if decoder_input_ids.dtype not in {torch.int32, torch.int64}:
-            raise TypeError("decoder_input_ids must use torch.int32 or torch.int64.")
-    if decoder_attention_mask is not None:
-        if not isinstance(decoder_attention_mask, Tensor):
-            raise TypeError("decoder_attention_mask must be a tensor.")
-        if decoder_attention_mask.is_meta:
-            raise ValueError("decoder_attention_mask cannot be a meta tensor.")
-        if decoder_attention_mask.is_complex() or not bool(
-            torch.isfinite(decoder_attention_mask).all()
-        ):
-            raise ValueError("decoder_attention_mask must contain finite binary values.")
-        if not bool(((decoder_attention_mask == 0) | (decoder_attention_mask == 1)).all()):
-            raise ValueError("decoder_attention_mask must contain finite binary values.")
-    pooling_names = (
-        (("mean",) if not full_embeddings else ())
-        if pooling is None
-        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
-    )
-    if full_embeddings and pooling is not None:
-        raise ValueError("full_embeddings=True cannot be combined with pooling.")
-    if not full_embeddings and not pooling_names:
-        raise ValueError("pooling is required unless full_embeddings=True.")
-    pooler = Pooler(pooling_names) if pooling_names else None
-
-    if batch_size <= 0:
-        raise ValueError("batch_size must be positive.")
-    if format == "pth" or (output is not None and Path(output).suffix.lower() == ".pth"):
-        raise ValueError("Writing pickle-based .pth embeddings is not supported.")
-    if format not in _SUPPORTED_STORAGE_FORMATS:
-        raise ValueError("format must be 'safetensors' or 'sqlite'.")
-    if max_length is not None and max_length <= 0:
-        raise ValueError("max_length must be positive when provided.")
-    if max_tokens_per_batch is not None and max_tokens_per_batch <= 0:
-        raise ValueError("max_tokens_per_batch must be positive when provided.")
-    if not isinstance(dtype, (torch.dtype, type(None))):
-        raise TypeError("dtype must be a torch.dtype or None.")
-    if batch_window_size is not None and batch_window_size <= 0:
-        raise ValueError("batch_window_size must be positive when provided.")
-    if _embedding_batch_fn is not None and not callable(_embedding_batch_fn):
-        raise TypeError("_embedding_batch_fn must be callable when provided.")
-    if _embedding_batch_fn is not None and _embedding_batch_identity is None:
-        raise ValueError(
-            "_embedding_batch_identity is required with _embedding_batch_fn so persisted "
-            "runs bind the family-specific embedding behavior."
-        )
-    if _embedding_batch_identity is not None and not isinstance(_embedding_batch_identity, Mapping):
-        raise TypeError("_embedding_batch_identity must be a mapping when provided.")
-    if isinstance(_allowed_unsupported_pooling, (str, bytes)) or not isinstance(
-        _allowed_unsupported_pooling, Sequence
-    ):
-        raise TypeError("_allowed_unsupported_pooling must be a sequence of pooler names.")
-    if not all(isinstance(name, str) for name in _allowed_unsupported_pooling):
-        raise TypeError("_allowed_unsupported_pooling must contain only strings.")
-    allowed_unsupported_pooling = frozenset(_allowed_unsupported_pooling)
-    if allowed_unsupported_pooling and _embedding_batch_fn is None:
-        raise ValueError(
-            "_allowed_unsupported_pooling is only valid with a family-specific _embedding_batch_fn."
-        )
-    resolved_batch_window_size = (
-        batch_size * _DEFAULT_BATCH_WINDOW_MULTIPLIER
-        if batch_window_size is None
-        else batch_window_size
-    )
-    if resolved_batch_window_size < batch_size:
-        raise ValueError("batch_window_size must be at least batch_size.")
-    records = _normalize_inputs(inputs, disk_backed=output is not None)
-    _validate_untruncated_lengths(
-        records,
-        max_length=max_length,
-        truncate=truncate,
-    )
-    pooling_names = (
-        (("mean",) if not full_embeddings else ())
-        if pooling is None
-        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
-    )
-    if full_embeddings:
-        if pooling is not None:
-            raise ValueError("full_embeddings=True cannot be combined with pooling.")
-    elif not pooling_names:
-        raise ValueError("pooling is required unless full_embeddings=True.")
-    store_all_hidden_states = bool(model_kwargs.get("store_all_hidden_states", False))
-    if store_all_hidden_states and not full_embeddings:
-        raise ValueError("store_all_hidden_states=True requires full_embeddings=True.")
-
-    unsupported = set(getattr(model, "embedding_unsupported_pooling", ()))
-    unknown_pooling_overrides = allowed_unsupported_pooling.difference(unsupported)
-    if unknown_pooling_overrides:
-        raise ValueError(
-            "_allowed_unsupported_pooling may only override poolers declared unsupported "
-            f"by the model; unknown overrides: {sorted(unknown_pooling_overrides)}."
-        )
-    unsupported.difference_update(allowed_unsupported_pooling)
-    requested_unsupported = unsupported.intersection(pooling_names)
-    if requested_unsupported:
-        raise ValueError(
-            f"{model.__class__.__name__} does not support pooling operations "
-            f"{sorted(requested_unsupported)}."
-        )
-
-    # Constructing the pooler validates names and duplicate operations before
-    # any checkpoint hashing, tokenization, or inference occurs.
-    pooler = Pooler(pooling_names) if pooling_names else None
-    embedding_context, normalized_decoder_inputs = _embedding_context(
-        model,
-        records,
-        hidden_state_source=hidden_state_source,
-        decoder_inputs=decoder_inputs,
-        decoder_input_ids=decoder_input_ids,
-        decoder_attention_mask=decoder_attention_mask,
-        model_kwargs=model_kwargs,
-    )
-    if _embedding_batch_identity is not None:
-        embedding_context["family_adapter"] = _fingerprint_jsonable(_embedding_batch_identity)
-        if allowed_unsupported_pooling:
-            embedding_context["family_adapter_pooling_override"] = sorted(
-                allowed_unsupported_pooling
-            )
-
-    # A pending automatic attention request settles here, inside the caller's
-    # autocast context, so the fingerprint records the backend that executes.
-    attention_resolution = getattr(model, "attention_resolution", None)
-    if attention_resolution is not None and attention_resolution.deferred:
-        model.resolve_attn_implementation()
-
-    tokenizer_metadata = _tokenizer_metadata(model, tokenizer)
-    (
-        input_fingerprint,
-        run_fingerprint,
-        resolved_model_state_fingerprint,
-        model_state_fingerprint_source,
-    ) = _run_fingerprint(
-        model,
-        records,
-        pooling=pooling_names,
-        full_embeddings=full_embeddings,
-        max_length=max_length,
-        truncate=truncate,
-        dtype=dtype,
-        model_kwargs=model_kwargs,
-        tokenizer_metadata=tokenizer_metadata,
-        model_state_fingerprint=model_state_fingerprint,
-        persist_output=output is not None,
-        embedding_context=embedding_context,
-        batch_size=batch_size,
-        batch_window_size=resolved_batch_window_size,
-        max_tokens_per_batch=max_tokens_per_batch,
-    )
-    destination = EmbeddingOutput(
-        records,
-        output=output,
-        format=format,
-        resume=resume,
-        shard_size=shard_size,
-        run_fingerprint=run_fingerprint,
-        input_fingerprint=input_fingerprint,
-        model_state_fingerprint=resolved_model_state_fingerprint,
-        model_state_fingerprint_source=model_state_fingerprint_source,
-        pooler=pooler,
-        pooling_names=pooling_names,
-    )
-    if destination.completed is not None:
-        return destination.completed
-
-    attention_backend = _attention_backend(model)
-    executor = BatchExecutor(
-        model=model,
-        batch_size=batch_size,
-        max_tokens_per_batch=max_tokens_per_batch,
-        max_length=max_length,
-        truncate=truncate,
-        model_kwargs=model_kwargs,
-        hidden_state_source=hidden_state_source,
-        normalized_decoder_inputs=normalized_decoder_inputs,
-        decoder_input_ids=decoder_input_ids,
-        decoder_attention_mask=decoder_attention_mask,
-        _embedding_batch_fn=_embedding_batch_fn,
-        tokenizer=tokenizer,
-        store_all_hidden_states=store_all_hidden_states,
-        full_embeddings=full_embeddings,
-        dtype=dtype,
-        pooler=pooler,
-        attention_backend=attention_backend,
-        need_attentions="parti" in pooling_names,
-    )
-    pool_slices = destination.pool_slices
-    with _temporary_eval(model), torch.inference_mode():
-        for window_start in range(
-            destination.start_position, len(records), resolved_batch_window_size
-        ):
-            window_stop = min(window_start + resolved_batch_window_size, len(records))
-            window_records = records[window_start:window_stop]
-            if not isinstance(window_records, Sequence):
-                raise RuntimeError("The immutable embedding spool returned a non-sequence window.")
-            new_records, pool_slices = executor.run_window(
-                window_records, window_start=window_start
-            )
-            destination.append(window_start, new_records)
-
+)
+from .inputs import (
+    _InputSpool,
+    _normalize_inputs,
+    _validate_untruncated_lengths,
+    iter_fasta as iter_fasta,
+    parse_fasta as parse_fasta,
+)
+from .output import EmbeddingOutput
+from .pooling import Pooler
+from .types import EmbeddingBatch, EmbeddingInput, EmbeddingResult
+
+
+_DEFAULT_BATCH_WINDOW_MULTIPLIER = 16
+_SUPPORTED_STORAGE_FORMATS = frozenset({"safetensors", "sqlite"})
+
+
+def embed_dataset(
+    model: Any,
+    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
+    *,
+    batch_size: int = 2,
+    pooling: str | Sequence[str] | None = None,
+    full_embeddings: bool = False,
+    output: str | Path | None = None,
+    format: str = "safetensors",
+    resume: bool = True,
+    tokenizer: Any | None = None,
+    max_length: int | None = None,
+    truncate: bool = True,
+    dtype: torch.dtype | None = torch.float32,
+    shard_size: int = 2 * 1024**3,
+    model_state_fingerprint: str | None = None,
+    batch_window_size: int | None = None,
+    max_tokens_per_batch: int | None = None,
+    hidden_state_source: str = "encoder",
+    decoder_inputs: Sequence[str] | None = None,
+    decoder_input_ids: Tensor | None = None,
+    decoder_attention_mask: Tensor | None = None,
+    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None = None,
+    _embedding_batch_identity: Mapping[str, Any] | None = None,
+    _allowed_unsupported_pooling: Sequence[str] = (),
+    **model_kwargs: Any,
+) -> EmbeddingResult:
+    """Embed protein sequences with stable ordering and residue-only pooling."""
+
+    for name, value in (
+        ("batch_size", batch_size),
+        ("shard_size", shard_size),
+    ):
+        if not isinstance(value, int) or isinstance(value, bool):
+            raise TypeError(f"{name} must be a positive integer.")
+        if value <= 0:
+            raise ValueError(f"{name} must be a positive integer.")
+    for optional_name, optional_value in (
+        ("max_length", max_length),
+        ("max_tokens_per_batch", max_tokens_per_batch),
+        ("batch_window_size", batch_window_size),
+    ):
+        if optional_value is not None and (
+            not isinstance(optional_value, int) or isinstance(optional_value, bool)
+        ):
+            raise TypeError(f"{optional_name} must be a positive integer when provided.")
+        if optional_value is not None and optional_value <= 0:
+            raise ValueError(f"{optional_name} must be a positive integer when provided.")
+    for name, value in (
+        ("full_embeddings", full_embeddings),
+        ("resume", resume),
+        ("truncate", truncate),
+    ):
+        if not isinstance(value, bool):
+            raise TypeError(f"{name} must be a boolean.")
+    if not isinstance(format, str):
+        raise TypeError("format must be a string.")
+    if output is not None and not isinstance(output, (str, Path)):
+        raise TypeError("output must be a path or None.")
+    if model_state_fingerprint is not None and (
+        not isinstance(model_state_fingerprint, str) or not model_state_fingerprint
+    ):
+        raise ValueError("model_state_fingerprint must be a non-empty string when provided.")
+    if hidden_state_source not in {"encoder", "decoder"}:
+        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
+    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
+    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
+        raise TypeError("hidden_state_index must be an integer.")
+    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
+    if not isinstance(store_all_hidden_states, bool):
+        raise TypeError("store_all_hidden_states must be a boolean.")
+    if decoder_input_ids is not None:
+        if not isinstance(decoder_input_ids, Tensor):
+            raise TypeError("decoder_input_ids must be a tensor.")
+        if decoder_input_ids.is_meta:
+            raise ValueError("decoder_input_ids cannot be a meta tensor.")
+        if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[1] == 0:
+            raise ValueError("decoder_input_ids must have non-empty shape (batch, sequence).")
+        if decoder_input_ids.dtype not in {torch.int32, torch.int64}:
+            raise TypeError("decoder_input_ids must use torch.int32 or torch.int64.")
+    if decoder_attention_mask is not None:
+        if not isinstance(decoder_attention_mask, Tensor):
+            raise TypeError("decoder_attention_mask must be a tensor.")
+        if decoder_attention_mask.is_meta:
+            raise ValueError("decoder_attention_mask cannot be a meta tensor.")
+        if decoder_attention_mask.is_complex() or not bool(
+            torch.isfinite(decoder_attention_mask).all()
+        ):
+            raise ValueError("decoder_attention_mask must contain finite binary values.")
+        if not bool(((decoder_attention_mask == 0) | (decoder_attention_mask == 1)).all()):
+            raise ValueError("decoder_attention_mask must contain finite binary values.")
+    pooling_names = (
+        (("mean",) if not full_embeddings else ())
+        if pooling is None
+        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
+    )
+    if full_embeddings and pooling is not None:
+        raise ValueError("full_embeddings=True cannot be combined with pooling.")
+    if not full_embeddings and not pooling_names:
+        raise ValueError("pooling is required unless full_embeddings=True.")
+    pooler = Pooler(pooling_names) if pooling_names else None
+
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive.")
+    if format == "pth" or (output is not None and Path(output).suffix.lower() == ".pth"):
+        raise ValueError("Writing pickle-based .pth embeddings is not supported.")
+    if format not in _SUPPORTED_STORAGE_FORMATS:
+        raise ValueError("format must be 'safetensors' or 'sqlite'.")
+    if max_length is not None and max_length <= 0:
+        raise ValueError("max_length must be positive when provided.")
+    if max_tokens_per_batch is not None and max_tokens_per_batch <= 0:
+        raise ValueError("max_tokens_per_batch must be positive when provided.")
+    if not isinstance(dtype, (torch.dtype, type(None))):
+        raise TypeError("dtype must be a torch.dtype or None.")
+    if batch_window_size is not None and batch_window_size <= 0:
+        raise ValueError("batch_window_size must be positive when provided.")
+    if _embedding_batch_fn is not None and not callable(_embedding_batch_fn):
+        raise TypeError("_embedding_batch_fn must be callable when provided.")
+    if _embedding_batch_fn is not None and _embedding_batch_identity is None:
+        raise ValueError(
+            "_embedding_batch_identity is required with _embedding_batch_fn so persisted "
+            "runs bind the family-specific embedding behavior."
+        )
+    if _embedding_batch_identity is not None and not isinstance(_embedding_batch_identity, Mapping):
+        raise TypeError("_embedding_batch_identity must be a mapping when provided.")
+    if isinstance(_allowed_unsupported_pooling, (str, bytes)) or not isinstance(
+        _allowed_unsupported_pooling, Sequence
+    ):
+        raise TypeError("_allowed_unsupported_pooling must be a sequence of pooler names.")
+    if not all(isinstance(name, str) for name in _allowed_unsupported_pooling):
+        raise TypeError("_allowed_unsupported_pooling must contain only strings.")
+    allowed_unsupported_pooling = frozenset(_allowed_unsupported_pooling)
+    if allowed_unsupported_pooling and _embedding_batch_fn is None:
+        raise ValueError(
+            "_allowed_unsupported_pooling is only valid with a family-specific _embedding_batch_fn."
+        )
+    resolved_batch_window_size = (
+        batch_size * _DEFAULT_BATCH_WINDOW_MULTIPLIER
+        if batch_window_size is None
+        else batch_window_size
+    )
+    if resolved_batch_window_size < batch_size:
+        raise ValueError("batch_window_size must be at least batch_size.")
+    records = _normalize_inputs(inputs, disk_backed=output is not None)
+    _validate_untruncated_lengths(
+        records,
+        max_length=max_length,
+        truncate=truncate,
+    )
+    pooling_names = (
+        (("mean",) if not full_embeddings else ())
+        if pooling is None
+        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
+    )
+    if full_embeddings:
+        if pooling is not None:
+            raise ValueError("full_embeddings=True cannot be combined with pooling.")
+    elif not pooling_names:
+        raise ValueError("pooling is required unless full_embeddings=True.")
+    store_all_hidden_states = bool(model_kwargs.get("store_all_hidden_states", False))
+    if store_all_hidden_states and not full_embeddings:
+        raise ValueError("store_all_hidden_states=True requires full_embeddings=True.")
+
+    unsupported = set(getattr(model, "embedding_unsupported_pooling", ()))
+    unknown_pooling_overrides = allowed_unsupported_pooling.difference(unsupported)
+    if unknown_pooling_overrides:
+        raise ValueError(
+            "_allowed_unsupported_pooling may only override poolers declared unsupported "
+            f"by the model; unknown overrides: {sorted(unknown_pooling_overrides)}."
+        )
+    unsupported.difference_update(allowed_unsupported_pooling)
+    requested_unsupported = unsupported.intersection(pooling_names)
+    if requested_unsupported:
+        raise ValueError(
+            f"{model.__class__.__name__} does not support pooling operations "
+            f"{sorted(requested_unsupported)}."
+        )
+
+    # Constructing the pooler validates names and duplicate operations before
+    # any checkpoint hashing, tokenization, or inference occurs.
+    pooler = Pooler(pooling_names) if pooling_names else None
+    embedding_context, normalized_decoder_inputs = _embedding_context(
+        model,
+        records,
+        hidden_state_source=hidden_state_source,
+        decoder_inputs=decoder_inputs,
+        decoder_input_ids=decoder_input_ids,
+        decoder_attention_mask=decoder_attention_mask,
+        model_kwargs=model_kwargs,
+    )
+    if _embedding_batch_identity is not None:
+        embedding_context["family_adapter"] = _fingerprint_jsonable(_embedding_batch_identity)
+        if allowed_unsupported_pooling:
+            embedding_context["family_adapter_pooling_override"] = sorted(
+                allowed_unsupported_pooling
+            )
+
+    # A pending automatic attention request settles here, inside the caller's
+    # autocast context, so the fingerprint records the backend that executes.
+    attention_resolution = getattr(model, "attention_resolution", None)
+    if attention_resolution is not None and attention_resolution.deferred:
+        model.resolve_attn_implementation()
+
+    tokenizer_metadata = _tokenizer_metadata(model, tokenizer)
+    (
+        input_fingerprint,
+        run_fingerprint,
+        resolved_model_state_fingerprint,
+        model_state_fingerprint_source,
+    ) = _run_fingerprint(
+        model,
+        records,
+        pooling=pooling_names,
+        full_embeddings=full_embeddings,
+        max_length=max_length,
+        truncate=truncate,
+        dtype=dtype,
+        model_kwargs=model_kwargs,
+        tokenizer_metadata=tokenizer_metadata,
+        model_state_fingerprint=model_state_fingerprint,
+        persist_output=output is not None,
+        embedding_context=embedding_context,
+        batch_size=batch_size,
+        batch_window_size=resolved_batch_window_size,
+        max_tokens_per_batch=max_tokens_per_batch,
+    )
+    destination = EmbeddingOutput(
+        records,
+        output=output,
+        format=format,
+        resume=resume,
+        shard_size=shard_size,
+        run_fingerprint=run_fingerprint,
+        input_fingerprint=input_fingerprint,
+        model_state_fingerprint=resolved_model_state_fingerprint,
+        model_state_fingerprint_source=model_state_fingerprint_source,
+        pooler=pooler,
+        pooling_names=pooling_names,
+    )
+    if destination.completed is not None:
+        return destination.completed
+
+    attention_backend = _attention_backend(model)
+    executor = BatchExecutor(
+        model=model,
+        batch_size=batch_size,
+        max_tokens_per_batch=max_tokens_per_batch,
+        max_length=max_length,
+        truncate=truncate,
+        model_kwargs=model_kwargs,
+        hidden_state_source=hidden_state_source,
+        normalized_decoder_inputs=normalized_decoder_inputs,
+        decoder_input_ids=decoder_input_ids,
+        decoder_attention_mask=decoder_attention_mask,
+        _embedding_batch_fn=_embedding_batch_fn,
+        tokenizer=tokenizer,
+        store_all_hidden_states=store_all_hidden_states,
+        full_embeddings=full_embeddings,
+        dtype=dtype,
+        pooler=pooler,
+        attention_backend=attention_backend,
+        need_attentions="parti" in pooling_names,
+    )
+    pool_slices = destination.pool_slices
+    with _temporary_eval(model), torch.inference_mode():
+        for window_start in range(
+            destination.start_position, len(records), resolved_batch_window_size
+        ):
+            window_stop = min(window_start + resolved_batch_window_size, len(records))
+            window_records = records[window_start:window_stop]
+            if not isinstance(window_records, Sequence):
+                raise RuntimeError("The immutable embedding spool returned a non-sequence window.")
+            new_records, pool_slices = executor.run_window(
+                window_records, window_start=window_start
+            )
+            destination.append(window_start, new_records)
+
     software_versions = identity._software_versions()
-    projection = getattr(model, "embedding_projection", None)
-    resolved_layer = getattr(
-        model,
-        "embedding_layer",
-        model_kwargs.get("hidden_state_index", -1),
-    )
-    token_policy = getattr(
-        model,
-        "embedding_token_policy",
-        {
-            "unit": "residue",
-            "include": ["biological residues"],
-            "exclude": [
-                "BOS",
-                "EOS",
-                "padding",
-                "chain delimiters",
-                "non-protein tokens",
-            ],
-        },
-    )
-    model_identity = _model_identity_metadata(model)
-    metadata: dict[str, Any] = {
-        "format_version": 1,
-        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-        "run_fingerprint": run_fingerprint,
-        "input_fingerprint": input_fingerprint,
-        "model_state_fingerprint": resolved_model_state_fingerprint,
-        "model_state_fingerprint_source": model_state_fingerprint_source,
-        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
-        **model_identity,
-        "dtype": str(dtype).removeprefix("torch.") if dtype is not None else "model",
-        "attention_backend": attention_backend,
-        "attention_kernel": _attention_kernel_metadata(attention_backend),
-        "layer": resolved_layer,
-        "projection": projection,
-        "esmc_source": getattr(model, "_esmc_source", None),
-        "esmc_revision": getattr(model, "_esmc_source_revision", None),
-        "esmc_files": getattr(model, "_esmc_source_files", None),
-        "token_policy": token_policy,
-        "tokenizer": tokenizer_metadata,
-        **embedding_context,
-        "pooling": list(pooling_names),
-        "pool_slices": pool_slices,
-        "full_embeddings": full_embeddings,
-        "max_length": max_length,
-        "truncate": truncate,
-        "truncation": {"enabled": truncate, "max_length": max_length},
-        "batching": {
-            "batch_size": batch_size,
-            "batch_window_size": resolved_batch_window_size,
-            "max_tokens_per_batch": max_tokens_per_batch,
-            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
-            "ordering": "bounded-length-bucketed-stable-output",
-            "resume_commit_granularity": (
-                "not-applicable"
-                if output is None
-                else "batch-window"
-                if format == "sqlite"
-                else "shard-flush"
-            ),
-        },
-        "residue_mask_policy": "biological-residues-only",
-        "record_count": len(records),
-        "descriptor_index": (
-            "memory-metadata"
-            if output is None
-            else "sqlite-records"
-            if format == "sqlite"
-            else "safetensors-generation-index"
-        ),
-        "storage_format": format if output is not None else "memory",
-        "software": software_versions,
-        "execution": _execution_identity_metadata(model),
-        "adapter": _adapter_identity_metadata(model),
-        "torch_version": software_versions["torch"],
-        "transformers_version": software_versions["transformers"],
-        "complete": True,
-    }
-    if destination.output_descriptors is not None:
-        metadata["outputs"] = destination.output_descriptors
-        metadata["tensor_hashes"] = [item["sha256"] for item in destination.output_descriptors]
-    status = getattr(model, "esmc_precision_status", None)
-    if status is not None:
-        metadata["esmc_precision"] = status.as_dict() if hasattr(status, "as_dict") else status
-    return destination.finish(metadata)
-
-
-class EmbeddingMixin:
-    """Small delegation mixin shared by FastPLMs model classes."""
-
-    def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
-        return embed_dataset(self, inputs, **kwargs)
-
-
-__all__ = [
-    "EmbeddingMixin",
-    "embed_dataset",
-    "iter_fasta",
-    "parse_fasta",
-    "select_hidden_state_embeddings",
-]
+    projection = getattr(model, "embedding_projection", None)
+    resolved_layer = getattr(
+        model,
+        "embedding_layer",
+        model_kwargs.get("hidden_state_index", -1),
+    )
+    token_policy = getattr(
+        model,
+        "embedding_token_policy",
+        {
+            "unit": "residue",
+            "include": ["biological residues"],
+            "exclude": [
+                "BOS",
+                "EOS",
+                "padding",
+                "chain delimiters",
+                "non-protein tokens",
+            ],
+        },
+    )
+    model_identity = _model_identity_metadata(model)
+    metadata: dict[str, Any] = {
+        "format_version": 1,
+        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+        "run_fingerprint": run_fingerprint,
+        "input_fingerprint": input_fingerprint,
+        "model_state_fingerprint": resolved_model_state_fingerprint,
+        "model_state_fingerprint_source": model_state_fingerprint_source,
+        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
+        **model_identity,
+        "dtype": str(dtype).removeprefix("torch.") if dtype is not None else "model",
+        "attention_backend": attention_backend,
+        "attention_kernel": _attention_kernel_metadata(attention_backend),
+        "layer": resolved_layer,
+        "projection": projection,
+        "esmc_source": getattr(model, "_esmc_source", None),
+        "esmc_revision": getattr(model, "_esmc_source_revision", None),
+        "esmc_files": getattr(model, "_esmc_source_files", None),
+        "token_policy": token_policy,
+        "tokenizer": tokenizer_metadata,
+        **embedding_context,
+        "pooling": list(pooling_names),
+        "pool_slices": pool_slices,
+        "full_embeddings": full_embeddings,
+        "max_length": max_length,
+        "truncate": truncate,
+        "truncation": {"enabled": truncate, "max_length": max_length},
+        "batching": {
+            "batch_size": batch_size,
+            "batch_window_size": resolved_batch_window_size,
+            "max_tokens_per_batch": max_tokens_per_batch,
+            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
+            "ordering": "bounded-length-bucketed-stable-output",
+            "resume_commit_granularity": (
+                "not-applicable"
+                if output is None
+                else "batch-window"
+                if format == "sqlite"
+                else "shard-flush"
+            ),
+        },
+        "residue_mask_policy": "biological-residues-only",
+        "record_count": len(records),
+        "descriptor_index": (
+            "memory-metadata"
+            if output is None
+            else "sqlite-records"
+            if format == "sqlite"
+            else "safetensors-generation-index"
+        ),
+        "storage_format": format if output is not None else "memory",
+        "software": software_versions,
+        "execution": _execution_identity_metadata(model),
+        "adapter": _adapter_identity_metadata(model),
+        "torch_version": software_versions["torch"],
+        "transformers_version": software_versions["transformers"],
+        "complete": True,
+    }
+    if destination.output_descriptors is not None:
+        metadata["outputs"] = destination.output_descriptors
+        metadata["tensor_hashes"] = [item["sha256"] for item in destination.output_descriptors]
+    status = getattr(model, "esmc_precision_status", None)
+    if status is not None:
+        metadata["esmc_precision"] = status.as_dict() if hasattr(status, "as_dict") else status
+    return destination.finish(metadata)
+
+
+class EmbeddingMixin:
+    """Small delegation mixin shared by FastPLMs model classes."""
+
+    def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
+        return embed_dataset(self, inputs, **kwargs)
+
+
+__all__ = [
+    "EmbeddingMixin",
+    "embed_dataset",
+    "iter_fasta",
+    "parse_fasta",
+    "select_hidden_state_embeddings",
+]

fastplms/embeddings/storage.py

@@ -9,6 +9,7 @@ import sqlite3
 import struct
 import numpy as np
 import torch
+
 from bisect import bisect_right
 from collections.abc import Iterable, Iterator, Sequence
 from pathlib import Path

fastplms/models/_esm_rotary.py

@@ -9,6 +9,7 @@ Transformers implementation.
 from __future__ import annotations
 
 import torch
+
 from torch import nn
 
 

fastplms/models/classification_probe.py

@@ -3,9 +3,9 @@
 from __future__ import annotations
 
 import math
-from typing import Any
-
 import torch
+
+from typing import Any
 from torch import nn
 from torch.nn import functional as F
 from transformers.modeling_outputs import (
@@ -14,6 +14,7 @@ from transformers.modeling_outputs import (
     TokenClassifierOutput,
 )
 
+
 try:
     from fastplms.attention import (
         AttentionBackend,
@@ -145,7 +146,7 @@ def token_classification_loss(
     if problem_type == "regression":
         targets = labels.to(logits.dtype)
         if num_labels == 1 and targets.ndim == logits.ndim - 1:
-            targets = targets.unsqueeze(-1)
+            targets = targets.unsqueeze(-1)  # (..., 1), matching single-target logits
         if targets.shape != logits.shape:
             raise ValueError(
                 "Token regression labels must match logits, except that the final "
@@ -210,7 +211,7 @@ class ProbeSelfAttention(nn.Module):
             sequence_length,
             self.num_heads,
             self.head_size,
-        ).transpose(1, 2)
+        ).transpose(1, 2)  # (b, h, l, d_h)
 
     def forward(
         self,
@@ -220,11 +221,11 @@ class ProbeSelfAttention(nn.Module):
         output_attentions: bool,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         batch_size, sequence_length, _ = hidden_states.shape
-        query, key, value = self.qkv(hidden_states).chunk(3, dim=-1)
-        query = self._reshape(query)
-        key = self._reshape(key)
-        value = self._reshape(value)
-        query, key = self.rotary(query, key)
+        query, key, value = self.qkv(hidden_states).chunk(3, dim=-1)  # each (b, l, d)
+        query = self._reshape(query)  # (b, h, l, d_h)
+        key = self._reshape(key)  # (b, h, l, d_h)
+        value = self._reshape(value)  # (b, h, l, d_h)
+        query, key = self.rotary(query, key)  # each (b, h, l, d_h)
         if output_attentions and self.backend != AttentionBackend.EAGER:
             raise ValueError(
                 f"output_attentions=True is unavailable for {self.backend.value!r}; "
@@ -241,11 +242,11 @@ class ProbeSelfAttention(nn.Module):
         dropout = self.dropout if self.training else 0.0
         attention_weights = None
         if self.backend == AttentionBackend.EAGER:
-            scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_size)
+            scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_size)  # (b, h, l, l)
             if attention_mask_4d is not None:
-                scores = scores.masked_fill(~attention_mask_4d, float("-inf"))
-            attention_weights = scores.softmax(dim=-1)
-            context = F.dropout(attention_weights, p=dropout, training=self.training) @ value
+                scores = scores.masked_fill(~attention_mask_4d, float("-inf"))  # (b, h, l, l)
+            attention_weights = scores.softmax(dim=-1)  # (b, h, l, l)
+            context = F.dropout(attention_weights, p=dropout, training=self.training) @ value  # (b, h, l, d_h)
         elif self.backend == AttentionBackend.SDPA:
             context = F.scaled_dot_product_attention(
                 query,
@@ -253,7 +254,7 @@ class ProbeSelfAttention(nn.Module):
                 value,
                 attn_mask=attention_mask_4d,
                 dropout_p=dropout,
-            )
+            )  # (b, h, l, d_h)
         elif self.backend == AttentionBackend.FLEX_ATTENTION:
             if flex_attention is None:
                 raise RuntimeError("'flex_attention' was requested but is unavailable.")
@@ -272,15 +273,15 @@ class ProbeSelfAttention(nn.Module):
                 block_mask=flex_block_mask,
                 scale=1.0 / math.sqrt(self.head_size),
                 kernel_options={"PRESCALE_QK": True, "BLOCK_N": 32},
-            )
+            )  # (b, h, l, d_h)
         else:
             raise AssertionError(f"Unhandled attention backend {self.backend.value!r}.")
         context = context.transpose(1, 2).contiguous().view(
             batch_size,
             sequence_length,
             self.hidden_size,
-        )
-        return self.output(context), attention_weights
+        )  # (b, l, d)
+        return self.output(context), attention_weights  # (b, l, d), optional (b, h, l, l)
 
 
 class ProteinTransformerProbe(nn.Module):
@@ -451,7 +452,7 @@ class SequenceClassificationProbe(_ClassificationProbe):
                 embeddings.shape[:2],
                 device=embeddings.device,
                 dtype=torch.bool,
-            )
+            )  # (b, l)
         outputs = self._forward_transformer(
             embeddings,
             attention_mask,
@@ -460,8 +461,8 @@ class SequenceClassificationProbe(_ClassificationProbe):
         )
         if self.pooler is None:
             raise AssertionError("Sequence classification requires a configured pooler.")
-        pooled = self.pooler(outputs.last_hidden_state, attention_mask)
-        logits = self.classifier(pooled)
+        pooled = self.pooler(outputs.last_hidden_state, attention_mask)  # (b, d)
+        logits = self.classifier(pooled)  # (b, num_labels)
         loss = None
         if labels is not None:
             problem_type = resolve_problem_type(self.config, labels, num_labels=self.num_labels)

fastplms/models/esm_plusplus/modeling_esm_plusplus.py

@@ -6,15 +6,15 @@ import importlib
 import importlib.metadata
 import math
 import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
 from collections.abc import Sequence
 from contextlib import contextmanager
 from dataclasses import asdict, dataclass
 from functools import partial
 from typing import Any, ClassVar
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
 from einops import rearrange
 from tokenizers import Tokenizer
 from tokenizers.models import BPE
@@ -378,7 +378,7 @@ class RotaryEmbedding(torch.nn.Module):
         inv_freq = self._compute_inv_freq(buffer_device)
         self._clear_cache()
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        arange = torch.arange(0, self.dim, 2, device=buffer_device, dtype=torch.float32)
+        arange = torch.arange(0, self.dim, 2, device=buffer_device, dtype=torch.float32)  # (d / 2,)
         scale = (
             (arange + 0.4 * self.dim) / (1.4 * self.dim) if self.scale_base is not None else None
         )
@@ -446,18 +446,18 @@ class RotaryEmbedding(torch.nn.Module):
         cos_angles = torch.cos(angles)  # (l, d / 2)
         sin_angles = torch.sin(angles)  # (l, d / 2)
         if self.scale is None:
-            self._cos_cached = cos_angles.to(dtype)
-            self._sin_cached = sin_angles.to(dtype)
-            self._cos_full_cached = torch.cat((self._cos_cached, self._cos_cached), dim=-1)
-            self._sin_full_cached = torch.cat((self._sin_cached, self._sin_cached), dim=-1)
+            self._cos_cached = cos_angles.to(dtype)  # (l, d / 2)
+            self._sin_cached = sin_angles.to(dtype)  # (l, d / 2)
+            self._cos_full_cached = torch.cat((self._cos_cached, self._cos_cached), dim=-1)  # (l, d)
+            self._sin_full_cached = torch.cat((self._sin_cached, self._sin_cached), dim=-1)  # (l, d)
             return
 
         centered_positions = (
             torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
-        ) / self.scale_base
-        scale = self.scale ** centered_positions.unsqueeze(-1)
-        self._cos_cached = (cos_angles * scale).to(dtype)
-        self._sin_cached = (sin_angles * scale).to(dtype)
+        ) / self.scale_base  # (l,)
+        scale = self.scale ** centered_positions.unsqueeze(-1)  # (l, d / 2)
+        self._cos_cached = (cos_angles * scale).to(dtype)  # (l, d / 2)
+        self._sin_cached = (sin_angles * scale).to(dtype)  # (l, d / 2)
         self._cos_k_cached = (cos_angles / scale).to(dtype)
         self._sin_k_cached = (sin_angles / scale).to(dtype)
 
@@ -574,12 +574,12 @@ class MultiHeadAttention(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
         # x: (b, l, d)
         qkv = self.layernorm_qkv(x)  # (b, l, 3 * d)
-        query_sequence, key_sequence, value_sequence = torch.chunk(qkv, 3, dim=-1)
+        query_sequence, key_sequence, value_sequence = torch.chunk(qkv, 3, dim=-1)  # each (b, l, d)
         query_sequence, key_sequence = (
             self.q_ln(query_sequence).to(query_sequence.dtype),
             self.k_ln(key_sequence).to(query_sequence.dtype),
-        )
-        query_sequence, key_sequence = self._apply_rotary(query_sequence, key_sequence)
+        )  # each (b, l, d)
+        query_sequence, key_sequence = self._apply_rotary(query_sequence, key_sequence)  # each (b, l, d)
         query_heads, key_heads, value_heads = map(
             self.reshaper, (query_sequence, key_sequence, value_sequence)
         )  # each (b, h, l, d_h)
@@ -596,7 +596,7 @@ class MultiHeadAttention(nn.Module):
             flash_padding_layout=flash_padding_layout,
         )
 
-        output = self.out_proj(attn_output)
+        output = self.out_proj(attn_output)  # (b, l, d)
         return output, attn_weights, s_max
 
     def _attn(
@@ -682,9 +682,9 @@ class MultiHeadAttention(nn.Module):
         attention_mask_2d: torch.Tensor | None = None,
         flash_padding_layout: FlashPaddingLayout | None = None,
     ) -> tuple[torch.Tensor, None]:
-        query_tokens = query_heads.transpose(1, 2).contiguous()
-        key_tokens = key_heads.transpose(1, 2).contiguous()
-        value_tokens = value_heads.transpose(1, 2).contiguous()
+        query_tokens = query_heads.transpose(1, 2).contiguous()  # (b, l, h, d_h)
+        key_tokens = key_heads.transpose(1, 2).contiguous()  # (b, l, h, d_h)
+        value_tokens = value_heads.transpose(1, 2).contiguous()  # (b, l, h, d_h)
         attn_output = kernels_flash_attention_func(
             query_states=query_tokens,
             key_states=key_tokens,
@@ -693,8 +693,8 @@ class MultiHeadAttention(nn.Module):
             causal=False,
             implementation=self.attn_backend.value,
             padding_layout=flash_padding_layout,
-        )
-        return rearrange(attn_output, "b s h d -> b s (h d)"), None
+        )  # (b, l, h, d_h)
+        return rearrange(attn_output, "b s h d -> b s (h d)"), None  # (b, l, h * d_h), None
 
     def _flex_attn(
         self,
@@ -1015,11 +1015,11 @@ class TransformerStack(nn.Module):
             # finite without allowing their states to enter residue attention.
             attention_mask_4d = (
                 mask_pattern[:, None, :, None] == mask_pattern[:, None, None, :]
-            )
+            )  # (b, 1, l, l)
         else:
             attention_mask_4d = (
                 mask_pattern.unsqueeze(-1) == mask_pattern.unsqueeze(-2)
-            ).unsqueeze(1)
+            ).unsqueeze(1)  # (b, 1, l, l)
         backend = (
             resolve_attention_backend_for_call(
                 self.attention_backend,
@@ -1837,7 +1837,7 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixi
                     inputs_embeds.shape[:2],
                     dtype=torch.bool,
                     device=inputs_embeds.device,
-                )
+                )  # (b, l)
 
         output = super().forward(
             input_ids=input_ids,

fastplms/models/esmfold2/__init__.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 from importlib import import_module
 from typing import TYPE_CHECKING, Any
 
+
 if TYPE_CHECKING:
     from .configuration_esmfold2 import ESMFold2Config as ESMFold2Config
     from .modeling_esmfold2 import ESMFold2Model as ESMFold2Model

fastplms/models/esmfold2/configuration_esmfold2.py

@@ -18,7 +18,6 @@ from __future__ import annotations
 
 from dataclasses import asdict, dataclass, field
 from typing import Any, TypeVar, cast
-
 from transformers.configuration_utils import PretrainedConfig
 
 from fastplms.attention import canonical_checkpoint_attention_backend

fastplms/models/esmfold2/embedding.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import torch
+
 from typing import Any, ClassVar
 from torch import Tensor
 

fastplms/models/esmfold2/esmfold2_affine3d.py

@@ -2,10 +2,11 @@
 
 from __future__ import annotations
 
+import torch
+
+from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any, Self
-
-import torch
 from torch.nn import functional as F
 
 from .esmfold2_misc import fp32_autocast_context
@@ -19,23 +20,26 @@ def _index_tuple(index: Any) -> tuple[Any, ...]:
 
 def _sqrt_subgradient(values: torch.Tensor) -> torch.Tensor:
     """Square root with a zero subgradient for non-positive inputs."""
+    # values: arbitrary shape; every element keeps its position.
 
-    result = torch.zeros_like(values)
-    positive = values > 0
-    result[positive] = torch.sqrt(values[positive])
-    return result
+    result = torch.zeros_like(values)  # values.shape
+    positive = values > 0  # values.shape
+    result[positive] = torch.sqrt(values[positive])  # (n_positive,) selected elements
+    return result  # values.shape
 
 
 def _quat_invert(quaternion: torch.Tensor) -> torch.Tensor:
-    conjugate_sign = torch.tensor([1, -1, -1, -1], device=quaternion.device)
-    return quaternion * conjugate_sign
+    # quaternion: (..., 4), real component first.
+    conjugate_sign = torch.tensor([1, -1, -1, -1], device=quaternion.device)  # (4,)
+    return quaternion * conjugate_sign  # (..., 4)
 
 
 def _quat_mult(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
     """Hamilton product for real-first quaternion tensors."""
+    # left, right: (..., 4); their leading dimensions must broadcast.
 
-    aw, ax, ay, az = torch.unbind(left, -1)
-    bw, bx, by, bz = torch.unbind(right, -1)
+    aw, ax, ay, az = torch.unbind(left, -1)  # each left.shape[:-1]
+    bw, bx, by, bz = torch.unbind(right, -1)  # each right.shape[:-1]
     return torch.stack(
         (
             aw * bw - ax * bx - ay * by - az * bz,
@@ -44,7 +48,7 @@ def _quat_mult(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
             aw * bz + ax * by - ay * bx + az * bw,
         ),
         -1,
-    )
+    )  # (..., 4), broadcast leading dimensions
 
 
 def _quat_rotation(
@@ -52,9 +56,10 @@ def _quat_rotation(
     points: torch.Tensor,
 ) -> torch.Tensor:
     """Rotate points using normalized real-first quaternions."""
+    # quaternion: (..., 4); points: (..., 3); leading dimensions broadcast.
 
-    aw, ax, ay, az = torch.unbind(quaternion, -1)
-    bx, by, bz = torch.unbind(points, -1)
+    aw, ax, ay, az = torch.unbind(quaternion, -1)  # each quaternion.shape[:-1]
+    bx, by, bz = torch.unbind(points, -1)  # each points.shape[:-1]
     product = torch.stack(
         (
             -ax * bx - ay * by - az * bz,
@@ -63,8 +68,8 @@ def _quat_rotation(
             aw * bz + ax * by - ay * bx,
         ),
         -1,
-    )
-    return _quat_mult(product, _quat_invert(quaternion))[..., 1:]
+    )  # (..., 4), broadcast leading dimensions
+    return _quat_mult(product, _quat_invert(quaternion))[..., 1:]  # (..., 3)
 
 
 def _graham_schmidt(
@@ -73,17 +78,18 @@ def _graham_schmidt(
     eps: float = 1e-12,
 ) -> torch.Tensor:
     """Construct a right-handed orthonormal frame from two directions."""
+    # x_axis, xy_plane: (..., 3); ... contains arbitrary frame batch axes.
 
     with fp32_autocast_context(x_axis.device.type):
-        e1 = xy_plane
-        denominator = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps)
-        x_axis = x_axis / denominator
-        projection = (x_axis * e1).sum(dim=-1, keepdim=True)
-        e1 = e1 - x_axis * projection
-        denominator = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps)
-        e1 = e1 / denominator
-        e2 = torch.cross(x_axis, e1, dim=-1)
-        return torch.stack([x_axis, e1, e2], dim=-1)
+        e1 = xy_plane  # (..., 3)
+        denominator = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps)  # (..., 1)
+        x_axis = x_axis / denominator  # (..., 3)
+        projection = (x_axis * e1).sum(dim=-1, keepdim=True)  # (..., 1)
+        e1 = e1 - x_axis * projection  # (..., 3)
+        denominator = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps)  # (..., 1)
+        e1 = e1 / denominator  # (..., 3)
+        e2 = torch.cross(x_axis, e1, dim=-1)  # (..., 3)
+        return torch.stack([x_axis, e1, e2], dim=-1)  # (..., 3, 3)
 
 
 class Rotation:
@@ -129,6 +135,7 @@ class Rotation:
 
     @classmethod
     def _from_tensor(cls, tensor: torch.Tensor) -> Self:
+        # tensor: (..., r), where r is the concrete rotation representation width.
         return cls(tensor)  # type: ignore[call-arg]
 
     def to(self, **kwargs) -> Self:
@@ -137,15 +144,17 @@ class Rotation:
     def detach(self, *args, **kwargs) -> Self:
         return self._from_tensor(self.tensor.detach(**kwargs))
 
-    def tensor_apply(self, func) -> Self:
-        transformed = [func(component) for component in self.tensor.unbind(dim=-1)]
-        return self._from_tensor(torch.stack(transformed, dim=-1))
+    def tensor_apply(self, func: Callable[[torch.Tensor], torch.Tensor]) -> Self:
+        # self.tensor: (..., r); func receives one (...) component and determines its output shape.
+        transformed = [func(component) for component in self.tensor.unbind(dim=-1)]  # one func-shaped array per rotation component
+        return self._from_tensor(torch.stack(transformed, dim=-1))  # rotation tensor: (*func_output_shape, n_components)
 
 
 class RotationQuat(Rotation):
     """A rotation represented by a real-first quaternion."""
 
-    def __init__(self, quats: torch.Tensor, normalized: bool = False):
+    def __init__(self, quats: torch.Tensor, normalized: bool = False) -> None:
+        # quats: (..., 4); ... is the rotation batch shape.
         if not isinstance(quats, torch.Tensor):
             raise TypeError("quats must be a Torch tensor.")
         if quats.ndim == 0 or quats.shape[-1] != 4:
@@ -156,32 +165,32 @@ class RotationQuat(Rotation):
             raise TypeError("normalized must be a boolean.")
         self._normalized = normalized
         if normalized:
-            quats = F.normalize(quats.to(torch.float32), dim=-1)
-            self._quats = quats.where(quats[..., :1] >= 0, -quats)
+            quats = F.normalize(quats.to(torch.float32), dim=-1)  # (..., 4)
+            self._quats = quats.where(quats[..., :1] >= 0, -quats)  # (..., 4)
         else:
-            self._quats = quats.to(torch.float32)
+            self._quats = quats.to(torch.float32)  # (..., 4)
 
     @property
     def tensor(self) -> torch.Tensor:
-        return self._quats
+        return self._quats  # (..., 4)
 
     @property
     def shape(self) -> torch.Size:
         return self._quats.shape[:-1]
 
     @classmethod
-    def identity(cls, shape, **tensor_kwargs) -> RotationQuat:
-        quaternions = torch.ones((*shape, 4), **tensor_kwargs)
-        selector = torch.tensor([1, 0, 0, 0], device=quaternions.device)
-        return cls(quaternions * selector)
+    def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> RotationQuat:
+        quaternions = torch.ones((*shape, 4), **tensor_kwargs)  # (*shape, 4)
+        selector = torch.tensor([1, 0, 0, 0], device=quaternions.device)  # (4,)
+        return cls(quaternions * selector)  # quaternion tensor: (*shape, 4)
 
     @classmethod
-    def random(cls, shape, **tensor_kwargs) -> RotationQuat:
-        return cls(torch.randn((*shape, 4), **tensor_kwargs), normalized=True)
+    def random(cls, shape: tuple[int, ...], **tensor_kwargs) -> RotationQuat:
+        return cls(torch.randn((*shape, 4), **tensor_kwargs), normalized=True)  # quaternion tensor: (*shape, 4)
 
     def __getitem__(self, idx: Any) -> RotationQuat:
         indices = _index_tuple(idx)
-        return RotationQuat(self._quats[(*indices, slice(None))])
+        return RotationQuat(self._quats[(*indices, slice(None))])  # quaternion tensor: (*indexed_shape, 4)
 
     def normalized(self) -> RotationQuat:
         if self._normalized:
@@ -192,9 +201,9 @@ class RotationQuat(Rotation):
         return self
 
     def as_matrix(self) -> RotationMatrix:
-        quaternion = self.normalized().tensor
-        r, i, j, k = torch.unbind(quaternion, -1)
-        scale = 2.0 / torch.linalg.norm(quaternion, dim=-1)
+        quaternion = self.normalized().tensor  # (..., 4)
+        r, i, j, k = torch.unbind(quaternion, -1)  # each (...)
+        scale = 2.0 / torch.linalg.norm(quaternion, dim=-1)  # (...)
         elements = torch.stack(
             (
                 1 - scale * (j * j + k * k),
@@ -208,54 +217,56 @@ class RotationQuat(Rotation):
                 1 - scale * (i * i + j * j),
             ),
             -1,
-        )
-        return RotationMatrix(elements.reshape((*quaternion.shape[:-1], 3, 3)))
+        )  # (..., 9)
+        return RotationMatrix(elements.reshape((*quaternion.shape[:-1], 3, 3)))  # rotation tensor: (..., 3, 3)
 
     def compose(self, other: RotationQuat) -> RotationQuat:
         with fp32_autocast_context(self.device.type):
-            return RotationQuat(_quat_mult(self._quats, other._quats))
+            return RotationQuat(_quat_mult(self._quats, other._quats))  # quaternion tensor: (..., 4), broadcast leading shapes
 
     def convert_compose(self, other: Rotation) -> RotationQuat:
         return self.compose(other.as_quat())
 
     def apply(self, points: torch.Tensor) -> torch.Tensor:
-        return _quat_rotation(self.normalized()._quats, points)
+        # points: (..., 3); quaternion and point leading shapes must broadcast.
+        return _quat_rotation(self.normalized()._quats, points)  # (..., 3), broadcast rotation/point leading shapes
 
     def invert(self) -> RotationQuat:
-        return RotationQuat(_quat_invert(self._quats))
+        return RotationQuat(_quat_invert(self._quats))  # quaternion tensor: (..., 4)
 
 
 class RotationMatrix(Rotation):
     """A rotation represented by a dense FP32 matrix."""
 
-    def __init__(self, rots: torch.Tensor):
+    def __init__(self, rots: torch.Tensor) -> None:
+        # rots: (..., 9) or (..., 3, 3); ... is the rotation batch shape.
         if not isinstance(rots, torch.Tensor):
             raise TypeError("rots must be a Torch tensor.")
         if rots.ndim > 0 and rots.shape[-1] == 9:
-            rots = rots.unflatten(-1, (3, 3))
+            rots = rots.unflatten(-1, (3, 3))  # (..., 3, 3)
         if rots.ndim < 2 or rots.shape[-2:] != (3, 3):
             raise ValueError(
                 "rots must have trailing shape (3, 3) or flattened width 9, got "
                 f"shape {tuple(rots.shape)}."
             )
-        self._rots = rots.to(torch.float32)
+        self._rots = rots.to(torch.float32)  # (..., 3, 3)
 
     @property
     def tensor(self) -> torch.Tensor:
-        return self._rots.flatten(-2)
+        return self._rots.flatten(-2)  # (..., 9)
 
     @property
     def shape(self) -> torch.Size:
         return self._rots.shape[:-2]
 
     @classmethod
-    def identity(cls, shape, **tensor_kwargs) -> RotationMatrix:
-        matrix = torch.eye(3, **tensor_kwargs)
-        matrix = matrix.view(*(1 for _ in shape), 3, 3)
-        return cls(matrix.expand(*shape, -1, -1))
+    def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> RotationMatrix:
+        matrix = torch.eye(3, **tensor_kwargs)  # (3, 3)
+        matrix = matrix.view(*(1 for _ in shape), 3, 3)  # (*singleton_shape, 3, 3); one singleton per requested axis
+        return cls(matrix.expand(*shape, -1, -1))  # rotation tensor: (*shape, 3, 3)
 
     @classmethod
-    def random(cls, shape, **tensor_kwargs) -> RotationMatrix:
+    def random(cls, shape: tuple[int, ...], **tensor_kwargs) -> RotationMatrix:
         return RotationQuat.random(shape, **tensor_kwargs).as_matrix()
 
     @staticmethod
@@ -264,23 +275,24 @@ class RotationMatrix(Rotation):
         xy_plane: torch.Tensor,
         eps: float = 1e-12,
     ) -> RotationMatrix:
-        return RotationMatrix(_graham_schmidt(x_axis, xy_plane, eps))
+        # x_axis, xy_plane: (..., 3).
+        return RotationMatrix(_graham_schmidt(x_axis, xy_plane, eps))  # rotation tensor: (..., 3, 3)
 
     def __getitem__(self, idx: Any) -> RotationMatrix:
         indices = _index_tuple(idx)
-        return RotationMatrix(self._rots[(*indices, slice(None), slice(None))])
+        return RotationMatrix(self._rots[(*indices, slice(None), slice(None))])  # rotation tensor: (*indexed_shape, 3, 3)
 
     def as_matrix(self) -> RotationMatrix:
         return self
 
     def to_3x3(self) -> torch.Tensor:
-        return self._rots
+        return self._rots  # (..., 3, 3)
 
     def as_quat(self, normalize: bool = False) -> RotationQuat:
         m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
             self._rots.flatten(-2),
             dim=-1,
-        )
+        )  # each (...)
         q_abs = _sqrt_subgradient(
             torch.stack(
                 (
@@ -291,7 +303,7 @@ class RotationMatrix(Rotation):
                 ),
                 dim=-1,
             )
-        )
+        )  # (..., 4)
         products = torch.stack(
             (
                 q_abs[..., 0] ** 2,
@@ -312,29 +324,30 @@ class RotationMatrix(Rotation):
                 q_abs[..., 3] ** 2,
             ),
             dim=-1,
-        ).unflatten(-1, (4, 4))
-        floor = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
-        candidates = products / (2.0 * q_abs[..., None].max(floor))
-        best = torch.zeros_like(q_abs, dtype=torch.bool)
-        best.scatter_(-1, q_abs.argmax(dim=-1, keepdim=True), True)
-        quaternion = candidates[best, :].reshape(q_abs.shape)
-        return RotationQuat(quaternion)
+        ).unflatten(-1, (4, 4))  # (..., 4, 4)
+        floor = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)  # ()
+        candidates = products / (2.0 * q_abs[..., None].max(floor))  # (..., 4, 4)
+        best = torch.zeros_like(q_abs, dtype=torch.bool)  # (..., 4)
+        best.scatter_(-1, q_abs.argmax(dim=-1, keepdim=True), True)  # (..., 4); selected index shape (..., 1)
+        quaternion = candidates[best, :].reshape(q_abs.shape)  # (..., 4)
+        return RotationQuat(quaternion)  # quaternion tensor: (..., 4)
 
     def compose(self, other: RotationMatrix) -> RotationMatrix:
         with fp32_autocast_context(self.device.type):
-            return RotationMatrix(self._rots @ other._rots)
+            return RotationMatrix(self._rots @ other._rots)  # rotation tensor: (..., 3, 3), broadcast leading shapes
 
     def convert_compose(self, other: Rotation) -> RotationMatrix:
         return self.compose(other.as_matrix())
 
     def apply(self, points: torch.Tensor) -> torch.Tensor:
+        # points: (..., 3); matrix leading shapes broadcast with point axes.
         with fp32_autocast_context(self.device.type):
             if self._rots.shape[-3] == 1:
-                return points @ self._rots.transpose(-1, -2).squeeze(-3)
-            return torch.einsum("...ij,...j", self._rots, points)
+                return points @ self._rots.transpose(-1, -2).squeeze(-3)  # (..., 3), with leading axes set by matmul broadcasting
+            return torch.einsum("...ij,...j", self._rots, points)  # (..., 3), broadcast leading shapes
 
     def invert(self) -> RotationMatrix:
-        return RotationMatrix(self._rots.transpose(-1, -2))
+        return RotationMatrix(self._rots.transpose(-1, -2))  # rotation tensor: (..., 3, 3)
 
 
 @dataclass(frozen=True)
@@ -378,7 +391,7 @@ class Affine3D:
 
     @property
     def tensor(self) -> torch.Tensor:
-        return torch.cat((self.rot.tensor, self.trans), dim=-1)
+        return torch.cat((self.rot.tensor, self.trans), dim=-1)  # (..., r + 3); r is 4 for quaternions or 9 for matrices
 
     @staticmethod
     def identity(
@@ -409,12 +422,13 @@ class Affine3D:
         rotation_type: type[Rotation] = RotationMatrix,
         **tensor_kwargs,
     ) -> Affine3D:
-        translation = torch.randn((*shape, 3), **tensor_kwargs).mul(std)
-        rotation = rotation_type.random(shape, **tensor_kwargs)
+        translation = torch.randn((*shape, 3), **tensor_kwargs).mul(std)  # (*shape, 3)
+        rotation = rotation_type.random(shape, **tensor_kwargs)  # rotation batch shape: shape
         return Affine3D(trans=translation, rot=rotation)
 
     @staticmethod
     def from_tensor(tensor: torch.Tensor) -> Affine3D:
+        # tensor: (..., 6/7/12), (..., 3, 4), or (..., 4, 4); ... is frame batch shape.
         if not isinstance(tensor, torch.Tensor):
             raise TypeError("tensor must be a Torch tensor.")
         if tensor.ndim == 0:
@@ -426,17 +440,17 @@ class Affine3D:
                     "matrix-form affine tensors must have trailing shape (3, 4) or "
                     f"(4, 4), got {tuple(tensor.shape)}."
                 )
-            translation = tensor[..., :3, 3]
-            rotation: Rotation = RotationMatrix(tensor[..., :3, :3])
+            translation = tensor[..., :3, 3]  # (..., 3)
+            rotation: Rotation = RotationMatrix(tensor[..., :3, :3])  # rotation tensor: (..., 3, 3)
         elif width == 6:
-            translation = tensor[..., -3:]
-            rotation = RotationQuat(F.pad(tensor[..., :3], (1, 0), value=1))
+            translation = tensor[..., -3:]  # (..., 3)
+            rotation = RotationQuat(F.pad(tensor[..., :3], (1, 0), value=1))  # quaternion tensor: (..., 4)
         elif width == 7:
-            translation = tensor[..., -3:]
-            rotation = RotationQuat(tensor[..., :4])
+            translation = tensor[..., -3:]  # (..., 3)
+            rotation = RotationQuat(tensor[..., :4])  # quaternion tensor: (..., 4)
         elif width == 12:
-            translation = tensor[..., -3:]
-            rotation = RotationMatrix(tensor[..., :-3].unflatten(-1, (3, 3)))
+            translation = tensor[..., -3:]  # (..., 3)
+            rotation = RotationMatrix(tensor[..., :-3].unflatten(-1, (3, 3)))  # rotation tensor: (..., 3, 3)
         else:
             raise RuntimeError(
                 f"Cannot detect rotation format from {tensor.shape[-1] - 3}-d flat vector"
@@ -448,6 +462,7 @@ class Affine3D:
         translation: torch.Tensor,
         rotation: torch.Tensor,
     ) -> Affine3D:
+        # translation: (..., 3); rotation: (..., 3, 3) or (..., 9).
         return Affine3D(translation, RotationMatrix(rotation))
 
     @staticmethod
@@ -457,13 +472,14 @@ class Affine3D:
         xy_plane: torch.Tensor,
         eps: float = 1e-10,
     ) -> Affine3D:
-        x_axis = origin - neg_x_axis
-        plane_direction = xy_plane - origin
+        # neg_x_axis, origin, xy_plane: (..., 3), with broadcast-compatible leading axes.
+        x_axis = origin - neg_x_axis  # (..., 3)
+        plane_direction = xy_plane - origin  # (..., 3)
         rotation = RotationMatrix.from_graham_schmidt(
             x_axis,
             plane_direction,
             eps,
-        )
+        )  # rotation tensor: (..., 3, 3)
         return Affine3D(trans=origin, rot=rotation)
 
     @staticmethod
@@ -478,7 +494,7 @@ class Affine3D:
 
     def __getitem__(self, idx: Any) -> Affine3D:
         indices = _index_tuple(idx)
-        translation = self.trans[(*indices, slice(None))]
+        translation = self.trans[(*indices, slice(None))]  # (*indexed_shape, 3)
         return Affine3D(trans=translation, rot=self.rot[idx])
 
     def to(self, **kwargs) -> Affine3D:
@@ -490,9 +506,10 @@ class Affine3D:
             self.rot.detach(**kwargs),
         )
 
-    def tensor_apply(self, func) -> Affine3D:
-        components = [func(value) for value in self.tensor.unbind(dim=-1)]
-        return Affine3D.from_tensor(torch.stack(components, dim=-1))
+    def tensor_apply(self, func: Callable[[torch.Tensor], torch.Tensor]) -> Affine3D:
+        # self.tensor: (..., r + 3); func maps each (...) component to a common output shape.
+        components = [func(value) for value in self.tensor.unbind(dim=-1)]  # one func-shaped array per affine component
+        return Affine3D.from_tensor(torch.stack(components, dim=-1))  # affine tensor: (*func_output_shape, r + 3)
 
     def as_matrix(self) -> Affine3D:
         return Affine3D(trans=self.trans, rot=self.rot.as_matrix())
@@ -509,8 +526,8 @@ class Affine3D:
         autoconvert: bool = False,
     ) -> Affine3D:
         compose_rotation = self.rot.convert_compose if autoconvert else self.rot.compose
-        rotation = compose_rotation(other.rot)
-        translation = self.rot.apply(other.trans) + self.trans
+        rotation = compose_rotation(other.rot)  # rotation batch shape: broadcast(self.shape, other.shape)
+        translation = self.rot.apply(other.trans) + self.trans  # (..., 3), broadcast transform shapes
         return Affine3D(trans=translation, rot=rotation)
 
     def compose_rotation(
@@ -522,28 +539,31 @@ class Affine3D:
         return Affine3D(trans=self.trans, rot=compose(other))
 
     def scale(self, value: torch.Tensor | float) -> Affine3D:
+        # value must broadcast with translation (..., 3); rotation shape is retained.
         return Affine3D(self.trans * value, self.rot)
 
     def mask(self, mask: torch.Tensor, with_zero: bool = False) -> Affine3D:
+        # mask: self.shape; affine tensor: (*self.shape, r + 3).
         if with_zero:
             masked = torch.zeros_like(self.tensor).where(
                 mask[..., None],
                 self.tensor,
-            )
+            )  # (..., r + 3)
             return Affine3D.from_tensor(masked)
         identity = self.identity(
             self.shape,
             rotation_type=type(self.rot),
             device=self.device,
             dtype=self.dtype,
-        ).tensor
-        return Affine3D.from_tensor(identity.where(mask[..., None], self.tensor))
+        ).tensor  # (..., r + 3)
+        return Affine3D.from_tensor(identity.where(mask[..., None], self.tensor))  # affine batch shape: self.shape
 
     def apply(self, points: torch.Tensor) -> torch.Tensor:
-        return self.rot.apply(points) + self.trans
+        # points: (..., 3); result uses broadcast transform/point leading axes.
+        return self.rot.apply(points) + self.trans  # (..., 3), broadcast transform/point leading shapes
 
     def invert(self) -> Affine3D:
-        rotation = self.rot.invert()
+        rotation = self.rot.invert()  # rotation batch shape: self.shape
         return Affine3D(trans=-rotation.apply(self.trans), rot=rotation)
 
 
@@ -551,6 +571,7 @@ def build_affine3d_from_coordinates(
     coords: torch.Tensor,
 ) -> tuple[Affine3D, torch.Tensor]:
     """Build residue frames from X with shape (b, l, 3, 3)."""
+    # coords: (b, l, 3, 3); final axes are N/CA/C atoms and xyz coordinates.
 
     if not isinstance(coords, torch.Tensor):
         raise TypeError("coords must be a Torch tensor.")
@@ -567,39 +588,40 @@ def build_affine3d_from_coordinates(
             dim=-1,
         ),
         dim=-1,
-    )
+    )  # (b, l)
 
     def backbone_affine(positions: torch.Tensor) -> Affine3D:
-        n, ca, c = positions.unbind(dim=-2)
-        return Affine3D.from_graham_schmidt(c, ca, n)
+        # positions: (..., 3, 3); final axes are N/CA/C atoms and xyz coordinates.
+        n, ca, c = positions.unbind(dim=-2)  # each (..., 3)
+        return Affine3D.from_graham_schmidt(c, ca, n)  # affine batch shape: positions.shape[:-2]
 
-    coords = coords.clone().float()
-    coords[~coord_mask] = 0
+    coords = coords.clone().float()  # (b, l, 3, 3)
+    coords[~coord_mask] = 0  # (n_missing, 3, 3) selected residue coordinates
     average = coords.masked_fill(~coord_mask[..., None, None], 0).sum(1) / (
         coord_mask.sum(-1)[..., None, None] + 1e-8
-    )
-    average_affine = backbone_affine(average.float()).as_matrix()
+    )  # (b, 3, 3)
+    average_affine = backbone_affine(average.float()).as_matrix()  # affine batch shape: (b,)
 
     b, length, _, _ = coords.shape
-    rotation = average_affine.rot.tensor[..., None, :].expand(b, length, 9)
-    translation = average_affine.trans[..., None, :].expand(b, length, 3)
+    rotation = average_affine.rot.tensor[..., None, :].expand(b, length, 9)  # (b, l, 9)
+    translation = average_affine.trans[..., None, :].expand(b, length, 3)  # (b, l, 3)
     identity = RotationMatrix.identity(
         (b, length),
         dtype=torch.float32,
         device=coords.device,
         requires_grad=False,
-    )
+    )  # rotation batch shape: (b, l)
     rotation = rotation.where(
         coord_mask.any(-1)[..., None, None],
         identity.tensor,
-    )
-    missing_frame = Affine3D(translation, RotationMatrix(rotation))
+    )  # (b, l, 9)
+    missing_frame = Affine3D(translation, RotationMatrix(rotation))  # affine batch shape: (b, l)
 
-    residue_frame = backbone_affine(coords.float())
+    residue_frame = backbone_affine(coords.float())  # affine batch shape: (b, l)
     residue_frame = Affine3D.from_tensor(
         residue_frame.tensor.where(
             coord_mask[..., None],
             missing_frame.tensor,
         )
-    )
-    return residue_frame, coord_mask
+    )  # affine batch shape: (b, l)
+    return residue_frame, coord_mask  # affine batch shape (b, l), mask (b, l)

fastplms/models/esmfold2/esmfold2_aligner.py

@@ -2,11 +2,11 @@
 
 from __future__ import annotations
 
-from dataclasses import Field, replace
-from typing import Any, ClassVar, Protocol, TypeVar
-
 import numpy as np
 import torch
+
+from dataclasses import Field, replace
+from typing import Any, ClassVar, Protocol, TypeVar
 from torch import Tensor
 
 from .esmfold2_protein_structure import compute_affine_and_rmsd
@@ -30,18 +30,19 @@ AlignableT = TypeVar("AlignableT", bound=Alignable)
 
 
 def _coordinate_batch(structure: Alignable) -> Tensor:
-    return torch.as_tensor(structure.atom37_positions, dtype=torch.double).unsqueeze(0)
+    # l is the residue count; the atom37 table has three Cartesian coordinates.
+    return torch.as_tensor(structure.atom37_positions, dtype=torch.double).unsqueeze(0)  # (1, l, 37, 3)
 
 
 def _shared_atom_mask(mobile: Alignable, target: Alignable, backbone_only: bool) -> Tensor:
     shared = np.asarray(mobile.atom37_mask, dtype=bool) & np.asarray(
         target.atom37_mask,
         dtype=bool,
-    )
+    )  # (l, 37)
     if backbone_only:
-        shared = shared.copy()
-        shared[:, 3:] = False
-    return torch.from_numpy(shared).unsqueeze(0)
+        shared = shared.copy()  # (l, 37)
+        shared[:, 3:] = False  # (l, 34); retain N, CA, C.
+    return torch.from_numpy(shared).unsqueeze(0)  # (1, l, 37)
 
 
 class Aligner:
@@ -57,16 +58,16 @@ class Aligner:
         if len(mobile) != len(target):
             raise AssertionError("mobile and target must contain the same residue count")
 
-        mobile_coordinates = _coordinate_batch(mobile)
-        target_coordinates = _coordinate_batch(target)
+        mobile_coordinates = _coordinate_batch(mobile)  # (1, l, 37, 3)
+        target_coordinates = _coordinate_batch(target)  # (1, l, 37, 3)
         if use_reflection:
-            target_coordinates = -target_coordinates
-        atom_mask = _shared_atom_mask(mobile, target, only_use_backbone)
+            target_coordinates = -target_coordinates  # (1, l, 37, 3)
+        atom_mask = _shared_atom_mask(mobile, target, only_use_backbone)  # (1, l, 37)
         self._affine3D, rmsd = compute_affine_and_rmsd(
             mobile_coordinates,
             target_coordinates,
             atom_exists_mask=atom_mask,
-        )
+        )  # affine shape: (1, 1); rmsd: ()
         self._rmsd = rmsd.item()
 
     @property
@@ -76,12 +77,12 @@ class Aligner:
     def apply(self, mobile: AlignableT) -> AlignableT:
         """Return a dataclass copy with all present atom coordinates aligned."""
 
-        present = np.asarray(mobile.atom37_mask, dtype=bool)
+        present = np.asarray(mobile.atom37_mask, dtype=bool)  # (l, 37)
         packed = torch.as_tensor(
             mobile.atom37_positions[present],
             dtype=torch.float32,
-        ).unsqueeze(0)
-        aligned = self._affine3D.apply(packed).squeeze(0).cpu().numpy()
-        atom37_positions = np.full_like(mobile.atom37_positions, np.nan)
-        atom37_positions[present] = aligned
+        ).unsqueeze(0)  # (1, n, 3), n is the number of present atoms.
+        aligned = self._affine3D.apply(packed).squeeze(0).cpu().numpy()  # (n, 3)
+        atom37_positions = np.full_like(mobile.atom37_positions, np.nan)  # (l, 37, 3)
+        atom37_positions[present] = aligned  # (n, 3)
         return replace(mobile, atom37_positions=atom37_positions)

fastplms/models/esmfold2/esmfold2_atom_indexer.py

@@ -2,11 +2,11 @@
 
 from __future__ import annotations
 
+import numpy as np
+
 from operator import attrgetter
 from typing import Any
 
-import numpy as np
-
 from .esmfold2_protein_structure import index_by_atom_name
 
 
@@ -19,7 +19,7 @@ class AtomIndexer:
 
     __slots__ = ("_get_property", "dim", "property", "structure")
 
-    def __init__(self, structure: Any, property: str, dim: int):
+    def __init__(self, structure: Any, property: str, dim: int) -> None:
         self.structure = structure
         self.property = property
         self.dim = dim

fastplms/models/esmfold2/esmfold2_conformers.py

@@ -11,21 +11,21 @@ import os
 import pickle
 import stat
 import tempfile
+import numpy as np
+
 from collections.abc import Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass
 from hashlib import file_digest
 from pathlib import Path
 from typing import Any, BinaryIO
-
-import numpy as np
 from huggingface_hub import hf_hub_download
 from huggingface_hub.constants import HF_HUB_CACHE
 
 from fastplms.registry import RuntimeAsset, get_model_registry
-
 from .esmfold2_constants import RES_TYPE_TO_CCD
 
+
 _CCD_ENVIRONMENT_VARIABLE = "ESMCFOLD_CCD_PATH"
 _CCD_ASSET_ID = "esmfold2_ccd"
 

fastplms/models/esmfold2/esmfold2_input_builder.py

@@ -2,14 +2,15 @@
 
 from __future__ import annotations
 
+import numpy as np
+
 from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Any, TypeAlias
 
-import numpy as np
-
 from .esmfold2_msa import MSA
 
+
 MSAInput: TypeAlias = MSA | None
 
 

fastplms/models/esmfold2/esmfold2_metrics.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import numpy as np
 import torch
 import torch.nn.functional as F
+
 from torch import Tensor
 from torch.amp import autocast  # type: ignore
 
@@ -18,8 +19,9 @@ from .esmfold2_protein_structure import (
 
 
 def _distance_matrix(positions: Tensor, eps: float) -> Tensor:
-    displacement = positions[..., None, :] - positions[..., None, :, :]
-    return torch.sqrt(eps + torch.sum(displacement**2, dim=-1))
+    # positions: (..., n, 3); n is the number of points in each distance matrix.
+    displacement = positions[..., None, :] - positions[..., None, :, :]  # (..., n, n, 3)
+    return torch.sqrt(eps + torch.sum(displacement**2, dim=-1))  # (..., n, n)
 
 
 def compute_lddt_from_dmat(
@@ -31,20 +33,21 @@ def compute_lddt_from_dmat(
     per_residue: bool = True,
 ) -> Tensor:
     """Score distance matrices ``D_pred`` and ``D_true`` with shape (..., l, l)."""
+    # dmat_pred, dmat_true, pairwise_mask: (..., l, l); cutoff broadcasts with distances.
 
     sequence_length = dmat_true.size(-1)
-    identity = torch.eye(sequence_length, device=dmat_true.device)
-    scored_pairs = (dmat_true < cutoff) * pairwise_mask * (1.0 - identity)
-    absolute_error = torch.abs(dmat_true - dmat_pred)
+    identity = torch.eye(sequence_length, device=dmat_true.device)  # (l, l)
+    scored_pairs = (dmat_true < cutoff) * pairwise_mask * (1.0 - identity)  # (..., l, l)
+    absolute_error = torch.abs(dmat_true - dmat_pred)  # (..., l, l)
     score = (
         (absolute_error < 0.5).type(absolute_error.dtype)
         + (absolute_error < 1.0).type(absolute_error.dtype)
         + (absolute_error < 2.0).type(absolute_error.dtype)
         + (absolute_error < 4.0).type(absolute_error.dtype)
-    ) * 0.25
+    ) * 0.25  # (..., l, l)
     dimensions = (-1,) if per_residue else (-2, -1)
-    normalization = 1.0 / (eps + scored_pairs.sum(dim=dimensions))
-    return normalization * (eps + (scored_pairs * score).sum(dim=dimensions))
+    normalization = 1.0 / (eps + scored_pairs.sum(dim=dimensions))  # (..., l) if per_residue, otherwise (...)
+    return normalization * (eps + (scored_pairs * score).sum(dim=dimensions))  # (..., l) if per_residue, otherwise (...)
 
 
 def compute_lddt(
@@ -58,16 +61,18 @@ def compute_lddt(
     sequence_id: Tensor | None = None,
 ) -> Tensor:
     """Compute lDDT from coordinate tensors and atom masks."""
+    # positions: (..., n, 3); all_atom_mask: (..., n); pairwise mask: (..., n, n).
+    # sequence_id, when supplied, follows the point axes (..., n).
 
-    expanded_mask = all_atom_mask[..., None]
-    true_distances = _distance_matrix(all_atom_positions, eps)
-    predicted_distances = _distance_matrix(all_atom_pred_pos, eps)
-    pair_mask = expanded_mask * expanded_mask.transpose(-2, -1)
+    expanded_mask = all_atom_mask[..., None]  # (..., n, 1)
+    true_distances = _distance_matrix(all_atom_positions, eps)  # (..., n, n)
+    predicted_distances = _distance_matrix(all_atom_pred_pos, eps)  # (..., n, n)
+    pair_mask = expanded_mask * expanded_mask.transpose(-2, -1)  # (..., n, n)
     if pairwise_all_atom_mask is not None:
-        pair_mask = pair_mask * pairwise_all_atom_mask
+        pair_mask = pair_mask * pairwise_all_atom_mask  # (..., n, n)
     if sequence_id is not None:
-        same_sequence = sequence_id[..., None] == sequence_id[..., None, :]
-        pair_mask = pair_mask * same_sequence.type_as(pair_mask)
+        same_sequence = sequence_id[..., None] == sequence_id[..., None, :]  # (..., n, n)
+        pair_mask = pair_mask * same_sequence.type_as(pair_mask)  # (..., n, n)
     return compute_lddt_from_dmat(
         predicted_distances,
         true_distances,
@@ -75,7 +80,7 @@ def compute_lddt(
         cutoff=cutoff,
         eps=eps,
         per_residue=per_residue,
-    )
+    )  # (..., n) if per_residue, otherwise (...)
 
 
 def compute_lddt_ca(
@@ -88,11 +93,13 @@ def compute_lddt_ca(
     sequence_id: Tensor | None = None,
 ) -> Tensor:
     """Compute lDDT using only C-alpha coordinates."""
+    # True coordinates/mask: (..., l, n_atoms, 3) / (..., l, n_atoms).
+    # Predicted rank-three input is treated as CA-only; otherwise its CA atom axis is selected.
 
     ca_index = residue_constants.atom_order["CA"]
     predicted_ca = (
         all_atom_pred_pos if all_atom_pred_pos.dim() == 3 else all_atom_pred_pos[..., ca_index, :]
-    )
+    )  # (..., l, 3)
     return compute_lddt(
         predicted_ca,
         all_atom_positions[..., ca_index, :],
@@ -101,7 +108,7 @@ def compute_lddt_ca(
         eps=eps,
         per_residue=per_residue,
         sequence_id=sequence_id,
-    )
+    )  # (..., l) if per_residue, otherwise (...)
 
 
 @torch.no_grad()
@@ -114,22 +121,24 @@ def compute_rmsd(
     reduction: str = "batch",
 ) -> Tensor:
     """Align ``X`` to ``Y`` and compute RMSD."""
+    # mobile/target: (b, n, 3) or (b, l, n_atoms, 3); masks omit xyz.
+    # b_eff counts unpacked sequences when sequence_id is provided; n counts flattened atoms.
 
     centered_mobile, _, centered_target, _, rotation, counts = compute_alignment_tensors(
         mobile,
         target,
         atom_exists_mask,
         sequence_id,
-    )
+    )  # coordinates (b_eff, n, 3), centroids (b_eff, 1, 3), rotation (b_eff, 3, 3), counts (b_eff, 1)
     rmsd = compute_rmsd_no_alignment(
         torch.matmul(centered_mobile, rotation),
         centered_target,
         counts,
         reduction=reduction,
-    )
+    )  # (b_eff, n / 3) per_residue; (b_eff,) per_sample; () batch
     if reduction == "per_residue" and sequence_id is not None:
-        return binpack(rmsd, sequence_id, pad_value=0)
-    return rmsd
+        return binpack(rmsd, sequence_id, pad_value=0)  # (b, packed_length)
+    return rmsd  # shape selected by reduction above
 
 
 def compute_gdt_ts(
@@ -140,36 +149,39 @@ def compute_gdt_ts(
     reduction: str = "per_sample",
 ) -> Tensor:
     """Align ``X`` to ``Y`` and compute GDT-TS."""
+    # mobile/target: batched xyz coordinates; masks omit xyz.
+    # b_eff counts unpacked sequences when sequence_id is provided; n counts flattened atoms.
 
     if atom_exists_mask is None:
-        atom_exists_mask = torch.isfinite(target).all(dim=-1)
+        atom_exists_mask = torch.isfinite(target).all(dim=-1)  # target.shape[:-1]
     centered_mobile, _, centered_target, _, rotation, _ = compute_alignment_tensors(
         mobile,
         target,
         atom_exists_mask,
         sequence_id,
-    )
+    )  # coordinates (b_eff, n, 3), centroids (b_eff, 1, 3), rotation (b_eff, 3, 3), counts (b_eff, 1)
     if sequence_id is not None:
-        atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=False)
+        atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=False)  # (b_eff, max_sequence_length, *original_mask.shape[2:])
     return compute_gdt_ts_no_alignment(
         torch.matmul(centered_mobile, rotation),
         centered_target,
         atom_exists_mask,
         reduction,
-    )
+    )  # (b_eff,) per_sample; () batch
 
 
 def _batched_contacts(predictions: Tensor, targets: Tensor) -> tuple[Tensor, Tensor]:
+    # predictions, targets: (l, l) or (b, l, l).
     if predictions.dim() == 2:
-        predictions = predictions.unsqueeze(0)
+        predictions = predictions.unsqueeze(0)  # (1, l, l)
     if targets.dim() == 2:
-        targets = targets.unsqueeze(0)
+        targets = targets.unsqueeze(0)  # (1, l, l)
     if predictions.size() != targets.size():
         raise ValueError(
             f"Size mismatch. Received predictions of size {predictions.size()}, "
             f"targets of size {targets.size()}"
         )
-    return predictions, targets
+    return predictions, targets  # each (b, l, l)
 
 
 def _valid_contact_mask(
@@ -178,14 +190,15 @@ def _valid_contact_mask(
     minsep: int,
     maxsep: int | None,
 ) -> Tensor:
+    # targets: (b, l, l); src_lengths: (b,).
     sequence_length = targets.shape[-1]
-    positions = torch.arange(sequence_length, device=targets.device)
-    separation = (positions.unsqueeze(0) - positions.unsqueeze(1)).unsqueeze(0)
-    valid = (separation >= minsep) & (targets >= 0)
+    positions = torch.arange(sequence_length, device=targets.device)  # (l,)
+    separation = (positions.unsqueeze(0) - positions.unsqueeze(1)).unsqueeze(0)  # (1, l, l)
+    valid = (separation >= minsep) & (targets >= 0)  # (b, l, l)
     if maxsep is not None:
-        valid &= separation < maxsep
-    within_length = positions.unsqueeze(0) < src_lengths.unsqueeze(1)
-    return valid & within_length.unsqueeze(1) & within_length.unsqueeze(2)
+        valid &= separation < maxsep  # (b, l, l)
+    within_length = positions.unsqueeze(0) < src_lengths.unsqueeze(1)  # (b, l)
+    return valid & within_length.unsqueeze(1) & within_length.unsqueeze(2)  # (b, l, l)
 
 
 def contact_precision(
@@ -197,8 +210,10 @@ def contact_precision(
     override_length: int | None = None,
 ) -> dict[str, Tensor]:
     """Compute P@L, P@L/5, and binned area for contact probabilities."""
+    # predictions, targets: (l, l) or (b, l, l); src_lengths: (b,) or None.
+    # n_upper_pairs counts the entries returned by triu_indices(l, minsep).
 
-    predictions, targets = _batched_contacts(predictions, targets)
+    predictions, targets = _batched_contacts(predictions, targets)  # each (b, l, l)
     batch_size, sequence_length, _ = predictions.shape
     if src_lengths is None:
         src_lengths = torch.full(
@@ -206,30 +221,30 @@ def contact_precision(
             sequence_length,
             dtype=torch.long,
             device=predictions.device,
-        )
-    valid = _valid_contact_mask(targets, src_lengths, minsep, maxsep)
-    masked_predictions = predictions.masked_fill(~valid, float("-inf"))
-    row_index, column_index = np.triu_indices(sequence_length, minsep)
-    upper_predictions = masked_predictions[:, row_index, column_index]
-    upper_targets = targets[:, row_index, column_index]
+        )  # (b,)
+    valid = _valid_contact_mask(targets, src_lengths, minsep, maxsep)  # (b, l, l)
+    masked_predictions = predictions.masked_fill(~valid, float("-inf"))  # (b, l, l)
+    row_index, column_index = np.triu_indices(sequence_length, minsep)  # each (n_upper_pairs,)
+    upper_predictions = masked_predictions[:, row_index, column_index]  # (b, n_upper_pairs)
+    upper_targets = targets[:, row_index, column_index]  # (b, n_upper_pairs)
 
     topk = sequence_length if override_length is None else max(sequence_length, override_length)
-    ranked_indices = upper_predictions.argsort(dim=-1, descending=True)[:, :topk]
-    batch_indices = torch.arange(batch_size, device=ranked_indices.device).unsqueeze(1)
-    ranked_targets = upper_targets[batch_indices, ranked_indices]
+    ranked_indices = upper_predictions.argsort(dim=-1, descending=True)[:, :topk]  # (b, min(topk, n_upper_pairs))
+    batch_indices = torch.arange(batch_size, device=ranked_indices.device).unsqueeze(1)  # (b, 1)
+    ranked_targets = upper_targets[batch_indices, ranked_indices]  # (b, min(topk, n_upper_pairs))
     if ranked_targets.size(1) < topk:
-        ranked_targets = F.pad(ranked_targets, [0, topk - ranked_targets.size(1)])
-    cumulative_contacts = ranked_targets.type_as(predictions).cumsum(dim=-1)
+        ranked_targets = F.pad(ranked_targets, [0, topk - ranked_targets.size(1)])  # (b, topk)
+    cumulative_contacts = ranked_targets.type_as(predictions).cumsum(dim=-1)  # (b, topk)
 
-    gather_lengths = src_lengths.unsqueeze(1)
+    gather_lengths = src_lengths.unsqueeze(1)  # (b, 1)
     if override_length is not None:
-        gather_lengths = override_length * torch.ones_like(gather_lengths)
-    fractions = torch.arange(0.1, 1.1, 0.1, device=predictions.device).unsqueeze(0)
-    gather_indices = (fractions * gather_lengths).type(torch.long).sub(1).clamp_min(0)
-    cumulative_bins = cumulative_contacts.gather(1, gather_indices)
-    precisions = cumulative_bins / (gather_indices + 1).type_as(cumulative_bins)
+        gather_lengths = override_length * torch.ones_like(gather_lengths)  # (b, 1)
+    fractions = torch.arange(0.1, 1.1, 0.1, device=predictions.device).unsqueeze(0)  # (1, 10)
+    gather_indices = (fractions * gather_lengths).type(torch.long).sub(1).clamp_min(0)  # (b, 10)
+    cumulative_bins = cumulative_contacts.gather(1, gather_indices)  # (b, 10)
+    precisions = cumulative_bins / (gather_indices + 1).type_as(cumulative_bins)  # (b, 10)
     return {
         "AUC": precisions.mean(dim=-1),
         "P@L": precisions[:, 9],
         "P@L5": precisions[:, 1],
-    }
+    }  # each metric: (b,)

fastplms/models/esmfold2/esmfold2_misc.py

@@ -6,6 +6,10 @@ module therefore performs no device selection, compilation, or remote access.
 
 from __future__ import annotations
 
+import numpy as np
+import torch
+import zstandard
+
 from collections import defaultdict
 from collections.abc import Generator, Iterable, Sequence
 from contextlib import AbstractContextManager, nullcontext
@@ -14,13 +18,10 @@ from io import BytesIO
 from typing import Any, Protocol, TypeVar, runtime_checkable
 from warnings import warn
 
-import numpy as np
-import torch
-import zstandard
-
 from .esmfold2_constants_esm3 import CHAIN_BREAK_STR
 from .esmfold2_utils_types import FunctionAnnotation
 
+
 MAX_SUPPORTED_DISTANCE = 1e6
 
 TSequence = TypeVar("TSequence", bound=Sequence)
@@ -50,44 +51,47 @@ def fp32_autocast_context(
 
 def maybe_tensor(value, convert_none_to_nan: bool = False) -> torch.Tensor | None:
     """Convert an optional array-like value to a tensor."""
+    # value: array-like arbitrary shape, or a list of identically shaped tensors.
 
     if value is None:
         return None
     if isinstance(value, torch.Tensor):
-        return value
+        return value  # value.shape
     if isinstance(value, list) and all(isinstance(element, torch.Tensor) for element in value):
-        return torch.stack(value)
+        return torch.stack(value)  # (n_values, *element_shape)
     if convert_none_to_nan:
-        value = np.asarray(value, dtype=np.float32)
-        value = np.where(value is None, np.nan, value)
-    return torch.tensor(value)
+        value = np.asarray(value, dtype=np.float32)  # shape inferred from the nested input
+        value = np.where(value is None, np.nan, value)  # value.shape
+    return torch.tensor(value)  # shape inferred from the array-like input
 
 
 def maybe_list(value, convert_nan_to_none: bool = False) -> list | None:
     """Convert an optional tensor or NumPy array to nested Python lists."""
+    # value: arbitrary-shaped array/tensor; element order and nesting are retained.
 
     if value is None:
         return None
     if not convert_nan_to_none:
         return value.tolist()
     if isinstance(value, torch.Tensor):
-        nan_mask = torch.isnan(value).cpu().numpy()
-        array = value.cpu().numpy().astype(object)
+        nan_mask = torch.isnan(value).cpu().numpy()  # value.shape
+        array = value.cpu().numpy().astype(object)  # value.shape
     elif isinstance(value, np.ndarray):
-        nan_mask = np.isnan(value)
-        array = value.astype(object)
+        nan_mask = np.isnan(value)  # value.shape
+        array = value.astype(object)  # value.shape
     else:
         raise TypeError("maybe_list can only work with torch.tensor or np.ndarray.")
-    array[nan_mask] = None
+    array[nan_mask] = None  # (n_nan,) selected elements
     return array.tolist()
 
 
 def replace_inf(data):
     """Replace infinite array values by the ESM API sentinel value."""
+    # data: array-like arbitrary shape; the returned list retains its nesting.
 
     if data is None:
         return None
-    array = np.asarray(data, dtype=np.float32)
+    array = np.asarray(data, dtype=np.float32)  # shape inferred from data
     return np.where(np.isinf(array), 1000, array).tolist()
 
 
@@ -118,9 +122,10 @@ def slice_any_object(
     idx: int | list[int] | slice | np.ndarray,
 ) -> TSequence:
     """Slice tensors, arrays, dataclasses, and ordinary Python sequences."""
+    # Array shape is caller-defined; idx follows that array type's indexing rules.
 
     if isinstance(obj, (np.ndarray, torch.Tensor)) or is_dataclass(obj):
-        return obj[idx]  # type: ignore[index,return-value]
+        return obj[idx]  # type: ignore[index,return-value]; shape determined by NumPy/Torch indexing when obj is an array
     return slice_python_object_as_numpy(obj, idx)
 
 
@@ -172,20 +177,21 @@ def concat_objects(objs: Sequence[Any], separator: Any | None = None):
             objs
             if separator is None
             else list(iterate_with_intermediate(objs, np.array([separator])))
-        )
-        return np.concatenate(pieces)
+        )  # arrays with a common trailing shape, interleaved with a separator if supplied
+        return np.concatenate(pieces)  # (sum of leading lengths, *common_trailing_shape)
     if isinstance(first, torch.Tensor):
         pieces = (
             objs
             if separator is None
             else list(iterate_with_intermediate(objs, torch.tensor([separator])))
-        )
-        return torch.cat(pieces)  # type: ignore[arg-type]
+        )  # tensors with a common trailing shape, interleaved with a separator if supplied
+        return torch.cat(pieces)  # type: ignore[arg-type]; (sum of leading lengths, *common_trailing_shape)
     raise TypeError(type(first))
 
 
-def rbf(values, v_min, v_max, n_bins=16):
+def rbf(values: torch.Tensor, v_min: float, v_max: float, n_bins: int = 16) -> torch.Tensor:
     """Encode values against evenly spaced radial basis centers."""
+    # values: arbitrary shape; the output appends a final n_bins axis.
 
     centers = torch.linspace(
         v_min,
@@ -193,34 +199,39 @@ def rbf(values, v_min, v_max, n_bins=16):
         n_bins,
         dtype=values.dtype,
         device=values.device,
-    )
-    centers = centers.reshape((1,) * values.ndim + (-1,))
-    standardized = (values.unsqueeze(-1) - centers) / ((v_max - v_min) / n_bins)
-    return torch.exp(-(standardized**2))
+    )  # (n_bins,)
+    centers = centers.reshape((1,) * values.ndim + (-1,))  # (1, ..., 1, n_bins); values.ndim leading singleton axes
+    standardized = (values.unsqueeze(-1) - centers) / ((v_max - v_min) / n_bins)  # (*values.shape, n_bins)
+    return torch.exp(-(standardized**2))  # (*values.shape, n_bins)
 
 
-def batched_gather(data, inds, dim=0, no_batch_dims=0):
+def batched_gather(
+    data: torch.Tensor, inds: torch.Tensor, dim: int = 0, no_batch_dims: int = 0
+) -> torch.Tensor:
     """Gather along one data dimension while retaining leading batch axes."""
+    # data/inds ranks are caller-defined. The first no_batch_dims axes index together;
+    # remaining advanced-index axes broadcast, while slice axes retain their data lengths.
 
     batch_indices = []
     index_rank = len(inds.shape)
     for axis, size in enumerate(data.shape[:no_batch_dims]):
         shape = (1,) * axis + (-1,) + (1,) * (index_rank - axis - 1)
-        batch_indices.append(torch.arange(size).view(*shape))
+        batch_indices.append(torch.arange(size).view(*shape))  # index-rank tensor; only the batch axis has length size
     tail = [slice(None)] * (len(data.shape) - no_batch_dims)
     tail[dim - no_batch_dims if dim >= 0 else dim] = inds
-    return data[tuple(batch_indices + tail)]
+    return data[tuple(batch_indices + tail)]  # broadcast advanced-index shape plus retained data slice axes
 
 
 def node_gather(s: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
     """Gather node features for each row of an edge-index tensor."""
+    # s: (..., n_nodes, d); edges: (..., n_nodes, n_neighbors).
 
     return batched_gather(
         s.unsqueeze(-3),
         edges,
         -2,
         no_batch_dims=len(s.shape) - 1,
-    )
+    )  # (..., n_nodes, n_neighbors, d)
 
 
 def knn_graph(
@@ -230,31 +241,32 @@ def knn_graph(
     sequence_id: torch.Tensor,
     *,
     no_knn: int,
-):
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Build nearest-neighbor edges, using sequence distance for missing geometry."""
+    # coords: (..., l, 3); masks: (..., l). With sequence_id, inputs use batch shape (b, l).
 
     length = coords.shape[-2]
-    coords = coords.nan_to_num()
-    missing_pair = ~(coord_mask[..., None, :] & coord_mask[..., :, None])
-    excluded_pair = padding_mask[..., None, :] | padding_mask[..., :, None]
+    coords = coords.nan_to_num()  # (..., l, 3)
+    missing_pair = ~(coord_mask[..., None, :] & coord_mask[..., :, None])  # (..., l, l)
+    excluded_pair = padding_mask[..., None, :] | padding_mask[..., :, None]  # (..., l, l)
     if sequence_id is not None:
-        excluded_pair |= sequence_id.unsqueeze(1) != sequence_id.unsqueeze(2)
+        excluded_pair |= sequence_id.unsqueeze(1) != sequence_id.unsqueeze(2)  # (b, l, l)
 
-    distances = (coords.unsqueeze(-2) - coords.unsqueeze(-3)).norm(dim=-1)
-    residue_index = torch.arange(length, device=coords.device)
-    sequence_distance = (residue_index.unsqueeze(-1) - residue_index.unsqueeze(-2)).abs()
+    distances = (coords.unsqueeze(-2) - coords.unsqueeze(-3)).norm(dim=-1)  # (..., l, l)
+    residue_index = torch.arange(length, device=coords.device)  # (l,)
+    sequence_distance = (residue_index.unsqueeze(-1) - residue_index.unsqueeze(-2)).abs()  # (l, l)
     if not (distances[~missing_pair] < MAX_SUPPORTED_DISTANCE).all():
         raise ValueError(
             "Coordinate pairwise distances exceed max supported distance "
             f"({MAX_SUPPORTED_DISTANCE}). "
         )
 
-    rank_distance = sequence_distance.to(distances.dtype).mul(1e2).add(MAX_SUPPORTED_DISTANCE)
-    rank_distance = rank_distance.where(missing_pair, distances)
-    rank_distance = rank_distance.masked_fill(excluded_pair, torch.inf)
-    sorted_distance, sorted_edge = rank_distance.sort(dim=-1, descending=False)
+    rank_distance = sequence_distance.to(distances.dtype).mul(1e2).add(MAX_SUPPORTED_DISTANCE)  # (l, l)
+    rank_distance = rank_distance.where(missing_pair, distances)  # (..., l, l)
+    rank_distance = rank_distance.masked_fill(excluded_pair, torch.inf)  # (..., l, l)
+    sorted_distance, sorted_edge = rank_distance.sort(dim=-1, descending=False)  # each (..., l, l)
     width = min(no_knn, length)
-    return sorted_edge[..., :width], sorted_distance[..., :width].isfinite()
+    return sorted_edge[..., :width], sorted_distance[..., :width].isfinite()  # each (..., l, min(no_knn, l))
 
 
 def stack_variable_length_tensors(
@@ -263,6 +275,7 @@ def stack_variable_length_tensors(
     dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Pad arbitrary tensor dimensions to their maxima, then stack."""
+    # Each sequence has the same rank; its axis lengths may differ. Padding uses each axis maximum.
 
     output_shape = [
         len(sequences),
@@ -273,56 +286,58 @@ def stack_variable_length_tensors(
         constant_value,
         dtype=sequences[0].dtype if dtype is None else dtype,
         device=sequences[0].device,
-    )
+    )  # (n_sequences, *axiswise_maximum_shape)
     for destination, source in zip(output, sequences, strict=True):
-        destination[tuple(slice(size) for size in source.shape)] = source
-    return output
+        destination[tuple(slice(size) for size in source.shape)] = source  # source.shape slice of destination
+    return output  # (n_sequences, *axiswise_maximum_shape)
 
 
 def binpack(
     tensor: torch.Tensor,
     sequence_id: torch.Tensor | None,
     pad_value: int | float,
-):
+) -> torch.Tensor:
     """Scatter a sequence-major tensor into the packed layout described by IDs."""
+    # tensor: (n_unpacked_sequences, max_length, ...); sequence_id: (b, packed_length) or None.
 
     if sequence_id is None:
-        return tensor
-    sequence_counts = sequence_id.max(dim=-1).values + 1
+        return tensor  # tensor.shape
+    sequence_counts = sequence_id.max(dim=-1).values + 1  # (b,)
     output = torch.full(
         sequence_id.shape + tensor.shape[2:],
         fill_value=pad_value,
         dtype=tensor.dtype,
         device=tensor.device,
-    )
+    )  # (b, packed_length, *tensor.shape[2:])
     source_index = 0
     for batch_index, (batch_ids, count) in enumerate(
         zip(sequence_id, sequence_counts, strict=True)
     ):
         for seqid in range(count):
-            selection = batch_ids == seqid
-            output[batch_index, selection] = tensor[source_index, : selection.sum()]
+            selection = batch_ids == seqid  # (packed_length,)
+            output[batch_index, selection] = tensor[source_index, : selection.sum()]  # (n_selected, *tensor.shape[2:]) selected rows
             source_index += 1
-    return output
+    return output  # (b, packed_length, *tensor.shape[2:])
 
 
 def unbinpack(
     tensor: torch.Tensor,
     sequence_id: torch.Tensor | None,
     pad_value: int | float,
-):
+) -> torch.Tensor:
     """Restore sequence-major rows from a packed tensor and its sequence IDs."""
+    # tensor: (b, packed_length, ...); sequence_id: (b, packed_length) or None.
 
     if sequence_id is None:
-        return tensor
+        return tensor  # tensor.shape
     rows = []
-    sequence_counts = sequence_id.max(dim=-1).values + 1
+    sequence_counts = sequence_id.max(dim=-1).values + 1  # (b,)
     for batch_index, (batch_ids, count) in enumerate(
         zip(sequence_id, sequence_counts, strict=True)
     ):
         for seqid in range(count):
-            rows.append(tensor[batch_index, batch_ids == seqid])
-    return stack_variable_length_tensors(rows, pad_value)
+            rows.append(tensor[batch_index, batch_ids == seqid])  # (selected_sequence_length, *tensor.shape[2:])
+    return stack_variable_length_tensors(rows, pad_value)  # (n_unpacked_sequences, max_sequence_length, *tensor.shape[2:])
 
 
 def merge_ranges(
@@ -386,7 +401,7 @@ def get_chainbreak_boundaries_from_sequence(
         boundaries.extend((index, index + 1))
     boundaries.append(len(sequence))
     assert len(boundaries) % 2 == 0
-    return np.asarray(boundaries).reshape(-1, 2)
+    return np.asarray(boundaries).reshape(-1, 2)  # (n_chains, 2)
 
 
 def deserialize_tensors(data: bytes) -> Any:

fastplms/models/esmfold2/esmfold2_mmcif_parsing.py

@@ -5,17 +5,18 @@ from __future__ import annotations
 import functools
 import io
 import os
-from contextlib import suppress
-from dataclasses import dataclass
-from datetime import datetime
-
 import biotite.structure as bs
 import biotite.structure.io.pdbx as pdbx
 import numpy as np
+
+from contextlib import suppress
+from dataclasses import dataclass
+from datetime import datetime
 from biotite.structure.io.pdbx import CIFColumn, CIFData, CIFFile
 
 from . import esmfold2_residue_constants as residue_constants
 
+
 PathOrBuffer = str | os.PathLike | io.StringIO
 
 PLDDT_B_FACTOR_SCALE = 100.0

fastplms/models/esmfold2/esmfold2_molecular_complex.py

@@ -11,18 +11,18 @@ from __future__ import annotations
 import io
 import os
 import re
-from dataclasses import asdict, dataclass
-from pathlib import Path
-from subprocess import check_output
-from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any
-
 import biotite.structure as bs
 import biotite.structure.io.pdbx as pdbx
 import brotli
 import msgpack
 import numpy as np
 import torch
+
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from subprocess import check_output
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, Any
 from biotite.structure.io.pdbx import (
     CIFCategory,
     CIFColumn,

fastplms/models/esmfold2/esmfold2_msa.py

@@ -4,13 +4,13 @@ from __future__ import annotations
 
 import dataclasses
 import string
+import numpy as np
+
 from collections.abc import Sequence
 from dataclasses import dataclass
 from functools import cached_property
 from itertools import islice
 from typing import Any
-
-import numpy as np
 from Bio import SeqIO
 from scipy.spatial.distance import cdist
 
@@ -20,6 +20,7 @@ from .esmfold2_parsing import FastaEntry, read_sequences, write_sequences
 from .esmfold2_sequential_dataclass import SequentialDataclass
 from .esmfold2_system import PathOrBuffer
 
+
 _A3M_INSERTION_DELETE_TABLE = str.maketrans(
     dict.fromkeys(string.ascii_lowercase + ".")
 )

fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py

@@ -4,22 +4,24 @@ from __future__ import annotations
 
 import os
 import tempfile
-from pathlib import Path
-
 import numpy as np
 
+from pathlib import Path
+
 from .esmfold2_system import run_subprocess_with_errorcheck
 
 
 def _byte_matrix(array: np.ndarray) -> np.ndarray:
     """Return a two-dimensional byte view used for Hamming comparisons."""
 
-    matrix = np.asarray(array).view(np.uint8)
-    return matrix.reshape(matrix.shape[0], -1)
+    # array: (n, l); w is its row width in bytes after viewing its dtype.
+    matrix = np.asarray(array).view(np.uint8)  # (n, w)
+    return matrix.reshape(matrix.shape[0], -1)  # (n, w)
 
 
 def _hamming_to_all(query: np.ndarray, sequences: np.ndarray) -> np.ndarray:
-    return np.not_equal(sequences, query).mean(axis=1, dtype=np.float64)
+    # query: (w,); sequences: (n, w).
+    return np.not_equal(sequences, query).mean(axis=1, dtype=np.float64)  # (n,)
 
 
 def greedy_select_indices(array: np.ndarray, num_seqs: int, mode: str = "max") -> list[int]:
@@ -48,20 +50,20 @@ def greedy_select_indices(array: np.ndarray, num_seqs: int, mode: str = "max") -
     if depth <= num_seqs:
         return list(range(depth))
 
-    sequences = _byte_matrix(array)
+    sequences = _byte_matrix(array)  # (n, w)
     selected = [0]
-    available = np.ones(depth, dtype=bool)
-    available[0] = False
-    distance_sum = _hamming_to_all(sequences[0], sequences)
+    available = np.ones(depth, dtype=bool)  # (n,)
+    available[0] = False  # ()
+    distance_sum = _hamming_to_all(sequences[0], sequences)  # (n,)
     choose = np.argmax if mode == "max" else np.argmin
 
     while len(selected) < num_seqs:
-        candidates = np.flatnonzero(available)
-        candidate_scores = distance_sum[candidates] / len(selected)
+        candidates = np.flatnonzero(available)  # (n_remaining,)
+        candidate_scores = distance_sum[candidates] / len(selected)  # (n_remaining,)
         next_index = int(candidates[int(choose(candidate_scores))])
         selected.append(next_index)
-        available[next_index] = False
-        distance_sum += _hamming_to_all(sequences[next_index], sequences)
+        available[next_index] = False  # ()
+        distance_sum += _hamming_to_all(sequences[next_index], sequences)  # (n,)
     return sorted(selected)
 
 

fastplms/models/esmfold2/esmfold2_normalize_coordinates.py

@@ -2,23 +2,25 @@
 
 from __future__ import annotations
 
-from typing import TypeVar
-
 import numpy as np
 import torch
+
+from typing import TypeVar
 from torch import Tensor
 
 from . import esmfold2_residue_constants as residue_constants
 from .esmfold2_affine3d import Affine3D
 
+
 ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor)
 
 
 def atom3_to_backbone_frames(bb_positions: Tensor) -> Affine3D:
     """Construct a frame from N, C-alpha, and C positions in ``X``."""
 
-    n_position, ca_position, c_position = bb_positions.unbind(dim=-2)
-    return Affine3D.from_graham_schmidt(c_position, ca_position, n_position)
+    # bb_positions: (..., 3, 3), ordered N, CA, C on the penultimate axis.
+    n_position, ca_position, c_position = bb_positions.unbind(dim=-2)  # each (..., 3)
+    return Affine3D.from_graham_schmidt(c_position, ca_position, n_position)  # affine shape: (...,)
 
 
 def index_by_atom_name(
@@ -26,42 +28,43 @@ def index_by_atom_name(
     atom_names: str | list[str],
     dim: int = -2,
 ) -> ArrayOrTensor:
-    """Select one or more named atoms along an atom37 axis."""
+    """Select named atoms, replacing axis ``dim`` by their count (or removing it)."""
 
     single_atom = isinstance(atom_names, str)
     names = [atom_names] if single_atom else atom_names
     indices = [residue_constants.atom_order[name] for name in names]
     axis = dim % atom37.ndim
     if isinstance(atom37, Tensor):
-        index = torch.tensor(indices, dtype=torch.long, device=atom37.device)
-        selected = torch.index_select(atom37, axis, index)
+        index = torch.tensor(indices, dtype=torch.long, device=atom37.device)  # (n_names,)
+        selected = torch.index_select(atom37, axis, index)  # atom37.shape with axis = n_names
     else:
-        selected = np.take(atom37, indices, axis=axis)
+        selected = np.take(atom37, indices, axis=axis)  # atom37.shape with axis = n_names
     return selected.squeeze(axis) if single_atom else selected  # type: ignore[return-value]
 
 
 def get_protein_normalization_frame(coords: Tensor) -> Affine3D:
-    """Build one frame from backbone coordinates ``X`` with shape (l, 37, 3)."""
+    """Build one frame per batch item from coordinates of shape (..., l, 37, 3)."""
 
-    backbone = index_by_atom_name(coords, ["N", "CA", "C"], dim=-2)
-    residue_is_valid = torch.isfinite(backbone).all(dim=-1).all(dim=-1)
-    weights = residue_is_valid[..., None, None]
-    coordinate_sum = backbone.masked_fill(~weights, 0).sum(dim=-3)
-    count = residue_is_valid.sum(dim=-1)[..., None, None]
-    mean_backbone = coordinate_sum / (count + 1e-8)
-    return atom3_to_backbone_frames(mean_backbone.float())
+    backbone = index_by_atom_name(coords, ["N", "CA", "C"], dim=-2)  # (..., l, 3, 3)
+    residue_is_valid = torch.isfinite(backbone).all(dim=-1).all(dim=-1)  # (..., l)
+    weights = residue_is_valid[..., None, None]  # (..., l, 1, 1)
+    coordinate_sum = backbone.masked_fill(~weights, 0).sum(dim=-3)  # (..., 3, 3)
+    count = residue_is_valid.sum(dim=-1)[..., None, None]  # (..., 1, 1)
+    mean_backbone = coordinate_sum / (count + 1e-8)  # (..., 3, 3)
+    return atom3_to_backbone_frames(mean_backbone.float())  # affine shape: (...,)
 
 
 def apply_frame_to_coords(coords: Tensor, frame: Affine3D) -> Tensor:
     """Express atom coordinates ``X`` in the inverse of ``frame``."""
 
-    transformed = frame[..., None, None].invert().apply(coords)
-    frame_is_valid = frame.trans.norm(dim=-1) > 0
-    normalized = torch.where(frame_is_valid[..., None, None, None], transformed, coords)
-    return normalized.masked_fill(torch.isinf(coords), torch.inf)
+    # coords: (..., l, 37, 3); frame shape: (...,).
+    transformed = frame[..., None, None].invert().apply(coords)  # (..., l, 37, 3)
+    frame_is_valid = frame.trans.norm(dim=-1) > 0  # (...,)
+    normalized = torch.where(frame_is_valid[..., None, None, None], transformed, coords)  # (..., l, 37, 3)
+    return normalized.masked_fill(torch.isinf(coords), torch.inf)  # (..., l, 37, 3)
 
 
 def normalize_coordinates(coords: Tensor) -> Tensor:
     """Normalize ``X`` with shape (..., l, 37, 3) to its backbone frame."""
 
-    return apply_frame_to_coords(coords, get_protein_normalization_frame(coords))
+    return apply_frame_to_coords(coords, get_protein_normalization_frame(coords))  # (..., l, 37, 3)

fastplms/models/esmfold2/esmfold2_output.py

@@ -2,14 +2,14 @@
 
 from __future__ import annotations
 
+import numpy as np
+import torch
+
 from collections.abc import Iterable
 from dataclasses import dataclass, field
 from itertools import groupby
 from typing import Any
 
-import numpy as np
-import torch
-
 from .esmfold2_constants import ELEMENT_NUMBER_TO_SYMBOL, MOL_TYPE_NONPOLYMER
 from .esmfold2_molecular_complex import MolecularComplex, MolecularComplexMetadata
 
@@ -92,17 +92,18 @@ def build_molecular_complex_from_features(
     tokens are collapsed into one non-polymer residue per chain.
     """
 
-    M = atom_mask.bool().cpu().numpy()
-    X = coords.float().cpu().numpy()
-    atom_names = ref_atom_name_chars.cpu().numpy()
-    elements = ref_element.cpu().numpy()
-    confidence = None if plddt is None else plddt.float().cpu().numpy()
+    # a counts padded atoms; l counts model tokens (ligand tokens may be atoms).
+    present_atoms = atom_mask.bool().cpu().numpy()  # (a,)
+    X = coords.float().cpu().numpy()  # (a, 3)
+    atom_names = ref_atom_name_chars.cpu().numpy()  # (a, 4)
+    elements = ref_element.cpu().numpy()  # (a,)
+    confidence = None if plddt is None else plddt.float().cpu().numpy()  # (l,) or None
     records = _ComplexRecords()
 
-    def decode_atoms(tokens: Iterable[Any]):
+    def decode_atoms(tokens: Iterable[Any]) -> Iterable[tuple[list[float], str, str]]:
         for token in tokens:
             for atom_index in range(token.atom_start, token.atom_start + token.atom_count):
-                if M[atom_index]:
+                if present_atoms[atom_index]:
                     yield (
                         X[atom_index].tolist(),
                         get_element_symbol(int(elements[atom_index])),

fastplms/models/esmfold2/esmfold2_paired_msa.py

@@ -3,10 +3,10 @@
 from __future__ import annotations
 
 import re
-from dataclasses import dataclass
-
 import numpy as np
 
+from dataclasses import dataclass
+
 from .esmfold2_constants import (
     MSA_GAP_TOKEN_ID,
     PROTEIN_3TO1,
@@ -15,6 +15,7 @@ from .esmfold2_constants import (
 )
 from .esmfold2_msa import MSA
 
+
 _TAXONOMY_PATTERN = re.compile(r"key=(-?\d+)")
 
 

fastplms/models/esmfold2/esmfold2_parsing.py

@@ -4,8 +4,9 @@ from __future__ import annotations
 
 import gzip
 import io
+
 from collections.abc import Generator, Iterable
-from contextlib import nullcontext
+from contextlib import AbstractContextManager, nullcontext
 from pathlib import Path
 from typing import NamedTuple, TextIO
 
@@ -45,7 +46,7 @@ def parse_fasta(text: str) -> Generator[FastaEntry, None, None]:
         raise ValueError("Found no sequences in input")
 
 
-def _open_reader(source: PathOrBuffer):
+def _open_reader(source: PathOrBuffer) -> AbstractContextManager[TextIO]:
     if isinstance(source, io.TextIOBase):
         return nullcontext(source)
     path = Path(source)
@@ -93,7 +94,7 @@ def append_fasta_sequence(header: str, sequence: str, path: str | Path) -> None:
         handle.write(f">{header}\n{sequence}\n")
 
 
-def _open_writer(destination: PathOrBuffer):
+def _open_writer(destination: PathOrBuffer) -> AbstractContextManager[TextIO]:
     if isinstance(destination, io.TextIOBase):
         return nullcontext(destination)
     path = Path(destination)

fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py

@@ -4,16 +4,19 @@ from __future__ import annotations
 
 import torch
 import torch.nn.functional as F
+
 from torch import Tensor
 
 from .esmfold2_affine3d import Affine3D
 
+
 _CPU_DEVICE = torch.device("cpu")
 
 
 def _compute_pae_masks(mask: Tensor) -> Tensor:
-    residue_mask = mask.bool()
-    return residue_mask.unsqueeze(-1) & residue_mask.unsqueeze(-2)
+    # mask: (..., l), where l is the residue count.
+    residue_mask = mask.bool()  # (..., l)
+    return residue_mask.unsqueeze(-1) & residue_mask.unsqueeze(-2)  # (..., l, l)
 
 
 def _pae_bins(
@@ -23,19 +26,20 @@ def _pae_bins(
 ) -> Tensor:
     """Return the representative distance for each PAE probability bin."""
 
-    boundaries = torch.linspace(0, max_bin, steps=num_bins - 1, device=device)
+    boundaries = torch.linspace(0, max_bin, steps=num_bins - 1, device=device)  # (n_bins - 1,)
     width = max_bin / (num_bins - 2)
-    centers = boundaries + width / 2
-    overflow_center = centers[-1:] + width
-    return torch.cat((centers, overflow_center))
+    centers = boundaries + width / 2  # (n_bins - 1,)
+    overflow_center = centers[-1:] + width  # (1,)
+    return torch.cat((centers, overflow_center))  # (n_bins,)
 
 
 def _masked_probabilities(logits: Tensor, pair_mask: Tensor) -> Tensor:
+    # logits: (..., l, l, n_bins); pair_mask: (..., l, l).
     masked_logits = logits.masked_fill(
         ~pair_mask.unsqueeze(-1),
         torch.finfo(logits.dtype).min,
-    )
-    return masked_logits.softmax(dim=-1)
+    )  # (..., l, l, n_bins)
+    return masked_logits.softmax(dim=-1)  # (..., l, l, n_bins)
 
 
 def masked_mean(
@@ -46,10 +50,11 @@ def masked_mean(
 ) -> Tensor:
     """Average values over true entries of a broadcast-compatible mask."""
 
-    weights = mask.expand_as(value)
-    weighted_sum = torch.sum(weights * value, dim=dim)
-    weight_sum = torch.sum(weights, dim=dim)
-    return weighted_sum / (weight_sum + eps)
+    # value has arbitrary shape; reduced_shape removes the axes named by dim.
+    weights = mask.expand_as(value)  # value.shape
+    weighted_sum = torch.sum(weights * value, dim=dim)  # reduced_shape
+    weight_sum = torch.sum(weights, dim=dim)  # reduced_shape
+    return weighted_sum / (weight_sum + eps)  # reduced_shape
 
 
 def compute_predicted_aligned_error(
@@ -61,30 +66,31 @@ def compute_predicted_aligned_error(
     """Convert PAE logits ``X`` with shape (..., l, l, n) to distances."""
 
     del sequence_id
-    pair_mask = _compute_pae_masks(aa_mask)
-    probabilities = _masked_probabilities(logits, pair_mask)
-    centers = _pae_bins(max_bin, logits.shape[-1], logits.device)
-    return torch.sum(probabilities * centers, dim=-1)
+    pair_mask = _compute_pae_masks(aa_mask)  # (..., l, l)
+    probabilities = _masked_probabilities(logits, pair_mask)  # (..., l, l, n_bins)
+    centers = _pae_bins(max_bin, logits.shape[-1], logits.device)  # (n_bins,)
+    return torch.sum(probabilities * centers, dim=-1)  # (..., l, l)
 
 
 @torch.no_grad()
 def compute_tm(logits: Tensor, aa_mask: Tensor, max_bin: float = 31.0) -> Tensor:
-    """Estimate TM score from pairwise PAE logits."""
+    """Estimate TM score from logits (..., l, l, n_bins) and residue mask (..., l)."""
 
-    pair_mask = _compute_pae_masks(aa_mask)
-    sequence_lengths = aa_mask.sum(dim=-1, keepdim=True)
-    centers = _pae_bins(max_bin, logits.shape[-1], logits.device)
-    distance_scale = 1.24 * (sequence_lengths.clamp_min(19) - 15) ** (1 / 3) - 1.8
-    tm_weights = 1.0 / (1 + (centers / distance_scale.unsqueeze(-1)) ** 2)
-    probabilities = _masked_probabilities(logits, pair_mask)
-    score_per_pair = torch.sum(probabilities * tm_weights.unsqueeze(-2), dim=-1)
-    score_per_anchor = masked_mean(pair_mask, score_per_pair, dim=-1)
-    return score_per_anchor.max(dim=-1).values
+    pair_mask = _compute_pae_masks(aa_mask)  # (..., l, l)
+    sequence_lengths = aa_mask.sum(dim=-1, keepdim=True)  # (..., 1)
+    centers = _pae_bins(max_bin, logits.shape[-1], logits.device)  # (n_bins,)
+    distance_scale = 1.24 * (sequence_lengths.clamp_min(19) - 15) ** (1 / 3) - 1.8  # (..., 1)
+    tm_weights = 1.0 / (1 + (centers / distance_scale.unsqueeze(-1)) ** 2)  # (..., 1, n_bins)
+    probabilities = _masked_probabilities(logits, pair_mask)  # (..., l, l, n_bins)
+    score_per_pair = torch.sum(probabilities * tm_weights.unsqueeze(-2), dim=-1)  # (..., l, l)
+    score_per_anchor = masked_mean(pair_mask, score_per_pair, dim=-1)  # (..., l)
+    return score_per_anchor.max(dim=-1).values  # (...,)
 
 
 def _local_coordinates(frames: Affine3D) -> Tensor:
-    origins = frames.trans[..., None, :, :]
-    return frames.invert()[..., None].apply(origins)
+    # frames.shape: (..., l); trans: (..., l, 3).
+    origins = frames.trans[..., None, :, :]  # (..., 1, l, 3)
+    return frames.invert()[..., None].apply(origins)  # (..., l, l, 3)
 
 
 def tm_loss(
@@ -96,32 +102,32 @@ def tm_loss(
     sequence_id: Tensor | None = None,
     max_bin: float = 31,
 ) -> Tensor:
-    """Cross-entropy loss for discretized aligned-position errors."""
+    """Cross-entropy loss for logits (b, l, l, n_bins) and residue frames (b, l)."""
 
     del sequence_id
-    predicted_frames = Affine3D.from_tensor(pred_affine)
-    target_frames = Affine3D.from_tensor(targ_affine)
+    predicted_frames = Affine3D.from_tensor(pred_affine)  # frame shape: (b, l)
+    target_frames = Affine3D.from_tensor(targ_affine)  # frame shape: (b, l)
     with torch.no_grad():
         squared_error = (
             (_local_coordinates(predicted_frames) - _local_coordinates(target_frames))
             .square()
             .sum(dim=-1)
-        )
+        )  # (b, l, l)
         boundaries = torch.linspace(
             0,
             max_bin,
             logits.shape[-1] - 1,
             device=logits.device,
-        ).square()
-        target_bins = (squared_error[..., None] > boundaries).sum(dim=-1).long()
+        ).square()  # (n_bins - 1,)
+        target_bins = (squared_error[..., None] > boundaries).sum(dim=-1).long()  # (b, l, l)
 
     cross_entropy = F.cross_entropy(
         logits.movedim(3, 1),
         target_bins,
         reduction="none",
-    )
-    pair_mask = _compute_pae_masks(targ_mask)
-    loss_per_sample = masked_mean(pair_mask, cross_entropy, dim=(-1, -2))
+    )  # (b, l, l)
+    pair_mask = _compute_pae_masks(targ_mask)  # (b, l, l)
+    loss_per_sample = masked_mean(pair_mask, cross_entropy, dim=(-1, -2))  # (b,)
     if tm_mask is None:
-        return loss_per_sample.mean()
-    return masked_mean(tm_mask, loss_per_sample)
+        return loss_per_sample.mean()  # ()
+    return masked_mean(tm_mask, loss_per_sample)  # ()

fastplms/models/esmfold2/esmfold2_prepare_input.py

@@ -10,15 +10,15 @@ from __future__ import annotations
 
 import math
 import warnings
+import numpy as np
+import torch
+
 from collections import defaultdict
 from contextlib import suppress
 from dataclasses import dataclass, field
 from itertools import combinations
 from typing import Any
 
-import numpy as np
-import torch
-
 from .esmfold2_conformers import (
     get_ccd_leaving_atoms,
     get_idealized_atom_pos,
@@ -62,6 +62,7 @@ from .esmfold2_types import (
     StructurePredictionInput,
 )
 
+
 _ZERO_POS = np.zeros(3, dtype=np.float32)
 _ENCODE_ATOM_NAME_CACHE: dict[str, list[int]] = {}
 _ELEMENT_ATOMIC_NUM_CACHE: dict[str, int] = {}

fastplms/models/esmfold2/esmfold2_processor.py

@@ -2,13 +2,13 @@
 
 from __future__ import annotations
 
+import numpy as np
+import torch
+
 from collections.abc import Mapping
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
-
-import numpy as np
-import torch
 from torch import Tensor
 from tqdm.auto import tqdm
 
@@ -20,6 +20,7 @@ from .esmfold2_types import MSA, Modification, ProteinInput, StructurePrediction
 from .modeling_esmfold2_common import MSA_CONDITIONING_INPUT_NAMES
 from .reproducibility import seed_context
 
+
 # Backward-compatible private alias for the pinned parity helpers. New callers
 # should import ``seed_context`` from the public ``fastplms.models.esmfold2``
 # package instead of reaching into implementation modules.

fastplms/models/esmfold2/esmfold2_protein_chain.py

@@ -4,18 +4,18 @@ from __future__ import annotations
 
 import io
 import warnings
-from collections.abc import Mapping, Sequence
-from dataclasses import asdict, dataclass, replace
-from functools import cached_property
-from pathlib import Path
-from typing import Any
-
 import biotite.structure as bs
 import brotli
 import msgpack
 import msgpack_numpy
 import numpy as np
 import torch
+
+from collections.abc import Mapping, Sequence
+from dataclasses import asdict, dataclass, replace
+from functools import cached_property
+from pathlib import Path
+from typing import Any
 from biotite.database import rcsb
 from biotite.structure.io.pdb import PDBFile
 from biotite.structure.io.pdbx import CIFCategory, CIFColumn, CIFData, CIFFile
@@ -42,6 +42,7 @@ from .esmfold2_normalize_coordinates import (
 from .esmfold2_protein_structure import index_by_atom_name
 from .esmfold2_utils_types import PathOrBuffer
 
+
 CHAIN_ID_CONST = "A"
 
 
@@ -70,27 +71,30 @@ def infer_cb(
     dihedral: float = -2.143,
 ):
     """Infer C-beta coordinates from C, N, and C-alpha coordinates."""
+    # C, N, Ca: (..., 3); xyz coordinates with broadcast-compatible leading axes.
 
     def normalize(X: np.ndarray) -> np.ndarray:
-        return X / np.sqrt(np.square(X).sum(-1, keepdims=True) + 1e-8)
+        # X: (..., 3); normalize each xyz vector independently.
+        return X / np.sqrt(np.square(X).sum(-1, keepdims=True) + 1e-8)  # X.shape; normalized along xyz
 
     with np.errstate(invalid="ignore"):
-        n_to_ca = N - Ca
-        n_to_c = N - C
-    axis = normalize(n_to_ca)
-    normal = normalize(np.cross(n_to_c, axis))
-    basis = (axis, np.cross(normal, axis), normal)
+        n_to_ca = N - Ca  # (..., 3)
+        n_to_c = N - C  # (..., 3)
+    axis = normalize(n_to_ca)  # (..., 3)
+    normal = normalize(np.cross(n_to_c, axis))  # (..., 3)
+    basis = (axis, np.cross(normal, axis), normal)  # three arrays (..., 3)
     offsets = (
         bond_length * np.cos(bond_angle),
         bond_length * np.sin(bond_angle) * np.cos(dihedral),
         -bond_length * np.sin(bond_angle) * np.sin(dihedral),
-    )
-    return Ca + sum(vector * offset for vector, offset in zip(basis, offsets, strict=False))
+    )  # three NumPy scalars, each ()
+    return Ca + sum(vector * offset for vector, offset in zip(basis, offsets, strict=False))  # (..., 3)
 
 
 def chain_to_ndarray(
     atom_array: bs.AtomArray, mmcif: MmcifWrapper, chain_id: str, is_predicted=False
 ):
+    # atom_array: n input atoms; l is the selected chain's sequence length.
     if not isinstance(atom_array, bs.AtomArray):
         raise TypeError("atom_array must be a biotite AtomArray.")
     if not isinstance(mmcif, MmcifWrapper):
@@ -106,14 +110,14 @@ def chain_to_ndarray(
     num_res = len(mmcif.chain_to_seqres[chain_id])
     sequence = mmcif.chain_to_seqres[chain_id]
 
-    atom_positions = np.full([num_res, residue_constants.atom_type_num, 3], np.nan)
-    atom_mask = np.full([num_res, residue_constants.atom_type_num], False, dtype=bool)
-    residue_index = np.full([num_res], -1, dtype=np.int64)
-    insertion_code = np.full([num_res], "", dtype="<U4")
+    atom_positions = np.full([num_res, residue_constants.atom_type_num, 3], np.nan)  # (l, 37, 3)
+    atom_mask = np.full([num_res, residue_constants.atom_type_num], False, dtype=bool)  # (l, 37)
+    residue_index = np.full([num_res], -1, dtype=np.int64)  # (l,)
+    insertion_code = np.full([num_res], "", dtype="<U4")  # (l,)
 
-    confidence = np.ones([num_res], dtype=np.float32)
+    confidence = np.ones([num_res], dtype=np.float32)  # (l,)
 
-    chain = atom_array[atom_array.chain_id == chain_id]
+    chain = atom_array[atom_array.chain_id == chain_id]  # AtomArray with n_chain_atoms entries
     if not isinstance(chain, bs.AtomArray):
         raise RuntimeError("Biotite selection did not return an AtomArray.")
     for res_index in range(num_res):
@@ -122,13 +126,13 @@ def chain_to_ndarray(
         if res_at_position.residue_number is None:
             continue
 
-        residue_index[res_index] = res_at_position.residue_number
-        insertion_code[res_index] = res_at_position.insertion_code
+        residue_index[res_index] = res_at_position.residue_number  # scalar element in (l,)
+        insertion_code[res_index] = res_at_position.insertion_code  # scalar element in (l,)
         res = chain[
             (chain.res_id == res_at_position.residue_number)
             & (chain.ins_code == res_at_position.insertion_code)
             & (chain.hetero == res_at_position.hetflag)
-        ]
+        ]  # AtomArray with n_residue_atoms entries
         if not isinstance(res, bs.AtomArray):
             raise RuntimeError("Biotite residue selection did not return an AtomArray.")
 
@@ -140,10 +144,10 @@ def chain_to_ndarray(
                 atom_name = "SD"
 
             if atom_name in residue_constants.atom_order:
-                atom_positions[res_index, residue_constants.atom_order[atom_name]] = atom.coord
-                atom_mask[res_index, residue_constants.atom_order[atom_name]] = True
+                atom_positions[res_index, residue_constants.atom_order[atom_name]] = atom.coord  # (3,) xyz vector in (l, 37, 3)
+                atom_mask[res_index, residue_constants.atom_order[atom_name]] = True  # scalar element in (l, 37)
                 if is_predicted and atom_name == "CA":
-                    confidence[res_index] = atom.b_factor / PLDDT_B_FACTOR_SCALE
+                    confidence[res_index] = atom.b_factor / PLDDT_B_FACTOR_SCALE  # scalar element in (l,)
 
     if not sequence or not all(sequence):
         raise ValueError("Some residue name was not specified correctly.")
@@ -155,22 +159,23 @@ def chain_to_ndarray(
         insertion_code,
         confidence,
         entity_id,
-    )
+    )  # array fields: (l, 37, 3), (l, 37), (l,), (l,), (l,)
 
 
 @dataclass(frozen=True)
 class ProteinChain:
     """Dataclass with atom37 representation of a single protein chain."""
 
+    # l is len(sequence), including separator rows when present.
     id: str
     sequence: str
     chain_id: str  # author chain id - mutable
     entity_id: int | None
-    residue_index: np.ndarray
-    insertion_code: np.ndarray
-    atom37_positions: np.ndarray
-    atom37_mask: np.ndarray
-    confidence: np.ndarray
+    residue_index: np.ndarray  # (l,)
+    insertion_code: np.ndarray  # (l,)
+    atom37_positions: np.ndarray  # (l, 37, 3)
+    atom37_mask: np.ndarray  # (l, 37)
+    confidence: np.ndarray  # (l,)
     mmcif: MmcifWrapper | None = None
     atom37_confidence: np.ndarray | None = None  # P has shape (l, 37).
 
@@ -186,7 +191,7 @@ class ProteinChain:
         """Yield every protein chain represented in an mmCIF structure."""
         mmcif = path if isinstance(path, MmcifWrapper) else MmcifWrapper.read(path, id)
         for chain in bs.chain_iter(mmcif.structure):
-            chain = chain[bs.filter_amino_acids(chain) & ~chain.hetero]
+            chain = chain[bs.filter_amino_acids(chain) & ~chain.hetero]  # AtomArray with n_protein_atoms entries
             if len(chain) == 0:
                 continue
             chain_id = chain.chain_id[0]
@@ -206,7 +211,7 @@ class ProteinChain:
                 insertion_code,
                 confidence,
                 _,
-            ) = chain_to_ndarray(chain, mmcif, chain_id, is_predicted)
+            ) = chain_to_ndarray(chain, mmcif, chain_id, is_predicted)  # array fields: (l, 37, 3), (l, 37), (l,), (l,), (l,)
             if not all(sequence):
                 raise ValueError("Some residue name was not specified correctly.")
 
@@ -280,7 +285,7 @@ class ProteinChain:
                 stacklevel=2,
             )
 
-        atom_array = mmcif.structure
+        atom_array = mmcif.structure  # AtomArray with n_structure_atoms entries
         (
             sequence,
             atom_positions,
@@ -289,7 +294,7 @@ class ProteinChain:
             insertion_code,
             confidence,
             _,
-        ) = chain_to_ndarray(atom_array, mmcif, chain_id, is_predicted)
+        ) = chain_to_ndarray(atom_array, mmcif, chain_id, is_predicted)  # array fields: (l, 37, 3), (l, 37), (l,), (l,), (l,)
         if not all(sequence):
             raise ValueError("Some residue name was not specified correctly.")
 
@@ -319,15 +324,17 @@ class ProteinChain:
         insertion_code: np.ndarray | None = None,
         confidence: np.ndarray | torch.Tensor | None = None,
     ):
+        # atom37_positions: (l, 37, 3) or (1, l, 37, 3).
+        # Optional residue indices and confidence are (l,) or (1, l).
         if isinstance(atom37_positions, torch.Tensor):
-            atom37_positions = atom37_positions.cpu().numpy()
+            atom37_positions = atom37_positions.cpu().numpy()  # same incoming shape, (l, 37, 3) or (1, l, 37, 3)
             if atom37_positions.ndim == 4:
                 if atom37_positions.shape[0] != 1:
                     raise ValueError(
                         "Cannot handle batched inputs, atom37_positions has shape "
                         f"{atom37_positions.shape}"
                     )
-                atom37_positions = atom37_positions[0]
+                atom37_positions = atom37_positions[0]  # (l, 37, 3)
 
         if not isinstance(atom37_positions, np.ndarray):
             raise TypeError("atom37_positions must be a NumPy array or Torch tensor.")
@@ -338,7 +345,7 @@ class ProteinChain:
             )
         seqlen = atom37_positions.shape[0]
 
-        atom_mask = np.isfinite(atom37_positions).all(-1)
+        atom_mask = np.isfinite(atom37_positions).all(-1)  # (l, 37)
 
         if id is None:
             id = ""
@@ -350,32 +357,32 @@ class ProteinChain:
             chain_id = "A"
 
         if residue_index is None:
-            residue_index = np.arange(1, seqlen + 1)
+            residue_index = np.arange(1, seqlen + 1)  # (l,)
         elif isinstance(residue_index, torch.Tensor):
-            residue_index = residue_index.cpu().numpy()
+            residue_index = residue_index.cpu().numpy()  # same incoming index shape
             if residue_index.ndim == 2:
                 if residue_index.shape[0] != 1:
                     raise ValueError(
                         "Cannot handle batched inputs, residue_index has shape "
                         f"{residue_index.shape}"
                     )
-                residue_index = residue_index[0]
+                residue_index = residue_index[0]  # (l,)
         if not isinstance(residue_index, np.ndarray):
             raise TypeError("residue_index must be a NumPy array or Torch tensor.")
 
         if insertion_code is None:
-            insertion_code = np.array(["" for _ in range(seqlen)])
+            insertion_code = np.array(["" for _ in range(seqlen)])  # (l,)
 
         if confidence is None:
-            confidence = np.ones(seqlen, dtype=np.float32)
+            confidence = np.ones(seqlen, dtype=np.float32)  # (l,)
         elif isinstance(confidence, torch.Tensor):
-            confidence = confidence.cpu().numpy()
+            confidence = confidence.cpu().numpy()  # same incoming confidence shape
             if confidence.ndim == 2:
                 if confidence.shape[0] != 1:
                     raise ValueError(
                         f"Cannot handle batched inputs, confidence has shape {confidence.shape}"
                     )
-                confidence = confidence[0]
+                confidence = confidence[0]  # (l,)
         if not isinstance(confidence, np.ndarray):
             raise TypeError("confidence must be a NumPy array or Torch tensor.")
 
@@ -404,15 +411,16 @@ class ProteinChain:
 
         This function passes all kwargs to from_atom37.
         """
+        # backbone_atom_coordinates: (l, 3, 3) or (1, l, 3, 3), atom then xyz axes.
         if isinstance(backbone_atom_coordinates, torch.Tensor):
-            backbone_atom_coordinates = backbone_atom_coordinates.cpu().numpy()
+            backbone_atom_coordinates = backbone_atom_coordinates.cpu().numpy()  # same incoming shape, (l, 3, 3) or (1, l, 3, 3)
             if backbone_atom_coordinates.ndim == 4:
                 if backbone_atom_coordinates.shape[0] != 1:
                     raise ValueError(
                         f"Cannot handle batched inputs, backbone_atom_coordinates has "
                         f"shape {backbone_atom_coordinates.shape}"
                     )
-                backbone_atom_coordinates = backbone_atom_coordinates[0]
+                backbone_atom_coordinates = backbone_atom_coordinates[0]  # (l, 3, 3)
 
         if not isinstance(backbone_atom_coordinates, np.ndarray):
             raise TypeError(
@@ -431,8 +439,8 @@ class ProteinChain:
             (backbone_atom_coordinates.shape[0], 37, 3),
             np.inf,
             dtype=backbone_atom_coordinates.dtype,
-        )
-        atom37_positions[:, :3, :] = backbone_atom_coordinates
+        )  # (l, 37, 3)
+        atom37_positions[:, :3, :] = backbone_atom_coordinates  # (l, 3, 3) backbone slice
 
         return cls.from_atom37(atom37_positions=atom37_positions, **kwargs)
 
@@ -464,7 +472,7 @@ class ProteinChain:
                 case _:
                     file_id = "null"
 
-        atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"])
+        atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"])  # AtomArray with n_file_atoms entries
         if len(atom_array) == 0:
             raise ValueError("PDB contains no atoms.")
         if chain_id == "detect":
@@ -473,7 +481,7 @@ class ProteinChain:
             bs.filter_amino_acids(atom_array)
             & ~atom_array.hetero
             & (atom_array.chain_id == chain_id)
-        ]
+        ]  # AtomArray with n_selected_protein_atoms entries
         if len(atom_array) == 0:
             raise ValueError(f"PDB contains no amino-acid atoms for chain {chain_id!r}.")
 
@@ -487,17 +495,17 @@ class ProteinChain:
 
         atom_positions = np.full(
             [num_res, residue_constants.atom_type_num, 3], np.nan, dtype=np.float32
-        )
-        atom_mask = np.full([num_res, residue_constants.atom_type_num], False, dtype=bool)
-        residue_index = np.full([num_res], -1, dtype=np.int64)
-        insertion_code = np.full([num_res], "", dtype="<U4")
+        )  # (l, 37, 3)
+        atom_mask = np.full([num_res, residue_constants.atom_type_num], False, dtype=bool)  # (l, 37)
+        residue_index = np.full([num_res], -1, dtype=np.int64)  # (l,)
+        insertion_code = np.full([num_res], "", dtype="<U4")  # (l,)
 
-        confidence = np.ones([num_res], dtype=np.float32)
+        confidence = np.ones([num_res], dtype=np.float32)  # (l,)
 
         for i, res in enumerate(bs.residue_iter(atom_array)):
             res_index = res[0].res_id
-            residue_index[i] = res_index
-            insertion_code[i] = res[0].ins_code
+            residue_index[i] = res_index  # scalar element in (l,)
+            insertion_code[i] = res[0].ins_code  # scalar element in (l,)
 
             # Atom level features
             for atom in res:
@@ -507,10 +515,10 @@ class ProteinChain:
                     atom_name = "SD"
 
                 if atom_name in residue_constants.atom_order:
-                    atom_positions[i, residue_constants.atom_order[atom_name]] = atom.coord
-                    atom_mask[i, residue_constants.atom_order[atom_name]] = True
+                    atom_positions[i, residue_constants.atom_order[atom_name]] = atom.coord  # (3,) xyz vector in (l, 37, 3)
+                    atom_mask[i, residue_constants.atom_order[atom_name]] = True  # scalar element in (l, 37)
                     if is_predicted and atom_name == "CA":
-                        confidence[i] = atom.b_factor / PLDDT_B_FACTOR_SCALE
+                        confidence[i] = atom.b_factor / PLDDT_B_FACTOR_SCALE  # scalar element in (l,)
 
         if not sequence or not all(sequence):
             raise ValueError("Some residue name was not specified correctly.")
@@ -567,7 +575,7 @@ class ProteinChain:
     ) -> ProteinChain:
         """A simple converter from bs.AtomArray -> ProteinChain.
         Uses PDB file format as intermediate."""
-        atom_array = atom_array.copy()
+        atom_array = atom_array.copy()  # AtomArray copy with unchanged atom count
         atom_array.box = None  # remove surrounding box, from_pdb won't handle this
         pdb_file = PDBFile()  # pyright: ignore
         pdb_file.set_structure(atom_array)
@@ -596,7 +604,7 @@ class ProteinChain:
             "residue_index": self.residue_index,
             "insertion_code": self.insertion_code,
             "confidence": self.confidence,
-        }
+        }  # arrays share l: positions (l, 37, 3), masks (l, 37), residue fields (l,)
         for name, values in aligned.items():
             if not isinstance(values, np.ndarray):
                 raise TypeError(f"{name} must be a NumPy array, got {type(values).__name__}.")
@@ -636,7 +644,7 @@ class ProteinChain:
                 )
         if not np.issubdtype(self.confidence.dtype, np.number):
             raise TypeError("confidence must use a numeric dtype.")
-        atom37_confidence = self.atom37_confidence
+        atom37_confidence = self.atom37_confidence  # (l, 37) or None
         if atom37_confidence is not None and not isinstance(atom37_confidence, np.ndarray):
             raise TypeError("atom37_confidence must be a NumPy array when provided.")
         if (
@@ -686,7 +694,7 @@ class ProteinChain:
                     self.atom37_confidence[res_idx_i, i]
                     if self.atom37_confidence is not None
                     else conf
-                )
+                )  # scalar atom confidence
                 atom = bs.Atom(
                     coord=pos,
                     chain_id="A" if self.chain_id is None else self.chain_id,
@@ -700,7 +708,7 @@ class ProteinChain:
                     occupancy=1.0,
                 )
                 atoms.append(atom)
-        return bs.array(atoms)
+        return bs.array(atoms)  # AtomArray with n_present_atoms entries
 
     # Coordinate transformations and dataset adapters
     def get_normalization_frame(self) -> Affine3D:
@@ -711,8 +719,8 @@ class ProteinChain:
         Returns:
             Affine3D: [] tensor of Affine3D frame
         """
-        coords = torch.from_numpy(self.atom37_positions)
-        frame = get_protein_normalization_frame(coords)
+        coords = torch.from_numpy(self.atom37_positions)  # (l, 37, 3)
+        frame = get_protein_normalization_frame(coords)  # one normalization frame
 
         return frame
 
@@ -725,9 +733,10 @@ class ProteinChain:
         Returns:
             ProteinChain: Transformed protein chain
         """
-        coords = torch.from_numpy(self.atom37_positions).to(frame.trans.dtype)
-        coords = apply_frame_to_coords(coords, frame)
-        atom37_positions = coords.numpy()
+        # frame is a rigid transform broadcast-compatible with (l, 37, 3) coordinates.
+        coords = torch.from_numpy(self.atom37_positions).to(frame.trans.dtype)  # (l, 37, 3)
+        coords = apply_frame_to_coords(coords, frame)  # (l, 37, 3)
+        atom37_positions = coords.numpy()  # (l, 37, 3)
         return replace(self, atom37_positions=atom37_positions)
 
     def normalize_coordinates(self) -> ProteinChain:
@@ -736,36 +745,36 @@ class ProteinChain:
 
     def infer_oxygen(self) -> ProteinChain:
         """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided."""
-        O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze()
+        O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze()  # (n_missing,) or () when exactly one oxygen is missing
 
-        O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32)
-        N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1)
-        N = torch.roll(N, -3)
-        N[..., -1, :] = torch.nan
+        O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32)  # (3,)
+        N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1)  # each (l, 3)
+        N = torch.roll(N, -3)  # (l, 3); torch.roll keeps the original shape
+        N[..., -1, :] = torch.nan  # (3,) xyz row
 
         # Get the frame defined by the CA-C-N atom
-        frames = Affine3D.from_graham_schmidt(CA, C, N)
-        oxygen_coordinates = frames.apply(O_vector)
-        atom37_positions = self.atom37_positions.copy()
-        atom37_mask = self.atom37_mask.copy()
+        frames = Affine3D.from_graham_schmidt(CA, C, N)  # affine batch shape: (l,)
+        oxygen_coordinates = frames.apply(O_vector)  # (l, 3)
+        atom37_positions = self.atom37_positions.copy()  # (l, 37, 3)
+        atom37_mask = self.atom37_mask.copy()  # (l, 37)
 
         atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = oxygen_coordinates[
             O_missing_indices
-        ].numpy()
+        ].numpy()  # (n_missing, 3) or (3,) selected oxygen coordinates
         atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan(
             atom37_positions[O_missing_indices, residue_constants.atom_order["O"]]
-        ).any(-1)
+        ).any(-1)  # (n_missing,) or () selected oxygen mask
         new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
         return new_chain
 
     @cached_property
     def inferred_cbeta(self) -> np.ndarray:
         """Infer cbeta positions based on N, C, CA."""
-        N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0)
+        N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0)  # each (l, 3)
         # See usage in trDesign codebase.
         # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140
-        CB = infer_cb(C, N, CA, 1.522, 1.927, -2.143)
-        return CB
+        CB = infer_cb(C, N, CA, 1.522, 1.927, -2.143)  # (l, 3)
+        return CB  # (l, 3)
 
     def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinChain:
         """Return a new chain with inferred CB atoms at all residues except GLY.
@@ -780,30 +789,30 @@ class ProteinChain:
                 calculation between two designs for a given structural template, w/
                 CB atoms.
         """
-        atom37_positions = self.atom37_positions.copy()
-        atom37_mask = self.atom37_mask.copy()
+        atom37_positions = self.atom37_positions.copy()  # (l, 37, 3)
+        atom37_mask = self.atom37_mask.copy()  # (l, 37)
 
-        inferred_cbeta_positions = self.inferred_cbeta
+        inferred_cbeta_positions = self.inferred_cbeta  # (l, 3)
         if not infer_cbeta_for_glycine:
-            inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan
+            inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan  # (n_glycine, 3) selected rows
 
-        atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions
+        atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions  # (l, 3) C-beta slice
         atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan(
             atom37_positions[:, residue_constants.atom_order["CB"]]
-        ).any(-1)
+        ).any(-1)  # (l,) C-beta mask
         new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
         return new_chain
 
     @cached_property
     def pdist_CA(self) -> np.ndarray:
-        CA = self.atoms["CA"]
-        pdist_CA = squareform(pdist(CA))
-        return pdist_CA
+        CA = self.atoms["CA"]  # (l, 3)
+        pdist_CA = squareform(pdist(CA))  # (l, l)
+        return pdist_CA  # (l, l)
 
     @cached_property
     def pdist_CB(self) -> np.ndarray:
-        pdist_CB = squareform(pdist(self.inferred_cbeta))
-        return pdist_CB
+        pdist_CB = squareform(pdist(self.inferred_cbeta))  # (l, l)
+        return pdist_CB  # (l, l)
 
     @classmethod
     def as_complex(cls, chains: Sequence[ProteinChain]):
@@ -824,23 +833,24 @@ class ProteinChain:
             "atom37_positions": np.full([1, 37, 3], np.inf),
             "atom37_mask": np.zeros([1, 37], dtype=bool),
             "confidence": np.array([0]),
-        }
+        }  # one-residue separator arrays: (1,), (1, 37, 3), or (1, 37)
 
         def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray):
+            # arrays: (l_i, *trailing_shape); separator: (1, *trailing_shape).
             if use_chainbreak:
                 full_array = []
                 for array in arrays:
                     full_array.append(array)
                     full_array.append(sep)
                 full_array = full_array[:-1]
-                return np.concatenate(full_array, 0)
+                return np.concatenate(full_array, 0)  # (sum(chain_lengths) + n_chains - 1, *trailing_shape)
             else:
-                return np.concatenate(arrays, 0)
+                return np.concatenate(arrays, 0)  # (sum(chain_lengths), *trailing_shape)
 
         array_args: dict[str, np.ndarray] = {
             name: join_arrays([getattr(chain, name) for chain in chains], sep)
             for name, sep in sep_tokens.items()
-        }
+        }  # each array retains its trailing atom/xyz axes
 
         chain_break = residue_constants.CHAIN_BREAK_TOKEN if use_chainbreak else ""
         return cls(
@@ -868,15 +878,15 @@ class ProteinChain:
                 raise ValueError(
                     f"Non-polymer {nonpolymer.comp_id!r} has no coordinate table."
                 )
-            chain_coords = self.atom37_positions[self.atom37_mask]
-            distance = cdist(nonpolymer_array.coord, chain_coords)
+            chain_coords = self.atom37_positions[self.atom37_mask]  # (n_present_atoms, 3)
+            distance = cdist(nonpolymer_array.coord, chain_coords)  # (n_ligand_atoms, n_present_atoms)
 
-            is_contact = distance < 5
+            is_contact = distance < 5  # (n_ligand_atoms, n_present_atoms)
             if not is_contact.any():
                 continue
-            contacting_atoms = np.where(is_contact.any(0))[0]
-            chain_index = np.where(self.atom37_mask)[0]
-            contacting_residues = np.unique(chain_index[contacting_atoms])
+            contacting_atoms = np.where(is_contact.any(0))[0]  # (n_contacting_atoms,)
+            chain_index = np.where(self.atom37_mask)[0]  # (n_present_atoms,)
+            contacting_residues = np.unique(chain_index[contacting_atoms])  # (n_contacting_residues,)
 
             result = {
                 "ligand": nonpolymer.name,
@@ -890,7 +900,7 @@ class ProteinChain:
         self, indices: list[int | str], ignore_x_mismatch: bool = False
     ) -> ProteinChain:
         numeric_indices = [idx if isinstance(idx, int) else int(idx[1:]) for idx in indices]
-        mask = np.isin(self.residue_index, numeric_indices)
+        mask = np.isin(self.residue_index, numeric_indices)  # (l,)
         new = self[mask]
         mismatches = []
         for aa, idx in zip(new.sequence, indices, strict=False):
@@ -922,20 +932,22 @@ class ProteinChain:
         # Convert to tensors and add batch dimension
         coordinates = (
             torch.from_numpy(self.atom37_positions).float().unsqueeze(0)
-        )  # X has shape (1, l, 37, 3).
-        plddt = torch.from_numpy(self.confidence).float().unsqueeze(0)  # P: (1, l)
+        )  # X has shape (1, l, 37, 3).; (1, l, 37, 3)
+        plddt = torch.from_numpy(self.confidence).float().unsqueeze(0)  # P: (1, l); (1, l)
         residue_index = (
             torch.from_numpy(self.residue_index).long().unsqueeze(0)
-        )  # R has shape (1, l).
+        )  # R has shape (1, l).; (1, l)
 
-        return coordinates, plddt, residue_index
+        return coordinates, plddt, residue_index  # (1, l, 37, 3), (1, l), (1, l)
 
     # Sequence access, interchange, and compact storage
     def __getitem__(self, idx: int | list[int] | slice | np.ndarray | torch.Tensor):
+        # idx selects residues; an integer is promoted to a length-one index.
+        # Returned fields retain atom/xyz trailing axes with the selected residue count.
         if isinstance(idx, int):
             idx = [idx]
         if isinstance(idx, torch.Tensor):
-            idx = idx.cpu().numpy()
+            idx = idx.cpu().numpy()  # same index shape
 
         sequence = slice_python_object_as_numpy(self.sequence, idx)
         return replace(
@@ -955,11 +967,11 @@ class ProteinChain:
         return len(self.sequence)
 
     def cbeta_contacts(self, distance_threshold: float = 8.0) -> np.ndarray:
-        distance = self.pdist_CB
-        contacts = (distance < distance_threshold).astype(np.int64)
-        contacts[np.isnan(distance)] = -1
+        distance = self.pdist_CB  # (l, l)
+        contacts = (distance < distance_threshold).astype(np.int64)  # (l, l)
+        contacts[np.isnan(distance)] = -1  # (n_missing_pairs,) selected entries
         np.fill_diagonal(contacts, -1)
-        return contacts
+        return contacts  # (l, l)
 
     def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True):
         """Dssp works better w/o insertions."""
@@ -989,7 +1001,7 @@ class ProteinChain:
                 "mode": CIFColumn(data=CIFData(array=np.array(["global", "local"]), dtype=np.str_)),
                 "name": CIFColumn(data=CIFData(array=np.array(["pLDDT", "pLDDT"]), dtype=np.str_)),
             },
-        )
+        )  # each CIF metric column has shape (2,)
 
         # table is a duplicate of data already in the atom array, but
         # needed by molstar to render pLDDT / confidence
@@ -1039,10 +1051,10 @@ class ProteinChain:
         need more than 2**32 residues..."""
         dct = {k: v for k, v in asdict(self).items() if k not in ["mmcif"]}
         if backbone_only:
-            dct["atom37_mask"][:, 3:] = False
-        dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]]
+            dct["atom37_mask"][:, 3:] = False  # (l, 34) mask slice for atoms beyond N/CA/C
+        dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]]  # (n_present_atoms, 3)
         if dct.get("atom37_confidence") is not None:
-            dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]]
+            dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]]  # (n_present_atoms,)
         else:
             dct.pop("atom37_confidence", None)
 
@@ -1050,9 +1062,9 @@ class ProteinChain:
             if isinstance(v, np.ndarray):
                 match v.dtype:
                     case np.int64:
-                        dct[k] = v.astype(np.int32)
+                        dct[k] = v.astype(np.int32)  # v.shape
                     case np.float64 | np.float32:
-                        dct[k] = v.astype(np.float16)
+                        dct[k] = v.astype(np.float16)  # v.shape
                     case _:
                         pass
                 if json_serializable:
@@ -1074,15 +1086,15 @@ class ProteinChain:
 
         for k, v in dct.items():
             if isinstance(v, list):
-                dct[k] = np.array(v)
+                dct[k] = np.array(v)  # shape inferred from serialized nested list
 
-        atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan)
-        atom37[dct["atom37_mask"]] = dct["atom37_positions"]
-        dct["atom37_positions"] = atom37
+        atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan)  # (l, 37, 3)
+        atom37[dct["atom37_mask"]] = dct["atom37_positions"]  # (n_present_atoms, 3) selected coordinates
+        dct["atom37_positions"] = atom37  # (l, 37, 3)
         if "atom37_confidence" in dct:
-            atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32)
-            atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"]
-            dct["atom37_confidence"] = atom37_conf
+            atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32)  # (l, 37)
+            atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"]  # (n_present_atoms,) selected confidence values
+            dct["atom37_confidence"] = atom37_conf  # (l, 37)
         dct = {
             k: (
                 v.astype(np.float32)
@@ -1091,7 +1103,7 @@ class ProteinChain:
             )
             for k, v in dct.items()
             if not (k == "atom37_confidence" and v is None)
-        }
+        }  # each converted array retains its serialized field shape
         return cls(**dct, mmcif=None)
 
     @classmethod
@@ -1111,10 +1123,10 @@ class ProteinChain:
 
     # Surface and structural comparison metrics
     def sasa(self, by_residue: bool = True):
-        arr = self.atom_array_no_insertions
+        arr = self.atom_array_no_insertions  # AtomArray with n_present_atoms entries
         if len(arr) == 0:
             raise ValueError("SASA requires at least one resolved atom.")
-        sasa_per_atom = bs.sasa(arr)  # type: ignore
+        sasa_per_atom = bs.sasa(arr)  # type: ignore; (n_present_atoms,)
         if by_residue:
             # Sum per-atom SASA into residue "bins", with np.bincount.
             if arr.res_id is None:
@@ -1127,12 +1139,12 @@ class ProteinChain:
                     np.bincount(arr.res_id, weights=sasa_per_atom)[1:],
                     np.zeros(num_trailing_residues),
                 ]
-            )
-            sasa_per_residue[~self.atom37_mask.any(-1)] = np.nan
+            )  # (l,)
+            sasa_per_residue[~self.atom37_mask.any(-1)] = np.nan  # (n_missing_residues,) selected entries
             if len(sasa_per_residue) != len(self):
                 raise RuntimeError("Residue SASA output does not align with the protein chain.")
-            return sasa_per_residue
-        return sasa_per_atom
+            return sasa_per_residue  # (l,)
+        return sasa_per_atom  # (n_present_atoms,)
 
     def sap_score(self, aggregation: str = "atom") -> np.ndarray:
         """Compute per-atom spatial aggregation propensity (SAP).
@@ -1141,7 +1153,7 @@ class ProteinChain:
         Protein aggregation sums positive atom scores, following Lauer et al. 2011.
         """
         sap_radius = 5.0
-        arr = self.atom_array_no_insertions
+        arr = self.atom_array_no_insertions  # AtomArray with n_present_atoms entries
         if len(arr) == 0:
             raise ValueError("SAP requires at least one resolved atom.")
 
@@ -1150,42 +1162,42 @@ class ProteinChain:
                 raise RuntimeError(f"Biotite AtomArray is missing required {name!r} data.")
 
         # compute SASA and residue-specific properties
-        sasa_per_atom = self.sasa(by_residue=False)
+        sasa_per_atom = self.sasa(by_residue=False)  # (n_present_atoms,)
         resid_to_resname = dict(zip(arr.res_id, arr.res_name, strict=False))
 
-        max_side_chain_asa = np.full(len(self), np.nan)
-        res_hydrophobicity = np.full(len(self), np.nan)
-        resolved_res_mask = self.atom37_mask.any(-1)
+        max_side_chain_asa = np.full(len(self), np.nan)  # (l,)
+        res_hydrophobicity = np.full(len(self), np.nan)  # (l,)
+        resolved_res_mask = self.atom37_mask.any(-1)  # (l,)
         num_trailing_residues = len(self) - arr.res_id.max()
 
         max_side_chain_asa[resolved_res_mask] = np.array(
             [residue_constants.side_chain_asa[resid_to_resname[i]] for i in np.unique(arr.res_id)]
-        )
+        )  # (n_resolved_residues,) selected entries
         res_hydrophobicity[resolved_res_mask] = np.array(
             [residue_constants.hydrophobicity[resid_to_resname[i]] for i in np.unique(arr.res_id)]
-        )
+        )  # (n_resolved_residues,) selected entries
 
         # compute SAP score
-        is_side_chain = ~bs.filter_peptide_backbone(arr)
-        sasa_per_atom[is_side_chain] = 0
+        is_side_chain = ~bs.filter_peptide_backbone(arr)  # (n_present_atoms,)
+        sasa_per_atom[is_side_chain] = 0  # (n_selected_atoms,) selected entries
         kdtree = KDTree(arr.coord)
         neighbors = kdtree.query_ball_tree(kdtree, sap_radius, p=2.0)
-        sap_by_atom = np.zeros_like(sasa_per_atom)
+        sap_by_atom = np.zeros_like(sasa_per_atom)  # (n_present_atoms,)
         for i, nn_list in enumerate(neighbors):
-            saa_nn = np.zeros_like(sasa_per_atom)
-            saa_nn[nn_list] = sasa_per_atom[nn_list]
+            saa_nn = np.zeros_like(sasa_per_atom)  # (n_present_atoms,)
+            saa_nn[nn_list] = sasa_per_atom[nn_list]  # (n_neighbors,) selected entries
             sasa_within_r = np.concatenate(
                 [
                     np.bincount(arr.res_id, weights=saa_nn)[1:],
                     np.zeros(num_trailing_residues),
                 ]
-            )
-            sap = np.nansum((sasa_within_r / max_side_chain_asa) * res_hydrophobicity)
-            sap_by_atom[i] = sap
+            )  # (l,)
+            sap = np.nansum((sasa_within_r / max_side_chain_asa) * res_hydrophobicity)  # scalar NumPy reduction
+            sap_by_atom[i] = sap  # scalar entry in (n_present_atoms,)
 
         match aggregation:
             case "atom":
-                return sap_by_atom
+                return sap_by_atom  # (n_present_atoms,)
             case "residue":
                 sap_by_residue = np.concatenate(
                     [
@@ -1195,11 +1207,11 @@ class ProteinChain:
                 ) / (
                     np.concatenate([np.bincount(arr.res_id)[1:], np.zeros(num_trailing_residues)])
                     + 1e-8
-                )
-                sap_by_residue[~resolved_res_mask] = np.nan
+                )  # (l,)
+                sap_by_residue[~resolved_res_mask] = np.nan  # (n_missing_residues,) selected entries
                 if len(sap_by_residue) != len(self):
                     raise RuntimeError("Residue SAP output does not align with the protein chain.")
-                return sap_by_residue
+                return sap_by_residue  # (l,)
             case "protein":
                 return sum(sap_by_atom[sap_by_atom > 0])  # pyright: ignore[reportReturnType]
             case _:
@@ -1216,19 +1228,19 @@ class ProteinChain:
 
         # https://www.mdpi.com/2073-4352/11/12/1539
         # The non-overlapping-atom approximation can produce globularity above one.
-        mask = self.atom37_mask.any(-1)
-        points = self.atom37_positions[self.atom37_mask]
+        mask = self.atom37_mask.any(-1)  # (l,)
+        points = self.atom37_positions[self.atom37_mask]  # (n_present_atoms, 3)
         sequence = [aa for aa, m in zip(self.sequence, mask, strict=False) if m]  # type: ignore
-        A, _ = self._mvee(points, tol=1e-3)
-        mvee_volume = (4 * np.pi) / (3 * np.sqrt(np.linalg.det(A)))
+        A, _ = self._mvee(points, tol=1e-3)  # A: (3, 3), center: (3, 1)
+        mvee_volume = (4 * np.pi) / (3 * np.sqrt(np.linalg.det(A)))  # NumPy scalar ()
         volume = sum(residue_constants.amino_acid_volumes[x] for x in sequence)
         ratio = volume / mvee_volume
 
         # The paper compares the ellipsoidal profile with scalar t, a measurement
         # of elongation. We want a single number, so we multiply by 1/(2t), so
         # that value is normalized between 0-1
-        eigenvalues = np.linalg.eigvals(A)
-        R = 1 / np.sqrt(eigenvalues)
+        eigenvalues = np.linalg.eigvals(A)  # (3,)
+        R = 1 / np.sqrt(eigenvalues)  # (3,)
         # ellipsoid radii length triangle inequality coefficient
         t = max(R[0] / (R[1] + R[2]), R[1] / (R[0] + R[2]), R[2] / (R[0] + R[1]))
         elongation_metric = 1 / max(t, 1)
@@ -1239,50 +1251,51 @@ class ProteinChain:
         # Finds minimum volume enclosing ellipsoid of a set of points.
         # Returns A, c where the ellipse is defined as:
         #    (x-c).T @ A @ (x-c) = 1
+        # P: (n_input_points, d); d is coordinate dimension. Hull selection changes only point count.
         hull = ConvexHull(P)
-        P = P[hull.vertices]
-        P = P.T
+        P = P[hull.vertices]  # (n_hull_vertices, d)
+        P = P.T  # (d, n); n = n_hull_vertices
 
         # Data points
         d, n = P.shape
-        Q = np.zeros((d + 1, n))
-        Q[:d, :] = P[:d, :n]
-        Q[d, :] = np.ones((1, n))
+        Q = np.zeros((d + 1, n))  # (d + 1, n)
+        Q[:d, :] = P[:d, :n]  # (d, n) slice
+        Q[d, :] = np.ones((1, n))  # (n,) row, assigned from (1, n)
 
         # Initializations
         count = 1
         err = 1.0
-        u = np.full((n, 1), 1 / n)  # First iteration.
+        u = np.full((n, 1), 1 / n)  # First iteration.; (n, 1)
 
         # Khachiyan Algorithm
         for _ in range(max_iter):
-            X = Q.dot(np.diag(u.squeeze())) @ Q.T
-            M = np.diag(Q.T @ np.linalg.inv(X) @ Q)
-            maximum, j = np.max(M), np.argmax(M)
+            X = Q.dot(np.diag(u.squeeze())) @ Q.T  # (d + 1, d + 1)
+            M = np.diag(Q.T @ np.linalg.inv(X) @ Q)  # (n,)
+            maximum, j = np.max(M), np.argmax(M)  # each NumPy scalar ()
             step_size = (maximum - d - 1) / ((d + 1) * (maximum - 1))
-            new_u = (1 - step_size) * u
-            new_u[j] += step_size
+            new_u = (1 - step_size) * u  # (n, 1)
+            new_u[j] += step_size  # (1,) selected row
             count += 1
-            err = np.linalg.norm(new_u - u)
-            u = new_u
+            err = np.linalg.norm(new_u - u)  # NumPy scalar ()
+            u = new_u  # (n, 1)
             if err < tol:
                 break
         else:
             raise ValueError("MVEE did not converge")
 
         d = P.shape[0]  # Fixed: use P.shape[0] instead of P.shape
-        U = np.diag(u.squeeze())
+        U = np.diag(u.squeeze())  # (n, n)
 
         # The A matrix for the ellipse
-        A = (1 / d) * np.linalg.inv(P @ U @ P.T - (P @ u) @ (P @ u).T)
+        A = (1 / d) * np.linalg.inv(P @ U @ P.T - (P @ u) @ (P @ u).T)  # (d, d)
 
         # Center of the ellipse
-        c = P @ u
+        c = P @ u  # (d, 1)
 
-        return A, c
+        return A, c  # (d, d), (d, 1)
 
     def radius_of_gyration(self):
-        arr = self.atom_array_no_insertions
+        arr = self.atom_array_no_insertions  # AtomArray with n_present_atoms entries
         return bs.gyration_radius(arr)
 
     def align(
@@ -1375,8 +1388,8 @@ class ProteinChain:
             torch.tensor(native.atom37_positions[target_inds]).unsqueeze(0),
             torch.tensor(native.atom37_mask[mobile_inds]).unsqueeze(0),
             **kwargs,
-        )
-        return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten()
+        )  # score shape follows selected coordinate axes and per_residue
+        return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten()  # scalar or (lddt.numel(),)
 
     def gdt_ts(
         self,
@@ -1409,12 +1422,12 @@ class ProteinChain:
                 & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1)
             ).unsqueeze(0),
             **kwargs,
-        )
-        return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten()
+        )  # () or (n_samples,), selected by reduction
+        return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten()  # scalar or (gdt_ts.numel(),)
 
     @cached_property
     def residue_index_no_insertions(self) -> np.ndarray:
-        return self.residue_index + np.cumsum(self.insertion_code != "")
+        return self.residue_index + np.cumsum(self.insertion_code != "")  # (l,)
 
     @cached_property
     def atom_array_no_insertions(self) -> bs.AtomArray:
@@ -1433,7 +1446,7 @@ class ProteinChain:
                     self.atom37_confidence[res_idx, i]
                     if self.atom37_confidence is not None
                     else conf
-                )
+                )  # scalar atom confidence
                 atom = bs.Atom(
                     coord=pos,
                     # hard coded to as we currently only support single chain structures
@@ -1447,4 +1460,4 @@ class ProteinChain:
                     occupancy=1.0,
                 )
                 atoms.append(atom)
-        return bs.array(atoms)
+        return bs.array(atoms)  # AtomArray with n_present_atoms entries

fastplms/models/esmfold2/esmfold2_protein_complex.py

@@ -7,6 +7,13 @@ import itertools
 import random
 import re
 import warnings
+import biotite.structure as bs
+import brotli
+import msgpack
+import msgpack_numpy
+import numpy as np
+import torch
+
 from collections.abc import Iterable, Sequence
 from dataclasses import asdict, dataclass, replace
 from functools import cached_property
@@ -14,13 +21,6 @@ from pathlib import Path
 from subprocess import check_output
 from tempfile import TemporaryDirectory
 from typing import Any
-
-import biotite.structure as bs
-import brotli
-import msgpack
-import msgpack_numpy
-import numpy as np
-import torch
 from biotite.database import rcsb
 from biotite.file import InvalidFileError
 from biotite.structure.io.pdb import PDBFile
@@ -50,6 +50,7 @@ from .esmfold2_protein_chain import (
 )
 from .esmfold2_utils_types import PathOrBuffer
 
+
 SINGLE_LETTER_CHAIN_IDS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
 
 
@@ -73,14 +74,15 @@ def _parse_operation_expression(expression: str) -> list[tuple[str, ...]]:
 
 def _apply_transformations_fast(chains, transformation_dict, operations):
     """Return transformed copies of each affected protein chain."""
+    # Each chain supplies coordinates (l, 37, 3); rotations are (3, 3), translations are (3,).
     transformed_chains = []
     for chain in chains:
         for operation in operations:
-            coordinates = chain.atom37_positions.copy()
+            coordinates = chain.atom37_positions.copy()  # (l, 37, 3)
             for op_step in operation:
                 transform = transformation_dict[op_step]
-                coordinates = matrix_rotate(coordinates, transform.rotation)
-                coordinates += transform.target_translation
+                coordinates = matrix_rotate(coordinates, transform.rotation)  # (l, 37, 3)
+                coordinates += transform.target_translation  # (l, 37, 3)
             transformed_chains.append(replace(chain, atom37_positions=coordinates))
     return transformed_chains
 
@@ -124,16 +126,17 @@ class DockQResult:
 class ProteinComplex:
     """Dataclass with atom37 representation of an entire protein complex."""
 
+    # l is len(sequence), including separator rows when present.
     id: str
     sequence: str
     entity_id: np.ndarray  # entities map to unique sequences
     chain_id: np.ndarray  # multiple chains might share an entity id
     sym_id: np.ndarray  # complexes might be copies of the same chain
-    residue_index: np.ndarray
-    insertion_code: np.ndarray
-    atom37_positions: np.ndarray
-    atom37_mask: np.ndarray
-    confidence: np.ndarray
+    residue_index: np.ndarray  # (l,)
+    insertion_code: np.ndarray  # (l,)
+    atom37_positions: np.ndarray  # (l, 37, 3)
+    atom37_mask: np.ndarray  # (l, 37)
+    confidence: np.ndarray  # (l,)
     # This metadata is parsed from the MMCIF file. For synthetic data, we do a best effort.
     metadata: ProteinComplexMetadata
     atom37_confidence: np.ndarray | None = None  # P has shape (l, 37).
@@ -141,25 +144,25 @@ class ProteinComplex:
     # Coordinate completion, concatenation, and comparison
     def infer_oxygen(self) -> ProteinComplex:
         """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided."""
-        O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze()
+        O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze()  # (n_missing,) or () when exactly one oxygen is missing
 
-        O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32)
-        N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1)
-        N = torch.roll(N, -3)
-        N[..., -1, :] = torch.nan
+        O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32)  # (3,)
+        N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1)  # each (l, 3)
+        N = torch.roll(N, -3)  # (l, 3); torch.roll keeps the original shape
+        N[..., -1, :] = torch.nan  # (3,) xyz row
 
         # Get the frame defined by the CA-C-N atom
-        frames = Affine3D.from_graham_schmidt(CA, C, N)
-        oxygen_coordinates = frames.apply(O_vector)
-        atom37_positions = self.atom37_positions.copy()
-        atom37_mask = self.atom37_mask.copy()
+        frames = Affine3D.from_graham_schmidt(CA, C, N)  # affine batch shape: (l,)
+        oxygen_coordinates = frames.apply(O_vector)  # (l, 3)
+        atom37_positions = self.atom37_positions.copy()  # (l, 37, 3)
+        atom37_mask = self.atom37_mask.copy()  # (l, 37)
 
         atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = oxygen_coordinates[
             O_missing_indices
-        ].numpy()
+        ].numpy()  # (n_missing, 3) or (3,) selected oxygen coordinates
         atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan(
             atom37_positions[O_missing_indices, residue_constants.atom_order["O"]]
-        ).any(-1)
+        ).any(-1)  # (n_missing,) or () selected oxygen mask
         new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
         return new_chain
 
@@ -176,20 +179,20 @@ class ProteinComplex:
                 calculation between two designs for a given structural template, w/
                 CB atoms.
         """
-        atom37_positions = self.atom37_positions.copy()
-        atom37_mask = self.atom37_mask.copy()
+        atom37_positions = self.atom37_positions.copy()  # (l, 37, 3)
+        atom37_mask = self.atom37_mask.copy()  # (l, 37)
 
-        N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0)
+        N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0)  # each (l, 3)
         # See usage in trDesign codebase.
         # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140
-        inferred_cbeta_positions = infer_cb(C, N, CA, 1.522, 1.927, -2.143)
+        inferred_cbeta_positions = infer_cb(C, N, CA, 1.522, 1.927, -2.143)  # (l, 3)
         if not infer_cbeta_for_glycine:
-            inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan
+            inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan  # (n_glycine, 3) selected rows
 
-        atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions
+        atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions  # (l, 3) C-beta slice
         atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan(
             atom37_positions[:, residue_constants.atom_order["CB"]]
-        ).any(-1)
+        ).any(-1)  # (l,) C-beta mask
         new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
         return new_chain
 
@@ -281,8 +284,8 @@ class ProteinComplex:
             torch.tensor(target.atom37_positions[target_inds]).unsqueeze(0),
             torch.tensor(aligned.atom37_mask[mobile_inds]).unsqueeze(0),
             **kwargs,
-        )
-        return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten()
+        )  # score shape follows selected coordinate axes and per_residue
+        return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten()  # scalar or (lddt.numel(),)
 
     def gdt_ts(
         self,
@@ -317,8 +320,8 @@ class ProteinComplex:
                 & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1)
             ).unsqueeze(0),
             **kwargs,
-        )
-        return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten()
+        )  # () or (n_samples,), selected by reduction
+        return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten()  # scalar or (gdt_ts.numel(),)
 
     def dockq(self, native: ProteinComplex):
         # This function uses dockqv2 to compute the DockQ score. Because it does a mapping
@@ -452,7 +455,7 @@ class ProteinComplex:
             "entity_id": self.entity_id,
             "chain_id": self.chain_id,
             "sym_id": self.sym_id,
-        }
+        }  # arrays share l: positions (l, 37, 3), masks (l, 37), residue fields (l,)
         for name, values in aligned.items():
             if not isinstance(values, np.ndarray):
                 raise TypeError(f"{name} must be a NumPy array, got {type(values).__name__}.")
@@ -489,7 +492,7 @@ class ProteinComplex:
                 )
         if not np.issubdtype(self.confidence.dtype, np.number):
             raise TypeError("confidence must use a numeric dtype.")
-        atom37_confidence = self.atom37_confidence
+        atom37_confidence = self.atom37_confidence  # (l, 37) or None
         if atom37_confidence is not None and not isinstance(atom37_confidence, np.ndarray):
             raise TypeError("atom37_confidence must be a NumPy array when provided.")
         if (
@@ -506,6 +509,7 @@ class ProteinComplex:
         NOTE: When slicing with a boolean mask, it's possible that the output array won't
         be the expected length. This is because we do our best to preserve chainbreak tokens.
         """
+        # idx selects residue rows; masks retain separator rows before repeated separators are removed.
 
         if isinstance(idx, int):
             idx = [idx]
@@ -513,8 +517,8 @@ class ProteinComplex:
             raise ValueError("ProteinComplex doesn't supports indexing with lists of indices")
 
         if isinstance(idx, np.ndarray):
-            is_chainbreak = np.asarray([s == "|" for s in self.sequence])
-            idx = idx.astype(bool) | is_chainbreak
+            is_chainbreak = np.asarray([s == "|" for s in self.sequence])  # (l,)
+            idx = idx.astype(bool) | is_chainbreak  # (l,)
 
         complex = self._unsafe_slice(idx)
         if len(complex) == 0:
@@ -524,17 +528,18 @@ class ProteinComplex:
         chainbreak_runs = np.asarray(
             [complex.sequence[i : i + 2] == "||" for i in range(len(complex.sequence) - 1)]
             + [complex.sequence[-1] == "|"]
-        )
+        )  # (selected_length,)
         # We should remove as many chainbreaks as possible from the start of the sequence
         for i in range(len(chainbreak_runs)):
             if complex.sequence[i] == "|":
-                chainbreak_runs[i] = True
+                chainbreak_runs[i] = True  # scalar mask element
             else:
                 break
         complex = complex._unsafe_slice(~chainbreak_runs)
         return complex
 
     def _unsafe_slice(self, idx: int | list[int] | slice | np.ndarray):
+        # idx selects residue rows; atom/xyz trailing axes and aligned field lengths are retained.
         sequence = slice_python_object_as_numpy(self.sequence, idx)
         return replace(
             self,
@@ -569,7 +574,7 @@ class ProteinComplex:
 
     @cached_property
     def chain_lengths(self) -> np.ndarray:
-        return np.diff(self.chain_boundaries, axis=1).flatten()
+        return np.diff(self.chain_boundaries, axis=1).flatten()  # (n_chains,)
 
     @cached_property
     def chain_boundaries(self) -> list[tuple[int, int]]:
@@ -664,11 +669,11 @@ class ProteinComplex:
         # Iterate over chains, build KDTree for each chain
         kdtrees = []
 
-        CA = self.atoms["CA"]
+        CA = self.atoms["CA"]  # (l, 3)
 
         for start, end in self.chain_boundaries:
-            chain_CA = CA[start:end]
-            chain_CA = chain_CA[np.isfinite(chain_CA).all(axis=-1)]
+            chain_CA = CA[start:end]  # (chain_length, 3)
+            chain_CA = chain_CA[np.isfinite(chain_CA).all(axis=-1)]  # (n_finite_ca, 3)
             kdtrees.append(KDTree(chain_CA))
 
         return kdtrees
@@ -676,25 +681,25 @@ class ProteinComplex:
     def chain_adjacency(self, cutoff: float = 8.0) -> np.ndarray:
         # Compute adjacency matrix for protein complex
         num_chains = self.num_chains
-        adjacency = np.zeros((num_chains, num_chains), dtype=bool)
+        adjacency = np.zeros((num_chains, num_chains), dtype=bool)  # (n_chains, n_chains)
         for (i, kdtree), (j, kdtree2) in itertools.combinations(
             enumerate(self.per_chain_kd_trees), 2
         ):
             adj = kdtree.query_ball_tree(kdtree2, cutoff)
             any_is_adjacent = any(len(a) > 0 for a in adj)
-            adjacency[i, j] = any_is_adjacent
-            adjacency[j, i] = any_is_adjacent
-        return adjacency
+            adjacency[i, j] = any_is_adjacent  # scalar matrix entry
+            adjacency[j, i] = any_is_adjacent  # scalar matrix entry
+        return adjacency  # (n_chains, n_chains)
 
     def chain_adjacency_by_index(self, index: int, cutoff: float = 8.0) -> np.ndarray:
         num_chains = len(self.chain_boundaries)
-        adjacency = np.zeros(num_chains, dtype=bool)
+        adjacency = np.zeros(num_chains, dtype=bool)  # (n_chains,)
         for i, kdtree in enumerate(self.per_chain_kd_trees):
             if i == index:
                 continue
             adj = kdtree.query_ball_tree(self.per_chain_kd_trees[index], cutoff)
-            adjacency[i] = any(len(a) > 0 for a in adj)
-        return adjacency
+            adjacency[i] = any(len(a) > 0 for a in adj)  # scalar vector entry
+        return adjacency  # (n_chains,)
 
     def add_prefix_to_chain_ids(self, prefix: str) -> ProteinComplex:
         """Rename all chains in the complex with a given prefix.
@@ -715,7 +720,7 @@ class ProteinComplex:
 
     def sasa(self, by_residue: bool = True):
         chain = self.as_chain(force_conversion=True)
-        return chain.sasa(by_residue=by_residue)
+        return chain.sasa(by_residue=by_residue)  # (l,) if by_residue, otherwise (n_present_atoms,)
 
     def to_mmcif_string(self) -> str:
         """Convert the ProteinComplex to mmCIF format.
@@ -727,7 +732,7 @@ class ProteinComplex:
         # Collect all atoms from all chains
         all_atoms = []
         for chain in self.chain_iter():
-            chain_atom_array = chain.atom_array
+            chain_atom_array = chain.atom_array  # AtomArray with n_chain_atoms entries
             # Convert AtomArray to list of atoms and add to collection
             all_atoms.extend(chain_atom_array)
 
@@ -735,7 +740,7 @@ class ProteinComplex:
         if not all_atoms:
             raise ValueError("No atoms found in protein complex")
 
-        atom_array = bs.array(all_atoms)
+        atom_array = bs.array(all_atoms)  # AtomArray with total n_present_atoms entries
 
         # Create CIF file
         f = CIFFile()
@@ -786,7 +791,7 @@ class ProteinComplex:
                     data=CIFData(array=np.array(entity_descriptions), dtype=np.str_)
                 ),
             },
-        )
+        )  # each entity column has shape (n_entities,)
 
         # Create _entity_poly section
         poly_entity_ids = []
@@ -814,7 +819,7 @@ class ProteinComplex:
                     data=CIFData(array=np.array(poly_sequences), dtype=np.str_)
                 ),
             },
-        )
+        )  # each polymer column has shape (n_polymer_entities,)
 
         # Create _struct_asym section
         asym_ids = []
@@ -835,18 +840,18 @@ class ProteinComplex:
                 ),
                 "details": CIFColumn(data=CIFData(array=np.array(asym_details), dtype=np.str_)),
             },
-        )
+        )  # each asym column has shape (n_chains,)
 
     # Construction, PDB interchange, and compact storage
     @classmethod
     def from_pdb(
         cls, path: PathOrBuffer, id: str | None = None, is_predicted: bool = False
     ) -> ProteinComplex:
-        atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"])
+        atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"])  # AtomArray with n_file_atoms entries
 
         chains = []
         for chain in bs.chain_iter(atom_array):
-            chain = chain[~chain.hetero]
+            chain = chain[~chain.hetero]  # AtomArray with n_nonhetero_atoms entries
             if len(chain) == 0:
                 continue
             chains.append(ProteinChain.from_atomarray(chain, id, is_predicted))
@@ -855,8 +860,8 @@ class ProteinComplex:
     def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True):
         atom_array = None
         for chain in self.chain_iter():
-            carr = chain.atom_array if include_insertions else chain.atom_array_no_insertions
-            atom_array = carr if atom_array is None else atom_array + carr
+            carr = chain.atom_array if include_insertions else chain.atom_array_no_insertions  # AtomArray with n_chain_atoms entries
+            atom_array = carr if atom_array is None else atom_array + carr  # AtomArray containing accumulated chain atoms
         f = PDBFile()
         f.set_structure(atom_array)
         f.write(path)
@@ -909,21 +914,21 @@ class ProteinComplex:
             # Frozen dataclasses do not make their NumPy members immutable.  Work on a
             # private mask so requesting a compact backbone payload cannot clear the
             # caller's side-chain atoms in-place.
-            atom37_mask = dct["atom37_mask"].copy()
-            atom37_mask[:, 3:] = False
-            dct["atom37_mask"] = atom37_mask
-        dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]]
+            atom37_mask = dct["atom37_mask"].copy()  # (l, 37)
+            atom37_mask[:, 3:] = False  # (l, 34) mask slice for atoms beyond N/CA/C
+            dct["atom37_mask"] = atom37_mask  # (l, 37)
+        dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]]  # (n_present_atoms, 3)
         if dct.get("atom37_confidence") is not None:
-            dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]]
+            dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]]  # (n_present_atoms,)
         else:
             dct.pop("atom37_confidence", None)
         for k, v in dct.items():
             if isinstance(v, np.ndarray):
                 match v.dtype:
                     case np.int64:
-                        dct[k] = v.astype(np.int32)
+                        dct[k] = v.astype(np.int32)  # v.shape
                     case np.float64 | np.float32:
-                        dct[k] = v.astype(np.float16)
+                        dct[k] = v.astype(np.float16)  # v.shape
                     case _:
                         pass
                 if json_serializable:
@@ -948,15 +953,15 @@ class ProteinComplex:
 
         for k, v in dct.items():
             if isinstance(v, list):
-                dct[k] = np.array(v)
+                dct[k] = np.array(v)  # shape inferred from serialized nested list
 
-        atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan)
-        atom37[dct["atom37_mask"]] = dct["atom37_positions"]
-        dct["atom37_positions"] = atom37
+        atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan)  # (l, 37, 3)
+        atom37[dct["atom37_mask"]] = dct["atom37_positions"]  # (n_present_atoms, 3) selected coordinates
+        dct["atom37_positions"] = atom37  # (l, 37, 3)
         if "atom37_confidence" in dct:
-            atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32)
-            atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"]
-            dct["atom37_confidence"] = atom37_conf
+            atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32)  # (l, 37)
+            atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"]  # (n_present_atoms,) selected confidence values
+            dct["atom37_confidence"] = atom37_conf  # (l, 37)
         dct = {
             k: (
                 v.astype(np.float32)
@@ -964,7 +969,7 @@ class ProteinComplex:
                 else v
             )
             for k, v in dct.items()
-        }
+        }  # each converted array retains its serialized field shape
         if "chain_boundaries" in dct:
             del dct["chain_boundaries"]
         if "chain_boundaries" in dct["metadata"]:
@@ -1030,12 +1035,13 @@ class ProteinComplex:
 
         # TODO(roshan): Make a proper protein complex class
         def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray):
+            # arrays: (l_i, *trailing_shape); sep: (1, *trailing_shape).
             full_array = []
             for array in arrays:
                 full_array.append(array)
                 full_array.append(sep)
             full_array = full_array[:-1]
-            return np.concatenate(full_array, 0)
+            return np.concatenate(full_array, 0)  # (sum(chain_lengths) + n_chains - 1, *trailing_shape)
 
         sep_tokens = {
             "residue_index": np.array([-1]),
@@ -1043,22 +1049,22 @@ class ProteinComplex:
             "atom37_positions": np.full([1, 37, 3], np.nan),
             "atom37_mask": np.zeros([1, 37], dtype=bool),
             "confidence": np.array([0]),
-        }
+        }  # one-residue separator arrays: (1,), (1, 37, 3), or (1, 37)
 
         any_has_atom37_conf = any(c.atom37_confidence is not None for c in chains)
         if any_has_atom37_conf:
-            sep_tokens["atom37_confidence"] = np.full([1, 37], np.nan, dtype=np.float32)
+            sep_tokens["atom37_confidence"] = np.full([1, 37], np.nan, dtype=np.float32)  # (1, 37)
 
         def _get_chain_attr(chain: ProteinChain, name: str) -> np.ndarray:
-            val = getattr(chain, name)
+            val = getattr(chain, name)  # (chain_length, *field_trailing_shape) or None
             if val is None and name == "atom37_confidence":
-                return np.full([len(chain), 37], np.nan, dtype=np.float32)
-            return val
+                return np.full([len(chain), 37], np.nan, dtype=np.float32)  # (chain_length, 37)
+            return val  # (chain_length, *field_trailing_shape)
 
         array_args: dict[str, np.ndarray] = {
             name: join_arrays([_get_chain_attr(chain, name) for chain in chains], sep)
             for name, sep in sep_tokens.items()
-        }
+        }  # fields retain atom/xyz axes; first axis includes chain separators
 
         multimer_arrays = []
         chain2num_max = -1
@@ -1070,7 +1076,7 @@ class ProteinComplex:
             num_res = c.residue_index.shape[0]
             if c.chain_id not in chain2num:
                 chain2num[c.chain_id] = (chain2num_max := chain2num_max + 1)
-            chain_id_array = np.full([num_res], chain2num[c.chain_id], dtype=np.int64)
+            chain_id_array = np.full([num_res], chain2num[c.chain_id], dtype=np.int64)  # (chain_length,)
 
             if c.entity_id is None:
                 entity_num = (ent2num_max := ent2num_max + 1)
@@ -1078,9 +1084,9 @@ class ProteinComplex:
                 if c.entity_id not in ent2num:
                     ent2num[c.entity_id] = (ent2num_max := ent2num_max + 1)
                 entity_num = ent2num[c.entity_id]
-            entity_id_array = np.full([num_res], entity_num, dtype=np.int64)
+            entity_id_array = np.full([num_res], entity_num, dtype=np.int64)  # (chain_length,)
 
-            sym_id_array = np.full([num_res], i, dtype=np.int64)
+            sym_id_array = np.full([num_res], i, dtype=np.int64)  # (chain_length,)
 
             multimer_arrays.append(
                 {
@@ -1092,11 +1098,11 @@ class ProteinComplex:
 
             total_index += num_res + 1
 
-        sep = np.array([-1])
+        sep = np.array([-1])  # (1,)
         update = {
             name: join_arrays([dct[name] for dct in multimer_arrays], sep=sep)
             for name in ["chain_id", "entity_id", "sym_id"]
-        }
+        }  # each field: (sum(chain_lengths) + n_chains - 1,)
         array_args.update(update)
 
         metadata = ProteinComplexMetadata(
@@ -1153,7 +1159,7 @@ def get_assembly_fast(
     ]
     if len(structure) == 0:
         raise NoProteinError
-    unique_asym_ids = np.unique(structure.label_asym_id)  # type: ignore
+    unique_asym_ids = np.unique(structure.label_asym_id)  # type: ignore; (n_unique_asym_ids,)
     asym2chain = {}
     asym2auth = {}
     for asym_id in unique_asym_ids:
@@ -1167,7 +1173,7 @@ def get_assembly_fast(
             insertion_code,
             confidence,
             entity_id,
-        ) = chain_to_ndarray(sub_structure, mmcif, chain_id, False)
+        ) = chain_to_ndarray(sub_structure, mmcif, chain_id, False)  # array fields: (l, 37, 3), (l, 37), (l,), (l,), (l,)
 
         asym2chain[asym_id] = ProteinChain(
             id=mmcif.id or "unknown",
@@ -1217,12 +1223,13 @@ def get_assembly_fast(
 
 
 def protein_chain_to_protein_complex(chain: ProteinChain) -> ProteinComplex:
+    # chain fields share residue axis l; splitting removes chain-break separator rows.
     if "|" not in chain.sequence:
         return ProteinComplex.from_chains([chain])
-    chain_breaks = np.array(list(chain.sequence)) == "|"
-    chain_break_inds = np.where(chain_breaks)[0]
-    chain_break_inds = np.concatenate([[0], chain_break_inds, [len(chain)]])
-    chain_break_inds = np.array(list(itertools.pairwise(chain_break_inds)))
+    chain_breaks = np.array(list(chain.sequence)) == "|"  # (l,)
+    chain_break_inds = np.where(chain_breaks)[0]  # (n_chainbreaks,)
+    chain_break_inds = np.concatenate([[0], chain_break_inds, [len(chain)]])  # (n_chainbreaks + 2,)
+    chain_break_inds = np.array(list(itertools.pairwise(chain_break_inds)))  # (n_chainbreaks + 1, 2)
     complex_chains = []
     for start, end in chain_break_inds:
         if start != 0:

fastplms/models/esmfold2/esmfold2_protein_structure.py

@@ -2,12 +2,12 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
-from typing import TypeVar
-
 import numpy as np
 import torch
 import torch.nn.functional as F
+
+from collections.abc import Callable
+from typing import TypeVar
 from torch import Tensor
 from torch.amp import autocast  # type: ignore
 
@@ -15,6 +15,7 @@ from .esmfold2_affine3d import Affine3D
 from .esmfold2_misc import unbinpack
 from .esmfold2_normalize_coordinates import index_by_atom_name
 
+
 ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor)
 
 
@@ -24,7 +25,7 @@ def _coordinate_operations(
     if isinstance(coordinates, np.ndarray):
 
         def normalize(X: ArrayOrTensor) -> ArrayOrTensor:
-            return X / np.linalg.norm(X, axis=-1, keepdims=True)
+            return X / np.linalg.norm(X, axis=-1, keepdims=True)  # X.shape; normalize last axis.
 
         return normalize, np.cross
     return F.normalize, torch.cross  # type: ignore[return-value]
@@ -42,16 +43,17 @@ def infer_cbeta_from_atom37(
     dihedral in radians used by the checkpoint's training geometry.
     """
 
-    n_position = index_by_atom_name(atom37, "N", dim=-2)
-    ca_position = index_by_atom_name(atom37, "CA", dim=-2)
-    c_position = index_by_atom_name(atom37, "C", dim=-2)
+    # atom37: (..., 37, 3); ... contains optional batch and residue axes.
+    n_position = index_by_atom_name(atom37, "N", dim=-2)  # (..., 3)
+    ca_position = index_by_atom_name(atom37, "CA", dim=-2)  # (..., 3)
+    c_position = index_by_atom_name(atom37, "C", dim=-2)  # (..., 3)
     normalize, cross = _coordinate_operations(atom37)
     with np.errstate(invalid="ignore"):
-        n_to_ca = n_position - ca_position
-        n_to_c = n_position - c_position
-    unit_n_to_ca = normalize(n_to_ca)
-    normal = normalize(cross(n_to_c, unit_n_to_ca))
-    basis = [unit_n_to_ca, cross(normal, unit_n_to_ca), normal]
+        n_to_ca = n_position - ca_position  # (..., 3)
+        n_to_c = n_position - c_position  # (..., 3)
+    unit_n_to_ca = normalize(n_to_ca)  # (..., 3)
+    normal = normalize(cross(n_to_c, unit_n_to_ca))  # (..., 3)
+    basis = [unit_n_to_ca, cross(normal, unit_n_to_ca), normal]  # three arrays (..., 3)
     coefficients = [
         bond_length * np.cos(bond_angle),
         bond_length * np.sin(bond_angle) * np.cos(dihedral),
@@ -59,8 +61,8 @@ def infer_cbeta_from_atom37(
     ]
     offset = sum(
         vector * coefficient for vector, coefficient in zip(basis, coefficients, strict=True)
-    )
-    return ca_position + offset
+    )  # (..., 3)
+    return ca_position + offset  # (..., 3)
 
 
 def _unpack_alignment_inputs(
@@ -71,12 +73,13 @@ def _unpack_alignment_inputs(
 ) -> tuple[Tensor, Tensor, Tensor | None]:
     if sequence_id is None:
         return mobile, target, atom_mask
-    unpacked_mobile = unbinpack(mobile, sequence_id, pad_value=torch.nan)
-    unpacked_target = unbinpack(target, sequence_id, pad_value=torch.nan)
+    # Packed (b, l, ..., 3) coordinates become (n_sequences, max_length, ..., 3).
+    unpacked_mobile = unbinpack(mobile, sequence_id, pad_value=torch.nan)  # unpacked coordinate shape
+    unpacked_target = unbinpack(target, sequence_id, pad_value=torch.nan)  # unpacked coordinate shape
     if atom_mask is None:
-        unpacked_mask = torch.isfinite(unpacked_target).all(dim=-1)
+        unpacked_mask = torch.isfinite(unpacked_target).all(dim=-1)  # unpacked shape without xyz
     else:
-        unpacked_mask = unbinpack(atom_mask, sequence_id, pad_value=0)
+        unpacked_mask = unbinpack(atom_mask, sequence_id, pad_value=0)  # unpacked shape without xyz
     return unpacked_mobile, unpacked_target, unpacked_mask
 
 
@@ -85,13 +88,14 @@ def _flatten_atom_axes(
     target: Tensor,
     atom_mask: Tensor | None,
 ) -> tuple[Tensor, Tensor, Tensor | None]:
+    # Inputs: (b, l, n_atoms, 3) or (b, n, 3); n = l * n_atoms after flattening.
     b = mobile.shape[0]
-    flat_mobile = mobile.view(b, -1, 3) if mobile.dim() == 4 else mobile
-    flat_target = target.view(b, -1, 3) if target.dim() == 4 else target
-    flat_mask = atom_mask
+    flat_mobile = mobile.view(b, -1, 3) if mobile.dim() == 4 else mobile  # (b, n, 3)
+    flat_target = target.view(b, -1, 3) if target.dim() == 4 else target  # (b, n, 3)
+    flat_mask = atom_mask  # (b, l, n_atoms), (b, n), or None
     if flat_mask is not None and flat_mask.dim() == 3:
-        flat_mask = flat_mask.view(b, -1)
-    return flat_mobile, flat_target, flat_mask
+        flat_mask = flat_mask.view(b, -1)  # (b, n)
+    return flat_mobile, flat_target, flat_mask  # (b, n, 3), (b, n, 3), (b, n) or None
 
 
 def _masked_coordinates(
@@ -104,12 +108,12 @@ def _masked_coordinates(
             mobile.shape[:2],
             dtype=torch.bool,
             device=mobile.device,
-        )
+        )  # (b, n)
         return mobile, target, atom_mask
-    expanded_mask = atom_mask.unsqueeze(-1)
+    expanded_mask = atom_mask.unsqueeze(-1)  # (b, n, 1)
     return (
-        mobile.masked_fill(~expanded_mask, 0),
-        target.masked_fill(~expanded_mask, 0),
+        mobile.masked_fill(~expanded_mask, 0),  # (b, n, 3)
+        target.masked_fill(~expanded_mask, 0),  # (b, n, 3)
         atom_mask,
     )
 
@@ -147,18 +151,19 @@ def compute_alignment_tensors(
         atom_exists_mask,
     )
 
-    num_valid_atoms = atom_exists_mask.sum(dim=-1, keepdim=True)
-    centroid_mobile = mobile.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)
-    centroid_target = target.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)
-    centroid_mobile[num_valid_atoms == 0] = 0
-    centroid_target[num_valid_atoms == 0] = 0
-
-    expanded_mask = atom_exists_mask.unsqueeze(-1)
-    centered_mobile = (mobile - centroid_mobile).masked_fill(~expanded_mask, 0)
-    centered_target = (target - centroid_target).masked_fill(~expanded_mask, 0)
-    covariance = torch.matmul(centered_mobile.transpose(1, 2), centered_target)
-    left_vectors, _, right_vectors = torch.svd(covariance)
-    rotation = torch.matmul(left_vectors, right_vectors.transpose(1, 2))
+    # b now counts unpacked sequences if sequence_id was supplied; n counts atoms.
+    num_valid_atoms = atom_exists_mask.sum(dim=-1, keepdim=True)  # (b, 1)
+    centroid_mobile = mobile.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)  # (b, 1, 3)
+    centroid_target = target.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)  # (b, 1, 3)
+    centroid_mobile[num_valid_atoms == 0] = 0  # (n_empty, 3)
+    centroid_target[num_valid_atoms == 0] = 0  # (n_empty, 3)
+
+    expanded_mask = atom_exists_mask.unsqueeze(-1)  # (b, n, 1)
+    centered_mobile = (mobile - centroid_mobile).masked_fill(~expanded_mask, 0)  # (b, n, 3)
+    centered_target = (target - centroid_target).masked_fill(~expanded_mask, 0)  # (b, n, 3)
+    covariance = torch.matmul(centered_mobile.transpose(1, 2), centered_target)  # (b, 3, 3)
+    left_vectors, _, right_vectors = torch.svd(covariance)  # (b, 3, 3), (b, 3), (b, 3, 3)
+    rotation = torch.matmul(left_vectors, right_vectors.transpose(1, 2))  # (b, 3, 3)
     return (
         centered_mobile,
         centroid_mobile,
@@ -185,16 +190,17 @@ def compute_rmsd_no_alignment(
     """Measure RMSD after alignment using a declared reduction."""
 
     _validate_reduction(reduction, ("per_residue", "per_sample", "batch"))
-    difference = aligned - target
+    # aligned/target: (b, n, 3); num_valid_atoms: (b, 1).
+    difference = aligned - target  # (b, n, 3)
     if reduction == "per_residue":
-        mean_squared_error = difference.square().view(difference.size(0), -1, 9).mean(-1)
+        mean_squared_error = difference.square().view(difference.size(0), -1, 9).mean(-1)  # (b, n / 3)
     else:
-        mean_squared_error = difference.square().sum(dim=(1, 2)) / num_valid_atoms.squeeze(-1)
-    rmsd = torch.sqrt(mean_squared_error)
+        mean_squared_error = difference.square().sum(dim=(1, 2)) / num_valid_atoms.squeeze(-1)  # (b,)
+    rmsd = torch.sqrt(mean_squared_error)  # (b, n / 3) for per_residue, otherwise (b,)
     if reduction in {"per_residue", "per_sample"}:
         return rmsd
-    valid_samples = num_valid_atoms.squeeze(-1) > 0
-    return rmsd.masked_fill(~valid_samples, 0).sum() / (valid_samples.sum() + 1e-8)
+    valid_samples = num_valid_atoms.squeeze(-1) > 0  # (b,)
+    return rmsd.masked_fill(~valid_samples, 0).sum() / (valid_samples.sum() + 1e-8)  # ()
 
 
 @torch.no_grad()
@@ -215,19 +221,19 @@ def compute_affine_and_rmsd(
         rotation,
         num_valid_atoms,
     ) = compute_alignment_tensors(mobile, target, atom_exists_mask, sequence_id)
-    translation = torch.matmul(-centroid_mobile, rotation) + centroid_target
+    translation = torch.matmul(-centroid_mobile, rotation) + centroid_target  # (b, 1, 3)
     affine = Affine3D.from_tensor_pair(
         translation,
-        rotation.unsqueeze(dim=-3).transpose(-2, -1),
+        rotation.unsqueeze(dim=-3).transpose(-2, -1),  # (b, 1, 3, 3)
     )
-    rotated_mobile = torch.matmul(centered_mobile, rotation)
+    rotated_mobile = torch.matmul(centered_mobile, rotation)  # (b, n, 3)
     rmsd = compute_rmsd_no_alignment(
         rotated_mobile,
         centered_target,
         num_valid_atoms,
         reduction="batch",
     )
-    return affine, rmsd
+    return affine, rmsd  # affine shape: (b, 1); rmsd: ()
 
 
 def compute_gdt_ts_no_alignment(
@@ -240,12 +246,13 @@ def compute_gdt_ts_no_alignment(
 
     _validate_reduction(reduction, ("per_sample", "batch"))
     if atom_exists_mask is None:
-        atom_exists_mask = torch.isfinite(target).all(dim=-1)
-    deviation = torch.linalg.vector_norm(aligned - target, dim=-1)
-    counts = atom_exists_mask.sum(dim=-1)
-    score_1 = ((deviation < 1) * atom_exists_mask).sum(dim=-1) / counts
-    score_2 = ((deviation < 2) * atom_exists_mask).sum(dim=-1) / counts
-    score_4 = ((deviation < 4) * atom_exists_mask).sum(dim=-1) / counts
-    score_8 = ((deviation < 8) * atom_exists_mask).sum(dim=-1) / counts
-    score = (score_1 + score_2 + score_4 + score_8) * 0.25
-    return score.mean() if reduction == "batch" else score
+        atom_exists_mask = torch.isfinite(target).all(dim=-1)  # (b, n)
+    # aligned/target: (b, n, 3); atom_exists_mask: (b, n).
+    deviation = torch.linalg.vector_norm(aligned - target, dim=-1)  # (b, n)
+    counts = atom_exists_mask.sum(dim=-1)  # (b,)
+    score_1 = ((deviation < 1) * atom_exists_mask).sum(dim=-1) / counts  # (b,)
+    score_2 = ((deviation < 2) * atom_exists_mask).sum(dim=-1) / counts  # (b,)
+    score_4 = ((deviation < 4) * atom_exists_mask).sum(dim=-1) / counts  # (b,)
+    score_8 = ((deviation < 8) * atom_exists_mask).sum(dim=-1) / counts  # (b,)
+    score = (score_1 + score_2 + score_4 + score_8) * 0.25  # (b,)
+    return score.mean() if reduction == "batch" else score  # () for batch, otherwise (b,)

fastplms/models/esmfold2/esmfold2_residue_constants.py

@@ -24,12 +24,13 @@ exactly against the pinned Biohub implementation.
 from __future__ import annotations
 
 import functools
+import numpy as np
+
 from collections import defaultdict, namedtuple
 from collections.abc import Mapping
 from pathlib import Path
 from typing import Any, cast
 
-import numpy as np
 
 ca_ca = 3.80209737096
 chi_angles_atoms = {

fastplms/models/esmfold2/esmfold2_sequential_dataclass.py

@@ -2,15 +2,16 @@
 
 from __future__ import annotations
 
+import numpy as np
+
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from dataclasses import Field, dataclass, fields, replace
 from typing import Any, Self
 
-import numpy as np
-
 from .esmfold2_misc import concat_objects, slice_any_object
 
+
 Index = int | list[int] | slice | np.ndarray
 
 

fastplms/models/esmfold2/esmfold2_system.py

@@ -4,9 +4,11 @@ from __future__ import annotations
 
 import io
 import subprocess
+
 from pathlib import Path
 from typing import Any, TypeAlias
 
+
 PathLike: TypeAlias = str | Path
 PathOrBuffer: TypeAlias = PathLike | io.StringIO
 

fastplms/models/esmfold2/esmfold2_types.py

@@ -6,6 +6,7 @@ from . import esmfold2_input_builder as _input_schema
 from .esmfold2_msa import MSA
 from .esmfold2_parsing import FastaEntry
 
+
 Modification = _input_schema.Modification
 ProteinInput = _input_schema.ProteinInput
 RNAInput = _input_schema.RNAInput

fastplms/models/esmfold2/esmfold2_utils_types.py

@@ -9,9 +9,11 @@ from __future__ import annotations
 
 import io
 import os
+
 from dataclasses import dataclass
 from typing import TypeAlias
 
+
 PathLike: TypeAlias = str | os.PathLike[str]
 PathOrBuffer: TypeAlias = PathLike | io.TextIOBase
 

fastplms/models/esmfold2/modeling_esmfold2.py

@@ -2,10 +2,12 @@
 
 Quickstart::
 
-    from transformers import ESMFold2Model
+    from pathlib import Path
+    from transformers import AutoModel
 
-    model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval()
-    open("ubq.pdb", "w").write(model.infer_protein_as_pdb("MQIFVKTLTGKT..."))
+    model = AutoModel.from_pretrained("Synthyra/ESMFold2", trust_remote_code=True).cuda().eval()
+    structure = model.infer_protein_as_pdb("MQIFVKTLTGKT")
+    Path("structure.pdb").write_text(structure, encoding="utf-8")
 
 For multi-chain, ligand, and MSA inputs, use ``model.input_types`` together
 with ``model.fold(...)`` or ``model.prepare_structure_input(...)``.
@@ -17,15 +19,15 @@ import gc
 import importlib
 import importlib.metadata
 import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
 from collections.abc import Mapping
 from contextlib import contextmanager
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any, ClassVar, Literal, cast
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
 from torch import Tensor
 from tqdm.auto import tqdm
 from transformers.modeling_outputs import ModelOutput
@@ -33,6 +35,7 @@ from transformers.modeling_utils import PreTrainedModel
 
 from ...attention import get_attn_implementation, set_config_attn_implementation
 
+
 try:
     from fastplms.models.ttt import FastPLMTestTimeTrainingMixin, TTTConfig
 except ModuleNotFoundError as error:
@@ -195,6 +198,7 @@ class _ESMFold2ESMplusplusAdapter(nn.Module):
         compute_sae: bool = True,
         normalize_sae: bool = False,
     ):
+        # input_ids and optional masks/sequence IDs: (b, t).
         del return_dict, compute_sae, normalize_sae
         output = self.model(
             input_ids=input_ids,
@@ -206,14 +210,14 @@ class _ESMFold2ESMplusplusAdapter(nn.Module):
             esmfold2_hidden_states=True,
         )
         if output_hidden_states:
-            hidden_states = output.hidden_states
+            hidden_states = output.hidden_states  # Tensor (n_states, b, t, d_lm), or sequence of (b, t, d_lm) tensors
             if hidden_states is None:
                 raise RuntimeError("ESM++ did not return requested hidden states.")
             if isinstance(hidden_states, torch.Tensor):
-                output.hidden_states = hidden_states
+                output.hidden_states = hidden_states  # (n_states, b, t, d_lm)
             else:
-                output.hidden_states = torch.stack(tuple(hidden_states), dim=0)
-        return output
+                output.hidden_states = torch.stack(tuple(hidden_states), dim=0)  # (n_states, b, t, d_lm)
+        return output  # model output; hidden_states stacked on the leading state axis when requested
 
 
 def _load_fastplms_esmplusplus_for_esmfold2(
@@ -547,14 +551,15 @@ class PairTransition(nn.Module):
         self._chunk_size = chunk_size
 
     def forward(self, x: Tensor) -> Tensor:
+        # x: (b, l, ..., d_model); l_c is the current chunk width.
         if self._chunk_size is None or x.shape[1] <= self._chunk_size:
-            return self.ffn(self.norm(x))
+            return self.ffn(self.norm(x))  # x.shape
         out: list[Tensor] = []
         for s in range(0, x.shape[1], self._chunk_size):
             e = min(s + self._chunk_size, x.shape[1])
-            sl = x[:, s:e]
+            sl = x[:, s:e]  # (b, l_c, ..., d_model)
             out.append(self.ffn(self.norm(sl)))
-        return torch.cat(out, dim=1)
+        return torch.cat(out, dim=1)  # x.shape
 
 
 class ConfidenceHead(nn.Module):
@@ -569,8 +574,8 @@ class ConfidenceHead(nn.Module):
         d_pair = config.d_pair
         d_inputs = config.inputs.d_inputs
 
-        boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1)
-        self.register_buffer("boundaries", boundaries)
+        boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1)  # (distogram_bins - 1,)
+        self.register_buffer("boundaries", boundaries)  # (distogram_bins - 1,)
         self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair)
 
         self.s_norm = nn.LayerNorm(d_single)
@@ -594,7 +599,7 @@ class ConfidenceHead(nn.Module):
         max_atoms_per_token = 23
         self.plddt_weight = nn.Parameter(
             torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins)
-        )
+        )  # (23, d_single, n_plddt_bins)
 
         self.pae_ln = nn.LayerNorm(d_pair)
         self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False)
@@ -604,7 +609,7 @@ class ConfidenceHead(nn.Module):
 
         self.resolved_ln = nn.LayerNorm(d_single)
         # 2 = resolved logits ([unresolved, resolved]).
-        self.resolved_weight = nn.Parameter(torch.zeros(max_atoms_per_token, d_single, 2))
+        self.resolved_weight = nn.Parameter(torch.zeros(max_atoms_per_token, d_single, 2))  # (23, d_single, 2)
 
     def set_kernel_backend(self, backend: str | None) -> None:
         self.folding_trunk.set_kernel_backend(backend)
@@ -614,14 +619,16 @@ class ConfidenceHead(nn.Module):
 
     @staticmethod
     def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor:
-        return x if num_diffusion_samples == 1 else x.repeat_interleave(num_diffusion_samples, 0)
+        # x: (b, ...); output repeats the batch axis by samples.
+        return x if num_diffusion_samples == 1 else x.repeat_interleave(num_diffusion_samples, 0)  # (b * samples, ...), including samples = 1
 
     @staticmethod
     def _flatten_sample_axis(x: Tensor) -> Tensor:
+        # x: (b, samples, n, c) or an already flattened tensor.
         if x.ndim == 4:
             b, mult, n, c = x.shape
-            return x.reshape(b * mult, n, c)
-        return x
+            return x.reshape(b * mult, n, c)  # (b * samples, n, c) for 4D input; otherwise x.shape
+        return x  # (b * samples, n, c) for 4D input; otherwise x.shape
 
     def forward(
         self,
@@ -638,55 +645,56 @@ class ConfidenceHead(nn.Module):
         relative_position_encoding: Tensor | None = None,
         token_bonds_encoding: Tensor | None = None,
     ) -> dict[str, Tensor]:
-        s_inputs_normed = self.s_inputs_norm(s_inputs)
+        # s_inputs: (b, l, d_inputs); z: (b, l, l, d_pair); x_pred: (bs, a, 3) or (b, samples, a, 3). bs = b * samples.
+        s_inputs_normed = self.s_inputs_norm(s_inputs)  # (b, l, d_inputs)
 
-        z_base = self.z_norm(z)
+        z_base = self.z_norm(z)  # (b, l, l, d_pair)
         if relative_position_encoding is not None:
-            z_base = z_base + relative_position_encoding
+            z_base = z_base + relative_position_encoding  # (b, l, l, d_pair)
         if token_bonds_encoding is not None:
-            z_base = z_base + token_bonds_encoding
-        z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2)
-        z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1)
+            z_base = z_base + token_bonds_encoding  # (b, l, l, d_pair)
+        z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2)  # (b, l, l, d_pair)
+        z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1)  # (b, l, l, d_pair)
         z_base = z_base + self.s_to_z_prod_out(
             self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :]
             * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :]
-        )
-
-        pair = self._repeat_batch(z_base, num_diffusion_samples)
-        x_pred_flat = self._flatten_sample_axis(x_pred)
-        atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples)
-        atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples)
-        rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long()
-        mask = self._repeat_batch(token_attention_mask, num_diffusion_samples)
+        )  # (b, l, l, d_pair)
+
+        pair = self._repeat_batch(z_base, num_diffusion_samples)  # (bs, l, l, d_pair)
+        x_pred_flat = self._flatten_sample_axis(x_pred)  # (bs, a, 3)
+        atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples)  # (bs, a)
+        atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples)  # (bs, a)
+        rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long()  # (bs, l)
+        mask = self._repeat_batch(token_attention_mask, num_diffusion_samples)  # (bs, l)
         expanded_batch_size = pair.shape[0]
 
-        rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m)
+        rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m)  # (bs, l, 3)
         rep_distances = torch.cdist(
             rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist"
-        )
-        distogram_bins = (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()
-        pair = pair + self.dist_bin_pairwise_embed(distogram_bins)
+        )  # (bs, l, l)
+        distogram_bins = (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()  # (bs, l, l)
+        pair = pair + self.dist_bin_pairwise_embed(distogram_bins)  # (bs, l, l, d_pair)
 
-        pair_mask = mask[:, :, None].float() * mask[:, None, :].float()
+        pair_mask = mask[:, :, None].float() * mask[:, None, :].float()  # (bs, l, l)
 
         # FoldingTrunk handles the bf16 cast internally during inference so
         # each block's fused trimul engages. In-place residual avoids an
         # extra fp32 pair allocation.
         with torch.amp.autocast("cuda", enabled=pair.is_cuda, dtype=torch.bfloat16):
-            pair_delta = self.folding_trunk(pair, pair_attention_mask=pair_mask)
-        pair.add_(pair_delta.float())
+            pair_delta = self.folding_trunk(pair, pair_attention_mask=pair_mask)  # (bs, l, l, d_pair)
+        pair.add_(pair_delta.float())  # (bs, l, l, d_pair)
         del pair_delta
-        single = self.row_attention_pooling(pair, mask)
+        single = self.row_attention_pooling(pair, mask)  # (bs, l, d_single)
 
-        atom_mask_f = atom_mask_m.float()
-        s_at_atoms = gather_token_to_atom(single, atom_to_token_m)
-        s_at_atoms_ln = self.plddt_ln(s_at_atoms)
+        atom_mask_f = atom_mask_m.float()  # (bs, a)
+        s_at_atoms = gather_token_to_atom(single, atom_to_token_m)  # (bs, a, d_single)
+        s_at_atoms_ln = self.plddt_ln(s_at_atoms)  # (bs, a, d_single)
 
-        intra_idx = _compute_intra_token_idx(atom_to_token_m)
-        intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1)
-        w_plddt = self.plddt_weight[intra_idx]
-        plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms_ln, w_plddt)
-        plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0)
+        intra_idx = _compute_intra_token_idx(atom_to_token_m)  # (bs, a)
+        intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1)  # (bs, a)
+        w_plddt = self.plddt_weight[intra_idx]  # (bs, a, d_single, n_plddt_bins)
+        plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms_ln, w_plddt)  # (bs, a, n_plddt_bins)
+        plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0)  # (bs, a)
 
         sequence_length = single.shape[1]
         plddt_sum = torch.zeros(
@@ -694,80 +702,80 @@ class ConfidenceHead(nn.Module):
             sequence_length,
             device=single.device,
             dtype=plddt_per_atom.dtype,
-        )
+        )  # (bs, l)
         atom_count = torch.zeros(
             expanded_batch_size,
             sequence_length,
             device=single.device,
             dtype=plddt_per_atom.dtype,
-        )
-        atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype)
-        plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t)
-        atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t)
-        plddt = plddt_sum / atom_count.clamp(min=1e-6)
+        )  # (bs, l)
+        atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype)  # (bs, a)
+        plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t)  # (bs, l)
+        atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t)  # (bs, l)
+        plddt = plddt_sum / atom_count.clamp(min=1e-6)  # (bs, l)
 
         complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / (
             atom_mask_f.sum(dim=-1) + _EPS
-        )
+        )  # (bs,)
 
-        expanded_type = self._repeat_batch(mol_type, num_diffusion_samples)
-        expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples)
-        is_ligand = (expanded_type == _NONPOLYMER_ID).float()
-        inter_chain = (expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)).float()
-        near_contact = (rep_distances < 8).float()
+        expanded_type = self._repeat_batch(mol_type, num_diffusion_samples)  # (bs, l)
+        expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples)  # (bs, l)
+        is_ligand = (expanded_type == _NONPOLYMER_ID).float()  # (bs, l)
+        inter_chain = (expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)).float()  # (bs, l, l)
+        near_contact = (rep_distances < 8).float()  # (bs, l, l)
         interface_per_token = (near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)).amax(
             dim=-1
-        )
+        )  # (bs, l)
         iplddt_weight = torch.where(
             is_ligand.bool(),
             torch.full_like(interface_per_token, 2.0),
             interface_per_token,
-        )
+        )  # (bs, l)
         iplddt_weight_atoms = gather_token_to_atom(
             iplddt_weight.unsqueeze(-1), atom_to_token_m
-        ).squeeze(-1)
-        atom_iplddt_w = atom_mask_f * iplddt_weight_atoms
+        ).squeeze(-1)  # (bs, a)
+        atom_iplddt_w = atom_mask_f * iplddt_weight_atoms  # (bs, a)
         complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / (
             atom_iplddt_w.sum(dim=-1) + _EPS
-        )
+        )  # (bs,)
 
-        plddt_ca = plddt_per_atom.gather(1, rep_idx_m)
+        plddt_ca = plddt_per_atom.gather(1, rep_idx_m)  # (bs, l)
 
         # PAE
-        pae_logits = self.pae_head(self.pae_ln(pair))
-        pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach()
+        pae_logits = self.pae_head(self.pae_ln(pair))  # (bs, l, l, n_pae_bins)
+        pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach()  # (bs, l, l)
 
         # PDE
-        pde_logits = self.pde_head(self.pde_ln(pair))
-        pde = _categorical_mean(pde_logits, start=0.0, end=32.0).detach()
+        pde_logits = self.pde_head(self.pde_ln(pair))  # (bs, l, l, n_pde_bins)
+        pde = _categorical_mean(pde_logits, start=0.0, end=32.0).detach()  # (bs, l, l)
 
         # Resolved (per-atom binary).
-        s_at_atoms_res = self.resolved_ln(s_at_atoms)
-        w_res = self.resolved_weight[intra_idx]
-        resolved_logits = torch.einsum("...c,...cb->...b", s_at_atoms_res, w_res)
+        s_at_atoms_res = self.resolved_ln(s_at_atoms)  # (bs, a, d_single)
+        w_res = self.resolved_weight[intra_idx]  # (bs, a, d_single, 2)
+        resolved_logits = torch.einsum("...c,...cb->...b", s_at_atoms_res, w_res)  # (bs, a, 2)
 
         # pTM / ipTM from pae_logits.
         n_bins = pae_logits.shape[-1]
         bin_width = 32.0 / n_bins
-        bin_centers = torch.arange(0.5 * bin_width, 32.0, bin_width, device=pae_logits.device)
-        mask_f = mask.float()
-        n_residues = mask_f.sum(dim=-1, keepdim=True)
-        d0 = 1.24 * (n_residues.clamp(min=19) - 15) ** (1 / 3) - 1.8
-        tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2)
-        pae_probs = F.softmax(pae_logits, dim=-1)
-        tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1)
-
-        pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2)
-        ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (pair_mask_2d.sum(dim=-1) + _EPS)
-        ptm = ptm_per_row.max(dim=-1).values
+        bin_centers = torch.arange(0.5 * bin_width, 32.0, bin_width, device=pae_logits.device)  # (n_pae_bins,)
+        mask_f = mask.float()  # (bs, l)
+        n_residues = mask_f.sum(dim=-1, keepdim=True)  # (bs, 1)
+        d0 = 1.24 * (n_residues.clamp(min=19) - 15) ** (1 / 3) - 1.8  # (bs, 1)
+        tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2)  # (bs, n_pae_bins)
+        pae_probs = F.softmax(pae_logits, dim=-1)  # (bs, l, l, n_pae_bins)
+        tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1)  # (bs, l, l)
+
+        pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2)  # (bs, l, l)
+        ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (pair_mask_2d.sum(dim=-1) + _EPS)  # (bs, l)
+        ptm = ptm_per_row.max(dim=-1).values  # (bs,)
 
         inter_chain_mask = (
             expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
-        ).float() * pair_mask_2d
+        ).float() * pair_mask_2d  # (bs, l, l)
         iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / (
             inter_chain_mask.sum(dim=-1) + _EPS
-        )
-        iptm = iptm_per_row.max(dim=-1).values
+        )  # (bs, l)
+        iptm = iptm_per_row.max(dim=-1).values  # (bs,)
 
         max_chain_id = int(expanded_asym.max().item()) if expanded_batch_size > 0 else 0
         n_chains = max_chain_id + 1
@@ -777,16 +785,16 @@ class ConfidenceHead(nn.Module):
             n_chains,
             device=tm_expected.device,
             dtype=tm_expected.dtype,
-        )
+        )  # (bs, n_chains, n_chains)
         for c1 in range(n_chains):
-            chain_c1 = (expanded_asym == c1).float() * mask_f
+            chain_c1 = (expanded_asym == c1).float() * mask_f  # (bs, l)
             if chain_c1.sum() == 0:
                 continue
             for c2 in range(n_chains):
-                chain_c2 = (expanded_asym == c2).float() * mask_f
-                pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2)
-                denom = pair_m.sum(dim=(-1, -2)) + _EPS
-                pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(dim=(-1, -2)) / denom
+                chain_c2 = (expanded_asym == c2).float() * mask_f  # (bs, l)
+                pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2)  # (bs, l, l)
+                denom = pair_m.sum(dim=(-1, -2)) + _EPS  # (bs,)
+                pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(dim=(-1, -2)) / denom  # (bs,)
 
         return {
             "plddt_logits": plddt_logits,
@@ -803,7 +811,7 @@ class ConfidenceHead(nn.Module):
             "ptm": ptm.detach(),
             "iptm": iptm.detach(),
             "pair_chains_iptm": pair_chains_iptm.detach(),
-        }
+        }  # mapping of confidence tensors with shapes traced above
 
 
 def _inverse_softplus(value: float) -> float:
@@ -834,9 +842,9 @@ def _convert_esmc_attention_outputs_to_te(module: nn.Module) -> tuple[str, ...]:
                     device=child.weight.device,
                 )
                 with torch.no_grad():
-                    replacement.weight.copy_(child.weight)
+                    replacement.weight.copy_(child.weight)  # child.weight.shape
                     if child.bias is not None:
-                        replacement.bias.copy_(child.bias)
+                        replacement.bias.copy_(child.bias)  # child.bias.shape
                 replacement.eval().requires_grad_(False)
                 setattr(owner, name, replacement)
                 converted.append(path)
@@ -953,15 +961,15 @@ class ESMFold2Model(
             self.lm_encoder = None
 
         self.parcae_input_norm = nn.LayerNorm(d_pair)
-        self.parcae_log_a = nn.Parameter(torch.zeros(d_pair))
+        self.parcae_log_a = nn.Parameter(torch.zeros(d_pair))  # (d_pair,)
         parcae_decay_init = math.sqrt(1.0 / 5.0)
         parcae_delta_init = -math.log(parcae_decay_init)
         self.parcae_log_delta = nn.Parameter(
             torch.full((d_pair,), _inverse_softplus(parcae_delta_init), dtype=torch.float32)
-        )
-        self.parcae_b_cont = nn.Parameter(torch.eye(d_pair))
+        )  # (d_pair,)
+        self.parcae_b_cont = nn.Parameter(torch.eye(d_pair))  # (d_pair, d_pair)
         self.parcae_readout = nn.Linear(d_pair, d_pair, bias=False)
-        nn.init.eye_(self.parcae_readout.weight)
+        nn.init.eye_(self.parcae_readout.weight)  # (d_pair, d_pair)
         self.parcae_coda = FoldingTrunk(
             n_layers=config.parcae.coda_n_layers, d_pair=d_pair, expansion_ratio=4
         )
@@ -1101,9 +1109,10 @@ class ESMFold2Model(
         input_ids: torch.Tensor | None = None,
         **kwargs,
     ) -> torch.Tensor:
+        # Encoded batch: b sequences, padded token width t including BOS/EOS.
         del kwargs
         if input_ids is not None:
-            return input_ids
+            return input_ids  # (b, t), or caller input_ids.shape
         if seq is None:
             raise ValueError("Pass either seq or input_ids for ESMFold2 TTT.")
         sequences = [seq] if isinstance(seq, str) else seq
@@ -1122,13 +1131,13 @@ class ESMFold2Model(
             (len(encoded), max_len),
             SEQUENCE_PAD_TOKEN,
             dtype=torch.long,
-        )
+        )  # (b, t)
         for row, token_ids in enumerate(encoded):
             input_tensor[row, : len(token_ids)] = torch.tensor(
                 token_ids,
                 dtype=torch.long,
-            )
-        return input_tensor
+            )  # (t_i,)
+        return input_tensor  # (b, t), or caller input_ids.shape
 
     def _ttt_mask_token(self) -> int:
         return SEQUENCE_MASK_TOKEN
@@ -1142,18 +1151,20 @@ class ESMFold2Model(
             SEQUENCE_STANDARD_AA_MAX_TOKEN,
             device=input_ids.device,
             dtype=input_ids.dtype,
-        )
+        )  # (SEQUENCE_STANDARD_AA_MAX_TOKEN - SEQUENCE_STANDARD_AA_MIN_TOKEN,)
 
     def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
+        # input_ids: arbitrary token-ID shape.
         return (input_ids >= SEQUENCE_STANDARD_AA_MIN_TOKEN) & (
             input_ids < SEQUENCE_STANDARD_AA_MAX_TOKEN
-        )
+        )  # input_ids.shape
 
     def _ttt_predict_logits(
         self,
         batch: torch.Tensor | dict[str, torch.Tensor],
         **kwargs,
     ) -> torch.Tensor:
+        # batch: (b, t) token IDs; backbone output last_hidden_state: (b, t, d_lm).
         del kwargs
         if not isinstance(batch, torch.Tensor):
             raise TypeError("ESMFold2 TTT expects input_ids tensors.")
@@ -1163,14 +1174,14 @@ class ESMFold2Model(
         self._ensure_ttt_lm_head()
         if self._ttt_lm_head is None:
             raise RuntimeError("ESMFold2 TTT MLM head initialization failed.")
-        attention_mask = batch.ne(SEQUENCE_PAD_TOKEN)
+        attention_mask = batch.ne(SEQUENCE_PAD_TOKEN)  # (b, t)
         output = self._esmc(
             input_ids=batch,
             attention_mask=attention_mask,
             return_dict=True,
             compute_sae=False,
         )
-        return self._ttt_lm_head(output.last_hidden_state)
+        return self._ttt_lm_head(output.last_hidden_state)  # (b, t, vocab_size)
 
     @classmethod
     def from_pretrained(
@@ -1287,6 +1298,7 @@ class ESMFold2Model(
         lm_mask_pct: float = 0.0,
         verbose: bool = False,
     ) -> Tensor:
+        # Input tensors: (b, l); n_states and d_lm come from the loaded backbone.
         if self._esmc_fp8 and torch.is_grad_enabled():
             _reload_esmc_bf16_for_gradients(
                 self,
@@ -1312,9 +1324,9 @@ class ESMFold2Model(
                         pad_to_multiple=pad_to,
                         lm_mask_pct=lm_mask_pct,
                         mask_token_id=SEQUENCE_MASK_TOKEN,
-                    )
+                    )  # (b, l, n_states, d_lm)
                     progress.update()
-                return result
+                return result  # (b, l, n_states, d_lm)
             return compute_lm_hidden_states(
                 self._esmc,
                 input_ids,
@@ -1325,19 +1337,20 @@ class ESMFold2Model(
                 pad_to_multiple=pad_to,
                 lm_mask_pct=lm_mask_pct,
                 mask_token_id=SEQUENCE_MASK_TOKEN,
-            )
+            )  # (b, l, n_states, d_lm)
 
     def _discretized_dynamics(self) -> tuple[Tensor, Tensor]:
-        delta = F.softplus(self.parcae_log_delta)
-        a = torch.exp(-delta * torch.exp(self.parcae_log_a))
-        b = delta[:, None] * self.parcae_b_cont
-        return a, b
+        delta = F.softplus(self.parcae_log_delta)  # (d_pair,)
+        a = torch.exp(-delta * torch.exp(self.parcae_log_a))  # (d_pair,)
+        b = delta[:, None] * self.parcae_b_cont  # (d_pair, d_pair)
+        return a, b  # (d_pair,), (d_pair, d_pair)
 
     def _init_pair_state(self, ref: Tensor) -> Tensor:
+        # ref: (b, l, l, d_pair).
         std = math.sqrt(2.0 / (5.0 * ref.shape[-1]))
-        state = torch.empty_like(ref, dtype=torch.float32)
-        nn.init.trunc_normal_(state, mean=0.0, std=std, a=-3 * std, b=3 * std)
-        return state.to(dtype=ref.dtype)
+        state = torch.empty_like(ref, dtype=torch.float32)  # ref.shape
+        nn.init.trunc_normal_(state, mean=0.0, std=std, a=-3 * std, b=3 * std)  # ref.shape
+        return state.to(dtype=ref.dtype)  # ref.shape
 
     def _run_one_loop(
         self,
@@ -1356,6 +1369,7 @@ class ESMFold2Model(
         # otherwise leaks about 2 GB of l^2 * c_z data into distogram/sample scope.
         # training=True forces dropout under eval(), matching the per-loop
         # dropout strategy used at train time.
+        # Pair states: (b, l, l, d_pair); pair_mask: (b, l, l); tok_mask: (b, l). MSA depth m may be subsampled.
         lm_cfg = self.config.lm_encoder
         _per_loop_lm_dropout = (
             lm_z is not None
@@ -1377,19 +1391,19 @@ class ESMFold2Model(
             if _per_loop_lm_dropout:
                 if lm_z is None:
                     raise RuntimeError("Per-loop LM dropout requires LM pair features.")
-                lm_z_i: Tensor | None = F.dropout(lm_z, p=_lm_dropout_p, training=True)
+                lm_z_i: Tensor | None = F.dropout(lm_z, p=_lm_dropout_p, training=True)  # (b, l, l, d_pair) or None
             else:
-                lm_z_i = lm_z
+                lm_z_i = lm_z  # (b, l, l, d_pair) or None
 
-            refined_lm_z: Tensor | None = None
+            refined_lm_z: Tensor | None = None  # (b, l, l, d_pair) or None
             if lm_z_i is not None and self.lm_encoder is not None:
                 refined_lm_z = self.lm_encoder(
                     lm_z_i.to(z_init.dtype), pair_attention_mask=pair_mask
-                )
+                )  # (b, l, l, d_pair) or None
 
-            z_inject_pair = z_init
+            z_inject_pair = z_init  # (b, l, l, d_pair)
             if lm_z_i is not None and self.lm_encoder is None:
-                z_inject_pair = z_inject_pair + lm_z_i.to(z_inject_pair.dtype)
+                z_inject_pair = z_inject_pair + lm_z_i.to(z_inject_pair.dtype)  # (b, l, l, d_pair)
 
             if self.msa_encoder is not None and _msa_inputs is not None:
                 msa_i, mask_i, hd_i, dv_i = maybe_subsample_msa(
@@ -1399,26 +1413,26 @@ class ESMFold2Model(
                     _msa_inputs["deletion_value"],
                     max_depth=_msa_inputs["max_depth"],
                     enabled=_msa_inputs["subsample_enabled"],
-                )
+                )  # each (b, m, l); masks/deletion tensors may be None
                 b_msa, m, l_msa = msa_i.shape
-                msa_oh = F.one_hot(msa_i.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES).float()
+                msa_oh = F.one_hot(msa_i.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES).float()  # (b, l, m, 33)
                 msa_attn = (
                     mask_i.permute(0, 2, 1).float()
                     if mask_i is not None
                     else tok_mask[:, :, None].expand(-1, -1, m).float()
-                )
+                )  # (b, l, m)
                 # Bias-free MSAEncoder.embed requires zeroed padding.
-                msa_oh = msa_oh * msa_attn.unsqueeze(-1)
+                msa_oh = msa_oh * msa_attn.unsqueeze(-1)  # (b, l, m, 33)
                 hd = (
                     hd_i.permute(0, 2, 1).float()
                     if hd_i is not None
                     else torch.zeros(b_msa, l_msa, m, device=msa_i.device)
-                )
+                )  # (b, l, m)
                 dv = (
                     dv_i.permute(0, 2, 1).float()
                     if dv_i is not None
                     else torch.zeros(b_msa, l_msa, m, device=msa_i.device)
-                )
+                )  # (b, l, m)
                 msa_pair = self.msa_encoder(
                     x_pair=z_inject_pair,
                     x_inputs=_msa_inputs["x_inputs"],
@@ -1426,19 +1440,19 @@ class ESMFold2Model(
                     has_deletion=hd,
                     deletion_value=dv,
                     msa_attention_mask=msa_attn,
-                ).to(z_inject_pair.dtype)
+                ).to(z_inject_pair.dtype)  # (b, l, l, d_pair)
                 z_inject_pair = (
                     msa_pair if self.config.msa_encoder_overwrite else (z_inject_pair + msa_pair)
-                )
+                )  # (b, l, l, d_pair)
 
             if refined_lm_z is not None:
-                z_inject_pair = z_inject_pair + refined_lm_z.to(z_inject_pair.dtype)
+                z_inject_pair = z_inject_pair + refined_lm_z.to(z_inject_pair.dtype)  # (b, l, l, d_pair)
 
-            injected_pair = self.parcae_input_norm(z_inject_pair)
-            z = a * z + F.linear(injected_pair.to(z.dtype), b_mat)
-            z = self.folding_trunk(z, pair_attention_mask=pair_mask)
+            injected_pair = self.parcae_input_norm(z_inject_pair)  # (b, l, l, d_pair)
+            z = a * z + F.linear(injected_pair.to(z.dtype), b_mat)  # (b, l, l, d_pair)
+            z = self.folding_trunk(z, pair_attention_mask=pair_mask)  # (b, l, l, d_pair)
 
-        return z
+        return z  # (b, l, l, d_pair)
 
     def forward(
         self,
@@ -1488,6 +1502,7 @@ class ESMFold2Model(
         disto_cond_mask: Tensor | None = None,
         verbose: bool = False,
     ) -> ESMFold2Output | tuple[Any, ...]:
+        # Token IDs/masks: (b, l); atom IDs/masks: (b, a); ref_pos: (b, a, 3); chars: (b, a, 4); MSA: (b, m, l); bs = b * samples.
         output_hidden_states, return_dict = _resolve_structure_output_controls(
             self.config,
             output_attentions=output_attentions,
@@ -1508,9 +1523,9 @@ class ESMFold2Model(
             disto_cond_mask=disto_cond_mask,
         )
         del gt_coords, is_resolved, frames_idx
-        tok_mask = token_attention_mask
-        atm_mask = atom_attention_mask
-        disto_idx = distogram_atom_idx
+        tok_mask = token_attention_mask  # (b, l)
+        atm_mask = atom_attention_mask  # (b, a)
+        disto_idx = distogram_atom_idx  # (b, l)
 
         n_loops: int = num_loops if num_loops is not None else self.config.num_loops
         n_samples: int = (
@@ -1521,37 +1536,37 @@ class ESMFold2Model(
         total_steps = max(1, n_loops + 1)
 
         if res_type.dim() == 2:
-            res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float()
-            res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float()
+            res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float()  # (b, l, 33)
+            res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float()  # (b, l, 33)
         else:
-            res_type_oh = res_type.float()
+            res_type_oh = res_type.float()  # (b, l, 33)
 
         if msa is not None:
-            msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float()
+            msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float()  # (b, m, l, 33)
             if msa_attention_mask is not None:
-                mask_f = msa_attention_mask.float().unsqueeze(-1)
-                msa_oh_profile = msa_oh_profile * mask_f
-                valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1)
-                profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1)
+                mask_f = msa_attention_mask.float().unsqueeze(-1)  # (b, m, l, 1)
+                msa_oh_profile = msa_oh_profile * mask_f  # (b, m, l, 33)
+                valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1)  # (b, l)
+                profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1)  # (b, l, 33)
             else:
-                profile = msa_oh_profile.mean(dim=1)
+                profile = msa_oh_profile.mean(dim=1)  # (b, l, 33)
         else:
-            profile = res_type_oh
+            profile = res_type_oh  # (b, l, 33)
 
         if deletion_mean is None:
             deletion_mean = torch.zeros(
                 res_type.shape[0], res_type.shape[1], device=res_type.device
-            )
+            )  # (b, l)
 
-        ref_element_oh = F.one_hot(ref_element.long(), num_classes=MAX_ATOMIC_NUMBER).float()
+        ref_element_oh = F.one_hot(ref_element.long(), num_classes=MAX_ATOMIC_NUMBER).float()  # (b, a, 128)
         ref_atom_name_chars_oh = F.one_hot(
             ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE
-        ).float()
+        ).float()  # (b, a, 4, 64)
         # Bias-free downstream Linears require zeroed padding.
-        atm_mask_f = atm_mask.float()
-        ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1)
-        ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(-1).unsqueeze(-1)
-        atom_to_token = atom_to_token * atm_mask.long()
+        atm_mask_f = atm_mask.float()  # (b, a)
+        ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1)  # (b, a, 128)
+        ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(-1).unsqueeze(-1)  # (b, a, 4, 64)
+        atom_to_token = atom_to_token * atm_mask.long()  # (b, a)
 
         use_amp = ref_pos.device.type == "cuda"
         with torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16):
@@ -1566,9 +1581,9 @@ class ESMFold2Model(
                 ref_element=ref_element_oh,
                 ref_atom_name_chars=ref_atom_name_chars_oh,
                 atom_to_token=atom_to_token,
-            )
+            )  # (b, l, d_inputs)
 
-            z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(x_inputs).unsqueeze(1)
+            z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(x_inputs).unsqueeze(1)  # (b, l, l, d_pair)
 
             relative_position_encoding = self.rel_pos(
                 residue_index=residue_index,
@@ -1576,9 +1591,9 @@ class ESMFold2Model(
                 sym_id=sym_id,
                 entity_id=entity_id,
                 token_index=token_index,
-            )
-            token_bonds_encoding = self.token_bonds(token_bonds.float())
-            z_init = z_init + relative_position_encoding + token_bonds_encoding
+            )  # (b, l, l, d_pair)
+            token_bonds_encoding = self.token_bonds(token_bonds.float())  # (b, l, l, d_pair)
+            z_init = z_init + relative_position_encoding + token_bonds_encoding  # (b, l, l, d_pair)
 
             if lm_hidden_states is None and input_ids is not None and self._esmc is not None:
                 lm_hidden_states = self._compute_lm_hidden_states(
@@ -1589,26 +1604,26 @@ class ESMFold2Model(
                     tok_mask,
                     lm_mask_pct=(self.config.lm_mask_pct if lm_mask_pct is None else lm_mask_pct),
                     verbose=verbose,
-                )
-            lm_z: Tensor | None = None
+                )  # (b, l, n_states, d_lm)
+            lm_z: Tensor | None = None  # (b, l, l, d_pair) or None
             if lm_hidden_states is not None:
-                lm_z = self.language_model(lm_hidden_states.detach())
+                lm_z = self.language_model(lm_hidden_states.detach())  # (b, l, l, d_pair) or None
             del lm_hidden_states
 
-            pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float()
+            pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float()  # (b, l, l)
 
-            z = self._init_pair_state(z_init)
+            z = self._init_pair_state(z_init)  # (b, l, l, d_pair)
 
-            a, b = self._discretized_dynamics()
-            a = a.view(1, 1, 1, -1).to(device=z.device, dtype=z.dtype)
-            b_mat = b.to(device=z.device, dtype=z.dtype)
+            a, b = self._discretized_dynamics()  # (d_pair,), (d_pair, d_pair)
+            a = a.view(1, 1, 1, -1).to(device=z.device, dtype=z.dtype)  # (1, 1, 1, d_pair)
+            b_mat = b.to(device=z.device, dtype=z.dtype)  # (d_pair, d_pair)
 
             _msa_inputs: dict | None = None
             if self.msa_encoder is not None and msa is not None:
                 msa_attention_mask = maybe_apply_msa_column_masking(
                     msa_attention_mask,
                     msa_column_mask_rate,
-                )
+                )  # (b, m, l)
                 _msa_inputs = dict(
                     x_inputs=x_inputs,
                     msa=msa,
@@ -1631,14 +1646,14 @@ class ESMFold2Model(
                 tok_mask=tok_mask,
                 total_steps=total_steps,
                 verbose=verbose,
-            )
+            )  # (b, l, l, d_pair)
             del z_init, lm_z, _msa_inputs, a, b_mat
 
-            z = self.parcae_readout(z)
-            z = self.parcae_coda(z, pair_attention_mask=pair_mask)
+            z = self.parcae_readout(z)  # (b, l, l, d_pair)
+            z = self.parcae_coda(z, pair_attention_mask=pair_mask)  # (b, l, l, d_pair)
 
-            z = z.float()
-        distogram_logits = self.distogram_head(z + z.transpose(-2, -3))
+            z = z.float()  # (b, l, l, d_pair)
+        distogram_logits = self.distogram_head(z + z.transpose(-2, -3))  # (b, l, l, n_distogram_bins)
 
         structure_output = self.structure_head.sample(
             z_trunk=z,
@@ -1666,13 +1681,13 @@ class ESMFold2Model(
             return_atom_repr=False,
             denoising_early_exit_rmsd=(0.10 if early_exit else None),
             verbose=verbose,
-        )
+        )  # tensor mapping follows the called head's shape contract
 
-        sample_coords = structure_output["sample_atom_coords"]
+        sample_coords = structure_output["sample_atom_coords"]  # (bs, a, 3), or explicit (b, samples, a, 3)
         if sample_coords is None:
             raise RuntimeError("ESMFold2 structure sampling did not return coordinates.")
         output: dict[str, Tensor] = {"distogram_logits": distogram_logits}
-        output["sample_atom_coords"] = sample_coords
+        output["sample_atom_coords"] = sample_coords  # sample_coords.shape
 
         confidence_config = self.config.confidence_head
         confidence_enabled = (
@@ -1697,7 +1712,7 @@ class ESMFold2Model(
                         num_diffusion_samples=n_samples,
                         relative_position_encoding=relative_position_encoding.detach(),
                         token_bonds_encoding=token_bonds_encoding.detach(),
-                    )
+                    )  # tensor mapping follows the called head's shape contract
                     progress.update()
             else:
                 confidence_output = self.confidence_head(
@@ -1713,18 +1728,18 @@ class ESMFold2Model(
                     num_diffusion_samples=n_samples,
                     relative_position_encoding=relative_position_encoding.detach(),
                     token_bonds_encoding=token_bonds_encoding.detach(),
-                )
+                )  # tensor mapping follows the called head's shape contract
         output.update(confidence_output)
-        output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask
-        output["residue_index"] = residue_index
-        output["entity_id"] = entity_id
+        output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask  # (b, a)
+        output["residue_index"] = residue_index  # (b, l)
+        output["entity_id"] = entity_id  # (b, l)
         return _finalize_structure_output(
             output,
             token_input_state=x_inputs,
             pair_state=z,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
-        )
+        )  # ESMFold2Output/tuple retaining the traced tensor shapes
 
     @torch.no_grad()
     def infer_protein(self, seq: str, **forward_kwargs) -> ESMFold2Output:
@@ -1743,7 +1758,7 @@ class ESMFold2Model(
         if not self.config.msa_conditioning:
             for name in MSA_CONDITIONING_INPUT_NAMES:
                 features.pop(name, None)
-        features = {k: v.to(self.device) for k, v in features.items()}
+        features = {k: v.to(self.device) for k, v in features.items()}  # every feature retains its shape
         return self(**features, **forward_kwargs, return_dict=True)
 
     @property
@@ -2056,14 +2071,15 @@ class MSAEncoderBlock(nn.Module):
         msa_attention_mask: Tensor,
         pair_attention_mask: Tensor,
     ) -> tuple[Tensor, Tensor]:
-        pair = pair + self.outer_product_mean(m, msa_attention_mask)
+        # m: (b, l, m_depth, d_msa); pair: (b, l, l, d_pair); corresponding masks omit the feature axis.
+        pair = pair + self.outer_product_mean(m, msa_attention_mask)  # (b, l, l, d_pair)
         if not self.is_final_block:
-            m = m + self.msa_pair_weighted_averaging(m, pair, pair_attention_mask)
-            m = m + self.msa_transition(m)
-        pair = pair + self.tri_mul_out(pair, mask=pair_attention_mask)
-        pair = pair + self.tri_mul_in(pair, mask=pair_attention_mask)
-        pair = pair + self.pair_transition(pair)
-        return m, pair
+            m = m + self.msa_pair_weighted_averaging(m, pair, pair_attention_mask)  # (b, l, m_depth, d_msa)
+            m = m + self.msa_transition(m)  # (b, l, m_depth, d_msa)
+        pair = pair + self.tri_mul_out(pair, mask=pair_attention_mask)  # (b, l, l, d_pair)
+        pair = pair + self.tri_mul_in(pair, mask=pair_attention_mask)  # (b, l, l, d_pair)
+        pair = pair + self.pair_transition(pair)  # (b, l, l, d_pair)
+        return m, pair  # (b, l, m_depth, d_msa), (b, l, l, d_pair)
 
 
 class MSAEncoder(nn.Module):
@@ -2110,12 +2126,13 @@ class MSAEncoder(nn.Module):
         msa_attention_mask: Tensor,
     ) -> Tensor:
         # Every input tensor is pre-transposed to shape (b, l, m, ...) before this call.
+        # x_pair: (b, l, l, d_pair); x_inputs: (b, l, d_inputs); MSA features: (b, l, m, 33), deletion/mask: (b, l, m).
         m_feat = torch.cat(
             [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)], dim=-1
-        )
-        m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2)
-        tok_mask = msa_attention_mask[:, :, 0].bool()
-        pair_attention_mask = tok_mask.unsqueeze(2) & tok_mask.unsqueeze(1)
+        )  # (b, l, m, 35)
+        m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2)  # (b, l, m, d_msa)
+        tok_mask = msa_attention_mask[:, :, 0].bool()  # (b, l)
+        pair_attention_mask = tok_mask.unsqueeze(2) & tok_mask.unsqueeze(1)  # (b, l, l)
         for block in self.blocks:
-            m, x_pair = block(m, x_pair, msa_attention_mask, pair_attention_mask)
-        return x_pair
+            m, x_pair = block(m, x_pair, msa_attention_mask, pair_attention_mask)  # (b, l, m, d_msa), (b, l, l, d_pair)
+        return x_pair  # (b, l, l, d_pair)

fastplms/models/esmfold2/modeling_esmfold2_classification.py

@@ -2,10 +2,10 @@
 
 from __future__ import annotations
 
-from typing import Any, Literal
-
 import torch
 import torch.nn as nn
+
+from typing import Any, Literal
 from torch import Tensor
 
 from ..classification_probe import SequenceClassificationProbe, TokenClassificationProbe
@@ -94,46 +94,47 @@ class _ESMFold2ClassificationMixin:
         if not sequences:
             raise ValueError("prepare_classifier_inputs requires at least one sequence.")
         encoded = [_encode_single_chain(sequence) for sequence in sequences]
-        sequence_length = max(map(len, encoded))
+        sequence_length = max(map(len, encoded))  # l; b = len(encoded).
         input_ids = torch.full(
             (len(encoded), sequence_length),
             SEQUENCE_PAD_TOKEN,
             dtype=torch.long,
             device=self.device,
-        )
-        attention_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+        )  # (b, l)
+        attention_mask = torch.zeros_like(input_ids, dtype=torch.bool)  # (b, l)
         for batch_index, token_ids in enumerate(encoded):
-            length = len(token_ids)
+            length = len(token_ids)  # l_i
             input_ids[batch_index, :length] = torch.tensor(
                 token_ids, dtype=torch.long, device=self.device
-            )
-            attention_mask[batch_index, :length] = True
-        return {"input_ids": input_ids, "attention_mask": attention_mask}
+            )  # (l_i,)
+            attention_mask[batch_index, :length] = True  # (l_i,)
+        return {"input_ids": input_ids, "attention_mask": attention_mask}  # both (b, l)
 
     def _classifier_embeddings(
         self, input_ids: Tensor, attention_mask: Tensor | None
     ) -> tuple[Tensor, Tensor]:
+        # input_ids: (b, l); attention_mask: (b, l) or None.
         if input_ids.ndim != 2:
             raise ValueError(
                 "ESMFold2 classifier input_ids must have shape (batch, residue), "
                 f"got {tuple(input_ids.shape)}."
             )
         if attention_mask is None:
-            attention_mask = input_ids.ne(SEQUENCE_PAD_TOKEN)
+            attention_mask = input_ids.ne(SEQUENCE_PAD_TOKEN)  # (b, l)
         elif attention_mask.shape != input_ids.shape:
             raise ValueError(
                 "ESMFold2 classifier attention_mask must match input_ids, got "
                 f"{tuple(attention_mask.shape)} and {tuple(input_ids.shape)}."
             )
-        residue_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
+        residue_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)  # (b, l)
         if not residue_mask.any(dim=1).all():
             raise ValueError("Every ESMFold2 classifier input must contain a protein residue.")
         if input_ids.masked_select(residue_mask).eq(SEQUENCE_PAD_TOKEN).any():
             raise ValueError("ESMFold2 classifier padding tokens cannot be attended residues.")
-        residue_ids = input_ids.masked_select(residue_mask)
+        residue_ids = input_ids.masked_select(residue_mask)  # (n_present,)
         valid_residue_ids = torch.tensor(
             sorted(_VALID_RESIDUE_IDS), dtype=input_ids.dtype, device=input_ids.device
-        )
+        )  # (n_valid_ids,)
         if not torch.isin(residue_ids, valid_residue_ids).all():
             raise ValueError(
                 "ESMFold2 classifiers accept residue-only single-chain protein inputs."
@@ -142,9 +143,9 @@ class _ESMFold2ClassificationMixin:
         batch_size, sequence_length = input_ids.shape
         residue_index = torch.arange(sequence_length, device=input_ids.device).expand(
             batch_size, -1
-        )
-        asym_id = torch.zeros_like(input_ids)
-        mol_type = torch.zeros_like(input_ids)
+        )  # (b, l)
+        asym_id = torch.zeros_like(input_ids)  # (b, l)
+        mol_type = torch.zeros_like(input_ids)  # (b, l)
         with torch.no_grad():
             hidden_states = self._compute_lm_hidden_states(
                 input_ids,
@@ -152,9 +153,9 @@ class _ESMFold2ClassificationMixin:
                 residue_index,
                 mol_type,
                 residue_mask,
-            )
-        embeddings = self.project_esmc_hidden_states(hidden_states, residue_mask)
-        return embeddings, residue_mask
+            )  # (b, l, n_states, d_lm)
+        embeddings = self.project_esmc_hidden_states(hidden_states, residue_mask)  # (b, l, d_pair)
+        return embeddings, residue_mask  # (b, l, d_pair), (b, l)
 
     def _classifier_forward(
         self,
@@ -165,7 +166,7 @@ class _ESMFold2ClassificationMixin:
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
     ):
-        embeddings, residue_mask = self._classifier_embeddings(input_ids, attention_mask)
+        embeddings, residue_mask = self._classifier_embeddings(input_ids, attention_mask)  # (b, l, d_pair), (b, l)
         return self.classifier(
             embeddings,
             attention_mask=residue_mask,

fastplms/models/esmfold2/modeling_esmfold2_common.py

@@ -10,13 +10,13 @@
 from __future__ import annotations
 
 import importlib
-from functools import partial
-from importlib.util import find_spec
-from typing import Any, ClassVar, cast
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+
+from functools import partial
+from importlib.util import find_spec
+from typing import Any, ClassVar, cast
 from torch import Tensor
 from torch.utils.checkpoint import checkpoint
 from tqdm.auto import tqdm
@@ -24,6 +24,7 @@ from tqdm.auto import tqdm
 from .configuration_esmfold2 import ESMFold2Config
 from .reproducibility import seed_context
 
+
 _seed_context = seed_context
 
 try:
@@ -197,15 +198,16 @@ class DropoutResidual(nn.Module):
             self._impl = nn.Dropout(r)
 
     def forward(self, residual: Tensor, delta: Tensor) -> Tensor:
+        # residual, delta: same shape; dropout shares delta axis self._batch_dim.
         if self._use_fused_kernels:
-            return self._impl(residual, delta)
-        # The unfused path broadcasts a row/column-shared mask M with shape (1, ...).
+            return self._impl(residual, delta)  # delta.shape
+        # The unfused mask shares the selected row/column axis and retains the other dimensions.
         if self._r == 0.0 or not self.training:
-            return residual + delta
+            return residual + delta  # delta.shape
         shape = list(delta.shape)
         shape[self._batch_dim] = 1
-        mask = self._impl(delta.new_ones(shape))
-        return residual + delta * mask
+        mask = self._impl(delta.new_ones(shape))  # delta.shape with shared axis set to 1
+        return residual + delta * mask  # delta.shape
 
 
 # ---------------------------------------------------------------------------
@@ -217,7 +219,7 @@ XYZ_DIMS: int = 3
 MAX_ATOMIC_NUMBER: int = 128
 
 # Input feature dim = 3 + 1 + 1 + 128 + 64*4 = 389
-ATOM_FEATURE_DIM: int = XYZ_DIMS + 1 + 1 + MAX_ATOMIC_NUMBER + CHAR_VOCAB_SIZE * MAX_CHARS
+ATOM_FEATURE_DIM: int = XYZ_DIMS + 1 + 1 + MAX_ATOMIC_NUMBER + CHAR_VOCAB_SIZE * MAX_CHARS  # d_atom_features
 
 
 NUM_RES_TYPES: int = 33
@@ -246,39 +248,41 @@ def maybe_subsample_msa(
     max_depth: int | None,
     enabled: bool,
 ) -> tuple[Tensor, Tensor | None, Tensor | None, Tensor | None]:
+    # MSA tensors: (b, m, l); k = max_depth for subsampled rows.
     if not enabled or max_depth is None:
-        return msa, msa_attention_mask, has_deletion, deletion_value
+        return msa, msa_attention_mask, has_deletion, deletion_value  # MSA tensors retain (b, selected_depth, l); optional tensors remain None
 
     depth = msa.size(1)
     if depth <= 1 or depth <= max_depth:
-        return msa, msa_attention_mask, has_deletion, deletion_value
+        return msa, msa_attention_mask, has_deletion, deletion_value  # MSA tensors retain (b, selected_depth, l); optional tensors remain None
 
-    indices = torch.zeros(max_depth, dtype=torch.long, device=msa.device)
-    indices[1:] = torch.randperm(depth - 1, device=msa.device)[: max_depth - 1] + 1
-    indices = indices.sort().values
+    indices = torch.zeros(max_depth, dtype=torch.long, device=msa.device)  # (k,)
+    indices[1:] = torch.randperm(depth - 1, device=msa.device)[: max_depth - 1] + 1  # (k - 1,)
+    indices = indices.sort().values  # (k,)
 
-    msa = msa[:, indices]
+    msa = msa[:, indices]  # (b, k, l)
     if msa_attention_mask is not None:
-        msa_attention_mask = msa_attention_mask[:, indices]
+        msa_attention_mask = msa_attention_mask[:, indices]  # (b, k, l)
     if has_deletion is not None:
-        has_deletion = has_deletion[:, indices]
+        has_deletion = has_deletion[:, indices]  # (b, k, l)
     if deletion_value is not None:
-        deletion_value = deletion_value[:, indices]
-    return msa, msa_attention_mask, has_deletion, deletion_value
+        deletion_value = deletion_value[:, indices]  # (b, k, l)
+    return msa, msa_attention_mask, has_deletion, deletion_value  # MSA tensors retain (b, selected_depth, l); optional tensors remain None
 
 
 def maybe_apply_msa_column_masking(
     msa_attention_mask: Tensor | None,
     rate: float,
 ) -> Tensor | None:
+    # msa_attention_mask: (b, m, l) or None.
     if msa_attention_mask is None or rate <= 0.0 or msa_attention_mask.size(1) <= 1:
-        return msa_attention_mask
+        return msa_attention_mask  # (b, m, l) or None
 
     batch_size, _, length = msa_attention_mask.shape
-    col_keep = torch.rand(batch_size, length, device=msa_attention_mask.device) >= rate
-    col_keep = col_keep.unsqueeze(1).expand_as(msa_attention_mask).clone()
-    col_keep[:, 0, :] = True
-    return msa_attention_mask.bool() & col_keep
+    col_keep = torch.rand(batch_size, length, device=msa_attention_mask.device) >= rate  # (b, l)
+    col_keep = col_keep.unsqueeze(1).expand_as(msa_attention_mask).clone()  # (b, m, l)
+    col_keep[:, 0, :] = True  # (b, l)
+    return msa_attention_mask.bool() & col_keep  # (b, m, l) or None
 
 
 # ===========================================================================
@@ -296,8 +300,8 @@ def gather_token_to_atom(token_features: Tensor, atom_to_token_idx: Tensor) -> T
     Returns:
         X with shape (b, a, d).
     """
-    idx = atom_to_token_idx.unsqueeze(-1).expand(-1, -1, token_features.size(-1))
-    return torch.gather(token_features, 1, idx)
+    idx = atom_to_token_idx.unsqueeze(-1).expand(-1, -1, token_features.size(-1))  # (b, a, d)
+    return torch.gather(token_features, 1, idx)  # (b, a, d)
 
 
 def scatter_atom_to_token(
@@ -319,20 +323,20 @@ def scatter_atom_to_token(
     """
     batch_size, n_atoms, d_model = atom_features.shape
     n_out = n_tokens
-    idx = atom_to_token_idx
+    idx = atom_to_token_idx  # (b, a)
     if atom_mask is not None:
-        idx = torch.where(atom_mask, atom_to_token_idx, n_tokens)
+        idx = torch.where(atom_mask, atom_to_token_idx, n_tokens)  # (b, a)
         n_out = n_tokens + 1
-    idx_expanded = idx.unsqueeze(-1).expand(batch_size, n_atoms, d_model)
+    idx_expanded = idx.unsqueeze(-1).expand(batch_size, n_atoms, d_model)  # (b, a, d)
     out = torch.zeros(
         batch_size,
         n_out,
         d_model,
         device=atom_features.device,
         dtype=atom_features.dtype,
-    )
-    out.scatter_reduce_(1, idx_expanded, atom_features, reduce="mean", include_self=False)
-    return out[:, :n_tokens, :]
+    )  # (b, n_out, d)
+    out.scatter_reduce_(1, idx_expanded, atom_features, reduce="mean", include_self=False)  # (b, n_out, d)
+    return out[:, :n_tokens, :]  # (b, l, d)
 
 
 def gather_rep_atom_coords(coords: Tensor, rep_atom_idx: Tensor) -> Tensor:
@@ -345,8 +349,8 @@ def gather_rep_atom_coords(coords: Tensor, rep_atom_idx: Tensor) -> Tensor:
     Returns:
         X with shape (b, l, 3).
     """
-    idx = rep_atom_idx.unsqueeze(-1).expand(-1, -1, coords.size(-1))
-    return torch.gather(coords, 1, idx)
+    idx = rep_atom_idx.unsqueeze(-1).expand(-1, -1, coords.size(-1))  # (b, l, 3)
+    return torch.gather(coords, 1, idx)  # (b, l, 3)
 
 
 def _compute_intra_token_idx(atom_to_token: Tensor) -> Tensor:
@@ -362,12 +366,12 @@ def _compute_intra_token_idx(atom_to_token: Tensor) -> Tensor:
         Index tensor I with shape (b, a) and values from zero through
         ``max_atoms_per_token - 1``.
     """
-    same_as_prev = F.pad(atom_to_token[:, 1:] == atom_to_token[:, :-1], (1, 0), value=False)
-    ones = torch.ones_like(atom_to_token)
-    cumsum = torch.cumsum(ones, dim=-1)
-    group_start = cumsum.masked_fill(same_as_prev, 0)
-    group_start = torch.cummax(group_start, dim=-1).values
-    return cumsum - group_start
+    same_as_prev = F.pad(atom_to_token[:, 1:] == atom_to_token[:, :-1], (1, 0), value=False)  # (b, a)
+    ones = torch.ones_like(atom_to_token)  # (b, a)
+    cumsum = torch.cumsum(ones, dim=-1)  # (b, a)
+    group_start = cumsum.masked_fill(same_as_prev, 0)  # (b, a)
+    group_start = torch.cummax(group_start, dim=-1).values  # (b, a)
+    return cumsum - group_start  # (b, a)
 
 
 def _categorical_mean(logits: Tensor, start: float, end: float) -> Tensor:
@@ -384,9 +388,9 @@ def _categorical_mean(logits: Tensor, start: float, end: float) -> Tensor:
         Expected value tensor Y with shape (...).
     """
     n_bins = logits.shape[-1]
-    edges = torch.linspace(start, end, n_bins + 1, device=logits.device, dtype=torch.float32)
+    edges = torch.linspace(start, end, n_bins + 1, device=logits.device, dtype=torch.float32)  # (n_bins + 1,)
     v_bins = (edges[:-1] + edges[1:]) / 2  # V_bin has shape (n_bins,).
-    return (logits.float().softmax(-1) @ v_bins.unsqueeze(1)).squeeze(-1)
+    return (logits.float().softmax(-1) @ v_bins.unsqueeze(1)).squeeze(-1)  # logits.shape[:-1]
 
 
 # ===========================================================================
@@ -403,16 +407,17 @@ class RowAttentionPooling(nn.Module):
         self.out_proj = nn.Linear(d_pair, d_single, bias=False)
 
     def forward(self, z: Tensor, mask: Tensor) -> Tensor:
-        scores = self.attn_proj(z).squeeze(-1)
+        # z: (b, l, l, d_pair); mask: (b, l).
+        scores = self.attn_proj(z).squeeze(-1)  # (b, l, l)
         mask_bias = torch.where(
             mask[:, None, :].bool(),
             torch.zeros_like(scores),
             torch.full_like(scores, -1e9),
-        )
-        scores = scores + mask_bias
-        weights = F.softmax(scores, dim=-1)
-        pooled = torch.einsum("bnm,bnmd->bnd", weights, z)
-        return self.out_proj(pooled)
+        )  # (b, l, l)
+        scores = scores + mask_bias  # (b, l, l)
+        weights = F.softmax(scores, dim=-1)  # (b, l, l)
+        pooled = torch.einsum("bnm,bnmd->bnd", weights, z)  # (b, l, d_pair)
+        return self.out_proj(pooled)  # (b, l, d_single)
 
 
 # ===========================================================================
@@ -460,6 +465,7 @@ class InputsEmbedder(nn.Module):
             X with shape (b, l, d_inputs), concatenating atom encoding,
             aatype, profile, and deletion mean.
         """
+        # aatype/profile: (b, l, 33); deletion_mean: (b, l); atom features use a atoms.
         a, _q, _c, _attn_params, _intermediates = self.atom_attention_encoder(
             ref_pos=ref_pos,
             atom_attention_mask=atom_attention_mask,
@@ -468,8 +474,8 @@ class InputsEmbedder(nn.Module):
             ref_element=ref_element,
             ref_atom_name_chars=ref_atom_name_chars,
             atom_to_token=atom_to_token,
-        )
-        return torch.cat([a, aatype, profile, deletion_mean.unsqueeze(-1)], dim=-1)
+        )  # a: (b, l, d_token / 2); _q/_c: (b, a, d_atom)
+        return torch.cat([a, aatype, profile, deletion_mean.unsqueeze(-1)], dim=-1)  # (b, l, d_token / 2 + 67)
 
 
 # ===========================================================================
@@ -511,38 +517,39 @@ class ResIdxAsymIdSymIdEntityIdEncoding(nn.Module):
         entity_id: Tensor,
         token_index: Tensor,
     ) -> Tensor:
-        bij_same_chain = asym_id.unsqueeze(2) == asym_id.unsqueeze(1)
-        bij_same_residue = residue_index.unsqueeze(2) == residue_index.unsqueeze(1)
-        bij_same_entity = entity_id.unsqueeze(2) == entity_id.unsqueeze(1)
+        # All input IDs: (b, l); r/c are relative residue/chain bin counts.
+        bij_same_chain = asym_id.unsqueeze(2) == asym_id.unsqueeze(1)  # (b, l, l)
+        bij_same_residue = residue_index.unsqueeze(2) == residue_index.unsqueeze(1)  # (b, l, l)
+        bij_same_entity = entity_id.unsqueeze(2) == entity_id.unsqueeze(1)  # (b, l, l)
 
-        dij_residue = residue_index.unsqueeze(2) - residue_index.unsqueeze(1)
+        dij_residue = residue_index.unsqueeze(2) - residue_index.unsqueeze(1)  # (b, l, l)
         dij_residue = torch.clip(
             dij_residue + self.n_relative_residx_bins,
             0,
             2 * self.n_relative_residx_bins,
-        )
-        dij_residue = torch.where(bij_same_chain, dij_residue, 2 * self.n_relative_residx_bins + 1)
-        aij_rel_pos = F.one_hot(dij_residue, 2 * self.n_relative_residx_bins + 2)
+        )  # (b, l, l)
+        dij_residue = torch.where(bij_same_chain, dij_residue, 2 * self.n_relative_residx_bins + 1)  # (b, l, l)
+        aij_rel_pos = F.one_hot(dij_residue, 2 * self.n_relative_residx_bins + 2)  # (b, l, l, 2 * r + 2)
 
         dij_token = torch.clip(
             token_index.unsqueeze(2) - token_index.unsqueeze(1) + self.n_relative_residx_bins,
             0,
             2 * self.n_relative_residx_bins,
-        )
+        )  # (b, l, l)
         dij_token = torch.where(
             bij_same_chain & bij_same_residue,
             dij_token,
             2 * self.n_relative_residx_bins + 1,
-        )
-        aij_rel_token = F.one_hot(dij_token, 2 * self.n_relative_residx_bins + 2)
+        )  # (b, l, l)
+        aij_rel_token = F.one_hot(dij_token, 2 * self.n_relative_residx_bins + 2)  # (b, l, l, 2 * r + 2)
 
         dij_chain = torch.clip(
             sym_id.unsqueeze(2) - sym_id.unsqueeze(1) + self.n_relative_chain_bins,
             0,
             2 * self.n_relative_chain_bins,
-        )
-        dij_chain = torch.where(bij_same_chain, 2 * self.n_relative_chain_bins + 1, dij_chain)
-        aij_rel_chain = F.one_hot(dij_chain, 2 * self.n_relative_chain_bins + 2)
+        )  # (b, l, l)
+        dij_chain = torch.where(bij_same_chain, 2 * self.n_relative_chain_bins + 1, dij_chain)  # (b, l, l)
+        aij_rel_chain = F.one_hot(dij_chain, 2 * self.n_relative_chain_bins + 2)  # (b, l, l, 2 * c + 2)
 
         feats = torch.cat(
             [
@@ -552,9 +559,9 @@ class ResIdxAsymIdSymIdEntityIdEncoding(nn.Module):
                 aij_rel_chain.float(),
             ],
             dim=-1,
-        )
+        )  # (b, l, l, 2 * (2 * r + 2) + 1 + 2 * c + 2)
 
-        return self.embed(feats)
+        return self.embed(feats)  # (b, l, l, d_pair)
 
 
 # ===========================================================================
@@ -575,12 +582,13 @@ class SingleToPair(nn.Module):
         )
 
     def forward(self, x: Tensor) -> Tensor:
-        x = self.downproject(x)
+        # x: (b, l, input_dim); d_down is downproject.out_features.
+        x = self.downproject(x)  # (b, l, d_down)
         x = torch.cat(
             [(x.unsqueeze(2) * x.unsqueeze(1)), (x.unsqueeze(2) - x.unsqueeze(1))],
             dim=3,
-        )
-        return self.output_mlp(x)
+        )  # (b, l, l, 2 * d_down)
+        return self.output_mlp(x)  # (b, l, l, output_dim)
 
 
 # ===========================================================================
@@ -604,7 +612,7 @@ class LanguageModelShim(nn.Module):
         self.base_z_linear = nn.Sequential(
             nn.LayerNorm(d_model), nn.Linear(d_model, d_z, bias=False)
         )
-        self.base_z_combine = nn.Parameter(torch.zeros(num_layers + 1))
+        self.base_z_combine = nn.Parameter(torch.zeros(num_layers + 1))  # (num_layers + 1,)
 
     def project_sequence(
         self,
@@ -641,12 +649,12 @@ class LanguageModelShim(nn.Module):
         # Match the learned projection parameters at this explicit boundary;
         # this preserves the official BF16 path and leaves FP32 models exact.
         projection_dtype = cast(nn.LayerNorm, self.base_z_linear[0]).weight.dtype
-        hidden_states = hidden_states.to(dtype=projection_dtype)
-        projected_states = self.base_z_linear(hidden_states)
-        layer_weights = self.base_z_combine.softmax(dim=0)
+        hidden_states = hidden_states.to(dtype=projection_dtype)  # (b, l, n_layers + 1, d_model)
+        projected_states = self.base_z_linear(hidden_states)  # (b, l, n_layers + 1, d_z)
+        layer_weights = self.base_z_combine.softmax(dim=0)  # (n_layers + 1,)
         # Preserve Biohub's matmul path exactly so checkpoint inference does
         # not change through a different reduction order.
-        projected = layer_weights @ projected_states
+        projected = layer_weights @ projected_states  # (b, l, d_z)
         if residue_mask is not None:
             if residue_mask.shape != hidden_states.shape[:2]:
                 raise ValueError(
@@ -655,8 +663,8 @@ class LanguageModelShim(nn.Module):
                 )
             projected = projected * residue_mask.to(
                 device=projected.device, dtype=projected.dtype
-            ).unsqueeze(-1)
-        return projected
+            ).unsqueeze(-1)  # (b, l, d_z)
+        return projected  # (b, l, d_z)
 
     def forward(self, hidden_states: Tensor, *, lm_dropout: float = 0.0) -> Tensor:
         """Project pre-computed ESMC hidden states to pair representation.
@@ -669,11 +677,11 @@ class LanguageModelShim(nn.Module):
         Returns:
             Z_pair with shape ``(b, l, l, d_pair)``.
         """
-        lm_z = self.project_sequence(hidden_states)
-        lm_z = self.base_z_mlp(lm_z)
+        lm_z = self.project_sequence(hidden_states)  # (b, l, d_z)
+        lm_z = self.base_z_mlp(lm_z)  # (b, l, l, d_z)
         if lm_dropout > 0:
-            lm_z = F.dropout(lm_z, p=lm_dropout, training=True)
-        return lm_z
+            lm_z = F.dropout(lm_z, p=lm_dropout, training=True)  # (b, l, l, d_z)
+        return lm_z  # (b, l, l, d_z)
 
 
 # ===========================================================================
@@ -700,62 +708,63 @@ def compute_lm_hidden_states(
     was trained on per-residue inputs, not per-atom), then scatter the
     hidden states back to the per-token layout.
     """
+    # Input IDs/masks: (b, l). Per sample p protein tokens collapse to u residues; t is padded LM length.
     b_size, l_size = input_ids.shape
     device = input_ids.device
-    protein_mask = (mol_type == 0) & token_mask
+    protein_mask = (mol_type == 0) & token_mask  # (b, l)
 
     lm_input_list = []
     lm_lengths = []
     # Per-batch maps from (original protein-token index) to (LM input position).
     expand_maps: list[Tensor] = []
     for batch_index in range(b_size):
-        mask_b = protein_mask[batch_index]
-        ids_b = input_ids[batch_index][mask_b]
-        asym_b = asym_id[batch_index][mask_b]
-        res_b = residue_index[batch_index][mask_b]
+        mask_b = protein_mask[batch_index]  # (l,)
+        ids_b = input_ids[batch_index][mask_b]  # (p,)
+        asym_b = asym_id[batch_index][mask_b]  # (p,)
+        res_b = residue_index[batch_index][mask_b]  # (p,)
 
         # Collapse: keep first token per (asym_id, residue_index) key, in
         # input order. ``inverse`` maps each original protein-token to its
         # collapsed residue index.
-        keys = torch.stack((asym_b, res_b), dim=1)
-        unique_keys, inverse = torch.unique(keys, dim=0, return_inverse=True)
+        keys = torch.stack((asym_b, res_b), dim=1)  # (p, 2)
+        unique_keys, inverse = torch.unique(keys, dim=0, return_inverse=True)  # (u, 2), (p,)
         n_unique = unique_keys.size(0)
-        token_positions = torch.arange(keys.size(0), device=device, dtype=torch.long)
-        first_pos = torch.full((n_unique,), keys.size(0), device=device, dtype=torch.long)
-        first_pos.scatter_reduce_(0, inverse, token_positions, reduce="amin", include_self=True)
-        ordered = torch.argsort(first_pos)
-        first_pos_ordered = first_pos[ordered]
-        ids_collapsed = ids_b[first_pos_ordered]
-        asym_collapsed = asym_b[first_pos_ordered]
-        remap = torch.empty_like(ordered)
-        remap[ordered] = torch.arange(n_unique, device=device, dtype=torch.long)
-        inverse_ordered = remap[inverse]
-
-        chain_ids = asym_collapsed.unique(sorted=True)
+        token_positions = torch.arange(keys.size(0), device=device, dtype=torch.long)  # (p,)
+        first_pos = torch.full((n_unique,), keys.size(0), device=device, dtype=torch.long)  # (u,)
+        first_pos.scatter_reduce_(0, inverse, token_positions, reduce="amin", include_self=True)  # (u,)
+        ordered = torch.argsort(first_pos)  # (u,)
+        first_pos_ordered = first_pos[ordered]  # (u,)
+        ids_collapsed = ids_b[first_pos_ordered]  # (u,)
+        asym_collapsed = asym_b[first_pos_ordered]  # (u,)
+        remap = torch.empty_like(ordered)  # (u,)
+        remap[ordered] = torch.arange(n_unique, device=device, dtype=torch.long)  # (u,)
+        inverse_ordered = remap[inverse]  # (p,)
+
+        chain_ids = asym_collapsed.unique(sorted=True)  # (n_chains,)
         # [BOS] chain1 [EOS BOS] chain2 ... [EOS]
-        parts: list[Tensor] = [torch.tensor([0], device=device, dtype=ids_b.dtype)]
+        parts: list[Tensor] = [torch.tensor([0], device=device, dtype=ids_b.dtype)]  # list of 1D token tensors
         # Per-chain LM positions accumulate; track them for the expand map.
-        per_token_lm_pos = torch.empty(n_unique, device=device, dtype=torch.long)
+        per_token_lm_pos = torch.empty(n_unique, device=device, dtype=torch.long)  # (u,)
         cursor = 1  # position 0 is the leading BOS
         for i, cid in enumerate(chain_ids):
-            in_chain = (asym_collapsed == cid).nonzero(as_tuple=True)[0]
+            in_chain = (asym_collapsed == cid).nonzero(as_tuple=True)[0]  # (u_chain,)
             parts.append(ids_collapsed[in_chain])
             per_token_lm_pos[in_chain] = torch.arange(
                 cursor, cursor + in_chain.shape[0], device=device, dtype=torch.long
-            )
+            )  # (u_chain,)
             cursor += in_chain.shape[0]
             if i < len(chain_ids) - 1:
                 parts.append(torch.tensor([2, 0], device=device, dtype=ids_b.dtype))
                 cursor += 2  # EOS + BOS
         parts.append(torch.tensor([2], device=device, dtype=ids_b.dtype))
-        lm_seq = torch.cat(parts)
+        lm_seq = torch.cat(parts)  # (t_i,)
         lm_input_list.append(lm_seq)
         lm_lengths.append(lm_seq.shape[0])
 
         # Map each original protein-token position to its LM input position.
-        prot_pos_b = mask_b.nonzero(as_tuple=True)[0]
-        expand_map = torch.full((l_size,), -1, device=device, dtype=torch.long)
-        expand_map[prot_pos_b] = per_token_lm_pos[inverse_ordered]
+        prot_pos_b = mask_b.nonzero(as_tuple=True)[0]  # (p,)
+        expand_map = torch.full((l_size,), -1, device=device, dtype=torch.long)  # (l,)
+        expand_map[prot_pos_b] = per_token_lm_pos[inverse_ordered]  # (p,)
         expand_maps.append(expand_map)
 
     # Pad the language-model input to its longest sequence. FP8 callers round
@@ -768,32 +777,32 @@ def compute_lm_hidden_states(
         1,
         device=device,
         dtype=input_ids.dtype,  # PAD=1
-    )
+    )  # (b, t)
     for batch_index in range(b_size):
-        lm_input_ids[batch_index, : lm_lengths[batch_index]] = lm_input_list[batch_index]
+        lm_input_ids[batch_index, : lm_lengths[batch_index]] = lm_input_list[batch_index]  # (t_i,)
 
     # sequence_id for chain-aware attention; PAD tokens get -1 (no attention).
-    sequence_id = (lm_input_ids == 0).cumsum(dim=1) - 1  # BOS=0
-    sequence_id = sequence_id.masked_fill(lm_input_ids == 1, -1)  # PAD=1
+    sequence_id = (lm_input_ids == 0).cumsum(dim=1) - 1  # BOS=0; (b, t)
+    sequence_id = sequence_id.masked_fill(lm_input_ids == 1, -1)  # PAD=1; (b, t)
 
     if lm_mask_pct > 0.0:
-        special = (lm_input_ids == 0) | (lm_input_ids == 1) | (lm_input_ids == 2)
-        do_mask = (torch.rand(lm_input_ids.shape, device=device) < lm_mask_pct) & ~special
-        lm_input_ids = lm_input_ids.masked_fill(do_mask, mask_token_id)
+        special = (lm_input_ids == 0) | (lm_input_ids == 1) | (lm_input_ids == 2)  # (b, t)
+        do_mask = (torch.rand(lm_input_ids.shape, device=device) < lm_mask_pct) & ~special  # (b, t)
+        lm_input_ids = lm_input_ids.masked_fill(do_mask, mask_token_id)  # (b, t)
 
     with torch.inference_mode():
         esmc_out = esmc(input_ids=lm_input_ids, sequence_id=sequence_id, output_hidden_states=True)
 
-    hidden_stack = esmc_out.hidden_states
+    hidden_stack = esmc_out.hidden_states  # (n_states, b, t, d_model)
     n_states, _, _, d_model = hidden_stack.shape
-    result = torch.zeros(b_size, l_size, n_states, d_model, device=device, dtype=hidden_stack.dtype)
+    result = torch.zeros(b_size, l_size, n_states, d_model, device=device, dtype=hidden_stack.dtype)  # (b, l, n_states, d_model)
     for batch_index in range(b_size):
-        M_i = protein_mask[batch_index]
-        positions = expand_maps[batch_index][M_i]
-        gathered = hidden_stack[:, batch_index, positions, :].permute(1, 0, 2)
-        result[batch_index, M_i.nonzero(as_tuple=True)[0]] = gathered
+        M_i = protein_mask[batch_index]  # (l,)
+        positions = expand_maps[batch_index][M_i]  # (p,)
+        gathered = hidden_stack[:, batch_index, positions, :].permute(1, 0, 2)  # (p, n_states, d_model)
+        result[batch_index, M_i.nonzero(as_tuple=True)[0]] = gathered  # (p, n_states, d_model)
 
-    return result.detach()
+    return result.detach()  # (b, l, n_states, d_model)
 
 
 # ===========================================================================
@@ -844,7 +853,8 @@ class TriangleMultiplicativeBlock(nn.Module):
         return self.flow
 
     def _triangular_contract(self, left_stream: Tensor, right_stream: Tensor) -> Tensor:
-        return torch.einsum(self._einsum_equation, left_stream, right_stream)
+        # Streams: (b, l, l, d_latent); equation chooses incoming/outgoing contraction.
+        return torch.einsum(self._einsum_equation, left_stream, right_stream)  # (b, l, l, d_latent)
 
     def _triangular_contract_chunked(
         self, left_stream: Tensor, right_stream: Tensor, chunk_size: int
@@ -869,17 +879,18 @@ class TriangleMultiplicativeBlock(nn.Module):
         chunks = []
         for start in range(0, length, chunk_size):
             rows = left_rows[:, :, start : start + chunk_size]  # (b, d, i_c, k)
-            product = torch.bmm(rows.reshape(batch_size * channels, -1, inner), right_columns)
+            product = torch.bmm(rows.reshape(batch_size * channels, -1, inner), right_columns)  # (b * d, i_c, j)
             product = product.view(batch_size, channels, rows.shape[2], -1)  # (b, d, i_c, j)
             chunks.append(product.permute(0, 2, 3, 1))  # (b, i_c, j, d)
-        return torch.cat(chunks, dim=1)
+        return torch.cat(chunks, dim=1)  # (b, l, l, d)
 
     def forward(self, pair_grid: Tensor, visibility: Tensor | None = None) -> Tensor:
+        # pair_grid: (b, l, l, d_input); visibility: (b, l, l); d = latent_channels.
         if visibility is None:
-            visibility = pair_grid.new_ones(pair_grid.shape[:-1])
+            visibility = pair_grid.new_ones(pair_grid.shape[:-1])  # (b, l, l)
 
         if self._use_kernels:
-            p_in_weight, g_in_weight = self.split_kernel_weights()
+            p_in_weight, g_in_weight = self.split_kernel_weights()  # each (2 * d, d_input)
             return _cue_tri_mul(  # type: ignore[misc]
                 pair_grid,
                 direction=self._kernel_flow_direction(),
@@ -893,38 +904,38 @@ class TriangleMultiplicativeBlock(nn.Module):
                 p_out_weight=self.proj_emit.weight,
                 g_out_weight=self.proj_gate.weight,
                 eps=_EPS,
-            )
+            )  # (b, l, l, d_input)
 
         # Every tensor below is as large as the pair representation or larger, and this
         # block sets the peak memory of a fold. Each name is dropped once it is dead, so
         # the allocator can reuse its buffer. No value changes.
-        normalized_grid = self.norm_start(pair_grid)
+        normalized_grid = self.norm_start(pair_grid)  # (b, l, l, d_input)
         bundled = self.proj_bundle(normalized_grid)  # (b, l, l, 4 * d)
-        signal, gate_logits = bundled.split(2 * self.latent_channels, dim=-1)
+        signal, gate_logits = bundled.split(2 * self.latent_channels, dim=-1)  # each (b, l, l, 2 * d)
         routed = signal * torch.sigmoid(gate_logits)  # (b, l, l, 2 * d)
         # The two views would keep the whole projection alive.
         del bundled, signal, gate_logits
-        routed = routed * visibility.unsqueeze(-1)
+        routed = routed * visibility.unsqueeze(-1)  # (b, l, l, 2 * d)
 
         left_stream, right_stream = routed.float().chunk(2, dim=-1)  # each (b, l, l, d)
         if torch.is_autocast_enabled(left_stream.device.type):
             # The contraction is an autocast operation. Casting its inputs here, as it
             # would, lets the full-precision product go before the contraction runs.
             autocast_dtype = torch.get_autocast_dtype(left_stream.device.type)
-            left_stream = left_stream.to(autocast_dtype)
-            right_stream = right_stream.to(autocast_dtype)
+            left_stream = left_stream.to(autocast_dtype)  # (b, l, l, d)
+            right_stream = right_stream.to(autocast_dtype)  # (b, l, l, d)
         del routed
         if self._chunk_size is not None:
             contracted = self._triangular_contract_chunked(
                 left_stream, right_stream, self._chunk_size
-            )
+            )  # (b, l, l, d)
         else:
-            contracted = self._triangular_contract(left_stream, right_stream)
+            contracted = self._triangular_contract(left_stream, right_stream)  # (b, l, l, d)
         del left_stream, right_stream
-        mixed = self.proj_emit(self.norm_mix(contracted))
+        mixed = self.proj_emit(self.norm_mix(contracted))  # (b, l, l, d_input)
         del contracted
-        output_gate = torch.sigmoid(self.proj_gate(normalized_grid))
-        return mixed * output_gate
+        output_gate = torch.sigmoid(self.proj_gate(normalized_grid))  # (b, l, l, d_input)
+        return mixed * output_gate  # (b, l, l, d_input)
 
 
 class TriangleMultiplicativeUpdate(nn.Module):
@@ -947,7 +958,8 @@ class TriangleMultiplicativeUpdate(nn.Module):
         self._engine.set_chunk_size(chunk_size)
 
     def forward(self, z: Tensor, mask: Tensor | None = None) -> Tensor:
-        return self._engine(z, visibility=mask)
+        # z: (b, l, l, d_pair); mask: (b, l, l) or None.
+        return self._engine(z, visibility=mask)  # (b, l, l, d_pair)
 
 
 # ===========================================================================
@@ -989,11 +1001,11 @@ class Transition(nn.Module):
                 dtype=dtype,
             )
             with torch.no_grad():
-                fused.LN_W.copy_(self.norm.weight)
+                fused.LN_W.copy_(self.norm.weight)  # (d_model,)
                 if has_ln_bias:
                     fused.LN_B.copy_(self.norm.bias)  # type: ignore[union-attr]
                 # FusedLNLinearSwiGLU.W12 is (d_model, 2*d_inner); transpose nn.Linear once.
-                fused.W12.copy_(self.ffn.w12.weight.t().contiguous())
+                fused.W12.copy_(self.ffn.w12.weight.t().contiguous())  # (d_model, 2 * d_inner)
             self._fused_swiglu = fused.eval().requires_grad_(False)
         else:
             self._fused_swiglu = None
@@ -1005,50 +1017,53 @@ class Transition(nn.Module):
 
     def _swiglu_pre_w3(self, x_normed: Tensor) -> Tensor:
         """SwiGLU through silu(x1)*x2, before the final w3."""
+        # x_normed: (..., d_model); d_inner = ffn.hidden_features.
         ffn = self.ffn
-        x12 = ffn.w12(x_normed)
-        x1, x2 = x12.split(ffn.hidden_features, dim=-1)
-        return F.silu(x1) * x2
+        x12 = ffn.w12(x_normed)  # (..., 2 * d_inner)
+        x1, x2 = x12.split(ffn.hidden_features, dim=-1)  # each (..., d_inner)
+        return F.silu(x1) * x2  # (..., d_inner)
 
     def _addmm_residual(self, x: Tensor, hidden: Tensor) -> Tensor:
         """x + w3(hidden) via single cuBLAS addmm: avoids transition-output allocation."""
+        # x: (..., d_model); hidden: (..., d_inner).
         ffn = self.ffn
         x_shape = x.shape
         out = torch.addmm(
             x.contiguous().view(-1, x_shape[-1]),
             hidden.view(-1, hidden.shape[-1]),
             ffn.w3.weight.t(),
-        )
-        return out.view(x_shape)
+        )  # (product(x.shape[:-1]), d_model)
+        return out.view(x_shape)  # x.shape
 
     def forward(self, x: Tensor) -> Tensor:
         # Inference-only fast path (addmm-fused residual + pre-alloc out)
         #: diverges bit-exactly from ``x + ffn(norm(x))`` so we only use
         # it when grad is disabled (binder-design / bit-exact tests run
         # with grad on and need the reference path).
+        # x: (b, l, ..., d_model); chunk width l_c <= _chunk_size.
         if not torch.is_grad_enabled() and self._can_use_fused_path(x):
             fused = self._fused_swiglu
             assert fused is not None
             pre_w3 = fused
             if self._chunk_size is None or x.shape[1] <= self._chunk_size:
-                hidden = pre_w3(x)
-                return self._addmm_residual(x, hidden)
-            out = torch.empty_like(x)
+                hidden = pre_w3(x)  # (b, l, ..., d_inner)
+                return self._addmm_residual(x, hidden)  # x.shape
+            out = torch.empty_like(x)  # x.shape
             for s in range(0, x.shape[1], self._chunk_size):
                 e = min(s + self._chunk_size, x.shape[1])
-                sl = x[:, s:e]
-                hidden = pre_w3(sl)
-                out[:, s:e] = self._addmm_residual(sl, hidden)
-            return out
+                sl = x[:, s:e]  # (b, l_c, ..., d_model)
+                hidden = pre_w3(sl)  # (b, l_c, ..., d_inner)
+                out[:, s:e] = self._addmm_residual(sl, hidden)  # (b, l_c, ..., d_model)
+            return out  # x.shape
         # Reference path: bit-exact with main: x + ffn(norm(x)).
         if self._chunk_size is None or x.shape[1] <= self._chunk_size:
-            return x + self.ffn(self.norm(x))
+            return x + self.ffn(self.norm(x))  # x.shape
         out_list: list[Tensor] = []
         for s in range(0, x.shape[1], self._chunk_size):
             e = min(s + self._chunk_size, x.shape[1])
-            sl = x[:, s:e]
+            sl = x[:, s:e]  # (b, l_c, ..., d_model)
             out_list.append(sl + self.ffn(self.norm(sl)))
-        return torch.cat(out_list, dim=1)
+        return torch.cat(out_list, dim=1)  # x.shape
 
 
 class PairUpdateBlock(nn.Module):
@@ -1087,12 +1102,13 @@ class PairUpdateBlock(nn.Module):
         self, pair: Tensor, direction: str, pair_attention_mask: Tensor | None
     ) -> Tensor:
         """Fused TriMul+residual call; weights from the corresponding engine."""
+        # pair: (b, l, l, d_pair); pair_attention_mask: (b, l, l) or None.
         tri = self.tri_mul_out if direction == "outgoing" else self.tri_mul_in
         engine: TriangleMultiplicativeBlock = tri._engine  # type: ignore[assignment]
-        p_in_weight, g_in_weight = engine.split_kernel_weights()
+        p_in_weight, g_in_weight = engine.split_kernel_weights()  # each (2 * d_pair, d_pair)
 
         def _bf16(t: Tensor) -> Tensor:
-            return t if t.dtype == torch.bfloat16 else t.to(torch.bfloat16)
+            return t if t.dtype == torch.bfloat16 else t.to(torch.bfloat16)  # t.shape
 
         return _fused_trimul_with_residual(  # type: ignore[misc]
             pair,
@@ -1109,17 +1125,18 @@ class PairUpdateBlock(nn.Module):
             g_out_weight=_bf16(engine.proj_gate.weight),
             mask=pair_attention_mask,
             eps=_EPS,
-        )
+        )  # pair.shape
 
     def forward(self, pair: Tensor, pair_attention_mask: Tensor | None = None) -> Tensor:
+        # pair: (b, l, l, d_pair); pair_attention_mask: (b, l, l) or None.
         if self._can_use_fused_trimul_with_residual(pair):
-            pair = self._fused_trimul_with_residual(pair, "outgoing", pair_attention_mask)
-            pair = self._fused_trimul_with_residual(pair, "incoming", pair_attention_mask)
+            pair = self._fused_trimul_with_residual(pair, "outgoing", pair_attention_mask)  # (b, l, l, d_pair)
+            pair = self._fused_trimul_with_residual(pair, "incoming", pair_attention_mask)  # (b, l, l, d_pair)
         else:
-            pair = self.row_drop(pair, self.tri_mul_out(pair, mask=pair_attention_mask))
-            pair = self.row_drop(pair, self.tri_mul_in(pair, mask=pair_attention_mask))
-        pair = self.pair_transition(pair)
-        return pair
+            pair = self.row_drop(pair, self.tri_mul_out(pair, mask=pair_attention_mask))  # (b, l, l, d_pair)
+            pair = self.row_drop(pair, self.tri_mul_in(pair, mask=pair_attention_mask))  # (b, l, l, d_pair)
+        pair = self.pair_transition(pair)  # (b, l, l, d_pair)
+        return pair  # (b, l, l, d_pair)
 
 
 class FoldingTrunk(nn.Module):
@@ -1145,22 +1162,23 @@ class FoldingTrunk(nn.Module):
     def forward(self, pair: Tensor, pair_attention_mask: Tensor | None = None) -> Tensor:
         # Cast the pair tensor to BF16 when the fused triangle backend is enabled
         # (its bwd kernel requires bf16). Other backends keep the input dtype.
+        # pair: (b, l, l, d_pair); pair_attention_mask: (b, l, l) or None.
         orig_dtype = pair.dtype
         fused_on = (
             len(self.blocks) > 0
             and getattr(self.blocks[0], "_kernel_backend", None) == BACKEND_FUSED
         )
         if pair.is_cuda and fused_on and orig_dtype != torch.bfloat16:
-            pair = pair.to(torch.bfloat16)
+            pair = pair.to(torch.bfloat16)  # (b, l, l, d_pair)
         for block in self.blocks:
             fn = partial(block, pair_attention_mask=pair_attention_mask)
             if torch.is_grad_enabled():
-                pair = checkpoint(fn, pair, use_reentrant=False)  # pyright: ignore
+                pair = checkpoint(fn, pair, use_reentrant=False)  # pyright: ignore; (b, l, l, d_pair)
             else:
-                pair = fn(pair)
+                pair = fn(pair)  # (b, l, l, d_pair)
         if pair.dtype != orig_dtype:
-            pair = pair.to(orig_dtype)
-        return pair
+            pair = pair.to(orig_dtype)  # (b, l, l, d_pair)
+        return pair  # (b, l, l, d_pair)
 
 
 # ===========================================================================
@@ -1201,28 +1219,29 @@ class OuterProductMean(nn.Module):
         self._chunk_size = chunk_size
 
     def forward(self, m: Tensor, msa_attention_mask: Tensor) -> Tensor:
-        m_norm = self.norm(m)
-        x = self.W(m_norm) * msa_attention_mask.unsqueeze(-1).to(m_norm.dtype)
-        a, b = x.chunk(2, dim=-1)
-        mask_f = msa_attention_mask.to(a.dtype)
-        n_valid = (mask_f @ mask_f.transpose(-1, -2)).unsqueeze(-1).clamp(min=1.0)
+        # m: (b, l, m_depth, d_msa); msa_attention_mask: (b, l, m_depth); d_h = d_hidden.
+        m_norm = self.norm(m)  # (b, l, m_depth, d_msa)
+        x = self.W(m_norm) * msa_attention_mask.unsqueeze(-1).to(m_norm.dtype)  # (b, l, m_depth, 2 * d_h)
+        a, b = x.chunk(2, dim=-1)  # each (batch, l, m_depth, d_h)
+        mask_f = msa_attention_mask.to(a.dtype)  # (batch, l, m_depth)
+        n_valid = (mask_f @ mask_f.transpose(-1, -2)).unsqueeze(-1).clamp(min=1.0)  # (batch, l, l, 1)
         if self._chunk_size is None:
-            outer = torch.einsum("bimc,bjmd->bijcd", a, b).flatten(-2)
+            outer = torch.einsum("bimc,bjmd->bijcd", a, b).flatten(-2)  # (batch, l, l, d_h * d_h)
             if self.divide_outer_before_proj:
-                return self.Wout(outer / n_valid)
-            return self.Wout(outer) / n_valid
+                return self.Wout(outer / n_valid)  # (batch, l, l, d_pair)
+            return self.Wout(outer) / n_valid  # (batch, l, l, d_pair)
         # Chunk along the left (i) axis so the peak einsum intermediate is
         # X uses shape (b, chunk, l, c, d) instead of (b, l, l, c, d).
         length = a.shape[1]
         out_chunks: list[Tensor] = []
         for start in range(0, length, self._chunk_size):
             end = min(start + self._chunk_size, length)
-            outer_chunk = torch.einsum("bimc,bjmd->bijcd", a[:, start:end], b).flatten(-2)
+            outer_chunk = torch.einsum("bimc,bjmd->bijcd", a[:, start:end], b).flatten(-2)  # (batch, l_c, l, d_h * d_h)
             if self.divide_outer_before_proj:
                 out_chunks.append(self.Wout(outer_chunk / n_valid[:, start:end]))
             else:
                 out_chunks.append(self.Wout(outer_chunk) / n_valid[:, start:end])
-        return torch.cat(out_chunks, dim=1)
+        return torch.cat(out_chunks, dim=1)  # (batch, l, l, d_pair)
 
 
 class MSAPairWeightedAveraging(nn.Module):
@@ -1252,18 +1271,18 @@ class MSAPairWeightedAveraging(nn.Module):
         batch_size, length, depth, _ = msa_repr.shape
         n_heads, head_width = self.n_heads, self.head_width
 
-        msa_normed = self.norm_single(msa_repr)
+        msa_normed = self.norm_single(msa_repr)  # (b, l, m, d_msa)
         bias = self.compute_bias(pair_repr)  # A has shape (b, l, l, n_heads).
-        bias.masked_fill_(~pair_attention_mask.unsqueeze(-1).bool(), -1e5)
-        attn = torch.softmax(bias, dim=-2)  # softmax over j
+        bias.masked_fill_(~pair_attention_mask.unsqueeze(-1).bool(), -1e5)  # (b, l, l, n_heads)
+        attn = torch.softmax(bias, dim=-2)  # softmax over j; (b, l, l, n_heads)
 
-        v = self.Wv(msa_normed).reshape(batch_size, length, depth, n_heads, head_width)
+        v = self.Wv(msa_normed).reshape(batch_size, length, depth, n_heads, head_width)  # (b, l, m, n_heads, head_width)
         gate = torch.sigmoid(self.Wgate(msa_normed)).reshape(
             batch_size, length, depth, n_heads, head_width
-        )
+        )  # (b, l, m, n_heads, head_width)
 
-        output = torch.einsum("bijh,bjmhd,bimhd->bimhd", attn, v, gate)
-        return self.Wout(output.reshape(batch_size, length, depth, n_heads * head_width))
+        output = torch.einsum("bijh,bjmhd,bimhd->bimhd", attn, v, gate)  # (b, l, m, n_heads, head_width)
+        return self.Wout(output.reshape(batch_size, length, depth, n_heads * head_width))  # (b, l, m, d_msa)
 
 
 # ===========================================================================
@@ -1283,10 +1302,11 @@ class TransitionLayer(nn.Module):
         self.out_proj = nn.Linear(hidden, d_model, bias=False)
 
     def forward(self, x: Tensor) -> Tensor:
-        x = self.norm(x)
-        a = self.a_proj(x)
-        b = self.b_proj(x)
-        return self.out_proj(F.silu(a) * b)
+        # x: (..., d_model); d_hidden = n * d_model.
+        x = self.norm(x)  # (..., d_model)
+        a = self.a_proj(x)  # (..., d_hidden)
+        b = self.b_proj(x)  # (..., d_hidden)
+        return self.out_proj(F.silu(a) * b)  # (..., d_model)
 
 
 # ===========================================================================
@@ -1302,14 +1322,15 @@ class AdaptiveLayerNorm(nn.Module):
         self.d_model = d_model
         self.d_cond = d_cond
         self.eps = eps
-        self.s_scale = nn.Parameter(torch.ones(d_cond))
+        self.s_scale = nn.Parameter(torch.ones(d_cond))  # (d_cond,)
         self.s_gate = nn.Linear(d_cond, d_model, bias=True)
         self.s_shift = nn.Linear(d_cond, d_model, bias=False)
 
     def forward(self, a: Tensor, s: Tensor) -> Tensor:
-        a_norm = F.layer_norm(a, (self.d_model,), None, None, self.eps)
-        s_norm = F.layer_norm(s, (self.d_cond,), self.s_scale, None, self.eps)
-        return torch.sigmoid(self.s_gate(s_norm)) * a_norm + self.s_shift(s_norm)
+        # a: (..., d_model); s: (..., d_cond) with broadcast-compatible leading axes.
+        a_norm = F.layer_norm(a, (self.d_model,), None, None, self.eps)  # a.shape
+        s_norm = F.layer_norm(s, (self.d_cond,), self.s_scale, None, self.eps)  # s.shape
+        return torch.sigmoid(self.s_gate(s_norm)) * a_norm + self.s_shift(s_norm)  # broadcast leading shape + (d_model,)
 
 
 # ===========================================================================
@@ -1326,12 +1347,13 @@ class FourierEmbedding(nn.Module):
     def __init__(self, c: int) -> None:
         super().__init__()
         self.c = c
-        self.register_buffer("w", torch.randn(c))
-        self.register_buffer("b", torch.randn(c))
+        self.register_buffer("w", torch.randn(c))  # (c,)
+        self.register_buffer("b", torch.randn(c))  # (c,)
 
     def forward(self, t_hat: Tensor) -> Tensor:
-        t = torch.as_tensor(t_hat, device=self.w.device, dtype=self.w.dtype).reshape(-1)
-        return torch.cos(2.0 * torch.pi * (t[:, None] * self.w[None, :] + self.b[None, :]))
+        # t_hat: scalar or arbitrary noise-time tensor; n = t_hat.numel().
+        t = torch.as_tensor(t_hat, device=self.w.device, dtype=self.w.dtype).reshape(-1)  # (n,)
+        return torch.cos(2.0 * torch.pi * (t[:, None] * self.w[None, :] + self.b[None, :]))  # (n, c)
 
 
 # ===========================================================================
@@ -1360,14 +1382,15 @@ class SwiGLU(nn.Module):
         self.hidden_features = hidden_features
 
     def forward(self, x: Tensor) -> Tensor:
-        x12 = self.w12(x)
-        x1, x2 = x12.split(self.hidden_features, dim=-1)
-        hidden = F.silu(x1)
+        # x: (..., in_features); d_hidden = hidden_features.
+        x12 = self.w12(x)  # (..., 2 * d_hidden)
+        x1, x2 = x12.split(self.hidden_features, dim=-1)  # each (..., d_hidden)
+        hidden = F.silu(x1)  # (..., d_hidden)
         # Without autograd the product can reuse the activation's buffer. On a pair tensor
         # that buffer is twice the pair representation. The values are the same either way.
-        hidden = hidden * x2 if torch.is_grad_enabled() else hidden.mul_(x2)
+        hidden = hidden * x2 if torch.is_grad_enabled() else hidden.mul_(x2)  # (..., d_hidden)
         del x12, x1, x2
-        return self.w3(hidden)
+        return self.w3(hidden)  # (..., out_features)
 
 
 class SwiGLUMLP(SwiGLU):
@@ -1386,8 +1409,9 @@ class SwiGLUMLP(SwiGLU):
 
 
 def _rotate_half(x: Tensor) -> Tensor:
-    x1, x2 = x.chunk(2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
+    # x: (..., d_rot), with an even final width.
+    x1, x2 = x.chunk(2, dim=-1)  # each (..., d_rot / 2)
+    return torch.cat((-x2, x1), dim=-1)  # x.shape
 
 
 def apply_rotary_emb_3d(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
@@ -1399,12 +1423,12 @@ def apply_rotary_emb_3d(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
         sin: S with shape (b, l, d / 2).
     """
     ro_dim = cos.shape[-1] * 2
-    cos = cos.unsqueeze(2).repeat(1, 1, 1, 2)
-    sin = sin.unsqueeze(2).repeat(1, 1, 1, 2)
+    cos = cos.unsqueeze(2).repeat(1, 1, 1, 2)  # (b, l, 1, ro_dim)
+    sin = sin.unsqueeze(2).repeat(1, 1, 1, 2)  # (b, l, 1, ro_dim)
     return torch.cat(
         [x[..., :ro_dim] * cos + _rotate_half(x[..., :ro_dim]) * sin, x[..., ro_dim:]],
         dim=-1,
-    )
+    )  # (b, l, h, d)
 
 
 @torch.compiler.disable
@@ -1418,6 +1442,7 @@ def build_3d_rope(
     uid_base_freq: float = 10.0,
 ) -> tuple[Tensor, Tensor]:
     """Build cos/sin for 3D RoPE + UID RoPE."""
+    # ref_pos: (b, a, 3); ref_space_uid: (b, a); s/u = spatial/UID pair counts.
     device = ref_pos.device
     batch_size, n_atoms = ref_pos.shape[:2]
     half_dim = head_dim // 2
@@ -1429,21 +1454,21 @@ def build_3d_rope(
             torch.arange(0, n_spatial_per_axis, dtype=torch.float32, device=device)
             / n_spatial_per_axis
         )
-    )
+    )  # (s,)
     uid_inv_freq = 1.0 / (
         uid_base_freq
         ** (torch.arange(0, n_uid_pairs, dtype=torch.float32, device=device) / n_uid_pairs)
-    )
+    )  # (u,)
 
-    pos_f32 = ref_pos.float()
-    spatial_freqs = torch.einsum("bna,k->bnak", pos_f32, spatial_inv_freq)
-    spatial_freqs = spatial_freqs.reshape(batch_size, n_atoms, n_spatial_total)
+    pos_f32 = ref_pos.float()  # (b, a, 3)
+    spatial_freqs = torch.einsum("bna,k->bnak", pos_f32, spatial_inv_freq)  # (b, a, 3, s)
+    spatial_freqs = spatial_freqs.reshape(batch_size, n_atoms, n_spatial_total)  # (b, a, 3 * s)
 
-    uid_f32 = ref_space_uid.float()
-    uid_freqs = torch.einsum("bn,k->bnk", uid_f32, uid_inv_freq)
+    uid_f32 = ref_space_uid.float()  # (b, a)
+    uid_freqs = torch.einsum("bn,k->bnk", uid_f32, uid_inv_freq)  # (b, a, u)
 
     n_active = n_spatial_total + n_uid_pairs
-    freqs = torch.cat([spatial_freqs, uid_freqs], dim=-1)
+    freqs = torch.cat([spatial_freqs, uid_freqs], dim=-1)  # (b, a, 3 * s + u)
 
     if n_active < half_dim:
         padding = torch.zeros(
@@ -1452,16 +1477,17 @@ def build_3d_rope(
             half_dim - n_active,
             device=device,
             dtype=torch.float32,
-        )
-        freqs = torch.cat([freqs, padding], dim=-1)
+        )  # (b, a, half_dim - n_active)
+        freqs = torch.cat([freqs, padding], dim=-1)  # (b, a, half_dim)
 
-    cos = freqs.cos().to(torch.bfloat16)
-    sin = freqs.sin().to(torch.bfloat16)
-    return cos, sin
+    cos = freqs.cos().to(torch.bfloat16)  # freqs.shape
+    sin = freqs.sin().to(torch.bfloat16)  # freqs.shape
+    return cos, sin  # each (b, a, max(3 * s + u, half_dim))
 
 
 def qk_norm(x: Tensor) -> Tensor:
-    return F.rms_norm(x, (x.size(-1),)).to(x.dtype)
+    # x: arbitrary leading dimensions and a final head-width axis.
+    return F.rms_norm(x, (x.size(-1),)).to(x.dtype)  # x.shape
 
 
 # ===========================================================================
@@ -1479,9 +1505,10 @@ class SwiGLUFFN(nn.Module):
         self.w_down = nn.Linear(hidden_size, d_model, bias=False)
 
     def forward(self, x: Tensor) -> Tensor:
-        x = x.to(self.w_up.weight.dtype)
-        x1, x2 = self.w_up(x).chunk(2, dim=-1)
-        return self.w_down(F.silu(x1) * x2)
+        # x: (..., d_model); d_hidden = w_down.in_features.
+        x = x.to(self.w_up.weight.dtype)  # (..., d_model)
+        x1, x2 = self.w_up(x).chunk(2, dim=-1)  # each (..., d_hidden)
+        return self.w_down(F.silu(x1) * x2)  # (..., d_model)
 
 
 # ===========================================================================
@@ -1524,7 +1551,7 @@ class SWA3DRoPEAttention(nn.Module):
         # indices: (t,) flat positions of real atoms; cu_seqlens: (b + 1,) int32 row offsets.
         indices, cu_seqlens, max_seqlen = attention_params[2:5]
         flat_shape = (batch_size * n_atoms, self.n_heads, self.head_dim)
-        q, k, v = q.reshape(flat_shape), k.reshape(flat_shape), v.reshape(flat_shape)
+        q, k, v = q.reshape(flat_shape), k.reshape(flat_shape), v.reshape(flat_shape)  # each (b * n_atoms, h, d_h)
         has_padding = indices.shape[0] != batch_size * n_atoms
         if has_padding:
             q, k, v = q[indices], k[indices], v[indices]  # each (t, h, d_h)
@@ -1545,46 +1572,47 @@ class SWA3DRoPEAttention(nn.Module):
         )  # (t, h, d_h)
         if has_padding:
             out = attended.new_zeros(flat_shape)  # (b * n_atoms, h, d_h)
-            out[indices] = attended
+            out[indices] = attended  # (t, h, d_h)
         else:
-            out = attended
-        return out.view(batch_size, n_atoms, self.n_heads, self.head_dim)
+            out = attended  # (b * n_atoms, h, d_h)
+        return out.view(batch_size, n_atoms, self.n_heads, self.head_dim)  # (b, n_atoms, h, d_h)
 
     def forward(self, x: Tensor, attention_params: tuple) -> Tensor:
+        # x: (b, a, d_model); h = n_heads, d_h = head_dim; r is rotary-pair count.
         batch_size, n_atoms = x.shape[:2]
-        cos, sin = attention_params[0], attention_params[1]
+        cos, sin = attention_params[0], attention_params[1]  # each (b, a, r)
 
-        x_input = x
-        qkv = self.Wqkv(x)
-        qkv = qkv.view(batch_size, n_atoms, 3, self.n_heads, self.head_dim).permute(2, 0, 1, 3, 4)
-        q, k, v = qkv.unbind(0)
-        q, k = qk_norm(q), qk_norm(k)
+        x_input = x  # (b, a, d_model)
+        qkv = self.Wqkv(x)  # (b, a, 3 * d_model)
+        qkv = qkv.view(batch_size, n_atoms, 3, self.n_heads, self.head_dim).permute(2, 0, 1, 3, 4)  # (3, b, a, h, d_h)
+        q, k, v = qkv.unbind(0)  # each (b, a, h, d_h)
+        q, k = qk_norm(q), qk_norm(k)  # each (b, a, h, d_h)
 
-        q = apply_rotary_emb_3d(q, cos, sin)
-        k = apply_rotary_emb_3d(k, cos, sin)
+        q = apply_rotary_emb_3d(q, cos, sin)  # (b, a, h, d_h)
+        k = apply_rotary_emb_3d(k, cos, sin)  # (b, a, h, d_h)
 
         input_dtype = q.dtype
         if q.dtype not in (torch.float16, torch.bfloat16):
-            q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16()
+            q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16()  # each (b, a, h, d_h)
 
         # ESMFold2 does not advertise FlashAttention. Keep this atom path on
         # PyTorch. Models that advertise FlashAttention dispatch through the
         # precompiled Hugging Face kernels interface in fastplms.attention.
         if self._atom_attention == ATOM_ATTENTION_WINDOWED:
-            out = self._windowed_attention(q, k, v, attention_params)
+            out = self._windowed_attention(q, k, v, attention_params)  # (b, a, h, d_h)
         else:
-            q_t = q.transpose(1, 2)
-            k_t = k.transpose(1, 2)
-            v_t = v.transpose(1, 2)
-            attn = torch.matmul(q_t, k_t.transpose(-2, -1)) * self.scale
-            attn = F.softmax(attn, dim=-1)
-            out = torch.matmul(attn, v_t).transpose(1, 2)
+            q_t = q.transpose(1, 2)  # (b, h, a, d_h)
+            k_t = k.transpose(1, 2)  # (b, h, a, d_h)
+            v_t = v.transpose(1, 2)  # (b, h, a, d_h)
+            attn = torch.matmul(q_t, k_t.transpose(-2, -1)) * self.scale  # (b, h, a, a)
+            attn = F.softmax(attn, dim=-1)  # (b, h, a, a)
+            out = torch.matmul(attn, v_t).transpose(1, 2)  # (b, a, h, d_h)
 
         out = out.to(input_dtype).reshape(  # type: ignore[union-attr]
             batch_size, n_atoms, -1
-        )
-        out = out * torch.sigmoid(self.gate_proj(x_input))
-        return self.out_proj(out)
+        )  # (b, a, d_model)
+        out = out * torch.sigmoid(self.gate_proj(x_input))  # (b, a, d_model)
+        return self.out_proj(out)  # (b, a, d_model)
 
 
 # ===========================================================================
@@ -1593,11 +1621,13 @@ class SWA3DRoPEAttention(nn.Module):
 
 
 def _rms_adaln_raw(x: Tensor, scale: Tensor, shift: Tensor) -> Tensor:
-    return F.rms_norm(x, (x.shape[-1],)) * (1 + scale) + shift
+    # x, scale, shift: broadcast-compatible arrays; normalize x final axis.
+    return F.rms_norm(x, (x.shape[-1],)) * (1 + scale) + shift  # broadcast(x.shape, scale.shape, shift.shape)
 
 
 def _gated_residual_raw(x: Tensor, gate: Tensor, y: Tensor) -> Tensor:
-    return x + gate * y
+    # x, gate, y: broadcast-compatible arrays.
+    return x + gate * y  # broadcast(x.shape, gate.shape, y.shape)
 
 
 class SWAAtomBlock(nn.Module):
@@ -1619,7 +1649,7 @@ class SWAAtomBlock(nn.Module):
         self.ffn_norm = nn.RMSNorm(d_atom, elementwise_affine=False)
 
         adaln_linear = nn.Linear(d_atom, 6 * d_atom, bias=False)
-        nn.init.zeros_(adaln_linear.weight)
+        nn.init.zeros_(adaln_linear.weight)  # (6 * d_atom, d_atom)
         self.adaln_modulation = nn.Sequential(nn.SiLU(), adaln_linear)
 
         self.attn = SWA3DRoPEAttention(d_atom, n_heads, half_window=half_window)
@@ -1631,19 +1661,20 @@ class SWAAtomBlock(nn.Module):
         )
 
     def forward(self, x: Tensor, c_l: Tensor, attention_params: tuple) -> Tensor:
-        mod = self.adaln_modulation(c_l)
+        # x: (b, a, d_atom); c_l: (b, d_atom) or (b, a, d_atom).
+        mod = self.adaln_modulation(c_l)  # c_l.shape[:-1] + (6 * d_atom,)
         if mod.dim() == 2:
-            mod = mod.unsqueeze(1)
-        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.chunk(6, dim=-1)
+            mod = mod.unsqueeze(1)  # (b, 1, 6 * d_atom)
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.chunk(6, dim=-1)  # each (b, 1 or a, d_atom)
 
-        attn_input = self._rms_adaln(x, scale_a, shift_a)
-        attn_out = self.attn(attn_input, attention_params)
-        x = self._gated_residual(x, gate_a, attn_out)
+        attn_input = self._rms_adaln(x, scale_a, shift_a)  # (b, a, d_atom)
+        attn_out = self.attn(attn_input, attention_params)  # (b, a, d_atom)
+        x = self._gated_residual(x, gate_a, attn_out)  # (b, a, d_atom)
 
-        ffn_input = self._rms_adaln(x, scale_f, shift_f)
-        ffn_out = self.ffn(ffn_input)
-        x = self._gated_residual(x, gate_f, ffn_out)
-        return x
+        ffn_input = self._rms_adaln(x, scale_f, shift_f)  # (b, a, d_atom)
+        ffn_out = self.ffn(ffn_input)  # (b, a, d_atom)
+        x = self._gated_residual(x, gate_f, ffn_out)  # (b, a, d_atom)
+        return x  # (b, a, d_atom)
 
 
 class SWAAtomTransformer(nn.Module):
@@ -1699,14 +1730,15 @@ class SWAAtomTransformer(nn.Module):
         attention_params: tuple,
         return_intermediates: bool = False,
     ) -> Tensor | tuple[Tensor, list[Tensor]]:
+        # q_l/c_l: (b, a, d_atom); each saved intermediate has the same shape.
         intermediates: list[Tensor] = []
         for block in self.blocks:
-            q_l = block(q_l, c_l, attention_params)
+            q_l = block(q_l, c_l, attention_params)  # (b, a, d_atom)
             if return_intermediates:
                 intermediates.append(q_l)
         if return_intermediates:
-            return q_l, intermediates
-        return q_l
+            return q_l, intermediates  # q_l: (b, a, d_atom); optional list contains tensors of that shape
+        return q_l  # q_l: (b, a, d_atom); optional list contains tensors of that shape
 
 
 # ===========================================================================
@@ -1721,13 +1753,14 @@ def _prepare_atom_encoder_metadata(
     num_diffusion_samples: int,
 ) -> tuple[Tensor, Tensor, Tensor, int, int]:
     """Prepare mask-derived atom metadata outside compiled diffusion graphs."""
-    mask_exp = atom_attention_mask.repeat_interleave(num_diffusion_samples, 0)
-    seqlens = mask_exp.sum(dim=-1, dtype=torch.int32)
-    indices = torch.nonzero(mask_exp.flatten(), as_tuple=False).flatten()
+    # atom_attention_mask/atom_to_token: (b, a); bs = b * num_diffusion_samples.
+    mask_exp = atom_attention_mask.repeat_interleave(num_diffusion_samples, 0)  # (bs, a)
+    seqlens = mask_exp.sum(dim=-1, dtype=torch.int32)  # (bs,)
+    indices = torch.nonzero(mask_exp.flatten(), as_tuple=False).flatten()  # (n_present,)
     max_seqlen = int(seqlens.max().item())
-    cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
+    cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))  # (bs + 1,)
     n_tokens = int(atom_to_token.max().item()) + 1
-    return mask_exp, indices, cu_seqlens, max_seqlen, n_tokens
+    return mask_exp, indices, cu_seqlens, max_seqlen, n_tokens  # (bs, a), (n_present,), (bs + 1,), scalar integers
 
 
 class ESMFold2AtomEncoder(nn.Module):
@@ -1805,6 +1838,8 @@ class ESMFold2AtomEncoder(nn.Module):
         ``inference_cache`` caches step-invariant tensors (c_base, 3D RoPE,
         attention indices, n_tokens) across diffusion steps.
         """
+        # ref_pos: (b, a, 3); ref_element: (b, a, 128); chars: (b, a, 4, 64); other atom inputs: (b, a).
+        # bs = b * samples; d_out = d_token for structure prediction, otherwise d_token / 2.
         batch_size, n_atoms = ref_pos.shape[:2]
 
         layer_cache = None
@@ -1821,46 +1856,46 @@ class ESMFold2AtomEncoder(nn.Module):
                     ref_atom_name_chars.reshape(batch_size, n_atoms, MAX_CHARS * CHAR_VOCAB_SIZE),
                 ],
                 dim=-1,
-            )
-            c_base = self.atom_norm(self.atom_linear(atom_feats))
-            cos, sin = self.atom_transformer._build_3d_rope(ref_pos, ref_space_uid)
-            cos = cos.repeat_interleave(num_diffusion_samples, 0)
-            sin = sin.repeat_interleave(num_diffusion_samples, 0)
+            )  # (b, a, d_atom_features)
+            c_base = self.atom_norm(self.atom_linear(atom_feats))  # (b, a, d_atom)
+            cos, sin = self.atom_transformer._build_3d_rope(ref_pos, ref_space_uid)  # each (b, a, rotary_pairs)
+            cos = cos.repeat_interleave(num_diffusion_samples, 0)  # (bs, a, rotary_pairs)
+            sin = sin.repeat_interleave(num_diffusion_samples, 0)  # (bs, a, rotary_pairs)
             mask_exp, indices, cu_seqlens, max_seqlen, n_tokens = (
                 _prepare_atom_encoder_metadata(
                     atom_attention_mask,
                     atom_to_token,
                     num_diffusion_samples,
                 )
-            )
+            )  # (bs, a), (n_present,), (bs + 1,), integers
             attention_params = (cos, sin, indices, cu_seqlens, max_seqlen)
             if layer_cache is not None:
-                layer_cache["c_base"] = c_base
+                layer_cache["c_base"] = c_base  # (b, a, d_atom)
                 layer_cache["attention_params"] = attention_params
-                layer_cache["mask_exp"] = mask_exp
+                layer_cache["mask_exp"] = mask_exp  # (bs, a)
                 layer_cache["n_tokens"] = n_tokens
                 layer_cache["atom_to_token_exp"] = atom_to_token.repeat_interleave(
                     num_diffusion_samples, 0
-                )
+                )  # (bs, a)
         else:
-            c_base = layer_cache["c_base"]
+            c_base = layer_cache["c_base"]  # (b, a, d_atom)
             attention_params = layer_cache["attention_params"]
-            mask_exp = layer_cache["mask_exp"]
+            mask_exp = layer_cache["mask_exp"]  # (bs, a)
             n_tokens = layer_cache["n_tokens"]
 
-        c = c_base
+        c = c_base  # (b, a, d_atom)
 
-        q = c
+        q = c  # (b, a, d_atom)
 
         if self.structure_prediction and r_l is not None:
-            q = q.repeat_interleave(num_diffusion_samples, 0)
+            q = q.repeat_interleave(num_diffusion_samples, 0)  # (bs, a, d_atom)
             if pred_r1 is None:
-                pred_r1 = torch.zeros_like(r_l)
-            r_input = torch.cat([r_l, pred_r1], dim=-1)
-            r_to_q = self.coords_linear(r_input)
-            q = q + r_to_q
+                pred_r1 = torch.zeros_like(r_l)  # r_l.shape = (bs, a, 3)
+            r_input = torch.cat([r_l, pred_r1], dim=-1)  # (bs, a, 6)
+            r_to_q = self.coords_linear(r_input)  # (bs, a, d_atom)
+            q = q + r_to_q  # (bs, a, d_atom)
 
-        c = c.repeat_interleave(num_diffusion_samples, 0)
+        c = c.repeat_interleave(num_diffusion_samples, 0)  # (bs, a, d_atom)
 
         result = self.atom_transformer(
             q_l=q,
@@ -1869,19 +1904,19 @@ class ESMFold2AtomEncoder(nn.Module):
             return_intermediates=return_intermediates,
         )
         if return_intermediates:
-            q, intermediates = result
+            q, intermediates = result  # q: (bs, a, d_atom); list of same-shaped tensors
         else:
-            q = result
+            q = result  # (bs, a, d_atom)
             intermediates = []
 
-        q_to_a = F.relu(self.atom_to_token_linear(q))
+        q_to_a = F.relu(self.atom_to_token_linear(q))  # (bs, a, d_out)
         if layer_cache is not None and "atom_to_token_exp" in layer_cache:
-            atom_to_token_exp = layer_cache["atom_to_token_exp"]
+            atom_to_token_exp = layer_cache["atom_to_token_exp"]  # (bs, a)
         else:
-            atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0)
-        a = scatter_atom_to_token(q_to_a, atom_to_token_exp, n_tokens, atom_mask=mask_exp.bool())
+            atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0)  # (bs, a)
+        a = scatter_atom_to_token(q_to_a, atom_to_token_exp, n_tokens, atom_mask=mask_exp.bool())  # (bs, n_tokens, d_out)
 
-        return a, q, c, attention_params, intermediates
+        return a, q, c, attention_params, intermediates  # a: (bs, l, d_out); q/c: (bs, a, d_atom); metadata tuple and intermediate list
 
 
 # ===========================================================================
@@ -1935,10 +1970,11 @@ class ESMFold2AtomDecoder(nn.Module):
         return_intermediates: bool = False,
     ) -> tuple[Tensor, list[Tensor]]:
         """Returns (r_update, intermediates)."""
-        atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0)
-        a_to_q = self.token_to_atom_linear(a_i)
-        a_to_q = gather_token_to_atom(a_to_q, atom_to_token_exp)
-        q_l = q_l + a_to_q
+        # a_i: (bs, l, d_token); q_l/c_l: (bs, a, d_atom); atom_to_token: (b, a); bs = b * samples.
+        atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0)  # (bs, a)
+        a_to_q = self.token_to_atom_linear(a_i)  # (bs, l, d_atom)
+        a_to_q = gather_token_to_atom(a_to_q, atom_to_token_exp)  # (bs, a, d_atom)
+        q_l = q_l + a_to_q  # (bs, a, d_atom)
 
         result = self.atom_transformer(
             q_l=q_l,
@@ -1947,13 +1983,13 @@ class ESMFold2AtomDecoder(nn.Module):
             return_intermediates=return_intermediates,
         )
         if return_intermediates:
-            q_l, intermediates = result
+            q_l, intermediates = result  # q_l: (bs, a, d_atom); list of same-shaped tensors
         else:
-            q_l = result
+            q_l = result  # (bs, a, d_atom)
             intermediates = []
 
-        r_l = self.output_linear(self.norm(q_l))
-        return r_l, intermediates
+        r_l = self.output_linear(self.norm(q_l))  # (bs, a, 3)
+        return r_l, intermediates  # r_l: (bs, a, 3); intermediate tensors: (bs, a, d_atom)
 
 
 # ===========================================================================
@@ -1983,8 +2019,8 @@ class AttentionPairBias(nn.Module):
             self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5)
             self.out_gate = nn.Linear(d_cond, d_model, bias=True)
             # adaln init: weight=0, bias=-2
-            nn.init.zeros_(self.out_gate.weight)
-            nn.init.constant_(self.out_gate.bias, -2.0)
+            nn.init.zeros_(self.out_gate.weight)  # (d_model, d_cond)
+            nn.init.constant_(self.out_gate.bias, -2.0)  # (d_model,)
         else:
             self.pre_norm = nn.LayerNorm(d_model, eps=1e-5)
 
@@ -2043,22 +2079,23 @@ class AttentionPairBias(nn.Module):
         conditions every denoising step on the same ``z``, so the PyTorch path
         projects the bias on the first step and reuses that tensor afterwards.
         """
+        # a: (bs, l, d_model); s: (bs, l, d_cond) or None; z: (b, l, l, d_pair) or (bs, l, l); bs = b * samples.
         bsz, n_queries, d_model = a.shape
 
-        x = self.adaln(a, s) if s is not None else self.pre_norm(a)
+        x = self.adaln(a, s) if s is not None else self.pre_norm(a)  # (bs, l, d_model)
 
         n_keys = x.shape[1]
-        q = self.q_proj(x).view(bsz, n_queries, self.num_heads, self.head_dim)
-        kv = self.kv_proj(x)
-        k, v = kv.chunk(2, dim=-1)
-        k = k.view(bsz, n_keys, self.num_heads, self.head_dim)
-        v = v.view(bsz, n_keys, self.num_heads, self.head_dim)
+        q = self.q_proj(x).view(bsz, n_queries, self.num_heads, self.head_dim)  # (bs, l, h, d_h)
+        kv = self.kv_proj(x)  # (bs, l, 2 * d_model)
+        k, v = kv.chunk(2, dim=-1)  # each (bs, l, d_model)
+        k = k.view(bsz, n_keys, self.num_heads, self.head_dim)  # (bs, l, h, d_h)
+        v = v.view(bsz, n_keys, self.num_heads, self.head_dim)  # (bs, l, h, d_h)
 
         use_fused_kernel = self._can_use_fused_pair_bias(z, n_queries, beta)
         use_cueq_kernel = not use_fused_kernel and self._can_use_cueq_pair_bias(z, n_queries, beta)
-        cached_pair_bias = None
+        cached_pair_bias = None  # (bs, l, l, h) or None
         if step_cache is not None and not use_fused_kernel and not use_cueq_kernel:
-            cached_pair_bias = step_cache.get("pair_bias")
+            cached_pair_bias = step_cache.get("pair_bias")  # (bs, l, l, h) or None
 
         # Expand z for num_diffusion_samples, unless its projection is already cached.
         if (
@@ -2073,21 +2110,21 @@ class AttentionPairBias(nn.Module):
             and attention_mask.shape[0] != bsz
             and num_diffusion_samples > 1
         ):
-            attention_mask = attention_mask.repeat_interleave(num_diffusion_samples, dim=0)
+            attention_mask = attention_mask.repeat_interleave(num_diffusion_samples, dim=0)  # (bs, l)
 
         if use_fused_kernel:
             kernel_mask = (
                 attention_mask
                 if attention_mask is not None
                 else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool)
-            )
-            pair_norm_w = self.pair_norm.weight
+            )  # (bs, l)
+            pair_norm_w = self.pair_norm.weight  # (d_pair,)
             pair_norm_b = (
                 self.pair_norm.bias
                 if self.pair_norm.bias is not None
                 else torch.zeros_like(pair_norm_w)
-            )
-            z_bf = z if z.dtype == torch.bfloat16 else z.to(torch.bfloat16)
+            )  # (d_pair,)
+            z_bf = z if z.dtype == torch.bfloat16 else z.to(torch.bfloat16)  # (bs, l, l, d_pair)
             bias = _fused_pair_bias(  # type: ignore[misc]
                 z_bf,
                 kernel_mask,
@@ -2095,26 +2132,26 @@ class AttentionPairBias(nn.Module):
                 num_heads=self.num_heads,
                 pair_norm_w=pair_norm_w,
                 pair_norm_b=pair_norm_b,
-            )  # A has shape (b, h, q, k).
-            q_bhqd = q.transpose(1, 2)
-            k_bhqd = k.transpose(1, 2)
-            v_bhqd = v.transpose(1, 2)
+            )  # bias: (bs, h, l, l)
+            q_bhqd = q.transpose(1, 2)  # (bs, h, l, d_h)
+            k_bhqd = k.transpose(1, 2)  # (bs, h, l, d_h)
+            v_bhqd = v.transpose(1, 2)  # (bs, h, l, d_h)
             attn_out = F.scaled_dot_product_attention(
                 q_bhqd, k_bhqd, v_bhqd, attn_mask=bias.to(q_bhqd.dtype)
-            )
-            g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim)
-            ctx = g * attn_out.transpose(1, 2)
-            out = self.out_proj(ctx.reshape(bsz, n_queries, d_model))
+            )  # (bs, h, l, d_h)
+            g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim)  # (bs, l, h, d_h)
+            ctx = g * attn_out.transpose(1, 2)  # (bs, l, h, d_h)
+            out = self.out_proj(ctx.reshape(bsz, n_queries, d_model))  # (bs, l, d_model)
             if s is not None:
-                out = torch.sigmoid(self.out_gate(s)) * out
-            return out
+                out = torch.sigmoid(self.out_gate(s)) * out  # (bs, l, d_model)
+            return out  # (bs, l, d_model)
 
         if use_cueq_kernel:
             kernel_mask = (
                 attention_mask
                 if attention_mask is not None
                 else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool)
-            )
+            )  # (bs, l)
             out, _ = _cue_attn_pair_bias(  # type: ignore[misc]
                 s=x,
                 q=q.transpose(1, 2),
@@ -2130,36 +2167,36 @@ class AttentionPairBias(nn.Module):
                 b_ln_z=self.pair_norm.bias,
                 return_z_proj=False,
                 is_cached_z_proj=False,
-            )
+            )  # out: (bs, l, d_model); unused kernel auxiliary
         else:
             # Standard attention with pair bias
-            g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim)
+            g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim)  # (bs, l, h, d_h)
 
-            logits = torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale
+            logits = torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale  # (bs, l, l, h)
 
             if cached_pair_bias is not None:
                 pair_bias = cached_pair_bias  # (b * samples, n, n, h)
             elif z.dim() == 4:
                 pair_bias = self.pair_bias_proj(self.pair_norm(z))  # (b * samples, n, n, h)
                 if step_cache is not None:
-                    step_cache["pair_bias"] = pair_bias
+                    step_cache["pair_bias"] = pair_bias  # (bs, l, l, h)
             else:
                 pair_bias = z.unsqueeze(-1)  # (b * samples, n, n, 1), a precomputed bias
-            logits = logits + pair_bias.to(dtype=logits.dtype)
+            logits = logits + pair_bias.to(dtype=logits.dtype)  # (bs, l, l, h)
 
             if attention_mask is not None:
                 min_val = torch.finfo(logits.dtype).min
-                mask_bias = torch.where(attention_mask.bool()[:, None, :, None], 0.0, min_val)
-                logits = logits + mask_bias.to(dtype=logits.dtype)
+                mask_bias = torch.where(attention_mask.bool()[:, None, :, None], 0.0, min_val)  # (bs, 1, l, 1)
+                logits = logits + mask_bias.to(dtype=logits.dtype)  # (bs, l, l, h)
 
-            attn = torch.softmax(logits, dim=-2).to(dtype=v.dtype)
-            ctx = torch.einsum("... i j h, ... j h d -> ... i h d", attn, v)
-            ctx = g * ctx
-            out = self.out_proj(ctx.reshape(bsz, n_queries, d_model))
+            attn = torch.softmax(logits, dim=-2).to(dtype=v.dtype)  # (bs, l, l, h)
+            ctx = torch.einsum("... i j h, ... j h d -> ... i h d", attn, v)  # (bs, l, h, d_h)
+            ctx = g * ctx  # (bs, l, h, d_h)
+            out = self.out_proj(ctx.reshape(bsz, n_queries, d_model))  # (bs, l, d_model)
 
         if s is not None:
-            out = torch.sigmoid(self.out_gate(s)) * out
-        return out
+            out = torch.sigmoid(self.out_gate(s)) * out  # (bs, l, d_model)
+        return out  # (bs, l, d_model)
 
 
 # ===========================================================================
@@ -2184,8 +2221,8 @@ class ConditionedTransitionBlock(nn.Module):
         if use_conditioning:
             self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5)
             self.output_gate = nn.Linear(d_cond, d_model, bias=True)
-            nn.init.zeros_(self.output_gate.weight)
-            nn.init.constant_(self.output_gate.bias, -2.0)
+            nn.init.zeros_(self.output_gate.weight)  # (d_model, d_cond)
+            nn.init.constant_(self.output_gate.bias, -2.0)  # (d_model,)
         else:
             self.pre_norm = nn.LayerNorm(d_model, eps=1e-5)
 
@@ -2193,15 +2230,16 @@ class ConditionedTransitionBlock(nn.Module):
         self.lin_out = nn.Linear(hidden, d_model, bias=False)
 
     def forward(self, a: Tensor, s: Tensor | None) -> Tensor:
-        x = self.adaln(a, s) if s is not None else self.pre_norm(a)
+        # a: (..., d_model); s: (..., d_cond) or None; d_hidden = lin_out.in_features.
+        x = self.adaln(a, s) if s is not None else self.pre_norm(a)  # (..., d_model)
 
-        swish_a, swish_b = self.lin_swish(x).chunk(2, dim=-1)
-        b = F.silu(swish_a) * swish_b
-        out = self.lin_out(b)
+        swish_a, swish_b = self.lin_swish(x).chunk(2, dim=-1)  # each (..., d_hidden)
+        b = F.silu(swish_a) * swish_b  # (..., d_hidden)
+        out = self.lin_out(b)  # (..., d_model)
 
         if s is not None:
-            out = torch.sigmoid(self.output_gate(s)) * out
-        return out
+            out = torch.sigmoid(self.output_gate(s)) * out  # (..., d_model)
+        return out  # (..., d_model)
 
 
 # ===========================================================================
@@ -2269,11 +2307,12 @@ class DiffusionTransformer(nn.Module):
         ``inference_cache`` must span only calls that share ``z``, as one
         ``sample`` call does; each block then keeps its pair bias across steps.
         """
+        # a: (bs, l, d_model); s: (bs, l, d_cond) or None; z follows AttentionPairBias contract.
         intermediates: list[Tensor] = []
         block_caches: dict[int, dict[str, Tensor]] | None = None
         if inference_cache is not None:
             block_caches = inference_cache.setdefault("token_pair_bias", {})
-        x = a
+        x = a  # (bs, l, d_model)
         for block_index, (attn, transition) in enumerate(
             zip(self.attn_blocks, self.transition_blocks, strict=True)
         ):
@@ -2286,11 +2325,11 @@ class DiffusionTransformer(nn.Module):
                 attention_mask=attention_mask,
                 num_diffusion_samples=num_diffusion_samples,
                 step_cache=step_cache,
-            )
-            x = x + transition(x, s)
+            )  # (bs, l, d_model)
+            x = x + transition(x, s)  # (bs, l, d_model)
             if return_intermediates:
                 intermediates.append(x)
-        return x, intermediates
+        return x, intermediates  # x: (bs, l, d_model); list of same-shaped intermediate tensors
 
 
 # ===========================================================================
@@ -2343,45 +2382,46 @@ class DiffusionConditioning(nn.Module):
         num_diffusion_samples: int = 1,
         inference_cache: dict[str, Tensor] | None = None,
     ) -> tuple[Tensor, Tensor]:
+        # z_trunk/relative_position_encoding: (b, l, l, c_z); s_inputs: (b or bs, l, c_s_inputs); bs = b * samples.
         sigma = self.sigma_data if sigma_data is None else float(sigma_data)
         base_batch = z_trunk.shape[0]
         target_batch = base_batch * num_diffusion_samples
 
         # z conditioning (cached across diffusion steps: independent of t_hat)
         if inference_cache is not None and "z" in inference_cache:
-            z = inference_cache["z"]
+            z = inference_cache["z"]  # (b, l, l, c_z)
         else:
-            z_rel = relative_position_encoding.to(dtype=torch.float32)
-            z = torch.cat([z_trunk.to(dtype=torch.float32), z_rel], dim=-1)
-            z = self.z_proj(self.z_input_norm(z))
+            z_rel = relative_position_encoding.to(dtype=torch.float32)  # (b, l, l, c_z)
+            z = torch.cat([z_trunk.to(dtype=torch.float32), z_rel], dim=-1)  # (b, l, l, 2 * c_z)
+            z = self.z_proj(self.z_input_norm(z))  # (b, l, l, c_z)
             with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                 for block in self.z_transitions:
-                    z = z + block(z)
+                    z = z + block(z)  # (b, l, l, c_z)
             if inference_cache is not None:
-                inference_cache["z"] = z
+                inference_cache["z"] = z  # (b, l, l, c_z)
 
         # s conditioning
-        s_inputs_eff = s_inputs
+        s_inputs_eff = s_inputs  # (b or bs, l, c_s_inputs)
         if s_inputs_eff.shape[0] != target_batch:
-            s_inputs_eff = s_inputs_eff.repeat_interleave(num_diffusion_samples, 0)
+            s_inputs_eff = s_inputs_eff.repeat_interleave(num_diffusion_samples, 0)  # (bs, l, c_s_inputs)
 
-        s = self.s_proj(self.s_input_norm(s_inputs_eff.to(dtype=torch.float32)))
+        s = self.s_proj(self.s_input_norm(s_inputs_eff.to(dtype=torch.float32)))  # (bs, l, c_s)
 
         # Noise embedding
-        t = torch.as_tensor(t_hat, dtype=torch.float32, device=s.device).reshape(-1)
+        t = torch.as_tensor(t_hat, dtype=torch.float32, device=s.device).reshape(-1)  # (t_hat.numel(),)
         if t.numel() == 1:
-            t = t.expand(target_batch)
+            t = t.expand(target_batch)  # (bs,)
         elif t.shape[0] != target_batch:
-            t = t.repeat_interleave(num_diffusion_samples, 0)
-        t_noise = 0.25 * torch.log((t / sigma).clamp(min=1e-20))
-        n = self.fourier(t_noise)
-        n = self.noise_proj(self.noise_norm(n))
-        s = s + n.unsqueeze(1)
+            t = t.repeat_interleave(num_diffusion_samples, 0)  # (bs,)
+        t_noise = 0.25 * torch.log((t / sigma).clamp(min=1e-20))  # (bs,)
+        n = self.fourier(t_noise)  # (bs, fourier_dim)
+        n = self.noise_proj(self.noise_norm(n))  # (bs, c_s)
+        s = s + n.unsqueeze(1)  # (bs, l, c_s)
 
         for block in self.s_transitions:
-            s = s + block(s)
+            s = s + block(s)  # (bs, l, c_s)
 
-        return s, z
+        return s, z  # s: (bs, l, c_s); z: (b, l, l, c_z)
 
 
 # ===========================================================================
@@ -2453,7 +2493,7 @@ class DiffusionModule(nn.Module):
         )
 
         self.s_to_token = nn.Linear(c_token, c_token, bias=False)
-        nn.init.zeros_(self.s_to_token.weight)
+        nn.init.zeros_(self.s_to_token.weight)  # (c_token, c_token)
 
         # Token transformer (DiffusionTransformer with pair bias)
         self.token_transformer = DiffusionTransformer(
@@ -2499,11 +2539,12 @@ class DiffusionModule(nn.Module):
         return_atom_repr: bool = False,
         inference_cache: dict[str, Tensor] | None = None,
     ) -> dict[str, Tensor | None]:
+        # x_noisy: (bs, a, 3); ref/ID tensors retain base batch b; bs = b * samples; l tokens.
         bsz = x_noisy.shape[0]
         sigma = self.sigma_data if sigma_data is None else float(sigma_data)
-        t = torch.as_tensor(t_hat, dtype=torch.float32, device=x_noisy.device).reshape(-1)
+        t = torch.as_tensor(t_hat, dtype=torch.float32, device=x_noisy.device).reshape(-1)  # (t_hat.numel(),)
         if t.numel() == 1:
-            t = t.expand(bsz)
+            t = t.expand(bsz)  # (bs,)
 
         # Step 1: conditioning (pair z is cached across diffusion steps)
         s, z = self.conditioning(
@@ -2515,11 +2556,11 @@ class DiffusionModule(nn.Module):
             sigma_data=sigma,
             num_diffusion_samples=num_diffusion_samples,
             inference_cache=inference_cache,
-        )
+        )  # (bs, l, c_token), (b, l, l, c_z)
 
         # Step 2: normalize noisy coords
-        denom = torch.sqrt(t * t + sigma * sigma)
-        r_noisy = x_noisy / denom[:, None, None]
+        denom = torch.sqrt(t * t + sigma * sigma)  # (bs,)
+        r_noisy = x_noisy / denom[:, None, None]  # (bs, a, 3)
 
         # Step 3: atom encoder
         a, q_skip, c_skip, p_skip, enc_intermediates = self.atom_encoder(
@@ -2535,10 +2576,10 @@ class DiffusionModule(nn.Module):
             num_diffusion_samples=num_diffusion_samples,
             return_intermediates=return_atom_repr,
             inference_cache=inference_cache,
-        )
+        )  # a: (bs, l, c_token); q/c: (bs, a, c_atom); metadata and atom intermediate list
 
         # Step 4: add conditioned s
-        a = a + self.s_to_token(self.s_step_norm(s))
+        a = a + self.s_to_token(self.s_step_norm(s))  # (bs, l, c_token)
 
         # Step 5: token transformer
         a, _ = self.token_transformer(
@@ -2549,10 +2590,10 @@ class DiffusionModule(nn.Module):
             attention_mask=token_attention_mask,
             num_diffusion_samples=num_diffusion_samples,
             inference_cache=inference_cache,
-        )
+        )  # a: (bs, l, c_token); unused intermediate list
 
         # Step 6: token norm
-        a = self.token_norm(a)
+        a = self.token_norm(a)  # (bs, l, c_token)
 
         # Step 7: atom decoder
         r_update, dec_intermediates = self.atom_decoder(
@@ -2564,26 +2605,26 @@ class DiffusionModule(nn.Module):
             atom_attention_mask=ref_mask,
             num_diffusion_samples=num_diffusion_samples,
             return_intermediates=return_atom_repr,
-        )
+        )  # r_update: (bs, a, 3); atom intermediate list
 
         # Step 8: compute denoised output
         sigma2 = sigma * sigma
-        t2 = t * t
-        out = (sigma2 / (sigma2 + t2))[:, None, None] * x_noisy
-        out = out + ((sigma * t) / torch.sqrt(sigma2 + t2))[:, None, None] * r_update
+        t2 = t * t  # (bs,)
+        out = (sigma2 / (sigma2 + t2))[:, None, None] * x_noisy  # (bs, a, 3)
+        out = out + ((sigma * t) / torch.sqrt(sigma2 + t2))[:, None, None] * r_update  # (bs, a, 3)
 
         # Collect atom intermediates from encoder + decoder
-        atom_intermediates: Tensor | None = None
+        atom_intermediates: Tensor | None = None  # (bs, a, n_atom_blocks, c_atom) or None
         if return_atom_repr:
             all_ints = enc_intermediates + dec_intermediates
             if all_ints:
-                atom_intermediates = torch.stack(all_ints, dim=2)
+                atom_intermediates = torch.stack(all_ints, dim=2)  # (bs, a, n_atom_blocks, c_atom) or None
 
         return {
             "x_denoised": out,
             "token_repr": a if return_token_repr else None,
             "atom_intermediates": atom_intermediates,
-        }
+        }  # mapping: x_denoised (bs, a, 3); token_repr (bs, l, c_token) or None; atom intermediates as above
 
 
 # ===========================================================================
@@ -2647,24 +2688,24 @@ class DiffusionStructureHead(nn.Module):
                 [self.inference_s_max * self.sigma_data, 0.0],
                 device=device,
                 dtype=torch.float32,
-            )
+            )  # (steps + 1,)
         p = float(self.inference_p)
         inv_p = 1.0 / p
-        k = torch.arange(steps, device=device, dtype=torch.float32)
+        k = torch.arange(steps, device=device, dtype=torch.float32)  # (steps,)
         base = self.inference_s_max**inv_p + (k / (steps - 1)) * (
             self.inference_s_min**inv_p - self.inference_s_max**inv_p
-        )
-        schedule = self.sigma_data * base.pow(p)
-        return F.pad(schedule, (0, 1), value=0.0)
+        )  # (steps,)
+        schedule = self.sigma_data * base.pow(p)  # (steps,)
+        return F.pad(schedule, (0, 1), value=0.0)  # (steps + 1,)
 
     @staticmethod
     def _random_rotations(n: int, dtype: torch.dtype, device: torch.device) -> Tensor:
-        q = torch.randn((n, 4), dtype=dtype, device=device)
-        scale = torch.sqrt((q * q).sum(dim=1))
-        signs = torch.where(q[:, 0] < 0, -scale, scale)
-        q = q / signs[:, None]
-        r, i, j, k = torch.unbind(q, dim=-1)
-        two_s = 2.0 / (q * q).sum(dim=-1)
+        q = torch.randn((n, 4), dtype=dtype, device=device)  # (n, 4)
+        scale = torch.sqrt((q * q).sum(dim=1))  # (n,)
+        signs = torch.where(q[:, 0] < 0, -scale, scale)  # (n,)
+        q = q / signs[:, None]  # (n, 4)
+        r, i, j, k = torch.unbind(q, dim=-1)  # each (n,)
+        two_s = 2.0 / (q * q).sum(dim=-1)  # (n,)
         return torch.stack(
             (
                 1 - two_s * (j * j + k * k),
@@ -2678,51 +2719,53 @@ class DiffusionStructureHead(nn.Module):
                 1 - two_s * (i * i + j * j),
             ),
             dim=-1,
-        ).reshape(n, 3, 3)
+        ).reshape(n, 3, 3)  # (n, 3, 3)
 
     def _center_random_augmentation(
         self, x: Tensor, atom_mask: Tensor, second_coords: Tensor | None = None
     ) -> tuple[Tensor, Tensor | None]:
         """Algorithm 19: center + random rotation + translation."""
+        # x/second_coords: (b, a, 3); atom_mask: (b, a).
         bsz = x.shape[0]
         mask = atom_mask.unsqueeze(-1)  # M has shape (b, a, 1).
-        denom = mask.sum(dim=1, keepdim=True).clamp(min=1)
-        mean = (x * mask).sum(dim=1, keepdim=True) / denom
-        x = x - mean
+        denom = mask.sum(dim=1, keepdim=True).clamp(min=1)  # (b, 1, 1)
+        mean = (x * mask).sum(dim=1, keepdim=True) / denom  # (b, 1, 3)
+        x = x - mean  # (b, a, 3)
         if second_coords is not None:
-            second_coords = second_coords - mean
+            second_coords = second_coords - mean  # (b, a, 3)
 
-        r = self._random_rotations(bsz, x.dtype, x.device)
-        x = torch.einsum("bmd,bds->bms", x, r)
+        r = self._random_rotations(bsz, x.dtype, x.device)  # (b, 3, 3)
+        x = torch.einsum("bmd,bds->bms", x, r)  # (b, a, 3)
         if second_coords is not None:
-            second_coords = torch.einsum("bmd,bds->bms", second_coords, r)
+            second_coords = torch.einsum("bmd,bds->bms", second_coords, r)  # (b, a, 3)
 
-        t = torch.randn_like(x[:, 0:1, :])
-        x = x + t
+        t = torch.randn_like(x[:, 0:1, :])  # (b, 1, 3)
+        x = x + t  # (b, a, 3)
         if second_coords is not None:
-            second_coords = second_coords + t
-        return x, second_coords
+            second_coords = second_coords + t  # (b, a, 3)
+        return x, second_coords  # each (b, a, 3), or second_coords None
 
     @staticmethod
     def _weighted_rigid_align(x: Tensor, x_gt: Tensor, w: Tensor, mask: Tensor) -> Tensor:
         """Kabsch alignment: align x to x_gt with weights w."""
+        # x/x_gt: (b, n, 3); w/mask: (b, n).
         w = (mask * w).unsqueeze(-1)  # W has shape (b, n, 1).
-        denom = w.sum(dim=-2, keepdim=True).clamp(min=1e-8)
-        mu = (x * w).sum(dim=-2, keepdim=True) / denom
-        mu_gt = (x_gt * w).sum(dim=-2, keepdim=True) / denom
-        x_c = x - mu
-        xgt_c = x_gt - mu_gt
-        covariance = torch.einsum("bni,bnj->bij", w * xgt_c, x_c)
-        covariance_f32 = covariance.float()
+        denom = w.sum(dim=-2, keepdim=True).clamp(min=1e-8)  # (b, 1, 1)
+        mu = (x * w).sum(dim=-2, keepdim=True) / denom  # (b, 1, 3)
+        mu_gt = (x_gt * w).sum(dim=-2, keepdim=True) / denom  # (b, 1, 3)
+        x_c = x - mu  # (b, n, 3)
+        xgt_c = x_gt - mu_gt  # (b, n, 3)
+        covariance = torch.einsum("bni,bnj->bij", w * xgt_c, x_c)  # (b, 3, 3)
+        covariance_f32 = covariance.float()  # (b, 3, 3)
         u, _, vh = torch.linalg.svd(
             covariance_f32, driver="gesvd" if covariance_f32.is_cuda else None
-        )
-        det = torch.linalg.det(u @ vh)
-        ones = torch.ones_like(det)
+        )  # (b, 3, 3), (b, 3), (b, 3, 3)
+        det = torch.linalg.det(u @ vh)  # (b,)
+        ones = torch.ones_like(det)  # (b,)
         rotation = (u @ torch.diag_embed(torch.stack([ones, ones, det], dim=-1)) @ vh).to(
             covariance.dtype
-        )
-        return x_c @ rotation.transpose(-1, -2) + mu_gt
+        )  # (b, 3, 3)
+        return x_c @ rotation.transpose(-1, -2) + mu_gt  # (b, n, 3)
 
     # ------------------------------------------------------------------
     # Sampling
@@ -2766,6 +2809,7 @@ class DiffusionStructureHead(nn.Module):
         so we inflate the underlying schedule length here to land back at the
         requested step count post-truncation.
         """
+        # z_trunk: (b, l, l, c_z); s_inputs: (b, l, c_s_inputs); atom features: b by a; bs = b * samples.
         n_atoms = tok_idx.shape[1]
         device = s_inputs.device
         target_batch = s_inputs.shape[0] * num_diffusion_samples
@@ -2774,26 +2818,26 @@ class DiffusionStructureHead(nn.Module):
 
         steps = self.inference_num_steps if num_sampling_steps is None else int(num_sampling_steps)
 
-        schedule = self.inference_noise_schedule(steps, device)
+        schedule = self.inference_noise_schedule(steps, device)  # (steps + 1,)
         if max_inference_sigma is not None:
-            schedule = schedule[schedule <= float(max_inference_sigma)]
-            schedule = F.pad(schedule, (1, 0), value=float(max_inference_sigma))
+            schedule = schedule[schedule <= float(max_inference_sigma)]  # (n_below_cap,)
+            schedule = F.pad(schedule, (1, 0), value=float(max_inference_sigma))  # (n_below_cap + 1,)
 
         lam = self.noise_scale if noise_scale is None else float(noise_scale)
         eta = self.step_scale if step_scale is None else float(step_scale)
 
-        x = schedule[0] * torch.randn(target_batch, n_atoms, 3, device=device, dtype=torch.float32)
-        atom_mask = ref_mask.repeat_interleave(num_diffusion_samples, 0).float()
+        x = schedule[0] * torch.randn(target_batch, n_atoms, 3, device=device, dtype=torch.float32)  # (bs, a, 3)
+        atom_mask = ref_mask.repeat_interleave(num_diffusion_samples, 0).float()  # (bs, a)
 
         gammas = torch.where(
             schedule > self.gamma_min,
             torch.full_like(schedule, self.gamma_0),
             torch.zeros_like(schedule),
-        )
+        )  # schedule.shape
 
-        x_denoised_prev: Tensor | None = None
-        token_repr: Tensor | None = None
-        diff_atom_intermediates: Tensor | None = None
+        x_denoised_prev: Tensor | None = None  # (bs, a, 3) or None
+        token_repr: Tensor | None = None  # (bs, l, c_token) or None
+        diff_atom_intermediates: Tensor | None = None  # (bs, a, n_blocks, c_atom) or None
 
         step_pairs = list(zip(schedule[:-1], schedule[1:], gammas[1:], strict=True))
         num_steps = len(step_pairs)
@@ -2810,12 +2854,12 @@ class DiffusionStructureHead(nn.Module):
         for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(step_iterator):
             x, x_denoised_prev = self._center_random_augmentation(
                 x, atom_mask, second_coords=x_denoised_prev
-            )
+            )  # each (bs, a, 3), second may be None
 
             sigma_tm_val = float(sigma_tm.item())
             t_hat_val = sigma_tm_val * (1.0 + float(gamma.item()))
             eps_std = lam * max(t_hat_val**2 - sigma_tm_val**2, 0.0) ** 0.5
-            x_noisy = x + eps_std * torch.randn_like(x)
+            x_noisy = x + eps_std * torch.randn_like(x)  # (bs, a, 3)
 
             is_last_step = step_idx == num_steps - 1
             request_atom_repr = return_atom_repr and (
@@ -2846,24 +2890,24 @@ class DiffusionStructureHead(nn.Module):
                 return_token_repr=True,
                 return_atom_repr=request_atom_repr,
                 inference_cache=inference_cache,
-            )
+            )  # tensor mapping follows the called head's shape contract
 
-            x_denoised = dm_out["x_denoised"]
-            token_repr = dm_out["token_repr"]
+            x_denoised = dm_out["x_denoised"]  # (bs, a, 3)
+            token_repr = dm_out["token_repr"]  # (bs, l, c_token) or None
             if request_atom_repr:
-                diff_atom_intermediates = dm_out.get("atom_intermediates")
+                diff_atom_intermediates = dm_out.get("atom_intermediates")  # (bs, a, n_blocks, c_atom) or None
 
             # Reverse diffusion alignment (Kabsch)
             with torch.autocast(device_type="cuda", enabled=False):
                 x_noisy = self._weighted_rigid_align(
                     x_noisy.float(), x_denoised.float(), atom_mask, atom_mask
-                )
-            x_noisy = x_noisy.to(dtype=x_denoised.dtype)
+                )  # (bs, a, 3)
+            x_noisy = x_noisy.to(dtype=x_denoised.dtype)  # (bs, a, 3)
 
             # ODE/SDE step
             sigma_t_val = float(sigma_t.item())
-            denoised_over_sigma = (x_noisy - x_denoised) / t_hat_val
-            x = x_noisy + eta * (sigma_t_val - t_hat_val) * denoised_over_sigma
+            denoised_over_sigma = (x_noisy - x_denoised) / t_hat_val  # (bs, a, 3)
+            x = x_noisy + eta * (sigma_t_val - t_hat_val) * denoised_over_sigma  # (bs, a, 3)
 
             # Denoising early-exit: stop when consecutive predictions converge
             if (
@@ -2877,17 +2921,17 @@ class DiffusionStructureHead(nn.Module):
                         x_denoised.float(),
                         atom_mask,
                         atom_mask,
-                    )
-                diff = (x_denoised.float() - aligned) * atom_mask.unsqueeze(-1)
+                    )  # (bs, a, 3)
+                diff = (x_denoised.float() - aligned) * atom_mask.unsqueeze(-1)  # (bs, a, 3)
                 per_sample_rmsd = (
                     diff.pow(2).sum(dim=(-1, -2)) / atom_mask.sum(dim=-1).clamp(min=1)
-                ).sqrt()
+                ).sqrt()  # (bs,)
                 if per_sample_rmsd.max().item() < denoising_early_exit_rmsd:
-                    x = x_denoised
-                    x_denoised_prev = x_denoised
+                    x = x_denoised  # (bs, a, 3)
+                    x_denoised_prev = x_denoised  # (bs, a, 3) or None
                     break
 
-            x_denoised_prev = x_denoised
+            x_denoised_prev = x_denoised  # (bs, a, 3) or None
 
         result: dict[str, Tensor | None] = {
             "sample_atom_coords": x,
@@ -2895,4 +2939,4 @@ class DiffusionStructureHead(nn.Module):
         }
         if return_atom_repr:
             result["diff_atom_intermediates"] = diff_atom_intermediates
-        return result
+        return result  # coordinate/token/optional atom-intermediate mapping with the shapes above

fastplms/models/esmfold2/modeling_esmfold2_experimental.py

@@ -9,13 +9,13 @@ re-injection and a different confidence/MSA stack.
 from __future__ import annotations
 
 import gc
-from collections.abc import Mapping
-from pathlib import Path
-from typing import Any, ClassVar, cast
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any, ClassVar, cast
 from torch import Tensor
 from tqdm.auto import tqdm
 from transformers.modeling_utils import PreTrainedModel
@@ -66,6 +66,7 @@ from .modeling_esmfold2_common import (
     validate_prepared_auxiliary_inputs,
 )
 
+
 _EPS = 1e-5
 _NONPOLYMER_ID = 3
 
@@ -82,8 +83,8 @@ class ConfidenceHead(nn.Module):
         d_pair = config.d_pair
         d_inputs = config.inputs.d_inputs
 
-        boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1)
-        self.register_buffer("boundaries", boundaries)
+        boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1)  # (distogram_bins - 1,)
+        self.register_buffer("boundaries", boundaries)  # (distogram_bins - 1,)
         self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair)
 
         self.s_norm = nn.LayerNorm(d_single)
@@ -105,7 +106,7 @@ class ConfidenceHead(nn.Module):
         max_atoms_per_token = 23
         self.plddt_weight = nn.Parameter(
             torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins)
-        )
+        )  # (23, d_single, n_plddt_bins)
         self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False)
 
     def set_kernel_backend(self, backend: str | None) -> None:
@@ -117,16 +118,18 @@ class ConfidenceHead(nn.Module):
 
     @staticmethod
     def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor:
+        # x: (b, ...); output repeats the batch axis by samples.
         if num_diffusion_samples == 1:
-            return x
-        return x.repeat_interleave(num_diffusion_samples, 0)
+            return x  # (b * samples, ...), including samples = 1
+        return x.repeat_interleave(num_diffusion_samples, 0)  # (b * samples, ...), including samples = 1
 
     @staticmethod
     def _flatten_sample_axis(x: Tensor) -> Tensor:
+        # x: (b, samples, n, c) or an already flattened tensor.
         if x.ndim == 4:
             b, mult, n, c = x.shape
-            return x.reshape(b * mult, n, c)
-        return x
+            return x.reshape(b * mult, n, c)  # (b * samples, n, c) for 4D input; otherwise x.shape
+        return x  # (b * samples, n, c) for 4D input; otherwise x.shape
 
     def forward(
         self,
@@ -143,109 +146,110 @@ class ConfidenceHead(nn.Module):
         relative_position_encoding: Tensor | None = None,
         token_bonds_encoding: Tensor | None = None,
     ) -> dict[str, Tensor]:
-        s_inputs_normed = self.s_inputs_norm(s_inputs)
-        z_base = self.z_norm(z)
+        # s_inputs: (b, l, d_inputs); z: (b, l, l, d_pair); x_pred: (bs, a, 3) or (b, samples, a, 3). bs = b * samples.
+        s_inputs_normed = self.s_inputs_norm(s_inputs)  # (b, l, d_inputs)
+        z_base = self.z_norm(z)  # (b, l, l, d_pair)
         if relative_position_encoding is not None:
-            z_base = z_base + relative_position_encoding
+            z_base = z_base + relative_position_encoding  # (b, l, l, d_pair)
         if token_bonds_encoding is not None:
-            z_base = z_base + token_bonds_encoding
-        z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2)
-        z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1)
+            z_base = z_base + token_bonds_encoding  # (b, l, l, d_pair)
+        z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2)  # (b, l, l, d_pair)
+        z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1)  # (b, l, l, d_pair)
         z_base = z_base + self.s_to_z_prod_out(
             self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :]
             * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :]
-        )
-
-        pair = self._repeat_batch(z_base, num_diffusion_samples)
-        x_pred_flat = self._flatten_sample_axis(x_pred)
-        atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples)
-        atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples)
-        rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long()
-        mask = self._repeat_batch(token_attention_mask, num_diffusion_samples)
+        )  # (b, l, l, d_pair)
+
+        pair = self._repeat_batch(z_base, num_diffusion_samples)  # (bs, l, l, d_pair)
+        x_pred_flat = self._flatten_sample_axis(x_pred)  # (bs, a, 3)
+        atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples)  # (bs, a)
+        atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples)  # (bs, a)
+        rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long()  # (bs, l)
+        mask = self._repeat_batch(token_attention_mask, num_diffusion_samples)  # (bs, l)
         batch_mult = pair.shape[0]
 
-        rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m)
+        rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m)  # (bs, l, 3)
         rep_distances = torch.cdist(
             rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist"
-        )
-        distogram_bins = (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()
-        pair = pair + self.dist_bin_pairwise_embed(distogram_bins)
-
-        pair_mask = mask[:, :, None].float() * mask[:, None, :].float()
-        pair = pair + self.folding_trunk(pair, pair_attention_mask=pair_mask)
-        single = self.row_attention_pooling(pair, mask)
-
-        atom_mask_f = atom_mask_m.float()
-        s_at_atoms = gather_token_to_atom(single, atom_to_token_m)
-        s_at_atoms = self.plddt_ln(s_at_atoms)
-        intra_idx = _compute_intra_token_idx(atom_to_token_m)
-        intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1)
-        plddt_weight = self.plddt_weight[intra_idx]
-        plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms, plddt_weight)
-        plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0)
+        )  # (bs, l, l)
+        distogram_bins = (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()  # (bs, l, l)
+        pair = pair + self.dist_bin_pairwise_embed(distogram_bins)  # (bs, l, l, d_pair)
+
+        pair_mask = mask[:, :, None].float() * mask[:, None, :].float()  # (bs, l, l)
+        pair = pair + self.folding_trunk(pair, pair_attention_mask=pair_mask)  # (bs, l, l, d_pair)
+        single = self.row_attention_pooling(pair, mask)  # (bs, l, d_single)
+
+        atom_mask_f = atom_mask_m.float()  # (bs, a)
+        s_at_atoms = gather_token_to_atom(single, atom_to_token_m)  # (bs, a, d_single)
+        s_at_atoms = self.plddt_ln(s_at_atoms)  # (bs, a, d_single)
+        intra_idx = _compute_intra_token_idx(atom_to_token_m)  # (bs, a)
+        intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1)  # (bs, a)
+        plddt_weight = self.plddt_weight[intra_idx]  # (bs, a, d_single, n_plddt_bins)
+        plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms, plddt_weight)  # (bs, a, n_plddt_bins)
+        plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0)  # (bs, a)
 
         length = single.shape[1]
         plddt_sum = torch.zeros(
             batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype
-        )
+        )  # (bs, l)
         atom_count = torch.zeros(
             batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype
-        )
-        atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype)
-        plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t)
-        atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t)
-        plddt = plddt_sum / atom_count.clamp(min=1e-6)
+        )  # (bs, l)
+        atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype)  # (bs, a)
+        plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t)  # (bs, l)
+        atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t)  # (bs, l)
+        plddt = plddt_sum / atom_count.clamp(min=1e-6)  # (bs, l)
 
         complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / (
             atom_mask_f.sum(dim=-1) + _EPS
-        )
+        )  # (bs,)
 
-        expanded_type = self._repeat_batch(mol_type, num_diffusion_samples)
-        expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples)
-        is_ligand = (expanded_type == _NONPOLYMER_ID).float()
-        inter_chain = (expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)).float()
-        near_contact = (rep_distances < 8).float()
+        expanded_type = self._repeat_batch(mol_type, num_diffusion_samples)  # (bs, l)
+        expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples)  # (bs, l)
+        is_ligand = (expanded_type == _NONPOLYMER_ID).float()  # (bs, l)
+        inter_chain = (expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)).float()  # (bs, l, l)
+        near_contact = (rep_distances < 8).float()  # (bs, l, l)
         interface_per_token = (near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)).amax(
             dim=-1
-        )
+        )  # (bs, l)
         iplddt_weight = torch.where(
             is_ligand.bool(),
             torch.full_like(interface_per_token, 2.0),
             interface_per_token,
-        )
+        )  # (bs, l)
         iplddt_weight_atoms = gather_token_to_atom(
             iplddt_weight.unsqueeze(-1), atom_to_token_m
-        ).squeeze(-1)
-        atom_iplddt_w = atom_mask_f * iplddt_weight_atoms
+        ).squeeze(-1)  # (bs, a)
+        atom_iplddt_w = atom_mask_f * iplddt_weight_atoms  # (bs, a)
         complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / (
             atom_iplddt_w.sum(dim=-1) + _EPS
-        )
-        plddt_ca = plddt_per_atom.gather(1, rep_idx_m)
+        )  # (bs,)
+        plddt_ca = plddt_per_atom.gather(1, rep_idx_m)  # (bs, l)
 
-        pae_logits = self.pae_head(pair)
-        pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach()
+        pae_logits = self.pae_head(pair)  # (bs, l, l, n_pae_bins)
+        pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach()  # (bs, l, l)
 
         n_bins = pae_logits.shape[-1]
         bin_width = 32.0 / n_bins
-        bin_centers = torch.arange(0.5 * bin_width, 32.0, bin_width, device=pae_logits.device)
-        mask_f = mask.float()
-        n_res = mask_f.sum(dim=-1, keepdim=True)
-        d0 = 1.24 * (n_res.clamp(min=19) - 15) ** (1 / 3) - 1.8
-        tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2)
-        pae_probs = F.softmax(pae_logits, dim=-1)
-        tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1)
-
-        pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2)
-        ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (pair_mask_2d.sum(dim=-1) + _EPS)
-        ptm = ptm_per_row.max(dim=-1).values
+        bin_centers = torch.arange(0.5 * bin_width, 32.0, bin_width, device=pae_logits.device)  # (n_pae_bins,)
+        mask_f = mask.float()  # (bs, l)
+        n_res = mask_f.sum(dim=-1, keepdim=True)  # (bs, 1)
+        d0 = 1.24 * (n_res.clamp(min=19) - 15) ** (1 / 3) - 1.8  # (bs, 1)
+        tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2)  # (bs, n_pae_bins)
+        pae_probs = F.softmax(pae_logits, dim=-1)  # (bs, l, l, n_pae_bins)
+        tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1)  # (bs, l, l)
+
+        pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2)  # (bs, l, l)
+        ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (pair_mask_2d.sum(dim=-1) + _EPS)  # (bs, l)
+        ptm = ptm_per_row.max(dim=-1).values  # (bs,)
 
         inter_chain_mask = (
             expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
-        ).float() * pair_mask_2d
+        ).float() * pair_mask_2d  # (bs, l, l)
         iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / (
             inter_chain_mask.sum(dim=-1) + _EPS
-        )
-        iptm = iptm_per_row.max(dim=-1).values
+        )  # (bs, l)
+        iptm = iptm_per_row.max(dim=-1).values  # (bs,)
 
         max_chain_id = int(expanded_asym.max().item()) if batch_mult > 0 else 0
         n_chains = max_chain_id + 1
@@ -255,16 +259,16 @@ class ConfidenceHead(nn.Module):
             n_chains,
             device=tm_expected.device,
             dtype=tm_expected.dtype,
-        )
+        )  # (bs, n_chains, n_chains)
         for c1 in range(n_chains):
-            chain_c1 = (expanded_asym == c1).float() * mask_f
+            chain_c1 = (expanded_asym == c1).float() * mask_f  # (bs, l)
             if chain_c1.sum() == 0:
                 continue
             for c2 in range(n_chains):
-                chain_c2 = (expanded_asym == c2).float() * mask_f
-                pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2)
-                denom = pair_m.sum(dim=(-1, -2)) + _EPS
-                pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(dim=(-1, -2)) / denom
+                chain_c2 = (expanded_asym == c2).float() * mask_f  # (bs, l)
+                pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2)  # (bs, l, l)
+                denom = pair_m.sum(dim=(-1, -2)) + _EPS  # (bs,)
+                pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(dim=(-1, -2)) / denom  # (bs,)
 
         return {
             "plddt_logits": plddt_logits,
@@ -278,7 +282,7 @@ class ConfidenceHead(nn.Module):
             "ptm": ptm.detach(),
             "iptm": iptm.detach(),
             "pair_chains_iptm": pair_chains_iptm.detach(),
-        }
+        }  # mapping of confidence tensors with shapes traced above
 
 
 class _TransitionFFN(nn.Module):
@@ -288,7 +292,8 @@ class _TransitionFFN(nn.Module):
         self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False)
 
     def forward(self, x: Tensor) -> Tensor:
-        return self.ffn(self.norm(x))
+        # x: (..., d_model).
+        return self.ffn(self.norm(x))  # x.shape
 
 
 class MSAEncoderBlock(nn.Module):
@@ -327,44 +332,45 @@ class MSAEncoderBlock(nn.Module):
         pair_attention_mask: Tensor,
         msa_track_mask: Tensor | None = None,
     ) -> tuple[Tensor, Tensor]:
+        # msa_repr: (b, l, m, d_msa); pair_repr: (b, l, l, d_pair); msa_track_mask: (b,) or None.
         mask4d = (
             msa_track_mask[:, None, None, None].to(dtype=msa_repr.dtype)
             if msa_track_mask is not None
             else None
-        )
+        )  # (b, 1, 1, 1) or None
 
-        pair_mask4d = mask4d[:, :, :1] if mask4d is not None else None
+        pair_mask4d = mask4d[:, :, :1] if mask4d is not None else None  # (b, 1, 1, 1) or None
 
-        msa_update = self.msa_pair_weighted_averaging(msa_repr, pair_repr, pair_attention_mask)
+        msa_update = self.msa_pair_weighted_averaging(msa_repr, pair_repr, pair_attention_mask)  # (b, l, m, d_msa)
         if mask4d is not None:
-            msa_update = msa_update * mask4d
-        msa_repr = msa_repr + msa_update
+            msa_update = msa_update * mask4d  # (b, l, m, d_msa)
+        msa_repr = msa_repr + msa_update  # (b, l, m, d_msa)
 
-        msa_transition = self.msa_transition(msa_repr)
+        msa_transition = self.msa_transition(msa_repr)  # (b, l, m, d_msa)
         if mask4d is not None:
-            msa_transition = msa_transition * mask4d
-        msa_repr = msa_repr + msa_transition
+            msa_transition = msa_transition * mask4d  # (b, l, m, d_msa)
+        msa_repr = msa_repr + msa_transition  # (b, l, m, d_msa)
 
-        pair_opm = self.outer_product_mean(msa_repr, msa_attention_mask)
+        pair_opm = self.outer_product_mean(msa_repr, msa_attention_mask)  # (b, l, l, d_pair)
         if pair_mask4d is not None:
-            pair_opm = pair_opm * pair_mask4d
-        pair_repr = pair_repr + pair_opm
+            pair_opm = pair_opm * pair_mask4d  # (b, l, l, d_pair)
+        pair_repr = pair_repr + pair_opm  # (b, l, l, d_pair)
 
-        pair_out = self.tri_mul_out(pair_repr, mask=pair_attention_mask)
+        pair_out = self.tri_mul_out(pair_repr, mask=pair_attention_mask)  # (b, l, l, d_pair)
         if pair_mask4d is not None:
-            pair_out = pair_out * pair_mask4d
-        pair_repr = pair_repr + pair_out
+            pair_out = pair_out * pair_mask4d  # (b, l, l, d_pair)
+        pair_repr = pair_repr + pair_out  # (b, l, l, d_pair)
 
-        pair_in = self.tri_mul_in(pair_repr, mask=pair_attention_mask)
+        pair_in = self.tri_mul_in(pair_repr, mask=pair_attention_mask)  # (b, l, l, d_pair)
         if pair_mask4d is not None:
-            pair_in = pair_in * pair_mask4d
-        pair_repr = pair_repr + pair_in
+            pair_in = pair_in * pair_mask4d  # (b, l, l, d_pair)
+        pair_repr = pair_repr + pair_in  # (b, l, l, d_pair)
 
-        pair_transition = self.pair_transition(pair_repr)
+        pair_transition = self.pair_transition(pair_repr)  # (b, l, l, d_pair)
         if pair_mask4d is not None:
-            pair_transition = pair_transition * pair_mask4d
-        pair_repr = pair_repr + pair_transition
-        return msa_repr, pair_repr
+            pair_transition = pair_transition * pair_mask4d  # (b, l, l, d_pair)
+        pair_repr = pair_repr + pair_transition  # (b, l, l, d_pair)
+        return msa_repr, pair_repr  # (b, l, m, d_msa), (b, l, l, d_pair)
 
 
 class MSAEncoder(nn.Module):
@@ -407,18 +413,19 @@ class MSAEncoder(nn.Module):
         deletion_value: Tensor,
         msa_attention_mask: Tensor,
     ) -> Tensor:
+        # x_pair: (b, l, l, d_pair); x_inputs: (b, l, d_inputs); MSA features: (b, l, m, 33), deletion/mask: (b, l, m).
         batch_size, _, depth = msa_attention_mask.shape
         m_feat = torch.cat(
             [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)],
             dim=-1,
-        )
-        m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2)
+        )  # (b, l, m, 35)
+        m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2)  # (b, l, m, d_msa)
         if depth > 1:
-            msa_track_mask = msa_attention_mask[:, :, 1:].any(dim=(1, 2))
+            msa_track_mask = msa_attention_mask[:, :, 1:].any(dim=(1, 2))  # (b,)
         else:
-            msa_track_mask = torch.zeros(batch_size, dtype=torch.bool, device=x_pair.device)
-        tok_mask = msa_attention_mask[:, :, 0]
-        pair_attention_mask = tok_mask.unsqueeze(2) * tok_mask.unsqueeze(1)
+            msa_track_mask = torch.zeros(batch_size, dtype=torch.bool, device=x_pair.device)  # (b,)
+        tok_mask = msa_attention_mask[:, :, 0]  # (b, l)
+        pair_attention_mask = tok_mask.unsqueeze(2) * tok_mask.unsqueeze(1)  # (b, l, l)
         for block in self.blocks:
             m, x_pair = cast(MSAEncoderBlock, block)(
                 m,
@@ -426,8 +433,8 @@ class MSAEncoder(nn.Module):
                 msa_attention_mask,
                 pair_attention_mask,
                 msa_track_mask,
-            )
-        return x_pair * msa_track_mask[:, None, None, None].to(dtype=x_pair.dtype)
+            )  # (b, l, m, d_msa), (b, l, l, d_pair)
+        return x_pair * msa_track_mask[:, None, None, None].to(dtype=x_pair.dtype)  # (b, l, l, d_pair)
 
 
 class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin, PreTrainedModel):
@@ -477,7 +484,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         self.pair_loop_proj = nn.Sequential(
             nn.LayerNorm(d_pair), nn.Linear(d_pair, d_pair, bias=False)
         )
-        nn.init.zeros_(cast(nn.Linear, self.pair_loop_proj[1]).weight)
+        nn.init.zeros_(cast(nn.Linear, self.pair_loop_proj[1]).weight)  # (d_pair, d_pair)
 
         self.structure_head = DiffusionStructureHead(config)
         self.distogram_head = nn.Linear(d_pair, config.structure_head.distogram_bins, bias=True)
@@ -648,6 +655,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         tok_mask: Tensor,
         verbose: bool = False,
     ) -> Tensor:
+        # Input tensors: (b, l); n_states and d_lm come from the loaded backbone.
         if self._esmc_fp8 and torch.is_grad_enabled():
             _reload_esmc_bf16_for_gradients(
                 self,
@@ -670,9 +678,9 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                         mol_type,
                         tok_mask,
                         pad_to_multiple=pad_to,
-                    )
+                    )  # (b, l, n_states, d_lm)
                     progress.update()
-                return result
+                return result  # (b, l, n_states, d_lm)
             return compute_lm_hidden_states(
                 self._esmc,
                 input_ids,
@@ -681,7 +689,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                 mol_type,
                 tok_mask,
                 pad_to_multiple=pad_to,
-            )
+            )  # (b, l, n_states, d_lm)
 
     def forward(
         self,
@@ -731,6 +739,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         disto_cond_mask: Tensor | None = None,
         verbose: bool = False,
     ) -> ESMFold2Output | tuple[Any, ...]:
+        # Token IDs/masks: (b, l); atom IDs/masks: (b, a); ref_pos: (b, a, 3); chars: (b, a, 4); MSA: (b, m, l); bs = b * samples.
         output_hidden_states, return_dict = _resolve_structure_output_controls(
             self.config,
             output_attentions=output_attentions,
@@ -751,8 +760,8 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             disto_cond_mask=disto_cond_mask,
         )
         del gt_coords, is_resolved, frames_idx
-        tok_mask = token_attention_mask
-        atm_mask = atom_attention_mask
+        tok_mask = token_attention_mask  # (b, l)
+        atm_mask = atom_attention_mask  # (b, a)
         n_loops = num_loops if num_loops is not None else self.config.num_loops
         n_samples = (
             num_diffusion_samples
@@ -761,46 +770,46 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         )
 
         if res_type.dim() == 2:
-            res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float()
-            res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float()
+            res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float()  # (b, l, 33)
+            res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float()  # (b, l, 33)
         else:
-            res_type_oh = res_type.float()
+            res_type_oh = res_type.float()  # (b, l, 33)
 
         if msa is not None:
-            msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float()
+            msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float()  # (b, m, l, 33)
             if msa_attention_mask is not None:
-                mask_f = msa_attention_mask.float().unsqueeze(-1)
-                msa_oh_profile = msa_oh_profile * mask_f
-                valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1)
-                profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1)
+                mask_f = msa_attention_mask.float().unsqueeze(-1)  # (b, m, l, 1)
+                msa_oh_profile = msa_oh_profile * mask_f  # (b, m, l, 33)
+                valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1)  # (b, l)
+                profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1)  # (b, l, 33)
             else:
-                profile = msa_oh_profile.mean(dim=1)
+                profile = msa_oh_profile.mean(dim=1)  # (b, l, 33)
         else:
-            profile = res_type_oh
+            profile = res_type_oh  # (b, l, 33)
 
         if res_type_soft is not None:
-            res_type_oh = res_type_soft.float()
+            res_type_oh = res_type_soft.float()  # (b, l, 33)
             if not self.config.disable_msa_features and provide_soft_sequence_to_msa_and_profile:
-                profile = res_type_oh
-                msa = res_type_oh.unsqueeze(1)
-                msa_attention_mask = tok_mask.unsqueeze(1)
+                profile = res_type_oh  # (b, l, 33)
+                msa = res_type_oh.unsqueeze(1)  # (b, 1, l, 33)
+                msa_attention_mask = tok_mask.unsqueeze(1)  # (b, 1, l)
 
         if deletion_mean is None:
             deletion_mean = torch.zeros(
                 res_type.shape[0], res_type.shape[1], device=res_type.device
-            )
+            )  # (b, l)
         if self.config.disable_msa_features:
-            profile = torch.zeros_like(profile)
-            deletion_mean = torch.zeros_like(deletion_mean)
+            profile = torch.zeros_like(profile)  # (b, l, 33)
+            deletion_mean = torch.zeros_like(deletion_mean)  # (b, l)
 
-        ref_element_oh = F.one_hot(ref_element.long(), num_classes=MAX_ATOMIC_NUMBER).float()
+        ref_element_oh = F.one_hot(ref_element.long(), num_classes=MAX_ATOMIC_NUMBER).float()  # (b, a, 128)
         ref_atom_name_chars_oh = F.one_hot(
             ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE
-        ).float()
-        atm_mask_f = atm_mask.float()
-        ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1)
-        ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(-1).unsqueeze(-1)
-        atom_to_token = atom_to_token * atm_mask.long()
+        ).float()  # (b, a, 4, 64)
+        atm_mask_f = atm_mask.float()  # (b, a)
+        ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1)  # (b, a, 128)
+        ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(-1).unsqueeze(-1)  # (b, a, 4, 64)
+        atom_to_token = atom_to_token * atm_mask.long()  # (b, a)
 
         use_amp = ref_pos.device.type == "cuda"
         with torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16):
@@ -815,58 +824,58 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                 ref_element=ref_element_oh,
                 ref_atom_name_chars=ref_atom_name_chars_oh,
                 atom_to_token=atom_to_token,
-            )
+            )  # (b, l, d_inputs)
 
-            z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(x_inputs).unsqueeze(1)
+            z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(x_inputs).unsqueeze(1)  # (b, l, l, d_pair)
             relative_position_encoding = self.rel_pos(
                 residue_index=residue_index,
                 asym_id=asym_id,
                 sym_id=sym_id,
                 entity_id=entity_id,
                 token_index=token_index,
-            )
-            token_bonds_encoding = self.token_bonds(token_bonds.float())
-            z_init = z_init + relative_position_encoding + token_bonds_encoding
+            )  # (b, l, l, d_pair)
+            token_bonds_encoding = self.token_bonds(token_bonds.float())  # (b, l, l, d_pair)
+            z_init = z_init + relative_position_encoding + token_bonds_encoding  # (b, l, l, d_pair)
 
             if lm_hidden_states is None and input_ids is not None and self._esmc is not None:
                 lm_hidden_states = self._compute_lm_hidden_states(
                     input_ids, asym_id, residue_index, mol_type, tok_mask, verbose=verbose
-                )
+                )  # (b, l, n_states, d_lm)
             if lm_hidden_states is not None:
                 lm_dropout = (
                     self.config.lm_dropout
                     if self.config.force_lm_dropout_during_inference or self.training
                     else 0.0
                 )
-                lm_z = self.language_model(lm_hidden_states.detach(), lm_dropout=lm_dropout)
-                z_init = z_init + lm_z.to(z_init.dtype)
+                lm_z = self.language_model(lm_hidden_states.detach(), lm_dropout=lm_dropout)  # (b, l, l, d_pair) or None
+                z_init = z_init + lm_z.to(z_init.dtype)  # (b, l, l, d_pair)
 
             msa_kwargs: dict[str, Tensor] | None = None
             if self.msa_encoder is not None and msa is not None:
                 if msa.dim() == 4:
                     batch_msa, depth, length_msa, _ = msa.shape
-                    msa_oh = msa.permute(0, 2, 1, 3).float()
+                    msa_oh = msa.permute(0, 2, 1, 3).float()  # (b, l, m, 33)
                 else:
                     batch_msa, depth, length_msa = msa.shape
                     msa_oh = F.one_hot(
                         msa.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES
-                    ).float()
+                    ).float()  # (b, l, m, 33)
                 msa_attn = (
                     msa_attention_mask.permute(0, 2, 1).float()
                     if msa_attention_mask is not None
                     else tok_mask[:, :, None].expand(-1, -1, depth).float()
-                )
-                msa_oh = msa_oh * msa_attn.unsqueeze(-1)
+                )  # (b, l, m)
+                msa_oh = msa_oh * msa_attn.unsqueeze(-1)  # (b, l, m, 33)
                 hd = (
                     has_deletion.permute(0, 2, 1).float()
                     if has_deletion is not None
                     else torch.zeros(batch_msa, length_msa, depth, device=msa.device)
-                )
+                )  # (b, l, m)
                 dv = (
                     deletion_value.permute(0, 2, 1).float()
                     if deletion_value is not None
                     else torch.zeros(batch_msa, length_msa, depth, device=msa.device)
-                )
+                )  # (b, l, m)
                 msa_kwargs = {
                     "x_inputs": x_inputs,
                     "msa_oh": msa_oh,
@@ -875,10 +884,10 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                     "msa_attention_mask": msa_attn,
                 }
 
-            pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float()
-            z = torch.zeros_like(z_init)
-            prev_pair: Tensor | None = None
-            prev_disto_probs: Tensor | None = None
+            pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float()  # (b, l, l)
+            z = torch.zeros_like(z_init)  # (b, l, l, d_pair)
+            prev_pair: Tensor | None = None  # (b, l, l, d_pair) or None
+            prev_disto_probs: Tensor | None = None  # (b, l, l, n_distogram_bins) or None
             loop_iterator = range(n_loops + 1)
             if verbose:
                 loop_iterator = tqdm(
@@ -889,32 +898,32 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                 )
 
             for loop_num in loop_iterator:
-                z = z_init + self.pair_loop_proj(z)
+                z = z_init + self.pair_loop_proj(z)  # (b, l, l, d_pair)
                 if msa_kwargs is not None and self.msa_encoder is not None:
-                    z = z + self.msa_encoder(x_pair=z, **msa_kwargs).to(z.dtype)
-                z = self.folding_trunk(z, pair_attention_mask=pair_mask)
+                    z = z + self.msa_encoder(x_pair=z, **msa_kwargs).to(z.dtype)  # (b, l, l, d_pair)
+                z = self.folding_trunk(z, pair_attention_mask=pair_mask)  # (b, l, l, d_pair)
 
                 if early_exit and loop_num < n_loops:
                     l2_converged = False
                     if prev_pair is not None and loop_num > 0:
                         rel_l2 = (
                             z.float() - prev_pair.float()
-                        ).norm() / prev_pair.float().norm().clamp(min=1e-8)
+                        ).norm() / prev_pair.float().norm().clamp(min=1e-8)  # ()
                         l2_converged = rel_l2.item() < 0.25
-                    prev_pair = z.detach().clone()
-                    sym_z = z.float() + z.float().transpose(-2, -3)
-                    cur_probs = F.softmax(self.distogram_head(sym_z).float(), dim=-1)
+                    prev_pair = z.detach().clone()  # (b, l, l, d_pair) or None
+                    sym_z = z.float() + z.float().transpose(-2, -3)  # (b, l, l, d_pair)
+                    cur_probs = F.softmax(self.distogram_head(sym_z).float(), dim=-1)  # (b, l, l, n_distogram_bins)
                     if prev_disto_probs is not None and loop_num > 0:
                         kl_per_pair = (
                             cur_probs
                             * (cur_probs.clamp(min=1e-8) / prev_disto_probs.clamp(min=1e-8)).log()
-                        ).sum(-1)
-                        kl = (kl_per_pair + kl_per_pair.transpose(-1, -2)).mean() / 2
+                        ).sum(-1)  # (b, l, l)
+                        kl = (kl_per_pair + kl_per_pair.transpose(-1, -2)).mean() / 2  # ()
                         if l2_converged or kl.item() < 0.05:
                             break
-                    prev_disto_probs = cur_probs.detach()
+                    prev_disto_probs = cur_probs.detach()  # (b, l, l, n_distogram_bins) or None
 
-            distogram_logits = self.distogram_head(z + z.transpose(-2, -3))
+            distogram_logits = self.distogram_head(z + z.transpose(-2, -3))  # (b, l, l, n_distogram_bins)
 
         with torch.no_grad(), _seed_context(seed):
             structure_output = self.structure_head.sample(
@@ -943,8 +952,8 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                 return_atom_repr=False,
                 denoising_early_exit_rmsd=(0.10 if early_exit else None),
                 verbose=verbose,
-            )
-        sample_coords = structure_output["sample_atom_coords"]
+            )  # tensor mapping follows the called head's shape contract
+        sample_coords = structure_output["sample_atom_coords"]  # (bs, a, 3), or explicit (b, samples, a, 3)
         if sample_coords is None:
             raise RuntimeError("ESMFold2 structure sampling did not return coordinates.")
         if sample_coords.ndim == 4:
@@ -953,15 +962,15 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                 batch * sample_count,
                 atom_count,
                 coord_dim,
-            )
-            rep_idx = distogram_atom_idx.repeat_interleave(sample_count, 0).long()
+            )  # (bs, a, 3)
+            rep_idx = distogram_atom_idx.repeat_interleave(sample_count, 0).long()  # (bs, l) for explicit sample axis, otherwise distogram_atom_idx.shape
         else:
-            sample_coords_for_gather = sample_coords
-            rep_idx = distogram_atom_idx.long()
+            sample_coords_for_gather = sample_coords  # (bs, a, 3)
+            rep_idx = distogram_atom_idx.long()  # (bs, l) for explicit sample axis, otherwise distogram_atom_idx.shape
         representative_atom_coords = gather_rep_atom_coords(
             sample_coords_for_gather,
             rep_idx,
-        )
+        )  # (*rep_idx.shape, 3); gather retains the index batch size
 
         output: dict[str, Tensor] = {
             "distogram_logits": distogram_logits,
@@ -990,7 +999,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                         num_diffusion_samples=n_samples,
                         relative_position_encoding=relative_position_encoding.detach(),
                         token_bonds_encoding=token_bonds_encoding.detach(),
-                    )
+                    )  # tensor mapping follows the called head's shape contract
                     progress.update()
             else:
                 confidence_output = self.confidence_head(
@@ -1006,18 +1015,18 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                     num_diffusion_samples=n_samples,
                     relative_position_encoding=relative_position_encoding.detach(),
                     token_bonds_encoding=token_bonds_encoding.detach(),
-                )
+                )  # tensor mapping follows the called head's shape contract
             output.update(confidence_output)
-        output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask
-        output["residue_index"] = residue_index
-        output["entity_id"] = entity_id
+        output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask  # (b, a)
+        output["residue_index"] = residue_index  # (b, l)
+        output["entity_id"] = entity_id  # (b, l)
         return _finalize_structure_output(
             output,
             token_input_state=x_inputs,
             pair_state=z,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
-        )
+        )  # ESMFold2Output/tuple retaining the traced tensor shapes
 
     @property
     def input_builder(self):
@@ -1058,7 +1067,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         if not self.config.msa_conditioning:
             for name in MSA_CONDITIONING_INPUT_NAMES:
                 features.pop(name, None)
-        features = {name: tensor.to(self.device) for name, tensor in features.items()}
+        features = {name: tensor.to(self.device) for name, tensor in features.items()}  # every feature retains its shape
         output = self(**features, **forward_kwargs, return_dict=True)
         for name in (
             "res_type",

fastplms/models/esmfold2/protein_utils.py

@@ -9,11 +9,11 @@ lazily from a provenance-bearing declarative package asset.
 from __future__ import annotations
 
 import json
+import torch
+
 from functools import cache
 from importlib.resources import files
 from typing import Any
-
-import torch
 from torch import Tensor
 
 from .esmfold2_constants import (
@@ -27,6 +27,7 @@ from .esmfold2_constants import (
     PROTEIN_UNK_RES_TYPE,
 )
 
+
 _GEOMETRY_ASSET = "protein_reference_geometry.json"
 _GEOMETRY_SCHEMA = "fastplms.esmfold2.reference_geometry.v1"
 
@@ -122,57 +123,57 @@ def prepare_protein_features(sequence: str) -> dict[str, Tensor]:
         raise ValueError("sequence must be non-empty")
 
     atoms, residue_types, input_ids, representative_atoms = _residue_records(sequence)
-    sequence_length = len(sequence)
+    sequence_length = len(sequence)  # l
     n_atoms = _padded_atom_count(len(atoms))
 
-    ref_pos = torch.zeros((n_atoms, 3), dtype=torch.float32)
-    ref_element = torch.zeros(n_atoms, dtype=torch.int64)
-    ref_charge = torch.zeros(n_atoms, dtype=torch.int8)
-    ref_atom_name_chars = torch.zeros((n_atoms, 4), dtype=torch.int64)
-    ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64)
-    atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool)
-    atom_to_token = torch.zeros(n_atoms, dtype=torch.int64)
+    ref_pos = torch.zeros((n_atoms, 3), dtype=torch.float32)  # (n_atoms, 3)
+    ref_element = torch.zeros(n_atoms, dtype=torch.int64)  # (n_atoms,)
+    ref_charge = torch.zeros(n_atoms, dtype=torch.int8)  # (n_atoms,)
+    ref_atom_name_chars = torch.zeros((n_atoms, 4), dtype=torch.int64)  # (n_atoms, 4)
+    ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64)  # (n_atoms,)
+    atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool)  # (n_atoms,)
+    atom_to_token = torch.zeros(n_atoms, dtype=torch.int64)  # (n_atoms,)
 
     for atom_index, atom in enumerate(atoms):
         token_index = atom["token_index"]
-        ref_pos[atom_index] = torch.tensor(atom["position"], dtype=torch.float32)
-        ref_element[atom_index] = ELEMENT_TO_ATOMIC_NUM[atom["element"]]
-        ref_charge[atom_index] = atom["charge"]
+        ref_pos[atom_index] = torch.tensor(atom["position"], dtype=torch.float32)  # (3,)
+        ref_element[atom_index] = ELEMENT_TO_ATOMIC_NUM[atom["element"]]  # ()
+        ref_charge[atom_index] = atom["charge"]  # ()
         ref_atom_name_chars[atom_index] = torch.tensor(
             _encode_atom_name(atom["name"]), dtype=torch.int64
-        )
-        ref_space_uid[atom_index] = token_index
-        atom_attention_mask[atom_index] = True
-        atom_to_token[atom_index] = token_index
+        )  # (4,)
+        ref_space_uid[atom_index] = token_index  # ()
+        atom_attention_mask[atom_index] = True  # ()
+        atom_to_token[atom_index] = token_index  # ()
 
-    residue_type_tensor = torch.tensor(residue_types, dtype=torch.int64)
-    msa = residue_type_tensor.unsqueeze(0)
+    residue_type_tensor = torch.tensor(residue_types, dtype=torch.int64)  # (l,)
+    msa = residue_type_tensor.unsqueeze(0)  # (1, l)
     features = {
-        "token_index": torch.arange(sequence_length, dtype=torch.int64),
-        "residue_index": torch.arange(sequence_length, dtype=torch.int64),
-        "asym_id": torch.zeros(sequence_length, dtype=torch.int64),
-        "sym_id": torch.zeros(sequence_length, dtype=torch.int64),
-        "entity_id": torch.ones(sequence_length, dtype=torch.int64),
-        "mol_type": torch.full((sequence_length,), MOL_TYPE_PROTEIN, dtype=torch.int64),
-        "res_type": residue_type_tensor,
-        "input_ids": torch.tensor(input_ids, dtype=torch.int64),
-        "token_bonds": torch.zeros((sequence_length, sequence_length, 1), dtype=torch.float32),
-        "token_attention_mask": torch.ones(sequence_length, dtype=torch.bool),
-        "ref_pos": ref_pos,
-        "ref_element": ref_element,
-        "ref_charge": ref_charge,
-        "ref_atom_name_chars": ref_atom_name_chars,
-        "ref_space_uid": ref_space_uid,
-        "atom_attention_mask": atom_attention_mask,
-        "atom_to_token": atom_to_token,
-        "distogram_atom_idx": torch.tensor(representative_atoms, dtype=torch.int64),
-        "msa": msa,
-        "msa_attention_mask": torch.ones_like(msa, dtype=torch.bool),
-        "has_deletion": torch.zeros_like(msa, dtype=torch.bool),
-        "deletion_value": torch.zeros_like(msa, dtype=torch.float32),
-        "deletion_mean": torch.zeros(sequence_length, dtype=torch.float32),
+        "token_index": torch.arange(sequence_length, dtype=torch.int64),  # (l,)
+        "residue_index": torch.arange(sequence_length, dtype=torch.int64),  # (l,)
+        "asym_id": torch.zeros(sequence_length, dtype=torch.int64),  # (l,)
+        "sym_id": torch.zeros(sequence_length, dtype=torch.int64),  # (l,)
+        "entity_id": torch.ones(sequence_length, dtype=torch.int64),  # (l,)
+        "mol_type": torch.full((sequence_length,), MOL_TYPE_PROTEIN, dtype=torch.int64),  # (l,)
+        "res_type": residue_type_tensor,  # (l,)
+        "input_ids": torch.tensor(input_ids, dtype=torch.int64),  # (l,)
+        "token_bonds": torch.zeros((sequence_length, sequence_length, 1), dtype=torch.float32),  # (l, l, 1)
+        "token_attention_mask": torch.ones(sequence_length, dtype=torch.bool),  # (l,)
+        "ref_pos": ref_pos,  # (n_atoms, 3)
+        "ref_element": ref_element,  # (n_atoms,)
+        "ref_charge": ref_charge,  # (n_atoms,)
+        "ref_atom_name_chars": ref_atom_name_chars,  # (n_atoms, 4)
+        "ref_space_uid": ref_space_uid,  # (n_atoms,)
+        "atom_attention_mask": atom_attention_mask,  # (n_atoms,)
+        "atom_to_token": atom_to_token,  # (n_atoms,)
+        "distogram_atom_idx": torch.tensor(representative_atoms, dtype=torch.int64),  # (l,)
+        "msa": msa,  # (1, l)
+        "msa_attention_mask": torch.ones_like(msa, dtype=torch.bool),  # (1, l)
+        "has_deletion": torch.zeros_like(msa, dtype=torch.bool),  # (1, l)
+        "deletion_value": torch.zeros_like(msa, dtype=torch.float32),  # (1, l)
+        "deletion_mean": torch.zeros(sequence_length, dtype=torch.float32),  # (l,)
     }
-    return {name: tensor.unsqueeze(0) for name, tensor in features.items()}
+    return {name: tensor.unsqueeze(0) for name, tensor in features.items()}  # each shape -> (1, *shape)
 
 
 __all__ = ["prepare_protein_features"]

fastplms/models/esmfold2/reproducibility.py

@@ -3,13 +3,13 @@
 from __future__ import annotations
 
 import random
+import numpy as np
+import torch
+
 from collections.abc import Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Any
-
-import numpy as np
-import torch
 from torch import Tensor
 
 

fastplms/models/ttt.py

@@ -6,6 +6,7 @@ import numbers
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+
 from collections.abc import Iterator, Mapping
 from dataclasses import asdict, dataclass, fields
 from typing import Any