Update FastPLMs runtime and model cards

Add-only FastPLMs files-only publication. Checkpoint weights and complete-artifact attestations are unchanged.

README.md

@@ -16,16 +16,35 @@ Accepted inputs are raw amino-acid sequences through the convenience API, or
 prepared model features.
 Supported Transformers entry points are `AutoConfig`, `AutoModel`.
 
+## Capabilities
+
+| Feature | Status |
+| --- | --- |
+| Sequence classification | Unavailable: no advertised AutoClass |
+| Token classification | Unavailable: no advertised AutoClass |
+| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model |
+| Embeddings | Unavailable for this structure-only checkpoint |
+| Test-time training | Unavailable for this inference-only checkpoint |
+| Attention variants | Supported: `eager` |
+| Compliance | Unavailable: this provisional family has no compliance tier |
+
+A supported interface is not a pretrained downstream predictor. Classification
+heads start untrained, and declared compliance metadata is not a claim that an
+arbitrary local build passed its release gate.
+
 ## Install and platform requirements
 
-Install the current FastPLMs package:
+Install the direct dependencies published with this model:
 
 ```bash
-python -m pip install \
-  "fastplms[structure] @ git+https://github.com/Synthyra/FastPLMs.git"
+python -m pip install -r \
+  "https://huggingface.co/Synthyra/Boltz2/resolve/main/requirements.txt"
 ```
 
-Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. Structure inference requires the `structure` extra and a CUDA device for the published execution contract. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network
+The FastPLMs implementation itself is embedded in the model repository and loaded
+by Transformers through `trust_remote_code=True`.
+
+Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct structure dependencies. The published execution contract requires a CUDA device. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network
 access on first download. For an air-gapped run, first build the manifest-pinned
 local artifact and use the offline form shown in the example.
 
@@ -38,22 +57,50 @@ model_id = "Synthyra/Boltz2"
 model = AutoModel.from_pretrained(
     model_id,
     trust_remote_code=True,
+    attn_implementation="eager",
 ).eval()
 ```
 
-This example uses the published Hub repository. For offline validation, build
-the manifest-pinned artifact and replace `model_id` with its local
-`dist/hub/Boltz2` path, then pass `local_files_only=True`.
+For offline validation, replace `model_id` with the manifest-built
+`dist/hub/Boltz2` path and pass `local_files_only=True`.
+
+## Attention and compliance
+
+The quick start selects `eager` explicitly. Declared variants are `eager`. An unavailable requested backend raises instead
+of silently switching implementations.
+`output_attentions=True` may use the documented, one-call eager fallback solely
+to materialize attention tensors; the configured backend remains unchanged.
 
-Leave attention unspecified for the Transformers default. Supported explicit
-choices are `eager`.
-Pass the selected name through `attn_implementation`.
-When an optimized backend cannot return full attention tensors,
-`output_attentions=True` emits one explicit runtime warning and uses a correctly
-masked eager implementation for that call only. The warning identifies the
-configured backend, effective backend, and reason. Configuration and later
-calls are unchanged.
-For BF16 execution, this family uses FP32 parameters with CUDA BF16 autocast.
+This family does not declare the `compliance` tier. Boltz2 remains provisional
+and its structure checks must not be broadened into parity claims.
+
+## PEFT fine-tuning
+
+Install the direct training dependencies, then attach LoRA to the loaded checkpoint:
+
+```bash
+python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
+```
+
+```python
+from peft import LoraConfig, get_peft_model
+
+peft_model = get_peft_model(
+    model,
+    LoraConfig(
+        r=8,
+        lora_alpha=16,
+        target_modules="all-linear",
+    ),
+)
+```
+
+This checkpoint has no advertised classifier. Supply the task-specific
+objective and preserve any new head through `modules_to_save`.
+All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and
+can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a
+support boundary. Record the target modules, base revision, data identity, and
+trainable parameter scope.
 
 ## Protein structure prediction
 
@@ -100,7 +147,7 @@ continues independently of the ESM++ and ESMFold2 release gates.
 - Precision policies: `default`
 - BF16 execution: `fp32_parameters_autocast`
 - Generation contract: `not_applicable`
-- Optional dependency group: `structure`
+- Artifact dependency set: `core + structure`
 - Weight publication allowed: `true`
 - Weight license status: `resolved`
 - Redistributable: `true`
@@ -111,28 +158,22 @@ continues independently of the ESM++ and ESMFold2 release gates.
 - FastPLMs weights: `Synthyra/Boltz2`
 - Runtime revision: recorded separately in the built artifact and published commit
 - Source-tree and runtime-bundle SHA-256: recorded in `provenance.json`
-- Generator/schema version and complete/runtime-only attestations: recorded in `provenance.json`
 - Official checkpoint: `boltz-community/boltz-2`
 - Artifact source: `fast`
 - State transform: `boltz2_inference_core_v1`
-- BF16 execution: `fp32_parameters_autocast`
 - Pinned upstreams: `boltz`
-- Reference container: `reference-boltz2`
 - Release tiers: `structure`, `artifact`, `benchmark`
 - Unresolved required file identities: `0`
 
-The local artifact records exact file identities, conversion details, source
-revisions, and legal texts in `provenance.json`. A nonzero unresolved count is a
-release blocker.
+`provenance.json` records exact file identities, conversion, source revisions,
+legal texts, schema, and attestations. A nonzero unresolved count blocks release.
 
 ## Validation boundary
 
-For tiers declared by the manifest, the release contract compares applicable
-semantic configuration, tokenizer behavior, state keys, shapes, dtypes,
-values, aliases, and representative inference with the pinned official
-implementation. This metadata does not by itself claim that a particular build
-passed, that one backend is faster, or that an output has biological or
-therapeutic validity.
+Declared tiers compare applicable configuration, tokenizer behavior, state,
+and representative inference with the pinned reference. Metadata alone does
+not claim a build passed, a backend is faster, or an output is biologically
+valid.
 
 ## License
 

config.json

@@ -137,11 +137,11 @@
   "fastplms_checkpoint_repo_id": "Synthyra/Boltz2",
   "fastplms_checkpoint_revision": "3b148fc5efea109c065ec82ba8683d024de7134e",
   "fastplms_model_id": "boltz2",
-  "fastplms_release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "fastplms_runtime_bundle_sha256": "e6e0f8c8632238f40a5112f22397a9a211eb02d3415c1c93d43305aef51f7728",
-  "fastplms_runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_source_tree_sha256": "78061b1bd37ba2dda853641fba681fdf943b859503e4fb21b1a7508909c369bf",
+  "fastplms_release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "fastplms_runtime_bundle_sha256": "3aebb85d11edb01329d2aafb4c8260bd67339785f0f252582246ae7c33206143",
+  "fastplms_runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_source_tree_sha256": "9d2fc9586cfbd1d2d2cc49d745311acb81d4e3740e59ead09e633917589806a8",
   "fastplms_weights_revision": "3b148fc5efea109c065ec82ba8683d024de7134e",
   "model_type": "boltz2_automodel",
   "num_bins": 64,

fastplms/__init__.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 from importlib import import_module
 from typing import Any
 
+
 __version__ = "1.0.0"
 
 _LAZY_EXPORTS = {

fastplms/models/__init__.py

@@ -7,4 +7,5 @@ tokenizers, compile kernels, or initialize an accelerator runtime.
 
 from __future__ import annotations
 
+
 __all__: tuple[str, ...] = ()

fastplms/models/boltz/__init__.py

@@ -5,6 +5,7 @@ from fastplms.models.boltz.modeling_boltz2 import (
     Boltz2StructureOutput,
 )
 
+
 __all__ = [
     "Boltz2Config",
     "Boltz2Model",

fastplms/models/boltz/_pair_attention.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import math
-
 import torch
 from torch import Tensor
 
@@ -11,9 +10,11 @@ from torch import Tensor
 def reshape_heads(states: Tensor, num_heads: int) -> Tensor:
     """Reshape X from ``(b, l, d)`` to ``(b, l, h, d_head)``."""
 
-    batch_size, sequence_length, width = states.shape
+    # states: (b, l, d); num_heads: h
+    batch_size, sequence_length, width = states.shape  # b, l, d
     if width % num_heads:
         raise ValueError(f"width {width} is not divisible by {num_heads} heads")
+    # states.view(...): (b, l, h, d_h)
     return states.view(batch_size, sequence_length, num_heads, width // num_heads)
 
 
@@ -27,20 +28,23 @@ def pair_biased_attention(
 ) -> Tensor:
     """Return pair-biased attention values with shape ``(b, l_q, h, d_head)``."""
 
-    head_dim = query_states.shape[-1]
+    # query_states: (b, l_q, h, d_h); key_states, value_states: (b, l_k, h, d_h)
+    # pair_bias: (b, h, l_q, l_k); key_mask: (b, l_k)
+    head_dim = query_states.shape[-1]  # d_h
     with torch.autocast("cuda", enabled=False):
-        # S is the FP32 attention-score tensor with shape (b, h, l_q, l_k).
         scores = torch.einsum(
             "bihd,bjhd->bhij",
             query_states.float(),
             key_states.float(),
-        )
-        scores = scores / math.sqrt(head_dim) + pair_bias.float()
-        scores = scores + (1 - key_mask[:, None, None].float()) * -mask_value
-        probabilities = scores.softmax(dim=-1)
+        )  # (b, h, l_q, l_k)
+        scores = scores / math.sqrt(head_dim) + pair_bias.float()  # (b, h, l_q, l_k)
+        scores = (
+            scores + (1 - key_mask[:, None, None].float()) * -mask_value
+        )  # (b, h, l_q, l_k)
+        probabilities = scores.softmax(dim=-1)  # (b, h, l_q, l_k)
         output = torch.einsum(
             "bhij,bjhd->bihd",
             probabilities,
             value_states.float(),
-        )
-    return output.to(value_states.dtype)
+        )  # (b, l_q, h, d_h)
+    return output.to(value_states.dtype)  # (b, l_q, h, d_h)

fastplms/models/boltz/cif_writer.py

@@ -1,7 +1,6 @@
-from pathlib import Path
-
 import numpy as np
 import torch
+from pathlib import Path
 
 from .minimal_structures import ProteinStructureTemplate
 
@@ -12,28 +11,30 @@ def _confidence_per_atom(
     num_atoms: int,
     sample_index: int,
 ) -> np.ndarray:
+    # plddt: (n_i,) or (n_s, n_i); atom_to_residue: length n_a; num_atoms: n_a
     if plddt is None:
-        return np.ones((num_atoms,), dtype=np.float32) * 100.0
+        return np.ones((num_atoms,), dtype=np.float32) * 100.0  # (n_a,)
 
-    values = plddt.detach().cpu()
+    values = plddt.detach().cpu()  # (n_i,) or (n_s, n_i)
     if values.ndim == 1:
-        values = values.unsqueeze(0)
+        values = values.unsqueeze(0)  # (1, n_i)
     if values.ndim != 2:
         raise ValueError("Expected pLDDT tensor S with shape (n_samples, n_items).")
+    # values: (n_s, n_i) beyond this point
     if not 0 <= sample_index < values.shape[0]:
         raise IndexError("sample_index out of range for pLDDT.")
 
-    selected = values[sample_index]
+    selected = values[sample_index]  # (n_i,)
     if selected.shape[0] == num_atoms:
-        return (selected.numpy() * 100.0).astype(np.float32)
+        return (selected.numpy() * 100.0).astype(np.float32)  # (n_a,)
 
-    num_residues = max(atom_to_residue) + 1
+    num_residues = max(atom_to_residue) + 1  # n_r
     if selected.shape[0] == num_residues:
-        expanded = np.zeros((num_atoms,), dtype=np.float32)
-        selected_np = selected.numpy()
+        expanded = np.zeros((num_atoms,), dtype=np.float32)  # (n_a,)
+        selected_np = selected.numpy()  # (n_r,)
         for atom_idx, residue_idx in enumerate(atom_to_residue):
-            expanded[atom_idx] = selected_np[residue_idx] * 100.0
-        return expanded
+            expanded[atom_idx] = selected_np[residue_idx] * 100.0  # () -> ()
+        return expanded  # (n_a,)
 
     raise ValueError(
         "pLDDT item count must match either atoms or residues: "
@@ -50,34 +51,39 @@ def write_cif(
     plddt: torch.Tensor | None = None,
     sample_index: int = 0,
 ) -> str:
-    coords = atom_coords.detach().cpu()
+    # atom_coords: (n_a, 3) or (n_s, n_a, 3)
+    # atom_mask: (n_a,) or (n_m, n_a); plddt: (n_i,) or (n_s, n_i)
+    coords = atom_coords.detach().cpu()  # (n_a, 3) or (n_s, n_a, 3)
     if coords.ndim == 2:
-        coords = coords.unsqueeze(0)
+        coords = coords.unsqueeze(0)  # (1, n_a, 3)
     if coords.ndim != 3 or coords.shape[-1] != 3:
         raise ValueError(
             "Expected coordinate tensor X with shape (n_samples, n_atoms, 3)."
         )
+    # coords: (n_s, n_a, 3) beyond this point
     if not 0 <= sample_index < coords.shape[0]:
         raise IndexError("sample_index out of range.")
-    selected_coords_tensor = coords[sample_index]
-    all_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor))
+    selected_coords_tensor = coords[sample_index]  # (n_a, 3)
+    all_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor))  # (n_a, 3)
     if torch.any(all_non_finite):
         raise ValueError(
             "CIF export received non-finite coordinates. "
             f"Non-finite count: {int(all_non_finite.sum().item())}"
         )
-    selected_coords = selected_coords_tensor.numpy()
+    selected_coords = selected_coords_tensor.numpy()  # (n_a, 3)
 
-    mask = atom_mask.detach().cpu()
+    mask = atom_mask.detach().cpu()  # (n_a,) or (n_m, n_a)
     if mask.ndim == 2:
-        mask = mask[0]
+        mask = mask[0]  # (n_a,)
     if mask.ndim != 1:
         raise ValueError("Expected atom mask M with shape (n_atoms,).")
+    # mask: (n_a,) beyond this point
     if mask.shape[0] != selected_coords.shape[0]:
         raise ValueError("Atom mask/coord size mismatch.")
     if not torch.any(mask > 0):
         raise ValueError("Atom mask has no valid atoms for CIF export.")
     valid_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor[mask > 0]))
+    # valid_non_finite: (n_v, 3), where n_v is the valid-atom count
     if torch.any(valid_non_finite):
         raise ValueError(
             "CIF export has non-finite coordinates in unmasked atoms. "
@@ -89,7 +95,7 @@ def write_cif(
         atom_to_residue=structure_template.atom_residue_index,
         num_atoms=structure_template.num_atoms,
         sample_index=sample_index,
-    )
+    )  # (n_a,)
     if b_iso.shape[0] != structure_template.num_atoms:
         raise RuntimeError("CIF confidence values do not match the structure atom count.")
 
@@ -122,8 +128,8 @@ def write_cif(
         atom_name = structure_template.atom_names[idx]
         element = structure_template.atom_elements[idx]
         chain_id = structure_template.atom_chain_id[idx]
-        x_val, y_val, z_val = selected_coords[idx].tolist()
-        b_factor = float(b_iso[idx])
+        x_val, y_val, z_val = selected_coords[idx].tolist()  # (3,) -> three scalars
+        b_factor = float(b_iso[idx])  # () -> Python float
 
         line = (
             f"ATOM {atom_id} {element} {atom_name} {residue_name} {chain_id} "

fastplms/models/boltz/minimal_featurizer.py

@@ -1,5 +1,4 @@
 import math
-
 import numpy as np
 import torch
 from torch.nn.functional import one_hot
@@ -7,6 +6,7 @@ from torch.nn.functional import one_hot
 from . import vb_const as const
 from .minimal_structures import ProteinStructureTemplate
 
+
 _ELEMENT_TO_Z = {
     "H": 1,
     "C": 6,
@@ -75,7 +75,7 @@ def _atom_name_to_codes(atom_name: str) -> torch.Tensor:
     vals = [ord(ch) - 32 for ch in clipped]
     while len(vals) < 4:
         vals.append(0)
-    out = torch.tensor(vals, dtype=torch.long)
+    out = torch.tensor(vals, dtype=torch.long)  # (4,)
     if not (torch.all(out >= 0) and torch.all(out < 64)):
         raise ValueError(f"Invalid atom-name encoding for '{atom_name}'.")
     return out
@@ -309,7 +309,7 @@ def _get_atom_position(res_name: str, atom_name: str, atom_idx: int) -> np.ndarr
     to a simple geometric placement for unknown residue/atom combinations.
     """
     if res_name in _RDKIT_CONFORMERS and atom_name in _RDKIT_CONFORMERS[res_name]:
-        return np.array(_RDKIT_CONFORMERS[res_name][atom_name], dtype=np.float32)
+        return np.array(_RDKIT_CONFORMERS[res_name][atom_name], dtype=np.float32)  # (3,)
     # Fallback for unknown atoms (should not happen for canonical AAs)
     angle = (atom_idx + 1) * 0.7
     radius = 1.4 + 0.03 * atom_idx
@@ -320,7 +320,7 @@ def _get_atom_position(res_name: str, atom_name: str, atom_idx: int) -> np.ndarr
             0.1 * ((atom_idx % 5) - 2),
         ],
         dtype=np.float32,
-    )
+    )  # (3,)
 
 
 def _build_template(
@@ -340,7 +340,7 @@ def _build_template(
     atom_elements: list[str] = []
     atom_residue_index: list[int] = []
     atom_chain_id: list[str] = []
-    atom_positions: list[np.ndarray] = []
+    atom_positions: list[np.ndarray] = []  # each: (3,)
     residue_center_atom_idx: list[int] = []
     residue_disto_atom_idx: list[int] = []
     residue_frame_atom_idx: list[int] = []
@@ -370,7 +370,7 @@ def _build_template(
             atom_residue_index.append(res_idx)
             atom_chain_id.append("A")
 
-            atom_pos = _get_atom_position(token_name, atom_name, local_idx)
+            atom_pos = _get_atom_position(token_name, atom_name, local_idx)  # (3,)
             atom_positions.append(atom_pos)
 
             if atom_name == center_atom_name:
@@ -422,14 +422,14 @@ def _build_template(
 
 def _random_rotation_matrix() -> torch.Tensor:
     """Sample a uniform random 3x3 rotation matrix (Algorithm 19 from AF2/Boltz)."""
-    quaternion = torch.randn((1, 4), dtype=torch.float32)
-    squared_norm = (quaternion * quaternion).sum(dim=1)
-    norm = torch.sqrt(squared_norm)
-    signed_norm = torch.where(quaternion[:, 0] < 0, -norm, norm)
-    quaternion = quaternion / signed_norm[:, None]
-
-    real, i, j, k = torch.unbind(quaternion, dim=-1)
-    two_s = 2.0 / (quaternion * quaternion).sum(dim=-1)
+    quaternion = torch.randn((1, 4), dtype=torch.float32)  # (1, 4)
+    squared_norm = (quaternion * quaternion).sum(dim=1)  # (1,)
+    norm = torch.sqrt(squared_norm)  # (1,)
+    signed_norm = torch.where(quaternion[:, 0] < 0, -norm, norm)  # (1,)
+    quaternion = quaternion / signed_norm[:, None]  # (1, 4)
+
+    real, i, j, k = torch.unbind(quaternion, dim=-1)  # each: (1,)
+    two_s = 2.0 / (quaternion * quaternion).sum(dim=-1)  # (1,)
     rotation = torch.stack(
         (
             1 - two_s * (j * j + k * k),
@@ -443,8 +443,8 @@ def _random_rotation_matrix() -> torch.Tensor:
             1 - two_s * (i * i + j * j),
         ),
         dim=-1,
-    )
-    return rotation.reshape(1, 3, 3)[0]
+    )  # (1, 9)
+    return rotation.reshape(1, 3, 3)[0]  # (3, 3)
 
 
 def _center_and_augment_atoms_per_residue(
@@ -457,25 +457,30 @@ def _center_and_augment_atoms_per_residue(
     The pinned implementation intentionally excludes the final residue because
     it iterates to the maximum reference-space identifier rather than through it.
     """
-    result = atom_positions.clone()
-    residue_index_tensor = torch.tensor(atom_residue_index, dtype=torch.long)
+    # atom_positions: (a, 3), where a is the unpadded atom count.
+    result = atom_positions.clone()  # (a, 3)
+    residue_index_tensor = torch.tensor(atom_residue_index, dtype=torch.long)  # (a,)
     for residue_idx in range(max(num_residues - 1, 0)):
-        residue_mask = residue_index_tensor == residue_idx
+        residue_mask = residue_index_tensor == residue_idx  # (a,)
         if not torch.any(residue_mask):
             raise RuntimeError(f"Residue index {residue_idx} has no atoms.")
-        residue_coords = result[residue_mask][None]
+        residue_coords = result[residue_mask][None]  # (1, a_r, 3)
         resolved_mask = torch.ones(
             residue_coords.shape[:2], dtype=torch.bool, device=residue_coords.device
-        )
+        )  # (1, a_r)
         residue_center = torch.sum(
             residue_coords * resolved_mask[:, :, None], dim=1, keepdim=True
-        ) / torch.sum(resolved_mask[:, :, None], dim=1, keepdim=True)
-        residue_coords = residue_coords - residue_center
-        rotation = _random_rotation_matrix()[None]
-        residue_coords = torch.einsum("bmd,bds->bms", residue_coords, rotation)
-        residue_coords = residue_coords + torch.randn_like(residue_coords[:, 0:1, :])
-        result[residue_mask] = residue_coords[0]
-    return result
+        ) / torch.sum(resolved_mask[:, :, None], dim=1, keepdim=True)  # (1, 1, 3)
+        residue_coords = residue_coords - residue_center  # (1, a_r, 3)
+        rotation = _random_rotation_matrix()[None]  # (1, 3, 3)
+        residue_coords = torch.einsum(
+            "bmd,bds->bms", residue_coords, rotation
+        )  # (1, a_r, 3)
+        residue_coords = (
+            residue_coords + torch.randn_like(residue_coords[:, 0:1, :])
+        )  # (1, a_r, 3)
+        result[residue_mask] = residue_coords[0]  # (a_r, 3)
+    return result  # (a, 3)
 
 
 def build_boltz2_features(
@@ -499,39 +504,44 @@ def build_boltz2_features(
     if num_tokens <= 0 or num_atoms <= 0:
         raise RuntimeError("Boltz2 feature construction produced an empty protein template.")
 
-    atom_positions = torch.tensor(np.asarray(atom_positions_np), dtype=torch.float32)
+    # t is the token count, a is the unpadded atom count, and a_p is the padded atom count.
+    atom_positions = torch.tensor(
+        np.asarray(atom_positions_np), dtype=torch.float32
+    )  # (a, 3)
     atom_positions = _center_and_augment_atoms_per_residue(
         atom_positions=atom_positions,
         atom_residue_index=template.atom_residue_index,
         num_residues=num_tokens,
-    )
+    )  # (a, 3)
 
-    token_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
-    residue_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
-    asym_id = torch.zeros((1, num_tokens), dtype=torch.long)
-    entity_id = torch.zeros((1, num_tokens), dtype=torch.long)
-    sym_id = torch.zeros((1, num_tokens), dtype=torch.long)
+    token_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)  # (1, t)
+    residue_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)  # (1, t)
+    asym_id = torch.zeros((1, num_tokens), dtype=torch.long)  # (1, t)
+    entity_id = torch.zeros((1, num_tokens), dtype=torch.long)  # (1, t)
+    sym_id = torch.zeros((1, num_tokens), dtype=torch.long)  # (1, t)
     mol_type = torch.full(
         (1, num_tokens),
         fill_value=const.chain_type_ids["PROTEIN"],
         dtype=torch.long,
-    )
+    )  # (1, t)
 
-    res_type_ids = torch.tensor(residue_token_ids, dtype=torch.long)
-    res_type = one_hot(res_type_ids, num_classes=const.num_tokens).unsqueeze(0)
+    res_type_ids = torch.tensor(residue_token_ids, dtype=torch.long)  # (t,)
+    res_type = one_hot(
+        res_type_ids, num_classes=const.num_tokens
+    ).unsqueeze(0)  # (1, t, n_res_type)
 
     # token_bonds encodes explicit covalent cross-links from structure bonds,
     # NOT backbone peptide bonds (those are implicit via residue_index + asym_id).
     # For a standard single-chain protein without cross-links, this is all zeros.
     # This matches the official Boltz2 featurizer (featurizerv2.py lines 696-705).
-    token_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.float32)
-    type_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.long)
-    token_bonds = token_bonds.unsqueeze(0).unsqueeze(-1)
-    type_bonds = type_bonds.unsqueeze(0)
+    token_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.float32)  # (t, t)
+    type_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.long)  # (t, t)
+    token_bonds = token_bonds.unsqueeze(0).unsqueeze(-1)  # (1, t, t, 1)
+    type_bonds = type_bonds.unsqueeze(0)  # (1, t, t)
 
-    token_pad_mask = torch.ones((1, num_tokens), dtype=torch.float32)
-    token_resolved_mask = torch.ones((1, num_tokens), dtype=torch.float32)
-    token_disto_mask = torch.ones((1, num_tokens), dtype=torch.float32)
+    token_pad_mask = torch.ones((1, num_tokens), dtype=torch.float32)  # (1, t)
+    token_resolved_mask = torch.ones((1, num_tokens), dtype=torch.float32)  # (1, t)
+    token_disto_mask = torch.ones((1, num_tokens), dtype=torch.float32)  # (1, t)
 
     num_contact_classes = len(const.contact_conditioning_info)
     unspecified_id = const.contact_conditioning_info["UNSPECIFIED"]
@@ -539,12 +549,14 @@ def build_boltz2_features(
         (num_tokens, num_tokens),
         fill_value=unspecified_id,
         dtype=torch.long,
-    )
+    )  # (t, t)
     contact_conditioning = one_hot(
         contact_ids,
         num_classes=num_contact_classes,
-    ).unsqueeze(0)
-    contact_threshold = torch.zeros((1, num_tokens, num_tokens), dtype=torch.float32)
+    ).unsqueeze(0)  # (1, t, t, n_contact)
+    contact_threshold = torch.zeros(
+        (1, num_tokens, num_tokens), dtype=torch.float32
+    )  # (1, t, t)
 
     if "x-ray diffraction" not in const.method_types_ids:
         raise RuntimeError("Boltz2 method metadata omits x-ray diffraction.")
@@ -552,20 +564,22 @@ def build_boltz2_features(
         (1, num_tokens),
         fill_value=const.method_types_ids["x-ray diffraction"],
         dtype=torch.long,
-    )
-    modified = torch.zeros((1, num_tokens), dtype=torch.long)
-    cyclic_period = torch.zeros((1, num_tokens), dtype=torch.float32)
-    affinity_token_mask = torch.zeros((1, num_tokens), dtype=torch.float32)
+    )  # (1, t)
+    modified = torch.zeros((1, num_tokens), dtype=torch.long)  # (1, t)
+    cyclic_period = torch.zeros((1, num_tokens), dtype=torch.float32)  # (1, t)
+    affinity_token_mask = torch.zeros((1, num_tokens), dtype=torch.float32)  # (1, t)
 
-    ref_pos = atom_positions.unsqueeze(0)
-    atom_pad_mask = torch.ones((1, num_atoms), dtype=torch.float32)
-    atom_resolved_mask = torch.ones((1, num_atoms), dtype=torch.bool)
+    ref_pos = atom_positions.unsqueeze(0)  # (1, a, 3)
+    atom_pad_mask = torch.ones((1, num_atoms), dtype=torch.float32)  # (1, a)
+    atom_resolved_mask = torch.ones((1, num_atoms), dtype=torch.bool)  # (1, a)
 
     atom_name_codes = torch.stack(
         [_atom_name_to_codes(atom_name) for atom_name in template.atom_names],
         dim=0,
-    )
-    ref_atom_name_chars = one_hot(atom_name_codes, num_classes=64).unsqueeze(0)
+    )  # (a, 4)
+    ref_atom_name_chars = one_hot(
+        atom_name_codes, num_classes=64
+    ).unsqueeze(0)  # (1, a, 4, 64)
 
     atomic_numbers = []
     for element in template.atom_elements:
@@ -578,7 +592,7 @@ def build_boltz2_features(
     ref_element = one_hot(
         torch.tensor(atomic_numbers, dtype=torch.long),
         num_classes=const.num_elements,
-    ).unsqueeze(0)
+    ).unsqueeze(0)  # (1, a, n_element)
 
     ref_charge_values = []
     ref_chirality_values = []
@@ -592,23 +606,27 @@ def build_boltz2_features(
         ref_chirality_values.append(
             2 if atom_name in _CHIRAL_ATOMS.get(residue_name, frozenset()) else 0
         )
-    ref_charge = torch.tensor(ref_charge_values, dtype=torch.float32).unsqueeze(0)
-    ref_chirality = torch.tensor(ref_chirality_values, dtype=torch.long).unsqueeze(0)
-    ref_space_uid = torch.tensor(template.atom_residue_index, dtype=torch.long).unsqueeze(0)
+    ref_charge = torch.tensor(ref_charge_values, dtype=torch.float32).unsqueeze(0)  # (1, a)
+    ref_chirality = torch.tensor(
+        ref_chirality_values, dtype=torch.long
+    ).unsqueeze(0)  # (1, a)
+    ref_space_uid = torch.tensor(
+        template.atom_residue_index, dtype=torch.long
+    ).unsqueeze(0)  # (1, a)
 
     atom_to_token = one_hot(
         torch.tensor(template.atom_residue_index, dtype=torch.long),
         num_classes=num_tokens,
-    ).unsqueeze(0)
+    ).unsqueeze(0)  # (1, a, t)
     token_to_rep_atom = one_hot(
         torch.tensor(residue_disto_atom_idx, dtype=torch.long),
         num_classes=num_atoms,
-    ).unsqueeze(0)
+    ).unsqueeze(0)  # (1, t, a)
     token_to_center_atom = one_hot(
         torch.tensor(residue_center_atom_idx, dtype=torch.long),
         num_classes=num_atoms,
-    ).unsqueeze(0)
-    r_set_to_rep_atom = token_to_center_atom.clone()
+    ).unsqueeze(0)  # (1, t, a)
+    r_set_to_rep_atom = token_to_center_atom.clone()  # (1, t, a)
 
     num_backbone_classes = (
         1 + len(const.protein_backbone_atom_index) + len(const.nucleic_backbone_atom_index)
@@ -622,17 +640,17 @@ def build_boltz2_features(
     atom_backbone_feat = one_hot(
         torch.tensor(backbone_ids, dtype=torch.long),
         num_classes=num_backbone_classes,
-    ).unsqueeze(0)
+    ).unsqueeze(0)  # (1, a, n_backbone)
 
     # X contains no observed coordinates for sequence-only inference.
-    coords = torch.zeros((1, 1, num_atoms, 3), dtype=torch.float32)
+    coords = torch.zeros((1, 1, num_atoms, 3), dtype=torch.float32)  # (1, 1, a, 3)
     disto_coords_ensemble = torch.zeros(
         (1, 1, num_tokens, 3),
         dtype=torch.float32,
-    )
+    )  # (1, 1, t, 3)
 
-    bfactor = torch.zeros((1, num_atoms), dtype=torch.float32)
-    atom_plddt = torch.ones((1, num_atoms), dtype=torch.float32)
+    bfactor = torch.zeros((1, num_atoms), dtype=torch.float32)  # (1, a)
+    atom_plddt = torch.ones((1, num_atoms), dtype=torch.float32)  # (1, a)
 
     if atoms_per_window_queries <= 0:
         raise ValueError("atoms_per_window_queries must be positive.")
@@ -640,60 +658,86 @@ def build_boltz2_features(
         (num_atoms - 1) // atoms_per_window_queries + 1
     ) * atoms_per_window_queries - num_atoms
     if pad_atoms > 0:
-        ref_pos = torch.nn.functional.pad(ref_pos, (0, 0, 0, pad_atoms), value=0.0)
-        atom_pad_mask = torch.nn.functional.pad(atom_pad_mask, (0, pad_atoms), value=0.0)
+        ref_pos = torch.nn.functional.pad(
+            ref_pos, (0, 0, 0, pad_atoms), value=0.0
+        )  # (1, a_p, 3)
+        atom_pad_mask = torch.nn.functional.pad(
+            atom_pad_mask, (0, pad_atoms), value=0.0
+        )  # (1, a_p)
         atom_resolved_mask = torch.nn.functional.pad(
             atom_resolved_mask,
             (0, pad_atoms),
             value=0.0,
-        )
+        )  # (1, a_p)
         ref_atom_name_chars = torch.nn.functional.pad(
             ref_atom_name_chars,
             (0, 0, 0, 0, 0, pad_atoms),
             value=0.0,
-        )
-        ref_element = torch.nn.functional.pad(ref_element, (0, 0, 0, pad_atoms), value=0.0)
-        ref_charge = torch.nn.functional.pad(ref_charge, (0, pad_atoms), value=0.0)
-        ref_chirality = torch.nn.functional.pad(ref_chirality, (0, pad_atoms), value=0)
+        )  # (1, a_p, 4, 64)
+        ref_element = torch.nn.functional.pad(
+            ref_element, (0, 0, 0, pad_atoms), value=0.0
+        )  # (1, a_p, n_element)
+        ref_charge = torch.nn.functional.pad(
+            ref_charge, (0, pad_atoms), value=0.0
+        )  # (1, a_p)
+        ref_chirality = torch.nn.functional.pad(
+            ref_chirality, (0, pad_atoms), value=0
+        )  # (1, a_p)
         atom_backbone_feat = torch.nn.functional.pad(
             atom_backbone_feat,
             (0, 0, 0, pad_atoms),
             value=0.0,
-        )
-        ref_space_uid = torch.nn.functional.pad(ref_space_uid, (0, pad_atoms), value=0)
-        coords = torch.nn.functional.pad(coords, (0, 0, 0, pad_atoms), value=0.0)
-        atom_to_token = torch.nn.functional.pad(atom_to_token, (0, 0, 0, pad_atoms), value=0.0)
+        )  # (1, a_p, n_backbone)
+        ref_space_uid = torch.nn.functional.pad(
+            ref_space_uid, (0, pad_atoms), value=0
+        )  # (1, a_p)
+        coords = torch.nn.functional.pad(
+            coords, (0, 0, 0, pad_atoms), value=0.0
+        )  # (1, 1, a_p, 3)
+        atom_to_token = torch.nn.functional.pad(
+            atom_to_token, (0, 0, 0, pad_atoms), value=0.0
+        )  # (1, a_p, t)
         token_to_rep_atom = torch.nn.functional.pad(
             token_to_rep_atom,
             (0, pad_atoms),
             value=0.0,
-        )
+        )  # (1, t, a_p)
         token_to_center_atom = torch.nn.functional.pad(
             token_to_center_atom,
             (0, pad_atoms),
             value=0.0,
-        )
+        )  # (1, t, a_p)
         r_set_to_rep_atom = torch.nn.functional.pad(
             r_set_to_rep_atom,
             (0, pad_atoms),
             value=0.0,
-        )
-        bfactor = torch.nn.functional.pad(bfactor, (0, pad_atoms), value=0.0)
-        atom_plddt = torch.nn.functional.pad(atom_plddt, (0, pad_atoms), value=0.0)
+        )  # (1, t, a_p)
+        bfactor = torch.nn.functional.pad(
+            bfactor, (0, pad_atoms), value=0.0
+        )  # (1, a_p)
+        atom_plddt = torch.nn.functional.pad(
+            atom_plddt, (0, pad_atoms), value=0.0
+        )  # (1, a_p)
 
     frames_idx = torch.tensor(
         residue_frame_atom_idx_flat,
         dtype=torch.long,
-    ).reshape(num_tokens, 3)
-    frames_idx = frames_idx.unsqueeze(0).unsqueeze(1)
-    frame_resolved_mask = torch.zeros((1, 1, num_tokens), dtype=torch.bool)
-
-    msa = torch.tensor(residue_token_ids, dtype=torch.long).unsqueeze(0).unsqueeze(0)
-    msa_paired = torch.ones((1, 1, num_tokens), dtype=torch.float32)
-    deletion_value = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
-    has_deletion = torch.zeros((1, 1, num_tokens), dtype=torch.bool)
-    msa_mask = torch.ones((1, 1, num_tokens), dtype=torch.long)
-    deletion_mean = torch.zeros((1, num_tokens), dtype=torch.float32)
+    ).reshape(num_tokens, 3)  # (t, 3)
+    frames_idx = frames_idx.unsqueeze(0).unsqueeze(1)  # (1, 1, t, 3)
+    frame_resolved_mask = torch.zeros(
+        (1, 1, num_tokens), dtype=torch.bool
+    )  # (1, 1, t)
+
+    msa = (
+        torch.tensor(residue_token_ids, dtype=torch.long).unsqueeze(0).unsqueeze(0)
+    )  # (1, 1, t)
+    msa_paired = torch.ones((1, 1, num_tokens), dtype=torch.float32)  # (1, 1, t)
+    deletion_value = torch.zeros(
+        (1, 1, num_tokens), dtype=torch.float32
+    )  # (1, 1, t)
+    has_deletion = torch.zeros((1, 1, num_tokens), dtype=torch.bool)  # (1, 1, t)
+    msa_mask = torch.ones((1, 1, num_tokens), dtype=torch.long)  # (1, 1, t)
+    deletion_mean = torch.zeros((1, num_tokens), dtype=torch.float32)  # (1, t)
     profile = (
         one_hot(
             torch.tensor(residue_token_ids, dtype=torch.long),
@@ -701,28 +745,48 @@ def build_boltz2_features(
         )
         .float()
         .unsqueeze(0)
-    )
+    )  # (1, t, n_res_type)
 
     template_restype = one_hot(
         torch.zeros((1, 1, num_tokens), dtype=torch.long),
         num_classes=const.num_tokens,
-    )
-    template_frame_rot = torch.zeros((1, 1, num_tokens, 3, 3), dtype=torch.float32)
-    template_frame_t = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
-    template_cb = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
-    template_ca = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
-    template_mask_cb = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
-    template_mask_frame = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
-    template_mask = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
-    query_to_template = torch.zeros((1, 1, num_tokens), dtype=torch.long)
-    visibility_ids = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    )  # (1, 1, t, n_res_type)
+    template_frame_rot = torch.zeros(
+        (1, 1, num_tokens, 3, 3), dtype=torch.float32
+    )  # (1, 1, t, 3, 3)
+    template_frame_t = torch.zeros(
+        (1, 1, num_tokens, 3), dtype=torch.float32
+    )  # (1, 1, t, 3)
+    template_cb = torch.zeros(
+        (1, 1, num_tokens, 3), dtype=torch.float32
+    )  # (1, 1, t, 3)
+    template_ca = torch.zeros(
+        (1, 1, num_tokens, 3), dtype=torch.float32
+    )  # (1, 1, t, 3)
+    template_mask_cb = torch.zeros(
+        (1, 1, num_tokens), dtype=torch.float32
+    )  # (1, 1, t)
+    template_mask_frame = torch.zeros(
+        (1, 1, num_tokens), dtype=torch.float32
+    )  # (1, 1, t)
+    template_mask = torch.zeros(
+        (1, 1, num_tokens), dtype=torch.float32
+    )  # (1, 1, t)
+    query_to_template = torch.zeros(
+        (1, 1, num_tokens), dtype=torch.long
+    )  # (1, 1, t)
+    visibility_ids = torch.zeros(
+        (1, 1, num_tokens), dtype=torch.float32
+    )  # (1, 1, t)
 
     disto_target = torch.zeros(
         (1, num_tokens, num_tokens, 1, num_bins),
         dtype=torch.float32,
-    )
-    disto_target[..., 0] = 1.0
-    disto_center = torch.zeros((1, num_tokens, 3), dtype=torch.float32)
+    )  # (1, t, t, 1, n_bin)
+    disto_target[..., 0] = 1.0  # (1, t, t, 1)
+    disto_center = torch.zeros(
+        (1, num_tokens, 3), dtype=torch.float32
+    )  # (1, t, 3)
 
     features: dict[str, torch.Tensor] = {
         "token_index": token_index,
@@ -783,4 +847,4 @@ def build_boltz2_features(
         "visibility_ids": visibility_ids,
     }
 
-    return features, template
+    return features, template  # each feature shape is traced at its binding above

fastplms/models/boltz/modeling_boltz2.py

@@ -1,14 +1,13 @@
 import copy
 import inspect
 import random
+import numpy as np
+import torch
+import torch.nn as nn
 from collections.abc import Iterator, Mapping, Sequence
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from typing import Any, ClassVar
-
-import numpy as np
-import torch
-import torch.nn as nn
 from torch import Tensor
 from transformers import PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import ModelOutput
@@ -113,6 +112,7 @@ def _require_key(mapping: dict[str, Any], key: str) -> Any:
 
 
 def _state_dict_without_wrappers(state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
+    # Every state tensor has its parameter-defined shape (...).
     cleaned: dict[str, Tensor] = {}
     for key, value in state_dict.items():
         if key.startswith("ema."):
@@ -122,13 +122,13 @@ def _state_dict_without_wrappers(state_dict: dict[str, Tensor]) -> dict[str, Ten
             new_key = new_key[len("model.") :]
         if new_key.startswith("module."):
             new_key = new_key[len("module.") :]
-        cleaned[new_key] = value
+        cleaned[new_key] = value  # (...)
     return cleaned
 
 
 def _to_cpu_detached(value: Any) -> Any:
     if torch.is_tensor(value):
-        return value.detach().cpu()
+        return value.detach().cpu()  # same shape (...)
     if isinstance(value, dict):
         out: dict[Any, Any] = {}
         for key, nested_value in value.items():
@@ -145,8 +145,8 @@ def _to_cpu_detached(value: Any) -> Any:
 class _RandomState:
     python: object
     numpy: tuple[Any, ...]
-    torch_cpu: Tensor
-    torch_cuda: list[Tensor] | None
+    torch_cpu: Tensor  # (n_cpu_rng_state,)
+    torch_cuda: list[Tensor] | None  # each: (n_cuda_rng_state,)
 
 
 @contextmanager
@@ -212,46 +212,46 @@ def _filtered_kwargs(target: Any, kwargs: dict[str, Any]) -> dict[str, Any]:
 class Boltz2ModelOutput(ModelOutput):
     """Raw Boltz2 inference output with standard AutoModel controls."""
 
-    last_hidden_state: torch.Tensor | None = None
-    hidden_states: tuple[torch.Tensor, ...] | None = None
-    attentions: tuple[torch.Tensor, ...] | None = None
-    pdistogram: torch.Tensor | None = None
-    s: torch.Tensor | None = None
-    z: torch.Tensor | None = None
-    sample_atom_coords: torch.Tensor | None = None
-    diff_token_repr: torch.Tensor | None = None
-    s_conf: torch.Tensor | None = None
-    z_conf: torch.Tensor | None = None
-    pde_logits: torch.Tensor | None = None
-    plddt_logits: torch.Tensor | None = None
-    resolved_logits: torch.Tensor | None = None
-    pde: torch.Tensor | None = None
-    plddt: torch.Tensor | None = None
-    complex_plddt: torch.Tensor | None = None
-    complex_iplddt: torch.Tensor | None = None
-    complex_pde: torch.Tensor | None = None
-    complex_ipde: torch.Tensor | None = None
-    pae_logits: torch.Tensor | None = None
-    pae: torch.Tensor | None = None
-    ptm: torch.Tensor | None = None
-    iptm: torch.Tensor | None = None
-    ligand_iptm: torch.Tensor | None = None
-    protein_iptm: torch.Tensor | None = None
+    last_hidden_state: torch.Tensor | None = None  # (b, t, d_s) or None
+    hidden_states: tuple[torch.Tensor, ...] | None = None  # each shape traced at source
+    attentions: tuple[torch.Tensor, ...] | None = None  # unsupported; always None
+    pdistogram: torch.Tensor | None = None  # (b, t, t, n_distogram, n_bin) or None
+    s: torch.Tensor | None = None  # (b, t, d_s) or None
+    z: torch.Tensor | None = None  # (b, t, t, d_z) or None
+    sample_atom_coords: torch.Tensor | None = None  # (b * m, a, 3) or None
+    diff_token_repr: torch.Tensor | None = None  # (...) or None
+    s_conf: torch.Tensor | None = None  # (...) or None
+    z_conf: torch.Tensor | None = None  # (...) or None
+    pde_logits: torch.Tensor | None = None  # (...) or None
+    plddt_logits: torch.Tensor | None = None  # (...) or None
+    resolved_logits: torch.Tensor | None = None  # (...) or None
+    pde: torch.Tensor | None = None  # (...) or None
+    plddt: torch.Tensor | None = None  # (...) or None
+    complex_plddt: torch.Tensor | None = None  # (...) or None
+    complex_iplddt: torch.Tensor | None = None  # (...) or None
+    complex_pde: torch.Tensor | None = None  # (...) or None
+    complex_ipde: torch.Tensor | None = None  # (...) or None
+    pae_logits: torch.Tensor | None = None  # (...) or None
+    pae: torch.Tensor | None = None  # (...) or None
+    ptm: torch.Tensor | None = None  # (...) or None
+    iptm: torch.Tensor | None = None  # (...) or None
+    ligand_iptm: torch.Tensor | None = None  # (...) or None
+    protein_iptm: torch.Tensor | None = None  # (...) or None
     pair_chains_iptm: Any = None
 
 
 @dataclass
 class Boltz2StructureOutput(ModelOutput):
-    sample_atom_coords: torch.Tensor | None = None
-    atom_pad_mask: torch.Tensor | None = None
-    plddt: torch.Tensor | None = None
-    confidence_score: torch.Tensor | None = None
-    complex_plddt: torch.Tensor | None = None
-    iptm: torch.Tensor | None = None
-    ptm: torch.Tensor | None = None
+    sample_atom_coords: torch.Tensor | None = None  # (m, a_p, 3) or None
+    atom_pad_mask: torch.Tensor | None = None  # (a_p,) or None
+    plddt: torch.Tensor | None = None  # (m, a_p) or None
+    confidence_score: torch.Tensor | None = None  # (m,) or None
+    complex_plddt: torch.Tensor | None = None  # (m,) or None
+    iptm: torch.Tensor | None = None  # (m,) or None
+    ptm: torch.Tensor | None = None  # (m,) or None
     sequence: str | None = None
     structure_template: ProteinStructureTemplate | None = None
-    raw_output: dict[str, torch.Tensor] | None = None
+    raw_output: dict[str, torch.Tensor] | None = None  # values keep model-output shapes
     seed: int | None = None
 
 
@@ -573,41 +573,49 @@ class Boltz2InferenceCore(nn.Module):
         run_confidence_sequentially: bool = True,
         detach_confidence: bool = True,
     ) -> dict[str, Tensor]:
-        s_inputs = self.input_embedder(feats)
-        s_init = self.s_init(s_inputs)
-
-        z_init = self.z_init_1(s_inputs)[:, :, None] + self.z_init_2(s_inputs)[:, None, :]
-        relative_position_encoding = self.rel_pos(feats)
-        z_init = z_init + relative_position_encoding
-        z_init = z_init + self.token_bonds(feats["token_bonds"].float())
+        # b is the batch size, t the token count, a the padded atom count,
+        # d_s the token width, d_z the pair width, and m the diffusion multiplicity.
+        s_inputs = self.input_embedder(feats)  # (b, t, d_s)
+        s_init = self.s_init(s_inputs)  # (b, t, d_s)
+
+        z_init = (
+            self.z_init_1(s_inputs)[:, :, None] + self.z_init_2(s_inputs)[:, None, :]
+        )  # (b, t, t, d_z)
+        relative_position_encoding = self.rel_pos(feats)  # (b, t, t, d_z)
+        z_init = z_init + relative_position_encoding  # (b, t, t, d_z)
+        z_init = z_init + self.token_bonds(
+            feats["token_bonds"].float()
+        )  # (b, t, t, d_z)
         if self.bond_type_feature:
-            z_init = z_init + self.token_bonds_type(feats["type_bonds"].long())
-        z_init = z_init + self.contact_conditioning(feats)
+            z_init = z_init + self.token_bonds_type(
+                feats["type_bonds"].long()
+            )  # (b, t, t, d_z)
+        z_init = z_init + self.contact_conditioning(feats)  # (b, t, t, d_z)
 
-        s = torch.zeros_like(s_init)
-        z = torch.zeros_like(z_init)
-        mask = feats["token_pad_mask"].float()
-        pair_mask = mask[:, :, None] * mask[:, None, :]
+        s = torch.zeros_like(s_init)  # (b, t, d_s)
+        z = torch.zeros_like(z_init)  # (b, t, t, d_z)
+        mask = feats["token_pad_mask"].float()  # (b, t)
+        pair_mask = mask[:, :, None] * mask[:, None, :]  # (b, t, t)
 
         if self.run_trunk_and_structure:
             for _ in range(recycling_steps + 1):
-                s = s_init + self.s_recycle(self.s_norm(s))
-                z = z_init + self.z_recycle(self.z_norm(z))
+                s = s_init + self.s_recycle(self.s_norm(s))  # (b, t, d_s)
+                z = z_init + self.z_recycle(self.z_norm(z))  # (b, t, t, d_z)
                 z = z + self.msa_module(
                     z,
                     s_inputs,
                     feats,
                     use_kernels=self.use_kernels,
-                )
+                )  # (b, t, t, d_z)
                 s, z = self.pairformer_module(
                     s,
                     z,
                     mask=mask,
                     pair_mask=pair_mask,
                     use_kernels=self.use_kernels,
-                )
+                )  # (b, t, d_s), (b, t, t, d_z)
 
-        pdistogram = self.distogram_module(z)
+        pdistogram = self.distogram_module(z)  # (b, t, t, n_distogram, n_bin)
         output: dict[str, Tensor] = {
             "pdistogram": pdistogram,
             "s": s,
@@ -622,7 +630,7 @@ class Boltz2InferenceCore(nn.Module):
                     relative_position_encoding=relative_position_encoding,
                     feats=feats,
                 )
-            )
+            )  # q/c: (b, a, d_a); biases retain their attention-specific axes
             diffusion_conditioning = {
                 "q": q,
                 "c": c,
@@ -642,29 +650,31 @@ class Boltz2InferenceCore(nn.Module):
                     max_parallel_samples=max_parallel_samples,
                     steering_args=self.steering_args,
                     diffusion_conditioning=diffusion_conditioning,
-                )
+                )  # tensor values include sample_atom_coords: (b * m, a, 3)
             output.update(struct_out)
 
         if self.confidence_prediction:
             if self.skip_run_structure:
-                x_pred = feats["coords"].repeat_interleave(diffusion_samples, 0)
+                x_pred = feats["coords"].repeat_interleave(
+                    diffusion_samples, 0
+                )  # (b * m, ..., a, 3)
             else:
                 if "sample_atom_coords" not in output:
                     raise RuntimeError("Structure sampling did not produce sample_atom_coords.")
-                x_pred = output["sample_atom_coords"]
+                x_pred = output["sample_atom_coords"]  # (b * m, a, 3)
 
             if detach_confidence:
-                s_inputs_c = s_inputs.detach()
-                s_c = s.detach()
-                z_c = z.detach()
-                x_pred_c = x_pred.detach()
-                pdist_c = output["pdistogram"][:, :, :, 0].detach()
+                s_inputs_c = s_inputs.detach()  # (b, t, d_s)
+                s_c = s.detach()  # (b, t, d_s)
+                z_c = z.detach()  # (b, t, t, d_z)
+                x_pred_c = x_pred.detach()  # (b * m, ..., a, 3)
+                pdist_c = output["pdistogram"][:, :, :, 0].detach()  # (b, t, t, n_bin)
             else:
-                s_inputs_c = s_inputs
-                s_c = s
-                z_c = z
-                x_pred_c = x_pred
-                pdist_c = output["pdistogram"][:, :, :, 0]
+                s_inputs_c = s_inputs  # (b, t, d_s)
+                s_c = s  # (b, t, d_s)
+                z_c = z  # (b, t, t, d_z)
+                x_pred_c = x_pred  # (b * m, ..., a, 3)
+                pdist_c = output["pdistogram"][:, :, :, 0]  # (b, t, t, n_bin)
 
             output.update(
                 self.confidence_module(
@@ -680,7 +690,7 @@ class Boltz2InferenceCore(nn.Module):
                 )
             )
 
-        return output
+        return output  # named tensors retain the shapes traced above
 
 
 class Boltz2Model(PreTrainedModel):
@@ -698,6 +708,7 @@ class Boltz2Model(PreTrainedModel):
         return
 
     def _detied_state_dict(self) -> dict[str, Tensor]:
+        # Every state tensor has its parameter-defined shape (...).
         raw_state = self.state_dict()
         seen_ptrs: dict[int, str] = {}
         out: dict[str, Tensor] = {}
@@ -705,13 +716,13 @@ class Boltz2Model(PreTrainedModel):
             if torch.is_tensor(tensor):
                 ptr = tensor.untyped_storage().data_ptr()
                 if ptr in seen_ptrs:
-                    out[key] = tensor.clone()
+                    out[key] = tensor.clone()  # (...)
                 else:
                     seen_ptrs[ptr] = key
-                    out[key] = tensor
+                    out[key] = tensor  # (...)
             else:
-                out[key] = tensor
-        return out
+                out[key] = tensor  # (...)
+        return out  # each tensor: (...)
 
     def save_pretrained(self, save_directory: str, **kwargs: Any) -> None:
         if "safe_serialization" not in kwargs:
@@ -757,7 +768,7 @@ class Boltz2Model(PreTrainedModel):
             checkpoint_path,
             map_location=map_location,
             weights_only=False,
-        )
+        )  # mapping values include state tensors with parameter-defined shapes (...)
         if not isinstance(checkpoint, dict):
             raise TypeError("Checkpoint must deserialize to a dictionary.")
         _require_key(checkpoint, "hyper_parameters")
@@ -766,7 +777,7 @@ class Boltz2Model(PreTrainedModel):
         hparams = checkpoint["hyper_parameters"]
         if not isinstance(hparams, dict):
             raise TypeError("Checkpoint hyper_parameters must be a dictionary.")
-        state_dict = checkpoint["state_dict"]
+        state_dict = checkpoint["state_dict"]  # each tensor: (...)
         if not isinstance(state_dict, dict):
             raise TypeError("Checkpoint state_dict must be a dictionary.")
 
@@ -778,7 +789,7 @@ class Boltz2Model(PreTrainedModel):
             default_diffusion_samples=default_diffusion_samples,
         )
         model = cls(config)
-        cleaned = _state_dict_without_wrappers(state_dict)
+        cleaned = _state_dict_without_wrappers(state_dict)  # each tensor: (...)
         target_keys = set(model.core.state_dict().keys())
         for key in target_keys:
             if ".attention.norm_s." in key:
@@ -789,7 +800,7 @@ class Boltz2Model(PreTrainedModel):
         filtered: dict[str, Tensor] = {}
         for key, value in cleaned.items():
             if key in target_keys:
-                filtered[key] = value
+                filtered[key] = value  # (...)
 
         missing = sorted(target_keys.difference(filtered.keys()))
         if missing:
@@ -826,6 +837,7 @@ class Boltz2Model(PreTrainedModel):
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
     ) -> Boltz2ModelOutput | tuple[Any, ...]:
+        # Feature shapes follow build_boltz2_features; b and t are batch and token counts.
         output_attentions = (
             self.config.output_attentions if output_attentions is None else output_attentions
         )
@@ -854,9 +866,9 @@ class Boltz2Model(PreTrainedModel):
             max_parallel_samples=max_parallel_samples,
             run_confidence_sequentially=run_confidence_sequentially,
             detach_confidence=detach_confidence,
-        )
-        token_state = raw_output.get("s")
-        pair_state = raw_output.get("z")
+        )  # named tensors use the Boltz2InferenceCore.forward shapes
+        token_state = raw_output.get("s")  # (b, t, d_s) or None
+        pair_state = raw_output.get("z")  # (b, t, t, d_z) or None
         model_output = Boltz2ModelOutput(
             last_hidden_state=token_state,
             hidden_states=(token_state, pair_state)
@@ -871,16 +883,19 @@ class Boltz2Model(PreTrainedModel):
         feats: dict[str, Tensor],
         float_dtype: torch.dtype = torch.float32,
     ) -> dict[str, Tensor]:
+        # Each feature tensor keeps its existing shape (...).
         moved: dict[str, Tensor] = {}
         for key, value in feats.items():
             if torch.is_tensor(value):
                 if value.is_floating_point():
-                    moved[key] = value.to(device=self.device, dtype=float_dtype)
+                    moved[key] = value.to(
+                        device=self.device, dtype=float_dtype
+                    )  # same shape (...)
                 else:
-                    moved[key] = value.to(device=self.device)
+                    moved[key] = value.to(device=self.device)  # same shape (...)
             else:
-                moved[key] = value
-        return moved
+                moved[key] = value  # same shape (...)
+        return moved  # each tensor: same shape (...)
 
     def predict_structure(
         self,
@@ -914,8 +929,10 @@ class Boltz2Model(PreTrainedModel):
                 atoms_per_window_queries=(
                     self.core.input_embedder.atom_encoder.atoms_per_window_queries
                 ),
-            )
-            feats = self._to_model_device(feats, float_dtype=torch.float32)
+            )  # feature shapes are traced in build_boltz2_features
+            feats = self._to_model_device(
+                feats, float_dtype=torch.float32
+            )  # unchanged feature shapes
             autocast_context = (
                 torch.autocast(device_type="cuda", dtype=torch.bfloat16)
                 if self.device.type == "cuda"
@@ -930,32 +947,42 @@ class Boltz2Model(PreTrainedModel):
                     max_parallel_samples=max_parallel_samples,
                     run_confidence_sequentially=run_confidence_sequentially,
                     return_dict=True,
-                )
+                )  # named tensors use Boltz2InferenceCore.forward shapes
 
-        sample_atom_coords_value = output.get("sample_atom_coords")
+        sample_atom_coords_value = output.get(
+            "sample_atom_coords"
+        )  # (m, a_p, 3) or absent
         if not isinstance(sample_atom_coords_value, torch.Tensor):
             raise RuntimeError("Boltz2 structure sampling did not return coordinate tensors.")
-        sample_atom_coords = sample_atom_coords_value.detach().cpu()
-        non_finite_mask = torch.logical_not(torch.isfinite(sample_atom_coords))
+        sample_atom_coords = sample_atom_coords_value.detach().cpu()  # (m, a_p, 3)
+        non_finite_mask = torch.logical_not(
+            torch.isfinite(sample_atom_coords)
+        )  # (m, a_p, 3)
         if torch.any(non_finite_mask):
             raise RuntimeError(
                 "sample_atom_coords contains non-finite values. "
                 f"Non-finite count: {int(non_finite_mask.sum().item())}"
             )
-        atom_pad_mask = feats["atom_pad_mask"][0].detach().cpu()
-        plddt = output["plddt"].detach().cpu() if "plddt" in output else None
+        atom_pad_mask = feats["atom_pad_mask"][0].detach().cpu()  # (a_p,)
+        plddt = (
+            output["plddt"].detach().cpu() if "plddt" in output else None
+        )  # (m, a_p) or None
         complex_plddt = (
             output["complex_plddt"].detach().cpu() if "complex_plddt" in output else None
-        )
-        iptm = output["iptm"].detach().cpu() if "iptm" in output else None
-        ptm = output["ptm"].detach().cpu() if "ptm" in output else None
-
-        confidence_score = None
+        )  # (m,) or None
+        iptm = (
+            output["iptm"].detach().cpu() if "iptm" in output else None
+        )  # (m,) or None
+        ptm = (
+            output["ptm"].detach().cpu() if "ptm" in output else None
+        )  # (m,) or None
+
+        confidence_score = None  # (m,) or None
         if (complex_plddt is not None) and (iptm is not None) and (ptm is not None):
             if torch.allclose(iptm, torch.zeros_like(iptm)):
-                confidence_score = (4 * complex_plddt + ptm) / 5
+                confidence_score = (4 * complex_plddt + ptm) / 5  # (m,)
             else:
-                confidence_score = (4 * complex_plddt + iptm) / 5
+                confidence_score = (4 * complex_plddt + iptm) / 5  # (m,)
 
         return Boltz2StructureOutput(
             sample_atom_coords=sample_atom_coords,

fastplms/models/boltz/vb_layers_attention.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 from collections.abc import Callable, MutableMapping
-
 from einops.layers.torch import Rearrange
 from torch import Tensor, nn
 
@@ -51,12 +50,13 @@ class AttentionPairBias(nn.Module):
         pair_states: Tensor,
         model_cache: MutableMapping[str, Tensor] | None,
     ) -> Tensor:
+        # pair_states: (b, l_q, l_k, d_z)
         if model_cache is not None and "z" in model_cache:
-            return model_cache["z"]
-        pair_bias = self.proj_z(pair_states)
+            return model_cache["z"]  # (b, h, l_q, l_k)
+        pair_bias = self.proj_z(pair_states)  # (b, h, l_q, l_k)
         if model_cache is not None:
             model_cache["z"] = pair_bias
-        return pair_bias
+        return pair_bias  # (b, h, l_q, l_k)
 
     def forward(
         self,
@@ -70,21 +70,26 @@ class AttentionPairBias(nn.Module):
     ) -> Tensor:
         """Transform S with shapes ``S: (b, l_q, d)`` and ``Z: (b, l, l, d_z)``."""
 
-        sequence_states = self.norm_s(s) if self.initial_norm else s
+        # s: (b, l_q, d); z: (b_z, l_q, l_k, d_z); mask: (b_k, l_k).
+        sequence_states = self.norm_s(s) if self.initial_norm else s  # (b, l_q, d)
         if to_keys is not None:
-            key_states = to_keys(sequence_states)
-            key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)
+            key_states = to_keys(sequence_states)  # (b, l_k, d)
+            key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)  # (b_k, l_k)
         else:
-            key_states = sequence_states if k_in is None else k_in
-            key_mask = mask
+            key_states = sequence_states if k_in is None else k_in  # (b, l_k, d)
+            key_mask = mask  # (b_k, l_k)
 
-        query = reshape_heads(self.proj_q(sequence_states), self.num_heads)
-        key = reshape_heads(self.proj_k(key_states), self.num_heads)
-        value = reshape_heads(self.proj_v(key_states), self.num_heads)
+        query = reshape_heads(
+            self.proj_q(sequence_states), self.num_heads
+        )  # (b, l_q, h, d_h)
+        key = reshape_heads(self.proj_k(key_states), self.num_heads)  # (b, l_k, h, d_h)
+        value = reshape_heads(
+            self.proj_v(key_states), self.num_heads
+        )  # (b, l_k, h, d_h)
         pair_bias = self._resolve_pair_bias(z, model_cache).repeat_interleave(
             multiplicity,
             dim=0,
-        )
+        )  # (b, h, l_q, l_k)
         attended = pair_biased_attention(
             query,
             key,
@@ -92,7 +97,7 @@ class AttentionPairBias(nn.Module):
             pair_bias,
             key_mask,
             self.inf,
-        )
-        attended = attended.reshape(s.shape[0], -1, self.c_s)
-        gate = self.proj_g(sequence_states).sigmoid()
-        return self.proj_o(gate * attended)
+        )  # (b, l_q, h, d_h)
+        attended = attended.reshape(s.shape[0], -1, self.c_s)  # (b, l_q, d)
+        gate = self.proj_g(sequence_states).sigmoid()  # (b, l_q, d)
+        return self.proj_o(gate * attended)  # (b, l_q, d)

fastplms/models/boltz/vb_layers_attentionv2.py

@@ -56,10 +56,13 @@ class AttentionPairBias(nn.Module):
     ) -> Tensor:
         """Transform S with K and pair input Z, returning shape ``(b, l_q, d)``."""
 
-        query = reshape_heads(self.proj_q(s), self.num_heads)
-        key = reshape_heads(self.proj_k(k_in), self.num_heads)
-        value = reshape_heads(self.proj_v(k_in), self.num_heads)
-        pair_bias = self.proj_z(z).repeat_interleave(multiplicity, dim=0)
+        # s: (b, l_q, d); k_in: (b, l_k, d); z: (b_z, l_q, l_k, d_z or h).
+        query = reshape_heads(self.proj_q(s), self.num_heads)  # (b, l_q, h, d_h)
+        key = reshape_heads(self.proj_k(k_in), self.num_heads)  # (b, l_k, h, d_h)
+        value = reshape_heads(self.proj_v(k_in), self.num_heads)  # (b, l_k, h, d_h)
+        pair_bias = self.proj_z(z).repeat_interleave(
+            multiplicity, dim=0
+        )  # (b, h, l_q, l_k)
         attended = pair_biased_attention(
             query,
             key,
@@ -67,7 +70,7 @@ class AttentionPairBias(nn.Module):
             pair_bias,
             mask,
             self.inf,
-        )
-        attended = attended.reshape(s.shape[0], -1, self.c_s)
-        gate = self.proj_g(s).sigmoid()
-        return self.proj_o(gate * attended)
+        )  # (b, l_q, h, d_h)
+        attended = attended.reshape(s.shape[0], -1, self.c_s)  # (b, l_q, d)
+        gate = self.proj_g(s).sigmoid()  # (b, l_q, d)
+        return self.proj_o(gate * attended)  # (b, l_q, d)

fastplms/models/boltz/vb_layers_confidence_utils.py

@@ -11,12 +11,15 @@ from . import vb_const as const
 def compute_collinear_mask(v1: torch.Tensor, v2: torch.Tensor) -> torch.Tensor:
     """Mark nondegenerate vector pairs whose angle defines a stable frame."""
 
-    norm1 = torch.norm(v1, dim=1, keepdim=True)
-    norm2 = torch.norm(v2, dim=1, keepdim=True)
-    unit1 = v1 / (norm1 + 1e-6)
-    unit2 = v2 / (norm2 + 1e-6)
-    separated = torch.abs(torch.sum(unit1 * unit2, dim=1)) < 0.9063
-    return separated & (norm1.reshape(-1) > 1e-2) & (norm2.reshape(-1) > 1e-2)
+    # v1/v2: (n, 3)
+    norm1 = torch.norm(v1, dim=1, keepdim=True)  # (n, 1)
+    norm2 = torch.norm(v2, dim=1, keepdim=True)  # (n, 1)
+    unit1 = v1 / (norm1 + 1e-6)  # (n, 3)
+    unit2 = v2 / (norm2 + 1e-6)  # (n, 3)
+    separated = torch.abs(torch.sum(unit1 * unit2, dim=1)) < 0.9063  # (n,)
+    return (
+        separated & (norm1.reshape(-1) > 1e-2) & (norm2.reshape(-1) > 1e-2)
+    )  # (n,)
 
 
 def _atom_chain_ids(feats: dict[str, torch.Tensor]) -> torch.Tensor:
@@ -24,7 +27,7 @@ def _atom_chain_ids(feats: dict[str, torch.Tensor]) -> torch.Tensor:
         return torch.bmm(
             feats["atom_to_token"].float(),
             feats["asym_id"].unsqueeze(-1).float(),
-        ).squeeze(-1)
+        ).squeeze(-1)  # (b, a)
 
 
 def _replace_nonpolymer_frames(
@@ -36,17 +39,19 @@ def _replace_nonpolymer_frames(
     *,
     inference: bool,
 ) -> None:
-    token_chain_ids = feats["asym_id"]
+    # coordinates: (b, m, a, 3); frame_indices: (b, m, t_or_a, 3).
+    token_chain_ids = feats["asym_id"]  # (b, t)
     for batch_index, sample_coordinates in enumerate(coordinates):
+        # sample_coordinates: (m, a, 3)
         token_offset = 0
         atom_offset = 0
         for chain_id in torch.unique(token_chain_ids[batch_index]):
-            token_mask = (token_chain_ids[batch_index] == chain_id) * feats["token_pad_mask"][
-                batch_index
-            ]
-            atom_mask = (atom_chain_ids[batch_index] == chain_id) * feats["atom_pad_mask"][
-                batch_index
-            ]
+            token_mask = (token_chain_ids[batch_index] == chain_id) * feats[
+                "token_pad_mask"
+            ][batch_index]  # (t,)
+            atom_mask = (atom_chain_ids[batch_index] == chain_id) * feats[
+                "atom_pad_mask"
+            ][batch_index]  # (a,)
             token_count = int(token_mask.sum().item())
             atom_count = int(atom_mask.sum().item())
             is_nonpolymer = (
@@ -62,21 +67,27 @@ def _replace_nonpolymer_frames(
                         else resolved_mask[batch_index]
                     )
                 )
-                chain_atom_mask = atom_mask.bool()
-                chain_coordinates = sample_coordinates[:, chain_atom_mask]
-                differences = chain_coordinates[:, None, :, :] - chain_coordinates[:, :, None, :]
-                distances = differences.square().sum(dim=-1) ** 0.5
-                valid = chain_resolved[chain_atom_mask]
-                invalid_pairs = 1 - (valid[None, :] * valid[:, None]).to(torch.float32)
-                invalid_pairs[invalid_pairs == 1] = torch.inf
-                nearest = torch.sort(distances + invalid_pairs, dim=2).indices
+                chain_atom_mask = atom_mask.bool()  # (a,)
+                chain_coordinates = sample_coordinates[:, chain_atom_mask]  # (m, a_c, 3)
+                differences = (
+                    chain_coordinates[:, None, :, :] - chain_coordinates[:, :, None, :]
+                )  # (m, a_c, a_c, 3)
+                distances = differences.square().sum(dim=-1) ** 0.5  # (m, a_c, a_c)
+                valid = chain_resolved[chain_atom_mask]  # (a_c,)
+                invalid_pairs = 1 - (valid[None, :] * valid[:, None]).to(
+                    torch.float32
+                )  # (a_c, a_c)
+                invalid_pairs[invalid_pairs == 1] = torch.inf  # (a_c, a_c)
+                nearest = torch.sort(
+                    distances + invalid_pairs, dim=2
+                ).indices  # (m, a_c, a_c)
                 frames = (
                     torch.cat(
                         (nearest[:, :, 1:2], nearest[:, :, 0:1], nearest[:, :, 2:3]),
                         dim=2,
                     )
                     + atom_offset
-                )
+                )  # (m, a_c, 3)
                 frame_indices[
                     batch_index,
                     :,
@@ -102,12 +113,18 @@ def compute_frame_pred(
     atom's three nearest resolved neighbors.
     """
 
-    atom_chain_ids = _atom_chain_ids(feats)
+    atom_chain_ids = _atom_chain_ids(feats)  # (b, a)
     expanded_batch, _, _ = pred_atom_coords.shape
     base_batch = expanded_batch // multiplicity
-    coordinates = pred_atom_coords.reshape(base_batch, multiplicity, -1, 3)
-    frame_indices = frames_idx_true.clone().repeat_interleave(multiplicity, 0)
-    frame_indices = frame_indices.reshape(base_batch, multiplicity, -1, 3)
+    coordinates = pred_atom_coords.reshape(
+        base_batch, multiplicity, -1, 3
+    )  # (b, m, a, 3)
+    frame_indices = frames_idx_true.clone().repeat_interleave(
+        multiplicity, 0
+    )  # (b * m, 1, t, 3) or (b * m, t, 3)
+    frame_indices = frame_indices.reshape(
+        base_batch, multiplicity, -1, 3
+    )  # (b, m, t, 3)
     _replace_nonpolymer_frames(
         coordinates,
         frame_indices,
@@ -117,18 +134,25 @@ def compute_frame_pred(
         inference=inference,
     )
 
-    batch_indices = torch.arange(base_batch, device=frame_indices.device)[:, None, None, None]
-    sample_indices = torch.arange(multiplicity, device=frame_indices.device)[None, :, None, None]
+    batch_indices = torch.arange(
+        base_batch, device=frame_indices.device
+    )[:, None, None, None]  # (b, 1, 1, 1)
+    sample_indices = torch.arange(
+        multiplicity, device=frame_indices.device
+    )[None, :, None, None]  # (1, m, 1, 1)
     frame_coordinates = coordinates[
         batch_indices,
         sample_indices,
         frame_indices,
-    ].reshape(-1, 3, 3)
+    ].reshape(-1, 3, 3)  # (b * m * t, 3, 3)
     valid_frames = compute_collinear_mask(
         frame_coordinates[:, 1] - frame_coordinates[:, 0],
         frame_coordinates[:, 1] - frame_coordinates[:, 2],
-    ).reshape(base_batch, multiplicity, -1)
-    return frame_indices, valid_frames * feats["token_pad_mask"][:, None, :]
+    ).reshape(base_batch, multiplicity, -1)  # (b, m, t)
+    return (
+        frame_indices,
+        valid_frames * feats["token_pad_mask"][:, None, :],
+    )  # (b, m, t, 3), (b, m, t)
 
 
 def compute_aggregated_metric(logits: torch.Tensor, end: float = 1.0) -> torch.Tensor:
@@ -141,25 +165,32 @@ def compute_aggregated_metric(logits: torch.Tensor, end: float = 1.0) -> torch.T
         end,
         bin_width,
         device=logits.device,
-    )
-    probabilities = nn.functional.softmax(logits, dim=-1)
+    )  # (n_bin,)
+    # logits: (..., n_bin)
+    probabilities = nn.functional.softmax(logits, dim=-1)  # (..., n_bin)
     broadcast_shape = (1,) * (probabilities.ndim - 1) + centers.shape
-    return torch.sum(probabilities * centers.view(broadcast_shape), dim=-1)
+    return torch.sum(
+        probabilities * centers.view(broadcast_shape), dim=-1
+    )  # (...)
 
 
 def tm_function(d: torch.Tensor, n_res: torch.Tensor) -> torch.Tensor:
     """Evaluate the TM-score distance kernel."""
 
-    distance_scale = 1.24 * (torch.clip(n_res, min=19) - 15) ** (1 / 3) - 1.8
-    return 1 / (1 + (d / distance_scale) ** 2)
+    # d and n_res broadcast to a shared shape (...).
+    distance_scale = 1.24 * (torch.clip(n_res, min=19) - 15) ** (1 / 3) - 1.8  # (...)
+    return 1 / (1 + (d / distance_scale) ** 2)  # (...)
 
 
 def _maximum_masked_tm(
     expected_tm: torch.Tensor,
     pair_mask: torch.Tensor,
 ) -> torch.Tensor:
-    per_anchor = torch.sum(expected_tm * pair_mask, dim=-1) / (torch.sum(pair_mask, dim=-1) + 1e-5)
-    return torch.max(per_anchor, dim=1).values
+    # expected_tm/pair_mask: (b * m, t, t).
+    per_anchor = torch.sum(expected_tm * pair_mask, dim=-1) / (
+        torch.sum(pair_mask, dim=-1) + 1e-5
+    )  # (b * m, t)
+    return torch.max(per_anchor, dim=1).values  # (b * m,)
 
 
 def compute_ptms(
@@ -176,18 +207,25 @@ def compute_ptms(
 ]:
     """Compute pTM, ipTM, interface-type, and chain-pair confidence scores."""
 
+    # logits: (b * m, t, t, n_bin); x_preds: (b * m, a, 3).
     _, frame_mask = compute_frame_pred(
         x_preds,
         feats["frames_idx"],
         feats,
         multiplicity,
         inference=True,
-    )
-    token_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
-    valid_anchor = frame_mask.reshape(-1, frame_mask.shape[-1])
-    base_pair_mask = valid_anchor[:, :, None] * token_mask[:, None, :] * token_mask[:, :, None]
-    asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)
-    interface_mask = base_pair_mask * (asym_id[:, None, :] != asym_id[:, :, None])
+    )  # frame_mask: (b, m, t)
+    token_mask = feats["token_pad_mask"].repeat_interleave(
+        multiplicity, 0
+    )  # (b * m, t)
+    valid_anchor = frame_mask.reshape(-1, frame_mask.shape[-1])  # (b * m, t)
+    base_pair_mask = (
+        valid_anchor[:, :, None] * token_mask[:, None, :] * token_mask[:, :, None]
+    )  # (b * m, t, t)
+    asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)  # (b * m, t)
+    interface_mask = base_pair_mask * (
+        asym_id[:, None, :] != asym_id[:, :, None]
+    )  # (b * m, t, t)
 
     num_bins = logits.shape[-1]
     pae_centers = torch.arange(
@@ -195,23 +233,31 @@ def compute_ptms(
         32.0,
         32.0 / num_bins,
         device=logits.device,
-    ).unsqueeze(0)
-    n_res = token_mask.sum(dim=-1, keepdim=True)
-    tm_values = tm_function(pae_centers, n_res).unsqueeze(1).unsqueeze(2)
-    expected_tm = torch.sum(nn.functional.softmax(logits, dim=-1) * tm_values, dim=-1)
+    ).unsqueeze(0)  # (1, n_bin)
+    n_res = token_mask.sum(dim=-1, keepdim=True)  # (b * m, 1)
+    tm_values = tm_function(pae_centers, n_res).unsqueeze(1).unsqueeze(
+        2
+    )  # (b * m, 1, 1, n_bin)
+    expected_tm = torch.sum(
+        nn.functional.softmax(logits, dim=-1) * tm_values, dim=-1
+    )  # (b * m, t, t)
 
-    ptm = _maximum_masked_tm(expected_tm, base_pair_mask)
-    iptm = _maximum_masked_tm(expected_tm, interface_mask)
+    ptm = _maximum_masked_tm(expected_tm, base_pair_mask)  # (b * m,)
+    iptm = _maximum_masked_tm(expected_tm, interface_mask)  # (b * m,)
 
-    token_type = feats["mol_type"].repeat_interleave(multiplicity, 0)
-    ligand = (token_type == const.chain_type_ids["NONPOLYMER"]).float()
-    protein = (token_type == const.chain_type_ids["PROTEIN"]).float()
+    token_type = feats["mol_type"].repeat_interleave(multiplicity, 0)  # (b * m, t)
+    ligand = (token_type == const.chain_type_ids["NONPOLYMER"]).float()  # (b * m, t)
+    protein = (token_type == const.chain_type_ids["PROTEIN"]).float()  # (b * m, t)
     ligand_protein = (
         ligand[:, :, None] * protein[:, None, :] + protein[:, :, None] * ligand[:, None, :]
-    )
-    protein_protein = protein[:, :, None] * protein[:, None, :]
-    ligand_iptm = _maximum_masked_tm(expected_tm, interface_mask * ligand_protein)
-    protein_iptm = _maximum_masked_tm(expected_tm, interface_mask * protein_protein)
+    )  # (b * m, t, t)
+    protein_protein = protein[:, :, None] * protein[:, None, :]  # (b * m, t, t)
+    ligand_iptm = _maximum_masked_tm(
+        expected_tm, interface_mask * ligand_protein
+    )  # (b * m,)
+    protein_iptm = _maximum_masked_tm(
+        expected_tm, interface_mask * protein_protein
+    )  # (b * m,)
 
     chain_pair_iptm: dict[int, dict[int, torch.Tensor]] = {}
     for first_chain in torch.unique(asym_id).tolist():

fastplms/models/boltz/vb_layers_dropout.py

@@ -36,5 +36,7 @@ def get_dropout_mask(
 
     probability = float(dropout) if training else 0.0
     sample_shape = _broadcast_mask_shape(z, columnwise=columnwise)
-    keep = torch.rand(sample_shape, dtype=torch.float32, device=z.device) >= probability
-    return keep * (1.0 / (1.0 - probability))
+    keep = (
+        torch.rand(sample_shape, dtype=torch.float32, device=z.device) >= probability
+    )  # (b, 1, n, 1) or (b, n, 1, 1)
+    return keep * (1.0 / (1.0 - probability))  # same broadcast-mask shape

fastplms/models/boltz/vb_layers_initialize.py

@@ -3,11 +3,11 @@
 from __future__ import annotations
 
 import math
-from typing import Literal
-
 import torch
+from typing import Literal
 from torch import Tensor
 
+
 FanMode = Literal["fan_in", "fan_out", "fan_avg"]
 
 
@@ -33,6 +33,7 @@ def trunc_normal_init_(
 ) -> None:
     """Fill W from a normal distribution truncated at two standard deviations."""
 
+    # weights: (d_out, d_in); initialization preserves this shape in place.
     variance = scale / max(1.0, _calculate_fan(weights.shape, fan))
     std = math.sqrt(variance)
     with torch.no_grad():
@@ -42,6 +43,7 @@ def trunc_normal_init_(
 def lecun_normal_init_(weights: Tensor) -> None:
     """Initialize W using fan-in-scaled truncated normal values."""
 
+    # weights: (d_out, d_in), mutated in place.
     trunc_normal_init_(weights)
 
 
@@ -58,6 +60,7 @@ def glorot_uniform_init_(weights: Tensor) -> None:
 
 
 def _fill_(tensor: Tensor, value: float) -> None:
+    # tensor: (...), mutated in place without changing shape.
     with torch.no_grad():
         tensor.fill_(value)
 
@@ -96,3 +99,11 @@ def ipa_point_weights_init_(weights: Tensor) -> None:
     """Initialize W so applying softplus yields one."""
 
     _fill_(weights, 0.541324854612918)
+    # weights: (d_out, d_in), mutated in place.
+    # weights: (d_out, d_in), mutated in place.
+    # weights: (...), mutated in place.
+    # weights: (...), mutated in place.
+    # bias: (...), mutated in place.
+    # bias: (...), mutated in place.
+    # weights: (d_out, d_in), mutated in place.
+    # weights: (...), mutated in place.

fastplms/models/boltz/vb_layers_outer_product_mean.py

@@ -25,19 +25,25 @@ class OuterProductMean(nn.Module):
     def _pair_counts(mask: Tensor) -> Tensor:
         """Count valid MSA observations for each residue pair."""
 
+        # mask: (b, s, l, 1)
         counts: Tensor | None = None
         for start in range(0, mask.shape[1], 64):
-            mask_slice = mask[:, start : start + 64]
-            contribution = (mask_slice[:, :, None, :] * mask_slice[:, :, :, None]).sum(dim=1)
-            counts = contribution if counts is None else counts + contribution
+            mask_slice = mask[:, start : start + 64]  # (b, s_chunk, l, 1)
+            contribution = (
+                mask_slice[:, :, None, :] * mask_slice[:, :, :, None]
+            ).sum(dim=1)  # (b, l, l, 1)
+            counts = (
+                contribution if counts is None else counts + contribution
+            )  # (b, l, l, 1)
         if counts is None:
             raise ValueError("MSA depth must be positive")
-        return counts.clamp(min=1)
+        return counts.clamp(min=1)  # (b, l, l, 1)
 
     @staticmethod
     def _outer_product(left: Tensor, right: Tensor, counts: Tensor) -> Tensor:
-        outer = torch.einsum("bsic,bsjd->bijcd", left, right)
-        return outer.reshape(*outer.shape[:3], -1) / counts
+        # left/right: (b, s, l, d_h); counts: (b, l, l, 1).
+        outer = torch.einsum("bsic,bsjd->bijcd", left, right)  # (b, l, l, d_h, d_h)
+        return outer.reshape(*outer.shape[:3], -1) / counts  # (b, l, l, d_h**2)
 
     def _chunked_projection(
         self,
@@ -47,32 +53,38 @@ class OuterProductMean(nn.Module):
         chunk_size: int,
         target: Tensor,
     ) -> Tensor:
+        # left/right: (b, s, l, d_h); counts: (b, l, l, 1).
         output: Tensor | None = None
         for start in range(0, self.c_hidden, chunk_size):
             stop = min(start + chunk_size, self.c_hidden)
-            outer = self._outer_product(left[..., start:stop], right, counts)
+            outer = self._outer_product(
+                left[..., start:stop], right, counts
+            )  # (b, l, l, d_chunk * d_h)
             weight = self.proj_o.weight[
                 :,
                 start * self.c_hidden : stop * self.c_hidden,
-            ]
-            contribution = outer.to(target) @ weight.T
-            output = contribution if output is None else output + contribution
+            ]  # (d_z, d_chunk * d_h)
+            contribution = outer.to(target) @ weight.T  # (b, l, l, d_z)
+            output = (
+                contribution if output is None else output + contribution
+            )  # (b, l, l, d_z)
         if output is None:
             raise ValueError("c_hidden must be positive")
-        return output + self.proj_o.bias
+        return output + self.proj_o.bias  # (b, l, l, d_z)
 
     def forward(self, m: Tensor, mask: Tensor, chunk_size: int | None = None) -> Tensor:
         """Return pair tensor Z with shape ``(b, l, l, d_z)``."""
 
-        expanded_mask = mask.unsqueeze(-1).to(m)
-        normalized = self.norm(m)
-        left = self.proj_a(normalized) * expanded_mask
-        right = self.proj_b(normalized) * expanded_mask
+        # m: (b, s, l, d_m); mask: (b, s, l).
+        expanded_mask = mask.unsqueeze(-1).to(m)  # (b, s, l, 1)
+        normalized = self.norm(m)  # (b, s, l, d_m)
+        left = self.proj_a(normalized) * expanded_mask  # (b, s, l, d_h)
+        right = self.proj_b(normalized) * expanded_mask  # (b, s, l, d_h)
 
         if chunk_size is not None and not self.training:
             if chunk_size <= 0:
                 raise ValueError("chunk_size must be positive")
-            counts = self._pair_counts(expanded_mask)
+            counts = self._pair_counts(expanded_mask)  # (b, l, l, 1)
             return self._chunked_projection(
                 left,
                 right,
@@ -81,7 +93,11 @@ class OuterProductMean(nn.Module):
                 normalized,
             )
 
-        pair_mask = expanded_mask[:, :, None, :] * expanded_mask[:, :, :, None]
-        counts = pair_mask.sum(dim=1).clamp(min=1)
-        outer = self._outer_product(left.float(), right.float(), counts)
-        return self.proj_o(outer.to(normalized))
+        pair_mask = (
+            expanded_mask[:, :, None, :] * expanded_mask[:, :, :, None]
+        )  # (b, s, l, l, 1)
+        counts = pair_mask.sum(dim=1).clamp(min=1)  # (b, l, l, 1)
+        outer = self._outer_product(
+            left.float(), right.float(), counts
+        )  # (b, l, l, d_h**2)
+        return self.proj_o(outer.to(normalized))  # (b, l, l, d_z)

fastplms/models/boltz/vb_layers_pair_averaging.py

@@ -34,39 +34,54 @@ class PairWeightedAveraging(nn.Module):
         init.final_init_(self.proj_o.weight)
 
     def _attention_weights(self, pair_states: Tensor, mask: Tensor) -> Tensor:
-        logits = self.proj_z(pair_states).permute(0, 3, 1, 2)
-        logits = logits + (1 - mask[:, None]) * -self.inf
-        return torch.softmax(logits, dim=-1)
+        # pair_states: (b, l, l, d_z); mask: (b, l).
+        logits = self.proj_z(pair_states).permute(0, 3, 1, 2)  # (b, h, l, l)
+        logits = logits + (1 - mask[:, None]) * -self.inf  # (b, h, l, l)
+        return torch.softmax(logits, dim=-1)  # (b, h, l, l)
 
     def _all_heads(self, msa_states: Tensor, pair_states: Tensor, mask: Tensor) -> Tensor:
+        # msa_states: (b, s, l, d_m); pair_states: (b, l, l, d_z).
         values = self.proj_m(msa_states).reshape(
             *msa_states.shape[:3],
             self.num_heads,
             self.c_h,
-        )
-        values = values.permute(0, 3, 1, 2, 4)
-        weights = self._attention_weights(pair_states, mask)
-        gate = self.proj_g(msa_states).sigmoid()
-        attended = torch.einsum("bhij,bhsjd->bhsid", weights, values)
-        attended = attended.permute(0, 2, 3, 1, 4)
-        attended = attended.reshape(*attended.shape[:3], self.num_heads * self.c_h)
-        return self.proj_o(gate * attended)
+        )  # (b, s, l, h, d_h)
+        values = values.permute(0, 3, 1, 2, 4)  # (b, h, s, l, d_h)
+        weights = self._attention_weights(pair_states, mask)  # (b, h, l, l)
+        gate = self.proj_g(msa_states).sigmoid()  # (b, s, l, h * d_h)
+        attended = torch.einsum(
+            "bhij,bhsjd->bhsid", weights, values
+        )  # (b, h, s, l, d_h)
+        attended = attended.permute(0, 2, 3, 1, 4)  # (b, s, l, h, d_h)
+        attended = attended.reshape(
+            *attended.shape[:3], self.num_heads * self.c_h
+        )  # (b, s, l, h * d_h)
+        return self.proj_o(gate * attended)  # (b, s, l, d_m)
 
     def _one_head(self, msa_states: Tensor, pair_states: Tensor, mask: Tensor, head: int) -> Tensor:
+        # msa_states: (b, s, l, d_m); pair_states: (b, l, l, d_z).
         start = head * self.c_h
         stop = start + self.c_h
-        value = msa_states @ self.proj_m.weight[start:stop].T
-        value = value.reshape(*value.shape[:3], 1, self.c_h).permute(0, 3, 1, 2, 4)
-        logits = pair_states @ self.proj_z.weight[head : head + 1].T
-        logits = logits.permute(0, 3, 1, 2)
-        weights = torch.softmax(logits + (1 - mask[:, None]) * -self.inf, dim=-1)
-        gate = (msa_states @ self.proj_g.weight[start:stop].T).sigmoid()
-        attended = torch.einsum("bhij,bhsjd->bhsid", weights, value)
+        value = msa_states @ self.proj_m.weight[start:stop].T  # (b, s, l, d_h)
+        value = value.reshape(*value.shape[:3], 1, self.c_h).permute(
+            0, 3, 1, 2, 4
+        )  # (b, 1, s, l, d_h)
+        logits = pair_states @ self.proj_z.weight[head : head + 1].T  # (b, l, l, 1)
+        logits = logits.permute(0, 3, 1, 2)  # (b, 1, l, l)
+        weights = torch.softmax(
+            logits + (1 - mask[:, None]) * -self.inf, dim=-1
+        )  # (b, 1, l, l)
+        gate = (
+            msa_states @ self.proj_g.weight[start:stop].T
+        ).sigmoid()  # (b, s, l, d_h)
+        attended = torch.einsum(
+            "bhij,bhsjd->bhsid", weights, value
+        )  # (b, 1, s, l, d_h)
         attended = attended.permute(0, 2, 3, 1, 4).reshape(
             *msa_states.shape[:3],
             self.c_h,
-        )
-        return (gate * attended) @ self.proj_o.weight[:, start:stop].T
+        )  # (b, s, l, d_h)
+        return (gate * attended) @ self.proj_o.weight[:, start:stop].T  # (b, s, l, d_m)
 
     def forward(
         self,
@@ -77,15 +92,20 @@ class PairWeightedAveraging(nn.Module):
     ) -> Tensor:
         """Return updated M with shape ``(b, s, l, d_m)``."""
 
-        msa_states = self.norm_m(m)
-        pair_states = self.norm_z(z)
+        # m: (b, s, l, d_m); z: (b, l, l, d_z); mask: (b, l).
+        msa_states = self.norm_m(m)  # (b, s, l, d_m)
+        pair_states = self.norm_z(z)  # (b, l, l, d_z)
         if not chunk_heads or self.training:
             return self._all_heads(msa_states, pair_states, mask)
 
         output: Tensor | None = None
         for head in range(self.num_heads):
-            contribution = self._one_head(msa_states, pair_states, mask, head)
-            output = contribution if output is None else output + contribution
+            contribution = self._one_head(
+                msa_states, pair_states, mask, head
+            )  # (b, s, l, d_m)
+            output = (
+                contribution if output is None else output + contribution
+            )  # (b, s, l, d_m)
         if output is None:
             raise ValueError("num_heads must be positive")
-        return output
+        return output  # (b, s, l, d_m)

fastplms/models/boltz/vb_layers_pairformer.py

@@ -2,9 +2,8 @@
 
 from __future__ import annotations
 
-from typing import Any, cast
-
 import torch
+from typing import Any, cast
 from torch import Tensor, nn
 from torch.utils.checkpoint import checkpoint
 
@@ -34,37 +33,42 @@ def _pair_update(
 ) -> Tensor:
     """Apply multiplicative, attention, and transition updates to pair tensor Z."""
 
+    # pair_states: (b, l, l, d_z); pair_mask: (b, l, l).
     multiplication_kernel = use_kernels or use_cuequiv_mul
     attention_kernel = use_kernels or use_cuequiv_attn
-    output = pair_states
+    output = pair_states  # (b, l, l, d_z)
     for update in (modules.tri_mul_out, modules.tri_mul_in):
-        dropout = get_dropout_mask(modules.dropout, output, modules.training)
+        dropout = get_dropout_mask(
+            modules.dropout, output, modules.training
+        )  # broadcastable to (b, l, l, d_z)
         output = output + dropout * update(
             output,
             mask=pair_mask,
             use_kernels=multiplication_kernel,
-        )
+        )  # (b, l, l, d_z)
 
-    dropout = get_dropout_mask(modules.dropout, output, modules.training)
+    dropout = get_dropout_mask(
+        modules.dropout, output, modules.training
+    )  # broadcastable to (b, l, l, d_z)
     output = output + dropout * modules.tri_att_start(
         output,
         mask=pair_mask,
         chunk_size=chunk_size,
         use_kernels=attention_kernel,
-    )
+    )  # (b, l, l, d_z)
     dropout = get_dropout_mask(
         modules.dropout,
         output,
         modules.training,
         columnwise=True,
-    )
+    )  # broadcastable to (b, l, l, d_z)
     output = output + dropout * modules.tri_att_end(
         output,
         mask=pair_mask,
         chunk_size=chunk_size,
         use_kernels=attention_kernel,
-    )
-    return cast(Tensor, output + modules.transition_z(output))
+    )  # (b, l, l, d_z)
+    return cast(Tensor, output + modules.transition_z(output))  # (b, l, l, d_z)
 
 
 def _triangle_chunk_size(pair_states: Tensor, training: bool) -> int | None:
@@ -134,18 +138,22 @@ class PairformerLayer(nn.Module):
             use_kernels,
             use_cuequiv_mul,
             use_cuequiv_attn,
-        )
+        )  # (b, l, l, d_z)
         with torch.autocast("cuda", enabled=False):
-            normalized = self.pre_norm_s(s.float())
+            normalized = self.pre_norm_s(s.float())  # (b, l, d_s)
             sequence_output = s.float() + self.attention(
                 s=normalized,
                 z=pair_output.float(),
                 mask=mask.float(),
                 k_in=normalized,
-            )
-            sequence_output = sequence_output + self.transition_s(sequence_output)
-            sequence_output = cast(Tensor, self.s_post_norm(sequence_output))
-        return sequence_output, pair_output
+            )  # (b, l, d_s)
+            sequence_output = (
+                sequence_output + self.transition_s(sequence_output)
+            )  # (b, l, d_s)
+            sequence_output = cast(
+                Tensor, self.s_post_norm(sequence_output)
+            )  # (b, l, d_s)
+        return sequence_output, pair_output  # (b, l, d_s), (b, l, l, d_z)
 
 
 class PairformerModule(nn.Module):
@@ -199,8 +207,9 @@ class PairformerModule(nn.Module):
     ) -> tuple[Tensor, Tensor]:
         """Return S and Z after every configured pairformer block."""
 
+        # s: (b, l, d_s); z: (b, l, l, d_z).
         chunk_size = _triangle_chunk_size(z, self.training)
-        sequence_output, pair_output = s, z
+        sequence_output, pair_output = s, z  # (b, l, d_s), (b, l, l, d_z)
         for layer in self.layers:
             if self.activation_checkpointing:
                 sequence_output, pair_output = checkpoint(
@@ -212,7 +221,7 @@ class PairformerModule(nn.Module):
                     chunk_size,
                     use_kernels,
                     use_reentrant=False,
-                )
+                )  # (b, l, d_s), (b, l, l, d_z)
             else:
                 sequence_output, pair_output = layer(
                     sequence_output,
@@ -221,8 +230,8 @@ class PairformerModule(nn.Module):
                     pair_mask,
                     chunk_size,
                     use_kernels,
-                )
-        return sequence_output, pair_output
+                )  # (b, l, d_s), (b, l, l, d_z)
+        return sequence_output, pair_output  # (b, l, d_s), (b, l, l, d_z)
 
 
 class PairformerNoSeqLayer(nn.Module):
@@ -320,8 +329,9 @@ class PairformerNoSeqModule(nn.Module):
     ) -> Tensor:
         """Return Z after every configured pair-only block."""
 
+        # z: (b, l, l, d_z); pair_mask: (b, l, l).
         chunk_size = _triangle_chunk_size(z, self.training)
-        output = z
+        output = z  # (b, l, l, d_z)
         for layer in self.layers:
             if self.activation_checkpointing:
                 output = checkpoint(
@@ -331,7 +341,9 @@ class PairformerNoSeqModule(nn.Module):
                     chunk_size,
                     use_kernels,
                     use_reentrant=False,
-                )
+                )  # (b, l, l, d_z)
             else:
-                output = layer(output, pair_mask, chunk_size, use_kernels)
-        return output
+                output = layer(
+                    output, pair_mask, chunk_size, use_kernels
+                )  # (b, l, l, d_z)
+        return output  # (b, l, l, d_z)

fastplms/models/boltz/vb_layers_transition.py

@@ -40,19 +40,26 @@ class Transition(nn.Module):
     def _project_hidden_slice(self, normalized: Tensor, start: int, stop: int) -> Tensor:
         """Return one hidden-width contribution to the output projection."""
 
-        gate = self.silu(torch.matmul(normalized, self.fc1.weight[start:stop].T))
-        value = torch.matmul(normalized, self.fc2.weight[start:stop].T)
-        return torch.matmul(gate * value, self.fc3.weight[:, start:stop].T)
+        # normalized: (..., d); the selected hidden width is d_chunk.
+        gate = self.silu(
+            torch.matmul(normalized, self.fc1.weight[start:stop].T)
+        )  # (..., d_chunk)
+        value = torch.matmul(
+            normalized, self.fc2.weight[start:stop].T
+        )  # (..., d_chunk)
+        return torch.matmul(
+            gate * value, self.fc3.weight[:, start:stop].T
+        )  # (..., d_out)
 
     def forward(self, x: Tensor, chunk_size: int | None = None) -> Tensor:
         """Transform X, optionally accumulating the hidden dimension in chunks."""
 
         # X is the normalized input tensor with shape (..., d).
-        normalized = self.norm(x)
+        normalized = self.norm(x)  # (..., d)
         if chunk_size is None or self.training:
             # H is the gated hidden tensor with shape (..., d_hidden).
             hidden_states = self.silu(self.fc1(normalized)) * self.fc2(normalized)
-            return self.fc3(hidden_states)
+            return self.fc3(hidden_states)  # (..., d_out)
         if chunk_size <= 0:
             raise ValueError("chunk_size must be positive")
 
@@ -62,8 +69,10 @@ class Transition(nn.Module):
                 normalized,
                 start,
                 min(start + chunk_size, self.hidden),
-            )
-            output = contribution if output is None else output + contribution
+            )  # (..., d_out)
+            output = (
+                contribution if output is None else output + contribution
+            )  # (..., d_out)
         if output is None:  # hidden is normally positive; keep malformed configs explicit.
             raise ValueError("hidden must be positive")
-        return output
+        return output  # (..., d_out)

fastplms/models/boltz/vb_layers_triangular_mult.py

@@ -3,14 +3,14 @@
 from __future__ import annotations
 
 import importlib
+import torch
 from importlib.util import find_spec
 from typing import Literal
-
-import torch
 from torch import Tensor, nn
 
 from . import vb_layers_initialize as init
 
+
 TriangleDirection = Literal["incoming", "outgoing"]
 
 
@@ -31,6 +31,7 @@ def kernel_triangular_mult(
 ) -> Tensor:
     """Dispatch the optional cuEquivariance triangle primitive lazily."""
 
+    # x: (b, l, l, d); mask: (b, l, l); returned tensor: (b, l, l, d).
     if (
         find_spec("cuequivariance_torch") is None
         or find_spec("cuequivariance_ops_torch") is None
@@ -81,6 +82,7 @@ class _TriangleMultiplication(nn.Module):
         init.gating_init_(self.g_out.weight)
 
     def _kernel_forward(self, pair_states: Tensor, mask: Tensor) -> Tensor:
+        # pair_states: (b, l, l, d); mask: (b, l, l).
         return kernel_triangular_mult(
             pair_states,
             direction=self.direction,
@@ -94,7 +96,7 @@ class _TriangleMultiplication(nn.Module):
             p_out_weight=self.p_out.weight,
             g_out_weight=self.g_out.weight,
             eps=1e-5,
-        )
+        )  # (b, l, l, d)
 
     def forward(self, x: Tensor, mask: Tensor, use_kernels: bool = False) -> Tensor:
         """Transform pair tensor X with shape ``(b, l, l, d)``."""
@@ -103,12 +105,16 @@ class _TriangleMultiplication(nn.Module):
             return self._kernel_forward(x, mask)
 
         # X_norm is the normalized pair tensor used by the output gate.
-        normalized = self.norm_in(x)
-        projected = self.p_in(normalized) * self.g_in(normalized).sigmoid()
-        projected = projected * mask.unsqueeze(-1)
-        left, right = torch.chunk(projected.float(), 2, dim=-1)
-        combined = torch.einsum(self.equation, left, right)
-        return self.p_out(self.norm_out(combined)) * self.g_out(normalized).sigmoid()
+        normalized = self.norm_in(x)  # (b, l, l, d)
+        projected = (
+            self.p_in(normalized) * self.g_in(normalized).sigmoid()
+        )  # (b, l, l, 2 * d)
+        projected = projected * mask.unsqueeze(-1)  # (b, l, l, 2 * d)
+        left, right = torch.chunk(projected.float(), 2, dim=-1)  # each: (b, l, l, d)
+        combined = torch.einsum(self.equation, left, right)  # (b, l, l, d)
+        return (
+            self.p_out(self.norm_out(combined)) * self.g_out(normalized).sigmoid()
+        )  # (b, l, l, d)
 
 
 class TriangleMultiplicationOutgoing(_TriangleMultiplication):

fastplms/models/boltz/vb_loss_diffusionv2.py

@@ -8,7 +8,6 @@ does not import an upstream runtime package.
 from __future__ import annotations
 
 import warnings
-
 import torch
 import torch.nn.functional as functional
 from einops import einsum
@@ -18,10 +17,11 @@ def _weighted_centroid(
     coordinates: torch.Tensor,
     weights: torch.Tensor,
 ) -> torch.Tensor:
+    # coordinates: (..., n, 3); weights: (..., n, 1).
     return (coordinates * weights).sum(dim=-2, keepdim=True) / weights.sum(
         dim=-2,
         keepdim=True,
-    )
+    )  # (..., 1, 3)
 
 
 def _warn_if_alignment_is_ambiguous(
@@ -62,25 +62,25 @@ def weighted_rigid_align(
 
     output_shape = torch.broadcast_shapes(true_coords.shape, pred_coords.shape)
     *batch_shape, num_points, coordinate_dim = output_shape
-    point_weights = (mask * weights).unsqueeze(-1)
+    point_weights = (mask * weights).unsqueeze(-1)  # (..., n, 1)
 
-    true_centroid = _weighted_centroid(true_coords, point_weights)
-    pred_centroid = _weighted_centroid(pred_coords, point_weights)
-    true_centered = true_coords - true_centroid
-    pred_centered = pred_coords - pred_centroid
+    true_centroid = _weighted_centroid(true_coords, point_weights)  # (..., 1, 3)
+    pred_centroid = _weighted_centroid(pred_coords, point_weights)  # (..., 1, 3)
+    true_centered = true_coords - true_centroid  # (..., n, 3)
+    pred_centered = pred_coords - pred_centroid  # (..., n, 3)
 
     covariance = einsum(
         point_weights * pred_centered,
         true_centered,
         "... n i, ... n j -> ... i j",
-    )
+    )  # (..., 3, 3)
     original_dtype = covariance.dtype
-    covariance_fp32 = covariance.to(torch.float32)
+    covariance_fp32 = covariance.to(torch.float32)  # (..., 3, 3)
     left_vectors, singular_values, right_vectors_h = torch.linalg.svd(
         covariance_fp32,
         driver="gesvd" if covariance_fp32.is_cuda else None,
-    )
-    right_vectors = right_vectors_h.mH
+    )  # left/right: (..., 3, 3); singular_values: (..., 3)
+    right_vectors = right_vectors_h.mH  # (..., 3, 3)
     _warn_if_alignment_is_ambiguous(
         mask,
         singular_values,
@@ -92,23 +92,25 @@ def weighted_rigid_align(
         "... i j, ... k j -> ... i k",
         left_vectors,
         right_vectors,
-    ).to(torch.float32)
+    ).to(torch.float32)  # (..., 3, 3)
     orientation = torch.eye(
         coordinate_dim,
         dtype=covariance_fp32.dtype,
         device=covariance.device,
-    )[None].repeat(*batch_shape, 1, 1)
-    orientation[..., -1, -1] = torch.det(preliminary_rotation)
+    )[None].repeat(*batch_shape, 1, 1)  # (..., 3, 3)
+    orientation[..., -1, -1] = torch.det(preliminary_rotation)  # (...)
     rotation = einsum(
         left_vectors,
         orientation,
         right_vectors,
         "... i j, ... j k, ... l k -> ... i l",
-    ).to(original_dtype)
+    ).to(original_dtype)  # (..., 3, 3)
 
-    aligned = einsum(true_centered, rotation, "... n i, ... j i -> ... n j") + pred_centroid
+    aligned = (
+        einsum(true_centered, rotation, "... n i, ... j i -> ... n j") + pred_centroid
+    )  # (..., n, 3)
     aligned.detach_()
-    return aligned
+    return aligned  # (..., n, 3)
 
 
 def _smooth_lddt_for_example(
@@ -120,33 +122,38 @@ def _smooth_lddt_for_example(
     nucleic_acid_cutoff: float,
     other_cutoff: float,
 ) -> torch.Tensor:
-    true_distances = torch.cdist(true_coords, true_coords)
-    nucleotide_rows = is_nucleotide.bool().unsqueeze(-1).expand_as(true_distances)
+    # pred_coords/true_coords: (n, 3); is_nucleotide/coords_mask: (n,).
+    true_distances = torch.cdist(true_coords, true_coords)  # (n, n)
+    nucleotide_rows = is_nucleotide.bool().unsqueeze(-1).expand_as(
+        true_distances
+    )  # (n, n)
     pair_mask = torch.where(
         nucleotide_rows,
         true_distances < nucleic_acid_cutoff,
         true_distances < other_cutoff,
-    )
+    )  # (n, n)
     pair_mask &= ~torch.eye(
         pred_coords.shape[0],
         dtype=torch.bool,
         device=pred_coords.device,
-    )
-    coordinate_rows = coords_mask.bool()
-    pair_mask &= coordinate_rows.unsqueeze(-1)
-    pair_mask &= coordinate_rows.unsqueeze(-2)
-
-    pair_indices = pair_mask.nonzero()
-    true_pair_distances = true_distances[pair_indices[:, 0], pair_indices[:, 1]]
+    )  # (n, n)
+    coordinate_rows = coords_mask.bool()  # (n,)
+    pair_mask &= coordinate_rows.unsqueeze(-1)  # (n, n)
+    pair_mask &= coordinate_rows.unsqueeze(-2)  # (n, n)
+
+    pair_indices = pair_mask.nonzero()  # (n_pair, 2)
+    true_pair_distances = true_distances[
+        pair_indices[:, 0], pair_indices[:, 1]
+    ]  # (n_pair,)
     pred_pair_distances = functional.pairwise_distance(
         pred_coords[pair_indices[:, 0]],
         pred_coords[pair_indices[:, 1]],
-    )
-    distance_error = torch.abs(true_pair_distances - pred_pair_distances)
+    )  # (n_pair,)
+    distance_error = torch.abs(true_pair_distances - pred_pair_distances)  # (n_pair,)
     smooth_agreement = (
         sum(torch.sigmoid(threshold - distance_error) for threshold in (0.5, 1.0, 2.0, 4.0)) / 4.0
-    )
-    return smooth_agreement.sum() / (pair_indices.shape[0] + 1e-5)
+    )  # (n_pair,)
+    return smooth_agreement.sum() / (pair_indices.shape[0] + 1e-5)  # ()
 
 
 def smooth_lddt_loss(
@@ -174,5 +181,5 @@ def smooth_lddt_loss(
             other_cutoff=other_cutoff,
         )
         for index in range(true_coords.shape[0])
-    ]
-    return 1.0 - torch.stack(agreements).mean(dim=0)
+    ]  # each: ()
+    return 1.0 - torch.stack(agreements).mean(dim=0)  # ()

fastplms/models/boltz/vb_modules_diffusion_conditioning.py

@@ -30,7 +30,10 @@ def _concatenate_biases(
     projections: nn.ModuleList,
     pair_features: torch.Tensor,
 ) -> torch.Tensor:
-    return torch.cat([projection(pair_features) for projection in projections], dim=-1)
+    # pair_features: (..., d_pair); each projection: (..., h).
+    return torch.cat(
+        [projection(pair_features) for projection in projections], dim=-1
+    )  # (..., depth * h)
 
 
 class DiffusionConditioning(nn.Module):
@@ -112,21 +115,26 @@ class DiffusionConditioning(nn.Module):
         order so downstream code can select them by layer.
         """
 
+        # b is batch size, t token count, a atom count, and k the atom-window count.
         z_conditioned = self.pairwise_conditioner(
             z_trunk,
             relative_position_encoding,
-        )
+        )  # (b, t, t, d_z)
         q, c, p, to_keys = self.atom_encoder(
             feats=feats,
             s_trunk=s_trunk,
             z=z_conditioned,
-        )
-        atom_encoder_bias = _concatenate_biases(self.atom_enc_proj_z, p)
-        atom_decoder_bias = _concatenate_biases(self.atom_dec_proj_z, p)
+        )  # q/c: (b, a, d_a); p: (b, k, w, h_k, d_p); to_keys: callable
+        atom_encoder_bias = _concatenate_biases(
+            self.atom_enc_proj_z, p
+        )  # (b, k, w, h_k, depth_enc * heads_enc)
+        atom_decoder_bias = _concatenate_biases(
+            self.atom_dec_proj_z, p
+        )  # (b, k, w, h_k, depth_dec * heads_dec)
         token_transformer_bias = _concatenate_biases(
             self.token_trans_proj_z,
             z_conditioned,
-        )
+        )  # (b, t, t, depth_token * heads_token)
         return (
             q,
             c,
@@ -134,4 +142,4 @@ class DiffusionConditioning(nn.Module):
             atom_encoder_bias,
             atom_decoder_bias,
             token_transformer_bias,
-        )
+        )  # tensor shapes are traced above; to_keys is the atom-key gatherer

fastplms/models/boltz/vb_modules_encodersv2.py

@@ -11,11 +11,10 @@ implementations under MIT terms.  See ``THIRD_PARTY_NOTICES.md``.
 
 from __future__ import annotations
 
+import torch
 from collections.abc import Callable
 from functools import partial
 from math import pi
-
-import torch
 from torch import nn
 from torch.nn.functional import one_hot
 
@@ -46,12 +45,13 @@ class FourierEmbedding(nn.Module):
     def forward(self, times: torch.Tensor) -> torch.Tensor:
         """Map times with shape ``(b,)`` to an embedding ``H: (b, d)``."""
 
-        random_phase = self.proj(times.reshape(-1, 1))
-        return torch.cos(2 * pi * random_phase)
+        random_phase = self.proj(times.reshape(-1, 1))  # (b, d)
+        return torch.cos(2 * pi * random_phase)  # (b, d)
 
 
 def _pairwise_difference(values: torch.Tensor) -> torch.Tensor:
-    return values[:, :, None] - values[:, None, :]
+    # values: (b, n)
+    return values[:, :, None] - values[:, None, :]  # (b, n, n)
 
 
 class RelativePositionEncoder(nn.Module):
@@ -76,54 +76,68 @@ class RelativePositionEncoder(nn.Module):
     def forward(self, feats: dict[str, torch.Tensor]) -> torch.Tensor:
         """Return relative pair features ``Z: (b, n, n, token_z)``."""
 
-        same_chain = feats["asym_id"][:, :, None] == feats["asym_id"][:, None, :]
-        same_residue = feats["residue_index"][:, :, None] == feats["residue_index"][:, None, :]
-        same_entity = feats["entity_id"][:, :, None] == feats["entity_id"][:, None, :]
-
-        residue_offset = _pairwise_difference(feats["residue_index"])
+        same_chain = (
+            feats["asym_id"][:, :, None] == feats["asym_id"][:, None, :]
+        )  # (b, n, n)
+        same_residue = (
+            feats["residue_index"][:, :, None] == feats["residue_index"][:, None, :]
+        )  # (b, n, n)
+        same_entity = (
+            feats["entity_id"][:, :, None] == feats["entity_id"][:, None, :]
+        )  # (b, n, n)
+
+        residue_offset = _pairwise_difference(feats["residue_index"])  # (b, n, n)
         if self.cyclic_pos_enc and torch.any(feats["cyclic_period"] > 0):
             period = torch.where(
                 feats["cyclic_period"] > 0,
                 feats["cyclic_period"],
                 torch.zeros_like(feats["cyclic_period"]) + 10000,
-            )
-            residue_offset = (residue_offset - period * torch.round(residue_offset / period)).long()
+            )  # (b, n)
+            residue_offset = (
+                residue_offset - period * torch.round(residue_offset / period)
+            ).long()  # (b, n, n)
         residue_offset = torch.clip(
             residue_offset + self.r_max,
             0,
             2 * self.r_max,
-        )
+        )  # (b, n, n)
         residue_offset = torch.where(
             same_chain,
             residue_offset,
             torch.zeros_like(residue_offset) + 2 * self.r_max + 1,
-        )
-        residue_features = one_hot(residue_offset, 2 * self.r_max + 2)
+        )  # (b, n, n)
+        residue_features = one_hot(
+            residue_offset, 2 * self.r_max + 2
+        )  # (b, n, n, 2 * r_max + 2)
 
         token_offset = torch.clip(
             _pairwise_difference(feats["token_index"]) + self.r_max,
             0,
             2 * self.r_max,
-        )
+        )  # (b, n, n)
         token_offset = torch.where(
             same_chain & same_residue,
             token_offset,
             torch.zeros_like(token_offset) + 2 * self.r_max + 1,
-        )
-        token_features = one_hot(token_offset, 2 * self.r_max + 2)
+        )  # (b, n, n)
+        token_features = one_hot(
+            token_offset, 2 * self.r_max + 2
+        )  # (b, n, n, 2 * r_max + 2)
 
         symmetry_offset = torch.clip(
             _pairwise_difference(feats["sym_id"]) + self.s_max,
             0,
             2 * self.s_max,
-        )
-        invalid_symmetry = ~same_entity if self.fix_sym_check else same_chain
+        )  # (b, n, n)
+        invalid_symmetry = ~same_entity if self.fix_sym_check else same_chain  # (b, n, n)
         symmetry_offset = torch.where(
             invalid_symmetry,
             torch.zeros_like(symmetry_offset) + 2 * self.s_max + 1,
             symmetry_offset,
-        )
-        symmetry_features = one_hot(symmetry_offset, 2 * self.s_max + 2)
+        )  # (b, n, n)
+        symmetry_features = one_hot(
+            symmetry_offset, 2 * self.s_max + 2
+        )  # (b, n, n, 2 * s_max + 2)
 
         pair_features = torch.cat(
             (
@@ -133,8 +147,8 @@ class RelativePositionEncoder(nn.Module):
                 symmetry_features.float(),
             ),
             dim=-1,
-        )
-        return self.linear_layer(pair_features)
+        )  # (b, n, n, d_rel)
+        return self.linear_layer(pair_features)  # (b, n, n, token_z)
 
 
 class SingleConditioning(nn.Module):
@@ -175,17 +189,20 @@ class SingleConditioning(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """Return conditioned token features ``S: (b, n, 2 * token_s)``."""
 
-        single = torch.cat((s_trunk, s_inputs), dim=-1)
-        single = self.single_embed(self.norm_single(single))
-        normalized_fourier = None
+        # times: (b,); s_trunk/s_inputs: (b, n, token_s).
+        single = torch.cat((s_trunk, s_inputs), dim=-1)  # (b, n, 2 * token_s)
+        single = self.single_embed(self.norm_single(single))  # (b, n, 2 * token_s)
+        normalized_fourier = None  # (b, d_fourier) or None
         if not self.disable_times:
-            fourier = self.fourier_embed(times)
-            normalized_fourier = self.norm_fourier(fourier)
-            time_condition = self.fourier_to_single(normalized_fourier)
-            single = time_condition[:, None, :] + single
+            fourier = self.fourier_embed(times)  # (b, d_fourier)
+            normalized_fourier = self.norm_fourier(fourier)  # (b, d_fourier)
+            time_condition = self.fourier_to_single(
+                normalized_fourier
+            )  # (b, 2 * token_s)
+            single = time_condition[:, None, :] + single  # (b, n, 2 * token_s)
         for transition in self.transitions:
-            single = transition(single) + single
-        return single, normalized_fourier
+            single = transition(single) + single  # (b, n, 2 * token_s)
+        return single, normalized_fourier  # (b, n, 2 * token_s), (b, d_fourier) or None
 
 
 class PairwiseConditioning(nn.Module):
@@ -217,10 +234,12 @@ class PairwiseConditioning(nn.Module):
     ) -> torch.Tensor:
         """Return conditioned pair features ``Z: (b, n, n, token_z)``."""
 
-        pair = self.dim_pairwise_init_proj(torch.cat((z_trunk, token_rel_pos_feats), dim=-1))
+        pair = self.dim_pairwise_init_proj(
+            torch.cat((z_trunk, token_rel_pos_feats), dim=-1)
+        )  # (b, n, n, token_z)
         for transition in self.transitions:
-            pair = transition(pair) + pair
-        return pair
+            pair = transition(pair) + pair  # (b, n, n, token_z)
+        return pair  # (b, n, n, token_z)
 
 
 def get_indexing_matrix(
@@ -247,13 +266,17 @@ def get_indexing_matrix(
             f"h must contain an even number of half-window key blocks; received {key_blocks}."
         )
 
-    positions = torch.arange(2 * k, device=device)
-    relative_blocks = ((positions.unsqueeze(0) - positions.unsqueeze(1)) + key_blocks // 2).clamp(
-        min=0, max=key_blocks + 1
-    )
-    relative_blocks = relative_blocks.view(k, 2, 2 * k)[:, 0, :]
-    selectors = one_hot(relative_blocks, num_classes=key_blocks + 2)[..., 1:-1].transpose(1, 0)
-    return selectors.reshape(2 * k, key_blocks * k).float()
+    positions = torch.arange(2 * k, device=device)  # (2 * k,)
+    relative_blocks = (
+        (positions.unsqueeze(0) - positions.unsqueeze(1)) + key_blocks // 2
+    ).clamp(min=0, max=key_blocks + 1)  # (2 * k, 2 * k)
+    relative_blocks = relative_blocks.view(k, 2, 2 * k)[:, 0, :]  # (k, 2 * k)
+    selectors = one_hot(relative_blocks, num_classes=key_blocks + 2)[
+        ..., 1:-1
+    ].transpose(1, 0)  # (2 * k, k, key_blocks)
+    return selectors.reshape(
+        2 * k, key_blocks * k
+    ).float()  # (2 * k, key_blocks * k)
 
 
 def single_to_keys(
@@ -266,9 +289,11 @@ def single_to_keys(
 
     b, n, d = single.shape
     k = n // w
-    half_windows = single.view(b, 2 * k, w // 2, d)
-    gathered = torch.einsum("b j i d, j k -> b k i d", half_windows, indexing_matrix)
-    return gathered.reshape(b, k, h, d)
+    half_windows = single.view(b, 2 * k, w // 2, d)  # (b, 2 * k, w / 2, d)
+    gathered = torch.einsum(
+        "b j i d, j k -> b k i d", half_windows, indexing_matrix
+    )  # (b, key_blocks * k, w / 2, d)
+    return gathered.reshape(b, k, h, d)  # (b, k, h, d)
 
 
 class AtomEncoder(nn.Module):
@@ -332,11 +357,12 @@ class AtomEncoder(nn.Module):
         b: int,
         n: int,
     ) -> torch.Tensor:
+        # Feature tensors share leading shape (b, n).
         feature_parts = [
             feats["ref_pos"],
             feats["ref_charge"].unsqueeze(-1),
             feats["ref_element"],
-        ]
+        ]  # each: (b, n, d_feature)
         if not self.use_no_atom_char:
             feature_parts.append(feats["ref_atom_name_chars"].reshape(b, n, 4 * 64))
         if self.use_atom_backbone_feat:
@@ -349,9 +375,11 @@ class AtomEncoder(nn.Module):
                     one_hot(feats["mol_type"], num_classes=4).float(),
                 ),
                 dim=-1,
-            )
-            feature_parts.append(torch.bmm(feats["atom_to_token"].float(), residue_features))
-        return torch.cat(feature_parts, dim=-1)
+            )  # (b, t, d_residue)
+            feature_parts.append(
+                torch.bmm(feats["atom_to_token"].float(), residue_features)
+            )  # (b, n, d_residue)
+        return torch.cat(feature_parts, dim=-1)  # (b, n, atom_feature_dim)
 
     def forward(
         self,
@@ -368,61 +396,81 @@ class AtomEncoder(nn.Module):
 
         with torch.autocast("cuda", enabled=False):
             b, n, _ = feats["ref_pos"].shape
-            atom_mask = feats["atom_pad_mask"].bool()
-            atom_positions = feats["ref_pos"]
-            atom_space = feats["ref_space_uid"]
-            atom_features = self._assemble_atom_features(feats, b, n)
-            conditioned_single = self.embed_atom_features(atom_features)
+            atom_mask = feats["atom_pad_mask"].bool()  # (b, n)
+            atom_positions = feats["ref_pos"]  # (b, n, 3)
+            atom_space = feats["ref_space_uid"]  # (b, n)
+            atom_features = self._assemble_atom_features(
+                feats, b, n
+            )  # (b, n, atom_feature_dim)
+            conditioned_single = self.embed_atom_features(
+                atom_features
+            )  # (b, n, atom_s)
 
             w = self.atoms_per_window_queries
             h = self.atoms_per_window_keys
             b, n = conditioned_single.shape[:2]
             k = n // w
-            indexing = get_indexing_matrix(k, w, h, conditioned_single.device)
+            indexing = get_indexing_matrix(
+                k, w, h, conditioned_single.device
+            )  # (2 * k, key_blocks * k)
             to_keys = partial(single_to_keys, indexing_matrix=indexing, w=w, h=h)
 
-            position_queries = atom_positions.view(b, k, w, 1, 3)
-            position_keys = to_keys(atom_positions).view(b, k, 1, h, 3)
-            displacement = position_keys - position_queries
+            position_queries = atom_positions.view(b, k, w, 1, 3)  # (b, k, w, 1, 3)
+            position_keys = to_keys(atom_positions).view(
+                b, k, 1, h, 3
+            )  # (b, k, 1, h, 3)
+            displacement = position_keys - position_queries  # (b, k, w, h, 3)
             inverse_squared_distance = 1 / (
                 1 + torch.sum(displacement * displacement, dim=-1, keepdim=True)
-            )
-
-            mask_queries = atom_mask.view(b, k, w, 1)
-            mask_keys = to_keys(atom_mask.unsqueeze(-1).float()).view(b, k, 1, h).bool()
-            space_queries = atom_space.view(b, k, w, 1)
-            space_keys = to_keys(atom_space.unsqueeze(-1).float()).view(b, k, 1, h).long()
+            )  # (b, k, w, h, 1)
+
+            mask_queries = atom_mask.view(b, k, w, 1)  # (b, k, w, 1)
+            mask_keys = (
+                to_keys(atom_mask.unsqueeze(-1).float()).view(b, k, 1, h).bool()
+            )  # (b, k, 1, h)
+            space_queries = atom_space.view(b, k, w, 1)  # (b, k, w, 1)
+            space_keys = (
+                to_keys(atom_space.unsqueeze(-1).float()).view(b, k, 1, h).long()
+            )  # (b, k, 1, h)
             valid_pair = (
                 (mask_queries & mask_keys & (space_queries == space_keys)).float().unsqueeze(-1)
-            )
+            )  # (b, k, w, h, 1)
 
-            pair = self.embed_atompair_ref_pos(displacement) * valid_pair
-            pair = pair + self.embed_atompair_ref_dist(inverse_squared_distance) * valid_pair
-            pair = pair + self.embed_atompair_mask(valid_pair) * valid_pair
-            query = conditioned_single
+            pair = self.embed_atompair_ref_pos(displacement) * valid_pair  # (b, k, w, h, atom_z)
+            pair = (
+                pair + self.embed_atompair_ref_dist(inverse_squared_distance) * valid_pair
+            )  # (b, k, w, h, atom_z)
+            pair = pair + self.embed_atompair_mask(valid_pair) * valid_pair  # (b, k, w, h, atom_z)
+            query = conditioned_single  # (b, n, atom_s)
 
             if self.structure_prediction:
                 if s_trunk is None or z is None:
                     raise ValueError("structure prediction requires S and Z trunk tensors")
-                atom_to_token = feats["atom_to_token"].float()
-                single_update = self.s_to_c_trans(s_trunk.float())
-                single_update = torch.bmm(atom_to_token, single_update)
-                conditioned_single = conditioned_single + single_update.to(conditioned_single)
-
-                token_queries = atom_to_token.view(b, k, w, atom_to_token.shape[-1])
-                token_keys = to_keys(atom_to_token)
-                pair_update = self.z_to_p_trans(z.float())
+                atom_to_token = feats["atom_to_token"].float()  # (b, n, t)
+                single_update = self.s_to_c_trans(s_trunk.float())  # (b, t, atom_s)
+                single_update = torch.bmm(
+                    atom_to_token, single_update
+                )  # (b, n, atom_s)
+                conditioned_single = conditioned_single + single_update.to(
+                    conditioned_single
+                )  # (b, n, atom_s)
+
+                token_queries = atom_to_token.view(
+                    b, k, w, atom_to_token.shape[-1]
+                )  # (b, k, w, t)
+                token_keys = to_keys(atom_to_token)  # (b, k, h, t)
+                pair_update = self.z_to_p_trans(z.float())  # (b, t, t, atom_z)
                 pair_update = torch.einsum(
                     "bijd,bwki,bwlj->bwkld",
                     pair_update,
                     token_queries,
                     token_keys,
-                )
-                pair = pair + pair_update.to(pair)
+                )  # (b, k, w, h, atom_z)
+                pair = pair + pair_update.to(pair)  # (b, k, w, h, atom_z)
 
             pair = pair + self.c_to_p_trans_q(
                 conditioned_single.view(b, k, w, 1, conditioned_single.shape[-1])
-            )
+            )  # (b, k, w, h, atom_z)
             pair = pair + self.c_to_p_trans_k(
                 to_keys(conditioned_single).view(
                     b,
@@ -431,9 +479,14 @@ class AtomEncoder(nn.Module):
                     h,
                     conditioned_single.shape[-1],
                 )
-            )
-            pair = pair + self.p_mlp(pair)
-        return query, conditioned_single, pair, to_keys
+            )  # (b, k, w, h, atom_z)
+            pair = pair + self.p_mlp(pair)  # (b, k, w, h, atom_z)
+        return (
+            query,
+            conditioned_single,
+            pair,
+            to_keys,
+        )  # (b, n, atom_s), (b, n, atom_s), (b, k, w, h, atom_z), callable
 
 
 class AtomAttentionEncoder(nn.Module):
@@ -484,14 +537,15 @@ class AtomAttentionEncoder(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Callable]:
         b, n, _ = feats["ref_pos"].shape
         del b, n
-        atom_mask = feats["atom_pad_mask"].bool()
+        # q/c: (b, a, atom_s); atom_enc_bias uses local-window attention axes.
+        atom_mask = feats["atom_pad_mask"].bool()  # (b, a)
         if self.structure_prediction:
             if r is None:
                 raise ValueError("structure prediction requires atom coordinates R")
-            q = q.repeat_interleave(multiplicity, 0)
-            q = q + self.r_to_q_trans(r)
-        c = c.repeat_interleave(multiplicity, 0)
-        atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
+            q = q.repeat_interleave(multiplicity, 0)  # (b * m, a, atom_s)
+            q = q + self.r_to_q_trans(r)  # (b * m, a, atom_s)
+        c = c.repeat_interleave(multiplicity, 0)  # (b * m, a, atom_s)
+        atom_mask = atom_mask.repeat_interleave(multiplicity, 0)  # (b * m, a)
         q = self.atom_encoder(
             q=q,
             mask=atom_mask,
@@ -499,15 +553,26 @@ class AtomAttentionEncoder(nn.Module):
             bias=atom_enc_bias,
             multiplicity=multiplicity,
             to_keys=to_keys,
-        )
+        )  # (b * m, a, atom_s)
 
         with torch.autocast("cuda", enabled=False):
-            atom_update = self.atom_to_token_trans(q).float()
-            atom_to_token = feats["atom_to_token"].float()
-            atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)
-            atom_to_token_mean = atom_to_token / (atom_to_token.sum(dim=1, keepdim=True) + 1e-6)
-            token_update = torch.bmm(atom_to_token_mean.transpose(1, 2), atom_update)
-        return token_update.to(q), q, c, to_keys
+            atom_update = self.atom_to_token_trans(q).float()  # (b * m, a, d_token)
+            atom_to_token = feats["atom_to_token"].float()  # (b, a, t)
+            atom_to_token = atom_to_token.repeat_interleave(
+                multiplicity, 0
+            )  # (b * m, a, t)
+            atom_to_token_mean = atom_to_token / (
+                atom_to_token.sum(dim=1, keepdim=True) + 1e-6
+            )  # (b * m, a, t)
+            token_update = torch.bmm(
+                atom_to_token_mean.transpose(1, 2), atom_update
+            )  # (b * m, t, d_token)
+        return (
+            token_update.to(q),
+            q,
+            c,
+            to_keys,
+        )  # (b * m, t, d_token), two (b * m, a, atom_s), callable
 
 
 class AtomAttentionDecoder(nn.Module):
@@ -560,12 +625,17 @@ class AtomAttentionDecoder(nn.Module):
         """Return atom-coordinate updates ``R_update: (b, a, 3)``."""
 
         with torch.autocast("cuda", enabled=False):
-            atom_to_token = feats["atom_to_token"].float()
-            atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)
-            token_update = self.a_to_q_trans(a.float())
-            atom_update = torch.bmm(atom_to_token, token_update)
-        q = q + atom_update.to(q)
-        atom_mask = feats["atom_pad_mask"].repeat_interleave(multiplicity, 0)
+            # a: (b * m, t, 2 * token_s); q/c: (b * m, a, atom_s).
+            atom_to_token = feats["atom_to_token"].float()  # (b, a, t)
+            atom_to_token = atom_to_token.repeat_interleave(
+                multiplicity, 0
+            )  # (b * m, a, t)
+            token_update = self.a_to_q_trans(a.float())  # (b * m, t, atom_s)
+            atom_update = torch.bmm(atom_to_token, token_update)  # (b * m, a, atom_s)
+        q = q + atom_update.to(q)  # (b * m, a, atom_s)
+        atom_mask = feats["atom_pad_mask"].repeat_interleave(
+            multiplicity, 0
+        )  # (b * m, a)
         q = self.atom_decoder(
             q=q,
             mask=atom_mask,
@@ -573,5 +643,5 @@ class AtomAttentionDecoder(nn.Module):
             bias=atom_dec_bias,
             multiplicity=multiplicity,
             to_keys=to_keys,
-        )
-        return self.atom_feat_to_atom_pos_update(q)
+        )  # (b * m, a, atom_s)
+        return self.atom_feat_to_atom_pos_update(q)  # (b * m, a, 3)

fastplms/models/boltz/vb_modules_transformersv2.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from typing import Any, cast
-
 from torch import Tensor, nn
 from torch.utils.checkpoint import checkpoint
 
@@ -23,10 +22,12 @@ class AdaLN(nn.Module):
         self.s_bias = LinearNoBias(dim_single_cond, dim)
 
     def forward(self, a: Tensor, s: Tensor) -> Tensor:
-        normalized = self.a_norm(a)
-        conditioning = self.s_norm(s)
+        # a: (..., d); s: (..., d_c).
+        normalized = self.a_norm(a)  # (..., d)
+        conditioning = self.s_norm(s)  # (..., d_c)
+        # (..., d)
         output = self.s_scale(conditioning).sigmoid() * normalized + self.s_bias(conditioning)
-        return cast(Tensor, output)
+        return cast(Tensor, output)  # (..., d)
 
 
 class ConditionedTransitionBlock(nn.Module):
@@ -49,13 +50,15 @@ class ConditionedTransitionBlock(nn.Module):
         self.b_to_a = LinearNoBias(inner_dim, dim_single)
 
         projection = nn.Linear(dim_single_cond, dim_single)
-        nn.init.zeros_(projection.weight)
-        nn.init.constant_(projection.bias, -2.0)
+        nn.init.zeros_(projection.weight)  # (d, d_c)
+        nn.init.constant_(projection.bias, -2.0)  # (d,)
         self.output_projection = nn.Sequential(projection, nn.Sigmoid())
 
     def forward(self, a: Tensor, s: Tensor) -> Tensor:
-        normalized = self.adaln(a, s)
-        hidden = self.swish_gate(normalized) * self.a_to_b(normalized)
+        # a: (..., d); s: (..., d_c).
+        normalized = self.adaln(a, s)  # (..., d)
+        hidden = self.swish_gate(normalized) * self.a_to_b(normalized)  # (..., d_inner)
+        # (..., d)
         return cast(Tensor, self.output_projection(s) * self.b_to_a(hidden))
 
 
@@ -89,14 +92,16 @@ class DiffusionTransformer(nn.Module):
         )
 
     def _split_pair_bias(self, bias: Tensor | None) -> Tensor | None:
+        # bias: (b_z, l_q, l_k, n_layer * h) or None.
         if not self.pair_bias_attn:
             return None
         if bias is None:
             raise ValueError("pair bias is required when pair_bias_attn=True")
-        batch_size, query_length, key_length, width = bias.shape
-        depth = len(self.layers)
+        batch_size, query_length, key_length, width = bias.shape  # b_z, l_q, l_k, n_layer * h
+        depth = len(self.layers)  # n_layer
         if depth == 0 or width % depth:
             raise ValueError("pair-bias width must be divisible by transformer depth")
+        # (b_z, l_q, l_k, n_layer, h)
         return bias.view(batch_size, query_length, key_length, depth, width // depth)
 
     def forward(
@@ -110,12 +115,16 @@ class DiffusionTransformer(nn.Module):
     ) -> Tensor:
         """Transform A and preserve shape ``(b * m, l, d)``."""
 
-        layer_biases = self._split_pair_bias(bias)
-        output = a
+        # a: (b * m, l_q, d); s: (b * m, l_q, d_c).
+        # bias: (b, l_q, l_k, n_layer * h) or None.
+        # mask: (b * m, l_q) before an optional to_keys mapping; to_keys preserves d.
+        layer_biases = self._split_pair_bias(bias)  # (b, l_q, l_k, n_layer, h) or None
+        output = a  # (b * m, l_q, d)
         for index, layer in enumerate(self.layers):
+            # (b, l_q, l_k, h) or None
             layer_bias = None if layer_biases is None else layer_biases[..., index, :]
             if self.activation_checkpointing:
-                output = checkpoint(
+                output = checkpoint(  # (b * m, l_q, d)
                     layer,
                     output,
                     s,
@@ -126,7 +135,7 @@ class DiffusionTransformer(nn.Module):
                     use_reentrant=False,
                 )
             else:
-                output = layer(
+                output = layer(  # (b * m, l_q, d)
                     output,
                     s,
                     layer_bias,
@@ -134,7 +143,7 @@ class DiffusionTransformer(nn.Module):
                     to_keys,
                     multiplicity,
                 )
-        return output
+        return output  # (b * m, l_q, d)
 
 
 class DiffusionTransformerLayer(nn.Module):
@@ -156,8 +165,8 @@ class DiffusionTransformerLayer(nn.Module):
             compute_pair_bias=False,
         )
         self.output_projection_linear = nn.Linear(conditioning_dim, dim)
-        nn.init.zeros_(self.output_projection_linear.weight)
-        nn.init.constant_(self.output_projection_linear.bias, -2.0)
+        nn.init.zeros_(self.output_projection_linear.weight)  # (d, d_c)
+        nn.init.constant_(self.output_projection_linear.bias, -2.0)  # (d,)
         self.output_projection = nn.Sequential(
             self.output_projection_linear,
             nn.Sigmoid(),
@@ -176,25 +185,28 @@ class DiffusionTransformerLayer(nn.Module):
     ) -> Tensor:
         """Update activation tensor A with conditioning S and pair bias."""
 
+        # a: (b * m, l_q, d); s: (b * m, l_q, d_c).
+        # bias: (b, l_q, l_k, h) or None.
+        # mask: (b * m, l_q) or None; to_keys maps l_q to l_k.
         if bias is None or mask is None:
             raise ValueError("diffusion attention requires pair bias and a key mask")
-        normalized = self.adaln(a, s)
-        key_states = normalized
-        key_mask = mask
+        normalized = self.adaln(a, s)  # (b * m, l_q, d)
+        key_states = normalized  # (b * m, l_q, d)
+        key_mask = mask  # (b * m, l_q)
         if to_keys is not None:
-            key_states = to_keys(normalized)
-            key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)
+            key_states = to_keys(normalized)  # (b * m, l_k, d)
+            key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)  # (b * m, l_k)
 
-        attended = self.pair_bias_attn(
+        attended = self.pair_bias_attn(  # (b * m, l_q, d)
             s=normalized,
             z=bias,
             mask=key_mask,
             multiplicity=multiplicity,
             k_in=key_states,
         )
-        output = a + self.output_projection(s) * attended
-        output = output + self.transition(output, s)
-        return cast(Tensor, self.post_lnorm(output))
+        output = a + self.output_projection(s) * attended  # (b * m, l_q, d)
+        output = output + self.transition(output, s)  # (b * m, l_q, d)
+        return cast(Tensor, self.post_lnorm(output))  # (b * m, l_q, d)
 
 
 class AtomTransformer(nn.Module):
@@ -224,15 +236,21 @@ class AtomTransformer(nn.Module):
     ) -> Tensor:
         """Transform atom tensor Q with shape ``(b, n_atoms, d)``."""
 
-        query_window = self.attn_window_queries
-        key_window = self.attn_window_keys
-        batch_size, atom_count, width = q.shape
-        window_count = atom_count // query_window
+        query_window = self.attn_window_queries  # w_q
+        key_window = self.attn_window_keys  # w_k
+        batch_size, atom_count, width = q.shape  # b * m, n_atom, d
+        window_count = atom_count // query_window  # k
+        # q, c: (b * m, n_atom, d); bias: (b, k, w_q, w_k, n_layer * h).
+        # mask: (b * m, n_atom); k = n_atom // w_q.
+        # (b * m * k, w_q, d)
         query_states = q.view(batch_size * window_count, query_window, -1)
+        # (b * m * k, w_q, d_c)
         conditioning = c.view(batch_size * window_count, query_window, -1)
+        # (b * m * k, w_q)
         query_mask = mask.view(batch_size * window_count, query_window)
+        # (b * m, k, w_q, w_k, n_layer * h)
         pair_bias = bias.repeat_interleave(multiplicity, dim=0)
-        pair_bias = pair_bias.view(
+        pair_bias = pair_bias.view(  # (b * m * k, w_q, w_k, n_layer * h)
             pair_bias.shape[0] * window_count,
             query_window,
             key_window,
@@ -240,10 +258,13 @@ class AtomTransformer(nn.Module):
         )
 
         def windowed_keys(states: Tensor) -> Tensor:
+            # states: (b * m * k, w_q, d_x).
+            # (b * m, n_atom, d_x)
             merged = states.view(batch_size, window_count * query_window, -1)
+            # (b * m * k, w_k, d_x)
             return to_keys(merged).view(batch_size * window_count, key_window, -1)
 
-        output = self.diffusion_transformer(
+        output = self.diffusion_transformer(  # (b * m * k, w_q, d)
             a=query_states,
             s=conditioning,
             bias=pair_bias,
@@ -251,4 +272,5 @@ class AtomTransformer(nn.Module):
             multiplicity=1,
             to_keys=windowed_keys,
         )
+        # (b * m, n_atom, d)
         return cast(Tensor, output.view(batch_size, window_count * query_window, width))

fastplms/models/boltz/vb_modules_trunkv2.py

@@ -2,9 +2,8 @@
 
 from __future__ import annotations
 
-from typing import Any, cast
-
 import torch
+from typing import Any, cast
 from torch import Tensor, nn
 from torch.nn import functional as F
 from torch.utils.checkpoint import checkpoint
@@ -30,8 +29,8 @@ class ContactConditioning(nn.Module):
         self.fourier_embedding = FourierEmbedding(token_z)
         input_width = token_z + len(const.contact_conditioning_info) - 1
         self.encoder = nn.Linear(input_width, token_z)
-        self.encoding_unspecified = nn.Parameter(torch.zeros(token_z))
-        self.encoding_unselected = nn.Parameter(torch.zeros(token_z))
+        self.encoding_unspecified = nn.Parameter(torch.zeros(token_z))  # (d_z,)
+        self.encoding_unselected = nn.Parameter(torch.zeros(token_z))  # (d_z,)
         self.cutoff_min = cutoff_min
         self.cutoff_max = cutoff_max
 
@@ -43,22 +42,25 @@ class ContactConditioning(nn.Module):
         if const.contact_conditioning_info["UNSELECTED"] != 1:
             raise ValueError("UNSELECTED contact conditioning must use channel one")
 
-        categories = feats["contact_conditioning"]
-        threshold = feats["contact_threshold"]
+        # c_contact is the number of contact-conditioning categories.
+        categories = feats["contact_conditioning"]  # (b, l, l, c_contact)
+        threshold = feats["contact_threshold"]  # (b, l, l)
+        # (b, l, l)
         normalized = (threshold - self.cutoff_min) / (self.cutoff_max - self.cutoff_min)
+        # (b, l, l, d_z)
         fourier = self.fourier_embedding(normalized.flatten()).reshape((*normalized.shape, -1))
-        selected_features = torch.cat(
+        selected_features = torch.cat(  # (b, l, l, c_contact - 1 + d_z)
             [categories[..., 2:], normalized.unsqueeze(-1), fourier],
             dim=-1,
         )
-        selected = self.encoder(selected_features)
-        special = categories[..., :2]
+        selected = self.encoder(selected_features)  # (b, l, l, d_z)
+        special = categories[..., :2]  # (b, l, l, 2)
         return cast(
             Tensor,
             selected * (1 - special.sum(dim=-1, keepdim=True))
             + self.encoding_unspecified * special[..., 0:1]
             + self.encoding_unselected * special[..., 1:2],
-        )
+        )  # (b, l, l, d_z)
 
 
 class InputEmbedder(nn.Module):
@@ -122,51 +124,56 @@ class InputEmbedder(nn.Module):
 
         if add_method_conditioning:
             self.method_conditioning_init = nn.Embedding(const.num_method_types, token_s)
-            self.method_conditioning_init.weight.data.fill_(0)
+            self.method_conditioning_init.weight.data.fill_(0)  # (n_method, d_s)
         if add_modified_flag:
             self.modified_conditioning_init = nn.Embedding(2, token_s)
-            self.modified_conditioning_init.weight.data.fill_(0)
+            self.modified_conditioning_init.weight.data.fill_(0)  # (2, d_s)
         if add_cyclic_flag:
             self.cyclic_conditioning_init = nn.Linear(1, token_s, bias=False)
-            self.cyclic_conditioning_init.weight.data.fill_(0)
+            self.cyclic_conditioning_init.weight.data.fill_(0)  # (d_s, 1)
         if add_mol_type_feat:
             self.mol_type_conditioning_init = nn.Embedding(
                 len(const.chain_type_ids),
                 token_s,
             )
-            self.mol_type_conditioning_init.weight.data.fill_(0)
+            self.mol_type_conditioning_init.weight.data.fill_(0)  # (n_mol_type, d_s)
 
     def forward(self, feats: dict[str, Tensor], affinity: bool = False) -> Tensor:
         """Return embedded sequence tensor S with shape ``(b, l, d_s)``."""
 
-        residue_type = feats["res_type"].float()
+        # n_atom is the padded atom count; k is the number of atom windows.
+        residue_type = feats["res_type"].float()  # (b, l, n_token_type)
         suffix = "_affinity" if affinity else ""
-        profile = feats[f"profile{suffix}"]
-        deletion_mean = feats[f"deletion_mean{suffix}"].unsqueeze(-1)
+        profile = feats[f"profile{suffix}"]  # (b, l, n_token_type)
+        deletion_mean = feats[f"deletion_mean{suffix}"].unsqueeze(-1)  # (b, l, 1)
+        # (b, n_atom, d_a), (b, n_atom, d_a), (b, k, w_q, w_k, d_az), callable
         atom_queries, atom_conditioning, atom_pairs, to_keys = self.atom_encoder(feats)
-        atom_bias = self.atom_enc_proj_z(atom_pairs)
+        atom_bias = self.atom_enc_proj_z(atom_pairs)  # (b, k, w_q, w_k, n_layer * h)
         atom_output, _, _, _ = self.atom_attention_encoder(
             feats=feats,
             q=atom_queries,
             c=atom_conditioning,
             atom_enc_bias=atom_bias,
             to_keys=to_keys,
-        )
-        output = (
+        )  # (b, l, d_s), (b, n_atom, d_a), (b, n_atom, d_a), callable
+        output = (  # (b, l, d_s)
             atom_output
             + self.res_type_encoding(residue_type)
             + self.msa_profile_encoding(torch.cat([profile, deletion_mean], dim=-1))
         )
         if self.add_method_conditioning:
+            # method_feature: (b, l); output: (b, l, d_s).
             output = output + self.method_conditioning_init(feats["method_feature"])
         if self.add_modified_flag:
+            # modified: (b, l); output: (b, l, d_s).
             output = output + self.modified_conditioning_init(feats["modified"])
         if self.add_cyclic_flag:
-            cyclic = feats["cyclic_period"].clamp(max=1.0).unsqueeze(-1)
-            output = output + self.cyclic_conditioning_init(cyclic)
+            cyclic = feats["cyclic_period"].clamp(max=1.0).unsqueeze(-1)  # (b, l, 1)
+            output = output + self.cyclic_conditioning_init(cyclic)  # (b, l, d_s)
         if self.add_mol_type_feat:
+            # mol_type: (b, l); output: (b, l, d_s).
             output = output + self.mol_type_conditioning_init(feats["mol_type"])
-        return cast(Tensor, output)
+        return cast(Tensor, output)  # (b, l, d_s)
 
 
 class _TemplateBase(nn.Module):
@@ -213,44 +220,54 @@ class _TemplateBase(nn.Module):
         feats: dict[str, Tensor],
         template_pair_mask: Tensor,
     ) -> Tensor:
-        residue_type = feats["template_restype"]
-        cb_mask = feats["template_mask_cb"]
-        frame_mask = feats["template_mask_frame"]
+        # residue_type: (b, t, l, n_token_type); template_pair_mask: (b, t, l, l).
+        residue_type = feats["template_restype"]  # (b, t, l, n_token_type)
+        cb_mask = feats["template_mask_cb"]  # (b, t, l)
+        frame_mask = feats["template_mask_frame"]  # (b, t, l)
+        # (b, t, l, l, 1)
         cb_pair_mask = (cb_mask[..., :, None] * cb_mask[..., None, :]).unsqueeze(-1)
+        # (b, t, l, l, 1)
         frame_pair_mask = (frame_mask[..., :, None] * frame_mask[..., None, :]).unsqueeze(-1)
         with torch.autocast(device_type="cuda", enabled=False):
+            # template_cb: (b, t, l, 3).
+            # (b, t, l, l)
             cb_distances = torch.cdist(feats["template_cb"], feats["template_cb"])
-            boundaries = torch.linspace(
+            boundaries = torch.linspace(  # (n_bin - 1,)
                 self.min_dist,
                 self.max_dist,
                 self.num_bins - 1,
             ).to(cb_distances.device)
-            bins = (cb_distances[..., None] > boundaries).sum(dim=-1).long()
-            distogram = F.one_hot(bins, num_classes=self.num_bins)
+            bins = (cb_distances[..., None] > boundaries).sum(dim=-1).long()  # (b, t, l, l)
+            distogram = F.one_hot(bins, num_classes=self.num_bins)  # (b, t, l, l, n_bin)
 
+            # (b, t, 1, l, 3, 3)
             rotations = feats["template_frame_rot"].unsqueeze(2).transpose(-1, -2)
+            # (b, t, 1, l, 3, 1)
             translations = feats["template_frame_t"].unsqueeze(2).unsqueeze(-1)
+            # (b, t, l, 1, 3, 1)
             ca_coordinates = feats["template_ca"].unsqueeze(3).unsqueeze(-1)
+            # (b, t, l, l, 3, 1)
             vectors = torch.matmul(rotations, ca_coordinates - translations)
-            norms = torch.norm(vectors, dim=-1, keepdim=True)
-            unit_vectors = torch.where(
+            norms = torch.norm(vectors, dim=-1, keepdim=True)  # (b, t, l, l, 3, 1)
+            unit_vectors = torch.where(  # (b, t, l, l, 3)
                 norms > 0,
                 vectors / norms,
                 torch.zeros_like(vectors),
             ).squeeze(-1)
-            pair_features = torch.cat(
+            pair_features = torch.cat(  # (b, t, l, l, n_bin + 5)
                 [distogram, cb_pair_mask, unit_vectors, frame_pair_mask],
                 dim=-1,
             )
+            # (b, t, l, l, n_bin + 5)
             pair_features = pair_features * template_pair_mask.unsqueeze(-1)
-            residue_i = residue_type[:, :, :, None].expand(
+            residue_i = residue_type[:, :, :, None].expand(  # (b, t, l, l, n_token_type)
                 -1,
                 -1,
                 -1,
                 residue_type.size(2),
                 -1,
             )
-            residue_j = residue_type[:, :, None, :].expand(
+            residue_j = residue_type[:, :, None, :].expand(  # (b, t, l, l, n_token_type)
                 -1,
                 -1,
                 residue_type.size(2),
@@ -260,7 +277,7 @@ class _TemplateBase(nn.Module):
             return cast(
                 Tensor,
                 self.a_proj(torch.cat([pair_features, residue_i, residue_j], dim=-1)),
-            )
+            )  # (b, t, l, l, d_t)
 
     def forward(
         self,
@@ -271,38 +288,40 @@ class _TemplateBase(nn.Module):
     ) -> Tensor:
         """Aggregate template pair tensor V into trunk update U."""
 
-        residue_type = feats["template_restype"]
+        # z: (b, l, l, d_z); pair_mask: (b, l, l).
+        residue_type = feats["template_restype"]  # (b, t, l, n_token_type)
         batch_size, template_count = residue_type.shape[:2]
-        template_present = feats["template_mask"].any(dim=2).float()
-        present_count = template_present.sum(dim=1).clamp(min=1)
-        features = self._template_features(
+        template_present = feats["template_mask"].any(dim=2).float()  # (b, t)
+        present_count = template_present.sum(dim=1).clamp(min=1)  # (b,)
+        features = self._template_features(  # (b, t, l, l, d_t)
             feats,
             self._template_pair_mask(feats, template_count),
         )
-        expanded_mask = pair_mask[:, None].expand(-1, template_count, -1, -1)
-        expanded_mask = expanded_mask.reshape(
+        expanded_mask = pair_mask[:, None].expand(-1, template_count, -1, -1)  # (b, t, l, l)
+        expanded_mask = expanded_mask.reshape(  # (b * t, l, l)
             batch_size * template_count,
             *expanded_mask.shape[2:],
         )
-        template_states = self.z_proj(self.z_norm(z[:, None])) + features
-        template_states = template_states.view(
+        template_states = self.z_proj(self.z_norm(z[:, None])) + features  # (b, t, l, l, d_t)
+        template_states = template_states.view(  # (b * t, l, l, d_t)
             batch_size * template_count,
             *template_states.shape[2:],
         )
-        template_states = template_states + self.pairformer(
+        template_states = template_states + self.pairformer(  # (b * t, l, l, d_t)
             template_states,
             expanded_mask,
             use_kernels=use_kernels,
         )
-        template_states = self.v_norm(template_states).view(
+        template_states = self.v_norm(template_states).view(  # (b, t, l, l, d_t)
             batch_size,
             template_count,
             *template_states.shape[1:],
         )
-        weights = template_present[:, :, None, None, None]
-        aggregate = (template_states * weights).sum(dim=1)
+        weights = template_present[:, :, None, None, None]  # (b, t, 1, 1, 1)
+        aggregate = (template_states * weights).sum(dim=1)  # (b, l, l, d_t)
+        # (b, l, l, d_t)
         aggregate = aggregate / present_count[:, None, None, None].to(template_states)
-        return cast(Tensor, self.u_proj(self.relu(aggregate)))
+        return cast(Tensor, self.u_proj(self.relu(aggregate)))  # (b, l, l, d_z)
 
 
 class TemplateModule(_TemplateBase):
@@ -339,9 +358,9 @@ class TemplateModule(_TemplateBase):
         )
 
     def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
-        asym_id = feats["asym_id"]
-        same_chain = (asym_id[:, :, None] == asym_id[:, None, :]).float()
-        return same_chain[:, None].expand(-1, count, -1, -1)
+        asym_id = feats["asym_id"]  # (b, l)
+        same_chain = (asym_id[:, :, None] == asym_id[:, None, :]).float()  # (b, l, l)
+        return same_chain[:, None].expand(-1, count, -1, -1)  # (b, t, l, l)
 
 
 class TemplateV2Module(_TemplateBase):
@@ -379,8 +398,8 @@ class TemplateV2Module(_TemplateBase):
 
     def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
         del count
-        visibility = feats["visibility_ids"]
-        return (visibility[..., :, None] == visibility[..., None, :]).float()
+        visibility = feats["visibility_ids"]  # (b, t, l)
+        return (visibility[..., :, None] == visibility[..., None, :]).float()  # (b, t, l, l)
 
 
 class MSAModule(nn.Module):
@@ -433,6 +452,7 @@ class MSAModule(nn.Module):
         pair_states: Tensor,
         training: bool,
     ) -> tuple[bool, int | None, int | None, int | None, int | None]:
+        # pair_states: (b, l, l, d_z).
         if training:
             return False, None, None, None, None
         if pair_states.shape[1] > const.chunk_size_threshold:
@@ -449,28 +469,32 @@ class MSAModule(nn.Module):
         """Return updated pair tensor Z after every MSA block."""
 
         chunking = self._chunk_configuration(z, self.training)
-        msa = feats["msa"]
+        # z: (b, l, l, d_z); emb: (b, l, d_s).
+        # s is MSA depth; n_token_type is the residue vocabulary size.
+        msa = feats["msa"]  # (b, s, l) or (b, s, l, n_token_type)
         if msa.dtype in (torch.long, torch.int32, torch.int64):
-            msa = F.one_hot(msa, num_classes=const.num_tokens).float()
-        msa_mask = feats["msa_mask"]
-        components = [
+            msa = F.one_hot(msa, num_classes=const.num_tokens).float()  # (b, s, l, n_token_type)
+        msa_mask = feats["msa_mask"]  # (b, s, l)
+        components = [  # (b, s, l, n_token_type), then two (b, s, l, 1) tensors
             msa,
-            feats["has_deletion"].unsqueeze(-1),
-            feats["deletion_value"].unsqueeze(-1),
+            feats["has_deletion"].unsqueeze(-1),  # (b, s, l, 1)
+            feats["deletion_value"].unsqueeze(-1),  # (b, s, l, 1)
         ]
         if self.use_paired_feature:
-            components.append(feats["msa_paired"].unsqueeze(-1))
-        msa_input = torch.cat(components, dim=-1)
+            components.append(feats["msa_paired"].unsqueeze(-1))  # (b, s, l, 1)
+        msa_input = torch.cat(components, dim=-1)  # (b, s, l, n_token_type + 2 or 3)
         if self.subsample_msa:
-            indices = torch.randperm(msa.shape[1])[: self.num_subsampled_msa]
-            msa_input = msa_input[:, indices]
-            msa_mask = msa_mask[:, indices]
-
-        msa_states = self.msa_proj(msa_input) + self.s_proj(emb).unsqueeze(1)
-        token_mask = feats["token_pad_mask"].float()
-        pair_mask = token_mask[:, :, None] * token_mask[:, None, :]
-        pair_states = z
+            indices = torch.randperm(msa.shape[1])[: self.num_subsampled_msa]  # (s_sub,)
+            msa_input = msa_input[:, indices]  # (b, s_sub, l, n_token_type + 2 or 3)
+            msa_mask = msa_mask[:, indices]  # (b, s_sub, l)
+
+        msa_states = self.msa_proj(msa_input) + self.s_proj(emb).unsqueeze(1)  # (b, s, l, d_m)
+        token_mask = feats["token_pad_mask"].float()  # (b, l)
+        pair_mask = token_mask[:, :, None] * token_mask[:, None, :]  # (b, l, l)
+        pair_states = z  # (b, l, l, d_z)
         for layer in self.layers:
+            # Tensor arguments: pair_states (b, l, l, d_z), msa_states (b, s, l, d_m),
+            # pair_mask (b, l, l), msa_mask (b, s, l).
             arguments = (
                 pair_states,
                 msa_states,
@@ -480,13 +504,13 @@ class MSAModule(nn.Module):
                 use_kernels,
             )
             if self.activation_checkpointing and self.training:
-                pair_states, msa_states = checkpoint(
+                pair_states, msa_states = checkpoint(  # (b, l, l, d_z), (b, s, l, d_m)
                     layer,
                     *arguments,
                 )
             else:
-                pair_states, msa_states = layer(*arguments)
-        return pair_states
+                pair_states, msa_states = layer(*arguments)  # (b, l, l, d_z), (b, s, l, d_m)
+        return pair_states  # (b, l, l, d_z)
 
 
 class MSALayer(nn.Module):
@@ -534,29 +558,31 @@ class MSALayer(nn.Module):
         """Return updated Z and M tensors."""
 
         del chunk_size_transition_z
-        dropout = get_dropout_mask(self.msa_dropout, m, self.training)
-        msa_states = m + dropout * self.pair_weighted_averaging(
+        # z: (b, l, l, d_z); m: (b, s, l, d_m).
+        # token_mask: (b, l, l); msa_mask: (b, s, l).
+        dropout = get_dropout_mask(self.msa_dropout, m, self.training)  # (b, s, 1, 1)
+        msa_states = m + dropout * self.pair_weighted_averaging(  # (b, s, l, d_m)
             m,
             z,
             token_mask,
             chunk_heads_pwa,
         )
-        msa_states = msa_states + self.msa_transition(
+        msa_states = msa_states + self.msa_transition(  # (b, s, l, d_m)
             msa_states,
             chunk_size_transition_msa,
         )
-        pair_states = z + self.outer_product_mean(
+        pair_states = z + self.outer_product_mean(  # (b, l, l, d_z)
             msa_states,
             msa_mask,
             chunk_size_outer_product,
         )
-        pair_states = self.pairformer_layer(
+        pair_states = self.pairformer_layer(  # (b, l, l, d_z)
             pair_states,
             token_mask,
             chunk_size_tri_attn,
             use_kernels=use_kernels,
         )
-        return pair_states, msa_states
+        return pair_states, msa_states  # (b, l, l, d_z), (b, s, l, d_m)
 
 
 class BFactorModule(nn.Module):
@@ -568,7 +594,8 @@ class BFactorModule(nn.Module):
         self.num_bins = num_bins
 
     def forward(self, s: Tensor) -> Tensor:
-        return cast(Tensor, self.bfactor(s))
+        # s: (..., d_s).
+        return cast(Tensor, self.bfactor(s))  # (..., n_bin)
 
 
 class DistogramModule(nn.Module):
@@ -581,8 +608,9 @@ class DistogramModule(nn.Module):
         self.num_bins = num_bins
 
     def forward(self, z: Tensor) -> Tensor:
-        symmetric = z + z.transpose(1, 2)
-        logits = self.distogram(symmetric)
+        # z: (b, l, l, d_z).
+        symmetric = z + z.transpose(1, 2)  # (b, l, l, d_z)
+        logits = self.distogram(symmetric)  # (b, l, l, n_distogram * n_bin)
         return cast(
             Tensor,
             logits.reshape(
@@ -592,4 +620,4 @@ class DistogramModule(nn.Module):
                 self.num_distograms,
                 self.num_bins,
             ),
-        )
+        )  # (b, l, l, n_distogram, n_bin)

fastplms/models/boltz/vb_modules_utils.py

@@ -6,14 +6,14 @@ and attribution records are retained in ``THIRD_PARTY_NOTICES.md``.
 
 from __future__ import annotations
 
+import torch
 from collections.abc import Iterable, Mapping, Sequence
 from functools import partial
 from typing import Any
-
-import torch
 from torch import Tensor, nn
 from torch.nn import functional as F
 
+
 LinearNoBias = partial(nn.Linear, bias=False)
 
 
@@ -32,19 +32,23 @@ def default(value: Any, fallback: Any) -> Any:
 def log(tensor: Tensor, eps: float = 1e-20) -> Tensor:
     """Compute a finite logarithm by applying a scalar lower bound."""
 
-    return torch.log(tensor.clamp(min=eps))
+    # tensor: (...).
+    return torch.log(tensor.clamp(min=eps))  # (...)
 
 
 class SwiGLU(nn.Module):
     """Split X in half and apply a SiLU-gated linear unit."""
 
     def forward(self, x: Tensor) -> Tensor:
-        values, gates = x.chunk(2, dim=-1)
-        return F.silu(gates) * values
+        # x: (..., 2 * d).
+        values, gates = x.chunk(2, dim=-1)  # each: (..., d)
+        return F.silu(gates) * values  # (..., d)
 
 
 def _masked_center(coordinates: Tensor, mask: Tensor) -> Tensor:
-    weights = mask[:, :, None]
+    # coordinates: (b, n_atom, 3); mask: (b, n_atom).
+    weights = mask[:, :, None]  # (b, n_atom, 1)
+    # (b, 1, 3)
     return (coordinates * weights).sum(dim=1, keepdim=True) / weights.sum(
         dim=1,
         keepdim=True,
@@ -54,21 +58,24 @@ def _masked_center(coordinates: Tensor, mask: Tensor) -> Tensor:
 def center(atom_coords: Tensor, atom_mask: Tensor) -> Tensor:
     """Center coordinate tensor X with shape ``(b, n_atoms, 3)``."""
 
-    return atom_coords - _masked_center(atom_coords, atom_mask)
+    # atom_coords: (b, n_atom, 3); atom_mask: (b, n_atom).
+    return atom_coords - _masked_center(atom_coords, atom_mask)  # (b, n_atom, 3)
 
 
 def _copysign(magnitudes: Tensor, signs: Tensor) -> Tensor:
     """Apply the elementwise sign of S to magnitude tensor M."""
 
-    return torch.where((magnitudes < 0) != (signs < 0), -magnitudes, magnitudes)
+    # magnitudes, signs: broadcast-compatible (...).
+    return torch.where((magnitudes < 0) != (signs < 0), -magnitudes, magnitudes)  # (...)
 
 
 def quaternion_to_matrix(quaternions: Tensor) -> Tensor:
     """Convert real-first quaternion tensor Q from ``(..., 4)`` to ``(..., 3, 3)``."""
 
-    real, i_axis, j_axis, k_axis = torch.unbind(quaternions, dim=-1)
-    scale = 2.0 / (quaternions * quaternions).sum(dim=-1)
-    entries = (
+    # quaternions: (..., 4).
+    real, i_axis, j_axis, k_axis = torch.unbind(quaternions, dim=-1)  # each: (...)
+    scale = 2.0 / (quaternions * quaternions).sum(dim=-1)  # (...)
+    entries = (  # nine tensors, each (...)
         1 - scale * (j_axis * j_axis + k_axis * k_axis),
         scale * (i_axis * j_axis - k_axis * real),
         scale * (i_axis * k_axis + j_axis * real),
@@ -79,6 +86,7 @@ def quaternion_to_matrix(quaternions: Tensor) -> Tensor:
         scale * (j_axis * k_axis + i_axis * real),
         1 - scale * (i_axis * i_axis + j_axis * j_axis),
     )
+    # (..., 3, 3)
     return torch.stack(entries, dim=-1).reshape((*quaternions.shape[:-1], 3, 3))
 
 
@@ -90,10 +98,10 @@ def random_quaternions(
     """Draw ``n`` uniformly distributed unit quaternions with nonnegative real part."""
 
     resolved_device = torch.device(device) if isinstance(device, str) else device
-    samples = torch.randn((n, 4), dtype=dtype, device=resolved_device)
-    squared_norm = (samples * samples).sum(dim=1)
-    signed_norm = _copysign(torch.sqrt(squared_norm), samples[:, 0])
-    return samples / signed_norm[:, None]
+    samples = torch.randn((n, 4), dtype=dtype, device=resolved_device)  # (n, 4)
+    squared_norm = (samples * samples).sum(dim=1)  # (n,)
+    signed_norm = _copysign(torch.sqrt(squared_norm), samples[:, 0])  # (n,)
+    return samples / signed_norm[:, None]  # (n, 4)
 
 
 def random_rotations(
@@ -103,6 +111,7 @@ def random_rotations(
 ) -> Tensor:
     """Draw rotation tensor R with shape ``(n, 3, 3)``."""
 
+    # (n, 3, 3)
     return quaternion_to_matrix(random_quaternions(n, dtype=dtype, device=device))
 
 
@@ -114,8 +123,8 @@ def compute_random_augmentation(
 ) -> tuple[Tensor, Tensor]:
     """Draw independent rotations R and translations T for a replicated batch."""
 
-    rotations = random_rotations(multiplicity, dtype=dtype, device=device)
-    translations = (
+    rotations = random_rotations(multiplicity, dtype=dtype, device=device)  # (m, 3, 3)
+    translations = (  # (m, 1, 3)
         torch.randn(
             (multiplicity, 1, 3),
             dtype=dtype,
@@ -123,7 +132,7 @@ def compute_random_augmentation(
         )
         * s_trans
     )
-    return rotations, translations
+    return rotations, translations  # (m, 3, 3), (m, 1, 3)
 
 
 def randomly_rotate(
@@ -133,14 +142,15 @@ def randomly_rotate(
 ) -> Tensor | tuple[Tensor, Tensor | None]:
     """Rotate X and optionally Y using the same sampled rotation tensor R."""
 
-    rotations = random_rotations(len(coords), coords.dtype, coords.device)
-    rotated = torch.einsum("bmd,bds->bms", coords, rotations)
+    # coords: (b, n, 3); second_coords: (b, n_second, 3) or None.
+    rotations = random_rotations(len(coords), coords.dtype, coords.device)  # (b, 3, 3)
+    rotated = torch.einsum("bmd,bds->bms", coords, rotations)  # (b, n, 3)
     if not return_second_coords:
-        return rotated
-    rotated_second = (
+        return rotated  # (b, n, 3)
+    rotated_second = (  # (b, n_second, 3) or None
         None if second_coords is None else torch.einsum("bmd,bds->bms", second_coords, rotations)
     )
-    return rotated, rotated_second
+    return rotated, rotated_second  # (b, n, 3), (b, n_second, 3) or None
 
 
 def center_random_augmentation(
@@ -154,25 +164,28 @@ def center_random_augmentation(
 ) -> Tensor | tuple[Tensor, Tensor | None]:
     """Center and rigidly augment coordinate tensors X and optional Y."""
 
-    primary = atom_coords
-    secondary = second_coords
+    # atom_coords: (b, n_atom, 3); atom_mask: (b, n_atom).
+    # second_coords: (b, n_second, 3) or None.
+    primary = atom_coords  # (b, n_atom, 3)
+    secondary = second_coords  # (b, n_second, 3) or None
     if centering:
-        centroid = _masked_center(primary, atom_mask)
-        primary = primary - centroid
+        centroid = _masked_center(primary, atom_mask)  # (b, 1, 3)
+        primary = primary - centroid  # (b, n_atom, 3)
         if secondary is not None:
-            secondary = secondary - centroid
+            secondary = secondary - centroid  # (b, n_second, 3)
 
     if augmentation:
-        primary, secondary = randomly_rotate(
+        primary, secondary = randomly_rotate(  # (b, n_atom, 3), (b, n_second, 3) or None
             primary,
             return_second_coords=True,
             second_coords=secondary,
         )
-        translation = torch.randn_like(primary[:, 0:1, :]) * s_trans
-        primary = primary + translation
+        translation = torch.randn_like(primary[:, 0:1, :]) * s_trans  # (b, 1, 3)
+        primary = primary + translation  # (b, n_atom, 3)
         if secondary is not None:
-            secondary = secondary + translation
+            secondary = secondary + translation  # (b, n_second, 3)
 
+    # (b, n_atom, 3), optionally paired with (b, n_second, 3) or None.
     return (primary, secondary) if return_second_coords else primary
 
 
@@ -185,21 +198,24 @@ class ExponentialMovingAverage:
         decay: float,
         use_num_updates: bool = True,
     ) -> None:
+        # parameters: one independently shaped tensor (...) per trainable parameter.
         if not 0.0 <= decay <= 1.0:
             raise ValueError("decay must lie in [0, 1]")
         self.decay = decay
         self.num_updates: int | None = 0 if use_num_updates else None
-        self.shadow_params = [
+        self.shadow_params = [  # each: same shape as one trainable parameter, (...)
             parameter.clone().detach() for parameter in parameters if parameter.requires_grad
         ]
         self.collected_params: list[Tensor] = []
 
     @staticmethod
     def _trainable(parameters: Iterable[nn.Parameter]) -> list[nn.Parameter]:
+        # Each item is one trainable parameter with an independently varying shape (...).
         return [parameter for parameter in parameters if parameter.requires_grad]
 
     def update(self, parameters: Iterable[nn.Parameter]) -> None:
-        trainable = self._trainable(parameters)
+        # parameters: one independently shaped tensor (...) per trainable parameter.
+        trainable = self._trainable(parameters)  # each: one trainable parameter, (...)
         if len(trainable) != len(self.shadow_params):
             raise ValueError("EMA parameter count changed")
         decay = self.decay
@@ -209,34 +225,40 @@ class ExponentialMovingAverage:
         update_weight = 1.0 - decay
         with torch.no_grad():
             for shadow, parameter in zip(self.shadow_params, trainable, strict=True):
-                shadow.sub_(update_weight * (shadow - parameter))
+                shadow.sub_(update_weight * (shadow - parameter))  # (...) unchanged in place
 
     def compatible(self, parameters: Sequence[Tensor]) -> bool:
         """Return whether parameter count and tensor shapes match the EMA state."""
 
+        # Each shadow and corresponding parameter has an independently varying shape (...).
         return len(self.shadow_params) == len(parameters) and all(
             shadow.shape == parameter.shape
             for shadow, parameter in zip(self.shadow_params, parameters, strict=True)
         )
 
     def copy_to(self, parameters: Iterable[nn.Parameter]) -> None:
-        trainable = self._trainable(parameters)
+        # parameters: one independently shaped tensor (...) per trainable parameter.
+        trainable = self._trainable(parameters)  # each: one trainable parameter, (...)
         if len(trainable) != len(self.shadow_params):
             raise ValueError("EMA parameter count changed")
         for shadow, parameter in zip(self.shadow_params, trainable, strict=True):
-            parameter.data.copy_(shadow.data)
+            parameter.data.copy_(shadow.data)  # (...) unchanged in place
 
     def store(self, parameters: Iterable[nn.Parameter]) -> None:
+        # parameters: one independently shaped tensor (...) per parameter.
+        # Each clone has the same shape (...) as its corresponding parameter.
         self.collected_params = [parameter.clone() for parameter in parameters]
 
     def restore(self, parameters: Iterable[nn.Parameter]) -> None:
-        current = list(parameters)
+        # parameters: one independently shaped tensor (...) per parameter.
+        current = list(parameters)  # each: one parameter, (...)
         if len(current) != len(self.collected_params):
             raise ValueError("stored EMA parameter count changed")
         for collected, parameter in zip(self.collected_params, current, strict=True):
-            parameter.data.copy_(collected.data)
+            parameter.data.copy_(collected.data)  # (...) unchanged in place
 
     def state_dict(self) -> dict[str, Any]:
+        # shadow_params contains independently shaped tensors (...).
         return {
             "decay": self.decay,
             "num_updates": self.num_updates,
@@ -248,10 +270,13 @@ class ExponentialMovingAverage:
         state_dict: Mapping[str, Any],
         device: torch.device | str,
     ) -> None:
+        # state_dict["shadow_params"] contains independently shaped tensors (...).
         self.decay = float(state_dict["decay"])
         updates = state_dict["num_updates"]
         self.num_updates = None if updates is None else int(updates)
+        # Each tensor retains its shape (...).
         self.shadow_params = [tensor.to(device) for tensor in state_dict["shadow_params"]]
 
     def to(self, device: torch.device | str) -> None:
+        # Each shadow parameter retains its independently varying shape (...).
         self.shadow_params = [tensor.to(device) for tensor in self.shadow_params]

fastplms/models/boltz/vb_potentials_potentials.py

@@ -9,12 +9,11 @@ oracle while preserving the public class names used by converted checkpoints.
 
 from __future__ import annotations
 
+import torch
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Any
 
-import torch
-
 from . import vb_const as const
 from .vb_loss_diffusionv2 import weighted_rigid_align
 from .vb_potentials_schedules import (
@@ -23,22 +22,23 @@ from .vb_potentials_schedules import (
     PiecewiseStepFunction,
 )
 
+
 Parameter = ParameterSchedule | float | int | bool
 ParameterMap = dict[str, Parameter]
 
 
 @dataclass(frozen=True, slots=True)
 class _PreparedPotential:
-    coordinates: torch.Tensor
-    index: torch.Tensor
+    coordinates: torch.Tensor  # (..., a, 3)
+    index: torch.Tensor  # (q, n)
     function_args: tuple[Any, ...]
-    com_index: torch.Tensor | None
-    atom_pad_mask: torch.Tensor | None
-    ref_coords: torch.Tensor | None
-    ref_mask: torch.Tensor | None
-    ref_token_index: torch.Tensor | None
-    negation_mask: torch.Tensor | None
-    union_index: torch.Tensor | None
+    com_index: torch.Tensor | None  # (a_input,)
+    atom_pad_mask: torch.Tensor | None  # (a_input,)
+    ref_coords: torch.Tensor | None  # (n_ref, n, 3)
+    ref_mask: torch.Tensor | None  # (n_ref, n)
+    ref_token_index: torch.Tensor | None  # (a_input,)
+    negation_mask: torch.Tensor | None  # (n,)
+    union_index: torch.Tensor | None  # (n,)
 
 
 def _reduce_atoms_to_centers(
@@ -46,15 +46,16 @@ def _reduce_atoms_to_centers(
     com_index: torch.Tensor,
     atom_pad_mask: torch.Tensor,
 ) -> torch.Tensor:
-    unpadded_index = com_index[atom_pad_mask]
-    unpadded_coordinates = coordinates[..., atom_pad_mask, :]
+    # coordinates: (..., a, 3); com_index/atom_pad_mask: (a,).
+    unpadded_index = com_index[atom_pad_mask]  # (a_valid,)
+    unpadded_coordinates = coordinates[..., atom_pad_mask, :]  # (..., a_valid, 3)
     center_shape = (*unpadded_coordinates.shape[:-2], unpadded_index.max() + 1, 3)
     return torch.zeros(center_shape, device=coordinates.device).scatter_reduce(
         -2,
         unpadded_index.unsqueeze(-1).expand_as(unpadded_coordinates),
         unpadded_coordinates,
         "mean",
-    )
+    )  # (..., n_center, 3)
 
 
 def _union_weights(
@@ -62,7 +63,8 @@ def _union_weights(
     union_index: torch.Tensor,
     union_lambda: float,
 ) -> torch.Tensor:
-    unnormalized = torch.exp(-union_lambda * energy)
+    # energy: (..., n); union_index: (n,).
+    unnormalized = torch.exp(-union_lambda * energy)  # (..., n)
     partition = torch.zeros(
         (*energy.shape[:-1], union_index.max() + 1),
         device=union_index.device,
@@ -71,10 +73,10 @@ def _union_weights(
         union_index.expand_as(unnormalized),
         unnormalized,
         "sum",
-    )
-    weights = unnormalized / partition[..., union_index]
-    weights[partition[..., union_index] == 0] = 0
-    return weights
+    )  # (..., n_union)
+    weights = unnormalized / partition[..., union_index]  # (..., n)
+    weights[partition[..., union_index] == 0] = 0  # (..., n)
+    return weights  # (..., n)
 
 
 def _scatter_constraint_gradients(
@@ -83,20 +85,22 @@ def _scatter_constraint_gradients(
     index: torch.Tensor,
     coordinates: torch.Tensor,
 ) -> torch.Tensor:
+    # coefficients: (..., n); variable_gradient: (..., q, n, 3).
+    # index: (q, n); coordinates: (b, a, 3).
     product = coefficients.tile(variable_gradient.shape[-3]).unsqueeze(
         -1
-    ) * variable_gradient.flatten(start_dim=-3, end_dim=-2)
+    ) * variable_gradient.flatten(start_dim=-3, end_dim=-2)  # (..., q * n, 3)
     if product.dim() > 3:
-        product = product.sum(dim=list(range(1, product.dim() - 2)))
+        product = product.sum(dim=list(range(1, product.dim() - 2)))  # (b, q * n, 3)
     scatter_index = (
         index.flatten(start_dim=0, end_dim=1).unsqueeze(-1).expand((*coordinates.shape[:-2], -1, 3))
-    )
+    )  # (b, q * n, 3)
     return torch.zeros_like(coordinates).scatter_reduce(
         -2,
         scatter_index,
         product,
         "sum",
-    )
+    )  # (b, a, 3)
 
 
 class Potential(ABC):
@@ -112,26 +116,29 @@ class Potential(ABC):
         parameters: dict[str, Any],
         computed_args: tuple[Any, ...] | None = None,
     ) -> _PreparedPotential:
+        # coordinates: (..., a_input, 3).
         if computed_args is None:
             computed_args = self.compute_args(feats, parameters)
-        index, args, com_args, ref_args, operator_args = computed_args
+        index, args, com_args, ref_args, operator_args = computed_args  # index: (q, n)
         com_index = atom_pad_mask = None
         if com_args is not None:
-            com_index, atom_pad_mask = com_args
+            com_index, atom_pad_mask = com_args  # each (a_input,)
             coordinates = _reduce_atoms_to_centers(
                 coordinates,
                 com_index,
                 atom_pad_mask,
-            )
+            )  # (..., n_center, 3)
 
         ref_coords = ref_mask = ref_token_index = None
         if ref_args is not None:
+            # ref_coords: (n_ref, n, 3); ref_mask: (n_ref, n).
+            # ref_atom_index: (n,); ref_token_index: (a_input,).
             ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args
-            coordinates = coordinates[..., ref_atom_index, :]
+            coordinates = coordinates[..., ref_atom_index, :]  # (..., n, 3)
 
         negation_mask = union_index = None
         if operator_args is not None:
-            negation_mask, union_index = operator_args
+            negation_mask, union_index = operator_args  # each (n,)
         return _PreparedPotential(
             coordinates=coordinates,
             index=index,
@@ -153,9 +160,10 @@ class Potential(ABC):
     ) -> torch.Tensor:
         """Evaluate one energy per coordinate sample."""
 
+        # coords: (..., a, 3).
         computed_args = self.compute_args(feats, parameters)
         if computed_args[0].shape[1] == 0:
-            return torch.zeros(coords.shape[:-2], device=coords.device)
+            return torch.zeros(coords.shape[:-2], device=coords.device)  # (...)
         prepared = self._prepare(coords, feats, parameters, computed_args)
         value = self.compute_variable(
             prepared.coordinates,
@@ -163,21 +171,21 @@ class Potential(ABC):
             ref_coords=prepared.ref_coords,
             ref_mask=prepared.ref_mask,
             compute_gradient=False,
-        )
+        )  # (..., n)
         energy = self.compute_function(
             value,
             *prepared.function_args,
             negation_mask=prepared.negation_mask,
             compute_derivative=False,
-        )
+        )  # (..., n)
         if prepared.union_index is not None:
             weights = _union_weights(
                 energy,
                 prepared.union_index,
                 parameters["union_lambda"],
-            )
-            return (energy * weights).sum(dim=-1)
-        return energy.sum(dim=tuple(range(1, energy.dim())))
+            )  # (..., n)
+            return (energy * weights).sum(dim=-1)  # (...)
+        return energy.sum(dim=tuple(range(1, energy.dim())))  # (b,)
 
     def compute_gradient(
         self,
@@ -187,9 +195,10 @@ class Potential(ABC):
     ) -> torch.Tensor:
         """Return the analytic coordinate gradient of the potential energy."""
 
+        # coords: (b, a, 3).
         computed_args = self.compute_args(feats, parameters)
         if computed_args[0].shape[1] == 0:
-            return torch.zeros_like(coords)
+            return torch.zeros_like(coords)  # (b, a, 3)
         prepared = self._prepare(coords, feats, parameters, computed_args)
         value, variable_gradient = self.compute_variable(
             prepared.coordinates,
@@ -197,19 +206,19 @@ class Potential(ABC):
             ref_coords=prepared.ref_coords,
             ref_mask=prepared.ref_mask,
             compute_gradient=True,
-        )
+        )  # value: (..., n); variable_gradient: (..., q, n, 3)
         energy, energy_derivative = self.compute_function(
             value,
             *prepared.function_args,
             negation_mask=prepared.negation_mask,
             compute_derivative=True,
-        )
+        )  # each (..., n)
         if prepared.union_index is not None:
             weights = _union_weights(
                 energy,
                 prepared.union_index,
                 parameters["union_lambda"],
-            )
+            )  # (..., n)
             union_energy = torch.zeros(
                 (*energy.shape[:-1], prepared.union_index.max() + 1),
                 device=prepared.union_index.device,
@@ -218,7 +227,7 @@ class Potential(ABC):
                 prepared.union_index.expand_as(energy),
                 energy * weights,
                 "sum",
-            )
+            )  # (..., n_union)
             coefficients = (
                 energy_derivative
                 * weights
@@ -227,21 +236,21 @@ class Potential(ABC):
                     + parameters["union_lambda"]
                     * (energy - union_energy[..., prepared.union_index])
                 )
-            )
+            )  # (..., n)
         else:
-            coefficients = energy_derivative
+            coefficients = energy_derivative  # (..., n)
 
         atom_gradient = _scatter_constraint_gradients(
             coefficients,
             variable_gradient,
             prepared.index,
             prepared.coordinates,
-        )
+        )  # (b, a_prepared, 3)
         if prepared.com_index is not None:
-            atom_gradient = atom_gradient[..., prepared.com_index, :]
+            atom_gradient = atom_gradient[..., prepared.com_index, :]  # (b, a, 3)
         elif prepared.ref_token_index is not None:
-            atom_gradient = atom_gradient[..., prepared.ref_token_index, :]
-        return atom_gradient
+            atom_gradient = atom_gradient[..., prepared.ref_token_index, :]  # (b, a, 3)
+        return atom_gradient  # (b, a, 3)
 
     def compute_parameters(self, t: float) -> dict[str, Any] | None:
         """Resolve scheduled parameters at diffusion time ``t``."""
@@ -261,6 +270,7 @@ class Potential(ABC):
         negation_mask: torch.Tensor | None = None,
         compute_derivative: bool = False,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # value: (..., n); negation_mask: (n,).
         raise NotImplementedError
 
     @abstractmethod
@@ -272,6 +282,7 @@ class Potential(ABC):
         ref_mask: torch.Tensor | None = None,
         compute_gradient: bool = False,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # coords: (..., a, 3); index: (q, n).
         raise NotImplementedError
 
     @abstractmethod
@@ -303,16 +314,18 @@ class FlatBottomPotential(Potential):
         negation_mask: torch.Tensor | None = None,
         compute_derivative: bool = False,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # value: (..., n); k/lower_bounds/upper_bounds: (n,) or (..., n).
+        # negation_mask: (n,) or a shape broadcastable to (..., n).
         lower = (
             (torch.full_like(value, float("-inf")) if lower_bounds is None else lower_bounds)
             .expand_as(value)
             .clone()
-        )
+        )  # (..., n)
         upper = (
             (torch.full_like(value, float("inf")) if upper_bounds is None else upper_bounds)
             .expand_as(value)
             .clone()
-        )
+        )  # (..., n)
 
         if negation_mask is not None:
             if not torch.is_tensor(negation_mask):
@@ -329,37 +342,38 @@ class FlatBottomPotential(Potential):
                     f"got {negation_mask.device} and {value.device}."
                 )
             try:
-                expanded_negation_mask = negation_mask.expand_as(value)
+                expanded_negation_mask = negation_mask.expand_as(value)  # (..., n)
             except RuntimeError as error:
                 raise ValueError(
                     "negation_mask must be broadcastable to value shape "
                     f"{tuple(value.shape)}, got {tuple(negation_mask.shape)}."
                 ) from error
-            unbounded_below = torch.isneginf(lower)
-            unbounded_above = torch.isposinf(upper)
+            unbounded_below = torch.isneginf(lower)  # (..., n)
+            unbounded_above = torch.isposinf(upper)  # (..., n)
+            # (..., n)
             valid_negation = unbounded_below | unbounded_above | expanded_negation_mask
             if not bool(torch.all(valid_negation).item()):
                 raise ValueError(
                     "negation_mask may be false only where at least one bound is infinite."
                 )
-            select_upper = ~unbounded_above & ~expanded_negation_mask
-            lower[select_upper] = upper[select_upper]
-            upper[select_upper] = float("inf")
-            select_lower = ~unbounded_below & ~expanded_negation_mask
-            upper[select_lower] = lower[select_lower]
-            lower[select_lower] = float("-inf")
-
-        below = value < lower
-        above = value > upper
-        energy = torch.zeros_like(value)
-        energy[below] = (k * (lower - value))[below]
-        energy[above] = (k * (value - upper))[above]
+            select_upper = ~unbounded_above & ~expanded_negation_mask  # (..., n)
+            lower[select_upper] = upper[select_upper]  # (..., n)
+            upper[select_upper] = float("inf")  # (..., n)
+            select_lower = ~unbounded_below & ~expanded_negation_mask  # (..., n)
+            upper[select_lower] = lower[select_lower]  # (..., n)
+            lower[select_lower] = float("-inf")  # (..., n)
+
+        below = value < lower  # (..., n)
+        above = value > upper  # (..., n)
+        energy = torch.zeros_like(value)  # (..., n)
+        energy[below] = (k * (lower - value))[below]  # (..., n)
+        energy[above] = (k * (value - upper))[above]  # (..., n)
         if not compute_derivative:
-            return energy
-        derivative = torch.zeros_like(value)
-        derivative[below] = -k.expand_as(below)[below]
-        derivative[above] = k.expand_as(above)[above]
-        return energy, derivative
+            return energy  # (..., n)
+        derivative = torch.zeros_like(value)  # (..., n)
+        derivative[below] = -k.expand_as(below)[below]  # (..., n)
+        derivative[above] = k.expand_as(above)[above]  # (..., n)
+        return energy, derivative  # each (..., n)
 
 
 class ReferencePotential(Potential):
@@ -373,19 +387,22 @@ class ReferencePotential(Potential):
         ref_mask: torch.Tensor,
         compute_gradient: bool = False,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # coords: (b, a, 3); index: (1, n).
+        # ref_coords: (n_ref, n, 3); ref_mask: (n_ref, n).
         aligned_reference = weighted_rigid_align(
             ref_coords.float(),
             coords[:, index].float(),
             ref_mask,
             ref_mask,
-        )
-        displacement = coords[:, index] - aligned_reference
-        distance = torch.linalg.norm(displacement, dim=-1)
+        )  # (b, n_ref, n, 3)
+        displacement = coords[:, index] - aligned_reference  # (b, n_ref, n, 3)
+        distance = torch.linalg.norm(displacement, dim=-1)  # (b, n_ref, n)
         if not compute_gradient:
-            return distance
-        unit_displacement = displacement / distance.unsqueeze(-1)
+            return distance  # (b, n_ref, n)
+        unit_displacement = displacement / distance.unsqueeze(-1)  # (b, n_ref, n, 3)
+        # (b, 1, n_ref, n, 3)
         gradient = (unit_displacement * ref_mask.unsqueeze(-1)).unsqueeze(1)
-        return distance, gradient
+        return distance, gradient  # (b, n_ref, n); (b, 1, n_ref, n, 3)
 
 
 class DistancePotential(Potential):
@@ -399,15 +416,17 @@ class DistancePotential(Potential):
         ref_mask: torch.Tensor | None = None,
         compute_gradient: bool = False,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # coords: (..., a, 3); index: (2, n).
         del ref_coords, ref_mask
         displacement = coords.index_select(-2, index[0]) - coords.index_select(
             -2,
             index[1],
-        )
-        distance = torch.linalg.norm(displacement, dim=-1)
+        )  # (..., n, 3)
+        distance = torch.linalg.norm(displacement, dim=-1)  # (..., n)
         if not compute_gradient:
-            return distance
-        unit_displacement = displacement / distance.unsqueeze(-1)
+            return distance  # (..., n)
+        unit_displacement = displacement / distance.unsqueeze(-1)  # (..., n, 3)
+        # Returns (..., n); (..., 2, n, 3).
         return distance, torch.stack((unit_displacement, -unit_displacement), dim=1)
 
 
@@ -422,36 +441,40 @@ class DihedralPotential(Potential):
         ref_mask: torch.Tensor | None = None,
         compute_gradient: bool = False,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # coords: (..., a, 3); index: (4, n).
         del ref_coords, ref_mask
+        # Each displacement has shape (..., n, 3).
         r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1])
         r_kj = coords.index_select(-2, index[2]) - coords.index_select(-2, index[1])
         r_kl = coords.index_select(-2, index[2]) - coords.index_select(-2, index[3])
-        n_ijk = torch.cross(r_ij, r_kj, dim=-1)
-        n_jkl = torch.cross(r_kj, r_kl, dim=-1)
-        r_kj_norm = torch.linalg.norm(r_kj, dim=-1)
-        n_ijk_norm = torch.linalg.norm(n_ijk, dim=-1)
-        n_jkl_norm = torch.linalg.norm(n_jkl, dim=-1)
+        n_ijk = torch.cross(r_ij, r_kj, dim=-1)  # (..., n, 3)
+        n_jkl = torch.cross(r_kj, r_kl, dim=-1)  # (..., n, 3)
+        r_kj_norm = torch.linalg.norm(r_kj, dim=-1)  # (..., n)
+        n_ijk_norm = torch.linalg.norm(n_ijk, dim=-1)  # (..., n)
+        n_jkl_norm = torch.linalg.norm(n_jkl, dim=-1)  # (..., n)
 
         orientation = torch.sign(
             r_kj.unsqueeze(-2) @ torch.cross(n_ijk, n_jkl, dim=-1).unsqueeze(-1)
-        ).squeeze(-1, -2)
+        ).squeeze(-1, -2)  # (..., n)
         cosine = (n_ijk.unsqueeze(-2) @ n_jkl.unsqueeze(-1)).squeeze(-1, -2) / (
             n_ijk_norm * n_jkl_norm
-        )
+        )  # (..., n)
+        # (..., n)
         angle = orientation * torch.arccos(torch.clamp(cosine, -1 + 1e-8, 1 - 1e-8))
         if not compute_gradient:
-            return angle
+            return angle  # (..., n)
 
         projection_i = (
             (r_ij.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
-        ).unsqueeze(-1)
+        ).unsqueeze(-1)  # (..., n, 1)
         projection_l = (
             (r_kl.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
-        ).unsqueeze(-1)
-        grad_i = n_ijk * (r_kj_norm / n_ijk_norm**2).unsqueeze(-1)
-        grad_l = -n_jkl * (r_kj_norm / n_jkl_norm**2).unsqueeze(-1)
-        grad_j = (projection_i - 1) * grad_i - projection_l * grad_l
-        grad_k = (projection_l - 1) * grad_l - projection_i * grad_i
+        ).unsqueeze(-1)  # (..., n, 1)
+        grad_i = n_ijk * (r_kj_norm / n_ijk_norm**2).unsqueeze(-1)  # (..., n, 3)
+        grad_l = -n_jkl * (r_kj_norm / n_jkl_norm**2).unsqueeze(-1)  # (..., n, 3)
+        grad_j = (projection_i - 1) * grad_i - projection_l * grad_l  # (..., n, 3)
+        grad_k = (projection_l - 1) * grad_l - projection_i * grad_i  # (..., n, 3)
+        # Returns (..., n); (..., 4, n, 3).
         return angle, torch.stack((grad_i, grad_j, grad_k, grad_l), dim=1)
 
 
@@ -466,33 +489,37 @@ class AbsDihedralPotential(DihedralPotential):
         ref_mask: torch.Tensor | None = None,
         compute_gradient: bool = False,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # coords: (..., a, 3); index: (4, n).
         del ref_coords, ref_mask
         if not compute_gradient:
-            return torch.abs(super().compute_variable(coords, index))
+            return torch.abs(super().compute_variable(coords, index))  # (..., n)
         angle, gradient = super().compute_variable(
             coords,
             index,
             compute_gradient=True,
-        )
-        gradient[(angle < 0)[..., None, :, None].expand_as(gradient)] *= -1
-        return torch.abs(angle), gradient
+        )  # (..., n); (..., 4, n, 3)
+        gradient[(angle < 0)[..., None, :, None].expand_as(gradient)] *= -1  # same
+        return torch.abs(angle), gradient  # (..., n); (..., 4, n, 3)
 
 
 def _element_radii(feats: dict[str, torch.Tensor]) -> torch.Tensor:
+    # feats["ref_element"]: (1, a, n_element).
     element_radii = torch.zeros(
         const.num_elements,
         dtype=torch.float32,
         device=feats["ref_element"].device,
-    )
+    )  # (n_element,)
     element_radii[1:119] = torch.tensor(
         const.vdw_radii,
         dtype=torch.float32,
         device=element_radii.device,
-    )
+    )  # (n_element,)
+    # (a,)
     return (feats["ref_element"].float() @ element_radii.unsqueeze(-1)).squeeze(-1)[0]
 
 
 def _atom_chain_ids(feats: dict[str, torch.Tensor]) -> torch.Tensor:
+    # atom_to_token: (1, a, l); asym_id: (1, l).
     return (
         torch.bmm(
             feats["atom_to_token"].float(),
@@ -500,94 +527,103 @@ def _atom_chain_ids(feats: dict[str, torch.Tensor]) -> torch.Tensor:
         )
         .squeeze(-1)
         .long()
-    )[0]
+    )[0]  # (a,)
 
 
 class PoseBustersPotential(FlatBottomPotential, DistancePotential):
     def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
-        pair_index = feats["rdkit_bounds_index"][0]
-        lower = feats["rdkit_lower_bounds"][0].clone()
-        upper = feats["rdkit_upper_bounds"][0].clone()
-        bond = feats["rdkit_bounds_bond_mask"][0]
-        angle = feats["rdkit_bounds_angle_mask"][0]
-        lower[bond * ~angle] *= 1.0 - parameters["bond_buffer"]
-        upper[bond * ~angle] *= 1.0 + parameters["bond_buffer"]
-        lower[~bond * angle] *= 1.0 - parameters["angle_buffer"]
-        upper[~bond * angle] *= 1.0 + parameters["angle_buffer"]
+        # rdkit_bounds_index: (1, 2, n); bound and mask features: (1, n).
+        pair_index = feats["rdkit_bounds_index"][0]  # (2, n)
+        lower = feats["rdkit_lower_bounds"][0].clone()  # (n,)
+        upper = feats["rdkit_upper_bounds"][0].clone()  # (n,)
+        bond = feats["rdkit_bounds_bond_mask"][0]  # (n,)
+        angle = feats["rdkit_bounds_angle_mask"][0]  # (n,)
+        lower[bond * ~angle] *= 1.0 - parameters["bond_buffer"]  # (n,)
+        upper[bond * ~angle] *= 1.0 + parameters["bond_buffer"]  # (n,)
+        lower[~bond * angle] *= 1.0 - parameters["angle_buffer"]  # (n,)
+        upper[~bond * angle] *= 1.0 + parameters["angle_buffer"]  # (n,)
         shared_buffer = min(parameters["bond_buffer"], parameters["angle_buffer"])
-        lower[bond * angle] *= 1.0 - shared_buffer
-        upper[bond * angle] *= 1.0 + shared_buffer
-        lower[~bond * ~angle] *= 1.0 - parameters["clash_buffer"]
-        upper[~bond * ~angle] = float("inf")
-
-        atom_radii = _element_radii(feats)
-        bond_cutoff = 0.35 + atom_radii[pair_index].mean(dim=0)
-        lower[~bond] = torch.max(lower[~bond], bond_cutoff[~bond])
-        upper[bond] = torch.min(upper[bond], bond_cutoff[bond])
+        lower[bond * angle] *= 1.0 - shared_buffer  # (n,)
+        upper[bond * angle] *= 1.0 + shared_buffer  # (n,)
+        lower[~bond * ~angle] *= 1.0 - parameters["clash_buffer"]  # (n,)
+        upper[~bond * ~angle] = float("inf")  # (n,)
+
+        atom_radii = _element_radii(feats)  # (a,)
+        bond_cutoff = 0.35 + atom_radii[pair_index].mean(dim=0)  # (n,)
+        lower[~bond] = torch.max(lower[~bond], bond_cutoff[~bond])  # (n,)
+        upper[bond] = torch.min(upper[bond], bond_cutoff[bond])  # (n,)
+        # Returns index (2, n) and three function arguments of shape (n,).
         return pair_index, (torch.ones_like(lower), lower, upper), None, None, None
 
 
 class ConnectionsPotential(FlatBottomPotential, DistancePotential):
     def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
-        pair_index = feats["connected_atom_index"][0]
+        # connected_atom_index: (1, 2, n).
+        pair_index = feats["connected_atom_index"][0]  # (2, n)
         upper = torch.full(
             (pair_index.shape[1],),
             parameters["buffer"],
             device=pair_index.device,
-        )
+        )  # (n,)
+        # Returns index (2, n), k (n,), and upper bound (n,).
         return pair_index, (torch.ones_like(upper), None, upper), None, None, None
 
 
 class VDWOverlapPotential(FlatBottomPotential, DistancePotential):
     def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
-        atom_chain_id = _atom_chain_ids(feats)
-        atom_pad_mask = feats["atom_pad_mask"][0].bool()
-        chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask])
-        nonion_atom = (chain_sizes > 1)[atom_chain_id]
-        atom_radii = _element_radii(feats)
+        atom_chain_id = _atom_chain_ids(feats)  # (a,)
+        atom_pad_mask = feats["atom_pad_mask"][0].bool()  # (a,)
+        chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask])  # (n_chain,)
+        nonion_atom = (chain_sizes > 1)[atom_chain_id]  # (a,)
+        atom_radii = _element_radii(feats)  # (a,)
         pair_index = torch.triu_indices(
             atom_chain_id.shape[0],
             atom_chain_id.shape[0],
             1,
             device=atom_chain_id.device,
-        )
-        pair_pad_mask = atom_pad_mask[pair_index].all(dim=0)
+        )  # (2, n_pair)
+        pair_pad_mask = atom_pad_mask[pair_index].all(dim=0)  # (n_pair,)
+        # (n_pair,)
         pair_ion_mask = nonion_atom[pair_index[0]] * nonion_atom[pair_index[1]]
 
-        num_chains = atom_chain_id.max() + 1
-        connected = feats["connected_chain_index"][0]
+        num_chains = atom_chain_id.max() + 1  # ()
+        connected = feats["connected_chain_index"][0]  # (2, n_connection)
         connected_matrix = torch.eye(
             num_chains,
             device=atom_chain_id.device,
             dtype=torch.bool,
-        )
-        connected_matrix[connected[0], connected[1]] = True
-        connected_matrix[connected[1], connected[0]] = True
+        )  # (n_chain, n_chain)
+        connected_matrix[connected[0], connected[1]] = True  # (n_chain, n_chain)
+        connected_matrix[connected[1], connected[0]] = True  # (n_chain, n_chain)
         connected_pair = connected_matrix[
             atom_chain_id[pair_index[0]],
             atom_chain_id[pair_index[1]],
-        ]
+        ]  # (n_pair,)
+        # (2, n)
         pair_index = pair_index[:, pair_pad_mask * pair_ion_mask * ~connected_pair]
+        # (n,)
         lower = atom_radii[pair_index].sum(dim=0) * (1.0 - parameters["buffer"])
+        # Returns index (2, n), k (n,), and lower bound (n,).
         return pair_index, (torch.ones_like(lower), lower, None), None, None, None
 
 
 class SymmetricChainCOMPotential(FlatBottomPotential, DistancePotential):
     def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
-        atom_chain_id = _atom_chain_ids(feats)
-        atom_pad_mask = feats["atom_pad_mask"][0].bool()
-        nonion_chain = torch.bincount(atom_chain_id[atom_pad_mask]) > 1
-        pair_index = feats["symmetric_chain_index"][0]
+        atom_chain_id = _atom_chain_ids(feats)  # (a,)
+        atom_pad_mask = feats["atom_pad_mask"][0].bool()  # (a,)
+        nonion_chain = torch.bincount(atom_chain_id[atom_pad_mask]) > 1  # (n_chain,)
+        pair_index = feats["symmetric_chain_index"][0]  # (2, n_candidate)
         pair_index = pair_index[
             :,
             nonion_chain[pair_index[0]] * nonion_chain[pair_index[1]],
-        ]
+        ]  # (2, n)
         lower = torch.full(
             (pair_index.shape[1],),
             parameters["buffer"],
             dtype=torch.float32,
             device=pair_index.device,
-        )
+        )  # (n,)
+        # Returns center index (2, n), k/lower (n,), and atom-to-center maps (a,).
         return (
             pair_index,
             (torch.ones_like(lower), lower, None),
@@ -602,52 +638,60 @@ def _oriented_bounds(
     positive_lower: float,
     negative_upper: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    lower = torch.zeros(orientations.shape, device=orientations.device)
-    upper = torch.zeros(orientations.shape, device=orientations.device)
-    lower[orientations] = positive_lower
-    upper[orientations] = float("inf")
-    lower[~orientations] = float("-inf")
-    upper[~orientations] = negative_upper
-    return lower, upper
+    # orientations: (n,).
+    lower = torch.zeros(orientations.shape, device=orientations.device)  # (n,)
+    upper = torch.zeros(orientations.shape, device=orientations.device)  # (n,)
+    lower[orientations] = positive_lower  # (n,)
+    upper[orientations] = float("inf")  # (n,)
+    lower[~orientations] = float("-inf")  # (n,)
+    upper[~orientations] = negative_upper  # (n,)
+    return lower, upper  # each (n,)
 
 
 class StereoBondPotential(FlatBottomPotential, AbsDihedralPotential):
     def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
-        index = feats["stereo_bond_index"][0]
-        orientation = feats["stereo_bond_orientations"][0].bool()
+        # stereo_bond_index: (1, 4, n); orientations: (1, n).
+        index = feats["stereo_bond_index"][0]  # (4, n)
+        orientation = feats["stereo_bond_orientations"][0].bool()  # (n,)
         lower, upper = _oriented_bounds(
             orientation,
             torch.pi - parameters["buffer"],
             parameters["buffer"],
-        )
+        )  # each (n,)
+        # Returns index (4, n) and three function arguments of shape (n,).
         return index, (torch.ones_like(lower), lower, upper), None, None, None
 
 
 class ChiralAtomPotential(FlatBottomPotential, DihedralPotential):
     def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
-        index = feats["chiral_atom_index"][0]
-        orientation = feats["chiral_atom_orientations"][0].bool()
+        # chiral_atom_index: (1, 4, n); orientations: (1, n).
+        index = feats["chiral_atom_index"][0]  # (4, n)
+        orientation = feats["chiral_atom_orientations"][0].bool()  # (n,)
         lower, upper = _oriented_bounds(
             orientation,
             parameters["buffer"],
             -parameters["buffer"],
-        )
+        )  # each (n,)
+        # Returns index (4, n) and three function arguments of shape (n,).
         return index, (torch.ones_like(lower), lower, upper), None, None, None
 
 
 class PlanarBondPotential(FlatBottomPotential, AbsDihedralPotential):
     def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
-        bond_index = feats["planar_bond_index"][0].T
+        # The feature stores two atom-index rows for each six-entry bond pattern.
+        bond_index = feats["planar_bond_index"][0].T  # (2, n_bond_entry)
         improper_pattern = torch.tensor(
             [[1, 2, 3, 0], [4, 5, 0, 3]],
             device=bond_index.device,
-        ).T
+        ).T  # (4, 2)
+        # (4, n_improper)
         improper_index = bond_index[:, improper_pattern].swapaxes(0, 1).flatten(start_dim=1)
         upper = torch.full(
             (improper_index.shape[1],),
             parameters["buffer"],
             device=improper_index.device,
-        )
+        )  # (n_improper,)
+        # Returns index (4, n_improper), k (n_improper,), and upper (n_improper,).
         return (
             improper_index,
             (torch.ones_like(upper), None, upper),
@@ -661,25 +705,28 @@ class TemplateReferencePotential(FlatBottomPotential, ReferencePotential):
     def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
         del parameters
         if "template_mask_cb" not in feats or "template_force" not in feats:
+            # Empty sentinel index: (1, 0).
             return torch.empty((1, 0)), None, None, None, None
-        template_mask = feats["template_mask_cb"][feats["template_force"]]
+        # template_mask_cb/template_force: (b, n_template, n)/(b, n_template).
+        template_mask = feats["template_mask_cb"][feats["template_force"]]  # (n_ref, n)
         if template_mask.shape[0] == 0:
+            # Empty sentinel index: (1, 0).
             return torch.empty((1, 0)), None, None, None, None
 
-        ref_coords = feats["template_cb"][feats["template_force"]].clone()
-        ref_mask = feats["template_mask_cb"][feats["template_force"]].clone()
+        ref_coords = feats["template_cb"][feats["template_force"]].clone()  # (n_ref, n, 3)
+        ref_mask = feats["template_mask_cb"][feats["template_force"]].clone()  # (n_ref, n)
         atom_indices = torch.arange(
             feats["atom_pad_mask"].shape[1],
             device=feats["atom_pad_mask"].device,
             dtype=torch.float32,
-        )[None, :, None]
+        )[None, :, None]  # (1, a, 1)
         ref_atom_index = (
             torch.bmm(
                 feats["token_to_rep_atom"].float(),
                 atom_indices,
             )
             .squeeze(-1)
-            .long()[0]
+            .long()[0]  # (n,)
         )
         ref_token_index = (
             torch.bmm(
@@ -687,24 +734,25 @@ class TemplateReferencePotential(FlatBottomPotential, ReferencePotential):
                 feats["token_index"].unsqueeze(-1).float(),
             )
             .squeeze(-1)
-            .long()[0]
+            .long()[0]  # (a,)
         )
 
         index = torch.arange(
             template_mask.shape[-1],
             dtype=torch.long,
             device=template_mask.device,
-        )[None]
+        )[None]  # (1, n)
         upper = torch.full(
             template_mask.shape,
             float("inf"),
             device=index.device,
             dtype=torch.float32,
-        )
-        reference_indices = torch.argwhere(template_mask).T
+        )  # (n_ref, n)
+        reference_indices = torch.argwhere(template_mask).T  # (2, n_active)
         upper[reference_indices.unbind()] = feats["template_force_threshold"][
             feats["template_force"]
-        ][reference_indices[0]]
+        ][reference_indices[0]]  # (n_ref, n)
+        # Returns index (1, n), function args (n_ref, n), and reference mappings.
         return (
             index,
             (torch.ones_like(upper), None, upper),
@@ -719,8 +767,10 @@ class ContactPotentital(FlatBottomPotential, DistancePotential):
 
     def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
         del parameters
-        index = feats["contact_pair_index"][0]
-        upper = feats["contact_thresholds"][0].clone()
+        # pair index: (1, 2, n); threshold/operator features: (1, n).
+        index = feats["contact_pair_index"][0]  # (2, n)
+        upper = feats["contact_thresholds"][0].clone()  # (n,)
+        # Returns index (2, n), k/upper (n,), and two operator tensors (n,).
         return (
             index,
             (torch.ones_like(upper), None, upper),

fastplms/models/boltz/vb_tri_attn_attention.py

@@ -2,9 +2,8 @@
 
 from __future__ import annotations
 
-from typing import cast
-
 import torch
+from typing import cast
 from torch import Tensor, nn
 
 from .vb_tri_attn_primitives import Attention, LayerNorm, Linear
@@ -44,6 +43,7 @@ class TriangleAttention(nn.Module):
     ) -> Tensor:
         """Evaluate attention in slices of the batch-like leading axes."""
 
+        # x: (..., l, l, d); biases/mask retain broadcastable triangular axes.
         def attention_call(**inputs: Tensor) -> Tensor:
             return cast(Tensor, self.mha(**inputs, use_kernels=use_kernels))
 
@@ -53,7 +53,7 @@ class TriangleAttention(nn.Module):
             "tri_bias": tri_bias,
             "mask_bias": mask_bias,
             "mask": mask,
-        }
+        }  # tensor values retain the input shapes above
         return cast(
             Tensor,
             chunk_layer(
@@ -74,6 +74,7 @@ class TriangleAttention(nn.Module):
         chunk_size: int | None,
         use_kernels: bool,
     ) -> Tensor:
+        # states: (..., l, l, d); returned tensor has the same shape.
         if chunk_size is not None and not use_kernels:
             return cast(
                 Tensor,
@@ -107,17 +108,19 @@ class TriangleAttention(nn.Module):
     ) -> Tensor:
         """Transform X with shape ``(..., l, l, d)`` under mask M."""
 
-        pair_mask = x.new_ones(x.shape[:-1]) if mask is None else mask
-        pair_states = x
+        pair_mask = x.new_ones(x.shape[:-1]) if mask is None else mask  # (..., l, l)
+        pair_states = x  # (..., l, l, d)
         if not self.starting:
-            pair_states = pair_states.transpose(-2, -3)
-            pair_mask = pair_mask.transpose(-1, -2)
-
-        normalized = self.layer_norm(pair_states)
-        expanded_mask = pair_mask[..., :, None, None, :]
-        mask_bias = self.inf * (expanded_mask - 1)
-        triangle_bias = permute_final_dims(self.linear(normalized), (2, 0, 1))
-        triangle_bias = triangle_bias.unsqueeze(-4)
+            pair_states = pair_states.transpose(-2, -3)  # (..., l, l, d)
+            pair_mask = pair_mask.transpose(-1, -2)  # (..., l, l)
+
+        normalized = self.layer_norm(pair_states)  # (..., l, l, d)
+        expanded_mask = pair_mask[..., :, None, None, :]  # (..., l, 1, 1, l)
+        mask_bias = self.inf * (expanded_mask - 1)  # (..., l, 1, 1, l)
+        triangle_bias = permute_final_dims(
+            self.linear(normalized), (2, 0, 1)
+        )  # (..., h, l, l)
+        triangle_bias = triangle_bias.unsqueeze(-4)  # (..., 1, h, l, l)
         output = self._run_attention(
             normalized,
             triangle_bias,
@@ -125,8 +128,10 @@ class TriangleAttention(nn.Module):
             expanded_mask,
             chunk_size,
             use_kernels,
-        )
-        return output if self.starting else output.transpose(-2, -3)
+        )  # (..., l, l, d)
+        return (
+            output if self.starting else output.transpose(-2, -3)
+        )  # (..., l, l, d)
 
 
 TriangleAttentionStartingNode = TriangleAttention

fastplms/models/boltz/vb_tri_attn_primitives.py

@@ -12,10 +12,9 @@ from __future__ import annotations
 
 import importlib
 import math
+import torch
 from collections.abc import Callable, Sequence
 from importlib.util import find_spec
-
-import torch
 from torch import nn
 
 from . import vb_layers_initialize as initialize
@@ -27,6 +26,7 @@ def _initialize_weight(
     bias: torch.Tensor | None,
     method: str,
 ) -> None:
+    # weight: (d_out, d_in); bias: (d_out,) or None; both retain shape in place.
     initializers = {
         "default": initialize.lecun_normal_init_,
         "relu": initialize.he_normal_init_,
@@ -68,21 +68,28 @@ class Linear(nn.Linear):
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         """Project ``input`` while preserving explicit precision policies."""
 
+        # input: (..., d_in); every branch returns (..., d_out).
         input_dtype = input.dtype
         if self.precision is not None:
             with torch.autocast("cuda", enabled=False):
-                bias = None if self.bias is None else self.bias.to(self.precision)
+                bias = (
+                    None if self.bias is None else self.bias.to(self.precision)
+                )  # (d_out,) or None
                 projected = nn.functional.linear(
                     input.to(self.precision),
                     self.weight.to(self.precision),
                     bias,
-                )
-            return projected.to(input_dtype)
+                )  # (..., d_out)
+            return projected.to(input_dtype)  # (..., d_out)
         if input_dtype is torch.bfloat16:
             with torch.autocast("cuda", enabled=False):
-                bias = None if self.bias is None else self.bias.to(input_dtype)
-                return nn.functional.linear(input, self.weight.to(input_dtype), bias)
-        return nn.functional.linear(input, self.weight, self.bias)
+                bias = (
+                    None if self.bias is None else self.bias.to(input_dtype)
+                )  # (d_out,) or None
+                return nn.functional.linear(
+                    input, self.weight.to(input_dtype), bias
+                )  # (..., d_out)
+        return nn.functional.linear(input, self.weight, self.bias)  # (..., d_out)
 
 
 class LayerNorm(nn.Module):
@@ -92,10 +99,11 @@ class LayerNorm(nn.Module):
         super().__init__()
         self.c_in = (c_in,)
         self.eps = eps
-        self.weight = nn.Parameter(torch.ones(c_in))
-        self.bias = nn.Parameter(torch.zeros(c_in))
+        self.weight = nn.Parameter(torch.ones(c_in))  # (c_in,)
+        self.bias = nn.Parameter(torch.zeros(c_in))  # (c_in,)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (..., c_in); normalization preserves shape.
         if x.dtype is torch.bfloat16:
             with torch.autocast("cuda", enabled=False):
                 return nn.functional.layer_norm(
@@ -126,11 +134,14 @@ def _attention(
 ) -> torch.Tensor:
     """Compute biased scaled attention after query scaling."""
 
-    scores = torch.matmul(query, permute_final_dims(key, (1, 0)))
+    # query/key/value: (..., h, n_q_or_k, d_h).
+    scores = torch.matmul(
+        query, permute_final_dims(key, (1, 0))
+    )  # (..., h, n_q, n_k)
     for bias in biases:
-        scores += bias
-    probabilities = softmax_no_cast(scores, dim=-1)
-    return torch.matmul(probabilities, value)
+        scores += bias  # (..., h, n_q, n_k)
+    probabilities = softmax_no_cast(scores, dim=-1)  # (..., h, n_q, n_k)
+    return torch.matmul(probabilities, value)  # (..., h, n_q, d_h)
 
 
 @torch.compiler.disable
@@ -144,6 +155,7 @@ def kernel_triangular_attn(
 ) -> torch.Tensor:
     """Call the optional cuEquivariance triangle-attention primitive."""
 
+    # q/k/v: (..., h, n, d_h); returned tensor uses the kernel's matching layout.
     if (
         find_spec("cuequivariance_torch") is None
         or find_spec("cuequivariance_ops_torch") is None
@@ -193,23 +205,32 @@ class Attention(nn.Module):
         """Project Q, K, and V to ``(..., h, n, d)`` tensors."""
 
         def split_heads(X: torch.Tensor) -> torch.Tensor:
-            return X.view(*X.shape[:-1], self.no_heads, -1).transpose(-2, -3)
-
-        q = split_heads(self.linear_q(q_x))
-        k = split_heads(self.linear_k(kv_x))
-        v = split_heads(self.linear_v(kv_x))
+            # X: (..., n, h * d_h)
+            return X.view(*X.shape[:-1], self.no_heads, -1).transpose(
+                -2, -3
+            )  # (..., h, n, d_h)
+
+        # q_x: (..., n_q, c_q); kv_x: (..., n_k, c_k).
+        q = split_heads(self.linear_q(q_x))  # (..., h, n_q, d_h)
+        k = split_heads(self.linear_k(kv_x))  # (..., h, n_k, d_h)
+        v = split_heads(self.linear_v(kv_x))  # (..., h, n_k, d_h)
         if apply_scale:
-            q /= math.sqrt(self.c_hidden)
-        return q, k, v
+            q /= math.sqrt(self.c_hidden)  # (..., h, n_q, d_h)
+        return q, k, v  # (..., h, n_q, d_h), two (..., h, n_k, d_h)
 
     def _wrap_up(self, output: torch.Tensor, q_x: torch.Tensor) -> torch.Tensor:
         """Apply the optional gate, merge heads, and project the update."""
 
+        # output: (..., n_q, h, d_h); q_x: (..., n_q, c_q).
         if self.linear_g is not None:
-            gate = self.sigmoid(self.linear_g(q_x))
-            gate = gate.view(*gate.shape[:-1], self.no_heads, -1)
-            output = output * gate
-        return self.linear_o(flatten_final_dims(output, 2))
+            gate = self.sigmoid(self.linear_g(q_x))  # (..., n_q, h * d_h)
+            gate = gate.view(
+                *gate.shape[:-1], self.no_heads, -1
+            )  # (..., n_q, h, d_h)
+            output = output * gate  # (..., n_q, h, d_h)
+        return self.linear_o(
+            flatten_final_dims(output, 2)
+        )  # (..., n_q, c_q)
 
     def forward(
         self,
@@ -222,7 +243,9 @@ class Attention(nn.Module):
     ) -> torch.Tensor:
         """Return a gated pair update with shape ``(..., n, c_q)``."""
 
-        q, k, v = self._prep_qkv(q_x, kv_x, apply_scale=not use_kernels)
+        q, k, v = self._prep_qkv(
+            q_x, kv_x, apply_scale=not use_kernels
+        )  # q: (..., h, n_q, d_h); k/v: (..., h, n_k, d_h)
         if use_kernels:
             output = kernel_triangular_attn(
                 q,
@@ -231,7 +254,12 @@ class Attention(nn.Module):
                 tri_bias=tri_bias,
                 mask=mask.bool(),
                 scale=1.0 / math.sqrt(self.c_hidden),
-            )
+            )  # (..., h, n_q, d_h)
         else:
-            output = _attention(q, k, v, (mask_bias, tri_bias))
-        return self._wrap_up(output.transpose(-2, -3), q_x)
+            output = _attention(
+                q, k, v, (mask_bias, tri_bias)
+            )  # (..., h, n_q, d_h)
+        return self._wrap_up(
+            output.transpose(-2, -3), q_x
+        )  # (..., n_q, c_q)
+    # tensor: (...); softmax preserves shape.

fastplms/models/boltz/vb_tri_attn_utils.py

@@ -7,21 +7,21 @@ the public function names expected by converted checkpoints.
 
 from __future__ import annotations
 
+import torch
 from collections.abc import Callable, Sequence
 from functools import partial
 from math import prod
 from typing import Any
 
-import torch
-
 
 def add(left: torch.Tensor, right: torch.Tensor, inplace: bool) -> torch.Tensor:
     """Add ``right`` to ``left``, optionally reusing ``left`` storage."""
 
+    # left/right: broadcast-compatible shapes; result has their broadcast shape.
     if inplace:
-        left += right
-        return left
-    return left + right
+        left += right  # same shape as left
+        return left  # same shape as left
+    return left + right  # broadcast shape
 
 
 def permute_final_dims(tensor: torch.Tensor, inds: Sequence[int]) -> torch.Tensor:
@@ -30,7 +30,7 @@ def permute_final_dims(tensor: torch.Tensor, inds: Sequence[int]) -> torch.Tenso
     final_count = len(inds)
     leading = list(range(tensor.ndim - final_count))
     final = [tensor.ndim - final_count + index for index in inds]
-    return tensor.permute(leading + final)
+    return tensor.permute(leading + final)  # (..., d_0, ..., d_n)
 
 
 def is_fp16_enabled() -> bool:
@@ -73,7 +73,8 @@ tensor_tree_map = partial(tree_map, leaf_type=torch.Tensor)
 def flatten_final_dims(tensor: torch.Tensor, no_dims: int) -> torch.Tensor:
     """Collapse the final ``no_dims`` axes into one axis."""
 
-    return tensor.reshape(*tensor.shape[:-no_dims], -1)
+    # tensor: (..., d_0, ..., d_n)
+    return tensor.reshape(*tensor.shape[:-no_dims], -1)  # (..., prod(final dims))
 
 
 def _fetch_dims(tree: Any) -> list[torch.Size]:
@@ -194,9 +195,13 @@ def _chunk_slice(
     batch_shape = tuple(tensor.shape[:no_batch_dims])
     start = _flat_idx_to_idx(flat_start, batch_shape)
     end = _flat_idx_to_idx(flat_end - 1, batch_shape)
-    pieces = [tensor[item] for item in _get_minimal_slice_set(start, end, batch_shape)]
+    pieces = [
+        tensor[item] for item in _get_minimal_slice_set(start, end, batch_shape)
+    ]  # each: (*covered_batch_shape, *feature_shape)
     feature_shape = tuple(tensor.shape[no_batch_dims:])
-    return torch.cat([piece.reshape(-1, *feature_shape) for piece in pieces])
+    return torch.cat(
+        [piece.reshape(-1, *feature_shape) for piece in pieces]
+    )  # (flat_end - flat_start, *feature_shape)
 
 
 def _prepare_input(
@@ -206,12 +211,15 @@ def _prepare_input(
     *,
     low_mem: bool,
 ) -> torch.Tensor:
+    # tensor: (*input_batch_shape, *feature_shape)
     feature_shape = tuple(tensor.shape[no_batch_dims:])
     if low_mem:
-        return tensor.expand(batch_shape + feature_shape)
+        return tensor.expand(batch_shape + feature_shape)  # (*batch_shape, *feature_shape)
     if any(size != 1 for size in tensor.shape[:no_batch_dims]):
-        tensor = tensor.expand(batch_shape + feature_shape)
-    return tensor.reshape(-1, *feature_shape)
+        tensor = tensor.expand(
+            batch_shape + feature_shape
+        )  # (*batch_shape, *feature_shape)
+    return tensor.reshape(-1, *feature_shape)  # (flat_batch, *feature_shape)
 
 
 def _write_chunk(
@@ -245,9 +253,9 @@ def _write_chunk(
     if not isinstance(source, torch.Tensor):
         raise ValueError(f"output type {type(source)!r} is not supported")
     if add_into_out:
-        destination[start:stop] += source
+        destination[start:stop] += source  # (stop - start, *feature_shape)
     else:
-        destination[start:stop] = source
+        destination[start:stop] = source  # (stop - start, *feature_shape)
 
 
 def chunk_layer(
@@ -274,14 +282,16 @@ def chunk_layer(
         no_batch_dims=no_batch_dims,
         low_mem=low_mem,
     )
-    prepared_inputs = tensor_tree_map(prepare, inputs)
+    prepared_inputs = tensor_tree_map(
+        prepare, inputs
+    )  # each tensor: (flat_batch, *feature_shape), or broadcast batch in low-memory mode
 
     output = None
     if _out is not None:
         output = tensor_tree_map(
             lambda tensor: tensor.reshape(-1, *tensor.shape[no_batch_dims:]),
             _out,
-        )
+        )  # each tensor: (flat_batch, *feature_shape)
 
     flat_batch_size = prod(batch_shape)
     for start in range(0, flat_batch_size, chunk_size):
@@ -302,14 +312,18 @@ def chunk_layer(
             ) -> torch.Tensor:
                 return tensor if tensor.shape[0] == 1 else tensor[start:stop]
 
-        input_chunk = tensor_tree_map(select, prepared_inputs)
-        output_chunk = layer(**input_chunk)
+        input_chunk = tensor_tree_map(
+            select, prepared_inputs
+        )  # each tensor: (chunk, *feature_shape)
+        output_chunk = layer(
+            **input_chunk
+        )  # each output tensor: (chunk, *output_feature_shape)
 
         if output is None:
             output = tensor_tree_map(
                 lambda tensor: tensor.new_zeros((flat_batch_size, *tensor.shape[1:])),
                 output_chunk,
-            )
+            )  # each tensor: (flat_batch, *output_feature_shape)
         _write_chunk(
             output,
             output_chunk,
@@ -321,4 +335,6 @@ def chunk_layer(
     return tensor_tree_map(
         lambda tensor: tensor.reshape(batch_shape + tuple(tensor.shape[1:])),
         output,
-    )
+    )  # each tensor: (*batch_shape, *output_feature_shape)
+    # tensor: (..., d_0, ..., d_n); only the named final dimensions are reordered.
+    # tensor: (*batch_shape, *feature_shape)

fastplms/registry.py

@@ -18,6 +18,7 @@ from types import MappingProxyType
 from typing import Any, Literal, cast
 from urllib.parse import urlparse
 
+
 _HEX_RE = re.compile(r"^[0-9a-f]+$")
 _IDENTIFIER_RE = re.compile(r"^[a-z0-9][a-z0-9_-]*$")
 _HUB_LICENSE_NAME_RE = re.compile(r"[^a-z0-9.]+")
@@ -559,20 +560,20 @@ def _require_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple
     value = table.get(key)
     if not isinstance(value, list) or not value or any(not isinstance(item, str) for item in value):
         raise RegistryError(f"{context}.{key} must be a non-empty string array.")
-    result = tuple(value)
-    if len(set(result)) != len(result):
+    strings = tuple(value)
+    if len(set(strings)) != len(strings):
         raise RegistryError(f"{context}.{key} contains duplicate values.")
-    return result
+    return strings
 
 
 def _optional_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]:
     value = table.get(key, [])
     if not isinstance(value, list) or any(not isinstance(item, str) for item in value):
         raise RegistryError(f"{context}.{key} must be a string array.")
-    result = tuple(value)
-    if len(set(result)) != len(result):
+    strings = tuple(value)
+    if len(set(strings)) != len(strings):
         raise RegistryError(f"{context}.{key} contains duplicate values.")
-    return result
+    return strings
 
 
 def _optional_str(table: Mapping[str, Any], key: str, context: str) -> str | None:
@@ -657,11 +658,11 @@ def _require_digest_list(
     table: Mapping[str, Any], key: str, context: str
 ) -> tuple[FileDigest, ...]:
     encoded = _require_str_list(table, key, context)
-    result = tuple(FileDigest.parse(value) for value in encoded)
-    paths = [item.path for item in result]
+    digests = tuple(FileDigest.parse(value) for value in encoded)
+    paths = [item.path for item in digests]
     if len(paths) != len(set(paths)):
         raise RegistryError(f"{context}.{key} contains duplicate paths.")
-    return result
+    return digests
 
 
 def _validate_revision(revision: str, context: str) -> None:
@@ -701,7 +702,7 @@ def _parse_oracle_assets(table: Mapping[str, Any], context: str) -> tuple[Oracle
     raw = table.get("oracle_assets", [])
     if not isinstance(raw, list):
         raise RegistryError(f"{context}.oracle_assets must be an array of tables.")
-    result: list[OracleAsset] = []
+    assets: list[OracleAsset] = []
     for index, value in enumerate(raw):
         asset_context = f"{context}.oracle_assets[{index}]"
         if not isinstance(value, dict):
@@ -736,7 +737,7 @@ def _parse_oracle_assets(table: Mapping[str, Any], context: str) -> tuple[Oracle
         size = value.get("size")
         if isinstance(size, bool) or not isinstance(size, int) or size <= 0:
             raise RegistryError(f"{asset_context}.size must be a positive byte count.")
-        result.append(
+        assets.append(
             OracleAsset(
                 role=role,
                 path=path,
@@ -745,16 +746,16 @@ def _parse_oracle_assets(table: Mapping[str, Any], context: str) -> tuple[Oracle
                 size=size,
             )
         )
-    roles = [asset.role for asset in result]
-    paths = [asset.path for asset in result]
-    urls = [asset.url for asset in result]
+    roles = [asset.role for asset in assets]
+    paths = [asset.path for asset in assets]
+    urls = [asset.url for asset in assets]
     if (
         len(roles) != len(set(roles))
         or len(paths) != len(set(paths))
         or len(urls) != len(set(urls))
     ):
         raise RegistryError(f"{context}.oracle_assets contains duplicate identities.")
-    return tuple(result)
+    return tuple(assets)
 
 
 def _parse_official_golden(
@@ -789,7 +790,7 @@ def _parse_official_golden(
 def _parse_attention_kernels(raw: object) -> dict[str, AttentionKernelSpec]:
     if not isinstance(raw, list) or not raw:
         raise RegistryError("The manifest must contain [[attention_kernels]] entries.")
-    result: dict[str, AttentionKernelSpec] = {}
+    kernels: dict[str, AttentionKernelSpec] = {}
     expected_variants = {
         "flash_attention_2": "flash_attn2",
         "flash_attention_3": "flash_attn3",
@@ -812,7 +813,7 @@ def _parse_attention_kernels(raw: object) -> dict[str, AttentionKernelSpec]:
         implementation = _require_str(value, "implementation", context)
         if implementation not in expected_variants:
             raise RegistryError(f"Unsupported attention kernel {implementation!r}.")
-        if implementation in result:
+        if implementation in kernels:
             raise RegistryError(f"Duplicate attention kernel {implementation!r}.")
         repository = _require_str(value, "repository", context)
         if _REPOSITORY_ID_RE.fullmatch(repository) is None:
@@ -834,7 +835,7 @@ def _parse_attention_kernels(raw: object) -> dict[str, AttentionKernelSpec]:
         dtypes = _require_str_list(value, "dtypes", context)
         if not set(dtypes).issubset(_ALLOWED_DTYPES):
             raise RegistryError(f"{context}.dtypes contains unsupported dtypes.")
-        result[implementation] = AttentionKernelSpec(
+        kernels[implementation] = AttentionKernelSpec(
             implementation=implementation,
             repository=repository,
             revision=revision,
@@ -842,15 +843,15 @@ def _parse_attention_kernels(raw: object) -> dict[str, AttentionKernelSpec]:
             expected_variant=expected_variant,
             dtypes=cast(tuple[DtypeName, ...], dtypes),
         )
-    if set(result) != set(expected_variants):
+    if set(kernels) != set(expected_variants):
         raise RegistryError("The manifest must pin both FlashAttention kernel versions.")
-    return result
+    return kernels
 
 
 def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
     if not isinstance(raw, list) or not raw:
         raise RegistryError("The manifest must contain at least one [[upstreams]] entry.")
-    result: dict[str, UpstreamSource] = {}
+    upstreams: dict[str, UpstreamSource] = {}
     paths: set[str] = set()
     for index, value in enumerate(raw):
         context = f"upstreams[{index}]"
@@ -860,7 +861,7 @@ def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
         source_id = _require_str(value, "id", context)
         if _IDENTIFIER_RE.fullmatch(source_id) is None:
             raise RegistryError(f"Invalid upstream ID: {source_id!r}")
-        if source_id in result:
+        if source_id in upstreams:
             raise RegistryError(f"Duplicate upstream ID: {source_id!r}")
         revision = _require_str(value, "revision", context)
         _validate_revision(revision, f"{context}.revision")
@@ -913,7 +914,7 @@ def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
             missing_e1 = sorted(required_e1.difference(distribution_map))
             if missing_e1:
                 raise RegistryError(f"{context} is missing E1 legal files: {missing_e1}")
-        result[source_id] = UpstreamSource(
+        upstreams[source_id] = UpstreamSource(
             id=source_id,
             path=path,
             url=url,
@@ -923,7 +924,7 @@ def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
             license_digests=license_digests,
             distribution_files=distribution_files,
         )
-    return result
+    return upstreams
 
 
 def _parse_families(
@@ -932,7 +933,7 @@ def _parse_families(
 ) -> dict[str, ModelFamily]:
     if not isinstance(raw, dict) or not raw:
         raise RegistryError("The manifest must contain [families.<id>] tables.")
-    result: dict[str, ModelFamily] = {}
+    families: dict[str, ModelFamily] = {}
     for family_id, value in raw.items():
         context = f"families.{family_id}"
         if _IDENTIFIER_RE.fullmatch(family_id) is None or not isinstance(value, dict):
@@ -1068,7 +1069,7 @@ def _parse_families(
                 f"{context}.conversion_provenance must identify {state_transform!r} and "
                 f"contain mechanism-first sections; missing {missing_sections}."
             )
-        result[family_id] = ModelFamily(
+        families[family_id] = ModelFamily(
             id=family_id,
             architecture=_require_str(value, "architecture", context),
             upstreams=source_ids,
@@ -1099,7 +1100,7 @@ def _parse_families(
             conversion_provenance=conversion_provenance,
             backbone_model=backbone_model,
         )
-    return result
+    return families
 
 
 def _parse_runtime_assets(
@@ -1108,7 +1109,7 @@ def _parse_runtime_assets(
 ) -> dict[str, RuntimeAsset]:
     if not isinstance(raw, list) or not raw:
         raise RegistryError("The manifest must contain at least one [[runtime_assets]] entry.")
-    result: dict[str, RuntimeAsset] = {}
+    runtime_assets: dict[str, RuntimeAsset] = {}
     identities: set[tuple[str, str, str]] = set()
     for index, value in enumerate(raw):
         context = f"runtime_assets[{index}]"
@@ -1118,7 +1119,7 @@ def _parse_runtime_assets(
         asset_id = _require_str(value, "id", context)
         if _IDENTIFIER_RE.fullmatch(asset_id) is None:
             raise RegistryError(f"Invalid runtime asset ID: {asset_id!r}")
-        if asset_id in result:
+        if asset_id in runtime_assets:
             raise RegistryError(f"Duplicate runtime asset ID: {asset_id!r}")
         repository = _require_str(value, "repository", context)
         if _REPOSITORY_ID_RE.fullmatch(repository) is None:
@@ -1164,7 +1165,7 @@ def _parse_runtime_assets(
         if identity in identities:
             raise RegistryError(f"Duplicate runtime asset identity: {identity!r}")
         identities.add(identity)
-        result[asset_id] = RuntimeAsset(
+        runtime_assets[asset_id] = RuntimeAsset(
             id=asset_id,
             repository=repository,
             revision=revision,
@@ -1176,7 +1177,7 @@ def _parse_runtime_assets(
             license_expression=license_expression,
             offline_behavior=offline_behavior,
         )
-    return result
+    return runtime_assets
 
 
 def _parse_models(
@@ -1185,7 +1186,7 @@ def _parse_models(
 ) -> dict[str, ModelSpec]:
     if not isinstance(raw, list) or not raw:
         raise RegistryError("The manifest must contain at least one [[models]] entry.")
-    result: dict[str, ModelSpec] = {}
+    models: dict[str, ModelSpec] = {}
     fast_repositories: set[str] = set()
     for index, value in enumerate(raw):
         context = f"models[{index}]"
@@ -1195,7 +1196,7 @@ def _parse_models(
         model_id = _require_str(value, "id", context)
         if _IDENTIFIER_RE.fullmatch(model_id) is None:
             raise RegistryError(f"Invalid model ID: {model_id!r}")
-        if model_id in result:
+        if model_id in models:
             raise RegistryError(f"Duplicate model ID: {model_id!r}")
         family_id = _require_str(value, "family", context)
         if family_id not in families:
@@ -1278,7 +1279,7 @@ def _parse_models(
                 if not class_path.startswith("fastplms.") or class_path.count(".") < 2:
                     raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}")
                 auto_map.append((auto_class, class_path))
-        result[model_id] = ModelSpec(
+        models[model_id] = ModelSpec(
             id=model_id,
             family=family,
             fast=fast,
@@ -1294,7 +1295,7 @@ def _parse_models(
             notes=notes,
             msa_conditioning=msa_conditioning,
         )
-    return result
+    return models
 
 
 def _validate_registry(
@@ -1409,21 +1410,21 @@ def _validate_registry(
 
 def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry:
     try:
-        data = tomllib.loads(raw_bytes.decode("utf-8"))
+        manifest = tomllib.loads(raw_bytes.decode("utf-8"))
     except (UnicodeDecodeError, tomllib.TOMLDecodeError) as error:
         raise RegistryError(f"Unable to parse model manifest: {error}") from error
-    _reject_unknown_fields(data, _ROOT_FIELDS, "manifest")
-    if data.get("schema_version") != 1:
+    _reject_unknown_fields(manifest, _ROOT_FIELDS, "manifest")
+    if manifest.get("schema_version") != 1:
         raise RegistryError("Unsupported model manifest schema_version; expected 1.")
-    legal_files = _require_digest_list(data, "legal_files", "manifest")
+    legal_files = _require_digest_list(manifest, "legal_files", "manifest")
     required_legal_paths = {"LICENSE", "THIRD_PARTY_NOTICES.md"}
     if {item.path for item in legal_files} != required_legal_paths:
         raise RegistryError("manifest.legal_files must contain LICENSE and THIRD_PARTY_NOTICES.md.")
-    attention_kernels = _parse_attention_kernels(data.get("attention_kernels"))
-    upstreams = _parse_upstreams(data.get("upstreams"))
-    families = _parse_families(data.get("families"), upstreams)
-    runtime_assets = _parse_runtime_assets(data.get("runtime_assets"), families)
-    models = _parse_models(data.get("models"), families)
+    attention_kernels = _parse_attention_kernels(manifest.get("attention_kernels"))
+    upstreams = _parse_upstreams(manifest.get("upstreams"))
+    families = _parse_families(manifest.get("families"), upstreams)
+    runtime_assets = _parse_runtime_assets(manifest.get("runtime_assets"), families)
+    models = _parse_models(manifest.get("models"), families)
     _validate_registry(upstreams, attention_kernels, families, models)
     return ModelRegistry(
         schema_version=1,

fastplms/runtime.py

@@ -11,6 +11,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Literal
 
+
 if TYPE_CHECKING:
     from collections.abc import Iterator
 

modeling_fastplms.py

@@ -1,4 +1,4 @@
-"""Generated bridge to the unchanged FastPLMs package sources."""
+"""Generated bridge to the embedded FastPLMs runtime sources."""
 
 import base64
 import hashlib
@@ -6,14 +6,13 @@ import importlib
 import importlib.util
 import sys
 import tempfile
-from importlib.metadata import PackageNotFoundError, distribution
 from io import BytesIO
 from pathlib import Path
 from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "e6e0f8c8632238f40a5112f22397a9a211eb02d3415c1c93d43305aef51f7728":
+if RUNTIME_HASH != "3aebb85d11edb01329d2aafb4c8260bd67339785f0f252582246ae7c33206143":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []
@@ -83,24 +82,6 @@ def _runtime_file_hashes(package_root):
         result[relative.as_posix()] = hashlib.sha256(path.read_bytes()).hexdigest()
     return result
 
-def _installed_runtime_digest(installed_root, relative):
-    candidate = installed_root / relative
-    if candidate.is_file():
-        return hashlib.sha256(candidate.read_bytes()).hexdigest()
-    if relative != "kernels.lock":
-        return None
-    try:
-        installed_distribution = distribution("fastplms")
-    except PackageNotFoundError:
-        return None
-    for entry in installed_distribution.files or ():
-        normalized = str(entry).replace("\\", "/")
-        if normalized.endswith(".dist-info/kernels.lock"):
-            lock_path = Path(installed_distribution.locate_file(entry))
-            if lock_path.is_file():
-                return hashlib.sha256(lock_path.read_bytes()).hexdigest()
-    return None
-
 def _extend_loaded_package_paths(package_root):
     for name, module in list(sys.modules.items()):
         if name != "fastplms" and not name.startswith("fastplms."):
@@ -114,29 +95,14 @@ def _extend_loaded_package_paths(package_root):
         if candidate.is_dir() and candidate_text not in paths:
             paths.append(candidate_text)
 
-def _merge_runtime(installed, package_root):
+def _merge_runtime(package, package_root):
     incoming = _runtime_file_hashes(package_root)
-    known = dict(getattr(installed, "__fastplms_artifact_runtime_files__", {}))
-    installed_root_text = getattr(
-        installed, "__fastplms_artifact_installed_root__", None
-    )
-    if not known:
-        installed_file = getattr(installed, "__file__", None)
-        if installed_file is None:
-            raise RuntimeError(
-                "The loaded fastplms package has no source path and cannot be verified "
-                "against the embedded artifact runtime."
-            )
-        installed_root = Path(installed_file).resolve().parent
-        for relative, digest in incoming.items():
-            if _installed_runtime_digest(installed_root, relative) != digest:
-                raise RuntimeError(
-                    "The installed FastPLMs runtime differs from this artifact at "
-                    f"{relative!r}. Install the artifact's matching FastPLMs release "
-                    "or use a separate Python process."
-                )
-        installed_root_text = str(installed_root)
-        installed.__fastplms_artifact_installed_root__ = installed_root_text
+    known = getattr(package, "__fastplms_artifact_runtime_files__", None)
+    if not isinstance(known, dict):
+        raise RuntimeError(
+            "A non-artifact fastplms module is already loaded. Load the Hub artifact "
+            "in a separate Python process."
+        )
     conflicts = sorted(
         relative
         for relative, digest in incoming.items()
@@ -148,35 +114,25 @@ def _merge_runtime(installed, package_root):
             + ", ".join(repr(path) for path in conflicts[:5])
             + ". Load incompatible releases in separate Python processes."
         )
-    if installed_root_text is not None:
-        installed_root = Path(installed_root_text)
-        for relative, digest in incoming.items():
-            if relative in known:
-                continue
-            if _installed_runtime_digest(installed_root, relative) != digest:
-                raise RuntimeError(
-                    "The installed FastPLMs runtime differs from this artifact at "
-                    f"{relative!r}. Install the artifact's matching FastPLMs release "
-                    "or use a separate Python process."
-                )
+    known = dict(known)
     known.update(incoming)
-    installed.__fastplms_artifact_runtime_files__ = known
-    roots = list(getattr(installed, "__fastplms_artifact_runtime_roots__", ()))
+    package.__fastplms_artifact_runtime_files__ = known
+    roots = list(getattr(package, "__fastplms_artifact_runtime_roots__", ()))
     if str(package_root) not in roots:
         roots.append(str(package_root))
-    installed.__fastplms_artifact_runtime_roots__ = tuple(roots)
+    package.__fastplms_artifact_runtime_roots__ = tuple(roots)
     temporaries = list(
-        getattr(installed, "__fastplms_artifact_runtime_temporaries__", ())
+        getattr(package, "__fastplms_artifact_runtime_temporaries__", ())
     )
     for temporary in _RUNTIME_TEMPORARIES:
         if temporary not in temporaries:
             temporaries.append(temporary)
-    installed.__fastplms_artifact_runtime_temporaries__ = tuple(temporaries)
-    hashes = set(getattr(installed, "__fastplms_artifact_runtime_hashes__", ()))
+    package.__fastplms_artifact_runtime_temporaries__ = tuple(temporaries)
+    hashes = set(getattr(package, "__fastplms_artifact_runtime_hashes__", ()))
     hashes.add(RUNTIME_HASH)
-    installed.__fastplms_artifact_runtime_hashes__ = frozenset(hashes)
+    package.__fastplms_artifact_runtime_hashes__ = frozenset(hashes)
     _extend_loaded_package_paths(package_root)
-    return installed
+    return package
 
 def _import_without_bytecode(module_name):
     previous = sys.dont_write_bytecode
@@ -187,13 +143,13 @@ def _import_without_bytecode(module_name):
         sys.dont_write_bytecode = previous
 
 def _install_runtime():
-    installed = sys.modules.get("fastplms")
-    hashes = getattr(installed, "__fastplms_artifact_runtime_hashes__", ())
+    package = sys.modules.get("fastplms")
+    hashes = getattr(package, "__fastplms_artifact_runtime_hashes__", ())
     if RUNTIME_HASH in hashes:
-        return installed
+        return package
     package_root = _ensure_runtime()
-    if installed is not None:
-        return _merge_runtime(installed, package_root)
+    if package is not None:
+        return _merge_runtime(package, package_root)
     spec = importlib.util.spec_from_file_location(
         "fastplms",
         package_root / "__init__.py",
@@ -223,8 +179,8 @@ def _install_runtime():
     return package
 
 _install_runtime()
-_module_225 = _import_without_bytecode("fastplms.models.boltz.modeling_boltz2")
-Boltz2Config = _module_225.Boltz2Config
+_module_181 = _import_without_bytecode("fastplms.models.boltz.modeling_boltz2")
+Boltz2Config = _module_181.Boltz2Config
 Boltz2Config.__module__ = __name__
-Boltz2Model = _module_225.Boltz2Model
+Boltz2Model = _module_181.Boltz2Model
 Boltz2Model.__module__ = __name__

requirements.txt

@@ -0,0 +1,20 @@
+# Direct runtime dependencies for Synthyra/Boltz2.
+# FastPLMs source is embedded in this model repository.
+torch>=2.13,<2.14
+transformers>=5.13,<5.14
+huggingface-hub>=0.34,<2
+tokenizers>=0.22,<0.23
+safetensors>=0.5,<1
+numpy>=1.26,<3
+einops>=0.8,<1
+tqdm>=4.67,<5
+accelerate>=1.10,<2
+biopython>=1.85,<2
+biotite>=1.4,<2
+brotli>=1.1,<2
+msgpack>=1.1,<2
+msgpack-numpy>=0.4.8,<1
+omegaconf>=2.3,<3
+rdkit>=2025.9,<2027
+scipy>=1.15,<2
+zstandard>=0.23,<1

runtime-attestation.json

@@ -2,56 +2,57 @@
   "files": {
     "LICENSES/FastPLMs-Apache-2.0.txt": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
     "LICENSES/boltz/LICENSE": "sha256:f0667fd5e66c51e1ba8ddaa0249c6d7225b30037e02c45782d8f2c2943ac2617",
-    "README.md": "sha256:f12be28984a29a0257a997ac948faa35766e2914613af35c653d9772325e737c",
+    "README.md": "sha256:4d8e31396f39d7b2e1f5b131a4d688c391091531477e896005b2b28f418ef8e5",
     "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
-    "config.json": "sha256:7f24cc83dab019d998f24fdeffe46526c01f52b041527c9ae17bf66570f7040d",
-    "fastplms/__init__.py": "sha256:4fb3196022ca8ec699d59d09bdbc5f0184195552b773698ab9b061fe3cd7df12",
+    "config.json": "sha256:cefe82f26d2a25cbff6958749f3e8727efdb42e8a5eae8c8677498482f45933a",
+    "fastplms/__init__.py": "sha256:8503e02debf24abc6d33c1913ff3eb9f245e63f7e62017e220d072c7393fdf2c",
     "fastplms/models.toml": "sha256:1b0d92911222f31b435bfee80251c9abe20b85392d65f8af8d285d1ccfbbc3c5",
-    "fastplms/models/__init__.py": "sha256:5e48c2cb3877aa6f42f3b5411d53b16bba2e32827bbde634f47f174c5cb36f86",
-    "fastplms/models/boltz/__init__.py": "sha256:54f7ad44ddbb3fd160382a50bb5d5a8a5b86cb66978b7d72445194fc21caa673",
-    "fastplms/models/boltz/_pair_attention.py": "sha256:13d16831664ac7d0bfc859e2d4b97d523f1c2ae0fe4a0b11f65b9a3197dffccd",
-    "fastplms/models/boltz/cif_writer.py": "sha256:095eeab50ee8b97fe0160b4c1f9c6df2be1412a1de13a1d8ad00f9847032a7ab",
-    "fastplms/models/boltz/minimal_featurizer.py": "sha256:35e5f4f85d24d5ba8d952ee64f1ddb98e9ce5245f364d231da185ae8a3b28ce1",
+    "fastplms/models/__init__.py": "sha256:d3dd084f5e8fafcf2d0fe4a26e4ab9284d3f9414b1f9a5bd86e308a406447c43",
+    "fastplms/models/boltz/__init__.py": "sha256:ef2bad434c2cb2dd171051ce13765ac364651839748670ad2affb5ac816256e9",
+    "fastplms/models/boltz/_pair_attention.py": "sha256:da466a398fd30427758c937d49c570dc5319f792bf954091e51ded1c0b51e0ff",
+    "fastplms/models/boltz/cif_writer.py": "sha256:f15329f21a0dc070b564dfd95114f145ef638ecd93024d62f342b3b1b26b0334",
+    "fastplms/models/boltz/minimal_featurizer.py": "sha256:e7647e2cb62e7608fdfce986d773609ae3a749c63818039bb1c4f77edb5f8415",
     "fastplms/models/boltz/minimal_structures.py": "sha256:df7f1ec765ddbb5a730edb9f304fe312a5d3c431ae564ca9b97e6537e0ebf476",
-    "fastplms/models/boltz/modeling_boltz2.py": "sha256:c0d00a63a23762853257a4a4c718f7a4ded358c900562ca260ed3596a113d934",
+    "fastplms/models/boltz/modeling_boltz2.py": "sha256:98fa1c9d307b9595ccd53021ac678293c7fecaf7ea12e851b107fc369dc172b5",
     "fastplms/models/boltz/vb_const.py": "sha256:67d3cfdb7f0facdb217fe9220404d598ac7924a672f1a2720f7d2d82b23270df",
-    "fastplms/models/boltz/vb_layers_attention.py": "sha256:ad09126f467b9c73f764b6258ae3f2bce9f9752f641c0347052f12548abdfe4b",
-    "fastplms/models/boltz/vb_layers_attentionv2.py": "sha256:407b1ce628d53c3e0b6f801d989b0ae59bd87d9f9e5d1155e61c1e1ec5a78e0f",
-    "fastplms/models/boltz/vb_layers_confidence_utils.py": "sha256:61b024671ca2eaaa08eb6ddfd75f374f8db97ef6ef75ffcfc8c10347f4fd239a",
-    "fastplms/models/boltz/vb_layers_dropout.py": "sha256:61db64dc4e6be29a41bad2705f73cc489dcf3227cc6e8914c1ea9e6f45ba57a5",
-    "fastplms/models/boltz/vb_layers_initialize.py": "sha256:c1f5a5d56843438bd11639778eed06352e3d4995d2d12c967865d1221208048c",
-    "fastplms/models/boltz/vb_layers_outer_product_mean.py": "sha256:60478bea7d9d6d7f518ab1461c37b91e42d9fa0ab8cfcd777be88c1ddbfff07e",
-    "fastplms/models/boltz/vb_layers_pair_averaging.py": "sha256:e626502b549348ecf26077d0cb1f88f37970b6b9987a6780b4599c5a0b1c0ada",
-    "fastplms/models/boltz/vb_layers_pairformer.py": "sha256:fb056b76231366a1a93373ff2effddde887768dcfca1c9ebc46a567ba39eacb2",
-    "fastplms/models/boltz/vb_layers_transition.py": "sha256:62d306f33a600d08f8bb60a1a8cb9a00a253f81c05ef8c2af3f4fb7936ea8c18",
-    "fastplms/models/boltz/vb_layers_triangular_mult.py": "sha256:ace40becff882828ae3aa2b4f77e1a3d854d2adf7180216013996e38dd0e6fdd",
-    "fastplms/models/boltz/vb_loss_diffusionv2.py": "sha256:3778aea2ac541a3c173d9cc89b0d1b2cff63555571e0b3776b98c746daa0720c",
+    "fastplms/models/boltz/vb_layers_attention.py": "sha256:c0fc4b0b5b9a38d56f18af85d0ed506e0c45826d8d4c797da4a281dc35801390",
+    "fastplms/models/boltz/vb_layers_attentionv2.py": "sha256:f0509a5be8ee19c3d01797317bdcb6caaf1d37bde695bd6e6be33d06e19e46d7",
+    "fastplms/models/boltz/vb_layers_confidence_utils.py": "sha256:ee802fa8c364a38b17dfd0c141bc0b87ac03c8731e8029129756511aa7194e5b",
+    "fastplms/models/boltz/vb_layers_dropout.py": "sha256:e962d7fd699c27808f0348ccfc57b719e801ac443328c2fcd799a9cd44753241",
+    "fastplms/models/boltz/vb_layers_initialize.py": "sha256:e05108acaf92497f02d6a5567862f4bfe53265ac8f0117d7263dfa0e396ddac3",
+    "fastplms/models/boltz/vb_layers_outer_product_mean.py": "sha256:e435b827c110ae54bde63b0a184b65fe3b87cebdd50f3cea6d39e4818ce324da",
+    "fastplms/models/boltz/vb_layers_pair_averaging.py": "sha256:73b93eaa790f77d03bf68200ed505d3b2a40afd0f2d806709ecc2c277fe58b13",
+    "fastplms/models/boltz/vb_layers_pairformer.py": "sha256:932566fbffd07ab817af1cd5d2b71de98f8ab2bf9c0f6b270e21bf6adca2bb99",
+    "fastplms/models/boltz/vb_layers_transition.py": "sha256:e3aad7419fbe86b833e5f2991a779d12fc99d329f797a3b33ce836d8a4abb2e7",
+    "fastplms/models/boltz/vb_layers_triangular_mult.py": "sha256:13f55152ad8478ca6fc1dcc5c56bb9028ba8f4a030e5402001399e9359cb643e",
+    "fastplms/models/boltz/vb_loss_diffusionv2.py": "sha256:77137ab5da4393af127c3022c43ca4721906cba0c3cdc1d2b6b73fd70352c573",
     "fastplms/models/boltz/vb_modules_confidencev2.py": "sha256:a80a6f711e9ea021c45d89967fd37f7e668eea9ab209f9c70f2562d695a15041",
-    "fastplms/models/boltz/vb_modules_diffusion_conditioning.py": "sha256:a3e94f4493c056e54691b8289cbc35de3408b41bcbc7a75a8786eb1c273c1c31",
+    "fastplms/models/boltz/vb_modules_diffusion_conditioning.py": "sha256:dcbabd880c2682e910973512b875fc43e08ddf1a98ee13ab0b5ada01ca5a71d4",
     "fastplms/models/boltz/vb_modules_diffusionv2.py": "sha256:0077b40a72d9bf19904e0ab04fdb78b2c642ec6169b688732a1d221aa20a3637",
-    "fastplms/models/boltz/vb_modules_encodersv2.py": "sha256:9e78b92db67911795eb7691c7538edd12153ba30f4918ff8abf9874a99d02943",
-    "fastplms/models/boltz/vb_modules_transformersv2.py": "sha256:3eb5790030c333452beddeb3fe3583e65529a70f0e443d9d44304f7447de48f6",
-    "fastplms/models/boltz/vb_modules_trunkv2.py": "sha256:355cf735d57393167273b3be1d802ef495f69f8d20900e4cd2ce9b2df2147c25",
-    "fastplms/models/boltz/vb_modules_utils.py": "sha256:7a8f348cb8c703c6fef98312bea1beb179ceae5ae00eadc6962780a722171d05",
-    "fastplms/models/boltz/vb_potentials_potentials.py": "sha256:f332824228c5fabca51dd4695334e4a2822c6c36cca99dfc0c7f659eaf6062a3",
+    "fastplms/models/boltz/vb_modules_encodersv2.py": "sha256:135a20ce20c7fda10b92f2770464cd81be4899a3960ae4f5ef0a96dc6f6df105",
+    "fastplms/models/boltz/vb_modules_transformersv2.py": "sha256:22a77bdb8af6ee1235369cbe0f0c987eba74ae5cbf10b2fdfcfc5ad1b081aae1",
+    "fastplms/models/boltz/vb_modules_trunkv2.py": "sha256:cbd222f1c939f7d1728c5de97535ab1b940205dbdc23bb1895e71d405568d192",
+    "fastplms/models/boltz/vb_modules_utils.py": "sha256:10a69588abb92424007971a5c5bb91294a07d3a0850932f7145276d132a73041",
+    "fastplms/models/boltz/vb_potentials_potentials.py": "sha256:ca2aacf0c2e5b8de13c5c6456967fb4d9b1cd72f74e13584e9248bbfc8b1a918",
     "fastplms/models/boltz/vb_potentials_schedules.py": "sha256:91997ee4b7a2a5ba1f40475c9e4fe6a58927f5ef7b51012743615b0f05e278d1",
-    "fastplms/models/boltz/vb_tri_attn_attention.py": "sha256:3c599f47575bfa762afd1b2fa70d6135ad66af72bfff8b703803d1f34f861080",
-    "fastplms/models/boltz/vb_tri_attn_primitives.py": "sha256:99b149ed2775b2d58baf4d5b4be414cfb2183234a0f46567ac6a37d9327909e7",
-    "fastplms/models/boltz/vb_tri_attn_utils.py": "sha256:df831665ee4bf485c5f95b3bc2a04dcb37aa002b03c42417036fc866e87622f1",
-    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
-    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
-    "fastplms_bundle.py": "sha256:f55c37e173eca40cfcd7e38ba7cdc2d0b15f6b846e9b478d7eff146ac15d3da7",
-    "modeling_fastplms.py": "sha256:f5defd6c4657ff805c1e835e51b2ef4becca9a21e255df767bd210ef34f808e5"
+    "fastplms/models/boltz/vb_tri_attn_attention.py": "sha256:70afafa17bb4c6fa1151a188dfe8a635b1224a6d814bf16b0a4d2949edbdd29a",
+    "fastplms/models/boltz/vb_tri_attn_primitives.py": "sha256:3d4beefb52ee69ee98422c3d9058f8e394c2125fc029e79eb959eee9d0eea4ce",
+    "fastplms/models/boltz/vb_tri_attn_utils.py": "sha256:bd0566f3c6d5e36d7e6c1a3f79a95fa2e06c1d18d554d4f9331a7a69ac7ab2e1",
+    "fastplms/registry.py": "sha256:cb37799dae5a4c0c780089ab2cca4cc56300c1c5f515b275892aa4e1c4b97eb6",
+    "fastplms/runtime.py": "sha256:9521c37fcd168bf9a4137277fb73672403ccf30d70d4845aa879fb22f42ddeb6",
+    "fastplms_bundle.py": "sha256:e23521271a92290a4da0a66314fba4a3c499826b802817c16340922f85e02421",
+    "modeling_fastplms.py": "sha256:39540ac66a63879fac53c2cf365f4d378bc9f81814f7ea6c4261c41a75e789e0",
+    "requirements.txt": "sha256:bf8969fb385af9448d39a9574a024d20e1e65609d60c28b5b3e58e6bf55bc351"
   },
   "model_id": "boltz2",
   "redistributable": true,
-  "release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "runtime_bundle_sha256": "e6e0f8c8632238f40a5112f22397a9a211eb02d3415c1c93d43305aef51f7728",
-  "runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
+  "release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "runtime_bundle_sha256": "3aebb85d11edb01329d2aafb4c8260bd67339785f0f252582246ae7c33206143",
+  "runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
   "schema_version": 2,
   "scope": "runtime-only",
-  "source_tree_sha256": "78061b1bd37ba2dda853641fba681fdf943b859503e4fb21b1a7508909c369bf",
+  "source_tree_sha256": "9d2fc9586cfbd1d2d2cc49d745311acb81d4e3740e59ead09e633917589806a8",
   "weights": {
     "repo_id": "Synthyra/Boltz2",
     "revision": "3b148fc5efea109c065ec82ba8683d024de7134e"