Update FastPLMs files from Synthyra/FastPLMs 32d9951 (PR 52: ESMFold2 noise cap, DPLM2 parity)

README.md

@@ -32,7 +32,7 @@ python -m pip install -r \
 The FastPLMs implementation itself is embedded in the model repository.
 Transformers loads it through `trust_remote_code=True`.
 
-This model requires Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13.
+This model requires Python 3.12-3.14, PyTorch 2.14, and Transformers 5.17.
 
 The artifact requirements include the structure dependencies.
 
@@ -74,7 +74,7 @@ materialize attention tensors. The configured backend does not change.
 Install the training dependencies. Then attach LoRA to the loaded checkpoint:
 
 ```bash
-python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
+python -m pip install "datasets>=5.0" "peft>=0.21"
 ```
 
 ```python

fastplms/atomic_files.py

@@ -0,0 +1,65 @@
+"""Replace a file in one step: readers see the old bytes or the new bytes, never a partial write."""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import IO, Any
+
+
+def write_bytes_atomically(path: Path, payload: bytes, *, create_parent: bool = False) -> None:
+    """Write ``payload`` to a temporary file beside ``path``, flush it to disk, then rename it over ``path``.
+
+    A failed write removes the temporary file and leaves ``path`` untouched. The parent directory must
+    exist unless ``create_parent`` is true. The temporary file sits in the same directory so the rename
+    never crosses a file system.
+    """
+
+    with _staged_handle(path, "wb", create_parent=create_parent) as handle:
+        handle.write(payload)
+
+
+def write_text_atomically(
+    path: Path,
+    text: str,
+    *,
+    encoding: str = "utf-8",
+    newline: str | None = None,
+    create_parent: bool = False,
+) -> None:
+    """Write ``text`` as ``write_bytes_atomically`` does.
+
+    ``newline`` is the argument of ``open``: ``None`` translates ``"\\n"`` to the platform separator and
+    ``"\\n"`` writes it unchanged.
+    """
+
+    with _staged_handle(
+        path, "w", create_parent=create_parent, encoding=encoding, newline=newline
+    ) as handle:
+        handle.write(text)
+
+
+@contextmanager
+def _staged_handle(
+    path: Path, mode: str, *, create_parent: bool, **open_arguments: Any
+) -> Iterator[IO[Any]]:
+    """Yield a handle to a temporary file; on a clean exit flush it and rename it over ``path``."""
+
+    if create_parent:
+        path.parent.mkdir(parents=True, exist_ok=True)
+    descriptor, temporary_name = tempfile.mkstemp(
+        prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
+    )
+    try:
+        with os.fdopen(descriptor, mode, **open_arguments) as handle:
+            yield handle
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temporary_name, path)
+    except BaseException:
+        Path(temporary_name).unlink(missing_ok=True)
+        raise

fastplms/digests.py

@@ -0,0 +1,35 @@
+"""SHA-256 digests of files and JSON values, the identities FastPLMs records and compares.
+
+This file exists twice, byte for byte: here and as ``features/digests.py``. ``features`` loads as a standalone
+package (``foundry.embedding.private_store``) and imports nothing outside itself, so it cannot reach this
+module. ``tests/tier1_unit/test_features_package_is_self_contained.py`` fails when the two differ.
+"""
+
+from __future__ import annotations
+
+import hashlib
+
+from pathlib import Path
+from typing import Any
+
+from .json_files import compact_json
+
+
+FILE_READ_BYTES = 1024 * 1024
+
+
+def file_sha256(path: str | Path) -> str:
+    """Return the SHA-256 of a file's bytes, read in blocks so a checkpoint never sits in memory."""
+
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as handle:
+        while block := handle.read(FILE_READ_BYTES):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def json_sha256(value: Any, *, ensure_ascii: bool = True, allow_nan: bool = True) -> str:
+    """Return the SHA-256 of ``value`` in its compact, key-sorted JSON form (``compact_json``)."""
+
+    encoded = compact_json(value, ensure_ascii=ensure_ascii, allow_nan=allow_nan).encode("utf-8")
+    return hashlib.sha256(encoded).hexdigest()

fastplms/json_files.py

@@ -0,0 +1,45 @@
+"""The two JSON text forms FastPLMs writes: compact for hashing, indented for files people read.
+
+This file exists twice, byte for byte: here and as ``features/json_files.py``. ``features`` loads as a
+standalone package (``foundry.embedding.private_store``) and imports nothing outside itself, so it cannot
+reach this module. ``tests/tier1_unit/test_features_package_is_self_contained.py`` fails when the two differ.
+"""
+
+from __future__ import annotations
+
+import json
+
+from typing import Any
+
+
+def compact_json(value: Any, *, ensure_ascii: bool = True, allow_nan: bool = True) -> str:
+    """Serialize with sorted keys and no whitespace, the form a digest or an identity is taken over."""
+
+    return json.dumps(
+        value,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=ensure_ascii,
+        allow_nan=allow_nan,
+    )
+
+
+def indented_json(
+    value: Any,
+    *,
+    ensure_ascii: bool = True,
+    allow_nan: bool = True,
+    sort_keys: bool = True,
+) -> str:
+    """Serialize with two-space indentation and one trailing newline, the form of a stored JSON file."""
+
+    return (
+        json.dumps(
+            value,
+            indent=2,
+            sort_keys=sort_keys,
+            ensure_ascii=ensure_ascii,
+            allow_nan=allow_nan,
+        )
+        + "\n"
+    )

fastplms/models.toml

@@ -4,11 +4,58 @@ legal_files = [
   "THIRD_PARTY_NOTICES.md=sha256:d35e506b728868b52290672d54a89c6aade09529af99dbc8a4db72d9e9ca3460",
 ]
 
+# Released hidden-state SAEs used by the canonical suite. The training-backbone revision
+# is not reported upstream; base_model names the selected inference base, not that history.
+[[sparse_autoencoders]]
+id = "esmc_small_sae_depth"
+base_model = "esmc_small"
+layer = 23
+input_width = 960
+k = 64
+codebook_dim = 16384
+input_kind = "hidden_state"
+checkpoint_repo = "biohub/ESMC-300M-sae-layer23-k64-codebook16384"
+checkpoint_revision = "71b054fc726a9f11153a601118f03657ac6e6e80"
+checkpoint_files = [
+  "config.json=sha256:e2864c5f9756c052f343310bafeddafc3edc9514ad85981a51a1a3411186434d",
+  "layer_23.safetensors=sha256:825257f785600a7a9d462882d5cf3431ebbbba94761434e33b0c4814a7cc6270",
+]
+
+[[sparse_autoencoders]]
+id = "esmc_large_sae_depth"
+base_model = "esmc_large"
+layer = 27
+input_width = 1152
+k = 64
+codebook_dim = 16384
+input_kind = "hidden_state"
+checkpoint_repo = "biohub/ESMC-600M-sae-layer27-k64-codebook16384"
+checkpoint_revision = "b96480999aca49b7684e093f1adfcb09b55a1720"
+checkpoint_files = [
+  "config.json=sha256:0e74c8af3ff4ac644f26be308e0fe6da6c75a4a52f0d3f05b0cbbdda27efc9ce",
+  "layer_27.safetensors=sha256:5e93c230aa200876cb300377c8e22a141a7e2e316c6a3b962b2a78c4518535db",
+]
+
+[[sparse_autoencoders]]
+id = "esmc_6b_sae_depth"
+base_model = "esmc_6b"
+layer = 60
+input_width = 2560
+k = 64
+codebook_dim = 16384
+input_kind = "hidden_state"
+checkpoint_repo = "biohub/ESMC-6B-sae-layer60-k64-codebook16384"
+checkpoint_revision = "99752fe6e4d25fbb26f887db2d9225e7a577da73"
+checkpoint_files = [
+  "config.json=sha256:42567d3f757bea5c6b618abf55ed905ce4304d1eed891a2c1feae8e1ef393fd0",
+  "layer_60.safetensors=sha256:ddc1417b42cffe2d7fc2ec31783a04ac88800e6a8f8c8542f1f92d8f4f16094c",
+]
+
 [[attention_kernels]]
 implementation = "flash_attention_2"
 repository = "kernels-community/flash-attn2"
-revision = "db6b51744f0cd7061386442c09df890fc6d9f47e"
-version = 2
+revision = "81fb77c12b2ad5d69380669b46739d5868614502"
+version = 3
 expected_variant = "flash_attn2"
 dtypes = ["bfloat16"]
 min_cuda_capability = [8, 0]
@@ -64,6 +111,9 @@ distribution_files = [
 id = "biohub-transformers"
 path = "vendor/upstream/biohub-transformers"
 url = "https://github.com/Biohub/transformers.git"
+# Biohub/transformers left GitHub by 2026-10-08. Its pinned commit stays reachable through the
+# fork network of huggingface/transformers, which serves the clone and the source archive.
+fetch_url = "https://github.com/huggingface/transformers.git"
 revision = "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf"
 license = "Apache-2.0"
 license_files = ["LICENSE"]
@@ -177,7 +227,7 @@ conversion_provenance = "Input: the pinned official ESM2 state dictionary. Trans
 representative = "esm2_8m"
 documentation = "docs/models.md#esm2"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/esm2", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/_esm_rotary.py", "models/esm2", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.esm2.modeling_fastesm.FastEsmConfig", AutoModel = "fastplms.models.esm2.modeling_fastesm.FastEsmModel", AutoModelForMaskedLM = "fastplms.models.esm2.modeling_fastesm.FastEsmForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.esm2.modeling_fastesm.FastEsmForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm2.modeling_fastesm.FastEsmForTokenClassification" }
 
 [families.esm_plusplus]
@@ -204,7 +254,7 @@ conversion_provenance = "Input: the pinned Biohub ESMC checkpoint. Transformatio
 representative = "esmc_small"
 documentation = "docs/models.md#esm-and-esmc"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm_plusplus", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/esm_plusplus", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusConfig", AutoModel = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusModel", AutoModelForMaskedLM = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForTokenClassification" }
 
 [families.esm3]
@@ -229,7 +279,7 @@ conversion_provenance = "Input: the pinned Biohub ESM3 checkpoint. Transformatio
 representative = "esm3_small"
 documentation = "docs/models.md#esm3"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm3", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/esm3", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.esm3.modeling_esm3.FastESM3Config", AutoModel = "fastplms.models.esm3.modeling_esm3.FastESM3Model", AutoModelForSequenceClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForTokenClassification" }
 
 [families.e1]
@@ -256,7 +306,7 @@ conversion_provenance = "Input: the pinned Profluent-E1 checkpoint and tokenizer
 representative = "e1_150m"
 documentation = "docs/models.md#e1"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/e1", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/e1", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.e1.modeling_e1.E1Config", AutoModel = "fastplms.models.e1.modeling_e1.E1Model", AutoModelForMaskedLM = "fastplms.models.e1.modeling_e1.E1ForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.e1.modeling_e1.E1ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.e1.modeling_e1.E1ForTokenClassification" }
 
 [families.dplm]
@@ -282,7 +332,7 @@ conversion_provenance = "Input: the pinned official DPLM1 checkpoint. Transforma
 representative = "dplm_150m"
 documentation = "docs/models.md#dplm"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.dplm.modeling_dplm.DPLMConfig", AutoModel = "fastplms.models.dplm.modeling_dplm.DPLMModel", AutoModelForMaskedLM = "fastplms.models.dplm.modeling_dplm.DPLMForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.dplm.modeling_dplm.DPLMForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.dplm.modeling_dplm.DPLMForTokenClassification" }
 
 [families.dplm2]
@@ -307,7 +357,7 @@ conversion_provenance = "Input: the pinned official DPLM2 checkpoint. Transforma
 representative = "dplm2_150m"
 documentation = "docs/models.md#dplm2"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm2", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm2", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.dplm2.modeling_dplm2.DPLM2Config", AutoModel = "fastplms.models.dplm2.modeling_dplm2.DPLM2Model", AutoModelForMaskedLM = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForTokenClassification" }
 tokenizer_class = "fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer"
 
@@ -334,7 +384,7 @@ representative = "ankh_base"
 documentation = "docs/models.md#ankh"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
 requires_complete_weight_publication = false
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/ankh", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/ankh", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.ankh.modeling_ankh.FastAnkhConfig", AutoModel = "fastplms.models.ankh.modeling_ankh.FastAnkhModel", AutoModelForMaskedLM = "fastplms.models.ankh.modeling_ankh.FastAnkhForMaskedLMExtension", AutoModelForSeq2SeqLM = "fastplms.models.ankh.modeling_ankh.FastAnkhForConditionalGeneration", AutoModelForSequenceClassification = "fastplms.models.ankh.modeling_ankh.FastAnkhForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.ankh.modeling_ankh.FastAnkhForTokenClassification" }
 
 [families.boltz2]
@@ -358,7 +408,7 @@ conversion_provenance = "Input: the pinned official Boltz2 checkpoint. Transform
 representative = "boltz2"
 documentation = "docs/models.md#boltz2"
 test_tiers = ["structure", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "models/boltz"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "models/boltz"]
 auto_map = { AutoConfig = "fastplms.models.boltz.modeling_boltz2.Boltz2Config", AutoModel = "fastplms.models.boltz.modeling_boltz2.Boltz2Model" }
 
 [families.esmfold]
@@ -383,7 +433,7 @@ conversion_provenance = "Input: the pinned native Meta ESMFold checkpoint plus i
 representative = "esmfold"
 documentation = "docs/models.md#esmfold"
 test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/classification_probe.py", "models/esmfold"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/_esm_rotary.py", "models/classification_probe.py", "models/esmfold"]
 auto_map = { AutoConfig = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig", AutoModel = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding", AutoModelForSequenceClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForTokenClassification" }
 
 [families.esmfold2]
@@ -410,7 +460,7 @@ conversion_provenance = "Input: each pinned Biohub ESMFold2 checkpoint and its s
 representative = "esmfold2"
 documentation = "docs/esmfold2.md"
 test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/classification_probe.py", "models/_esm_rotary.py", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/classification_probe.py", "models/_esm_rotary.py", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForTokenClassification" }
 
 [[models]]
@@ -418,7 +468,7 @@ id = "esm2_8m"
 family = "esm2"
 size_category = "small"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_8m.json=sha256:6975e86d1d8f27488bf2a676551feaa48cc19254c9d24b6acb09198122745609", tensors = "tests/goldens/esm2_8m.safetensors=sha256:b40217566c33c71988d28869de353be54a3b3ebfc21fdfd29056e88cf7e99f4c" }
+official_golden = { metadata = "tests/goldens/esm2_8m.json=sha256:06b3ccbf45a46a7aed3833b49503567f7534ffe60e6b3e49fcfcf17d04b7237e", tensors = "tests/goldens/esm2_8m.safetensors=sha256:d08a7572cbef20b8b19b545bcb0427b7e9ae19b986d015c559cd5a3a2cfc8aa4" }
 fast_repo = "Synthyra/ESM2-8M"
 fast_revision = "185ecbd45665d050a8dae326d91886d330c5f9d0"
 fast_files = [
@@ -457,7 +507,7 @@ id = "esm2_35m"
 family = "esm2"
 size_category = "small"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_35m.json=sha256:e919d3ce6d20b6a942d27d92323814ae7594a0129dc9c4de27c5053e96675bcd", tensors = "tests/goldens/esm2_35m.safetensors=sha256:c9b8bb616cf884fb7744521a2fcc6eed23586342d11241e6c9ef16454ec31e17" }
+official_golden = { metadata = "tests/goldens/esm2_35m.json=sha256:61cd24fd91ef2dc49cbc0f97b14c0d5c97849eca6f36ed8c294295147909cb38", tensors = "tests/goldens/esm2_35m.safetensors=sha256:4f82d10286e16041c2f23365dfdd8508b911864633287dd76580425527c2d922" }
 fast_repo = "Synthyra/ESM2-35M"
 fast_revision = "37ab9f56b41e365b3bd9e25d6fefe9150fd910f0"
 fast_files = [
@@ -496,7 +546,7 @@ id = "esm2_150m"
 family = "esm2"
 size_category = "medium"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_150m.json=sha256:c04c93486024ba0fa1c81fbfbe92ee79d1d4c7f1cfcc2c9886728522f752feab", tensors = "tests/goldens/esm2_150m.safetensors=sha256:c03fe9916dba137b452a6bbe944c7dc414db4019a6f0921e87b92d4bb6a8a42f" }
+official_golden = { metadata = "tests/goldens/esm2_150m.json=sha256:864e1f9c1d8e939e8da050b8617b8a7d6983fbb11a5911275d1571d08e460a79", tensors = "tests/goldens/esm2_150m.safetensors=sha256:20ad986b8e6e09f36d0158f5939d1914e0809b84329695028d767aab93338621" }
 fast_repo = "Synthyra/ESM2-150M"
 fast_revision = "979e0880dfc9e0c0080839b83d9d2dc05b92786a"
 fast_files = [
@@ -535,7 +585,7 @@ id = "esm2_650m"
 family = "esm2"
 size_category = "large"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_650m.json=sha256:f18332172fcb3abf5dd2485fd55f5b0d193ad3b93a44cc744e0d02817c927477", tensors = "tests/goldens/esm2_650m.safetensors=sha256:c3a66b75add03628e62e238cb63da6a9e4d321f8160e84bdf2a131c096977f86" }
+official_golden = { metadata = "tests/goldens/esm2_650m.json=sha256:b57f18c8803e68e6c621f479fd7a70596fbc416fc4a17d2ab91c639b480006c4", tensors = "tests/goldens/esm2_650m.safetensors=sha256:261a8b71c4b7c90b1f294031558ee0bf00d77a98352021eb78771200c8d40548" }
 fast_repo = "Synthyra/ESM2-650M"
 fast_revision = "ca0718a5d52b80d5c60dd76860e55e061a95fb0a"
 fast_files = [
@@ -574,7 +624,7 @@ id = "esm2_3b"
 family = "esm2"
 size_category = "xlarge"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_3b.json=sha256:5043b2333c57a34d54fac53916722d1acb4b6fd50395b9abafa805435b184a48", tensors = "tests/goldens/esm2_3b.safetensors=sha256:dfd5a8cb05d3e814a080185c4808c8e7ec2277f070f395562fcfbe4376789e4e" }
+official_golden = { metadata = "tests/goldens/esm2_3b.json=sha256:aa441786329889d43983811cac155cabd37b38cfd47bc90d415732a295c33c6d", tensors = "tests/goldens/esm2_3b.safetensors=sha256:61080e55e07db4a562a19e9b6b662e71a9fd3d8fee357aa3df00697d2fa58e33" }
 notes = "The pinned default SDPA BF16 path uses a checkpoint-specific numeric calibration: relative L2 target/hard limit 0.06/0.07, relative Q99.9 0.15/0.18, first-percentile residue cosine 0.994/0.992, and pooled cosine 0.998/0.997. Exact state identity and the global logits-distribution contract remain required."
 fast_repo = "Synthyra/ESM2-3B"
 fast_revision = "ff89d0180f414ab9c677219a25da79bf09185456"
@@ -617,7 +667,7 @@ id = "esmc_small"
 family = "esm_plusplus"
 size_category = "medium"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esmc_small.json=sha256:bb02652cf3cc484756b98ffa4ba55ed4c55870d2cea3342adb1d920ba9dfe10a", tensors = "tests/goldens/esmc_small.safetensors=sha256:03378d0f0fdd8161178ebb2c1f0da1b9776a726c8e8d3a10c009808a24de5654" }
+official_golden = { metadata = "tests/goldens/esmc_small.json=sha256:15cbbf909a12b5b81ff47b5f142babb75ebf9b3d85673905f2320c2a119ae347", tensors = "tests/goldens/esmc_small.safetensors=sha256:98219356e2845cbd10a80715b19c95632956d40567a7a993409c426b56040b06" }
 notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration."
 fast_repo = "Synthyra/ESMplusplus_small"
 fast_revision = "46c5f7d562e47d4c14165b424c71ab7db008e6fb"
@@ -643,7 +693,7 @@ id = "esmc_large"
 family = "esm_plusplus"
 size_category = "large"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esmc_large.json=sha256:7a4d614f67b6fde417f3fd89f61e7ec442ae284769734b2b73e14945a816a8fd", tensors = "tests/goldens/esmc_large.safetensors=sha256:e13302df4cf7e8381552f1043a8fd0f31f3e0d50b2ab6009fb86b7940ae8ff79" }
+official_golden = { metadata = "tests/goldens/esmc_large.json=sha256:86ffe179a9feafd067793f440d25e3606f184f0d32063aeaeb1842b14aa55d3b", tensors = "tests/goldens/esmc_large.safetensors=sha256:2b1ce3de5a7a27f171055f5c7aec4e7e50f0e809a1dc46508aaa0b472365015d" }
 notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration."
 fast_repo = "Synthyra/ESMplusplus_large"
 fast_revision = "f813401638b3fddab09748aec1ad2bf537aa4208"
@@ -669,7 +719,7 @@ id = "esmc_6b"
 family = "esm_plusplus"
 size_category = "xlarge"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esmc_6b.json=sha256:e229d938719782f280fab22dfc4c43e86109fdb0cc523631168c5a491afaace3", tensors = "tests/goldens/esmc_6b.safetensors=sha256:a948945e985c7deaca7be8b7eed09c0a9521a2af3f2b10fc2ec7a7d2a0f99ada" }
+official_golden = { metadata = "tests/goldens/esmc_6b.json=sha256:7e044e7e8d106d43169c4148521c54d0d09bb989624543d648247a110f1907d4", tensors = "tests/goldens/esmc_6b.safetensors=sha256:46139de6b28ab4f9e3244fb6ef3f259fc5e6518a3825f0fb4fd4a4fece6da8e6" }
 notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration."
 fast_repo = "Synthyra/ESMplusplus_6B"
 fast_revision = "0d579cce3b0f09efa6b3baddf6cc3fd8c9b616c8"
@@ -681,6 +731,7 @@ fast_files = [
   "model-00004-of-00006.safetensors=sha256:e46c6113c89c6f3e9b072c1bef02d763a625c37bcd8f9da2ed9363891c9a0758",
   "model-00005-of-00006.safetensors=sha256:6d92cb2bf9791de644de2ae86f8523d802ac3b4aaabfff0716ab6c2b97f6fb14",
   "model-00006-of-00006.safetensors=sha256:5fc1a8632490bb34162823c35d0d591337b9e4195b22cc0560741397a6e9d0b3",
+  "model.safetensors.index.json=git-sha1:f30f0b6b35e11d09a516c078b0847ee023bd1817",
   "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
   "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
   "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
@@ -706,7 +757,7 @@ family = "esm3"
 tokenizer_source = "esmc_small"
 size_category = "large"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm3_small.json=sha256:5470e8596cbba0e2882647eccbc53c36d8b48b0f3947d1fe0bcea68da1078c32", tensors = "tests/goldens/esm3_small.safetensors=sha256:d957922f810c9ab4c557d80d5aaaf6a3aab79a5a45e4638012a634a4134803b1" }
+official_golden = { metadata = "tests/goldens/esm3_small.json=sha256:e476bad7ccdfa2f908261a03e20116c0854eece292a553196f940b594088204d", tensors = "tests/goldens/esm3_small.safetensors=sha256:251e050926a1f2426401bac6b93b4cb00041d0b0a5c8a458d0ef0e6a5d0d87c3" }
 fast_repo = "Synthyra/ESM3_small"
 fast_revision = "7ddb5a740f9e5f93933eb6410c0ee8684bc63ec1"
 fast_files = [
@@ -732,7 +783,7 @@ id = "e1_150m"
 family = "e1"
 size_category = "small"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/e1_150m.json=sha256:701a64a6ab1a2fec5a427555b6af96232526c15cb3d5b4dc7fb253ac8f20b922", tensors = "tests/goldens/e1_150m.safetensors=sha256:6558bc8f1a7b20629eaaaa6f72601d0c2cdb859a5dc13595549b1773b6e2de41" }
+official_golden = { metadata = "tests/goldens/e1_150m.json=sha256:e7520331135fd4c98f52f21c470dbba5cb5f4defb418494c63199fb000621e65", tensors = "tests/goldens/e1_150m.safetensors=sha256:03a2e93e7b3e54b12f99eea7678ff80828922f92c3683d4909346a05d74815bc" }
 fast_repo = "Synthyra/Profluent-E1-150M"
 fast_revision = "7c5f3bbf697226a2e0900db7a100f9201774a907"
 fast_files = [
@@ -751,7 +802,7 @@ id = "e1_300m"
 family = "e1"
 size_category = "medium"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/e1_300m.json=sha256:d3478f3f5957a0e0377864074dde0107de890019f96cb63548ee17ffb8f3ec3a", tensors = "tests/goldens/e1_300m.safetensors=sha256:92778b9ef95a803ddc84b3e3ca764c59e045872a94bcff0eb0cd47647732c188" }
+official_golden = { metadata = "tests/goldens/e1_300m.json=sha256:d97637bf789fed61b4ad575f904f37befcea42e9436e8d841e8f8f2cff1d1b17", tensors = "tests/goldens/e1_300m.safetensors=sha256:d125ea866899d2788330d6877d77c48c69febbf7f01d629f4107f9647cfa13ff" }
 fast_repo = "Synthyra/Profluent-E1-300M"
 fast_revision = "5ef52c0ad2ae2578f40622696b763523810e8e26"
 fast_files = [
@@ -770,7 +821,7 @@ id = "e1_600m"
 family = "e1"
 size_category = "large"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/e1_600m.json=sha256:914be191c28141c1f84535cdb69ead0588a2057bb19d46c5bc7f3891a3d6739e", tensors = "tests/goldens/e1_600m.safetensors=sha256:22ed8417a4651ded255099f6d15c63c2c40552e700d2b0470d1adfde3a39c513" }
+official_golden = { metadata = "tests/goldens/e1_600m.json=sha256:0c9afc83b96ed8ad7f339a26df4ba5853e7994ad54dd3cd837deb9869c0cd2a9", tensors = "tests/goldens/e1_600m.safetensors=sha256:d470d6833f0b38bec070aa37e9f3c49717ea82d31c794af843550bb7e56253a6" }
 fast_repo = "Synthyra/Profluent-E1-600M"
 fast_revision = "6c8bf0ec83b0e0178677c528b101efffd0677742"
 fast_files = [
@@ -789,7 +840,7 @@ id = "dplm_150m"
 family = "dplm"
 size_category = "small"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm_150m.json=sha256:3228551fe3bed951db9ec97347143ec4462ce7c221ac240b7ce7730948c1dc1f", tensors = "tests/goldens/dplm_150m.safetensors=sha256:392992235195beed97ab8359b90a2e11e52f4326606f99a471447bed81d146bd" }
+official_golden = { metadata = "tests/goldens/dplm_150m.json=sha256:541167ffcd12b2d7e101046b2d17e800c70ba5db1b561212575301f2d1c407f8", tensors = "tests/goldens/dplm_150m.safetensors=sha256:f8e4ddd580d4708dad2d7da002735b2a01d69ece478c46c4932eb2f3933dc836" }
 fast_repo = "Synthyra/DPLM-150M"
 fast_revision = "90ba742754151a774f3b7ed580170d0a76b3e69d"
 fast_files = [
@@ -814,7 +865,7 @@ id = "dplm_650m"
 family = "dplm"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm_650m.json=sha256:bf58d0ce73aaac7e6fb1923ef3d9adad67122df2a3dd414c3229488ef9587a6d", tensors = "tests/goldens/dplm_650m.safetensors=sha256:073f0a6abea7e48f28c2d921ff8329a28e22627f01979277cb324908a01b3378" }
+official_golden = { metadata = "tests/goldens/dplm_650m.json=sha256:2bd482c94b1b6b3ac901b5ff9b06e2d189cda1a8fe2f8ff32e5d52b64c2c446b", tensors = "tests/goldens/dplm_650m.safetensors=sha256:70e2c46ed722942d3b628b556b92c9994688cb4853b278eb1e40a3dfbbd9dfc4" }
 fast_repo = "Synthyra/DPLM-650M"
 fast_revision = "05dc16d97c5c028aed924c9ed681cee4ab609760"
 fast_files = [
@@ -839,7 +890,7 @@ id = "dplm_3b"
 family = "dplm"
 size_category = "xlarge"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm_3b.json=sha256:a5b6df8b9c7b371976892ec1d6c45581a32ad3a6325c6c0a0b3267012848c8ed", tensors = "tests/goldens/dplm_3b.safetensors=sha256:75b0a0854fc391133920b0feaaeb8f69ab7568a88b3759627aca1556c4338c1e" }
+official_golden = { metadata = "tests/goldens/dplm_3b.json=sha256:0f8d8fb7df562c30abf730abe55db4a8c13d40877dd21772375fe1f34030159a", tensors = "tests/goldens/dplm_3b.safetensors=sha256:6942218079232a8185b1ae6e1578474c7784722e28f528b76ec89ca77777e2b9" }
 fast_repo = "Synthyra/DPLM-3B"
 fast_revision = "7d764dd3d70ecf1ac0e64693de64a0064aacac65"
 fast_files = [
@@ -869,7 +920,7 @@ id = "dplm2_150m"
 family = "dplm2"
 size_category = "small"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm2_150m.json=sha256:d269de779ea1503de72c77e7b2e6224afc9797bd945b40c571ff6faec782e4aa", tensors = "tests/goldens/dplm2_150m.safetensors=sha256:17fc26600938ba5364b8ecb96750786d33e9f92bcd4ea4df3e12a389340748eb" }
+official_golden = { metadata = "tests/goldens/dplm2_150m.json=sha256:5f2c496e4b557faf70bdde03b531f753fc7f20c5868bab79dcf378ee358781a1", tensors = "tests/goldens/dplm2_150m.safetensors=sha256:7390559bc54f27f08972c8f0adcf17954065450d937b3984303c1ae1ec003855" }
 artifact_source = "official"
 canonical_state_sha256 = "82e1751f59052b8de72b082517557db47947e8d9b4ac2f11278369e6c0cbf001"
 fast_repo = "Synthyra/DPLM2-150M"
@@ -896,7 +947,7 @@ id = "dplm2_650m"
 family = "dplm2"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm2_650m.json=sha256:d9a7548f9af657a72d441ca70f27379863724fcce8ddd3da4f672104b7bfb772", tensors = "tests/goldens/dplm2_650m.safetensors=sha256:c4e0e467c252c3ac813363d2d4b17a5e3bd99e75fad315e76d97689b4655ddac" }
+official_golden = { metadata = "tests/goldens/dplm2_650m.json=sha256:af7b8c4fab48773e1ddd2cfd941cf3d518ea27ed8519f467ddf0da788a0aa6be", tensors = "tests/goldens/dplm2_650m.safetensors=sha256:c45f4ca267183b11dd06244ed3c3f39cf5b6139f778733fa63a6e2b5a46b0678" }
 artifact_source = "official"
 canonical_state_sha256 = "cba76b6602d2258de9fffff953b608d93cb8ef4a9e89b0bbd27e160c81e78bb4"
 fast_repo = "Synthyra/DPLM2-650M"
@@ -925,7 +976,7 @@ size_category = "xlarge"
 # The pinned public sampler fails before generation because cls_token_id is None.
 # State, tokenizer, and inference parity remain required for this checkpoint.
 generation_contract = "official_unavailable"
-official_golden = { metadata = "tests/goldens/dplm2_3b.json=sha256:d6e0e02af53b13cb129192f06e264758aa21c9ebf4ee82411cf67037082d2329", tensors = "tests/goldens/dplm2_3b.safetensors=sha256:838b11824d08f83bcb0c0b3268e579f3a87dbfb965370cfe5c3f8793b96b1964" }
+official_golden = { metadata = "tests/goldens/dplm2_3b.json=sha256:ded7978cf84da1d56d8b4693418030ec3b0086bf6466be2242fb0ce424296b9f", tensors = "tests/goldens/dplm2_3b.safetensors=sha256:7490f771f8b8f8137b333f98d6b7dbd2cb48708c8649d4459279a6711b9e5870" }
 notes = "The pinned official DPLM2-3B sampler fails before generation, so live generation equivalence cannot be established for this checkpoint. State, tokenizer, and inference parity remain required."
 artifact_source = "official"
 canonical_state_sha256 = "8c46ec09115dbe6cbfb91d94ab5e906369d57e27fe620a7741c6f8cb1b6ca890"
@@ -958,7 +1009,7 @@ id = "ankh_base"
 family = "ankh"
 size_category = "medium"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh_base.json=sha256:ebce8d7de821827ee995789c9b38d79252d3b2f76888130b0a8a7eedafaefe2b", tensors = "tests/goldens/ankh_base.safetensors=sha256:f0e78aa15d11749e0c64ff57f9e88c51cec6538a0adf8951f839df70cc708b65" }
+official_golden = { metadata = "tests/goldens/ankh_base.json=sha256:21f129b7b71c026cbd3b3fdb712af06d5b6f256b487ef8b01fcc506d2bc6c7b6", tensors = "tests/goldens/ankh_base.safetensors=sha256:449db3638117d5b6027380fdf4b05698bb156024b2b3f7e0ceb295a86e5ec756" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
 artifact_source = "official"
 canonical_state_sha256 = "cdd8d30d88e5bf41f44e1eef4470d8e46607aba5f7c7c805b06c035b89c8c16f"
@@ -987,7 +1038,7 @@ id = "ankh_large"
 family = "ankh"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh_large.json=sha256:59492518b021de5cfaea87d672c9448c8558e99a3443ba2cc7ab544963196ecb", tensors = "tests/goldens/ankh_large.safetensors=sha256:3fb8d3ac27716d15a9ea92aeef6acf2b977bcc887d9b535000539e523673459b" }
+official_golden = { metadata = "tests/goldens/ankh_large.json=sha256:8f269215ee29887661abb6c139139dd1782f5112b94eb63be968282a7f050e1b", tensors = "tests/goldens/ankh_large.safetensors=sha256:4a7aea66e91cf880b0b839495d7704831725b672fe169f8e7258b5643d77aff9" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
 artifact_source = "official"
 canonical_state_sha256 = "e498a2e9aea76ef784cbe3e596c6b3f5e9a40e209ad837f7e3207099e4d74483"
@@ -1017,7 +1068,7 @@ id = "ankh2_large"
 family = "ankh"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh2_large.json=sha256:e8df38994ca1a1e0c598ace34a0b257b264937e4fdbb01bc41544985116b02a4", tensors = "tests/goldens/ankh2_large.safetensors=sha256:25fe1569f55c635fab8fa49c1d62a889a35a2a738bad921f5764a85b58fd4b5d" }
+official_golden = { metadata = "tests/goldens/ankh2_large.json=sha256:bdafcd179bc055ea223de43228801b1b6c35196741ac4cc4cca0571e3291e6dc", tensors = "tests/goldens/ankh2_large.safetensors=sha256:44d5fc5c74a0d166ceac54d12505a109f8338cb7090af2f37e39bd78545a3c52" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
 artifact_source = "official"
 canonical_state_sha256 = "597c4fe2fa8711f11a25317905f1d62fa92905e55fdd5c0a79614cd9c9d2bca3"
@@ -1049,7 +1100,7 @@ id = "ankh3_large"
 family = "ankh"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh3_large.json=sha256:2e5bb05b3baa5baa78f61fef7d2a2c669b0da5dbfaf6b50b12abd3e17253a961", tensors = "tests/goldens/ankh3_large.safetensors=sha256:e5c494ac418e0a2fe7bdad1376676d48960d58ec9e044d19bfffccb8c3288513" }
+official_golden = { metadata = "tests/goldens/ankh3_large.json=sha256:ab3a671260bdad481635d4a4be1b8072d5a2e02e3df1178202713c9a7e56af49", tensors = "tests/goldens/ankh3_large.safetensors=sha256:cc55274281e87d22cf7828117678df2e37a1ece543f5b401b65a0cd6f8b6fa27" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
 artifact_source = "official"
 canonical_state_sha256 = "60acb7ef86e85dc0c51fc1edf4c8e69a0480049723b6b2c95e6e9faa720c112a"
@@ -1083,7 +1134,7 @@ id = "ankh3_xl"
 family = "ankh"
 size_category = "xlarge"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh3_xl.json=sha256:66bb12e033e4163be225d636108a479393228a4f5061015c8af114e766c3c486", tensors = "tests/goldens/ankh3_xl.safetensors=sha256:72d34567d0228cb6f1ee701c578ed4039fead4346e3f161a52e0e74df28dc8ae" }
+official_golden = { metadata = "tests/goldens/ankh3_xl.json=sha256:a2d302227ee616698af18502d03e2b6b136d589134058d82b018520f18498431", tensors = "tests/goldens/ankh3_xl.safetensors=sha256:d18c086a4849b220761019f219590fe0c1c6f17698c365ccb0a5a8804446286c" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head. The official PyTorch shard index is deliberately excluded: the builder verifies every declared source shard directly and writes a new canonical safetensors index."
 artifact_source = "official"
 canonical_state_sha256 = "dd2188e0d2ca65232135714eef6de394239734d843ddae4928c7398685d858e7"
@@ -1176,7 +1227,7 @@ family = "esmfold2"
 size_category = "structure"
 generation_contract = "not_applicable"
 msa_conditioning = true
-official_golden = { metadata = "tests/goldens/esmfold2.json=sha256:f6e0ed1ec400b9a0fcc817db51774be968dc454b7a32645a07c479e42423ab20", tensors = "tests/goldens/esmfold2.safetensors=sha256:e4d6be4344c528e26b13f79a9303549e3de7e582da195c0078db3ce957fad420" }
+official_golden = { metadata = "tests/goldens/esmfold2.json=sha256:60b9e8827615a73acd89ac23f13e3d0a412c17879fbf3308444546e6bbb4ff03", tensors = "tests/goldens/esmfold2.safetensors=sha256:a1a9f3b7a9f7e36ef6a7077d76fa1b1c8b5529ed42ac2cfa485797162ce866fb" }
 fast_repo = "Synthyra/ESMFold2"
 fast_revision = "cd5a0927cec585a778d983b99a8db23d2e9b281e"
 fast_files = [
@@ -1196,7 +1247,7 @@ family = "esmfold2"
 size_category = "structure"
 generation_contract = "not_applicable"
 msa_conditioning = false
-official_golden = { metadata = "tests/goldens/esmfold2_fast.json=sha256:091b004c0b330217b59c12acd6da3d6edaf91e48d95f6d5f40fc20399cef9478", tensors = "tests/goldens/esmfold2_fast.safetensors=sha256:6e2e1cd07401538b4d9df994f82abe7a5b38a01e8d1ee26681e1216d44a81990" }
+official_golden = { metadata = "tests/goldens/esmfold2_fast.json=sha256:659bb338be4584787faff619d1cbcb8766aece7b5facf8a84470e850299b662f", tensors = "tests/goldens/esmfold2_fast.safetensors=sha256:8da91024a2cc63984585267a18f595892ebe489b2a05c4f79a994c3ac5de2f40" }
 fast_repo = "Synthyra/ESMFold2-Fast"
 fast_revision = "407875bfcaa42552bfcb25acd67ee1888b790170"
 fast_files = [
@@ -1216,7 +1267,7 @@ family = "esmfold2"
 size_category = "structure"
 generation_contract = "not_applicable"
 msa_conditioning = true
-official_golden = { metadata = "tests/goldens/esmfold2_experimental_cutoff2025.json=sha256:cfd0e35b2bc468a0dc4f614d3acfa2fce004f96e9ae2433256ed095b829d55cc", tensors = "tests/goldens/esmfold2_experimental_cutoff2025.safetensors=sha256:9347466bbe803b6f5dc82e3356ca6cbbf2c2edd8765f9fd273385bda255019f6" }
+official_golden = { metadata = "tests/goldens/esmfold2_experimental_cutoff2025.json=sha256:2e42d958f7a99ca8edaf0d9fac8e1a58b78200df652116b7bc102ac799123633", tensors = "tests/goldens/esmfold2_experimental_cutoff2025.safetensors=sha256:8168f8692c4932e1160e733b15210836120223615f6190fb0ecfa2f940c2420f" }
 fast_repo = "Synthyra/ESMFold2-Experimental-Cutoff2025"
 fast_revision = "632ff4a9e68f1de78ee956a613267bdcdb5b354d"
 fast_files = [
@@ -1237,7 +1288,7 @@ family = "esmfold2"
 size_category = "structure"
 generation_contract = "not_applicable"
 msa_conditioning = false
-official_golden = { metadata = "tests/goldens/esmfold2_experimental_fast_cutoff2025.json=sha256:1d0b2da4f1579243f37ae04bd4b834b747005cd8e8e7665e00d088123c43afd9", tensors = "tests/goldens/esmfold2_experimental_fast_cutoff2025.safetensors=sha256:516e216d05d7e6bee59e77126d3e595e2bb7821929433f00c259c5d5241964bb" }
+official_golden = { metadata = "tests/goldens/esmfold2_experimental_fast_cutoff2025.json=sha256:d41c54ca28270fc1c1b95485c5009990bd2c215d35ec5cc2c2fa611d096bad21", tensors = "tests/goldens/esmfold2_experimental_fast_cutoff2025.safetensors=sha256:97b58525aefd21c23ec6cb2defa3b3b41982c872dfe70d94306648edbe40a5b9" }
 fast_repo = "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025"
 fast_revision = "8f022c2514a6c32692aaca078a8391d6bc6c4bac"
 fast_files = [
@@ -1254,6 +1305,7 @@ auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFo
 
 [[models]]
 id = "esmfold2_300"
+official_golden = { metadata = "tests/goldens/esmfold2_300.json=sha256:d4fd12f27352c53582bb40a8d85eab76d6d98f647ce7660a6f891a1dfe68c039", tensors = "tests/goldens/esmfold2_300.safetensors=sha256:40a954ff14edc7c2bff95a252241f08db4cae3821ca9b92ae216775cceed8e8b" }
 confidence_adaptation = { release = "v1", head_sha256 = "40fd7f3d82fcefe8ad20ab2b32a37a68a84b54a527a4bce6eb9437bad2b77e31", base_weight_sha256 = "44d6797c5efebf24753d502b40950e0874871c96ceea14f2d7f7e39cebac67fd", donor_repo = "biohub/ESMFold2-Experimental-Fast-Cutoff2025", donor_revision = "74b88548bf19688b8727432db0d698cb2e1d8783", donor_weight_sha256 = "4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f", training_url = "https://wandb.ai/lhallee/fastplms-confidence/runs/9558b6d23daf", evaluation_url = "https://huggingface.co/datasets/Synthyra/FastPLMs-artifacts/tree/309f353b07e0e46de4d77a5266d4eddd695538e3/confidence-v2/v2-reproduction-20260922/public/evaluation/esmfold2_300", evidence_path = "docs/evidence/confidence/esmfold2_300-v1.json", frozen_base = { repo = "Synthyra/ESMFold2-300", revision = "a38a62ae930d157484b331c2bf4241684573adba", files = ["config.json=git-sha1:47ec20cf8b234c3b41d6f3ae1bdfe95d4eb4849e", "model.safetensors=sha256:44d6797c5efebf24753d502b40950e0874871c96ceea14f2d7f7e39cebac67fd"] } }
 family = "esmfold2"
 size_category = "structure"
@@ -1273,6 +1325,7 @@ backbone = { repo = "biohub/ESMC-300M-1500000", revision = "56803b6378b82e16c3b2
 
 [[models]]
 id = "esmfold2_600"
+official_golden = { metadata = "tests/goldens/esmfold2_600.json=sha256:ea249ff118975f29979143ea081cfa1c0fb60cb1fa393dd99a12caf413524230", tensors = "tests/goldens/esmfold2_600.safetensors=sha256:69bbb8d53a816b29e979e4469728df0a6b662e4d0be81a76068b36974d2b909e" }
 confidence_adaptation = { release = "v1", head_sha256 = "e84726a050722e1b722712c87d17a5388bd3699e2520e4a59abb1d828dfb8de7", base_weight_sha256 = "11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602", donor_repo = "biohub/ESMFold2-Experimental-Fast-Cutoff2025", donor_revision = "74b88548bf19688b8727432db0d698cb2e1d8783", donor_weight_sha256 = "4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f", training_url = "https://wandb.ai/lhallee/fastplms-confidence/runs/820d2cfa56c0", evaluation_url = "https://huggingface.co/datasets/Synthyra/FastPLMs-artifacts/tree/6e62186cd36b9047cc4691980076be9f76482192/confidence-v2/v2-reproduction-20260922/public/evaluation/esmfold2_600", evidence_path = "docs/evidence/confidence/esmfold2_600-v1.json", frozen_base = { repo = "Synthyra/ESMFold2-600", revision = "71c67d0b2b73dc245ea7c3cc0d0476439a882d08", files = ["config.json=git-sha1:8e271837cbdada96c4974c8e543f84065e0f06f1", "model.safetensors=sha256:11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602"] } }
 family = "esmfold2"
 size_category = "structure"
@@ -1289,3 +1342,277 @@ notes = "Experimental Fast model with a frozen 600M ESM++ backbone, 24 folding b
 auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" }
 backbone_model = "esmc_large"
 backbone = { repo = "biohub/ESMC-600M-1500000", revision = "21af9cc429af76ebda6c48074fb624db4735aaaf", files = ["config.json=git-sha1:ec29f6009b21d710f64bf1c058f3a9710833d692", "model.safetensors=sha256:d6869f5ae0f11e5dc829b195e062e87cfcc2f851a08a5edbaf5d1083ae7f76cc", "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", "tokenizer_config.json=git-sha1:f49f57b24a1c93bd544974811e8ecbd61b7fae89", "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b"] }
+
+# Parity goldens, uploaded 2026-09-22 so the reference tensors live somewhere durable
+# rather than only inside the GitHub repository. Each sha256 is the one already recorded
+# in that model's official_golden.tensors, so the existing integrity check is unchanged.
+
+[[golden_artifacts]]
+id = "esm2_8m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_8m.safetensors"
+sha256 = "d08a7572cbef20b8b19b545bcb0427b7e9ae19b986d015c559cd5a3a2cfc8aa4"
+size = 537589
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm2_35m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_35m.safetensors"
+sha256 = "4f82d10286e16041c2f23365dfdd8508b911864633287dd76580425527c2d922"
+size = 779509
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm2_150m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_150m.safetensors"
+sha256 = "20ad986b8e6e09f36d0158f5939d1914e0809b84329695028d767aab93338621"
+size = 1021437
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm2_650m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_650m.safetensors"
+sha256 = "261a8b71c4b7c90b1f294031558ee0bf00d77a98352021eb78771200c8d40548"
+size = 1989117
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm2_3b"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_3b.safetensors"
+sha256 = "61080e55e07db4a562a19e9b6b662e71a9fd3d8fee357aa3df00697d2fa58e33"
+size = 3924485
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmc_small"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmc_small.safetensors"
+sha256 = "98219356e2845cbd10a80715b19c95632956d40567a7a993409c426b56040b06"
+size = 1165354
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmc_large"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmc_large.safetensors"
+sha256 = "2b1ce3de5a7a27f171055f5c7aec4e7e50f0e809a1dc46508aaa0b472365015d"
+size = 1383082
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmc_6b"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmc_6b.safetensors"
+sha256 = "46139de6b28ab4f9e3244fb6ef3f259fc5e6518a3825f0fb4fd4a4fece6da8e6"
+size = 2979762
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm3_small"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm3_small.safetensors"
+sha256 = "251e050926a1f2426401bac6b93b4cb00041d0b0a5c8a458d0ef0e6a5d0d87c3"
+size = 2398877
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "e1_150m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/e1_150m.safetensors"
+sha256 = "03a2e93e7b3e54b12f99eea7678ff80828922f92c3683d4909346a05d74815bc"
+size = 958851
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "e1_300m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/e1_300m.safetensors"
+sha256 = "d125ea866899d2788330d6877d77c48c69febbf7f01d629f4107f9647cfa13ff"
+size = 1258379
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "e1_600m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/e1_600m.safetensors"
+sha256 = "d470d6833f0b38bec070aa37e9f3c49717ea82d31c794af843550bb7e56253a6"
+size = 1557899
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm_150m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm_150m.safetensors"
+sha256 = "f8e4ddd580d4708dad2d7da002735b2a01d69ece478c46c4932eb2f3933dc836"
+size = 1021437
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm_650m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm_650m.safetensors"
+sha256 = "70e2c46ed722942d3b628b556b92c9994688cb4853b278eb1e40a3dfbbd9dfc4"
+size = 1989117
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm_3b"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm_3b.safetensors"
+sha256 = "6942218079232a8185b1ae6e1578474c7784722e28f528b76ec89ca77777e2b9"
+size = 3924485
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm2_150m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm2_150m.safetensors"
+sha256 = "7390559bc54f27f08972c8f0adcf17954065450d937b3984303c1ae1ec003855"
+size = 26826946
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm2_650m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm2_650m.safetensors"
+sha256 = "c45f4ca267183b11dd06244ed3c3f39cf5b6139f778733fa63a6e2b5a46b0678"
+size = 28762314
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm2_3b"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm2_3b.safetensors"
+sha256 = "7490f771f8b8f8137b333f98d6b7dbd2cb48708c8649d4459279a6711b9e5870"
+size = 32633034
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh_base"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh_base.safetensors"
+sha256 = "449db3638117d5b6027380fdf4b05698bb156024b2b3f7e0ceb295a86e5ec756"
+size = 860722
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh_large"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh_large.safetensors"
+sha256 = "4a7aea66e91cf880b0b839495d7704831725b672fe169f8e7258b5643d77aff9"
+size = 1717818
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh2_large"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh2_large.safetensors"
+sha256 = "44d5fc5c74a0d166ceac54d12505a109f8338cb7090af2f37e39bd78545a3c52"
+size = 1717818
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh3_large"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh3_large.safetensors"
+sha256 = "cc55274281e87d22cf7828117678df2e37a1ece543f5b401b65a0cd6f8b6fa27"
+size = 1717818
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh3_xl"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh3_xl.safetensors"
+sha256 = "d18c086a4849b220761019f219590fe0c1c6f17698c365ccb0a5a8804446286c"
+size = 2860602
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold.safetensors"
+sha256 = "873b1b325a43d8e0f35f355c8914a2a9fe611cc48763875e9e6a22e09ec9ebcb"
+size = 179144
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2.safetensors"
+sha256 = "a1a9f3b7a9f7e36ef6a7077d76fa1b1c8b5529ed42ac2cfa485797162ce866fb"
+size = 2573976
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_fast"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_fast.safetensors"
+sha256 = "8da91024a2cc63984585267a18f595892ebe489b2a05c4f79a994c3ac5de2f40"
+size = 2573976
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_experimental_cutoff2025"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_experimental_cutoff2025.safetensors"
+sha256 = "8168f8692c4932e1160e733b15210836120223615f6190fb0ecfa2f940c2420f"
+size = 1771064
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_experimental_fast_cutoff2025"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_experimental_fast_cutoff2025.safetensors"
+sha256 = "97b58525aefd21c23ec6cb2defa3b3b41982c872dfe70d94306648edbe40a5b9"
+size = 1771064
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_300"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_300.safetensors"
+sha256 = "40a954ff14edc7c2bff95a252241f08db4cae3821ca9b92ae216775cceed8e8b"
+size = 865376
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_600"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_600.safetensors"
+sha256 = "69bbb8d53a816b29e979e4469728df0a6b662e4d0be81a76068b36974d2b909e"
+size = 865376
+offline_behavior = "requires_cached_verified_file"

fastplms/models/boltz/__init__.py

@@ -1,16 +1,16 @@
-"""Transformers-compatible Boltz2 configuration, model, and structure outputs."""
-
-from fastplms.models.boltz.modeling_boltz2 import (
-    Boltz2Config,
-    Boltz2Model,
-    Boltz2ModelOutput,
-    Boltz2StructureOutput,
-)
-
-
-__all__ = [
-    "Boltz2Config",
-    "Boltz2Model",
-    "Boltz2ModelOutput",
-    "Boltz2StructureOutput",
-]
+"""Transformers-compatible Boltz2 configuration, model, and structure outputs."""
+
+from fastplms.models.boltz.modeling_boltz2 import (
+    Boltz2Config,
+    Boltz2Model,
+    Boltz2ModelOutput,
+    Boltz2StructureOutput,
+)
+
+
+__all__ = [
+    "Boltz2Config",
+    "Boltz2Model",
+    "Boltz2ModelOutput",
+    "Boltz2StructureOutput",
+]

fastplms/models/boltz/_pair_attention.py

@@ -1,51 +1,51 @@
-"""Shared tensor operations for Boltz2 pair-biased attention layers."""
-
-from __future__ import annotations
-
-import math
-import torch
-
-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)``."""
-
-    # 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)
-
-
-def pair_biased_attention(
-    query_states: Tensor,
-    key_states: Tensor,
-    value_states: Tensor,
-    pair_bias: Tensor,
-    key_mask: Tensor,
-    mask_value: float,
-) -> Tensor:
-    """Return pair-biased attention values with shape ``(b, l_q, h, d_head)``."""
-
-    # 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):
-        scores = torch.einsum(
-            "bihd,bjhd->bhij",
-            query_states.float(),
-            key_states.float(),
-        )  # (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(),
-        )  # (b, l_q, h, d_h)
-    return output.to(value_states.dtype)  # (b, l_q, h, d_h)
+"""Shared tensor operations for Boltz2 pair-biased attention layers."""
+
+from __future__ import annotations
+
+import math
+import torch
+
+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)``."""
+
+    # 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)
+
+
+def pair_biased_attention(
+    query_states: Tensor,
+    key_states: Tensor,
+    value_states: Tensor,
+    pair_bias: Tensor,
+    key_mask: Tensor,
+    mask_value: float,
+) -> Tensor:
+    """Return pair-biased attention values with shape ``(b, l_q, h, d_head)``."""
+
+    # 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):
+        scores = torch.einsum(
+            "bihd,bjhd->bhij",
+            query_states.float(),
+            key_states.float(),
+        )  # (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(),
+        )  # (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,150 +1,150 @@
-"""Write Boltz2 coordinates and confidence values as an mmCIF atom table."""
-
-import numpy as np
-import torch
-
-from pathlib import Path
-
-from .minimal_structures import ProteinStructureTemplate
-
-
-def _confidence_per_atom(
-    plddt: torch.Tensor | None,
-    atom_to_residue: list[int],
-    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  # (n_a,)
-
-    values = plddt.detach().cpu()  # (n_i,) or (n_s, n_i)
-    if values.ndim == 1:
-        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]  # (n_i,)
-    if selected.shape[0] == num_atoms:
-        return (selected.numpy() * 100.0).astype(np.float32)  # (n_a,)
-
-    num_residues = max(atom_to_residue) + 1  # n_r
-    if selected.shape[0] == num_residues:
-        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  # (n_a,)
-
-    raise ValueError(
-        "pLDDT item count must match either atoms or residues: "
-        f"received {selected.shape[0]}, expected {num_atoms} atoms or "
-        f"{num_residues} residues."
-    )
-
-
-def write_cif(
-    structure_template: ProteinStructureTemplate,
-    atom_coords: torch.Tensor,
-    atom_mask: torch.Tensor,
-    output_path: str,
-    plddt: torch.Tensor | None = None,
-    sample_index: int = 0,
-) -> str:
-    # 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)  # (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]  # (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()  # (n_a, 3)
-
-    mask = atom_mask.detach().cpu()  # (n_a,) or (n_m, n_a)
-    if mask.ndim == 2:
-        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. "
-            f"Non-finite count: {int(valid_non_finite.sum().item())}"
-        )
-
-    b_iso = _confidence_per_atom(
-        plddt=plddt,
-        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.")
-
-    lines = [
-        "data_boltz2_prediction",
-        "#",
-        "loop_",
-        "_atom_site.group_PDB",
-        "_atom_site.id",
-        "_atom_site.type_symbol",
-        "_atom_site.label_atom_id",
-        "_atom_site.label_comp_id",
-        "_atom_site.label_asym_id",
-        "_atom_site.label_seq_id",
-        "_atom_site.Cartn_x",
-        "_atom_site.Cartn_y",
-        "_atom_site.Cartn_z",
-        "_atom_site.occupancy",
-        "_atom_site.B_iso_or_equiv",
-        "_atom_site.pdbx_PDB_model_num",
-    ]
-
-    atom_id = 1
-    for idx in range(structure_template.num_atoms):
-        if mask[idx] <= 0:
-            continue
-
-        residue_idx = structure_template.atom_residue_index[idx]
-        residue_name = structure_template.residue_names[residue_idx]
-        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()  # (3,) -> three scalars
-        b_factor = float(b_iso[idx])  # () -> Python float
-
-        line = (
-            f"ATOM {atom_id} {element} {atom_name} {residue_name} {chain_id} "
-            f"{residue_idx + 1} {x_val:.3f} {y_val:.3f} {z_val:.3f} 1.00 {b_factor:.2f} 1"
-        )
-        lines.append(line)
-        atom_id += 1
-
-    lines.append("#")
-    text = "\n".join(lines) + "\n"
-
-    out_path = Path(output_path)
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    out_path.write_text(text, encoding="utf-8")
-    return str(out_path)
+"""Write Boltz2 coordinates and confidence values as an mmCIF atom table."""
+
+import numpy as np
+import torch
+
+from pathlib import Path
+
+from .minimal_structures import ProteinStructureTemplate
+
+
+def _confidence_per_atom(
+    plddt: torch.Tensor | None,
+    atom_to_residue: list[int],
+    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  # (n_a,)
+
+    values = plddt.detach().cpu()  # (n_i,) or (n_s, n_i)
+    if values.ndim == 1:
+        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]  # (n_i,)
+    if selected.shape[0] == num_atoms:
+        return (selected.numpy() * 100.0).astype(np.float32)  # (n_a,)
+
+    num_residues = max(atom_to_residue) + 1  # n_r
+    if selected.shape[0] == num_residues:
+        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  # (n_a,)
+
+    raise ValueError(
+        "pLDDT item count must match either atoms or residues: "
+        f"received {selected.shape[0]}, expected {num_atoms} atoms or "
+        f"{num_residues} residues."
+    )
+
+
+def write_cif(
+    structure_template: ProteinStructureTemplate,
+    atom_coords: torch.Tensor,
+    atom_mask: torch.Tensor,
+    output_path: str,
+    plddt: torch.Tensor | None = None,
+    sample_index: int = 0,
+) -> str:
+    # 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)  # (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]  # (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()  # (n_a, 3)
+
+    mask = atom_mask.detach().cpu()  # (n_a,) or (n_m, n_a)
+    if mask.ndim == 2:
+        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. "
+            f"Non-finite count: {int(valid_non_finite.sum().item())}"
+        )
+
+    b_iso = _confidence_per_atom(
+        plddt=plddt,
+        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.")
+
+    lines = [
+        "data_boltz2_prediction",
+        "#",
+        "loop_",
+        "_atom_site.group_PDB",
+        "_atom_site.id",
+        "_atom_site.type_symbol",
+        "_atom_site.label_atom_id",
+        "_atom_site.label_comp_id",
+        "_atom_site.label_asym_id",
+        "_atom_site.label_seq_id",
+        "_atom_site.Cartn_x",
+        "_atom_site.Cartn_y",
+        "_atom_site.Cartn_z",
+        "_atom_site.occupancy",
+        "_atom_site.B_iso_or_equiv",
+        "_atom_site.pdbx_PDB_model_num",
+    ]
+
+    atom_id = 1
+    for idx in range(structure_template.num_atoms):
+        if mask[idx] <= 0:
+            continue
+
+        residue_idx = structure_template.atom_residue_index[idx]
+        residue_name = structure_template.residue_names[residue_idx]
+        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()  # (3,) -> three scalars
+        b_factor = float(b_iso[idx])  # () -> Python float
+
+        line = (
+            f"ATOM {atom_id} {element} {atom_name} {residue_name} {chain_id} "
+            f"{residue_idx + 1} {x_val:.3f} {y_val:.3f} {z_val:.3f} 1.00 {b_factor:.2f} 1"
+        )
+        lines.append(line)
+        atom_id += 1
+
+    lines.append("#")
+    text = "\n".join(lines) + "\n"
+
+    out_path = Path(output_path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text(text, encoding="utf-8")
+    return str(out_path)

fastplms/models/boltz/minimal_featurizer.py

@@ -1,853 +1,853 @@
-"""Construct Boltz2 protein features from canonical atom tables and conformers."""
-
-import math
-import numpy as np
-import torch
-
-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,
-    "N": 7,
-    "O": 8,
-    "P": 15,
-    "S": 16,
-}
-
-# Canonical formal charges and tetrahedral chirality values extracted from the
-# pinned Boltz2 molecule archive. Values use Boltz2's public feature encoding.
-_FORMAL_CHARGES = {
-    ("ARG", "NH1"): 1.0,
-    ("HIS", "ND1"): 1.0,
-    ("LYS", "NZ"): 1.0,
-}
-_CHIRAL_ATOMS = {
-    "ALA": frozenset({"CA"}),
-    "ARG": frozenset({"CA"}),
-    "ASN": frozenset({"CA"}),
-    "ASP": frozenset({"CA"}),
-    "CYS": frozenset({"CA"}),
-    "GLN": frozenset({"CA"}),
-    "GLU": frozenset({"CA"}),
-    "HIS": frozenset({"CA"}),
-    "ILE": frozenset({"CA", "CB"}),
-    "LEU": frozenset({"CA"}),
-    "LYS": frozenset({"CA"}),
-    "MET": frozenset({"CA"}),
-    "PHE": frozenset({"CA"}),
-    "PRO": frozenset({"CA"}),
-    "SER": frozenset({"CA"}),
-    "THR": frozenset({"CA", "CB"}),
-    "TRP": frozenset({"CA"}),
-    "TYR": frozenset({"CA"}),
-    "UNK": frozenset({"CA"}),
-    "VAL": frozenset({"CA"}),
-}
-
-
-def _normalize_sequence(sequence: str) -> str:
-    if not isinstance(sequence, str):
-        raise TypeError("Amino acid sequence must be a string.")
-    seq = sequence.strip().upper()
-    if not seq:
-        raise ValueError("Amino acid sequence must be non-empty.")
-    for aa in seq:
-        if aa not in const.prot_letter_to_token:
-            raise ValueError(f"Unsupported residue code '{aa}'.")
-    return seq
-
-
-def _atom_name_to_element(atom_name: str) -> str:
-    name = atom_name.strip().upper()
-    if len(name) == 0:
-        return "C"
-    if name[0].isdigit():
-        name = name[1:]
-    if len(name) >= 2 and name[0:2] in ("CL", "BR", "FE", "MG", "ZN", "NA", "CA"):
-        return name[0]
-    return name[0]
-
-
-def _atom_name_to_codes(atom_name: str) -> torch.Tensor:
-    clipped = atom_name.strip()[:4]
-    vals = [ord(ch) - 32 for ch in clipped]
-    while len(vals) < 4:
-        vals.append(0)
-    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
-
-
-# Raw first-conformer coordinates extracted at full float32 precision from the
-# hash-pinned Boltz2 molecule archive. The feature builder applies the official
-# seeded centering, rotation, and translation policy below.
-_RDKIT_CONFORMERS: dict[str, dict[str, list[float]]] = {
-    "ALA": {
-        "N": [-0.9241582155227661, 1.1821246147155762, 0.712748110294342],
-        "CA": [-0.2663755416870117, -0.08827890455722809, 0.4008508622646332],
-        "C": [1.1188693046569824, 0.1387452930212021, -0.14366498589515686],
-        "O": [1.2882297039031982, 0.8058497309684753, -1.2001261711120605],
-        "CB": [-1.1134333610534668, -0.8915302753448486, -0.5876986384391785],
-    },
-    "ARG": {
-        "N": [3.253326416015625, -1.699564814567566, -0.9852627515792847],
-        "CA": [2.191511392593384, -0.692007839679718, -0.9506462812423706],
-        "C": [2.80021333694458, 0.672307550907135, -1.1160222291946411],
-        "O": [2.465440034866333, 1.3987600803375244, -2.0899009704589844],
-        "CB": [1.3811039924621582, -0.7949594855308533, 0.36316436529159546],
-        "CG": [0.20344306528568268, 0.19287171959877014, 0.457282155752182],
-        "CD": [-0.902410626411438, -0.09953747689723969, -0.5671625137329102],
-        "NE": [-2.0867979526519775, 0.7146661281585693, -0.2911669909954071],
-        "CZ": [-3.018158197402954, 0.4804283082485199, 0.7831193804740906],
-        "NH1": [-4.080137729644775, 1.4154621362686157, 0.9938872456550598],
-        "NH2": [-2.9226150512695312, -0.5637124180793762, 1.5533623695373535],
-    },
-    "ASN": {
-        "N": [-1.5767980813980103, -1.7227835655212402, 0.2427234649658203],
-        "CA": [-0.7041534781455994, -0.5622008442878723, 0.42767149209976196],
-        "C": [-1.246688961982727, 0.6072975993156433, -0.3483567237854004],
-        "O": [-1.554918885231018, 0.47644123435020447, -1.5638835430145264],
-        "CB": [0.7263866662979126, -0.9032993316650391, -0.028588544577360153],
-        "CG": [1.6871562004089355, 0.20810098946094513, 0.28412675857543945],
-        "OD1": [2.19687557220459, 0.28916797041893005, 1.433493971824646],
-        "ND2": [1.996717095375061, 1.196706771850586, -0.6981719136238098],
-    },
-    "ASP": {
-        "N": [-0.11774874478578568, -1.6310220956802368, 0.374360054731369],
-        "CA": [-0.3847021758556366, -0.19290897250175476, 0.28121331334114075],
-        "C": [-1.806337594985962, 0.059408850967884064, -0.1472136378288269],
-        "O": [-2.2237648963928223, -0.36817988753318787, -1.257502555847168],
-        "CB": [0.5981687307357788, 0.4785478413105011, -0.6920451521873474],
-        "CG": [2.0088446140289307, 0.36018264293670654, -0.20304732024669647],
-        "OD1": [2.747267961502075, -0.5711832046508789, -0.6216948628425598],
-        "OD2": [2.49090313911438, 1.2577818632125854, 0.7443997859954834],
-    },
-    "CYS": {
-        "N": [-0.058702241629362106, 1.771048665046692, 0.2434082329273224],
-        "CA": [-0.06546340137720108, 0.47663936018943787, -0.4456723928451538],
-        "C": [-1.273919701576233, -0.3504071533679962, -0.08361759781837463],
-        "O": [-1.7324607372283936, -0.34313708543777466, 1.0909277200698853],
-        "CB": [1.2432398796081543, -0.28917407989501953, -0.19254730641841888],
-        "SG": [1.484184980392456, -0.7191365361213684, 1.5649693012237549],
-    },
-    "GLN": {
-        "N": [-1.8543237447738647, -1.0024770498275757, -1.6278940439224243],
-        "CA": [-1.292184591293335, -0.6786512732505798, -0.3153715133666992],
-        "C": [-2.2264492511749268, 0.25171151757240295, 0.4067918062210083],
-        "O": [-2.7255098819732666, -0.08834805339574814, 1.5129321813583374],
-        "CB": [0.11662524193525314, -0.06213820353150368, -0.45663803815841675],
-        "CG": [0.8140791058540344, 0.10461423546075821, 0.9002283811569214],
-        "CD": [2.195122003555298, 0.6569428443908691, 0.7154991030693054],
-        "OE1": [2.3870291709899902, 1.9003654718399048, 0.7825236916542053],
-        "NE2": [3.287371873855591, -0.21307705342769623, 0.41959717869758606],
-    },
-    "GLU": {
-        "N": [-1.3492857217788696, -1.114271640777588, -1.3739068508148193],
-        "CA": [-1.2765676975250244, -0.506334662437439, -0.04287439212203026],
-        "C": [-1.9195666313171387, 0.8547031879425049, -0.04214814677834511],
-        "O": [-1.9091347455978394, 1.5701956748962402, -1.080504298210144],
-        "CB": [0.18272103369235992, -0.43209096789360046, 0.45213374495506287],
-        "CG": [1.1058456897735596, 0.3659400939941406, -0.481784850358963],
-        "CD": [2.5090205669403076, 0.35135790705680847, 0.03646872192621231],
-        "OE1": [2.896414041519165, 1.2445733547210693, 0.8365193605422974],
-        "OE2": [3.3727896213531494, -0.6754258275032043, -0.33141785860061646],
-    },
-    "GLY": {
-        "N": [-1.291549801826477, 0.6080796122550964, -0.4228580892086029],
-        "CA": [-0.4895951449871063, -0.2882237136363983, 0.40191715955734253],
-        "C": [0.9350062608718872, -0.2543502449989319, -0.04786944016814232],
-        "O": [1.3473576307296753, -1.0836502313613892, -0.9022219777107239],
-    },
-    "HIS": {
-        "N": [1.0371263027191162, -1.5621215105056763, 0.4178937077522278],
-        "CA": [1.1980191469192505, -0.40241551399230957, -0.46312251687049866],
-        "C": [2.6561269760131836, -0.07123222947120667, -0.6407948136329651],
-        "O": [3.1795547008514404, -0.12371637672185898, -1.7859554290771484],
-        "CB": [0.44051393866539, 0.8116919994354248, 0.09850303083658218],
-        "CG": [-1.0356801748275757, 0.5519806742668152, 0.15254715085029602],
-        "ND1": [-1.9089158773422241, 0.6490958333015442, -0.968245267868042],
-        "CD2": [-1.7095766067504883, 0.1138235554099083, 1.2061808109283447],
-        "CE1": [-3.0751633644104004, 0.28129827976226807, -0.5511555075645447],
-        "NE2": [-3.0650453567504883, -0.07846951484680176, 0.8257263898849487],
-    },
-    "ILE": {
-        "N": [-1.2378733158111572, -1.8146690130233765, -0.16437499225139618],
-        "CA": [-1.2780015468597412, -0.3994033634662628, 0.22437156736850739],
-        "C": [-1.9879958629608154, 0.4497596025466919, -0.80372154712677],
-        "O": [-2.0291709899902344, 0.10288607329130173, -2.015406608581543],
-        "CB": [0.14225825667381287, 0.146738663315773, 0.5415176153182983],
-        "CG1": [1.1062071323394775, 0.04519597440958023, -0.6694174408912659],
-        "CG2": [0.7215384244918823, -0.5685611963272095, 1.7750537395477295],
-        "CD1": [2.371805191040039, 0.8852734565734863, -0.49045464396476746],
-    },
-    "LEU": {
-        "N": [1.6452248096466064, -1.017622470855713, -1.0626215934753418],
-        "CA": [1.3129304647445679, 0.08074887096881866, -0.15053576231002808],
-        "C": [2.431154727935791, 0.28310757875442505, 0.8357383608818054],
-        "O": [2.9231173992156982, -0.7013152837753296, 1.451209306716919],
-        "CB": [0.006050171796232462, -0.2141101360321045, 0.6198461651802063],
-        "CG": [-1.2534235715866089, -0.3772832751274109, -0.26637399196624756],
-        "CD1": [-2.455688714981079, -0.7604416012763977, 0.6070646643638611],
-        "CD2": [-1.5742294788360596, 0.8985711932182312, -1.059889554977417],
-    },
-    "LYS": {
-        "N": [-2.3938536643981934, -1.4751482009887695, -0.9336304068565369],
-        "CA": [-2.2094509601593018, -0.6203567981719971, 0.2421264797449112],
-        "C": [-3.4513463973999023, 0.1961250752210617, 0.4710204303264618],
-        "O": [-3.9688456058502197, 0.2503468096256256, 1.6188287734985352],
-        "CB": [-1.008726716041565, 0.3269270658493042, 0.0479682981967926],
-        "CG": [0.336457222700119, -0.41559070348739624, 0.047045741230249405],
-        "CD": [1.5088860988616943, 0.5716065168380737, -0.00632853340357542],
-        "CE": [2.853410005569458, -0.1678532361984253, -0.020898228511214256],
-        "NZ": [3.9690465927124023, 0.7765668630599976, -0.07007527351379395],
-    },
-    "MET": {
-        "N": [-1.743452548980713, -0.44039490818977356, 1.851780652999878],
-        "CA": [-1.2168668508529663, 0.4041755199432373, 0.7779996991157532],
-        "C": [-2.3118340969085693, 0.683431088924408, -0.2134547382593155],
-        "O": [-2.6660380363464355, 1.8705639839172363, -0.44476020336151123],
-        "CB": [-0.0005482881097123027, -0.262584388256073, 0.09933014959096909],
-        "CG": [0.6952047944068909, 0.6768104434013367, -0.8930339813232422],
-        "SD": [2.136805772781372, -0.14913176000118256, -1.6563159227371216],
-        "CE": [3.3404407501220703, 0.348065584897995, -0.38104408979415894],
-    },
-    "PHE": {
-        "N": [2.7976438999176025, -1.6209964752197266, -0.4779351055622101],
-        "CA": [1.5799994468688965, -0.8289614915847778, -0.6608061194419861],
-        "C": [1.9535551071166992, 0.5758142471313477, -1.0445665121078491],
-        "O": [1.5211102962493896, 1.072593331336975, -2.1190123558044434],
-        "CB": [0.7137221097946167, -0.8474348187446594, 0.6163292527198792],
-        "CG": [-0.6229045987129211, -0.18666480481624603, 0.3912639319896698],
-        "CD1": [-0.8259273767471313, 1.0978055000305176, 0.7381877303123474],
-        "CD2": [-1.7160141468048096, -0.9425657391548157, -0.2700555920600891],
-        "CE1": [-2.131312131881714, 1.741182565689087, 0.4800238013267517],
-        "CE2": [-2.8960046768188477, -0.3551671802997589, -0.505269467830658],
-        "CZ": [-3.11456298828125, 1.0511738061904907, -0.11009909957647324],
-    },
-    "PRO": {
-        "N": [-0.577339768409729, -0.5123927593231201, -1.150407314300537],
-        "CA": [0.5119893550872803, 0.23565764725208282, -0.5195383429527283],
-        "C": [1.8393173217773438, -0.46712079644203186, -0.6550371646881104],
-        "O": [2.0836586952209473, -1.166333556175232, -1.6751221418380737],
-        "CB": [0.1029478907585144, 0.45346882939338684, 0.9297074675559998],
-        "CG": [-1.4098080396652222, 0.5215330123901367, 0.8787456750869751],
-        "CD": [-1.7971025705337524, -0.06041542813181877, -0.47876954078674316],
-    },
-    "SER": {
-        "N": [0.8959873914718628, -1.423298716545105, -0.2692987024784088],
-        "CA": [-0.018787339329719543, -0.28751760721206665, -0.3863537311553955],
-        "C": [0.767615795135498, 0.9929561018943787, -0.42147380113601685],
-        "O": [0.5327092409133911, 1.8535151481628418, -1.3116248846054077],
-        "CB": [-0.9973848462104797, -0.2769756019115448, 0.7988247871398926],
-        "OG": [-1.8999865055084229, 0.7915746569633484, 0.691534698009491],
-    },
-    "THR": {
-        "N": [0.005045319441705942, 1.7655576467514038, 0.0639338344335556],
-        "CA": [0.4231603145599365, 0.3755212128162384, 0.2734847664833069],
-        "C": [1.7184523344039917, 0.11012311279773712, -0.4465082287788391],
-        "O": [1.9252793788909912, 0.5967074036598206, -1.5912789106369019],
-        "CB": [-0.6634359359741211, -0.615561306476593, -0.21173778176307678],
-        "OG1": [-1.0762907266616821, -0.29800114035606384, -1.5171915292739868],
-        "CG2": [-1.8869022130966187, -0.6176311373710632, 0.7084044218063354],
-    },
-    "TRP": {
-        "N": [-3.0439648628234863, 1.0628471374511719, 0.015692168846726418],
-        "CA": [-2.463679552078247, 0.024657616391777992, -0.8414031863212585],
-        "C": [-2.186711072921753, -1.2228965759277344, -0.04691409692168236],
-        "O": [-2.3364906311035156, -2.3535354137420654, -0.5830773115158081],
-        "CB": [-1.2062920331954956, 0.5195542573928833, -1.592403531074524],
-        "CG": [-0.06450498104095459, 0.884079098701477, -0.6835198402404785],
-        "CD1": [0.16109998524188995, 2.085163116455078, -0.14801783859729767],
-        "CD2": [0.9781660437583923, -0.005544683896005154, -0.17733941972255707],
-        "NE1": [1.3106837272644043, 2.0506503582000732, 0.698072075843811],
-        "CE2": [1.7528777122497559, 0.7044392228126526, 0.6196392774581909],
-        "CE3": [1.236093521118164, -1.4355946779251099, -0.4235689342021942],
-        "CZ2": [2.9070138931274414, 0.10196790844202042, 1.3070299625396729],
-        "CZ3": [2.2851555347442627, -2.007185697555542, 0.1949695348739624],
-        "CH2": [3.1585631370544434, -1.2034809589385986, 1.1002938747406006],
-    },
-    "TYR": {
-        "N": [-1.7950150966644287, 0.49119046330451965, -1.3951144218444824],
-        "CA": [-1.8436503410339355, -0.2694968581199646, -0.14338290691375732],
-        "C": [-3.240288019180298, -0.28215137124061584, 0.41974759101867676],
-        "O": [-3.8150112628936768, 0.7977103590965271, 0.7253130674362183],
-        "CB": [-0.8541494011878967, 0.30990728735923767, 0.8830214738845825],
-        "CG": [0.5672639012336731, 0.2125871181488037, 0.3916495442390442],
-        "CD1": [1.2694358825683594, -0.924761176109314, 0.5431669354438782],
-        "CD2": [1.1872504949569702, 1.3605931997299194, -0.31554359197616577],
-        "CE1": [2.6519992351531982, -1.0213464498519897, 0.029911600053310394],
-        "CE2": [2.4384164810180664, 1.2693941593170166, -0.7831515073776245],
-        "CZ": [3.211331367492676, 0.022552935406565666, -0.6021429300308228],
-        "OH": [4.513507843017578, -0.05260590463876724, -1.0929937362670898],
-    },
-    "VAL": {
-        "N": [0.9408224821090698, -1.2608877420425415, 0.652370810508728],
-        "CA": [0.7287879586219788, -0.3937721848487854, -0.5120788216590881],
-        "C": [1.7670493125915527, 0.6979415416717529, -0.5508439540863037],
-        "O": [2.2449657917022705, 1.1683598756790161, 0.5171348452568054],
-        "CB": [-0.7015937566757202, 0.21480272710323334, -0.5226467847824097],
-        "CG1": [-1.7662651538848877, -0.845522403717041, -0.8411717414855957],
-        "CG2": [-1.0500152111053467, 0.9374979138374329, 0.7892321944236755],
-    },
-    "UNK": {
-        "N": [1.6134154796600342, -1.304404377937317, -0.35241976380348206],
-        "CA": [0.6100443601608276, -0.23816797137260437, -0.3852081298828125],
-        "C": [1.0861350297927856, 0.9232651591300964, 0.4423142373561859],
-        "O": [1.374688982963562, 0.7629808187484741, 1.6591525077819824],
-        "CB": [-0.7607048153877258, -0.7552146315574646, 0.10266945511102676],
-    },
-}
-
-
-def _get_atom_position(res_name: str, atom_name: str, atom_idx: int) -> np.ndarray:
-    """Get the canonical RDKit conformer position for an atom.
-
-    Uses pre-extracted positions from official Boltz2 mol files. Falls back
-    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)  # (3,)
-    # Fallback for unknown atoms (should not happen for canonical AAs)
-    angle = (atom_idx + 1) * 0.7
-    radius = 1.4 + 0.03 * atom_idx
-    return np.array(
-        [
-            radius * math.cos(angle),
-            radius * math.sin(angle),
-            0.1 * ((atom_idx % 5) - 2),
-        ],
-        dtype=np.float32,
-    )  # (3,)
-
-
-def _build_template(
-    sequence: str,
-) -> tuple[
-    ProteinStructureTemplate,
-    list[str],
-    list[int],
-    list[int],
-    list[int],
-    list[np.ndarray],
-    list[int],
-]:
-    residue_names: list[str] = []
-    residue_token_ids: list[int] = []
-    atom_names: list[str] = []
-    atom_elements: list[str] = []
-    atom_residue_index: list[int] = []
-    atom_chain_id: list[str] = []
-    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] = []
-
-    global_atom_idx = 0
-    for res_idx, aa in enumerate(sequence):
-        token_name = const.prot_letter_to_token[aa]
-        residue_names.append(token_name)
-        residue_token_ids.append(const.token_ids[token_name])
-
-        residue_atoms = const.ref_atoms[token_name]
-        if not residue_atoms:
-            raise RuntimeError(f"No reference atoms for residue {token_name}.")
-        center_atom_name = const.res_to_center_atom[token_name]
-        disto_atom_name = const.res_to_disto_atom[token_name]
-
-        center_idx = -1
-        disto_idx = -1
-        n_idx = -1
-        ca_idx = -1
-        c_idx = -1
-
-        for local_idx, atom_name in enumerate(residue_atoms):
-            atom_names.append(atom_name)
-            element = _atom_name_to_element(atom_name)
-            atom_elements.append(element)
-            atom_residue_index.append(res_idx)
-            atom_chain_id.append("A")
-
-            atom_pos = _get_atom_position(token_name, atom_name, local_idx)  # (3,)
-            atom_positions.append(atom_pos)
-
-            if atom_name == center_atom_name:
-                center_idx = global_atom_idx
-            if atom_name == disto_atom_name:
-                disto_idx = global_atom_idx
-            if atom_name == "N":
-                n_idx = global_atom_idx
-            if atom_name == "CA":
-                ca_idx = global_atom_idx
-            if atom_name == "C":
-                c_idx = global_atom_idx
-            global_atom_idx += 1
-
-        if center_idx == -1:
-            center_idx = global_atom_idx - len(residue_atoms)
-        if disto_idx == -1:
-            disto_idx = center_idx
-        if n_idx == -1:
-            n_idx = center_idx
-        if ca_idx == -1:
-            ca_idx = center_idx
-        if c_idx == -1:
-            c_idx = center_idx
-
-        residue_center_atom_idx.append(center_idx)
-        residue_disto_atom_idx.append(disto_idx)
-        residue_frame_atom_idx.extend([n_idx, ca_idx, c_idx])
-
-    template = ProteinStructureTemplate(
-        sequence=sequence,
-        residue_names=residue_names,
-        atom_names=atom_names,
-        atom_elements=atom_elements,
-        atom_residue_index=atom_residue_index,
-        atom_chain_id=atom_chain_id,
-    )
-
-    return (
-        template,
-        residue_names,
-        residue_token_ids,
-        residue_center_atom_idx,
-        residue_disto_atom_idx,
-        atom_positions,
-        residue_frame_atom_idx,
-    )
-
-
-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)  # (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),
-            two_s * (i * j - k * real),
-            two_s * (i * k + j * real),
-            two_s * (i * j + k * real),
-            1 - two_s * (i * i + k * k),
-            two_s * (j * k - i * real),
-            two_s * (i * k - j * real),
-            two_s * (j * k + i * real),
-            1 - two_s * (i * i + j * j),
-        ),
-        dim=-1,
-    )  # (1, 9)
-    return rotation.reshape(1, 3, 3)[0]  # (3, 3)
-
-
-def _center_and_augment_atoms_per_residue(
-    atom_positions: torch.Tensor,
-    atom_residue_index: list[int],
-    num_residues: int,
-) -> torch.Tensor:
-    """Apply Boltz2's seeded conformer augmentation per residue.
-
-    The pinned implementation intentionally excludes the final residue because
-    it iterates to the maximum reference-space identifier rather than through it.
-    """
-    # 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  # (a,)
-        if not torch.any(residue_mask):
-            raise RuntimeError(f"Residue index {residue_idx} has no atoms.")
-        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)  # (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(
-    amino_acid_sequence: str,
-    num_bins: int = 64,
-    atoms_per_window_queries: int = 32,
-) -> tuple[dict[str, torch.Tensor], ProteinStructureTemplate]:
-    sequence = _normalize_sequence(amino_acid_sequence)
-    (
-        template,
-        residue_names,
-        residue_token_ids,
-        residue_center_atom_idx,
-        residue_disto_atom_idx,
-        atom_positions_np,
-        residue_frame_atom_idx_flat,
-    ) = _build_template(sequence)
-
-    num_tokens = len(residue_names)
-    num_atoms = len(atom_positions_np)
-    if num_tokens <= 0 or num_atoms <= 0:
-        raise RuntimeError("Boltz2 feature construction produced an empty protein template.")
-
-    # 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)  # (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)  # (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)  # (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)  # (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"]
-    contact_ids = torch.full(
-        (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)  # (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.")
-    method_feature = torch.full(
-        (1, num_tokens),
-        fill_value=const.method_types_ids["x-ray diffraction"],
-        dtype=torch.long,
-    )  # (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)  # (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,
-    )  # (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:
-        z_value = _ELEMENT_TO_Z[element] if element in _ELEMENT_TO_Z else _ELEMENT_TO_Z["C"]
-        if z_value >= const.num_elements:
-            raise RuntimeError(
-                f"Atomic number {z_value} exceeds the Boltz2 element vocabulary."
-            )
-        atomic_numbers.append(z_value)
-    ref_element = one_hot(
-        torch.tensor(atomic_numbers, dtype=torch.long),
-        num_classes=const.num_elements,
-    ).unsqueeze(0)  # (1, a, n_element)
-
-    ref_charge_values = []
-    ref_chirality_values = []
-    for atom_name, residue_idx in zip(
-        template.atom_names,
-        template.atom_residue_index,
-        strict=True,
-    ):
-        residue_name = residue_names[residue_idx]
-        ref_charge_values.append(_FORMAL_CHARGES.get((residue_name, atom_name), 0.0))
-        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)  # (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)  # (1, a, t)
-    token_to_rep_atom = one_hot(
-        torch.tensor(residue_disto_atom_idx, dtype=torch.long),
-        num_classes=num_atoms,
-    ).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)  # (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)
-    )
-    backbone_ids = []
-    for atom_name in template.atom_names:
-        if atom_name in const.protein_backbone_atom_index:
-            backbone_ids.append(const.protein_backbone_atom_index[atom_name] + 1)
-        else:
-            backbone_ids.append(0)
-    atom_backbone_feat = one_hot(
-        torch.tensor(backbone_ids, dtype=torch.long),
-        num_classes=num_backbone_classes,
-    ).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)  # (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)  # (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.")
-    pad_atoms = (
-        (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
-        )  # (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,
-        )  # (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,
-        )  # (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,
-        )  # (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)  # (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),
-            num_classes=const.num_tokens,
-        )
-        .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,
-    )  # (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,
-    )  # (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,
-        "residue_index": residue_index,
-        "asym_id": asym_id,
-        "entity_id": entity_id,
-        "sym_id": sym_id,
-        "mol_type": mol_type,
-        "res_type": res_type,
-        "disto_center": disto_center,
-        "token_bonds": token_bonds,
-        "type_bonds": type_bonds,
-        "token_pad_mask": token_pad_mask,
-        "token_resolved_mask": token_resolved_mask,
-        "token_disto_mask": token_disto_mask,
-        "contact_conditioning": contact_conditioning,
-        "contact_threshold": contact_threshold,
-        "method_feature": method_feature,
-        "modified": modified,
-        "cyclic_period": cyclic_period,
-        "affinity_token_mask": affinity_token_mask,
-        "ref_pos": ref_pos,
-        "atom_resolved_mask": atom_resolved_mask,
-        "ref_atom_name_chars": ref_atom_name_chars,
-        "ref_element": ref_element,
-        "ref_charge": ref_charge,
-        "ref_chirality": ref_chirality,
-        "atom_backbone_feat": atom_backbone_feat,
-        "ref_space_uid": ref_space_uid,
-        "coords": coords,
-        "atom_pad_mask": atom_pad_mask,
-        "atom_to_token": atom_to_token,
-        "token_to_rep_atom": token_to_rep_atom,
-        "r_set_to_rep_atom": r_set_to_rep_atom,
-        "token_to_center_atom": token_to_center_atom,
-        "disto_target": disto_target,
-        "disto_coords_ensemble": disto_coords_ensemble,
-        "bfactor": bfactor,
-        "plddt": atom_plddt,
-        "frames_idx": frames_idx,
-        "frame_resolved_mask": frame_resolved_mask,
-        "msa": msa,
-        "msa_paired": msa_paired,
-        "deletion_value": deletion_value,
-        "has_deletion": has_deletion,
-        "deletion_mean": deletion_mean,
-        "profile": profile,
-        "msa_mask": msa_mask,
-        "template_restype": template_restype,
-        "template_frame_rot": template_frame_rot,
-        "template_frame_t": template_frame_t,
-        "template_cb": template_cb,
-        "template_ca": template_ca,
-        "template_mask_cb": template_mask_cb,
-        "template_mask_frame": template_mask_frame,
-        "template_mask": template_mask,
-        "query_to_template": query_to_template,
-        "visibility_ids": visibility_ids,
-    }
-
-    return features, template  # each feature shape is traced at its binding above
+"""Construct Boltz2 protein features from canonical atom tables and conformers."""
+
+import math
+import numpy as np
+import torch
+
+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,
+    "N": 7,
+    "O": 8,
+    "P": 15,
+    "S": 16,
+}
+
+# Canonical formal charges and tetrahedral chirality values extracted from the
+# pinned Boltz2 molecule archive. Values use Boltz2's public feature encoding.
+_FORMAL_CHARGES = {
+    ("ARG", "NH1"): 1.0,
+    ("HIS", "ND1"): 1.0,
+    ("LYS", "NZ"): 1.0,
+}
+_CHIRAL_ATOMS = {
+    "ALA": frozenset({"CA"}),
+    "ARG": frozenset({"CA"}),
+    "ASN": frozenset({"CA"}),
+    "ASP": frozenset({"CA"}),
+    "CYS": frozenset({"CA"}),
+    "GLN": frozenset({"CA"}),
+    "GLU": frozenset({"CA"}),
+    "HIS": frozenset({"CA"}),
+    "ILE": frozenset({"CA", "CB"}),
+    "LEU": frozenset({"CA"}),
+    "LYS": frozenset({"CA"}),
+    "MET": frozenset({"CA"}),
+    "PHE": frozenset({"CA"}),
+    "PRO": frozenset({"CA"}),
+    "SER": frozenset({"CA"}),
+    "THR": frozenset({"CA", "CB"}),
+    "TRP": frozenset({"CA"}),
+    "TYR": frozenset({"CA"}),
+    "UNK": frozenset({"CA"}),
+    "VAL": frozenset({"CA"}),
+}
+
+
+def _normalize_sequence(sequence: str) -> str:
+    if not isinstance(sequence, str):
+        raise TypeError("Amino acid sequence must be a string.")
+    seq = sequence.strip().upper()
+    if not seq:
+        raise ValueError("Amino acid sequence must be non-empty.")
+    for aa in seq:
+        if aa not in const.prot_letter_to_token:
+            raise ValueError(f"Unsupported residue code '{aa}'.")
+    return seq
+
+
+def _atom_name_to_element(atom_name: str) -> str:
+    name = atom_name.strip().upper()
+    if len(name) == 0:
+        return "C"
+    if name[0].isdigit():
+        name = name[1:]
+    if len(name) >= 2 and name[0:2] in ("CL", "BR", "FE", "MG", "ZN", "NA", "CA"):
+        return name[0]
+    return name[0]
+
+
+def _atom_name_to_codes(atom_name: str) -> torch.Tensor:
+    clipped = atom_name.strip()[:4]
+    vals = [ord(ch) - 32 for ch in clipped]
+    while len(vals) < 4:
+        vals.append(0)
+    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
+
+
+# Raw first-conformer coordinates extracted at full float32 precision from the
+# hash-pinned Boltz2 molecule archive. The feature builder applies the official
+# seeded centering, rotation, and translation policy below.
+_RDKIT_CONFORMERS: dict[str, dict[str, list[float]]] = {
+    "ALA": {
+        "N": [-0.9241582155227661, 1.1821246147155762, 0.712748110294342],
+        "CA": [-0.2663755416870117, -0.08827890455722809, 0.4008508622646332],
+        "C": [1.1188693046569824, 0.1387452930212021, -0.14366498589515686],
+        "O": [1.2882297039031982, 0.8058497309684753, -1.2001261711120605],
+        "CB": [-1.1134333610534668, -0.8915302753448486, -0.5876986384391785],
+    },
+    "ARG": {
+        "N": [3.253326416015625, -1.699564814567566, -0.9852627515792847],
+        "CA": [2.191511392593384, -0.692007839679718, -0.9506462812423706],
+        "C": [2.80021333694458, 0.672307550907135, -1.1160222291946411],
+        "O": [2.465440034866333, 1.3987600803375244, -2.0899009704589844],
+        "CB": [1.3811039924621582, -0.7949594855308533, 0.36316436529159546],
+        "CG": [0.20344306528568268, 0.19287171959877014, 0.457282155752182],
+        "CD": [-0.902410626411438, -0.09953747689723969, -0.5671625137329102],
+        "NE": [-2.0867979526519775, 0.7146661281585693, -0.2911669909954071],
+        "CZ": [-3.018158197402954, 0.4804283082485199, 0.7831193804740906],
+        "NH1": [-4.080137729644775, 1.4154621362686157, 0.9938872456550598],
+        "NH2": [-2.9226150512695312, -0.5637124180793762, 1.5533623695373535],
+    },
+    "ASN": {
+        "N": [-1.5767980813980103, -1.7227835655212402, 0.2427234649658203],
+        "CA": [-0.7041534781455994, -0.5622008442878723, 0.42767149209976196],
+        "C": [-1.246688961982727, 0.6072975993156433, -0.3483567237854004],
+        "O": [-1.554918885231018, 0.47644123435020447, -1.5638835430145264],
+        "CB": [0.7263866662979126, -0.9032993316650391, -0.028588544577360153],
+        "CG": [1.6871562004089355, 0.20810098946094513, 0.28412675857543945],
+        "OD1": [2.19687557220459, 0.28916797041893005, 1.433493971824646],
+        "ND2": [1.996717095375061, 1.196706771850586, -0.6981719136238098],
+    },
+    "ASP": {
+        "N": [-0.11774874478578568, -1.6310220956802368, 0.374360054731369],
+        "CA": [-0.3847021758556366, -0.19290897250175476, 0.28121331334114075],
+        "C": [-1.806337594985962, 0.059408850967884064, -0.1472136378288269],
+        "O": [-2.2237648963928223, -0.36817988753318787, -1.257502555847168],
+        "CB": [0.5981687307357788, 0.4785478413105011, -0.6920451521873474],
+        "CG": [2.0088446140289307, 0.36018264293670654, -0.20304732024669647],
+        "OD1": [2.747267961502075, -0.5711832046508789, -0.6216948628425598],
+        "OD2": [2.49090313911438, 1.2577818632125854, 0.7443997859954834],
+    },
+    "CYS": {
+        "N": [-0.058702241629362106, 1.771048665046692, 0.2434082329273224],
+        "CA": [-0.06546340137720108, 0.47663936018943787, -0.4456723928451538],
+        "C": [-1.273919701576233, -0.3504071533679962, -0.08361759781837463],
+        "O": [-1.7324607372283936, -0.34313708543777466, 1.0909277200698853],
+        "CB": [1.2432398796081543, -0.28917407989501953, -0.19254730641841888],
+        "SG": [1.484184980392456, -0.7191365361213684, 1.5649693012237549],
+    },
+    "GLN": {
+        "N": [-1.8543237447738647, -1.0024770498275757, -1.6278940439224243],
+        "CA": [-1.292184591293335, -0.6786512732505798, -0.3153715133666992],
+        "C": [-2.2264492511749268, 0.25171151757240295, 0.4067918062210083],
+        "O": [-2.7255098819732666, -0.08834805339574814, 1.5129321813583374],
+        "CB": [0.11662524193525314, -0.06213820353150368, -0.45663803815841675],
+        "CG": [0.8140791058540344, 0.10461423546075821, 0.9002283811569214],
+        "CD": [2.195122003555298, 0.6569428443908691, 0.7154991030693054],
+        "OE1": [2.3870291709899902, 1.9003654718399048, 0.7825236916542053],
+        "NE2": [3.287371873855591, -0.21307705342769623, 0.41959717869758606],
+    },
+    "GLU": {
+        "N": [-1.3492857217788696, -1.114271640777588, -1.3739068508148193],
+        "CA": [-1.2765676975250244, -0.506334662437439, -0.04287439212203026],
+        "C": [-1.9195666313171387, 0.8547031879425049, -0.04214814677834511],
+        "O": [-1.9091347455978394, 1.5701956748962402, -1.080504298210144],
+        "CB": [0.18272103369235992, -0.43209096789360046, 0.45213374495506287],
+        "CG": [1.1058456897735596, 0.3659400939941406, -0.481784850358963],
+        "CD": [2.5090205669403076, 0.35135790705680847, 0.03646872192621231],
+        "OE1": [2.896414041519165, 1.2445733547210693, 0.8365193605422974],
+        "OE2": [3.3727896213531494, -0.6754258275032043, -0.33141785860061646],
+    },
+    "GLY": {
+        "N": [-1.291549801826477, 0.6080796122550964, -0.4228580892086029],
+        "CA": [-0.4895951449871063, -0.2882237136363983, 0.40191715955734253],
+        "C": [0.9350062608718872, -0.2543502449989319, -0.04786944016814232],
+        "O": [1.3473576307296753, -1.0836502313613892, -0.9022219777107239],
+    },
+    "HIS": {
+        "N": [1.0371263027191162, -1.5621215105056763, 0.4178937077522278],
+        "CA": [1.1980191469192505, -0.40241551399230957, -0.46312251687049866],
+        "C": [2.6561269760131836, -0.07123222947120667, -0.6407948136329651],
+        "O": [3.1795547008514404, -0.12371637672185898, -1.7859554290771484],
+        "CB": [0.44051393866539, 0.8116919994354248, 0.09850303083658218],
+        "CG": [-1.0356801748275757, 0.5519806742668152, 0.15254715085029602],
+        "ND1": [-1.9089158773422241, 0.6490958333015442, -0.968245267868042],
+        "CD2": [-1.7095766067504883, 0.1138235554099083, 1.2061808109283447],
+        "CE1": [-3.0751633644104004, 0.28129827976226807, -0.5511555075645447],
+        "NE2": [-3.0650453567504883, -0.07846951484680176, 0.8257263898849487],
+    },
+    "ILE": {
+        "N": [-1.2378733158111572, -1.8146690130233765, -0.16437499225139618],
+        "CA": [-1.2780015468597412, -0.3994033634662628, 0.22437156736850739],
+        "C": [-1.9879958629608154, 0.4497596025466919, -0.80372154712677],
+        "O": [-2.0291709899902344, 0.10288607329130173, -2.015406608581543],
+        "CB": [0.14225825667381287, 0.146738663315773, 0.5415176153182983],
+        "CG1": [1.1062071323394775, 0.04519597440958023, -0.6694174408912659],
+        "CG2": [0.7215384244918823, -0.5685611963272095, 1.7750537395477295],
+        "CD1": [2.371805191040039, 0.8852734565734863, -0.49045464396476746],
+    },
+    "LEU": {
+        "N": [1.6452248096466064, -1.017622470855713, -1.0626215934753418],
+        "CA": [1.3129304647445679, 0.08074887096881866, -0.15053576231002808],
+        "C": [2.431154727935791, 0.28310757875442505, 0.8357383608818054],
+        "O": [2.9231173992156982, -0.7013152837753296, 1.451209306716919],
+        "CB": [0.006050171796232462, -0.2141101360321045, 0.6198461651802063],
+        "CG": [-1.2534235715866089, -0.3772832751274109, -0.26637399196624756],
+        "CD1": [-2.455688714981079, -0.7604416012763977, 0.6070646643638611],
+        "CD2": [-1.5742294788360596, 0.8985711932182312, -1.059889554977417],
+    },
+    "LYS": {
+        "N": [-2.3938536643981934, -1.4751482009887695, -0.9336304068565369],
+        "CA": [-2.2094509601593018, -0.6203567981719971, 0.2421264797449112],
+        "C": [-3.4513463973999023, 0.1961250752210617, 0.4710204303264618],
+        "O": [-3.9688456058502197, 0.2503468096256256, 1.6188287734985352],
+        "CB": [-1.008726716041565, 0.3269270658493042, 0.0479682981967926],
+        "CG": [0.336457222700119, -0.41559070348739624, 0.047045741230249405],
+        "CD": [1.5088860988616943, 0.5716065168380737, -0.00632853340357542],
+        "CE": [2.853410005569458, -0.1678532361984253, -0.020898228511214256],
+        "NZ": [3.9690465927124023, 0.7765668630599976, -0.07007527351379395],
+    },
+    "MET": {
+        "N": [-1.743452548980713, -0.44039490818977356, 1.851780652999878],
+        "CA": [-1.2168668508529663, 0.4041755199432373, 0.7779996991157532],
+        "C": [-2.3118340969085693, 0.683431088924408, -0.2134547382593155],
+        "O": [-2.6660380363464355, 1.8705639839172363, -0.44476020336151123],
+        "CB": [-0.0005482881097123027, -0.262584388256073, 0.09933014959096909],
+        "CG": [0.6952047944068909, 0.6768104434013367, -0.8930339813232422],
+        "SD": [2.136805772781372, -0.14913176000118256, -1.6563159227371216],
+        "CE": [3.3404407501220703, 0.348065584897995, -0.38104408979415894],
+    },
+    "PHE": {
+        "N": [2.7976438999176025, -1.6209964752197266, -0.4779351055622101],
+        "CA": [1.5799994468688965, -0.8289614915847778, -0.6608061194419861],
+        "C": [1.9535551071166992, 0.5758142471313477, -1.0445665121078491],
+        "O": [1.5211102962493896, 1.072593331336975, -2.1190123558044434],
+        "CB": [0.7137221097946167, -0.8474348187446594, 0.6163292527198792],
+        "CG": [-0.6229045987129211, -0.18666480481624603, 0.3912639319896698],
+        "CD1": [-0.8259273767471313, 1.0978055000305176, 0.7381877303123474],
+        "CD2": [-1.7160141468048096, -0.9425657391548157, -0.2700555920600891],
+        "CE1": [-2.131312131881714, 1.741182565689087, 0.4800238013267517],
+        "CE2": [-2.8960046768188477, -0.3551671802997589, -0.505269467830658],
+        "CZ": [-3.11456298828125, 1.0511738061904907, -0.11009909957647324],
+    },
+    "PRO": {
+        "N": [-0.577339768409729, -0.5123927593231201, -1.150407314300537],
+        "CA": [0.5119893550872803, 0.23565764725208282, -0.5195383429527283],
+        "C": [1.8393173217773438, -0.46712079644203186, -0.6550371646881104],
+        "O": [2.0836586952209473, -1.166333556175232, -1.6751221418380737],
+        "CB": [0.1029478907585144, 0.45346882939338684, 0.9297074675559998],
+        "CG": [-1.4098080396652222, 0.5215330123901367, 0.8787456750869751],
+        "CD": [-1.7971025705337524, -0.06041542813181877, -0.47876954078674316],
+    },
+    "SER": {
+        "N": [0.8959873914718628, -1.423298716545105, -0.2692987024784088],
+        "CA": [-0.018787339329719543, -0.28751760721206665, -0.3863537311553955],
+        "C": [0.767615795135498, 0.9929561018943787, -0.42147380113601685],
+        "O": [0.5327092409133911, 1.8535151481628418, -1.3116248846054077],
+        "CB": [-0.9973848462104797, -0.2769756019115448, 0.7988247871398926],
+        "OG": [-1.8999865055084229, 0.7915746569633484, 0.691534698009491],
+    },
+    "THR": {
+        "N": [0.005045319441705942, 1.7655576467514038, 0.0639338344335556],
+        "CA": [0.4231603145599365, 0.3755212128162384, 0.2734847664833069],
+        "C": [1.7184523344039917, 0.11012311279773712, -0.4465082287788391],
+        "O": [1.9252793788909912, 0.5967074036598206, -1.5912789106369019],
+        "CB": [-0.6634359359741211, -0.615561306476593, -0.21173778176307678],
+        "OG1": [-1.0762907266616821, -0.29800114035606384, -1.5171915292739868],
+        "CG2": [-1.8869022130966187, -0.6176311373710632, 0.7084044218063354],
+    },
+    "TRP": {
+        "N": [-3.0439648628234863, 1.0628471374511719, 0.015692168846726418],
+        "CA": [-2.463679552078247, 0.024657616391777992, -0.8414031863212585],
+        "C": [-2.186711072921753, -1.2228965759277344, -0.04691409692168236],
+        "O": [-2.3364906311035156, -2.3535354137420654, -0.5830773115158081],
+        "CB": [-1.2062920331954956, 0.5195542573928833, -1.592403531074524],
+        "CG": [-0.06450498104095459, 0.884079098701477, -0.6835198402404785],
+        "CD1": [0.16109998524188995, 2.085163116455078, -0.14801783859729767],
+        "CD2": [0.9781660437583923, -0.005544683896005154, -0.17733941972255707],
+        "NE1": [1.3106837272644043, 2.0506503582000732, 0.698072075843811],
+        "CE2": [1.7528777122497559, 0.7044392228126526, 0.6196392774581909],
+        "CE3": [1.236093521118164, -1.4355946779251099, -0.4235689342021942],
+        "CZ2": [2.9070138931274414, 0.10196790844202042, 1.3070299625396729],
+        "CZ3": [2.2851555347442627, -2.007185697555542, 0.1949695348739624],
+        "CH2": [3.1585631370544434, -1.2034809589385986, 1.1002938747406006],
+    },
+    "TYR": {
+        "N": [-1.7950150966644287, 0.49119046330451965, -1.3951144218444824],
+        "CA": [-1.8436503410339355, -0.2694968581199646, -0.14338290691375732],
+        "C": [-3.240288019180298, -0.28215137124061584, 0.41974759101867676],
+        "O": [-3.8150112628936768, 0.7977103590965271, 0.7253130674362183],
+        "CB": [-0.8541494011878967, 0.30990728735923767, 0.8830214738845825],
+        "CG": [0.5672639012336731, 0.2125871181488037, 0.3916495442390442],
+        "CD1": [1.2694358825683594, -0.924761176109314, 0.5431669354438782],
+        "CD2": [1.1872504949569702, 1.3605931997299194, -0.31554359197616577],
+        "CE1": [2.6519992351531982, -1.0213464498519897, 0.029911600053310394],
+        "CE2": [2.4384164810180664, 1.2693941593170166, -0.7831515073776245],
+        "CZ": [3.211331367492676, 0.022552935406565666, -0.6021429300308228],
+        "OH": [4.513507843017578, -0.05260590463876724, -1.0929937362670898],
+    },
+    "VAL": {
+        "N": [0.9408224821090698, -1.2608877420425415, 0.652370810508728],
+        "CA": [0.7287879586219788, -0.3937721848487854, -0.5120788216590881],
+        "C": [1.7670493125915527, 0.6979415416717529, -0.5508439540863037],
+        "O": [2.2449657917022705, 1.1683598756790161, 0.5171348452568054],
+        "CB": [-0.7015937566757202, 0.21480272710323334, -0.5226467847824097],
+        "CG1": [-1.7662651538848877, -0.845522403717041, -0.8411717414855957],
+        "CG2": [-1.0500152111053467, 0.9374979138374329, 0.7892321944236755],
+    },
+    "UNK": {
+        "N": [1.6134154796600342, -1.304404377937317, -0.35241976380348206],
+        "CA": [0.6100443601608276, -0.23816797137260437, -0.3852081298828125],
+        "C": [1.0861350297927856, 0.9232651591300964, 0.4423142373561859],
+        "O": [1.374688982963562, 0.7629808187484741, 1.6591525077819824],
+        "CB": [-0.7607048153877258, -0.7552146315574646, 0.10266945511102676],
+    },
+}
+
+
+def _get_atom_position(res_name: str, atom_name: str, atom_idx: int) -> np.ndarray:
+    """Get the canonical RDKit conformer position for an atom.
+
+    Uses pre-extracted positions from official Boltz2 mol files. Falls back
+    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)  # (3,)
+    # Fallback for unknown atoms (should not happen for canonical AAs)
+    angle = (atom_idx + 1) * 0.7
+    radius = 1.4 + 0.03 * atom_idx
+    return np.array(
+        [
+            radius * math.cos(angle),
+            radius * math.sin(angle),
+            0.1 * ((atom_idx % 5) - 2),
+        ],
+        dtype=np.float32,
+    )  # (3,)
+
+
+def _build_template(
+    sequence: str,
+) -> tuple[
+    ProteinStructureTemplate,
+    list[str],
+    list[int],
+    list[int],
+    list[int],
+    list[np.ndarray],
+    list[int],
+]:
+    residue_names: list[str] = []
+    residue_token_ids: list[int] = []
+    atom_names: list[str] = []
+    atom_elements: list[str] = []
+    atom_residue_index: list[int] = []
+    atom_chain_id: list[str] = []
+    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] = []
+
+    global_atom_idx = 0
+    for res_idx, aa in enumerate(sequence):
+        token_name = const.prot_letter_to_token[aa]
+        residue_names.append(token_name)
+        residue_token_ids.append(const.token_ids[token_name])
+
+        residue_atoms = const.ref_atoms[token_name]
+        if not residue_atoms:
+            raise RuntimeError(f"No reference atoms for residue {token_name}.")
+        center_atom_name = const.res_to_center_atom[token_name]
+        disto_atom_name = const.res_to_disto_atom[token_name]
+
+        center_idx = -1
+        disto_idx = -1
+        n_idx = -1
+        ca_idx = -1
+        c_idx = -1
+
+        for local_idx, atom_name in enumerate(residue_atoms):
+            atom_names.append(atom_name)
+            element = _atom_name_to_element(atom_name)
+            atom_elements.append(element)
+            atom_residue_index.append(res_idx)
+            atom_chain_id.append("A")
+
+            atom_pos = _get_atom_position(token_name, atom_name, local_idx)  # (3,)
+            atom_positions.append(atom_pos)
+
+            if atom_name == center_atom_name:
+                center_idx = global_atom_idx
+            if atom_name == disto_atom_name:
+                disto_idx = global_atom_idx
+            if atom_name == "N":
+                n_idx = global_atom_idx
+            if atom_name == "CA":
+                ca_idx = global_atom_idx
+            if atom_name == "C":
+                c_idx = global_atom_idx
+            global_atom_idx += 1
+
+        if center_idx == -1:
+            center_idx = global_atom_idx - len(residue_atoms)
+        if disto_idx == -1:
+            disto_idx = center_idx
+        if n_idx == -1:
+            n_idx = center_idx
+        if ca_idx == -1:
+            ca_idx = center_idx
+        if c_idx == -1:
+            c_idx = center_idx
+
+        residue_center_atom_idx.append(center_idx)
+        residue_disto_atom_idx.append(disto_idx)
+        residue_frame_atom_idx.extend([n_idx, ca_idx, c_idx])
+
+    template = ProteinStructureTemplate(
+        sequence=sequence,
+        residue_names=residue_names,
+        atom_names=atom_names,
+        atom_elements=atom_elements,
+        atom_residue_index=atom_residue_index,
+        atom_chain_id=atom_chain_id,
+    )
+
+    return (
+        template,
+        residue_names,
+        residue_token_ids,
+        residue_center_atom_idx,
+        residue_disto_atom_idx,
+        atom_positions,
+        residue_frame_atom_idx,
+    )
+
+
+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)  # (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),
+            two_s * (i * j - k * real),
+            two_s * (i * k + j * real),
+            two_s * (i * j + k * real),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * real),
+            two_s * (i * k - j * real),
+            two_s * (j * k + i * real),
+            1 - two_s * (i * i + j * j),
+        ),
+        dim=-1,
+    )  # (1, 9)
+    return rotation.reshape(1, 3, 3)[0]  # (3, 3)
+
+
+def _center_and_augment_atoms_per_residue(
+    atom_positions: torch.Tensor,
+    atom_residue_index: list[int],
+    num_residues: int,
+) -> torch.Tensor:
+    """Apply Boltz2's seeded conformer augmentation per residue.
+
+    The pinned implementation intentionally excludes the final residue because
+    it iterates to the maximum reference-space identifier rather than through it.
+    """
+    # 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  # (a,)
+        if not torch.any(residue_mask):
+            raise RuntimeError(f"Residue index {residue_idx} has no atoms.")
+        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)  # (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(
+    amino_acid_sequence: str,
+    num_bins: int = 64,
+    atoms_per_window_queries: int = 32,
+) -> tuple[dict[str, torch.Tensor], ProteinStructureTemplate]:
+    sequence = _normalize_sequence(amino_acid_sequence)
+    (
+        template,
+        residue_names,
+        residue_token_ids,
+        residue_center_atom_idx,
+        residue_disto_atom_idx,
+        atom_positions_np,
+        residue_frame_atom_idx_flat,
+    ) = _build_template(sequence)
+
+    num_tokens = len(residue_names)
+    num_atoms = len(atom_positions_np)
+    if num_tokens <= 0 or num_atoms <= 0:
+        raise RuntimeError("Boltz2 feature construction produced an empty protein template.")
+
+    # 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)  # (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)  # (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)  # (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)  # (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"]
+    contact_ids = torch.full(
+        (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)  # (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.")
+    method_feature = torch.full(
+        (1, num_tokens),
+        fill_value=const.method_types_ids["x-ray diffraction"],
+        dtype=torch.long,
+    )  # (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)  # (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,
+    )  # (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:
+        z_value = _ELEMENT_TO_Z[element] if element in _ELEMENT_TO_Z else _ELEMENT_TO_Z["C"]
+        if z_value >= const.num_elements:
+            raise RuntimeError(
+                f"Atomic number {z_value} exceeds the Boltz2 element vocabulary."
+            )
+        atomic_numbers.append(z_value)
+    ref_element = one_hot(
+        torch.tensor(atomic_numbers, dtype=torch.long),
+        num_classes=const.num_elements,
+    ).unsqueeze(0)  # (1, a, n_element)
+
+    ref_charge_values = []
+    ref_chirality_values = []
+    for atom_name, residue_idx in zip(
+        template.atom_names,
+        template.atom_residue_index,
+        strict=True,
+    ):
+        residue_name = residue_names[residue_idx]
+        ref_charge_values.append(_FORMAL_CHARGES.get((residue_name, atom_name), 0.0))
+        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)  # (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)  # (1, a, t)
+    token_to_rep_atom = one_hot(
+        torch.tensor(residue_disto_atom_idx, dtype=torch.long),
+        num_classes=num_atoms,
+    ).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)  # (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)
+    )
+    backbone_ids = []
+    for atom_name in template.atom_names:
+        if atom_name in const.protein_backbone_atom_index:
+            backbone_ids.append(const.protein_backbone_atom_index[atom_name] + 1)
+        else:
+            backbone_ids.append(0)
+    atom_backbone_feat = one_hot(
+        torch.tensor(backbone_ids, dtype=torch.long),
+        num_classes=num_backbone_classes,
+    ).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)  # (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)  # (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.")
+    pad_atoms = (
+        (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
+        )  # (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,
+        )  # (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,
+        )  # (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,
+        )  # (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)  # (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),
+            num_classes=const.num_tokens,
+        )
+        .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,
+    )  # (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,
+    )  # (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,
+        "residue_index": residue_index,
+        "asym_id": asym_id,
+        "entity_id": entity_id,
+        "sym_id": sym_id,
+        "mol_type": mol_type,
+        "res_type": res_type,
+        "disto_center": disto_center,
+        "token_bonds": token_bonds,
+        "type_bonds": type_bonds,
+        "token_pad_mask": token_pad_mask,
+        "token_resolved_mask": token_resolved_mask,
+        "token_disto_mask": token_disto_mask,
+        "contact_conditioning": contact_conditioning,
+        "contact_threshold": contact_threshold,
+        "method_feature": method_feature,
+        "modified": modified,
+        "cyclic_period": cyclic_period,
+        "affinity_token_mask": affinity_token_mask,
+        "ref_pos": ref_pos,
+        "atom_resolved_mask": atom_resolved_mask,
+        "ref_atom_name_chars": ref_atom_name_chars,
+        "ref_element": ref_element,
+        "ref_charge": ref_charge,
+        "ref_chirality": ref_chirality,
+        "atom_backbone_feat": atom_backbone_feat,
+        "ref_space_uid": ref_space_uid,
+        "coords": coords,
+        "atom_pad_mask": atom_pad_mask,
+        "atom_to_token": atom_to_token,
+        "token_to_rep_atom": token_to_rep_atom,
+        "r_set_to_rep_atom": r_set_to_rep_atom,
+        "token_to_center_atom": token_to_center_atom,
+        "disto_target": disto_target,
+        "disto_coords_ensemble": disto_coords_ensemble,
+        "bfactor": bfactor,
+        "plddt": atom_plddt,
+        "frames_idx": frames_idx,
+        "frame_resolved_mask": frame_resolved_mask,
+        "msa": msa,
+        "msa_paired": msa_paired,
+        "deletion_value": deletion_value,
+        "has_deletion": has_deletion,
+        "deletion_mean": deletion_mean,
+        "profile": profile,
+        "msa_mask": msa_mask,
+        "template_restype": template_restype,
+        "template_frame_rot": template_frame_rot,
+        "template_frame_t": template_frame_t,
+        "template_cb": template_cb,
+        "template_ca": template_ca,
+        "template_mask_cb": template_mask_cb,
+        "template_mask_frame": template_mask_frame,
+        "template_mask": template_mask,
+        "query_to_template": query_to_template,
+        "visibility_ids": visibility_ids,
+    }
+
+    return features, template  # each feature shape is traced at its binding above

fastplms/models/boltz/minimal_structures.py

@@ -1,21 +1,21 @@
-"""Atom and residue metadata for sequence-only Boltz2 structure exports."""
-
-from dataclasses import dataclass
-
-
-@dataclass
-class ProteinStructureTemplate:
-    sequence: str
-    residue_names: list[str]
-    atom_names: list[str]
-    atom_elements: list[str]
-    atom_residue_index: list[int]
-    atom_chain_id: list[str]
-
-    @property
-    def num_atoms(self) -> int:
-        return len(self.atom_names)
-
-    @property
-    def num_residues(self) -> int:
-        return len(self.residue_names)
+"""Atom and residue metadata for sequence-only Boltz2 structure exports."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class ProteinStructureTemplate:
+    sequence: str
+    residue_names: list[str]
+    atom_names: list[str]
+    atom_elements: list[str]
+    atom_residue_index: list[int]
+    atom_chain_id: list[str]
+
+    @property
+    def num_atoms(self) -> int:
+        return len(self.atom_names)
+
+    @property
+    def num_residues(self) -> int:
+        return len(self.residue_names)

fastplms/models/boltz/modeling_boltz2.py

@@ -1,1061 +1,1066 @@
-"""Transformers interfaces and checkpoint conversion for the Boltz2 runtime."""
-
-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
-from torch import Tensor
-from tqdm.auto import tqdm
-from transformers import PretrainedConfig, PreTrainedModel
-from transformers.modeling_outputs import ModelOutput
-
-from . import vb_const as const
-from . import vb_layers_initialize as init
-from .cif_writer import write_cif
-from .minimal_featurizer import build_boltz2_features
-from .minimal_structures import ProteinStructureTemplate
-from .vb_const import bond_types as _vb_const_bond_types  # noqa: F401
-from .vb_layers_attention import AttentionPairBias as _vb_layers_attention_marker  # noqa: F401
-from .vb_layers_attentionv2 import AttentionPairBias as _vb_layers_attentionv2_marker  # noqa: F401
-from .vb_layers_confidence_utils import (
-    compute_ptms as _vb_layers_confidence_utils_marker,  # noqa: F401
-)
-from .vb_layers_dropout import get_dropout_mask as _vb_layers_dropout_marker  # noqa: F401
-from .vb_layers_initialize import gating_init_ as _vb_layers_initialize_marker  # noqa: F401
-from .vb_layers_outer_product_mean import (
-    OuterProductMean as _vb_layers_outer_product_mean_marker,  # noqa: F401
-)
-from .vb_layers_pair_averaging import (
-    PairWeightedAveraging as _vb_layers_pair_averaging_marker,  # noqa: F401
-)
-from .vb_layers_pairformer import PairformerModule
-from .vb_layers_transition import Transition as _vb_layers_transition_marker  # noqa: F401
-from .vb_layers_triangular_mult import (
-    TriangleMultiplicationIncoming as _vb_layers_triangular_mult_marker,  # noqa: F401
-)
-from .vb_loss_diffusionv2 import weighted_rigid_align as _vb_loss_diffusionv2_marker  # noqa: F401
-from .vb_modules_confidencev2 import ConfidenceModule
-from .vb_modules_diffusion_conditioning import DiffusionConditioning
-from .vb_modules_diffusionv2 import AtomDiffusion, DiffusionModule
-from .vb_modules_encodersv2 import RelativePositionEncoder
-from .vb_modules_transformersv2 import (
-    DiffusionTransformer as _vb_modules_transformersv2_marker,  # noqa: F401
-)
-from .vb_modules_trunkv2 import (
-    ContactConditioning,
-    DistogramModule,
-    InputEmbedder,
-    MSAModule,
-)
-from .vb_modules_utils import LinearNoBias as _vb_modules_utils_marker  # noqa: F401
-from .vb_potentials_potentials import (
-    get_potentials as _vb_potentials_potentials_marker,  # noqa: F401
-)
-from .vb_potentials_schedules import (
-    ParameterSchedule as _vb_potentials_schedules_marker,  # noqa: F401
-)
-from .vb_tri_attn_attention import (
-    TriangleAttentionStartingNode as _vb_tri_attn_attention_marker,  # noqa: F401
-)
-from .vb_tri_attn_primitives import Attention as _vb_tri_attn_primitives_marker  # noqa: F401
-from .vb_tri_attn_utils import permute_final_dims as _vb_tri_attn_utils_marker  # noqa: F401
-
-
-def _default_steering_args() -> dict[str, Any]:
-    return {
-        "fk_steering": False,
-        "num_particles": 3,
-        "fk_lambda": 4.0,
-        "fk_resampling_interval": 3,
-        "physical_guidance_update": False,
-        "contact_guidance_update": False,
-        "num_gd_steps": 16,
-    }
-
-
-def _boltz2_reference_diffusion_overrides() -> dict[str, Any]:
-    # Match Boltz2 CLI inference defaults from boltz.main/Boltz2DiffusionParams.
-    return {
-        "gamma_0": 0.8,
-        "gamma_min": 1.0,
-        "noise_scale": 1.003,
-        "rho": 7,
-        "step_scale": 1.5,
-        "sigma_min": 0.0001,
-        "sigma_max": 160.0,
-        "sigma_data": 16.0,
-        "P_mean": -1.2,
-        "P_std": 1.5,
-        "coordinate_augmentation": True,
-        "alignment_reverse_diff": True,
-        "synchronize_sigmas": True,
-    }
-
-
-def _enforce_pairformer_v2(pairformer_args: Mapping[str, Any], context: str) -> dict[str, Any]:
-    if not isinstance(pairformer_args, Mapping):
-        raise TypeError(f"Expected {context} pairformer_args to be a dictionary.")
-    out = _to_plain_python(copy.deepcopy(pairformer_args))
-    if "v2" in out and not out["v2"]:
-        raise ValueError(f"{context} pairformer_args['v2'] must be True for Boltz2.")
-    out["v2"] = True
-    return out
-
-
-def _require_key(mapping: dict[str, Any], key: str) -> Any:
-    if key not in mapping:
-        raise KeyError(f"Missing required key '{key}' in checkpoint hyperparameters.")
-    return mapping[key]
-
-
-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."):
-            continue
-        new_key = key
-        if new_key.startswith("model."):
-            new_key = new_key[len("model.") :]
-        if new_key.startswith("module."):
-            new_key = new_key[len("module.") :]
-        cleaned[new_key] = value  # (...)
-    return cleaned
-
-
-def _to_cpu_detached(value: Any) -> Any:
-    if torch.is_tensor(value):
-        return value.detach().cpu()  # same shape (...)
-    if isinstance(value, dict):
-        out: dict[Any, Any] = {}
-        for key, nested_value in value.items():
-            out[key] = _to_cpu_detached(nested_value)
-        return out
-    if isinstance(value, list):
-        return [_to_cpu_detached(item) for item in value]
-    if isinstance(value, tuple):
-        return tuple(_to_cpu_detached(item) for item in value)
-    return value
-
-
-@dataclass(frozen=True)
-class _RandomState:
-    python: object
-    numpy: tuple[Any, ...]
-    torch_cpu: Tensor  # (n_cpu_rng_state,)
-    torch_cuda: list[Tensor] | None  # each: (n_cuda_rng_state,)
-
-
-@contextmanager
-def _seed_context(seed: int | None) -> Iterator[None]:
-    """Temporarily seed every RNG used by Boltz2 and restore caller state."""
-
-    if seed is None:
-        yield
-        return
-    if isinstance(seed, bool) or not isinstance(seed, int):
-        raise TypeError(
-            "Boltz2 seed must be an int or None; bool and coerced values are not accepted."
-        )
-    state = _RandomState(
-        python=random.getstate(),
-        numpy=np.random.get_state(),
-        torch_cpu=torch.random.get_rng_state(),
-        torch_cuda=torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
-    )
-    normalized_seed = seed % (2**32)
-    random.seed(normalized_seed)
-    np.random.seed(normalized_seed)
-    torch.manual_seed(normalized_seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(normalized_seed)
-    try:
-        yield
-    finally:
-        random.setstate(state.python)
-        np.random.set_state(state.numpy)
-        torch.random.set_rng_state(state.torch_cpu)
-        if state.torch_cuda is not None:
-            torch.cuda.set_rng_state_all(state.torch_cuda)
-
-
-def _to_plain_python(value: Any) -> Any:
-    if isinstance(value, Mapping):
-        out: dict[Any, Any] = {}
-        for key, nested_value in value.items():
-            out[key] = _to_plain_python(nested_value)
-        return out
-    if isinstance(value, list):
-        return [_to_plain_python(item) for item in value]
-    if isinstance(value, tuple):
-        return [_to_plain_python(item) for item in value]
-    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
-        return [_to_plain_python(item) for item in value]
-    return value
-
-
-def _filtered_kwargs(target: Any, kwargs: dict[str, Any]) -> dict[str, Any]:
-    signature = inspect.signature(target.__init__)
-    allowed = set(signature.parameters.keys())
-    allowed.discard("self")
-    filtered: dict[str, Any] = {}
-    for key, value in kwargs.items():
-        if key in allowed:
-            filtered[key] = value
-    return filtered
-
-
-@dataclass
-class Boltz2ModelOutput(ModelOutput):
-    """Raw Boltz2 inference output with standard AutoModel controls."""
-
-    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  # (m, a_p, 3) or None
-    atom_pad_mask: torch.Tensor | None = None  # (a_p,) or None
-    plddt: torch.Tensor | None = None  # (m, t) or (m, a_p), depending on confidence resolution
-    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  # values keep model-output shapes
-    seed: int | None = None
-
-
-class Boltz2Config(PretrainedConfig):
-    model_type = "boltz2_automodel"
-
-    def __init__(
-        self,
-        core_kwargs: dict[str, Any] | None = None,
-        num_bins: int = 64,
-        default_recycling_steps: int = 3,
-        default_sampling_steps: int = 200,
-        default_diffusion_samples: int = 1,
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-        if core_kwargs is None:
-            core_kwargs = {}
-        self.core_kwargs = core_kwargs
-        self.num_bins = num_bins
-        self.default_recycling_steps = default_recycling_steps
-        self.default_sampling_steps = default_sampling_steps
-        self.default_diffusion_samples = default_diffusion_samples
-
-    @classmethod
-    def from_hyperparameters(
-        cls,
-        hparams: dict[str, Any],
-        use_kernels: bool = False,
-        default_recycling_steps: int | None = None,
-        default_sampling_steps: int | None = None,
-        default_diffusion_samples: int | None = None,
-    ) -> "Boltz2Config":
-        if not isinstance(hparams, dict):
-            raise TypeError("Expected checkpoint hyperparameters as a dictionary.")
-        required = [
-            "atom_s",
-            "atom_z",
-            "token_s",
-            "token_z",
-            "num_bins",
-            "embedder_args",
-            "msa_args",
-            "pairformer_args",
-            "score_model_args",
-            "diffusion_process_args",
-        ]
-        for key in required:
-            _require_key(hparams, key)
-
-        pairformer_args = _enforce_pairformer_v2(
-            hparams["pairformer_args"],
-            context="checkpoint",
-        )
-        diffusion_process_args = _to_plain_python(copy.deepcopy(hparams["diffusion_process_args"]))
-        diffusion_overrides = _boltz2_reference_diffusion_overrides()
-        for key in diffusion_overrides:
-            diffusion_process_args[key] = diffusion_overrides[key]
-
-        core_kwargs: dict[str, Any] = {
-            "atom_s": hparams["atom_s"],
-            "atom_z": hparams["atom_z"],
-            "token_s": hparams["token_s"],
-            "token_z": hparams["token_z"],
-            "num_bins": hparams["num_bins"],
-            "embedder_args": _to_plain_python(copy.deepcopy(hparams["embedder_args"])),
-            "msa_args": _to_plain_python(copy.deepcopy(hparams["msa_args"])),
-            "pairformer_args": pairformer_args,
-            "score_model_args": _to_plain_python(copy.deepcopy(hparams["score_model_args"])),
-            "diffusion_process_args": diffusion_process_args,
-            "use_kernels": use_kernels,
-        }
-
-        if "confidence_model_args" in hparams:
-            confidence_model_args = _to_plain_python(
-                copy.deepcopy(hparams["confidence_model_args"])
-            )
-            if "pairformer_args" in confidence_model_args:
-                confidence_model_args["pairformer_args"] = _enforce_pairformer_v2(
-                    confidence_model_args["pairformer_args"],
-                    context="confidence",
-                )
-            core_kwargs["confidence_model_args"] = confidence_model_args
-        else:
-            core_kwargs["confidence_model_args"] = None
-
-        core_kwargs["confidence_prediction"] = hparams.get("confidence_prediction", True)
-        core_kwargs["token_level_confidence"] = hparams.get("token_level_confidence", True)
-        core_kwargs["alpha_pae"] = hparams.get("alpha_pae", 0.0)
-        core_kwargs["atoms_per_window_queries"] = hparams.get("atoms_per_window_queries", 32)
-        core_kwargs["atoms_per_window_keys"] = hparams.get("atoms_per_window_keys", 128)
-        core_kwargs["atom_feature_dim"] = hparams.get("atom_feature_dim", 128)
-        core_kwargs["bond_type_feature"] = hparams.get("bond_type_feature", False)
-        core_kwargs["run_trunk_and_structure"] = hparams.get("run_trunk_and_structure", True)
-        core_kwargs["skip_run_structure"] = hparams.get("skip_run_structure", False)
-        core_kwargs["fix_sym_check"] = hparams.get("fix_sym_check", False)
-        core_kwargs["cyclic_pos_enc"] = hparams.get("cyclic_pos_enc", False)
-        core_kwargs["use_no_atom_char"] = hparams.get("use_no_atom_char", False)
-        core_kwargs["use_atom_backbone_feat"] = hparams.get("use_atom_backbone_feat", False)
-        core_kwargs["use_residue_feats_atoms"] = hparams.get("use_residue_feats_atoms", False)
-        core_kwargs["conditioning_cutoff_min"] = hparams.get("conditioning_cutoff_min", 4.0)
-        core_kwargs["conditioning_cutoff_max"] = hparams.get("conditioning_cutoff_max", 20.0)
-
-        if "steering_args" in hparams and hparams["steering_args"] is not None:
-            core_kwargs["steering_args"] = _to_plain_python(copy.deepcopy(hparams["steering_args"]))
-        else:
-            core_kwargs["steering_args"] = _default_steering_args()
-
-        if "validation_args" in hparams:
-            validation_args = hparams["validation_args"]
-            if not isinstance(validation_args, Mapping):
-                raise TypeError(
-                    "Expected 'validation_args' in checkpoint hyperparameters to be a mapping."
-                )
-            if default_recycling_steps is None and "recycling_steps" in validation_args:
-                default_recycling_steps = validation_args["recycling_steps"]
-            if default_sampling_steps is None and "sampling_steps" in validation_args:
-                default_sampling_steps = validation_args["sampling_steps"]
-            if default_diffusion_samples is None and "diffusion_samples" in validation_args:
-                default_diffusion_samples = validation_args["diffusion_samples"]
-
-        if default_recycling_steps is None:
-            default_recycling_steps = 3
-        if default_sampling_steps is None:
-            default_sampling_steps = 200
-        if default_diffusion_samples is None:
-            default_diffusion_samples = 1
-
-        return cls(
-            core_kwargs=core_kwargs,
-            num_bins=hparams["num_bins"],
-            default_recycling_steps=default_recycling_steps,
-            default_sampling_steps=default_sampling_steps,
-            default_diffusion_samples=default_diffusion_samples,
-        )
-
-
-class Boltz2InferenceCore(nn.Module):
-    def __init__(
-        self,
-        atom_s: int,
-        atom_z: int,
-        token_s: int,
-        token_z: int,
-        num_bins: int,
-        embedder_args: dict[str, Any],
-        msa_args: dict[str, Any],
-        pairformer_args: dict[str, Any],
-        score_model_args: dict[str, Any],
-        diffusion_process_args: dict[str, Any],
-        confidence_model_args: dict[str, Any] | None = None,
-        atom_feature_dim: int = 128,
-        confidence_prediction: bool = True,
-        token_level_confidence: bool = True,
-        alpha_pae: float = 0.0,
-        atoms_per_window_queries: int = 32,
-        atoms_per_window_keys: int = 128,
-        run_trunk_and_structure: bool = True,
-        skip_run_structure: bool = False,
-        bond_type_feature: bool = False,
-        fix_sym_check: bool = False,
-        cyclic_pos_enc: bool = False,
-        use_no_atom_char: bool = False,
-        use_atom_backbone_feat: bool = False,
-        use_residue_feats_atoms: bool = False,
-        conditioning_cutoff_min: float = 4.0,
-        conditioning_cutoff_max: float = 20.0,
-        use_kernels: bool = False,
-        steering_args: dict[str, Any] | None = None,
-    ) -> None:
-        super().__init__()
-        self.use_kernels = use_kernels
-        self.confidence_prediction = confidence_prediction
-        self.token_level_confidence = token_level_confidence
-        self.alpha_pae = alpha_pae
-        self.run_trunk_and_structure = run_trunk_and_structure
-        self.skip_run_structure = skip_run_structure
-        self.bond_type_feature = bond_type_feature
-        self.steering_args = (
-            steering_args if steering_args is not None else _default_steering_args()
-        )
-        if "v2" not in pairformer_args:
-            raise ValueError("Boltz2 requires pairformer_args['v2'].")
-        if not pairformer_args["v2"]:
-            raise ValueError("Boltz2 requires pairformer_args['v2']=True.")
-
-        full_embedder_args = {
-            "atom_s": atom_s,
-            "atom_z": atom_z,
-            "token_s": token_s,
-            "token_z": token_z,
-            "atoms_per_window_queries": atoms_per_window_queries,
-            "atoms_per_window_keys": atoms_per_window_keys,
-            "atom_feature_dim": atom_feature_dim,
-            "use_no_atom_char": use_no_atom_char,
-            "use_atom_backbone_feat": use_atom_backbone_feat,
-            "use_residue_feats_atoms": use_residue_feats_atoms,
-            **embedder_args,
-        }
-        full_embedder_args = _filtered_kwargs(InputEmbedder, full_embedder_args)
-        self.input_embedder = InputEmbedder(**full_embedder_args)
-
-        self.s_init = nn.Linear(token_s, token_s, bias=False)
-        self.z_init_1 = nn.Linear(token_s, token_z, bias=False)
-        self.z_init_2 = nn.Linear(token_s, token_z, bias=False)
-        self.rel_pos = RelativePositionEncoder(
-            token_z,
-            fix_sym_check=fix_sym_check,
-            cyclic_pos_enc=cyclic_pos_enc,
-        )
-        self.token_bonds = nn.Linear(1, token_z, bias=False)
-        if self.bond_type_feature:
-            self.token_bonds_type = nn.Embedding(len(const.bond_types) + 1, token_z)
-
-        self.contact_conditioning = ContactConditioning(
-            token_z=token_z,
-            cutoff_min=conditioning_cutoff_min,
-            cutoff_max=conditioning_cutoff_max,
-        )
-        self.s_norm = nn.LayerNorm(token_s)
-        self.z_norm = nn.LayerNorm(token_z)
-
-        self.s_recycle = nn.Linear(token_s, token_s, bias=False)
-        self.z_recycle = nn.Linear(token_z, token_z, bias=False)
-        init.gating_init_(self.s_recycle.weight)
-        init.gating_init_(self.z_recycle.weight)
-
-        msa_kwargs = _filtered_kwargs(
-            MSAModule,
-            {"token_z": token_z, "token_s": token_s, **msa_args},
-        )
-        self.msa_module = MSAModule(**msa_kwargs)
-
-        pairformer_kwargs = _filtered_kwargs(
-            PairformerModule,
-            {"token_s": token_s, "token_z": token_z, **pairformer_args},
-        )
-        if "token_s" not in pairformer_kwargs or "token_z" not in pairformer_kwargs:
-            raise RuntimeError("Boltz2 pairformer construction lost its token dimensions.")
-        pairformer_token_s = pairformer_kwargs.pop("token_s")
-        pairformer_token_z = pairformer_kwargs.pop("token_z")
-        self.pairformer_module = PairformerModule(
-            pairformer_token_s,
-            pairformer_token_z,
-            **pairformer_kwargs,
-        )
-
-        diffusion_conditioning_kwargs = {
-            "token_s": token_s,
-            "token_z": token_z,
-            "atom_s": atom_s,
-            "atom_z": atom_z,
-            "atoms_per_window_queries": atoms_per_window_queries,
-            "atoms_per_window_keys": atoms_per_window_keys,
-            "atom_encoder_depth": score_model_args["atom_encoder_depth"],
-            "atom_encoder_heads": score_model_args["atom_encoder_heads"],
-            "token_transformer_depth": score_model_args["token_transformer_depth"],
-            "token_transformer_heads": score_model_args["token_transformer_heads"],
-            "atom_decoder_depth": score_model_args["atom_decoder_depth"],
-            "atom_decoder_heads": score_model_args["atom_decoder_heads"],
-            "atom_feature_dim": atom_feature_dim,
-            "conditioning_transition_layers": score_model_args["conditioning_transition_layers"],
-            "use_no_atom_char": use_no_atom_char,
-            "use_atom_backbone_feat": use_atom_backbone_feat,
-            "use_residue_feats_atoms": use_residue_feats_atoms,
-        }
-        diffusion_conditioning_kwargs = _filtered_kwargs(
-            DiffusionConditioning,
-            diffusion_conditioning_kwargs,
-        )
-        self.diffusion_conditioning = DiffusionConditioning(**diffusion_conditioning_kwargs)
-
-        structure_score_model_args = {
-            "token_s": token_s,
-            "atom_s": atom_s,
-            "atoms_per_window_queries": atoms_per_window_queries,
-            "atoms_per_window_keys": atoms_per_window_keys,
-            **score_model_args,
-        }
-        structure_score_model_args = _filtered_kwargs(
-            DiffusionModule,
-            structure_score_model_args,
-        )
-        structure_module_kwargs = {
-            "score_model_args": structure_score_model_args,
-            "compile_score": False,
-            **diffusion_process_args,
-        }
-        structure_module_kwargs = _filtered_kwargs(AtomDiffusion, structure_module_kwargs)
-        self.structure_module = AtomDiffusion(**structure_module_kwargs)
-        self.distogram_module = DistogramModule(token_z, num_bins)
-
-        if self.confidence_prediction:
-            if confidence_model_args is None:
-                raise ValueError(
-                    "confidence_prediction=True requires confidence_model_args in config."
-                )
-            confidence_kwargs = {
-                "token_s": token_s,
-                "token_z": token_z,
-                "token_level_confidence": token_level_confidence,
-                "bond_type_feature": bond_type_feature,
-                "fix_sym_check": fix_sym_check,
-                "cyclic_pos_enc": cyclic_pos_enc,
-                "conditioning_cutoff_min": conditioning_cutoff_min,
-                "conditioning_cutoff_max": conditioning_cutoff_max,
-                **confidence_model_args,
-            }
-            confidence_kwargs = _filtered_kwargs(ConfidenceModule, confidence_kwargs)
-            self.confidence_module = ConfidenceModule(**confidence_kwargs)
-
-    def forward(
-        self,
-        feats: dict[str, Tensor],
-        recycling_steps: int = 3,
-        num_sampling_steps: int | None = None,
-        diffusion_samples: int = 1,
-        max_parallel_samples: int | None = None,
-        run_confidence_sequentially: bool = True,
-        detach_confidence: bool = True,
-        verbose: bool = False,
-    ) -> dict[str, Tensor]:
-        # 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.
-        with tqdm(
-            total=1,
-            desc="Boltz2: Input embeddings",
-            disable=not verbose,
-            dynamic_ncols=True,
-            leave=False,
-        ) as progress:
-            s_inputs = self.input_embedder(feats)  # (b, t, d_s)
-            progress.update()
-        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()
-            )  # (b, t, t, d_z)
-        z_init = z_init + self.contact_conditioning(feats)  # (b, t, t, d_z)
-
-        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 tqdm(
-                range(recycling_steps + 1),
-                desc="Boltz2: Recycling",
-                unit="cycle",
-                disable=not verbose,
-                dynamic_ncols=True,
-                leave=False,
-            ):
-                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)  # (b, t, t, n_distogram, n_bin)
-        output: dict[str, Tensor] = {
-            "pdistogram": pdistogram,
-            "s": s,
-            "z": z,
-        }
-
-        if self.run_trunk_and_structure and (not self.skip_run_structure):
-            q, c, to_keys, atom_enc_bias, atom_dec_bias, token_trans_bias = (
-                self.diffusion_conditioning(
-                    s_trunk=s,
-                    z_trunk=z,
-                    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,
-                "to_keys": to_keys,
-                "atom_enc_bias": atom_enc_bias,
-                "atom_dec_bias": atom_dec_bias,
-                "token_trans_bias": token_trans_bias,
-            }
-            with torch.autocast("cuda", enabled=False):
-                struct_out = self.structure_module.sample(
-                    s_trunk=s.float(),
-                    s_inputs=s_inputs.float(),
-                    feats=feats,
-                    num_sampling_steps=num_sampling_steps,
-                    atom_mask=feats["atom_pad_mask"].float(),
-                    multiplicity=diffusion_samples,
-                    max_parallel_samples=max_parallel_samples,
-                    steering_args=self.steering_args,
-                    diffusion_conditioning=diffusion_conditioning,
-                    verbose=verbose,
-                )  # 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
-                )  # (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"]  # (b * m, a, 3)
-
-            if detach_confidence:
-                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  # (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)
-
-            with tqdm(
-                total=1,
-                desc="Boltz2: Confidence",
-                disable=not verbose,
-                dynamic_ncols=True,
-                leave=False,
-            ) as progress:
-                output.update(
-                    self.confidence_module(
-                        s_inputs=s_inputs_c,
-                        s=s_c,
-                        z=z_c,
-                        x_pred=x_pred_c,
-                        feats=feats,
-                        pred_distogram_logits=pdist_c,
-                        multiplicity=diffusion_samples,
-                        run_sequentially=run_confidence_sequentially,
-                        use_kernels=self.use_kernels,
-                    )
-                )
-                progress.update()
-
-        return output  # named tensors retain the shapes traced above
-
-
-class Boltz2Model(PreTrainedModel):
-    config_class = Boltz2Config
-    base_model_prefix = "core"
-    all_tied_weights_keys: ClassVar[dict[str, str]] = {}
-
-    def __init__(self, config: Boltz2Config) -> None:
-        super().__init__(config)
-        if not isinstance(config.core_kwargs, dict):
-            raise TypeError("config.core_kwargs must be a dictionary.")
-        self.core = Boltz2InferenceCore(**config.core_kwargs)
-
-    def _init_weights(self, module: nn.Module) -> None:
-        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] = {}
-        for key, tensor in raw_state.items():
-            if torch.is_tensor(tensor):
-                ptr = tensor.untyped_storage().data_ptr()
-                if ptr in seen_ptrs:
-                    out[key] = tensor.clone()  # (...)
-                else:
-                    seen_ptrs[ptr] = key
-                    out[key] = tensor  # (...)
-            else:
-                out[key] = tensor  # (...)
-        return out  # each tensor: (...)
-
-    def save_pretrained(self, save_directory: str, **kwargs: Any) -> None:
-        if "safe_serialization" not in kwargs:
-            kwargs["safe_serialization"] = True
-        if "state_dict" not in kwargs:
-            kwargs["state_dict"] = self._detied_state_dict()
-        super().save_pretrained(save_directory, **kwargs)
-
-    @property
-    def device(self) -> torch.device:
-        return next(self.parameters()).device
-
-    @classmethod
-    def from_boltz_checkpoint(
-        cls,
-        checkpoint_path: str,
-        map_location: str | torch.device = "cpu",
-        use_kernels: bool = False,
-        default_recycling_steps: int | None = None,
-        default_sampling_steps: int | None = None,
-        default_diffusion_samples: int | None = None,
-        *,
-        allow_unsafe_pickle: bool = False,
-    ) -> "Boltz2Model":
-        """Import an official Boltz Lightning checkpoint into the inference model.
-
-        Lightning checkpoints require Python pickle deserialization. The caller
-        must opt in with the literal value ``True`` after independently verifying
-        that the checkpoint comes from a trusted source and matches its expected
-        cryptographic hash.
-        """
-
-        if allow_unsafe_pickle is not True:
-            raise ValueError(
-                "Boltz Lightning checkpoints contain Python pickle and cannot be loaded "
-                "safely. Pass allow_unsafe_pickle=True only for a trusted, hash-verified "
-                "checkpoint."
-            )
-
-        # The official Lightning checkpoints contain OmegaConf objects, so this explicit
-        # trusted-checkpoint boundary requires full Python unpickling.
-        checkpoint = torch.load(
-            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")
-        _require_key(checkpoint, "state_dict")
-
-        hparams = checkpoint["hyper_parameters"]
-        if not isinstance(hparams, dict):
-            raise TypeError("Checkpoint hyper_parameters must be a dictionary.")
-        state_dict = checkpoint["state_dict"]  # each tensor: (...)
-        if not isinstance(state_dict, dict):
-            raise TypeError("Checkpoint state_dict must be a dictionary.")
-
-        config = Boltz2Config.from_hyperparameters(
-            hparams,
-            use_kernels=use_kernels,
-            default_recycling_steps=default_recycling_steps,
-            default_sampling_steps=default_sampling_steps,
-            default_diffusion_samples=default_diffusion_samples,
-        )
-        model = cls(config)
-        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:
-                raise RuntimeError(
-                    "Boltz2 inference core unexpectedly uses v1 attention parameters. "
-                    "Expected pairformer v2 architecture."
-                )
-        filtered: dict[str, Tensor] = {}
-        for key, value in cleaned.items():
-            if key in target_keys:
-                filtered[key] = value  # (...)
-
-        missing = sorted(target_keys.difference(filtered.keys()))
-        if missing:
-            raise RuntimeError(
-                "Checkpoint is missing required parameters for Boltz2 inference core. "
-                f"Missing keys (first 20): {missing[:20]}"
-            )
-
-        load_result = model.core.load_state_dict(filtered, strict=False)
-        loaded_missing = sorted(load_result.missing_keys)
-        if loaded_missing:
-            raise RuntimeError(
-                "Model has unexpected missing keys after load_state_dict. "
-                f"Missing keys (first 20): {loaded_missing[:20]}"
-            )
-        if load_result.unexpected_keys:
-            raise RuntimeError(
-                "Model has unexpected checkpoint keys after load_state_dict: "
-                f"{sorted(load_result.unexpected_keys)[:20]}"
-            )
-        model.eval()
-        return model
-
-    def forward(
-        self,
-        feats: dict[str, Tensor],
-        recycling_steps: int | None = None,
-        num_sampling_steps: int | None = None,
-        diffusion_samples: int | None = None,
-        max_parallel_samples: int | None = None,
-        run_confidence_sequentially: bool = True,
-        detach_confidence: bool = True,
-        output_attentions: bool | None = None,
-        output_hidden_states: bool | None = None,
-        return_dict: bool | None = None,
-        verbose: bool = False,
-    ) -> 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
-        )
-        if output_attentions:
-            raise NotImplementedError(
-                "Boltz2 does not expose normalized attention tensors from its structure "
-                "modules. output_attentions=True is unsupported."
-            )
-        output_hidden_states = (
-            self.config.output_hidden_states
-            if output_hidden_states is None
-            else output_hidden_states
-        )
-        return_dict = self.config.use_return_dict if return_dict is None else return_dict
-        if recycling_steps is None:
-            recycling_steps = self.config.default_recycling_steps
-        if num_sampling_steps is None:
-            num_sampling_steps = self.config.default_sampling_steps
-        if diffusion_samples is None:
-            diffusion_samples = self.config.default_diffusion_samples
-        raw_output = self.core(
-            feats=feats,
-            recycling_steps=recycling_steps,
-            num_sampling_steps=num_sampling_steps,
-            diffusion_samples=diffusion_samples,
-            max_parallel_samples=max_parallel_samples,
-            run_confidence_sequentially=run_confidence_sequentially,
-            detach_confidence=detach_confidence,
-            verbose=verbose,
-        )  # 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)
-            if output_hidden_states and token_state is not None and pair_state is not None
-            else None,
-            **raw_output,
-        )
-        return model_output if return_dict else model_output.to_tuple()
-
-    def _to_model_device(
-        self,
-        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
-                    )  # same shape (...)
-                else:
-                    moved[key] = value.to(device=self.device)  # same shape (...)
-            else:
-                moved[key] = value  # same shape (...)
-        return moved  # each tensor: same shape (...)
-
-    def predict_structure(
-        self,
-        amino_acid_sequence: str,
-        recycling_steps: int | None = None,
-        num_sampling_steps: int | None = None,
-        diffusion_samples: int | None = None,
-        max_parallel_samples: int | None = None,
-        run_confidence_sequentially: bool = True,
-        float_dtype: torch.dtype = torch.float32,
-        seed: int | None = None,
-        verbose: bool = False,
-    ) -> Boltz2StructureOutput:
-        if float_dtype != torch.float32:
-            raise ValueError(
-                "Boltz2 predict_structure requires FP32 features and parameter storage; "
-                "CUDA compute runs under the declared BF16 autocast policy."
-            )
-        parameter_dtypes = {
-            parameter.dtype for parameter in self.parameters() if parameter.is_floating_point()
-        }
-        if parameter_dtypes and parameter_dtypes != {torch.float32}:
-            raise ValueError(
-                "Boltz2 predict_structure requires FP32 parameter storage before CUDA "
-                f"BF16 autocast; found {sorted(map(str, parameter_dtypes))}."
-            )
-
-        with _seed_context(seed):
-            with tqdm(
-                total=1,
-                desc="Boltz2: Prepare features",
-                disable=not verbose,
-                dynamic_ncols=True,
-                leave=False,
-            ) as progress:
-                feats, template = build_boltz2_features(
-                    amino_acid_sequence=amino_acid_sequence,
-                    num_bins=self.config.num_bins,
-                    atoms_per_window_queries=(
-                        self.core.input_embedder.atom_encoder.atoms_per_window_queries
-                    ),
-                )  # feature shapes are traced in build_boltz2_features
-                feats = self._to_model_device(
-                    feats, float_dtype=torch.float32
-                )  # unchanged feature shapes
-                progress.update()
-            autocast_context = (
-                torch.autocast(device_type="cuda", dtype=torch.bfloat16)
-                if self.device.type == "cuda"
-                else nullcontext()
-            )
-            with torch.no_grad(), autocast_context:
-                output = self.forward(
-                    feats=feats,
-                    recycling_steps=recycling_steps,
-                    num_sampling_steps=num_sampling_steps,
-                    diffusion_samples=diffusion_samples,
-                    max_parallel_samples=max_parallel_samples,
-                    run_confidence_sequentially=run_confidence_sequentially,
-                    return_dict=True,
-                    verbose=verbose,
-                )  # named tensors use Boltz2InferenceCore.forward shapes
-
-        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()  # (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()  # (a_p,)
-        plddt = (
-            output["plddt"].detach().cpu() if "plddt" in output else None
-        )  # (m, t) or (m, a_p), depending on confidence resolution
-        complex_plddt = (
-            output["complex_plddt"].detach().cpu() if "complex_plddt" in output else 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  # (m,)
-            else:
-                confidence_score = (4 * complex_plddt + iptm) / 5  # (m,)
-
-        return Boltz2StructureOutput(
-            sample_atom_coords=sample_atom_coords,
-            atom_pad_mask=atom_pad_mask,
-            plddt=plddt,
-            confidence_score=confidence_score,
-            complex_plddt=complex_plddt,
-            iptm=iptm,
-            ptm=ptm,
-            sequence=template.sequence,
-            structure_template=template,
-            raw_output={key: _to_cpu_detached(val) for key, val in output.items()},
-            seed=seed,
-        )
-
-    def save_as_cif(
-        self,
-        structure_output: Boltz2StructureOutput,
-        output_path: str,
-        sample_index: int = 0,
-    ) -> str:
-        if structure_output.structure_template is None:
-            raise ValueError("structure_output.structure_template is required for CIF export.")
-        if structure_output.sample_atom_coords is None:
-            raise ValueError("structure_output.sample_atom_coords is required for CIF export.")
-        if structure_output.atom_pad_mask is None:
-            raise ValueError("structure_output.atom_pad_mask is required for CIF export.")
-        return write_cif(
-            structure_template=structure_output.structure_template,
-            atom_coords=structure_output.sample_atom_coords,
-            atom_mask=structure_output.atom_pad_mask,
-            output_path=output_path,
-            plddt=structure_output.plddt,
-            sample_index=sample_index,
-        )
+"""Transformers interfaces and checkpoint conversion for the Boltz2 runtime."""
+
+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
+from torch import Tensor
+from tqdm.auto import tqdm
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from . import vb_const as const
+from . import vb_layers_initialize as init
+from .cif_writer import write_cif
+from .minimal_featurizer import build_boltz2_features
+from .minimal_structures import ProteinStructureTemplate
+from .vb_const import bond_types as _vb_const_bond_types  # noqa: F401
+from .vb_layers_attention import AttentionPairBias as _vb_layers_attention_marker  # noqa: F401
+from .vb_layers_attentionv2 import AttentionPairBias as _vb_layers_attentionv2_marker  # noqa: F401
+from .vb_layers_confidence_utils import (
+    compute_ptms as _vb_layers_confidence_utils_marker,  # noqa: F401
+)
+from .vb_layers_dropout import get_dropout_mask as _vb_layers_dropout_marker  # noqa: F401
+from .vb_layers_initialize import gating_init_ as _vb_layers_initialize_marker  # noqa: F401
+from .vb_layers_outer_product_mean import (
+    OuterProductMean as _vb_layers_outer_product_mean_marker,  # noqa: F401
+)
+from .vb_layers_pair_averaging import (
+    PairWeightedAveraging as _vb_layers_pair_averaging_marker,  # noqa: F401
+)
+from .vb_layers_pairformer import PairformerModule
+from .vb_layers_transition import Transition as _vb_layers_transition_marker  # noqa: F401
+from .vb_layers_triangular_mult import (
+    TriangleMultiplicationIncoming as _vb_layers_triangular_mult_marker,  # noqa: F401
+)
+from .vb_loss_diffusionv2 import weighted_rigid_align as _vb_loss_diffusionv2_marker  # noqa: F401
+from .vb_modules_confidencev2 import ConfidenceModule
+from .vb_modules_diffusion_conditioning import DiffusionConditioning
+from .vb_modules_diffusionv2 import AtomDiffusion, DiffusionModule
+from .vb_modules_encodersv2 import RelativePositionEncoder
+from .vb_modules_transformersv2 import (
+    DiffusionTransformer as _vb_modules_transformersv2_marker,  # noqa: F401
+)
+from .vb_modules_trunkv2 import (
+    ContactConditioning,
+    DistogramModule,
+    InputEmbedder,
+    MSAModule,
+)
+from .vb_modules_utils import LinearNoBias as _vb_modules_utils_marker  # noqa: F401
+from .vb_potentials_potentials import (
+    get_potentials as _vb_potentials_potentials_marker,  # noqa: F401
+)
+from .vb_potentials_schedules import (
+    ParameterSchedule as _vb_potentials_schedules_marker,  # noqa: F401
+)
+from .vb_tri_attn_attention import (
+    TriangleAttentionStartingNode as _vb_tri_attn_attention_marker,  # noqa: F401
+)
+from .vb_tri_attn_primitives import Attention as _vb_tri_attn_primitives_marker  # noqa: F401
+from .vb_tri_attn_utils import permute_final_dims as _vb_tri_attn_utils_marker  # noqa: F401
+
+
+def _default_steering_args() -> dict[str, Any]:
+    return {
+        "fk_steering": False,
+        "num_particles": 3,
+        "fk_lambda": 4.0,
+        "fk_resampling_interval": 3,
+        "physical_guidance_update": False,
+        "contact_guidance_update": False,
+        "num_gd_steps": 16,
+    }
+
+
+def _boltz2_reference_diffusion_overrides() -> dict[str, Any]:
+    # Match Boltz2 CLI inference defaults from boltz.main/Boltz2DiffusionParams.
+    return {
+        "gamma_0": 0.8,
+        "gamma_min": 1.0,
+        "noise_scale": 1.003,
+        "rho": 7,
+        "step_scale": 1.5,
+        "sigma_min": 0.0001,
+        "sigma_max": 160.0,
+        "sigma_data": 16.0,
+        "P_mean": -1.2,
+        "P_std": 1.5,
+        "coordinate_augmentation": True,
+        "alignment_reverse_diff": True,
+        "synchronize_sigmas": True,
+    }
+
+
+def _enforce_pairformer_v2(pairformer_args: Mapping[str, Any], context: str) -> dict[str, Any]:
+    if not isinstance(pairformer_args, Mapping):
+        raise TypeError(f"Expected {context} pairformer_args to be a dictionary.")
+    out = _to_plain_python(copy.deepcopy(pairformer_args))
+    if "v2" in out and not out["v2"]:
+        raise ValueError(f"{context} pairformer_args['v2'] must be True for Boltz2.")
+    out["v2"] = True
+    return out
+
+
+def _require_key(mapping: dict[str, Any], key: str) -> Any:
+    if key not in mapping:
+        raise KeyError(f"Missing required key '{key}' in checkpoint hyperparameters.")
+    return mapping[key]
+
+
+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."):
+            continue
+        new_key = key
+        if new_key.startswith("model."):
+            new_key = new_key[len("model.") :]
+        if new_key.startswith("module."):
+            new_key = new_key[len("module.") :]
+        cleaned[new_key] = value  # (...)
+    return cleaned
+
+
+def _to_cpu_detached(value: Any) -> Any:
+    if torch.is_tensor(value):
+        return value.detach().cpu()  # same shape (...)
+    if isinstance(value, dict):
+        out: dict[Any, Any] = {}
+        for key, nested_value in value.items():
+            out[key] = _to_cpu_detached(nested_value)
+        return out
+    if isinstance(value, list):
+        return [_to_cpu_detached(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_to_cpu_detached(item) for item in value)
+    return value
+
+
+@dataclass(frozen=True)
+class _RandomState:
+    python: object
+    numpy: tuple[Any, ...]
+    torch_cpu: Tensor  # (n_cpu_rng_state,)
+    torch_cuda: list[Tensor] | None  # each: (n_cuda_rng_state,)
+
+
+@contextmanager
+def _seed_context(seed: int | None) -> Iterator[None]:
+    """Temporarily seed every RNG used by Boltz2 and restore caller state."""
+
+    if seed is None:
+        yield
+        return
+    if isinstance(seed, bool) or not isinstance(seed, int):
+        raise TypeError(
+            "Boltz2 seed must be an int or None; bool and coerced values are not accepted."
+        )
+    state = _RandomState(
+        python=random.getstate(),
+        numpy=np.random.get_state(),
+        torch_cpu=torch.random.get_rng_state(),
+        torch_cuda=torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
+    )
+    normalized_seed = seed % (2**32)
+    random.seed(normalized_seed)
+    np.random.seed(normalized_seed)
+    torch.manual_seed(normalized_seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(normalized_seed)
+    try:
+        yield
+    finally:
+        random.setstate(state.python)
+        np.random.set_state(state.numpy)
+        torch.random.set_rng_state(state.torch_cpu)
+        if state.torch_cuda is not None:
+            torch.cuda.set_rng_state_all(state.torch_cuda)
+
+
+def _to_plain_python(value: Any) -> Any:
+    if isinstance(value, Mapping):
+        out: dict[Any, Any] = {}
+        for key, nested_value in value.items():
+            out[key] = _to_plain_python(nested_value)
+        return out
+    if isinstance(value, list):
+        return [_to_plain_python(item) for item in value]
+    if isinstance(value, tuple):
+        return [_to_plain_python(item) for item in value]
+    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
+        return [_to_plain_python(item) for item in value]
+    return value
+
+
+def _filtered_kwargs(target: Any, kwargs: dict[str, Any]) -> dict[str, Any]:
+    signature = inspect.signature(target.__init__)
+    allowed = set(signature.parameters.keys())
+    allowed.discard("self")
+    filtered: dict[str, Any] = {}
+    for key, value in kwargs.items():
+        if key in allowed:
+            filtered[key] = value
+    return filtered
+
+
+@dataclass
+class Boltz2ModelOutput(ModelOutput):
+    """Raw Boltz2 inference output with standard AutoModel controls."""
+
+    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  # (m, a_p, 3) or None
+    atom_pad_mask: torch.Tensor | None = None  # (a_p,) or None
+    plddt: torch.Tensor | None = None  # (m, t) or (m, a_p), depending on confidence resolution
+    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  # values keep model-output shapes
+    seed: int | None = None
+
+
+class Boltz2Config(PretrainedConfig):
+    model_type = "boltz2_automodel"
+
+    def __init__(
+        self,
+        core_kwargs: dict[str, Any] | None = None,
+        num_bins: int = 64,
+        default_recycling_steps: int = 3,
+        default_sampling_steps: int = 200,
+        default_diffusion_samples: int = 1,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        if core_kwargs is None:
+            core_kwargs = {}
+        self.core_kwargs = core_kwargs
+        self.num_bins = num_bins
+        self.default_recycling_steps = default_recycling_steps
+        self.default_sampling_steps = default_sampling_steps
+        self.default_diffusion_samples = default_diffusion_samples
+
+    @classmethod
+    def from_hyperparameters(
+        cls,
+        hparams: dict[str, Any],
+        use_kernels: bool = False,
+        default_recycling_steps: int | None = None,
+        default_sampling_steps: int | None = None,
+        default_diffusion_samples: int | None = None,
+    ) -> "Boltz2Config":
+        if not isinstance(hparams, dict):
+            raise TypeError("Expected checkpoint hyperparameters as a dictionary.")
+        required = [
+            "atom_s",
+            "atom_z",
+            "token_s",
+            "token_z",
+            "num_bins",
+            "embedder_args",
+            "msa_args",
+            "pairformer_args",
+            "score_model_args",
+            "diffusion_process_args",
+        ]
+        for key in required:
+            _require_key(hparams, key)
+
+        pairformer_args = _enforce_pairformer_v2(
+            hparams["pairformer_args"],
+            context="checkpoint",
+        )
+        diffusion_process_args = _to_plain_python(copy.deepcopy(hparams["diffusion_process_args"]))
+        diffusion_overrides = _boltz2_reference_diffusion_overrides()
+        for key in diffusion_overrides:
+            diffusion_process_args[key] = diffusion_overrides[key]
+
+        core_kwargs: dict[str, Any] = {
+            "atom_s": hparams["atom_s"],
+            "atom_z": hparams["atom_z"],
+            "token_s": hparams["token_s"],
+            "token_z": hparams["token_z"],
+            "num_bins": hparams["num_bins"],
+            "embedder_args": _to_plain_python(copy.deepcopy(hparams["embedder_args"])),
+            "msa_args": _to_plain_python(copy.deepcopy(hparams["msa_args"])),
+            "pairformer_args": pairformer_args,
+            "score_model_args": _to_plain_python(copy.deepcopy(hparams["score_model_args"])),
+            "diffusion_process_args": diffusion_process_args,
+            "use_kernels": use_kernels,
+        }
+
+        if "confidence_model_args" in hparams:
+            confidence_model_args = _to_plain_python(
+                copy.deepcopy(hparams["confidence_model_args"])
+            )
+            if "pairformer_args" in confidence_model_args:
+                confidence_model_args["pairformer_args"] = _enforce_pairformer_v2(
+                    confidence_model_args["pairformer_args"],
+                    context="confidence",
+                )
+            core_kwargs["confidence_model_args"] = confidence_model_args
+        else:
+            core_kwargs["confidence_model_args"] = None
+
+        core_kwargs["confidence_prediction"] = hparams.get("confidence_prediction", True)
+        core_kwargs["token_level_confidence"] = hparams.get("token_level_confidence", True)
+        core_kwargs["alpha_pae"] = hparams.get("alpha_pae", 0.0)
+        core_kwargs["atoms_per_window_queries"] = hparams.get("atoms_per_window_queries", 32)
+        core_kwargs["atoms_per_window_keys"] = hparams.get("atoms_per_window_keys", 128)
+        core_kwargs["atom_feature_dim"] = hparams.get("atom_feature_dim", 128)
+        core_kwargs["bond_type_feature"] = hparams.get("bond_type_feature", False)
+        core_kwargs["run_trunk_and_structure"] = hparams.get("run_trunk_and_structure", True)
+        core_kwargs["skip_run_structure"] = hparams.get("skip_run_structure", False)
+        core_kwargs["fix_sym_check"] = hparams.get("fix_sym_check", False)
+        core_kwargs["cyclic_pos_enc"] = hparams.get("cyclic_pos_enc", False)
+        core_kwargs["use_no_atom_char"] = hparams.get("use_no_atom_char", False)
+        core_kwargs["use_atom_backbone_feat"] = hparams.get("use_atom_backbone_feat", False)
+        core_kwargs["use_residue_feats_atoms"] = hparams.get("use_residue_feats_atoms", False)
+        core_kwargs["conditioning_cutoff_min"] = hparams.get("conditioning_cutoff_min", 4.0)
+        core_kwargs["conditioning_cutoff_max"] = hparams.get("conditioning_cutoff_max", 20.0)
+
+        if "steering_args" in hparams and hparams["steering_args"] is not None:
+            core_kwargs["steering_args"] = _to_plain_python(copy.deepcopy(hparams["steering_args"]))
+        else:
+            core_kwargs["steering_args"] = _default_steering_args()
+
+        if "validation_args" in hparams:
+            validation_args = hparams["validation_args"]
+            if not isinstance(validation_args, Mapping):
+                raise TypeError(
+                    "Expected 'validation_args' in checkpoint hyperparameters to be a mapping."
+                )
+            if default_recycling_steps is None and "recycling_steps" in validation_args:
+                default_recycling_steps = validation_args["recycling_steps"]
+            if default_sampling_steps is None and "sampling_steps" in validation_args:
+                default_sampling_steps = validation_args["sampling_steps"]
+            if default_diffusion_samples is None and "diffusion_samples" in validation_args:
+                default_diffusion_samples = validation_args["diffusion_samples"]
+
+        if default_recycling_steps is None:
+            default_recycling_steps = 3
+        if default_sampling_steps is None:
+            default_sampling_steps = 200
+        if default_diffusion_samples is None:
+            default_diffusion_samples = 1
+
+        return cls(
+            core_kwargs=core_kwargs,
+            num_bins=hparams["num_bins"],
+            default_recycling_steps=default_recycling_steps,
+            default_sampling_steps=default_sampling_steps,
+            default_diffusion_samples=default_diffusion_samples,
+        )
+
+
+class Boltz2InferenceCore(nn.Module):
+    def __init__(
+        self,
+        atom_s: int,
+        atom_z: int,
+        token_s: int,
+        token_z: int,
+        num_bins: int,
+        embedder_args: dict[str, Any],
+        msa_args: dict[str, Any],
+        pairformer_args: dict[str, Any],
+        score_model_args: dict[str, Any],
+        diffusion_process_args: dict[str, Any],
+        confidence_model_args: dict[str, Any] | None = None,
+        atom_feature_dim: int = 128,
+        confidence_prediction: bool = True,
+        token_level_confidence: bool = True,
+        alpha_pae: float = 0.0,
+        atoms_per_window_queries: int = 32,
+        atoms_per_window_keys: int = 128,
+        run_trunk_and_structure: bool = True,
+        skip_run_structure: bool = False,
+        bond_type_feature: bool = False,
+        fix_sym_check: bool = False,
+        cyclic_pos_enc: bool = False,
+        use_no_atom_char: bool = False,
+        use_atom_backbone_feat: bool = False,
+        use_residue_feats_atoms: bool = False,
+        conditioning_cutoff_min: float = 4.0,
+        conditioning_cutoff_max: float = 20.0,
+        use_kernels: bool = False,
+        steering_args: dict[str, Any] | None = None,
+    ) -> None:
+        super().__init__()
+        self.use_kernels = use_kernels
+        self.confidence_prediction = confidence_prediction
+        self.token_level_confidence = token_level_confidence
+        self.alpha_pae = alpha_pae
+        self.run_trunk_and_structure = run_trunk_and_structure
+        self.skip_run_structure = skip_run_structure
+        self.bond_type_feature = bond_type_feature
+        self.steering_args = (
+            steering_args if steering_args is not None else _default_steering_args()
+        )
+        if "v2" not in pairformer_args:
+            raise ValueError("Boltz2 requires pairformer_args['v2'].")
+        if not pairformer_args["v2"]:
+            raise ValueError("Boltz2 requires pairformer_args['v2']=True.")
+
+        full_embedder_args = {
+            "atom_s": atom_s,
+            "atom_z": atom_z,
+            "token_s": token_s,
+            "token_z": token_z,
+            "atoms_per_window_queries": atoms_per_window_queries,
+            "atoms_per_window_keys": atoms_per_window_keys,
+            "atom_feature_dim": atom_feature_dim,
+            "use_no_atom_char": use_no_atom_char,
+            "use_atom_backbone_feat": use_atom_backbone_feat,
+            "use_residue_feats_atoms": use_residue_feats_atoms,
+            **embedder_args,
+        }
+        full_embedder_args = _filtered_kwargs(InputEmbedder, full_embedder_args)
+        self.input_embedder = InputEmbedder(**full_embedder_args)
+
+        self.s_init = nn.Linear(token_s, token_s, bias=False)
+        self.z_init_1 = nn.Linear(token_s, token_z, bias=False)
+        self.z_init_2 = nn.Linear(token_s, token_z, bias=False)
+        self.rel_pos = RelativePositionEncoder(
+            token_z,
+            fix_sym_check=fix_sym_check,
+            cyclic_pos_enc=cyclic_pos_enc,
+        )
+        self.token_bonds = nn.Linear(1, token_z, bias=False)
+        if self.bond_type_feature:
+            self.token_bonds_type = nn.Embedding(len(const.bond_types) + 1, token_z)
+
+        self.contact_conditioning = ContactConditioning(
+            token_z=token_z,
+            cutoff_min=conditioning_cutoff_min,
+            cutoff_max=conditioning_cutoff_max,
+        )
+        self.s_norm = nn.LayerNorm(token_s)
+        self.z_norm = nn.LayerNorm(token_z)
+
+        self.s_recycle = nn.Linear(token_s, token_s, bias=False)
+        self.z_recycle = nn.Linear(token_z, token_z, bias=False)
+        init.gating_init_(self.s_recycle.weight)
+        init.gating_init_(self.z_recycle.weight)
+
+        msa_kwargs = _filtered_kwargs(
+            MSAModule,
+            {"token_z": token_z, "token_s": token_s, **msa_args},
+        )
+        self.msa_module = MSAModule(**msa_kwargs)
+
+        pairformer_kwargs = _filtered_kwargs(
+            PairformerModule,
+            {"token_s": token_s, "token_z": token_z, **pairformer_args},
+        )
+        if "token_s" not in pairformer_kwargs or "token_z" not in pairformer_kwargs:
+            raise RuntimeError("Boltz2 pairformer construction lost its token dimensions.")
+        pairformer_token_s = pairformer_kwargs.pop("token_s")
+        pairformer_token_z = pairformer_kwargs.pop("token_z")
+        self.pairformer_module = PairformerModule(
+            pairformer_token_s,
+            pairformer_token_z,
+            **pairformer_kwargs,
+        )
+
+        diffusion_conditioning_kwargs = {
+            "token_s": token_s,
+            "token_z": token_z,
+            "atom_s": atom_s,
+            "atom_z": atom_z,
+            "atoms_per_window_queries": atoms_per_window_queries,
+            "atoms_per_window_keys": atoms_per_window_keys,
+            "atom_encoder_depth": score_model_args["atom_encoder_depth"],
+            "atom_encoder_heads": score_model_args["atom_encoder_heads"],
+            "token_transformer_depth": score_model_args["token_transformer_depth"],
+            "token_transformer_heads": score_model_args["token_transformer_heads"],
+            "atom_decoder_depth": score_model_args["atom_decoder_depth"],
+            "atom_decoder_heads": score_model_args["atom_decoder_heads"],
+            "atom_feature_dim": atom_feature_dim,
+            "conditioning_transition_layers": score_model_args["conditioning_transition_layers"],
+            "use_no_atom_char": use_no_atom_char,
+            "use_atom_backbone_feat": use_atom_backbone_feat,
+            "use_residue_feats_atoms": use_residue_feats_atoms,
+        }
+        diffusion_conditioning_kwargs = _filtered_kwargs(
+            DiffusionConditioning,
+            diffusion_conditioning_kwargs,
+        )
+        self.diffusion_conditioning = DiffusionConditioning(**diffusion_conditioning_kwargs)
+
+        structure_score_model_args = {
+            "token_s": token_s,
+            "atom_s": atom_s,
+            "atoms_per_window_queries": atoms_per_window_queries,
+            "atoms_per_window_keys": atoms_per_window_keys,
+            **score_model_args,
+        }
+        structure_score_model_args = _filtered_kwargs(
+            DiffusionModule,
+            structure_score_model_args,
+        )
+        structure_module_kwargs = {
+            "score_model_args": structure_score_model_args,
+            "compile_score": False,
+            **diffusion_process_args,
+        }
+        structure_module_kwargs = _filtered_kwargs(AtomDiffusion, structure_module_kwargs)
+        self.structure_module = AtomDiffusion(**structure_module_kwargs)
+        self.distogram_module = DistogramModule(token_z, num_bins)
+
+        if self.confidence_prediction:
+            if confidence_model_args is None:
+                raise ValueError(
+                    "confidence_prediction=True requires confidence_model_args in config."
+                )
+            confidence_kwargs = {
+                "token_s": token_s,
+                "token_z": token_z,
+                "token_level_confidence": token_level_confidence,
+                "bond_type_feature": bond_type_feature,
+                "fix_sym_check": fix_sym_check,
+                "cyclic_pos_enc": cyclic_pos_enc,
+                "conditioning_cutoff_min": conditioning_cutoff_min,
+                "conditioning_cutoff_max": conditioning_cutoff_max,
+                **confidence_model_args,
+            }
+            confidence_kwargs = _filtered_kwargs(ConfidenceModule, confidence_kwargs)
+            self.confidence_module = ConfidenceModule(**confidence_kwargs)
+
+    def forward(
+        self,
+        feats: dict[str, Tensor],
+        recycling_steps: int = 3,
+        num_sampling_steps: int | None = None,
+        diffusion_samples: int = 1,
+        max_parallel_samples: int | None = None,
+        run_confidence_sequentially: bool = True,
+        detach_confidence: bool = True,
+        verbose: bool = False,
+    ) -> dict[str, Tensor]:
+        # 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.
+        with tqdm(
+            total=1,
+            desc="Boltz2: Input embeddings",
+            disable=not verbose,
+            dynamic_ncols=True,
+            leave=False,
+        ) as progress:
+            s_inputs = self.input_embedder(feats)  # (b, t, d_s)
+            progress.update()
+        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()
+            )  # (b, t, t, d_z)
+        z_init = z_init + self.contact_conditioning(feats)  # (b, t, t, d_z)
+
+        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 tqdm(
+                range(recycling_steps + 1),
+                desc="Boltz2: Recycling",
+                unit="cycle",
+                disable=not verbose,
+                dynamic_ncols=True,
+                leave=False,
+            ):
+                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)  # (b, t, t, n_distogram, n_bin)
+        output: dict[str, Tensor] = {
+            "pdistogram": pdistogram,
+            "s": s,
+            "z": z,
+        }
+
+        if self.run_trunk_and_structure and (not self.skip_run_structure):
+            q, c, to_keys, atom_enc_bias, atom_dec_bias, token_trans_bias = (
+                self.diffusion_conditioning(
+                    s_trunk=s,
+                    z_trunk=z,
+                    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,
+                "to_keys": to_keys,
+                "atom_enc_bias": atom_enc_bias,
+                "atom_dec_bias": atom_dec_bias,
+                "token_trans_bias": token_trans_bias,
+            }
+            with torch.autocast("cuda", enabled=False):
+                struct_out = self.structure_module.sample(
+                    s_trunk=s.float(),
+                    s_inputs=s_inputs.float(),
+                    feats=feats,
+                    num_sampling_steps=num_sampling_steps,
+                    atom_mask=feats["atom_pad_mask"].float(),
+                    multiplicity=diffusion_samples,
+                    max_parallel_samples=max_parallel_samples,
+                    steering_args=self.steering_args,
+                    diffusion_conditioning=diffusion_conditioning,
+                    verbose=verbose,
+                )  # 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
+                )  # (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"]  # (b * m, a, 3)
+
+            if detach_confidence:
+                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  # (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)
+
+            with tqdm(
+                total=1,
+                desc="Boltz2: Confidence",
+                disable=not verbose,
+                dynamic_ncols=True,
+                leave=False,
+            ) as progress:
+                output.update(
+                    self.confidence_module(
+                        s_inputs=s_inputs_c,
+                        s=s_c,
+                        z=z_c,
+                        x_pred=x_pred_c,
+                        feats=feats,
+                        pred_distogram_logits=pdist_c,
+                        multiplicity=diffusion_samples,
+                        run_sequentially=run_confidence_sequentially,
+                        use_kernels=self.use_kernels,
+                    )
+                )
+                progress.update()
+
+        return output  # named tensors retain the shapes traced above
+
+
+class Boltz2Model(PreTrainedModel):
+    config_class = Boltz2Config
+    base_model_prefix = "core"
+    all_tied_weights_keys: ClassVar[dict[str, str]] = {}
+
+    def __init__(self, config: Boltz2Config) -> None:
+        super().__init__(config)
+        if not isinstance(config.core_kwargs, dict):
+            raise TypeError("config.core_kwargs must be a dictionary.")
+        self.core = Boltz2InferenceCore(**config.core_kwargs)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        # Parameters keep their constructed or loaded values. Transformers 5 builds the
+        # model on the meta device and leaves a buffer no checkpoint holds uninitialized,
+        # so the diffusion module's non-persistent zero, which the loss breakdown
+        # reports when the smooth lDDT loss is off, is refilled here.
+        if isinstance(module, AtomDiffusion):
+            module.zero.zero_()
+
+    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] = {}
+        for key, tensor in raw_state.items():
+            if torch.is_tensor(tensor):
+                ptr = tensor.untyped_storage().data_ptr()
+                if ptr in seen_ptrs:
+                    out[key] = tensor.clone()  # (...)
+                else:
+                    seen_ptrs[ptr] = key
+                    out[key] = tensor  # (...)
+            else:
+                out[key] = tensor  # (...)
+        return out  # each tensor: (...)
+
+    def save_pretrained(self, save_directory: str, **kwargs: Any) -> None:
+        if "safe_serialization" not in kwargs:
+            kwargs["safe_serialization"] = True
+        if "state_dict" not in kwargs:
+            kwargs["state_dict"] = self._detied_state_dict()
+        super().save_pretrained(save_directory, **kwargs)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @classmethod
+    def from_boltz_checkpoint(
+        cls,
+        checkpoint_path: str,
+        map_location: str | torch.device = "cpu",
+        use_kernels: bool = False,
+        default_recycling_steps: int | None = None,
+        default_sampling_steps: int | None = None,
+        default_diffusion_samples: int | None = None,
+        *,
+        allow_unsafe_pickle: bool = False,
+    ) -> "Boltz2Model":
+        """Import an official Boltz Lightning checkpoint into the inference model.
+
+        Lightning checkpoints require Python pickle deserialization. The caller
+        must opt in with the literal value ``True`` after independently verifying
+        that the checkpoint comes from a trusted source and matches its expected
+        cryptographic hash.
+        """
+
+        if allow_unsafe_pickle is not True:
+            raise ValueError(
+                "Boltz Lightning checkpoints contain Python pickle and cannot be loaded "
+                "safely. Pass allow_unsafe_pickle=True only for a trusted, hash-verified "
+                "checkpoint."
+            )
+
+        # The official Lightning checkpoints contain OmegaConf objects, so this explicit
+        # trusted-checkpoint boundary requires full Python unpickling.
+        checkpoint = torch.load(
+            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")
+        _require_key(checkpoint, "state_dict")
+
+        hparams = checkpoint["hyper_parameters"]
+        if not isinstance(hparams, dict):
+            raise TypeError("Checkpoint hyper_parameters must be a dictionary.")
+        state_dict = checkpoint["state_dict"]  # each tensor: (...)
+        if not isinstance(state_dict, dict):
+            raise TypeError("Checkpoint state_dict must be a dictionary.")
+
+        config = Boltz2Config.from_hyperparameters(
+            hparams,
+            use_kernels=use_kernels,
+            default_recycling_steps=default_recycling_steps,
+            default_sampling_steps=default_sampling_steps,
+            default_diffusion_samples=default_diffusion_samples,
+        )
+        model = cls(config)
+        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:
+                raise RuntimeError(
+                    "Boltz2 inference core unexpectedly uses v1 attention parameters. "
+                    "Expected pairformer v2 architecture."
+                )
+        filtered: dict[str, Tensor] = {}
+        for key, value in cleaned.items():
+            if key in target_keys:
+                filtered[key] = value  # (...)
+
+        missing = sorted(target_keys.difference(filtered.keys()))
+        if missing:
+            raise RuntimeError(
+                "Checkpoint is missing required parameters for Boltz2 inference core. "
+                f"Missing keys (first 20): {missing[:20]}"
+            )
+
+        load_result = model.core.load_state_dict(filtered, strict=False)
+        loaded_missing = sorted(load_result.missing_keys)
+        if loaded_missing:
+            raise RuntimeError(
+                "Model has unexpected missing keys after load_state_dict. "
+                f"Missing keys (first 20): {loaded_missing[:20]}"
+            )
+        if load_result.unexpected_keys:
+            raise RuntimeError(
+                "Model has unexpected checkpoint keys after load_state_dict: "
+                f"{sorted(load_result.unexpected_keys)[:20]}"
+            )
+        model.eval()
+        return model
+
+    def forward(
+        self,
+        feats: dict[str, Tensor],
+        recycling_steps: int | None = None,
+        num_sampling_steps: int | None = None,
+        diffusion_samples: int | None = None,
+        max_parallel_samples: int | None = None,
+        run_confidence_sequentially: bool = True,
+        detach_confidence: bool = True,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        verbose: bool = False,
+    ) -> 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
+        )
+        if output_attentions:
+            raise NotImplementedError(
+                "Boltz2 does not expose normalized attention tensors from its structure "
+                "modules. output_attentions=True is unsupported."
+            )
+        output_hidden_states = (
+            self.config.output_hidden_states
+            if output_hidden_states is None
+            else output_hidden_states
+        )
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+        if recycling_steps is None:
+            recycling_steps = self.config.default_recycling_steps
+        if num_sampling_steps is None:
+            num_sampling_steps = self.config.default_sampling_steps
+        if diffusion_samples is None:
+            diffusion_samples = self.config.default_diffusion_samples
+        raw_output = self.core(
+            feats=feats,
+            recycling_steps=recycling_steps,
+            num_sampling_steps=num_sampling_steps,
+            diffusion_samples=diffusion_samples,
+            max_parallel_samples=max_parallel_samples,
+            run_confidence_sequentially=run_confidence_sequentially,
+            detach_confidence=detach_confidence,
+            verbose=verbose,
+        )  # 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)
+            if output_hidden_states and token_state is not None and pair_state is not None
+            else None,
+            **raw_output,
+        )
+        return model_output if return_dict else model_output.to_tuple()
+
+    def _to_model_device(
+        self,
+        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
+                    )  # same shape (...)
+                else:
+                    moved[key] = value.to(device=self.device)  # same shape (...)
+            else:
+                moved[key] = value  # same shape (...)
+        return moved  # each tensor: same shape (...)
+
+    def predict_structure(
+        self,
+        amino_acid_sequence: str,
+        recycling_steps: int | None = None,
+        num_sampling_steps: int | None = None,
+        diffusion_samples: int | None = None,
+        max_parallel_samples: int | None = None,
+        run_confidence_sequentially: bool = True,
+        float_dtype: torch.dtype = torch.float32,
+        seed: int | None = None,
+        verbose: bool = False,
+    ) -> Boltz2StructureOutput:
+        if float_dtype != torch.float32:
+            raise ValueError(
+                "Boltz2 predict_structure requires FP32 features and parameter storage; "
+                "CUDA compute runs under the declared BF16 autocast policy."
+            )
+        parameter_dtypes = {
+            parameter.dtype for parameter in self.parameters() if parameter.is_floating_point()
+        }
+        if parameter_dtypes and parameter_dtypes != {torch.float32}:
+            raise ValueError(
+                "Boltz2 predict_structure requires FP32 parameter storage before CUDA "
+                f"BF16 autocast; found {sorted(map(str, parameter_dtypes))}."
+            )
+
+        with _seed_context(seed):
+            with tqdm(
+                total=1,
+                desc="Boltz2: Prepare features",
+                disable=not verbose,
+                dynamic_ncols=True,
+                leave=False,
+            ) as progress:
+                feats, template = build_boltz2_features(
+                    amino_acid_sequence=amino_acid_sequence,
+                    num_bins=self.config.num_bins,
+                    atoms_per_window_queries=(
+                        self.core.input_embedder.atom_encoder.atoms_per_window_queries
+                    ),
+                )  # feature shapes are traced in build_boltz2_features
+                feats = self._to_model_device(
+                    feats, float_dtype=torch.float32
+                )  # unchanged feature shapes
+                progress.update()
+            autocast_context = (
+                torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+                if self.device.type == "cuda"
+                else nullcontext()
+            )
+            with torch.no_grad(), autocast_context:
+                output = self.forward(
+                    feats=feats,
+                    recycling_steps=recycling_steps,
+                    num_sampling_steps=num_sampling_steps,
+                    diffusion_samples=diffusion_samples,
+                    max_parallel_samples=max_parallel_samples,
+                    run_confidence_sequentially=run_confidence_sequentially,
+                    return_dict=True,
+                    verbose=verbose,
+                )  # named tensors use Boltz2InferenceCore.forward shapes
+
+        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()  # (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()  # (a_p,)
+        plddt = (
+            output["plddt"].detach().cpu() if "plddt" in output else None
+        )  # (m, t) or (m, a_p), depending on confidence resolution
+        complex_plddt = (
+            output["complex_plddt"].detach().cpu() if "complex_plddt" in output else 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  # (m,)
+            else:
+                confidence_score = (4 * complex_plddt + iptm) / 5  # (m,)
+
+        return Boltz2StructureOutput(
+            sample_atom_coords=sample_atom_coords,
+            atom_pad_mask=atom_pad_mask,
+            plddt=plddt,
+            confidence_score=confidence_score,
+            complex_plddt=complex_plddt,
+            iptm=iptm,
+            ptm=ptm,
+            sequence=template.sequence,
+            structure_template=template,
+            raw_output={key: _to_cpu_detached(val) for key, val in output.items()},
+            seed=seed,
+        )
+
+    def save_as_cif(
+        self,
+        structure_output: Boltz2StructureOutput,
+        output_path: str,
+        sample_index: int = 0,
+    ) -> str:
+        if structure_output.structure_template is None:
+            raise ValueError("structure_output.structure_template is required for CIF export.")
+        if structure_output.sample_atom_coords is None:
+            raise ValueError("structure_output.sample_atom_coords is required for CIF export.")
+        if structure_output.atom_pad_mask is None:
+            raise ValueError("structure_output.atom_pad_mask is required for CIF export.")
+        return write_cif(
+            structure_template=structure_output.structure_template,
+            atom_coords=structure_output.sample_atom_coords,
+            atom_mask=structure_output.atom_pad_mask,
+            output_path=output_path,
+            plddt=structure_output.plddt,
+            sample_index=sample_index,
+        )

fastplms/models/boltz/vb_const.py

@@ -1,353 +1,353 @@
-"""Minimal biological constants required by the local Boltz2 runtime.
-
-The upstream project contains additional training-time curation tables.  They
-are intentionally excluded here because FastPLMs neither trains Boltz2 nor
-uses upstream data-processing code.  Keeping only runtime inputs makes this
-module auditable and prevents parity-oracle data from becoming a package
-dependency.
-"""
-
-from __future__ import annotations
-
-
-def _index(values: list[str]) -> dict[str, int]:
-    return {value: index for index, value in enumerate(values)}
-
-
-chain_types = ["PROTEIN", "DNA", "RNA", "NONPOLYMER"]
-chain_type_ids = _index(chain_types)
-
-
-canonical_tokens = [
-    "ALA",
-    "ARG",
-    "ASN",
-    "ASP",
-    "CYS",
-    "GLN",
-    "GLU",
-    "GLY",
-    "HIS",
-    "ILE",
-    "LEU",
-    "LYS",
-    "MET",
-    "PHE",
-    "PRO",
-    "SER",
-    "THR",
-    "TRP",
-    "TYR",
-    "VAL",
-    "UNK",
-]
-tokens = [
-    "<pad>",
-    "-",
-    *canonical_tokens,
-    "A",
-    "G",
-    "C",
-    "U",
-    "N",
-    "DA",
-    "DG",
-    "DC",
-    "DT",
-    "DN",
-]
-token_ids = _index(tokens)
-num_tokens = len(tokens)
-
-prot_letter_to_token = dict(
-    zip(
-        "ARNDCEQGHILKMFPSTWYV",
-        [
-            "ALA",
-            "ARG",
-            "ASN",
-            "ASP",
-            "CYS",
-            "GLU",
-            "GLN",
-            "GLY",
-            "HIS",
-            "ILE",
-            "LEU",
-            "LYS",
-            "MET",
-            "PHE",
-            "PRO",
-            "SER",
-            "THR",
-            "TRP",
-            "TYR",
-            "VAL",
-        ],
-        strict=True,
-    )
-)
-prot_letter_to_token.update({letter: "UNK" for letter in "XJBZOU"} | {"-": "-"})
-
-
-def _parse_atom_rows(rows: str) -> dict[str, list[str]]:
-    table: dict[str, list[str]] = {}
-    for row in rows.strip().splitlines():
-        residue, *atom_names = row.split()
-        table[residue] = atom_names
-    return table
-
-
-ref_atoms = {"PAD": [], "-": []}
-ref_atoms.update(
-    _parse_atom_rows(
-        """
-UNK N CA C O CB
-ALA N CA C O CB
-ARG N CA C O CB CG CD NE CZ NH1 NH2
-ASN N CA C O CB CG OD1 ND2
-ASP N CA C O CB CG OD1 OD2
-CYS N CA C O CB SG
-GLN N CA C O CB CG CD OE1 NE2
-GLU N CA C O CB CG CD OE1 OE2
-GLY N CA C O
-HIS N CA C O CB CG ND1 CD2 CE1 NE2
-ILE N CA C O CB CG1 CG2 CD1
-LEU N CA C O CB CG CD1 CD2
-LYS N CA C O CB CG CD CE NZ
-MET N CA C O CB CG SD CE
-PHE N CA C O CB CG CD1 CD2 CE1 CE2 CZ
-PRO N CA C O CB CG CD
-SER N CA C O CB OG
-THR N CA C O CB OG1 CG2
-TRP N CA C O CB CG CD1 CD2 NE1 CE2 CE3 CZ2 CZ3 CH2
-TYR N CA C O CB CG CD1 CD2 CE1 CE2 CZ OH
-VAL N CA C O CB CG1 CG2
-"""
-    )
-)
-
-protein_backbone_atom_names = ["N", "CA", "C", "O"]
-nucleic_backbone_atom_names = [
-    "P",
-    "OP1",
-    "OP2",
-    "O5'",
-    "C5'",
-    "C4'",
-    "O4'",
-    "C3'",
-    "O3'",
-    "C2'",
-    "O2'",
-    "C1'",
-]
-protein_backbone_atom_index = _index(protein_backbone_atom_names)
-nucleic_backbone_atom_index = _index(nucleic_backbone_atom_names)
-
-_rna_backbone = nucleic_backbone_atom_names
-_dna_backbone = [atom for atom in _rna_backbone if atom != "O2'"]
-ref_atoms.update(
-    {
-        "A": [*_rna_backbone, "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"],
-        "G": [*_rna_backbone, "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"],
-        "C": [*_rna_backbone, "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"],
-        "U": [*_rna_backbone, "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C6"],
-        "N": list(_rna_backbone),
-        "DA": [*_dna_backbone, "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"],
-        "DG": [*_dna_backbone, "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"],
-        "DC": [*_dna_backbone, "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"],
-        "DT": [*_dna_backbone, "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C7", "C6"],
-        "DN": list(_dna_backbone),
-    }
-)
-
-_protein_tokens = ["UNK", *canonical_tokens[:-1]]
-_nucleic_tokens = ["A", "G", "C", "U", "N", "DA", "DG", "DC", "DT", "DN"]
-res_to_center_atom = {
-    **dict.fromkeys(_protein_tokens, "CA"),
-    **dict.fromkeys(_nucleic_tokens, "C1'"),
-}
-res_to_disto_atom = {
-    **dict.fromkeys(_protein_tokens, "CB"),
-    "GLY": "CA",
-    "A": "C4",
-    "G": "C4",
-    "C": "C2",
-    "U": "C2",
-    "N": "C1'",
-    "DA": "C4",
-    "DG": "C4",
-    "DC": "C2",
-    "DT": "C2",
-    "DN": "C1'",
-}
-
-
-num_elements = 128
-bond_types = ["OTHER", "SINGLE", "DOUBLE", "TRIPLE", "AROMATIC", "COVALENT"]
-contact_conditioning_info = {
-    "UNSPECIFIED": 0,
-    "UNSELECTED": 1,
-    "POCKET>BINDER": 2,
-    "BINDER>POCKET": 3,
-    "CONTACT": 4,
-}
-chunk_size_threshold = 384
-
-_method_groups = {
-    0: ("MD",),
-    1: ("X-RAY DIFFRACTION",),
-    2: ("ELECTRON MICROSCOPY",),
-    3: ("SOLUTION NMR",),
-    4: (
-        "SOLID-STATE NMR",
-        "NEUTRON DIFFRACTION",
-        "ELECTRON CRYSTALLOGRAPHY",
-        "FIBER DIFFRACTION",
-        "POWDER DIFFRACTION",
-        "INFRARED SPECTROSCOPY",
-        "FLUORESCENCE TRANSFER",
-        "EPR",
-        "THEORETICAL MODEL",
-        "SOLUTION SCATTERING",
-        "OTHER",
-    ),
-    5: ("AFDB",),
-    6: ("BOLTZ-1",),
-    7: ("FUTURE1",),
-    8: ("FUTURE2",),
-    9: ("FUTURE3",),
-    10: ("FUTURE4",),
-    11: ("FUTURE5",),
-}
-method_types_ids = {
-    method.lower(): identifier
-    for identifier, methods in _method_groups.items()
-    for method in methods
-}
-num_method_types = len(_method_groups)
-
-vdw_radii = [
-    float(radius)
-    for radius in [
-        "1.2",
-        "1.4",
-        "2.2",
-        "1.9",
-        "1.8",
-        "1.7",
-        "1.6",
-        "1.55",
-        "1.5",
-        "1.54",
-        "2.4",
-        "2.2",
-        "2.1",
-        "2.1",
-        "1.95",
-        "1.8",
-        "1.8",
-        "1.88",
-        "2.8",
-        "2.4",
-        "2.3",
-        "2.15",
-        "2.05",
-        "2.05",
-        "2.05",
-        "2.05",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.1",
-        "2.1",
-        "2.1",
-        "2.05",
-        "1.9",
-        "1.9",
-        "2.02",
-        "2.9",
-        "2.55",
-        "2.4",
-        "2.3",
-        "2.15",
-        "2.1",
-        "2.05",
-        "2.05",
-        "2.0",
-        "2.05",
-        "2.1",
-        "2.2",
-        "2.2",
-        "2.25",
-        "2.2",
-        "2.1",
-        "2.1",
-        "2.16",
-        "3.0",
-        "2.7",
-        "2.5",
-        "2.48",
-        "2.47",
-        "2.45",
-        "2.43",
-        "2.42",
-        "2.4",
-        "2.38",
-        "2.37",
-        "2.35",
-        "2.33",
-        "2.32",
-        "2.3",
-        "2.28",
-        "2.27",
-        "2.25",
-        "2.2",
-        "2.1",
-        "2.05",
-        "2.0",
-        "2.0",
-        "2.05",
-        "2.1",
-        "2.05",
-        "2.2",
-        "2.3",
-        "2.3",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.4",
-        "2.0",
-        "2.3",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-        "2.0",
-    ]
-]
+"""Minimal biological constants required by the local Boltz2 runtime.
+
+The upstream project contains additional training-time curation tables.  They
+are intentionally excluded here because FastPLMs neither trains Boltz2 nor
+uses upstream data-processing code.  Keeping only runtime inputs makes this
+module auditable and prevents parity-oracle data from becoming a package
+dependency.
+"""
+
+from __future__ import annotations
+
+
+def _index(values: list[str]) -> dict[str, int]:
+    return {value: index for index, value in enumerate(values)}
+
+
+chain_types = ["PROTEIN", "DNA", "RNA", "NONPOLYMER"]
+chain_type_ids = _index(chain_types)
+
+
+canonical_tokens = [
+    "ALA",
+    "ARG",
+    "ASN",
+    "ASP",
+    "CYS",
+    "GLN",
+    "GLU",
+    "GLY",
+    "HIS",
+    "ILE",
+    "LEU",
+    "LYS",
+    "MET",
+    "PHE",
+    "PRO",
+    "SER",
+    "THR",
+    "TRP",
+    "TYR",
+    "VAL",
+    "UNK",
+]
+tokens = [
+    "<pad>",
+    "-",
+    *canonical_tokens,
+    "A",
+    "G",
+    "C",
+    "U",
+    "N",
+    "DA",
+    "DG",
+    "DC",
+    "DT",
+    "DN",
+]
+token_ids = _index(tokens)
+num_tokens = len(tokens)
+
+prot_letter_to_token = dict(
+    zip(
+        "ARNDCEQGHILKMFPSTWYV",
+        [
+            "ALA",
+            "ARG",
+            "ASN",
+            "ASP",
+            "CYS",
+            "GLU",
+            "GLN",
+            "GLY",
+            "HIS",
+            "ILE",
+            "LEU",
+            "LYS",
+            "MET",
+            "PHE",
+            "PRO",
+            "SER",
+            "THR",
+            "TRP",
+            "TYR",
+            "VAL",
+        ],
+        strict=True,
+    )
+)
+prot_letter_to_token.update({letter: "UNK" for letter in "XJBZOU"} | {"-": "-"})
+
+
+def _parse_atom_rows(rows: str) -> dict[str, list[str]]:
+    table: dict[str, list[str]] = {}
+    for row in rows.strip().splitlines():
+        residue, *atom_names = row.split()
+        table[residue] = atom_names
+    return table
+
+
+ref_atoms = {"PAD": [], "-": []}
+ref_atoms.update(
+    _parse_atom_rows(
+        """
+UNK N CA C O CB
+ALA N CA C O CB
+ARG N CA C O CB CG CD NE CZ NH1 NH2
+ASN N CA C O CB CG OD1 ND2
+ASP N CA C O CB CG OD1 OD2
+CYS N CA C O CB SG
+GLN N CA C O CB CG CD OE1 NE2
+GLU N CA C O CB CG CD OE1 OE2
+GLY N CA C O
+HIS N CA C O CB CG ND1 CD2 CE1 NE2
+ILE N CA C O CB CG1 CG2 CD1
+LEU N CA C O CB CG CD1 CD2
+LYS N CA C O CB CG CD CE NZ
+MET N CA C O CB CG SD CE
+PHE N CA C O CB CG CD1 CD2 CE1 CE2 CZ
+PRO N CA C O CB CG CD
+SER N CA C O CB OG
+THR N CA C O CB OG1 CG2
+TRP N CA C O CB CG CD1 CD2 NE1 CE2 CE3 CZ2 CZ3 CH2
+TYR N CA C O CB CG CD1 CD2 CE1 CE2 CZ OH
+VAL N CA C O CB CG1 CG2
+"""
+    )
+)
+
+protein_backbone_atom_names = ["N", "CA", "C", "O"]
+nucleic_backbone_atom_names = [
+    "P",
+    "OP1",
+    "OP2",
+    "O5'",
+    "C5'",
+    "C4'",
+    "O4'",
+    "C3'",
+    "O3'",
+    "C2'",
+    "O2'",
+    "C1'",
+]
+protein_backbone_atom_index = _index(protein_backbone_atom_names)
+nucleic_backbone_atom_index = _index(nucleic_backbone_atom_names)
+
+_rna_backbone = nucleic_backbone_atom_names
+_dna_backbone = [atom for atom in _rna_backbone if atom != "O2'"]
+ref_atoms.update(
+    {
+        "A": [*_rna_backbone, "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"],
+        "G": [*_rna_backbone, "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"],
+        "C": [*_rna_backbone, "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"],
+        "U": [*_rna_backbone, "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C6"],
+        "N": list(_rna_backbone),
+        "DA": [*_dna_backbone, "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"],
+        "DG": [*_dna_backbone, "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"],
+        "DC": [*_dna_backbone, "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"],
+        "DT": [*_dna_backbone, "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C7", "C6"],
+        "DN": list(_dna_backbone),
+    }
+)
+
+_protein_tokens = ["UNK", *canonical_tokens[:-1]]
+_nucleic_tokens = ["A", "G", "C", "U", "N", "DA", "DG", "DC", "DT", "DN"]
+res_to_center_atom = {
+    **dict.fromkeys(_protein_tokens, "CA"),
+    **dict.fromkeys(_nucleic_tokens, "C1'"),
+}
+res_to_disto_atom = {
+    **dict.fromkeys(_protein_tokens, "CB"),
+    "GLY": "CA",
+    "A": "C4",
+    "G": "C4",
+    "C": "C2",
+    "U": "C2",
+    "N": "C1'",
+    "DA": "C4",
+    "DG": "C4",
+    "DC": "C2",
+    "DT": "C2",
+    "DN": "C1'",
+}
+
+
+num_elements = 128
+bond_types = ["OTHER", "SINGLE", "DOUBLE", "TRIPLE", "AROMATIC", "COVALENT"]
+contact_conditioning_info = {
+    "UNSPECIFIED": 0,
+    "UNSELECTED": 1,
+    "POCKET>BINDER": 2,
+    "BINDER>POCKET": 3,
+    "CONTACT": 4,
+}
+chunk_size_threshold = 384
+
+_method_groups = {
+    0: ("MD",),
+    1: ("X-RAY DIFFRACTION",),
+    2: ("ELECTRON MICROSCOPY",),
+    3: ("SOLUTION NMR",),
+    4: (
+        "SOLID-STATE NMR",
+        "NEUTRON DIFFRACTION",
+        "ELECTRON CRYSTALLOGRAPHY",
+        "FIBER DIFFRACTION",
+        "POWDER DIFFRACTION",
+        "INFRARED SPECTROSCOPY",
+        "FLUORESCENCE TRANSFER",
+        "EPR",
+        "THEORETICAL MODEL",
+        "SOLUTION SCATTERING",
+        "OTHER",
+    ),
+    5: ("AFDB",),
+    6: ("BOLTZ-1",),
+    7: ("FUTURE1",),
+    8: ("FUTURE2",),
+    9: ("FUTURE3",),
+    10: ("FUTURE4",),
+    11: ("FUTURE5",),
+}
+method_types_ids = {
+    method.lower(): identifier
+    for identifier, methods in _method_groups.items()
+    for method in methods
+}
+num_method_types = len(_method_groups)
+
+vdw_radii = [
+    float(radius)
+    for radius in [
+        "1.2",
+        "1.4",
+        "2.2",
+        "1.9",
+        "1.8",
+        "1.7",
+        "1.6",
+        "1.55",
+        "1.5",
+        "1.54",
+        "2.4",
+        "2.2",
+        "2.1",
+        "2.1",
+        "1.95",
+        "1.8",
+        "1.8",
+        "1.88",
+        "2.8",
+        "2.4",
+        "2.3",
+        "2.15",
+        "2.05",
+        "2.05",
+        "2.05",
+        "2.05",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.1",
+        "2.1",
+        "2.1",
+        "2.05",
+        "1.9",
+        "1.9",
+        "2.02",
+        "2.9",
+        "2.55",
+        "2.4",
+        "2.3",
+        "2.15",
+        "2.1",
+        "2.05",
+        "2.05",
+        "2.0",
+        "2.05",
+        "2.1",
+        "2.2",
+        "2.2",
+        "2.25",
+        "2.2",
+        "2.1",
+        "2.1",
+        "2.16",
+        "3.0",
+        "2.7",
+        "2.5",
+        "2.48",
+        "2.47",
+        "2.45",
+        "2.43",
+        "2.42",
+        "2.4",
+        "2.38",
+        "2.37",
+        "2.35",
+        "2.33",
+        "2.32",
+        "2.3",
+        "2.28",
+        "2.27",
+        "2.25",
+        "2.2",
+        "2.1",
+        "2.05",
+        "2.0",
+        "2.0",
+        "2.05",
+        "2.1",
+        "2.05",
+        "2.2",
+        "2.3",
+        "2.3",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.4",
+        "2.0",
+        "2.3",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+        "2.0",
+    ]
+]

fastplms/models/boltz/vb_layers_attention.py

@@ -1,103 +1,103 @@
-"""Pair-biased self-attention used in the Boltz2 diffusion stack."""
-
-from __future__ import annotations
-
-from collections.abc import Callable, MutableMapping
-from einops.layers.torch import Rearrange
-from torch import Tensor, nn
-
-from . import vb_layers_initialize as init
-from ._pair_attention import pair_biased_attention, reshape_heads
-
-
-class AttentionPairBias(nn.Module):
-    """Attend over sequence states while adding a learned pair representation."""
-
-    def __init__(
-        self,
-        c_s: int,
-        c_z: int,
-        num_heads: int,
-        inf: float = 1e6,
-        initial_norm: bool = True,
-    ) -> None:
-        super().__init__()
-        if c_s % num_heads:
-            raise ValueError(f"c_s={c_s} must be divisible by num_heads={num_heads}")
-
-        self.c_s = c_s
-        self.num_heads = num_heads
-        self.head_dim = c_s // num_heads
-        self.inf = inf
-        self.initial_norm = initial_norm
-        if initial_norm:
-            self.norm_s = nn.LayerNorm(c_s)
-
-        self.proj_q = nn.Linear(c_s, c_s)
-        self.proj_k = nn.Linear(c_s, c_s, bias=False)
-        self.proj_v = nn.Linear(c_s, c_s, bias=False)
-        self.proj_g = nn.Linear(c_s, c_s, bias=False)
-        self.proj_z = nn.Sequential(
-            nn.LayerNorm(c_z),
-            nn.Linear(c_z, num_heads, bias=False),
-            Rearrange("b ... h -> b h ..."),
-        )
-        self.proj_o = nn.Linear(c_s, c_s, bias=False)
-        init.final_init_(self.proj_o.weight)
-
-    def _resolve_pair_bias(
-        self,
-        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"]  # (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  # (b, h, l_q, l_k)
-
-    def forward(
-        self,
-        s: Tensor,
-        z: Tensor,
-        mask: Tensor,
-        k_in: Tensor | None = None,
-        multiplicity: int = 1,
-        to_keys: Callable[[Tensor], Tensor] | None = None,
-        model_cache: MutableMapping[str, Tensor] | None = None,
-    ) -> Tensor:
-        """Transform S with shapes ``S: (b, l_q, d)`` and ``Z: (b, l, l, d_z)``."""
-
-        # 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)  # (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  # (b, l_k, d)
-            key_mask = mask  # (b_k, l_k)
-
-        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,
-            value,
-            pair_bias,
-            key_mask,
-            self.inf,
-        )  # (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)
+"""Pair-biased self-attention used in the Boltz2 diffusion stack."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, MutableMapping
+from einops.layers.torch import Rearrange
+from torch import Tensor, nn
+
+from . import vb_layers_initialize as init
+from ._pair_attention import pair_biased_attention, reshape_heads
+
+
+class AttentionPairBias(nn.Module):
+    """Attend over sequence states while adding a learned pair representation."""
+
+    def __init__(
+        self,
+        c_s: int,
+        c_z: int,
+        num_heads: int,
+        inf: float = 1e6,
+        initial_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        if c_s % num_heads:
+            raise ValueError(f"c_s={c_s} must be divisible by num_heads={num_heads}")
+
+        self.c_s = c_s
+        self.num_heads = num_heads
+        self.head_dim = c_s // num_heads
+        self.inf = inf
+        self.initial_norm = initial_norm
+        if initial_norm:
+            self.norm_s = nn.LayerNorm(c_s)
+
+        self.proj_q = nn.Linear(c_s, c_s)
+        self.proj_k = nn.Linear(c_s, c_s, bias=False)
+        self.proj_v = nn.Linear(c_s, c_s, bias=False)
+        self.proj_g = nn.Linear(c_s, c_s, bias=False)
+        self.proj_z = nn.Sequential(
+            nn.LayerNorm(c_z),
+            nn.Linear(c_z, num_heads, bias=False),
+            Rearrange("b ... h -> b h ..."),
+        )
+        self.proj_o = nn.Linear(c_s, c_s, bias=False)
+        init.final_init_(self.proj_o.weight)
+
+    def _resolve_pair_bias(
+        self,
+        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"]  # (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  # (b, h, l_q, l_k)
+
+    def forward(
+        self,
+        s: Tensor,
+        z: Tensor,
+        mask: Tensor,
+        k_in: Tensor | None = None,
+        multiplicity: int = 1,
+        to_keys: Callable[[Tensor], Tensor] | None = None,
+        model_cache: MutableMapping[str, Tensor] | None = None,
+    ) -> Tensor:
+        """Transform S with shapes ``S: (b, l_q, d)`` and ``Z: (b, l, l, d_z)``."""
+
+        # 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)  # (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  # (b, l_k, d)
+            key_mask = mask  # (b_k, l_k)
+
+        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,
+            value,
+            pair_bias,
+            key_mask,
+            self.inf,
+        )  # (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

@@ -1,76 +1,76 @@
-"""Cross-attention variant of Boltz2 pair-biased attention."""
-
-from __future__ import annotations
-
-from einops.layers.torch import Rearrange
-from torch import Tensor, nn
-
-from . import vb_layers_initialize as init
-from ._pair_attention import pair_biased_attention, reshape_heads
-
-
-class AttentionPairBias(nn.Module):
-    """Attend from S to K while adding either learned or precomputed pair bias."""
-
-    def __init__(
-        self,
-        c_s: int,
-        c_z: int | None = None,
-        num_heads: int | None = None,
-        inf: float = 1e6,
-        compute_pair_bias: bool = True,
-    ) -> None:
-        super().__init__()
-        if num_heads is None or c_s % num_heads:
-            raise ValueError("num_heads must divide c_s")
-        if compute_pair_bias and c_z is None:
-            raise ValueError("c_z is required when pair bias is learned")
-
-        self.c_s = c_s
-        self.num_heads = num_heads
-        self.head_dim = c_s // num_heads
-        self.inf = inf
-        self.proj_q = nn.Linear(c_s, c_s)
-        self.proj_k = nn.Linear(c_s, c_s, bias=False)
-        self.proj_v = nn.Linear(c_s, c_s, bias=False)
-        self.proj_g = nn.Linear(c_s, c_s, bias=False)
-        self.compute_pair_bias = compute_pair_bias
-        if compute_pair_bias:
-            self.proj_z = nn.Sequential(
-                nn.LayerNorm(c_z),
-                nn.Linear(c_z, num_heads, bias=False),
-                Rearrange("b ... h -> b h ..."),
-            )
-        else:
-            self.proj_z = Rearrange("b ... h -> b h ...")
-        self.proj_o = nn.Linear(c_s, c_s, bias=False)
-        init.final_init_(self.proj_o.weight)
-
-    def forward(
-        self,
-        s: Tensor,
-        z: Tensor,
-        mask: Tensor,
-        k_in: Tensor,
-        multiplicity: int = 1,
-    ) -> Tensor:
-        """Transform S with K and pair input Z, returning shape ``(b, l_q, d)``."""
-
-        # 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,
-            value,
-            pair_bias,
-            mask,
-            self.inf,
-        )  # (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)
+"""Cross-attention variant of Boltz2 pair-biased attention."""
+
+from __future__ import annotations
+
+from einops.layers.torch import Rearrange
+from torch import Tensor, nn
+
+from . import vb_layers_initialize as init
+from ._pair_attention import pair_biased_attention, reshape_heads
+
+
+class AttentionPairBias(nn.Module):
+    """Attend from S to K while adding either learned or precomputed pair bias."""
+
+    def __init__(
+        self,
+        c_s: int,
+        c_z: int | None = None,
+        num_heads: int | None = None,
+        inf: float = 1e6,
+        compute_pair_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        if num_heads is None or c_s % num_heads:
+            raise ValueError("num_heads must divide c_s")
+        if compute_pair_bias and c_z is None:
+            raise ValueError("c_z is required when pair bias is learned")
+
+        self.c_s = c_s
+        self.num_heads = num_heads
+        self.head_dim = c_s // num_heads
+        self.inf = inf
+        self.proj_q = nn.Linear(c_s, c_s)
+        self.proj_k = nn.Linear(c_s, c_s, bias=False)
+        self.proj_v = nn.Linear(c_s, c_s, bias=False)
+        self.proj_g = nn.Linear(c_s, c_s, bias=False)
+        self.compute_pair_bias = compute_pair_bias
+        if compute_pair_bias:
+            self.proj_z = nn.Sequential(
+                nn.LayerNorm(c_z),
+                nn.Linear(c_z, num_heads, bias=False),
+                Rearrange("b ... h -> b h ..."),
+            )
+        else:
+            self.proj_z = Rearrange("b ... h -> b h ...")
+        self.proj_o = nn.Linear(c_s, c_s, bias=False)
+        init.final_init_(self.proj_o.weight)
+
+    def forward(
+        self,
+        s: Tensor,
+        z: Tensor,
+        mask: Tensor,
+        k_in: Tensor,
+        multiplicity: int = 1,
+    ) -> Tensor:
+        """Transform S with K and pair input Z, returning shape ``(b, l_q, d)``."""
+
+        # 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,
+            value,
+            pair_bias,
+            mask,
+            self.inf,
+        )  # (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

@@ -1,273 +1,273 @@
-"""Confidence geometry and aggregate metrics for Boltz2 outputs."""
-
-from __future__ import annotations
-
-import torch
-
-from torch import nn
-
-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."""
-
-    # 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:
-    with torch.amp.autocast("cuda", enabled=False):
-        return torch.bmm(
-            feats["atom_to_token"].float(),
-            feats["asym_id"].unsqueeze(-1).float(),
-        ).squeeze(-1)  # (b, a)
-
-
-def _replace_nonpolymer_frames(
-    coordinates: torch.Tensor,
-    frame_indices: torch.Tensor,
-    feats: dict[str, torch.Tensor],
-    atom_chain_ids: torch.Tensor,
-    resolved_mask: torch.Tensor | None,
-    *,
-    inference: bool,
-) -> None:
-    # 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]  # (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 = (
-                feats["mol_type"][batch_index, token_offset] == const.chain_type_ids["NONPOLYMER"]
-            )
-            if is_nonpolymer and atom_count >= 3:
-                chain_resolved = (
-                    feats["atom_pad_mask"][batch_index]
-                    if inference
-                    else (
-                        feats["atom_resolved_mask"][batch_index]
-                        if resolved_mask is None
-                        else resolved_mask[batch_index]
-                    )
-                )
-                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,
-                    :,
-                    token_offset : token_offset + atom_count,
-                    :,
-                ] = frames
-            token_offset += token_count
-            atom_offset += atom_count
-
-
-def compute_frame_pred(
-    pred_atom_coords: torch.Tensor,
-    frames_idx_true: torch.Tensor,
-    feats: dict[str, torch.Tensor],
-    multiplicity: int,
-    resolved_mask: torch.Tensor | None = None,
-    inference: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Construct predicted local frames and their validity mask.
-
-    ``pred_atom_coords`` has shape ``(b * multiplicity, a, 3)``.  Polymer
-    frames retain the supplied atom indices; nonpolymer frames use each
-    atom's three nearest resolved neighbors.
-    """
-
-    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
-    )  # (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,
-        feats,
-        atom_chain_ids,
-        resolved_mask,
-        inference=inference,
-    )
-
-    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)  # (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)  # (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:
-    """Convert categorical confidence logits to bin-center expectations."""
-
-    num_bins = logits.shape[-1]
-    bin_width = end / num_bins
-    centers = torch.arange(
-        0.5 * bin_width,
-        end,
-        bin_width,
-        device=logits.device,
-    )  # (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
-    )  # (...)
-
-
-def tm_function(d: torch.Tensor, n_res: torch.Tensor) -> torch.Tensor:
-    """Evaluate the TM-score distance kernel."""
-
-    # 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:
-    # 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(
-    logits: torch.Tensor,
-    x_preds: torch.Tensor,
-    feats: dict[str, torch.Tensor],
-    multiplicity: int,
-) -> tuple[
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    dict[int, dict[int, torch.Tensor]],
-]:
-    """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,
-    )  # 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(
-        0.5 * (32.0 / num_bins),
-        32.0,
-        32.0 / num_bins,
-        device=logits.device,
-    ).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)  # (b * m,)
-    iptm = _maximum_masked_tm(expected_tm, interface_mask)  # (b * m,)
-
-    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, :]
-    )  # (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():
-        scores: dict[int, torch.Tensor] = {}
-        for second_chain in torch.unique(asym_id).tolist():
-            chain_mask = base_pair_mask
-            chain_mask = chain_mask * (asym_id[:, None, :] == first_chain)
-            chain_mask = chain_mask * (asym_id[:, :, None] == second_chain)
-            scores[second_chain] = _maximum_masked_tm(expected_tm, chain_mask)
-        chain_pair_iptm[first_chain] = scores
-
-    return ptm, iptm, ligand_iptm, protein_iptm, chain_pair_iptm
+"""Confidence geometry and aggregate metrics for Boltz2 outputs."""
+
+from __future__ import annotations
+
+import torch
+
+from torch import nn
+
+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."""
+
+    # 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:
+    with torch.amp.autocast("cuda", enabled=False):
+        return torch.bmm(
+            feats["atom_to_token"].float(),
+            feats["asym_id"].unsqueeze(-1).float(),
+        ).squeeze(-1)  # (b, a)
+
+
+def _replace_nonpolymer_frames(
+    coordinates: torch.Tensor,
+    frame_indices: torch.Tensor,
+    feats: dict[str, torch.Tensor],
+    atom_chain_ids: torch.Tensor,
+    resolved_mask: torch.Tensor | None,
+    *,
+    inference: bool,
+) -> None:
+    # 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]  # (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 = (
+                feats["mol_type"][batch_index, token_offset] == const.chain_type_ids["NONPOLYMER"]
+            )
+            if is_nonpolymer and atom_count >= 3:
+                chain_resolved = (
+                    feats["atom_pad_mask"][batch_index]
+                    if inference
+                    else (
+                        feats["atom_resolved_mask"][batch_index]
+                        if resolved_mask is None
+                        else resolved_mask[batch_index]
+                    )
+                )
+                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,
+                    :,
+                    token_offset : token_offset + atom_count,
+                    :,
+                ] = frames
+            token_offset += token_count
+            atom_offset += atom_count
+
+
+def compute_frame_pred(
+    pred_atom_coords: torch.Tensor,
+    frames_idx_true: torch.Tensor,
+    feats: dict[str, torch.Tensor],
+    multiplicity: int,
+    resolved_mask: torch.Tensor | None = None,
+    inference: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Construct predicted local frames and their validity mask.
+
+    ``pred_atom_coords`` has shape ``(b * multiplicity, a, 3)``.  Polymer
+    frames retain the supplied atom indices; nonpolymer frames use each
+    atom's three nearest resolved neighbors.
+    """
+
+    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
+    )  # (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,
+        feats,
+        atom_chain_ids,
+        resolved_mask,
+        inference=inference,
+    )
+
+    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)  # (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)  # (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:
+    """Convert categorical confidence logits to bin-center expectations."""
+
+    num_bins = logits.shape[-1]
+    bin_width = end / num_bins
+    centers = torch.arange(
+        0.5 * bin_width,
+        end,
+        bin_width,
+        device=logits.device,
+    )  # (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
+    )  # (...)
+
+
+def tm_function(d: torch.Tensor, n_res: torch.Tensor) -> torch.Tensor:
+    """Evaluate the TM-score distance kernel."""
+
+    # 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:
+    # 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(
+    logits: torch.Tensor,
+    x_preds: torch.Tensor,
+    feats: dict[str, torch.Tensor],
+    multiplicity: int,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    dict[int, dict[int, torch.Tensor]],
+]:
+    """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,
+    )  # 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(
+        0.5 * (32.0 / num_bins),
+        32.0,
+        32.0 / num_bins,
+        device=logits.device,
+    ).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)  # (b * m,)
+    iptm = _maximum_masked_tm(expected_tm, interface_mask)  # (b * m,)
+
+    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, :]
+    )  # (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():
+        scores: dict[int, torch.Tensor] = {}
+        for second_chain in torch.unique(asym_id).tolist():
+            chain_mask = base_pair_mask
+            chain_mask = chain_mask * (asym_id[:, None, :] == first_chain)
+            chain_mask = chain_mask * (asym_id[:, :, None] == second_chain)
+            scores[second_chain] = _maximum_masked_tm(expected_tm, chain_mask)
+        chain_pair_iptm[first_chain] = scores
+
+    return ptm, iptm, ligand_iptm, protein_iptm, chain_pair_iptm

fastplms/models/boltz/vb_layers_dropout.py

@@ -1,43 +1,43 @@
-"""Broadcast dropout masks used by the Boltz pair stack.
-
-The mask is sampled in ``float32`` even when the pair tensor uses a lower
-precision dtype.  This matches the checkpoint implementation while keeping
-the broadcast axis explicit.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from torch import Tensor
-
-
-def _broadcast_mask_shape(Z: Tensor, *, columnwise: bool) -> tuple[int, ...]:
-    """Return the row-wise or column-wise pair-mask shape for ``Z``."""
-
-    if Z.ndim != 4:
-        raise ValueError(f"pair tensor must have four dimensions, got {Z.ndim}")
-    b, rows, columns, _ = Z.shape
-    return (b, 1, columns, 1) if columnwise else (b, rows, 1, 1)
-
-
-def get_dropout_mask(
-    dropout: float,
-    z: Tensor,
-    training: bool,
-    columnwise: bool = False,
-) -> Tensor:
-    """Sample an inverted-dropout mask that broadcasts over pair channels.
-
-    ``Z`` is the pair tensor with shape ``(b, n, n, d)``.  Row-wise masks
-    vary along the first residue axis; column-wise masks vary along the
-    second.  Evaluation returns an all-one mask while retaining the same RNG
-    call pattern as training.
-    """
-
-    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
-    )  # (b, 1, n, 1) or (b, n, 1, 1)
-    return keep * (1.0 / (1.0 - probability))  # same broadcast-mask shape
+"""Broadcast dropout masks used by the Boltz pair stack.
+
+The mask is sampled in ``float32`` even when the pair tensor uses a lower
+precision dtype.  This matches the checkpoint implementation while keeping
+the broadcast axis explicit.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from torch import Tensor
+
+
+def _broadcast_mask_shape(Z: Tensor, *, columnwise: bool) -> tuple[int, ...]:
+    """Return the row-wise or column-wise pair-mask shape for ``Z``."""
+
+    if Z.ndim != 4:
+        raise ValueError(f"pair tensor must have four dimensions, got {Z.ndim}")
+    b, rows, columns, _ = Z.shape
+    return (b, 1, columns, 1) if columnwise else (b, rows, 1, 1)
+
+
+def get_dropout_mask(
+    dropout: float,
+    z: Tensor,
+    training: bool,
+    columnwise: bool = False,
+) -> Tensor:
+    """Sample an inverted-dropout mask that broadcasts over pair channels.
+
+    ``Z`` is the pair tensor with shape ``(b, n, n, d)``.  Row-wise masks
+    vary along the first residue axis; column-wise masks vary along the
+    second.  Evaluation returns an all-one mask while retaining the same RNG
+    call pattern as training.
+    """
+
+    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
+    )  # (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

@@ -1,110 +1,110 @@
-"""Torch-only parameter initialization for the Boltz2 runtime."""
-
-from __future__ import annotations
-
-import math
-import torch
-
-from typing import Literal
-from torch import Tensor
-
-
-FanMode = Literal["fan_in", "fan_out", "fan_avg"]
-
-
-def _calculate_fan(shape: torch.Size | tuple[int, ...], fan: FanMode = "fan_in") -> float:
-    """Resolve the selected fan for a two-dimensional linear weight tensor."""
-
-    if len(shape) != 2:
-        raise ValueError(f"linear weights must be two-dimensional, received {tuple(shape)}")
-    fan_out, fan_in = shape
-    if fan == "fan_in":
-        return float(fan_in)
-    if fan == "fan_out":
-        return float(fan_out)
-    if fan == "fan_avg":
-        return (fan_in + fan_out) / 2
-    raise ValueError(f"invalid fan mode: {fan!r}")
-
-
-def trunc_normal_init_(
-    weights: Tensor,
-    scale: float = 1.0,
-    fan: FanMode = "fan_in",
-) -> 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():
-        torch.nn.init.trunc_normal_(weights, mean=0.0, std=std, a=-2 * std, b=2 * std)
-
-
-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)
-
-
-def he_normal_init_(weights: Tensor) -> None:
-    """Initialize W using twice the fan-in variance."""
-
-    # weights: (d_out, d_in), mutated in place.
-    trunc_normal_init_(weights, scale=2.0)
-
-
-def glorot_uniform_init_(weights: Tensor) -> None:
-    """Initialize W with Xavier uniform values."""
-
-    # weights: (...), rank at least two; mutated in place.
-    torch.nn.init.xavier_uniform_(weights, gain=1.0)
-
-
-def _fill_(tensor: Tensor, value: float) -> None:
-    # tensor: (...), mutated in place without changing shape.
-    with torch.no_grad():
-        tensor.fill_(value)
-
-
-def final_init_(weights: Tensor) -> None:
-    """Zero the final projection W."""
-
-    # weights: (...), mutated in place.
-    _fill_(weights, 0.0)
-
-
-def gating_init_(weights: Tensor) -> None:
-    """Zero the gating projection W."""
-
-    # weights: (...), mutated in place.
-    _fill_(weights, 0.0)
-
-
-def bias_init_zero_(bias: Tensor) -> None:
-    """Set the bias tensor to zero."""
-
-    # bias: (...), mutated in place.
-    _fill_(bias, 0.0)
-
-
-def bias_init_one_(bias: Tensor) -> None:
-    """Set the bias tensor to one."""
-
-    # bias: (...), mutated in place.
-    _fill_(bias, 1.0)
-
-
-def normal_init_(weights: Tensor) -> None:
-    """Initialize W with linear Kaiming-normal values."""
-
-    # weights: (...), rank at least two; mutated in place.
-    torch.nn.init.kaiming_normal_(weights, nonlinearity="linear")
-
-
-def ipa_point_weights_init_(weights: Tensor) -> None:
-    """Initialize W so applying softplus yields one."""
-
-    # weights: (...), mutated in place.
-    _fill_(weights, 0.541324854612918)
+"""Torch-only parameter initialization for the Boltz2 runtime."""
+
+from __future__ import annotations
+
+import math
+import torch
+
+from typing import Literal
+from torch import Tensor
+
+
+FanMode = Literal["fan_in", "fan_out", "fan_avg"]
+
+
+def _calculate_fan(shape: torch.Size | tuple[int, ...], fan: FanMode = "fan_in") -> float:
+    """Resolve the selected fan for a two-dimensional linear weight tensor."""
+
+    if len(shape) != 2:
+        raise ValueError(f"linear weights must be two-dimensional, received {tuple(shape)}")
+    fan_out, fan_in = shape
+    if fan == "fan_in":
+        return float(fan_in)
+    if fan == "fan_out":
+        return float(fan_out)
+    if fan == "fan_avg":
+        return (fan_in + fan_out) / 2
+    raise ValueError(f"invalid fan mode: {fan!r}")
+
+
+def trunc_normal_init_(
+    weights: Tensor,
+    scale: float = 1.0,
+    fan: FanMode = "fan_in",
+) -> 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():
+        torch.nn.init.trunc_normal_(weights, mean=0.0, std=std, a=-2 * std, b=2 * std)
+
+
+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)
+
+
+def he_normal_init_(weights: Tensor) -> None:
+    """Initialize W using twice the fan-in variance."""
+
+    # weights: (d_out, d_in), mutated in place.
+    trunc_normal_init_(weights, scale=2.0)
+
+
+def glorot_uniform_init_(weights: Tensor) -> None:
+    """Initialize W with Xavier uniform values."""
+
+    # weights: (...), rank at least two; mutated in place.
+    torch.nn.init.xavier_uniform_(weights, gain=1.0)
+
+
+def _fill_(tensor: Tensor, value: float) -> None:
+    # tensor: (...), mutated in place without changing shape.
+    with torch.no_grad():
+        tensor.fill_(value)
+
+
+def final_init_(weights: Tensor) -> None:
+    """Zero the final projection W."""
+
+    # weights: (...), mutated in place.
+    _fill_(weights, 0.0)
+
+
+def gating_init_(weights: Tensor) -> None:
+    """Zero the gating projection W."""
+
+    # weights: (...), mutated in place.
+    _fill_(weights, 0.0)
+
+
+def bias_init_zero_(bias: Tensor) -> None:
+    """Set the bias tensor to zero."""
+
+    # bias: (...), mutated in place.
+    _fill_(bias, 0.0)
+
+
+def bias_init_one_(bias: Tensor) -> None:
+    """Set the bias tensor to one."""
+
+    # bias: (...), mutated in place.
+    _fill_(bias, 1.0)
+
+
+def normal_init_(weights: Tensor) -> None:
+    """Initialize W with linear Kaiming-normal values."""
+
+    # weights: (...), rank at least two; mutated in place.
+    torch.nn.init.kaiming_normal_(weights, nonlinearity="linear")
+
+
+def ipa_point_weights_init_(weights: Tensor) -> None:
+    """Initialize W so applying softplus yields one."""
+
+    # weights: (...), mutated in place.
+    _fill_(weights, 0.541324854612918)

fastplms/models/boltz/vb_layers_outer_product_mean.py

@@ -1,105 +1,105 @@
-"""Outer-product aggregation from MSA states to pair states."""
-
-from __future__ import annotations
-
-import torch
-
-from torch import Tensor, nn
-
-from . import vb_layers_initialize as init
-
-
-class OuterProductMean(nn.Module):
-    """Aggregate M with shape ``(b, s, l, d_m)`` into pair tensor Z."""
-
-    def __init__(self, c_in: int, c_hidden: int, c_out: int) -> None:
-        super().__init__()
-        self.c_hidden = c_hidden
-        self.norm = nn.LayerNorm(c_in)
-        self.proj_a = nn.Linear(c_in, c_hidden, bias=False)
-        self.proj_b = nn.Linear(c_in, c_hidden, bias=False)
-        self.proj_o = nn.Linear(c_hidden * c_hidden, c_out)
-        init.final_init_(self.proj_o.weight)
-        init.final_init_(self.proj_o.bias)
-
-    @staticmethod
-    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]  # (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)  # (b, l, l, 1)
-
-    @staticmethod
-    def _outer_product(left: Tensor, right: Tensor, counts: Tensor) -> Tensor:
-        # left: (b, s, l, d_left); right: (b, s, l, d_right).
-        # counts: (b, l, l, 1); chunked callers may use d_left < d_right.
-        outer = torch.einsum("bsic,bsjd->bijcd", left, right)  # (b, l, l, d_left, d_right)
-        return outer.reshape(*outer.shape[:3], -1) / counts  # (b, l, l, d_left * d_right)
-
-    def _chunked_projection(
-        self,
-        left: Tensor,
-        right: Tensor,
-        counts: Tensor,
-        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
-            )  # (b, l, l, d_chunk * d_h)
-            weight = self.proj_o.weight[
-                :,
-                start * self.c_hidden : stop * self.c_hidden,
-            ]  # (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  # (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)``."""
-
-        # 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)  # (b, l, l, 1)
-            return self._chunked_projection(
-                left,
-                right,
-                counts,
-                chunk_size,
-                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)
+"""Outer-product aggregation from MSA states to pair states."""
+
+from __future__ import annotations
+
+import torch
+
+from torch import Tensor, nn
+
+from . import vb_layers_initialize as init
+
+
+class OuterProductMean(nn.Module):
+    """Aggregate M with shape ``(b, s, l, d_m)`` into pair tensor Z."""
+
+    def __init__(self, c_in: int, c_hidden: int, c_out: int) -> None:
+        super().__init__()
+        self.c_hidden = c_hidden
+        self.norm = nn.LayerNorm(c_in)
+        self.proj_a = nn.Linear(c_in, c_hidden, bias=False)
+        self.proj_b = nn.Linear(c_in, c_hidden, bias=False)
+        self.proj_o = nn.Linear(c_hidden * c_hidden, c_out)
+        init.final_init_(self.proj_o.weight)
+        init.final_init_(self.proj_o.bias)
+
+    @staticmethod
+    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]  # (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)  # (b, l, l, 1)
+
+    @staticmethod
+    def _outer_product(left: Tensor, right: Tensor, counts: Tensor) -> Tensor:
+        # left: (b, s, l, d_left); right: (b, s, l, d_right).
+        # counts: (b, l, l, 1); chunked callers may use d_left < d_right.
+        outer = torch.einsum("bsic,bsjd->bijcd", left, right)  # (b, l, l, d_left, d_right)
+        return outer.reshape(*outer.shape[:3], -1) / counts  # (b, l, l, d_left * d_right)
+
+    def _chunked_projection(
+        self,
+        left: Tensor,
+        right: Tensor,
+        counts: Tensor,
+        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
+            )  # (b, l, l, d_chunk * d_h)
+            weight = self.proj_o.weight[
+                :,
+                start * self.c_hidden : stop * self.c_hidden,
+            ]  # (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  # (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)``."""
+
+        # 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)  # (b, l, l, 1)
+            return self._chunked_projection(
+                left,
+                right,
+                counts,
+                chunk_size,
+                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

@@ -1,112 +1,112 @@
-"""Pair-weighted aggregation of Boltz2 MSA states."""
-
-from __future__ import annotations
-
-import torch
-
-from torch import Tensor, nn
-
-from . import vb_layers_initialize as init
-
-
-class PairWeightedAveraging(nn.Module):
-    """Update M with pair-derived attention weights while preserving its shape."""
-
-    def __init__(
-        self,
-        c_m: int,
-        c_z: int,
-        c_h: int,
-        num_heads: int,
-        inf: float = 1e6,
-    ) -> None:
-        super().__init__()
-        self.c_m = c_m
-        self.c_z = c_z
-        self.c_h = c_h
-        self.num_heads = num_heads
-        self.inf = inf
-        self.norm_m = nn.LayerNorm(c_m)
-        self.norm_z = nn.LayerNorm(c_z)
-        self.proj_m = nn.Linear(c_m, c_h * num_heads, bias=False)
-        self.proj_g = nn.Linear(c_m, c_h * num_heads, bias=False)
-        self.proj_z = nn.Linear(c_z, num_heads, bias=False)
-        self.proj_o = nn.Linear(c_h * num_heads, c_m, bias=False)
-        init.final_init_(self.proj_o.weight)
-
-    def _attention_weights(self, pair_states: Tensor, mask: Tensor) -> Tensor:
-        # pair_states: (b, l, l, d_z); mask: (b, l, 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,
-        )  # (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  # (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,
-        )  # (b, s, l, d_h)
-        return (gate * attended) @ self.proj_o.weight[:, start:stop].T  # (b, s, l, d_m)
-
-    def forward(
-        self,
-        m: Tensor,
-        z: Tensor,
-        mask: Tensor,
-        chunk_heads: bool = True,
-    ) -> Tensor:
-        """Return updated M with shape ``(b, s, l, d_m)``."""
-
-        # m: (b, s, l, d_m); z: (b, l, l, d_z); mask: (b, l, 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
-            )  # (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  # (b, s, l, d_m)
+"""Pair-weighted aggregation of Boltz2 MSA states."""
+
+from __future__ import annotations
+
+import torch
+
+from torch import Tensor, nn
+
+from . import vb_layers_initialize as init
+
+
+class PairWeightedAveraging(nn.Module):
+    """Update M with pair-derived attention weights while preserving its shape."""
+
+    def __init__(
+        self,
+        c_m: int,
+        c_z: int,
+        c_h: int,
+        num_heads: int,
+        inf: float = 1e6,
+    ) -> None:
+        super().__init__()
+        self.c_m = c_m
+        self.c_z = c_z
+        self.c_h = c_h
+        self.num_heads = num_heads
+        self.inf = inf
+        self.norm_m = nn.LayerNorm(c_m)
+        self.norm_z = nn.LayerNorm(c_z)
+        self.proj_m = nn.Linear(c_m, c_h * num_heads, bias=False)
+        self.proj_g = nn.Linear(c_m, c_h * num_heads, bias=False)
+        self.proj_z = nn.Linear(c_z, num_heads, bias=False)
+        self.proj_o = nn.Linear(c_h * num_heads, c_m, bias=False)
+        init.final_init_(self.proj_o.weight)
+
+    def _attention_weights(self, pair_states: Tensor, mask: Tensor) -> Tensor:
+        # pair_states: (b, l, l, d_z); mask: (b, l, 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,
+        )  # (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  # (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,
+        )  # (b, s, l, d_h)
+        return (gate * attended) @ self.proj_o.weight[:, start:stop].T  # (b, s, l, d_m)
+
+    def forward(
+        self,
+        m: Tensor,
+        z: Tensor,
+        mask: Tensor,
+        chunk_heads: bool = True,
+    ) -> Tensor:
+        """Return updated M with shape ``(b, s, l, d_m)``."""
+
+        # m: (b, s, l, d_m); z: (b, l, l, d_z); mask: (b, l, 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
+            )  # (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  # (b, s, l, d_m)

fastplms/models/boltz/vb_layers_pairformer.py

@@ -1,350 +1,350 @@
-"""Pairformer layers for joint or pair-only Boltz2 representations."""
-
-from __future__ import annotations
-
-import torch
-
-from typing import Any, cast
-from torch import Tensor, nn
-from torch.utils.checkpoint import checkpoint
-
-from . import vb_const as const
-from .vb_layers_attention import AttentionPairBias
-from .vb_layers_attentionv2 import AttentionPairBias as AttentionPairBiasV2
-from .vb_layers_dropout import get_dropout_mask
-from .vb_layers_transition import Transition
-from .vb_layers_triangular_mult import (
-    TriangleMultiplicationIncoming,
-    TriangleMultiplicationOutgoing,
-)
-from .vb_tri_attn_attention import (
-    TriangleAttentionEndingNode,
-    TriangleAttentionStartingNode,
-)
-
-
-def _pair_update(
-    modules: Any,
-    pair_states: Tensor,
-    pair_mask: Tensor,
-    chunk_size: int | None,
-    use_kernels: bool,
-    use_cuequiv_mul: bool,
-    use_cuequiv_attn: bool,
-) -> 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  # (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
-        )  # 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
-    )  # 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,
-    )  # (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:
-    if training:
-        return None
-    return 128 if pair_states.shape[1] > const.chunk_size_threshold else 512
-
-
-class PairformerLayer(nn.Module):
-    """Update sequence tensor S and pair tensor Z once."""
-
-    def __init__(
-        self,
-        token_s: int,
-        token_z: int,
-        num_heads: int = 16,
-        dropout: float = 0.25,
-        pairwise_head_width: int = 32,
-        pairwise_num_heads: int = 4,
-        post_layer_norm: bool = False,
-        v2: bool = False,
-    ) -> None:
-        super().__init__()
-        self.token_z = token_z
-        self.dropout = dropout
-        self.num_heads = num_heads
-        self.post_layer_norm = post_layer_norm
-        self.pre_norm_s = nn.LayerNorm(token_s)
-        attention_class = AttentionPairBiasV2 if v2 else AttentionPairBias
-        self.attention = attention_class(token_s, token_z, num_heads)
-        self.tri_mul_out = TriangleMultiplicationOutgoing(token_z)
-        self.tri_mul_in = TriangleMultiplicationIncoming(token_z)
-        self.tri_att_start = TriangleAttentionStartingNode(
-            token_z,
-            pairwise_head_width,
-            pairwise_num_heads,
-            inf=1e9,
-        )
-        self.tri_att_end = TriangleAttentionEndingNode(
-            token_z,
-            pairwise_head_width,
-            pairwise_num_heads,
-            inf=1e9,
-        )
-        self.transition_s = Transition(token_s, token_s * 4)
-        self.transition_z = Transition(token_z, token_z * 4)
-        self.s_post_norm = nn.LayerNorm(token_s) if post_layer_norm else nn.Identity()
-
-    def forward(
-        self,
-        s: Tensor,
-        z: Tensor,
-        mask: Tensor,
-        pair_mask: Tensor,
-        chunk_size_tri_attn: int | None = None,
-        use_kernels: bool = False,
-        use_cuequiv_mul: bool = False,
-        use_cuequiv_attn: bool = False,
-    ) -> tuple[Tensor, Tensor]:
-        """Return updated ``S: (b, l, d_s)`` and ``Z: (b, l, l, d_z)``."""
-
-        pair_output = _pair_update(
-            self,
-            z,
-            pair_mask,
-            chunk_size_tri_attn,
-            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())  # (b, l, d_s)
-            sequence_output = s.float() + self.attention(
-                s=normalized,
-                z=pair_output.float(),
-                mask=mask.float(),
-                k_in=normalized,
-            )  # (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):
-    """Run a stack of joint sequence and pair updates."""
-
-    def __init__(
-        self,
-        token_s: int,
-        token_z: int,
-        num_blocks: int,
-        num_heads: int = 16,
-        dropout: float = 0.25,
-        pairwise_head_width: int = 32,
-        pairwise_num_heads: int = 4,
-        post_layer_norm: bool = False,
-        activation_checkpointing: bool = False,
-        v2: bool = False,
-        **kwargs: Any,
-    ) -> None:
-        del kwargs
-        super().__init__()
-        self.token_z = token_z
-        self.num_blocks = num_blocks
-        self.dropout = dropout
-        self.num_heads = num_heads
-        self.post_layer_norm = post_layer_norm
-        self.activation_checkpointing = activation_checkpointing
-        self.layers = nn.ModuleList(
-            [
-                PairformerLayer(
-                    token_s,
-                    token_z,
-                    num_heads,
-                    dropout,
-                    pairwise_head_width,
-                    pairwise_num_heads,
-                    post_layer_norm,
-                    v2,
-                )
-                for _ in range(num_blocks)
-            ]
-        )
-
-    def forward(
-        self,
-        s: Tensor,
-        z: Tensor,
-        mask: Tensor,
-        pair_mask: Tensor,
-        use_kernels: bool = False,
-    ) -> 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  # (b, l, d_s), (b, l, l, d_z)
-        for layer in self.layers:
-            if self.activation_checkpointing:
-                sequence_output, pair_output = checkpoint(
-                    layer,
-                    sequence_output,
-                    pair_output,
-                    mask,
-                    pair_mask,
-                    chunk_size,
-                    use_kernels,
-                    use_reentrant=False,
-                )  # (b, l, d_s), (b, l, l, d_z)
-            else:
-                sequence_output, pair_output = layer(
-                    sequence_output,
-                    pair_output,
-                    mask,
-                    pair_mask,
-                    chunk_size,
-                    use_kernels,
-                )  # (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):
-    """Update pair tensor Z without a sequence track."""
-
-    def __init__(
-        self,
-        token_z: int,
-        dropout: float = 0.25,
-        pairwise_head_width: int = 32,
-        pairwise_num_heads: int = 4,
-        post_layer_norm: bool = False,
-    ) -> None:
-        super().__init__()
-        self.token_z = token_z
-        self.dropout = dropout
-        self.post_layer_norm = post_layer_norm
-        self.tri_mul_out = TriangleMultiplicationOutgoing(token_z)
-        self.tri_mul_in = TriangleMultiplicationIncoming(token_z)
-        self.tri_att_start = TriangleAttentionStartingNode(
-            token_z,
-            pairwise_head_width,
-            pairwise_num_heads,
-            inf=1e9,
-        )
-        self.tri_att_end = TriangleAttentionEndingNode(
-            token_z,
-            pairwise_head_width,
-            pairwise_num_heads,
-            inf=1e9,
-        )
-        self.transition_z = Transition(token_z, token_z * 4)
-
-    def forward(
-        self,
-        z: Tensor,
-        pair_mask: Tensor,
-        chunk_size_tri_attn: int | None = None,
-        use_kernels: bool = False,
-        use_cuequiv_mul: bool = False,
-        use_cuequiv_attn: bool = False,
-    ) -> Tensor:
-        """Return updated pair tensor Z with shape ``(b, l, l, d_z)``."""
-
-        return _pair_update(
-            self,
-            z,
-            pair_mask,
-            chunk_size_tri_attn,
-            use_kernels,
-            use_cuequiv_mul,
-            use_cuequiv_attn,
-        )
-
-
-class PairformerNoSeqModule(nn.Module):
-    """Run a stack of pair-only updates."""
-
-    def __init__(
-        self,
-        token_z: int,
-        num_blocks: int,
-        dropout: float = 0.25,
-        pairwise_head_width: int = 32,
-        pairwise_num_heads: int = 4,
-        post_layer_norm: bool = False,
-        activation_checkpointing: bool = False,
-        **kwargs: Any,
-    ) -> None:
-        del kwargs
-        super().__init__()
-        self.token_z = token_z
-        self.num_blocks = num_blocks
-        self.dropout = dropout
-        self.post_layer_norm = post_layer_norm
-        self.activation_checkpointing = activation_checkpointing
-        self.layers = nn.ModuleList(
-            [
-                PairformerNoSeqLayer(
-                    token_z,
-                    dropout,
-                    pairwise_head_width,
-                    pairwise_num_heads,
-                    post_layer_norm,
-                )
-                for _ in range(num_blocks)
-            ]
-        )
-
-    def forward(
-        self,
-        z: Tensor,
-        pair_mask: Tensor,
-        use_kernels: bool = False,
-    ) -> 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  # (b, l, l, d_z)
-        for layer in self.layers:
-            if self.activation_checkpointing:
-                output = checkpoint(
-                    layer,
-                    output,
-                    pair_mask,
-                    chunk_size,
-                    use_kernels,
-                    use_reentrant=False,
-                )  # (b, l, l, d_z)
-            else:
-                output = layer(
-                    output, pair_mask, chunk_size, use_kernels
-                )  # (b, l, l, d_z)
-        return output  # (b, l, l, d_z)
+"""Pairformer layers for joint or pair-only Boltz2 representations."""
+
+from __future__ import annotations
+
+import torch
+
+from typing import Any, cast
+from torch import Tensor, nn
+from torch.utils.checkpoint import checkpoint
+
+from . import vb_const as const
+from .vb_layers_attention import AttentionPairBias
+from .vb_layers_attentionv2 import AttentionPairBias as AttentionPairBiasV2
+from .vb_layers_dropout import get_dropout_mask
+from .vb_layers_transition import Transition
+from .vb_layers_triangular_mult import (
+    TriangleMultiplicationIncoming,
+    TriangleMultiplicationOutgoing,
+)
+from .vb_tri_attn_attention import (
+    TriangleAttentionEndingNode,
+    TriangleAttentionStartingNode,
+)
+
+
+def _pair_update(
+    modules: Any,
+    pair_states: Tensor,
+    pair_mask: Tensor,
+    chunk_size: int | None,
+    use_kernels: bool,
+    use_cuequiv_mul: bool,
+    use_cuequiv_attn: bool,
+) -> 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  # (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
+        )  # 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
+    )  # 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,
+    )  # (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:
+    if training:
+        return None
+    return 128 if pair_states.shape[1] > const.chunk_size_threshold else 512
+
+
+class PairformerLayer(nn.Module):
+    """Update sequence tensor S and pair tensor Z once."""
+
+    def __init__(
+        self,
+        token_s: int,
+        token_z: int,
+        num_heads: int = 16,
+        dropout: float = 0.25,
+        pairwise_head_width: int = 32,
+        pairwise_num_heads: int = 4,
+        post_layer_norm: bool = False,
+        v2: bool = False,
+    ) -> None:
+        super().__init__()
+        self.token_z = token_z
+        self.dropout = dropout
+        self.num_heads = num_heads
+        self.post_layer_norm = post_layer_norm
+        self.pre_norm_s = nn.LayerNorm(token_s)
+        attention_class = AttentionPairBiasV2 if v2 else AttentionPairBias
+        self.attention = attention_class(token_s, token_z, num_heads)
+        self.tri_mul_out = TriangleMultiplicationOutgoing(token_z)
+        self.tri_mul_in = TriangleMultiplicationIncoming(token_z)
+        self.tri_att_start = TriangleAttentionStartingNode(
+            token_z,
+            pairwise_head_width,
+            pairwise_num_heads,
+            inf=1e9,
+        )
+        self.tri_att_end = TriangleAttentionEndingNode(
+            token_z,
+            pairwise_head_width,
+            pairwise_num_heads,
+            inf=1e9,
+        )
+        self.transition_s = Transition(token_s, token_s * 4)
+        self.transition_z = Transition(token_z, token_z * 4)
+        self.s_post_norm = nn.LayerNorm(token_s) if post_layer_norm else nn.Identity()
+
+    def forward(
+        self,
+        s: Tensor,
+        z: Tensor,
+        mask: Tensor,
+        pair_mask: Tensor,
+        chunk_size_tri_attn: int | None = None,
+        use_kernels: bool = False,
+        use_cuequiv_mul: bool = False,
+        use_cuequiv_attn: bool = False,
+    ) -> tuple[Tensor, Tensor]:
+        """Return updated ``S: (b, l, d_s)`` and ``Z: (b, l, l, d_z)``."""
+
+        pair_output = _pair_update(
+            self,
+            z,
+            pair_mask,
+            chunk_size_tri_attn,
+            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())  # (b, l, d_s)
+            sequence_output = s.float() + self.attention(
+                s=normalized,
+                z=pair_output.float(),
+                mask=mask.float(),
+                k_in=normalized,
+            )  # (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):
+    """Run a stack of joint sequence and pair updates."""
+
+    def __init__(
+        self,
+        token_s: int,
+        token_z: int,
+        num_blocks: int,
+        num_heads: int = 16,
+        dropout: float = 0.25,
+        pairwise_head_width: int = 32,
+        pairwise_num_heads: int = 4,
+        post_layer_norm: bool = False,
+        activation_checkpointing: bool = False,
+        v2: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        del kwargs
+        super().__init__()
+        self.token_z = token_z
+        self.num_blocks = num_blocks
+        self.dropout = dropout
+        self.num_heads = num_heads
+        self.post_layer_norm = post_layer_norm
+        self.activation_checkpointing = activation_checkpointing
+        self.layers = nn.ModuleList(
+            [
+                PairformerLayer(
+                    token_s,
+                    token_z,
+                    num_heads,
+                    dropout,
+                    pairwise_head_width,
+                    pairwise_num_heads,
+                    post_layer_norm,
+                    v2,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+
+    def forward(
+        self,
+        s: Tensor,
+        z: Tensor,
+        mask: Tensor,
+        pair_mask: Tensor,
+        use_kernels: bool = False,
+    ) -> 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  # (b, l, d_s), (b, l, l, d_z)
+        for layer in self.layers:
+            if self.activation_checkpointing:
+                sequence_output, pair_output = checkpoint(
+                    layer,
+                    sequence_output,
+                    pair_output,
+                    mask,
+                    pair_mask,
+                    chunk_size,
+                    use_kernels,
+                    use_reentrant=False,
+                )  # (b, l, d_s), (b, l, l, d_z)
+            else:
+                sequence_output, pair_output = layer(
+                    sequence_output,
+                    pair_output,
+                    mask,
+                    pair_mask,
+                    chunk_size,
+                    use_kernels,
+                )  # (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):
+    """Update pair tensor Z without a sequence track."""
+
+    def __init__(
+        self,
+        token_z: int,
+        dropout: float = 0.25,
+        pairwise_head_width: int = 32,
+        pairwise_num_heads: int = 4,
+        post_layer_norm: bool = False,
+    ) -> None:
+        super().__init__()
+        self.token_z = token_z
+        self.dropout = dropout
+        self.post_layer_norm = post_layer_norm
+        self.tri_mul_out = TriangleMultiplicationOutgoing(token_z)
+        self.tri_mul_in = TriangleMultiplicationIncoming(token_z)
+        self.tri_att_start = TriangleAttentionStartingNode(
+            token_z,
+            pairwise_head_width,
+            pairwise_num_heads,
+            inf=1e9,
+        )
+        self.tri_att_end = TriangleAttentionEndingNode(
+            token_z,
+            pairwise_head_width,
+            pairwise_num_heads,
+            inf=1e9,
+        )
+        self.transition_z = Transition(token_z, token_z * 4)
+
+    def forward(
+        self,
+        z: Tensor,
+        pair_mask: Tensor,
+        chunk_size_tri_attn: int | None = None,
+        use_kernels: bool = False,
+        use_cuequiv_mul: bool = False,
+        use_cuequiv_attn: bool = False,
+    ) -> Tensor:
+        """Return updated pair tensor Z with shape ``(b, l, l, d_z)``."""
+
+        return _pair_update(
+            self,
+            z,
+            pair_mask,
+            chunk_size_tri_attn,
+            use_kernels,
+            use_cuequiv_mul,
+            use_cuequiv_attn,
+        )
+
+
+class PairformerNoSeqModule(nn.Module):
+    """Run a stack of pair-only updates."""
+
+    def __init__(
+        self,
+        token_z: int,
+        num_blocks: int,
+        dropout: float = 0.25,
+        pairwise_head_width: int = 32,
+        pairwise_num_heads: int = 4,
+        post_layer_norm: bool = False,
+        activation_checkpointing: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        del kwargs
+        super().__init__()
+        self.token_z = token_z
+        self.num_blocks = num_blocks
+        self.dropout = dropout
+        self.post_layer_norm = post_layer_norm
+        self.activation_checkpointing = activation_checkpointing
+        self.layers = nn.ModuleList(
+            [
+                PairformerNoSeqLayer(
+                    token_z,
+                    dropout,
+                    pairwise_head_width,
+                    pairwise_num_heads,
+                    post_layer_norm,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+
+    def forward(
+        self,
+        z: Tensor,
+        pair_mask: Tensor,
+        use_kernels: bool = False,
+    ) -> 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  # (b, l, l, d_z)
+        for layer in self.layers:
+            if self.activation_checkpointing:
+                output = checkpoint(
+                    layer,
+                    output,
+                    pair_mask,
+                    chunk_size,
+                    use_kernels,
+                    use_reentrant=False,
+                )  # (b, l, l, d_z)
+            else:
+                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

@@ -1,79 +1,79 @@
-"""Gated feed-forward transition used by Boltz2 pair representations."""
-
-from __future__ import annotations
-
-import torch
-
-from torch import Tensor, nn
-
-from . import vb_layers_initialize as init
-
-
-class Transition(nn.Module):
-    """Apply a normalized SwiGLU-style transition.
-
-    The checkpoint-facing module names remain ``norm`` and ``fc1`` through
-    ``fc3``. For an input tensor X with shape ``(..., d)``, the module returns
-    a tensor with shape ``(..., d_out)``.
-    """
-
-    def __init__(
-        self,
-        dim: int = 128,
-        hidden: int = 512,
-        out_dim: int | None = None,
-    ) -> None:
-        super().__init__()
-        output_dim = dim if out_dim is None else out_dim
-        self.norm = nn.LayerNorm(dim, eps=1e-5)
-        self.fc1 = nn.Linear(dim, hidden, bias=False)
-        self.fc2 = nn.Linear(dim, hidden, bias=False)
-        self.fc3 = nn.Linear(hidden, output_dim, bias=False)
-        self.silu = nn.SiLU()
-        self.hidden = hidden
-
-        init.bias_init_one_(self.norm.weight)
-        init.bias_init_zero_(self.norm.bias)
-        init.lecun_normal_init_(self.fc1.weight)
-        init.lecun_normal_init_(self.fc2.weight)
-        init.final_init_(self.fc3.weight)
-
-    def _project_hidden_slice(self, normalized: Tensor, start: int, stop: int) -> Tensor:
-        """Return one hidden-width contribution to the output projection."""
-
-        # 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)  # (..., 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)  # (..., d_out)
-        if chunk_size <= 0:
-            raise ValueError("chunk_size must be positive")
-
-        output: Tensor | None = None
-        for start in range(0, self.hidden, chunk_size):
-            contribution = self._project_hidden_slice(
-                normalized,
-                start,
-                min(start + chunk_size, self.hidden),
-            )  # (..., 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  # (..., d_out)
+"""Gated feed-forward transition used by Boltz2 pair representations."""
+
+from __future__ import annotations
+
+import torch
+
+from torch import Tensor, nn
+
+from . import vb_layers_initialize as init
+
+
+class Transition(nn.Module):
+    """Apply a normalized SwiGLU-style transition.
+
+    The checkpoint-facing module names remain ``norm`` and ``fc1`` through
+    ``fc3``. For an input tensor X with shape ``(..., d)``, the module returns
+    a tensor with shape ``(..., d_out)``.
+    """
+
+    def __init__(
+        self,
+        dim: int = 128,
+        hidden: int = 512,
+        out_dim: int | None = None,
+    ) -> None:
+        super().__init__()
+        output_dim = dim if out_dim is None else out_dim
+        self.norm = nn.LayerNorm(dim, eps=1e-5)
+        self.fc1 = nn.Linear(dim, hidden, bias=False)
+        self.fc2 = nn.Linear(dim, hidden, bias=False)
+        self.fc3 = nn.Linear(hidden, output_dim, bias=False)
+        self.silu = nn.SiLU()
+        self.hidden = hidden
+
+        init.bias_init_one_(self.norm.weight)
+        init.bias_init_zero_(self.norm.bias)
+        init.lecun_normal_init_(self.fc1.weight)
+        init.lecun_normal_init_(self.fc2.weight)
+        init.final_init_(self.fc3.weight)
+
+    def _project_hidden_slice(self, normalized: Tensor, start: int, stop: int) -> Tensor:
+        """Return one hidden-width contribution to the output projection."""
+
+        # 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)  # (..., 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)  # (..., d_out)
+        if chunk_size <= 0:
+            raise ValueError("chunk_size must be positive")
+
+        output: Tensor | None = None
+        for start in range(0, self.hidden, chunk_size):
+            contribution = self._project_hidden_slice(
+                normalized,
+                start,
+                min(start + chunk_size, self.hidden),
+            )  # (..., 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  # (..., d_out)

fastplms/models/boltz/vb_layers_triangular_mult.py

@@ -1,132 +1,132 @@
-"""Incoming and outgoing triangular multiplicative pair updates."""
-
-from __future__ import annotations
-
-import importlib
-import torch
-
-from importlib.util import find_spec
-from typing import Literal
-from torch import Tensor, nn
-
-from . import vb_layers_initialize as init
-
-
-TriangleDirection = Literal["incoming", "outgoing"]
-
-
-@torch.compiler.disable
-def kernel_triangular_mult(
-    x: Tensor,
-    direction: TriangleDirection,
-    mask: Tensor,
-    norm_in_weight: Tensor,
-    norm_in_bias: Tensor,
-    p_in_weight: Tensor,
-    g_in_weight: Tensor,
-    norm_out_weight: Tensor,
-    norm_out_bias: Tensor,
-    p_out_weight: Tensor,
-    g_out_weight: Tensor,
-    eps: float,
-) -> 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
-    ):
-        raise RuntimeError(
-            "Boltz2 use_kernels=True requires cuequivariance_torch and the CUDA 13 "
-            "cuequivariance_ops_torch runtime from the 'structure,cueq' extras."
-        )
-    cueq = importlib.import_module("cuequivariance_torch")
-    return cueq.triangle_multiplicative_update(
-        x,
-        direction=direction,
-        mask=mask,
-        norm_in_weight=norm_in_weight,
-        norm_in_bias=norm_in_bias,
-        p_in_weight=p_in_weight,
-        g_in_weight=g_in_weight,
-        norm_out_weight=norm_out_weight,
-        norm_out_bias=norm_out_bias,
-        p_out_weight=p_out_weight,
-        g_out_weight=g_out_weight,
-        eps=eps,
-    )
-
-
-class _TriangleMultiplication(nn.Module):
-    direction: TriangleDirection
-    equation: str
-
-    def __init__(self, dim: int, direction: TriangleDirection, equation: str) -> None:
-        super().__init__()
-        self.direction = direction
-        self.equation = equation
-        self.norm_in = nn.LayerNorm(dim, eps=1e-5)
-        self.p_in = nn.Linear(dim, 2 * dim, bias=False)
-        self.g_in = nn.Linear(dim, 2 * dim, bias=False)
-        self.norm_out = nn.LayerNorm(dim)
-        self.p_out = nn.Linear(dim, dim, bias=False)
-        self.g_out = nn.Linear(dim, dim, bias=False)
-
-        init.bias_init_one_(self.norm_in.weight)
-        init.bias_init_zero_(self.norm_in.bias)
-        init.lecun_normal_init_(self.p_in.weight)
-        init.gating_init_(self.g_in.weight)
-        init.bias_init_one_(self.norm_out.weight)
-        init.bias_init_zero_(self.norm_out.bias)
-        init.final_init_(self.p_out.weight)
-        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,
-            mask=mask,
-            norm_in_weight=self.norm_in.weight,
-            norm_in_bias=self.norm_in.bias,
-            p_in_weight=self.p_in.weight,
-            g_in_weight=self.g_in.weight,
-            norm_out_weight=self.norm_out.weight,
-            norm_out_bias=self.norm_out.bias,
-            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)``."""
-
-        if use_kernels:
-            return self._kernel_forward(x, mask)
-
-        # X_norm is the normalized pair tensor used by the output gate.
-        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):
-    """Aggregate pair paths that share their destination index."""
-
-    def __init__(self, dim: int = 128) -> None:
-        super().__init__(dim, direction="outgoing", equation="bikd,bjkd->bijd")
-
-
-class TriangleMultiplicationIncoming(_TriangleMultiplication):
-    """Aggregate pair paths that share their source index."""
-
-    def __init__(self, dim: int = 128) -> None:
-        super().__init__(dim, direction="incoming", equation="bkid,bkjd->bijd")
+"""Incoming and outgoing triangular multiplicative pair updates."""
+
+from __future__ import annotations
+
+import importlib
+import torch
+
+from importlib.util import find_spec
+from typing import Literal
+from torch import Tensor, nn
+
+from . import vb_layers_initialize as init
+
+
+TriangleDirection = Literal["incoming", "outgoing"]
+
+
+@torch.compiler.disable
+def kernel_triangular_mult(
+    x: Tensor,
+    direction: TriangleDirection,
+    mask: Tensor,
+    norm_in_weight: Tensor,
+    norm_in_bias: Tensor,
+    p_in_weight: Tensor,
+    g_in_weight: Tensor,
+    norm_out_weight: Tensor,
+    norm_out_bias: Tensor,
+    p_out_weight: Tensor,
+    g_out_weight: Tensor,
+    eps: float,
+) -> 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
+    ):
+        raise RuntimeError(
+            "Boltz2 use_kernels=True requires cuequivariance_torch and the CUDA 13 "
+            "cuequivariance_ops_torch runtime from the 'structure,cueq' extras."
+        )
+    cueq = importlib.import_module("cuequivariance_torch")
+    return cueq.triangle_multiplicative_update(
+        x,
+        direction=direction,
+        mask=mask,
+        norm_in_weight=norm_in_weight,
+        norm_in_bias=norm_in_bias,
+        p_in_weight=p_in_weight,
+        g_in_weight=g_in_weight,
+        norm_out_weight=norm_out_weight,
+        norm_out_bias=norm_out_bias,
+        p_out_weight=p_out_weight,
+        g_out_weight=g_out_weight,
+        eps=eps,
+    )
+
+
+class _TriangleMultiplication(nn.Module):
+    direction: TriangleDirection
+    equation: str
+
+    def __init__(self, dim: int, direction: TriangleDirection, equation: str) -> None:
+        super().__init__()
+        self.direction = direction
+        self.equation = equation
+        self.norm_in = nn.LayerNorm(dim, eps=1e-5)
+        self.p_in = nn.Linear(dim, 2 * dim, bias=False)
+        self.g_in = nn.Linear(dim, 2 * dim, bias=False)
+        self.norm_out = nn.LayerNorm(dim)
+        self.p_out = nn.Linear(dim, dim, bias=False)
+        self.g_out = nn.Linear(dim, dim, bias=False)
+
+        init.bias_init_one_(self.norm_in.weight)
+        init.bias_init_zero_(self.norm_in.bias)
+        init.lecun_normal_init_(self.p_in.weight)
+        init.gating_init_(self.g_in.weight)
+        init.bias_init_one_(self.norm_out.weight)
+        init.bias_init_zero_(self.norm_out.bias)
+        init.final_init_(self.p_out.weight)
+        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,
+            mask=mask,
+            norm_in_weight=self.norm_in.weight,
+            norm_in_bias=self.norm_in.bias,
+            p_in_weight=self.p_in.weight,
+            g_in_weight=self.g_in.weight,
+            norm_out_weight=self.norm_out.weight,
+            norm_out_bias=self.norm_out.bias,
+            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)``."""
+
+        if use_kernels:
+            return self._kernel_forward(x, mask)
+
+        # X_norm is the normalized pair tensor used by the output gate.
+        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):
+    """Aggregate pair paths that share their destination index."""
+
+    def __init__(self, dim: int = 128) -> None:
+        super().__init__(dim, direction="outgoing", equation="bikd,bjkd->bijd")
+
+
+class TriangleMultiplicationIncoming(_TriangleMultiplication):
+    """Aggregate pair paths that share their source index."""
+
+    def __init__(self, dim: int = 128) -> None:
+        super().__init__(dim, direction="incoming", equation="bkid,bkjd->bijd")

fastplms/models/boltz/vb_loss_diffusionv2.py

@@ -1,186 +1,186 @@
-"""Geometry objectives shared by Boltz diffusion and steering code.
-
-The rigid-alignment mechanism is based on the Kabsch formulation used by
-AlphaFold 3 implementations.  The implementation is maintained locally and
-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
-
-
-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(
-    mask: torch.Tensor,
-    singular_values: torch.Tensor,
-    *,
-    num_points: int,
-    coordinate_dim: int,
-) -> None:
-    if torch.any(mask.sum(dim=-1) < coordinate_dim + 1):
-        warnings.warn(
-            "The size of one of the point clouds is <= dim+1. "
-            "`WeightedRigidAlign` cannot return a unique rotation.",
-            RuntimeWarning,
-            stacklevel=3,
-        )
-    if (singular_values.abs() <= 1e-15).any() and num_points >= coordinate_dim + 1:
-        warnings.warn(
-            "Excessively low rank of cross-correlation between aligned "
-            "point clouds. `WeightedRigidAlign` cannot return a unique rotation.",
-            RuntimeWarning,
-            stacklevel=3,
-        )
-
-
-def weighted_rigid_align(
-    true_coords: torch.Tensor,
-    pred_coords: torch.Tensor,
-    weights: torch.Tensor,
-    mask: torch.Tensor,
-) -> torch.Tensor:
-    """Align true coordinates to predicted coordinates with weighted Kabsch.
-
-    ``true_coords`` and ``pred_coords`` have shape ``(..., n, 3)``.  The
-    returned tensor is detached because alignment defines a fixed target for
-    the diffusion loss.
-    """
-
-    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)  # (..., n, 1)
-
-    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)  # (..., 3, 3)
-    left_vectors, singular_values, right_vectors_h = torch.linalg.svd(
-        covariance_fp32,
-        driver="gesvd" if covariance_fp32.is_cuda else None,
-    )  # left/right: (..., 3, 3); singular_values: (..., 3)
-    right_vectors = right_vectors_h.mH  # (..., 3, 3)
-    _warn_if_alignment_is_ambiguous(
-        mask,
-        singular_values,
-        num_points=num_points,
-        coordinate_dim=coordinate_dim,
-    )
-
-    preliminary_rotation = torch.einsum(
-        "... i j, ... k j -> ... i k",
-        left_vectors,
-        right_vectors,
-    ).to(torch.float32)  # (..., 3, 3)
-    orientation = torch.eye(
-        coordinate_dim,
-        dtype=covariance_fp32.dtype,
-        device=covariance.device,
-    )[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)  # (..., 3, 3)
-
-    aligned = (
-        einsum(true_centered, rotation, "... n i, ... j i -> ... n j") + pred_centroid
-    )  # (..., n, 3)
-    aligned.detach_()
-    return aligned  # (..., n, 3)
-
-
-def _smooth_lddt_for_example(
-    pred_coords: torch.Tensor,
-    true_coords: torch.Tensor,
-    is_nucleotide: torch.Tensor,
-    coords_mask: torch.Tensor,
-    *,
-    nucleic_acid_cutoff: float,
-    other_cutoff: float,
-) -> torch.Tensor:
-    # 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,
-    )  # (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]],
-    )  # (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
-    )  # (n_pair,)
-    return smooth_agreement.sum() / (pair_indices.shape[0] + 1e-5)  # ()
-
-
-def smooth_lddt_loss(
-    pred_coords: torch.Tensor,
-    true_coords: torch.Tensor,
-    is_nucleotide: torch.Tensor,
-    coords_mask: torch.Tensor,
-    nucleic_acid_cutoff: float = 30.0,
-    other_cutoff: float = 15.0,
-    multiplicity: int = 1,
-) -> torch.Tensor:
-    """Return one minus the smooth local-distance agreement.
-
-    Coordinate tensors have shape ``(b, n, 3)``.  Sequence-level masks may
-    be shared across repeated diffusion samples through ``multiplicity``.
-    """
-
-    agreements = [
-        _smooth_lddt_for_example(
-            pred_coords[index],
-            true_coords[index],
-            is_nucleotide[index // multiplicity],
-            coords_mask[index // multiplicity],
-            nucleic_acid_cutoff=nucleic_acid_cutoff,
-            other_cutoff=other_cutoff,
-        )
-        for index in range(true_coords.shape[0])
-    ]  # each: ()
-    return 1.0 - torch.stack(agreements).mean(dim=0)  # ()
+"""Geometry objectives shared by Boltz diffusion and steering code.
+
+The rigid-alignment mechanism is based on the Kabsch formulation used by
+AlphaFold 3 implementations.  The implementation is maintained locally and
+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
+
+
+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(
+    mask: torch.Tensor,
+    singular_values: torch.Tensor,
+    *,
+    num_points: int,
+    coordinate_dim: int,
+) -> None:
+    if torch.any(mask.sum(dim=-1) < coordinate_dim + 1):
+        warnings.warn(
+            "The size of one of the point clouds is <= dim+1. "
+            "`WeightedRigidAlign` cannot return a unique rotation.",
+            RuntimeWarning,
+            stacklevel=3,
+        )
+    if (singular_values.abs() <= 1e-15).any() and num_points >= coordinate_dim + 1:
+        warnings.warn(
+            "Excessively low rank of cross-correlation between aligned "
+            "point clouds. `WeightedRigidAlign` cannot return a unique rotation.",
+            RuntimeWarning,
+            stacklevel=3,
+        )
+
+
+def weighted_rigid_align(
+    true_coords: torch.Tensor,
+    pred_coords: torch.Tensor,
+    weights: torch.Tensor,
+    mask: torch.Tensor,
+) -> torch.Tensor:
+    """Align true coordinates to predicted coordinates with weighted Kabsch.
+
+    ``true_coords`` and ``pred_coords`` have shape ``(..., n, 3)``.  The
+    returned tensor is detached because alignment defines a fixed target for
+    the diffusion loss.
+    """
+
+    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)  # (..., n, 1)
+
+    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)  # (..., 3, 3)
+    left_vectors, singular_values, right_vectors_h = torch.linalg.svd(
+        covariance_fp32,
+        driver="gesvd" if covariance_fp32.is_cuda else None,
+    )  # left/right: (..., 3, 3); singular_values: (..., 3)
+    right_vectors = right_vectors_h.mH  # (..., 3, 3)
+    _warn_if_alignment_is_ambiguous(
+        mask,
+        singular_values,
+        num_points=num_points,
+        coordinate_dim=coordinate_dim,
+    )
+
+    preliminary_rotation = torch.einsum(
+        "... i j, ... k j -> ... i k",
+        left_vectors,
+        right_vectors,
+    ).to(torch.float32)  # (..., 3, 3)
+    orientation = torch.eye(
+        coordinate_dim,
+        dtype=covariance_fp32.dtype,
+        device=covariance.device,
+    )[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)  # (..., 3, 3)
+
+    aligned = (
+        einsum(true_centered, rotation, "... n i, ... j i -> ... n j") + pred_centroid
+    )  # (..., n, 3)
+    aligned.detach_()
+    return aligned  # (..., n, 3)
+
+
+def _smooth_lddt_for_example(
+    pred_coords: torch.Tensor,
+    true_coords: torch.Tensor,
+    is_nucleotide: torch.Tensor,
+    coords_mask: torch.Tensor,
+    *,
+    nucleic_acid_cutoff: float,
+    other_cutoff: float,
+) -> torch.Tensor:
+    # 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,
+    )  # (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]],
+    )  # (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
+    )  # (n_pair,)
+    return smooth_agreement.sum() / (pair_indices.shape[0] + 1e-5)  # ()
+
+
+def smooth_lddt_loss(
+    pred_coords: torch.Tensor,
+    true_coords: torch.Tensor,
+    is_nucleotide: torch.Tensor,
+    coords_mask: torch.Tensor,
+    nucleic_acid_cutoff: float = 30.0,
+    other_cutoff: float = 15.0,
+    multiplicity: int = 1,
+) -> torch.Tensor:
+    """Return one minus the smooth local-distance agreement.
+
+    Coordinate tensors have shape ``(b, n, 3)``.  Sequence-level masks may
+    be shared across repeated diffusion samples through ``multiplicity``.
+    """
+
+    agreements = [
+        _smooth_lddt_for_example(
+            pred_coords[index],
+            true_coords[index],
+            is_nucleotide[index // multiplicity],
+            coords_mask[index // multiplicity],
+            nucleic_acid_cutoff=nucleic_acid_cutoff,
+            other_cutoff=other_cutoff,
+        )
+        for index in range(true_coords.shape[0])
+    ]  # each: ()
+    return 1.0 - torch.stack(agreements).mean(dim=0)  # ()

fastplms/models/boltz/vb_modules_confidencev2.py

@@ -1,548 +1,548 @@
-"""Predict atom and token confidence from Boltz2 sequence and pair states."""
-
-import torch
-
-from typing import Any
-from torch import nn
-
-from . import vb_const as const
-from . import vb_layers_initialize as init
-from .vb_layers_confidence_utils import (
-    compute_aggregated_metric,
-    compute_ptms,
-)
-from .vb_layers_pairformer import PairformerModule
-from .vb_modules_encodersv2 import RelativePositionEncoder
-from .vb_modules_trunkv2 import (
-    ContactConditioning,
-)
-from .vb_modules_utils import LinearNoBias
-
-
-def _token_slot_logits_to_atom_logits(
-    token_logits: torch.Tensor,
-    atom_to_token: torch.Tensor,
-    atom_pad_mask: torch.Tensor,
-    *,
-    multiplicity: int,
-) -> torch.Tensor:
-    """Gather per-token atom-slot logits onto each example's atom table.
-
-    ``token_logits`` is ordered as ``(batch * multiplicity, token, slot,
-    channel)``. ``atom_to_token`` determines both the owning token and the
-    within-token slot of every atom, so heterogeneous atom counts do not leak
-    across batch rows.
-    """
-
-    # b: batch; m: multiplicity; a: atoms; t: tokens; c: channels.
-    if token_logits.ndim != 4:
-        raise ValueError(
-            "token_logits must have shape (batch * multiplicity, token, slot, channel), "
-            f"got {tuple(token_logits.shape)}."
-        )
-    if atom_to_token.ndim != 3:
-        raise ValueError(
-            "atom_to_token must have shape (batch, atom, token), "
-            f"got {tuple(atom_to_token.shape)}."
-        )
-    if atom_pad_mask.shape != atom_to_token.shape[:2]:
-        raise ValueError(
-            "atom_pad_mask must match the batch and atom axes of atom_to_token; "
-            f"got {tuple(atom_pad_mask.shape)} and {tuple(atom_to_token.shape)}."
-        )
-    if multiplicity < 1:
-        raise ValueError(f"multiplicity must be positive, got {multiplicity}.")
-
-    batch_size, atom_count, token_count = atom_to_token.shape
-    if token_logits.shape[0] != batch_size * multiplicity:
-        raise ValueError(
-            "token_logits batch axis must equal batch * multiplicity; "
-            f"got {token_logits.shape[0]} and {batch_size} * {multiplicity}."
-        )
-    if token_logits.shape[1] != token_count:
-        raise ValueError(
-            "token_logits and atom_to_token disagree on token count; "
-            f"got {token_logits.shape[1]} and {token_count}."
-        )
-
-    valid_atoms = atom_pad_mask.bool()  # (b, a)
-    assignments = atom_to_token.bool() & valid_atoms.unsqueeze(-1)  # (b, a, t)
-    token_index = assignments.to(dtype=torch.int64).argmax(dim=-1)  # (b, a)
-    # Cumulative one-hot counts give each atom its ordinal within its owning
-    # token without assuming that atoms from different tokens are contiguous.
-    cumulative_slots = assignments.to(dtype=torch.int64).cumsum(dim=1) - 1  # (b, a, t)
-    slot_index = (cumulative_slots * assignments).sum(dim=-1)  # (b, a)
-
-    slots_per_token = token_logits.shape[2]
-    flattened_index = token_index * slots_per_token + slot_index  # (b, a)
-    flattened_index = flattened_index.masked_fill(~valid_atoms, 0)  # (b, a)
-    flattened_index = flattened_index.repeat_interleave(multiplicity, dim=0)  # (b * m, a)
-    expanded_atom_mask = valid_atoms.repeat_interleave(multiplicity, dim=0)  # (b * m, a)
-
-    flattened_logits = token_logits.flatten(1, 2)  # (b * m, t * n_slot, c)
-    gather_index = flattened_index.unsqueeze(-1).expand(
-        -1,
-        atom_count,
-        flattened_logits.shape[-1],
-    )  # (b * m, a, c)
-    atom_logits = torch.gather(flattened_logits, dim=1, index=gather_index)  # (b * m, a, c)
-    return atom_logits * expanded_atom_mask.unsqueeze(-1).to(dtype=atom_logits.dtype)  # (b * m, a, c)
-
-
-class ConfidenceModule(nn.Module):
-    """Update confidence features with predicted geometry and a pairformer stack."""
-
-    def __init__(
-        self,
-        token_s: int,
-        token_z: int,
-        pairformer_args: dict[str, Any],
-        num_dist_bins: int = 64,
-        token_level_confidence: bool = True,
-        max_dist: float = 22,
-        add_s_to_z_prod: bool = False,
-        add_s_input_to_s: bool = False,
-        add_z_input_to_z: bool = False,
-        maximum_bond_distance: int = 0,
-        bond_type_feature: bool = False,
-        confidence_args: dict[str, Any] | None = None,
-        compile_pairformer: bool = False,
-        fix_sym_check: bool = False,
-        cyclic_pos_enc: bool = False,
-        return_latent_feats: bool = False,
-        conditioning_cutoff_min: float | None = None,
-        conditioning_cutoff_max: float | None = None,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__()
-        self.max_num_atoms_per_token = 23
-        self.no_update_s = pairformer_args.get("no_update_s", False)
-        boundaries = torch.linspace(2, max_dist, num_dist_bins - 1)  # (n_dist_bin - 1,)
-        self.register_buffer("boundaries", boundaries)
-        self.dist_bin_pairwise_embed = nn.Embedding(num_dist_bins, token_z)
-        init.gating_init_(self.dist_bin_pairwise_embed.weight)
-        self.token_level_confidence = token_level_confidence
-
-        self.s_to_z = LinearNoBias(token_s, token_z)
-        self.s_to_z_transpose = LinearNoBias(token_s, token_z)
-        init.gating_init_(self.s_to_z.weight)
-        init.gating_init_(self.s_to_z_transpose.weight)
-
-        self.add_s_to_z_prod = add_s_to_z_prod
-        if add_s_to_z_prod:
-            self.s_to_z_prod_in1 = LinearNoBias(token_s, token_z)
-            self.s_to_z_prod_in2 = LinearNoBias(token_s, token_z)
-            self.s_to_z_prod_out = LinearNoBias(token_z, token_z)
-            init.gating_init_(self.s_to_z_prod_out.weight)
-
-        self.s_inputs_norm = nn.LayerNorm(token_s)
-        if not self.no_update_s:
-            self.s_norm = nn.LayerNorm(token_s)
-        self.z_norm = nn.LayerNorm(token_z)
-
-        self.add_s_input_to_s = add_s_input_to_s
-        if add_s_input_to_s:
-            self.s_input_to_s = LinearNoBias(token_s, token_s)
-            init.gating_init_(self.s_input_to_s.weight)
-
-        self.add_z_input_to_z = add_z_input_to_z
-        if add_z_input_to_z:
-            self.rel_pos = RelativePositionEncoder(
-                token_z, fix_sym_check=fix_sym_check, cyclic_pos_enc=cyclic_pos_enc
-            )
-            self.token_bonds = nn.Linear(
-                1 if maximum_bond_distance == 0 else maximum_bond_distance + 2,
-                token_z,
-                bias=False,
-            )
-            self.bond_type_feature = bond_type_feature
-            if bond_type_feature:
-                self.token_bonds_type = nn.Embedding(len(const.bond_types) + 1, token_z)
-
-            self.contact_conditioning = ContactConditioning(
-                token_z=token_z,
-                cutoff_min=conditioning_cutoff_min,
-                cutoff_max=conditioning_cutoff_max,
-            )
-        pairformer_args["v2"] = True
-        self.pairformer_stack = PairformerModule(
-            token_s,
-            token_z,
-            **pairformer_args,
-        )
-        self.return_latent_feats = return_latent_feats
-
-        self.confidence_heads = ConfidenceHeads(
-            token_s,
-            token_z,
-            token_level_confidence=token_level_confidence,
-            **confidence_args,
-        )
-
-    def forward(
-        self,
-        s_inputs: torch.Tensor,  # (b, t, d_s)
+"""Predict atom and token confidence from Boltz2 sequence and pair states."""
+
+import torch
+
+from typing import Any
+from torch import nn
+
+from . import vb_const as const
+from . import vb_layers_initialize as init
+from .vb_layers_confidence_utils import (
+    compute_aggregated_metric,
+    compute_ptms,
+)
+from .vb_layers_pairformer import PairformerModule
+from .vb_modules_encodersv2 import RelativePositionEncoder
+from .vb_modules_trunkv2 import (
+    ContactConditioning,
+)
+from .vb_modules_utils import LinearNoBias
+
+
+def _token_slot_logits_to_atom_logits(
+    token_logits: torch.Tensor,
+    atom_to_token: torch.Tensor,
+    atom_pad_mask: torch.Tensor,
+    *,
+    multiplicity: int,
+) -> torch.Tensor:
+    """Gather per-token atom-slot logits onto each example's atom table.
+
+    ``token_logits`` is ordered as ``(batch * multiplicity, token, slot,
+    channel)``. ``atom_to_token`` determines both the owning token and the
+    within-token slot of every atom, so heterogeneous atom counts do not leak
+    across batch rows.
+    """
+
+    # b: batch; m: multiplicity; a: atoms; t: tokens; c: channels.
+    if token_logits.ndim != 4:
+        raise ValueError(
+            "token_logits must have shape (batch * multiplicity, token, slot, channel), "
+            f"got {tuple(token_logits.shape)}."
+        )
+    if atom_to_token.ndim != 3:
+        raise ValueError(
+            "atom_to_token must have shape (batch, atom, token), "
+            f"got {tuple(atom_to_token.shape)}."
+        )
+    if atom_pad_mask.shape != atom_to_token.shape[:2]:
+        raise ValueError(
+            "atom_pad_mask must match the batch and atom axes of atom_to_token; "
+            f"got {tuple(atom_pad_mask.shape)} and {tuple(atom_to_token.shape)}."
+        )
+    if multiplicity < 1:
+        raise ValueError(f"multiplicity must be positive, got {multiplicity}.")
+
+    batch_size, atom_count, token_count = atom_to_token.shape
+    if token_logits.shape[0] != batch_size * multiplicity:
+        raise ValueError(
+            "token_logits batch axis must equal batch * multiplicity; "
+            f"got {token_logits.shape[0]} and {batch_size} * {multiplicity}."
+        )
+    if token_logits.shape[1] != token_count:
+        raise ValueError(
+            "token_logits and atom_to_token disagree on token count; "
+            f"got {token_logits.shape[1]} and {token_count}."
+        )
+
+    valid_atoms = atom_pad_mask.bool()  # (b, a)
+    assignments = atom_to_token.bool() & valid_atoms.unsqueeze(-1)  # (b, a, t)
+    token_index = assignments.to(dtype=torch.int64).argmax(dim=-1)  # (b, a)
+    # Cumulative one-hot counts give each atom its ordinal within its owning
+    # token without assuming that atoms from different tokens are contiguous.
+    cumulative_slots = assignments.to(dtype=torch.int64).cumsum(dim=1) - 1  # (b, a, t)
+    slot_index = (cumulative_slots * assignments).sum(dim=-1)  # (b, a)
+
+    slots_per_token = token_logits.shape[2]
+    flattened_index = token_index * slots_per_token + slot_index  # (b, a)
+    flattened_index = flattened_index.masked_fill(~valid_atoms, 0)  # (b, a)
+    flattened_index = flattened_index.repeat_interleave(multiplicity, dim=0)  # (b * m, a)
+    expanded_atom_mask = valid_atoms.repeat_interleave(multiplicity, dim=0)  # (b * m, a)
+
+    flattened_logits = token_logits.flatten(1, 2)  # (b * m, t * n_slot, c)
+    gather_index = flattened_index.unsqueeze(-1).expand(
+        -1,
+        atom_count,
+        flattened_logits.shape[-1],
+    )  # (b * m, a, c)
+    atom_logits = torch.gather(flattened_logits, dim=1, index=gather_index)  # (b * m, a, c)
+    return atom_logits * expanded_atom_mask.unsqueeze(-1).to(dtype=atom_logits.dtype)  # (b * m, a, c)
+
+
+class ConfidenceModule(nn.Module):
+    """Update confidence features with predicted geometry and a pairformer stack."""
+
+    def __init__(
+        self,
+        token_s: int,
+        token_z: int,
+        pairformer_args: dict[str, Any],
+        num_dist_bins: int = 64,
+        token_level_confidence: bool = True,
+        max_dist: float = 22,
+        add_s_to_z_prod: bool = False,
+        add_s_input_to_s: bool = False,
+        add_z_input_to_z: bool = False,
+        maximum_bond_distance: int = 0,
+        bond_type_feature: bool = False,
+        confidence_args: dict[str, Any] | None = None,
+        compile_pairformer: bool = False,
+        fix_sym_check: bool = False,
+        cyclic_pos_enc: bool = False,
+        return_latent_feats: bool = False,
+        conditioning_cutoff_min: float | None = None,
+        conditioning_cutoff_max: float | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self.max_num_atoms_per_token = 23
+        self.no_update_s = pairformer_args.get("no_update_s", False)
+        boundaries = torch.linspace(2, max_dist, num_dist_bins - 1)  # (n_dist_bin - 1,)
+        self.register_buffer("boundaries", boundaries)
+        self.dist_bin_pairwise_embed = nn.Embedding(num_dist_bins, token_z)
+        init.gating_init_(self.dist_bin_pairwise_embed.weight)
+        self.token_level_confidence = token_level_confidence
+
+        self.s_to_z = LinearNoBias(token_s, token_z)
+        self.s_to_z_transpose = LinearNoBias(token_s, token_z)
+        init.gating_init_(self.s_to_z.weight)
+        init.gating_init_(self.s_to_z_transpose.weight)
+
+        self.add_s_to_z_prod = add_s_to_z_prod
+        if add_s_to_z_prod:
+            self.s_to_z_prod_in1 = LinearNoBias(token_s, token_z)
+            self.s_to_z_prod_in2 = LinearNoBias(token_s, token_z)
+            self.s_to_z_prod_out = LinearNoBias(token_z, token_z)
+            init.gating_init_(self.s_to_z_prod_out.weight)
+
+        self.s_inputs_norm = nn.LayerNorm(token_s)
+        if not self.no_update_s:
+            self.s_norm = nn.LayerNorm(token_s)
+        self.z_norm = nn.LayerNorm(token_z)
+
+        self.add_s_input_to_s = add_s_input_to_s
+        if add_s_input_to_s:
+            self.s_input_to_s = LinearNoBias(token_s, token_s)
+            init.gating_init_(self.s_input_to_s.weight)
+
+        self.add_z_input_to_z = add_z_input_to_z
+        if add_z_input_to_z:
+            self.rel_pos = RelativePositionEncoder(
+                token_z, fix_sym_check=fix_sym_check, cyclic_pos_enc=cyclic_pos_enc
+            )
+            self.token_bonds = nn.Linear(
+                1 if maximum_bond_distance == 0 else maximum_bond_distance + 2,
+                token_z,
+                bias=False,
+            )
+            self.bond_type_feature = bond_type_feature
+            if bond_type_feature:
+                self.token_bonds_type = nn.Embedding(len(const.bond_types) + 1, token_z)
+
+            self.contact_conditioning = ContactConditioning(
+                token_z=token_z,
+                cutoff_min=conditioning_cutoff_min,
+                cutoff_max=conditioning_cutoff_max,
+            )
+        pairformer_args["v2"] = True
+        self.pairformer_stack = PairformerModule(
+            token_s,
+            token_z,
+            **pairformer_args,
+        )
+        self.return_latent_feats = return_latent_feats
+
+        self.confidence_heads = ConfidenceHeads(
+            token_s,
+            token_z,
+            token_level_confidence=token_level_confidence,
+            **confidence_args,
+        )
+
+    def forward(
+        self,
+        s_inputs: torch.Tensor,  # (b, t, d_s)
         s: torch.Tensor,  # (b, t, d_s)
         z: torch.Tensor,  # (b, t, t, d_z)
-        x_pred: torch.Tensor,  # (b * m, a, 3) or (b, m, a, 3)
-        feats: dict[str, torch.Tensor],
+        x_pred: torch.Tensor,  # (b * m, a, 3) or (b, m, a, 3)
+        feats: dict[str, torch.Tensor],
         pred_distogram_logits: torch.Tensor,  # (b, t, t, n_bin)
-        multiplicity: int = 1,
-        run_sequentially: bool = False,
-        use_kernels: bool = False,
-    ) -> dict[str, Any]:
-        # b: base batch; m: diffusion multiplicity; t: tokens; a: atoms.
-        if run_sequentially and multiplicity > 1:
-            batch_size = z.shape[0]
-            expected_shape = (batch_size, multiplicity)
-            if x_pred.ndim >= 4 and x_pred.shape[:2] == expected_shape:
-                sample_coordinates = x_pred  # (b, m, a, 3)
-            elif x_pred.shape[0] == batch_size * multiplicity:
-                sample_coordinates = x_pred.reshape(
-                    batch_size,
-                    multiplicity,
-                    *x_pred.shape[1:],
-                )  # (b, m, a, 3)
-            else:
-                raise ValueError(
-                    "Sequential confidence expected coordinates with leading shape "
-                    f"{expected_shape} or {batch_size * multiplicity}, got "
-                    f"{tuple(x_pred.shape)}."
-                )
-            out_dicts = []
-            for sample_idx in range(multiplicity):
-                out_dicts.append(
-                    self.forward(
-                        s_inputs,
-                        s,
-                        z,
-                        sample_coordinates[:, sample_idx],
-                        feats,
-                        pred_distogram_logits,
-                        multiplicity=1,
-                        run_sequentially=False,
-                        use_kernels=use_kernels,
-                    )
-                )
-
-            out_dict = {}
-            for key in out_dicts[0]:
-                if key != "pair_chains_iptm":
-                    values = [out[key] for out in out_dicts]  # list of sample tensors, each (b, *feature_shape)
-                    out_dict[key] = torch.stack(values, dim=1).flatten(0, 1)  # (b * m, *feature_shape)
-                else:
-                    pair_chains_iptm = {}
-                    for chain_idx1 in out_dicts[0][key]:
-                        chains_iptm = {}
-                        for chain_idx2 in out_dicts[0][key][chain_idx1]:
-                            values = [
-                                out[key][chain_idx1][chain_idx2]
-                                for out in out_dicts
+        multiplicity: int = 1,
+        run_sequentially: bool = False,
+        use_kernels: bool = False,
+    ) -> dict[str, Any]:
+        # b: base batch; m: diffusion multiplicity; t: tokens; a: atoms.
+        if run_sequentially and multiplicity > 1:
+            batch_size = z.shape[0]
+            expected_shape = (batch_size, multiplicity)
+            if x_pred.ndim >= 4 and x_pred.shape[:2] == expected_shape:
+                sample_coordinates = x_pred  # (b, m, a, 3)
+            elif x_pred.shape[0] == batch_size * multiplicity:
+                sample_coordinates = x_pred.reshape(
+                    batch_size,
+                    multiplicity,
+                    *x_pred.shape[1:],
+                )  # (b, m, a, 3)
+            else:
+                raise ValueError(
+                    "Sequential confidence expected coordinates with leading shape "
+                    f"{expected_shape} or {batch_size * multiplicity}, got "
+                    f"{tuple(x_pred.shape)}."
+                )
+            out_dicts = []
+            for sample_idx in range(multiplicity):
+                out_dicts.append(
+                    self.forward(
+                        s_inputs,
+                        s,
+                        z,
+                        sample_coordinates[:, sample_idx],
+                        feats,
+                        pred_distogram_logits,
+                        multiplicity=1,
+                        run_sequentially=False,
+                        use_kernels=use_kernels,
+                    )
+                )
+
+            out_dict = {}
+            for key in out_dicts[0]:
+                if key != "pair_chains_iptm":
+                    values = [out[key] for out in out_dicts]  # list of sample tensors, each (b, *feature_shape)
+                    out_dict[key] = torch.stack(values, dim=1).flatten(0, 1)  # (b * m, *feature_shape)
+                else:
+                    pair_chains_iptm = {}
+                    for chain_idx1 in out_dicts[0][key]:
+                        chains_iptm = {}
+                        for chain_idx2 in out_dicts[0][key][chain_idx1]:
+                            values = [
+                                out[key][chain_idx1][chain_idx2]
+                                for out in out_dicts
                             ]  # list of sample chain-pair scores, each (b,)
-                            chains_iptm[chain_idx2] = torch.stack(
-                                values,
-                                dim=1,
-                            ).flatten(0, 1)  # (b * m,)
-                        pair_chains_iptm[chain_idx1] = chains_iptm
-                    out_dict[key] = pair_chains_iptm  # mapping of chain-pair scores, each (b * m,)
-            return out_dict
-
-        s_inputs = self.s_inputs_norm(s_inputs)  # (b, t, d_s)
-        if not self.no_update_s:
-            s = self.s_norm(s)  # (b, t, d_s)
-
-        if self.add_s_input_to_s:
-            s = s + self.s_input_to_s(s_inputs)  # (b, t, d_s)
-
-        z = self.z_norm(z)  # (b, t, t, d_z)
-
-        if self.add_z_input_to_z:
-            relative_position_encoding = self.rel_pos(feats)  # (b, t, t, d_z)
-            z = z + relative_position_encoding  # (b, t, t, d_z)
-            z = z + self.token_bonds(feats["token_bonds"].float())  # (b, t, t, d_z)
-            if self.bond_type_feature:
-                z = z + self.token_bonds_type(feats["type_bonds"].long())  # (b, t, t, d_z)
-            z = z + self.contact_conditioning(feats)  # (b, t, t, d_z)
-
-        s = s.repeat_interleave(multiplicity, 0)  # (b * m, t, d_s)
-
-        z = (
-            z
-            + self.s_to_z(s_inputs)[:, :, None, :]
-            + self.s_to_z_transpose(s_inputs)[:, None, :, :]
-        )  # (b, t, t, d_z)
-        if self.add_s_to_z_prod:
-            z = z + self.s_to_z_prod_out(
-                self.s_to_z_prod_in1(s_inputs)[:, :, None, :]
-                * self.s_to_z_prod_in2(s_inputs)[:, None, :, :]
-            )  # (b, t, t, d_z)
-
-        z = z.repeat_interleave(multiplicity, 0)  # (b * m, t, t, d_z)
-        s_inputs = s_inputs.repeat_interleave(multiplicity, 0)  # (b * m, t, d_s)
-
+                            chains_iptm[chain_idx2] = torch.stack(
+                                values,
+                                dim=1,
+                            ).flatten(0, 1)  # (b * m,)
+                        pair_chains_iptm[chain_idx1] = chains_iptm
+                    out_dict[key] = pair_chains_iptm  # mapping of chain-pair scores, each (b * m,)
+            return out_dict
+
+        s_inputs = self.s_inputs_norm(s_inputs)  # (b, t, d_s)
+        if not self.no_update_s:
+            s = self.s_norm(s)  # (b, t, d_s)
+
+        if self.add_s_input_to_s:
+            s = s + self.s_input_to_s(s_inputs)  # (b, t, d_s)
+
+        z = self.z_norm(z)  # (b, t, t, d_z)
+
+        if self.add_z_input_to_z:
+            relative_position_encoding = self.rel_pos(feats)  # (b, t, t, d_z)
+            z = z + relative_position_encoding  # (b, t, t, d_z)
+            z = z + self.token_bonds(feats["token_bonds"].float())  # (b, t, t, d_z)
+            if self.bond_type_feature:
+                z = z + self.token_bonds_type(feats["type_bonds"].long())  # (b, t, t, d_z)
+            z = z + self.contact_conditioning(feats)  # (b, t, t, d_z)
+
+        s = s.repeat_interleave(multiplicity, 0)  # (b * m, t, d_s)
+
+        z = (
+            z
+            + self.s_to_z(s_inputs)[:, :, None, :]
+            + self.s_to_z_transpose(s_inputs)[:, None, :, :]
+        )  # (b, t, t, d_z)
+        if self.add_s_to_z_prod:
+            z = z + self.s_to_z_prod_out(
+                self.s_to_z_prod_in1(s_inputs)[:, :, None, :]
+                * self.s_to_z_prod_in2(s_inputs)[:, None, :, :]
+            )  # (b, t, t, d_z)
+
+        z = z.repeat_interleave(multiplicity, 0)  # (b * m, t, t, d_z)
+        s_inputs = s_inputs.repeat_interleave(multiplicity, 0)  # (b * m, t, d_s)
+
         token_to_rep_atom = feats["token_to_rep_atom"]  # (b, t, a)
         token_to_rep_atom = token_to_rep_atom.repeat_interleave(multiplicity, 0)  # (b * m, t, a)
-        if len(x_pred.shape) == 4:
-            b, multiplicity, n, _ = x_pred.shape
-            x_pred = x_pred.reshape(b * multiplicity, n, -1)  # (b * m, a, 3)
-        else:
-            _, n, _ = x_pred.shape
-        x_pred_repr = torch.bmm(token_to_rep_atom.float(), x_pred)  # (b * m, t, 3)
-        d = torch.cdist(x_pred_repr, x_pred_repr)  # (b * m, t, t)
-        distogram = (d.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()  # (b * m, t, t)
-        distogram = self.dist_bin_pairwise_embed(distogram)  # (b * m, t, t, d_z)
+        if len(x_pred.shape) == 4:
+            b, multiplicity, n, _ = x_pred.shape
+            x_pred = x_pred.reshape(b * multiplicity, n, -1)  # (b * m, a, 3)
+        else:
+            _, n, _ = x_pred.shape
+        x_pred_repr = torch.bmm(token_to_rep_atom.float(), x_pred)  # (b * m, t, 3)
+        d = torch.cdist(x_pred_repr, x_pred_repr)  # (b * m, t, t)
+        distogram = (d.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()  # (b * m, t, t)
+        distogram = self.dist_bin_pairwise_embed(distogram)  # (b * m, t, t, d_z)
         z = z + distogram  # (b * m, t, t, d_z)
-
-        mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, t)
-        pair_mask = mask[:, :, None] * mask[:, None, :]  # (b * m, t, t)
-
-        s_t, z_t = self.pairformer_stack(
-            s, z, mask=mask, pair_mask=pair_mask, use_kernels=use_kernels
-        )  # (b * m, t, d_s), (b * m, t, t, d_z)
-
-        # Confidence heads consume the pairformer states without an additional residual.
-        s = s_t  # (b * m, t, d_s)
-        z = z_t  # (b * m, t, t, d_z)
-
-        out_dict = {}
-
-        if self.return_latent_feats:
-            out_dict["s_conf"] = s  # (b * m, t, d_s)
-            out_dict["z_conf"] = z  # (b * m, t, t, d_z)
-
-        out_dict.update(
-            self.confidence_heads(
-                s=s,
-                z=z,
-                x_pred=x_pred,
-                d=d,
-                feats=feats,
-                multiplicity=multiplicity,
-                pred_distogram_logits=pred_distogram_logits,
-            )
-        )
-        return out_dict
-
-
-class ConfidenceHeads(nn.Module):
-    def __init__(
-        self,
-        token_s: int,
-        token_z: int,
-        num_plddt_bins: int = 50,
-        num_pde_bins: int = 64,
-        num_pae_bins: int = 64,
-        token_level_confidence: bool = True,
-        use_separate_heads: bool = False,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__()
-        self.max_num_atoms_per_token = 23
-        self.token_level_confidence = token_level_confidence
-        self.use_separate_heads = use_separate_heads
-
-        if self.use_separate_heads:
-            self.to_pae_intra_logits = LinearNoBias(token_z, num_pae_bins)
-            self.to_pae_inter_logits = LinearNoBias(token_z, num_pae_bins)
-        else:
-            self.to_pae_logits = LinearNoBias(token_z, num_pae_bins)
-
-        if self.use_separate_heads:
-            self.to_pde_intra_logits = LinearNoBias(token_z, num_pde_bins)
-            self.to_pde_inter_logits = LinearNoBias(token_z, num_pde_bins)
-        else:
-            self.to_pde_logits = LinearNoBias(token_z, num_pde_bins)
-
-        if self.token_level_confidence:
-            self.to_plddt_logits = LinearNoBias(token_s, num_plddt_bins)
-            self.to_resolved_logits = LinearNoBias(token_s, 2)
-        else:
-            self.to_plddt_logits = LinearNoBias(
-                token_s, num_plddt_bins * self.max_num_atoms_per_token
-            )
-            self.to_resolved_logits = LinearNoBias(token_s, 2 * self.max_num_atoms_per_token)
-
-    def forward(
-        self,
+
+        mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, t)
+        pair_mask = mask[:, :, None] * mask[:, None, :]  # (b * m, t, t)
+
+        s_t, z_t = self.pairformer_stack(
+            s, z, mask=mask, pair_mask=pair_mask, use_kernels=use_kernels
+        )  # (b * m, t, d_s), (b * m, t, t, d_z)
+
+        # Confidence heads consume the pairformer states without an additional residual.
+        s = s_t  # (b * m, t, d_s)
+        z = z_t  # (b * m, t, t, d_z)
+
+        out_dict = {}
+
+        if self.return_latent_feats:
+            out_dict["s_conf"] = s  # (b * m, t, d_s)
+            out_dict["z_conf"] = z  # (b * m, t, t, d_z)
+
+        out_dict.update(
+            self.confidence_heads(
+                s=s,
+                z=z,
+                x_pred=x_pred,
+                d=d,
+                feats=feats,
+                multiplicity=multiplicity,
+                pred_distogram_logits=pred_distogram_logits,
+            )
+        )
+        return out_dict
+
+
+class ConfidenceHeads(nn.Module):
+    def __init__(
+        self,
+        token_s: int,
+        token_z: int,
+        num_plddt_bins: int = 50,
+        num_pde_bins: int = 64,
+        num_pae_bins: int = 64,
+        token_level_confidence: bool = True,
+        use_separate_heads: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self.max_num_atoms_per_token = 23
+        self.token_level_confidence = token_level_confidence
+        self.use_separate_heads = use_separate_heads
+
+        if self.use_separate_heads:
+            self.to_pae_intra_logits = LinearNoBias(token_z, num_pae_bins)
+            self.to_pae_inter_logits = LinearNoBias(token_z, num_pae_bins)
+        else:
+            self.to_pae_logits = LinearNoBias(token_z, num_pae_bins)
+
+        if self.use_separate_heads:
+            self.to_pde_intra_logits = LinearNoBias(token_z, num_pde_bins)
+            self.to_pde_inter_logits = LinearNoBias(token_z, num_pde_bins)
+        else:
+            self.to_pde_logits = LinearNoBias(token_z, num_pde_bins)
+
+        if self.token_level_confidence:
+            self.to_plddt_logits = LinearNoBias(token_s, num_plddt_bins)
+            self.to_resolved_logits = LinearNoBias(token_s, 2)
+        else:
+            self.to_plddt_logits = LinearNoBias(
+                token_s, num_plddt_bins * self.max_num_atoms_per_token
+            )
+            self.to_resolved_logits = LinearNoBias(token_s, 2 * self.max_num_atoms_per_token)
+
+    def forward(
+        self,
         s: torch.Tensor,  # (b * m, t, d_s)
         z: torch.Tensor,  # (b * m, t, t, d_z)
         x_pred: torch.Tensor,  # (b * m, a, 3)
         d: torch.Tensor,  # (b * m, t, t)
-        feats: dict[str, torch.Tensor],
+        feats: dict[str, torch.Tensor],
         pred_distogram_logits: torch.Tensor,  # (b, t, t, n_bin)
-        multiplicity: int = 1,
-    ) -> dict[str, Any]:
-        # b: base batch; m: diffusion multiplicity; t: tokens; a: atoms.
-        if self.use_separate_heads:
-            asym_id_token = feats["asym_id"]  # (b, t)
-            is_same_chain = asym_id_token.unsqueeze(-1) == asym_id_token.unsqueeze(-2)  # (b, t, t)
-            is_different_chain = ~is_same_chain  # (b, t, t)
-
-        if self.use_separate_heads:
-            pae_intra_logits = self.to_pae_intra_logits(z)  # (b * m, t, t, n_pae)
-            pae_intra_logits = pae_intra_logits * is_same_chain.float().unsqueeze(-1)  # (b * m, t, t, n_pae)
-
-            pae_inter_logits = self.to_pae_inter_logits(z)  # (b * m, t, t, n_pae)
-            pae_inter_logits = pae_inter_logits * is_different_chain.float().unsqueeze(-1)  # (b * m, t, t, n_pae)
-
-            pae_logits = pae_inter_logits + pae_intra_logits  # (b * m, t, t, n_pae)
-        else:
-            pae_logits = self.to_pae_logits(z)  # (b * m, t, t, n_pae)
-
-        if self.use_separate_heads:
-            pde_intra_logits = self.to_pde_intra_logits(z + z.transpose(1, 2))  # (b * m, t, t, n_pde)
-            pde_intra_logits = pde_intra_logits * is_same_chain.float().unsqueeze(-1)  # (b * m, t, t, n_pde)
-
-            pde_inter_logits = self.to_pde_inter_logits(z + z.transpose(1, 2))  # (b * m, t, t, n_pde)
-            pde_inter_logits = pde_inter_logits * is_different_chain.float().unsqueeze(-1)  # (b * m, t, t, n_pde)
-
-            pde_logits = pde_inter_logits + pde_intra_logits  # (b * m, t, t, n_pde)
-        else:
-            pde_logits = self.to_pde_logits(z + z.transpose(1, 2))  # (b * m, t, t, n_pde)
-        resolved_logits = self.to_resolved_logits(s)  # (b * m, t, c_resolved); c_resolved is 2 or 2 * n_slot
-        plddt_logits = self.to_plddt_logits(s)  # (b * m, t, c_plddt); c_plddt is n_plddt or n_slot * n_plddt
-
-        ligand_weight = 20
-        non_interface_weight = 1
-        interface_weight = 10
-
+        multiplicity: int = 1,
+    ) -> dict[str, Any]:
+        # b: base batch; m: diffusion multiplicity; t: tokens; a: atoms.
+        if self.use_separate_heads:
+            asym_id_token = feats["asym_id"]  # (b, t)
+            is_same_chain = asym_id_token.unsqueeze(-1) == asym_id_token.unsqueeze(-2)  # (b, t, t)
+            is_different_chain = ~is_same_chain  # (b, t, t)
+
+        if self.use_separate_heads:
+            pae_intra_logits = self.to_pae_intra_logits(z)  # (b * m, t, t, n_pae)
+            pae_intra_logits = pae_intra_logits * is_same_chain.float().unsqueeze(-1)  # (b * m, t, t, n_pae)
+
+            pae_inter_logits = self.to_pae_inter_logits(z)  # (b * m, t, t, n_pae)
+            pae_inter_logits = pae_inter_logits * is_different_chain.float().unsqueeze(-1)  # (b * m, t, t, n_pae)
+
+            pae_logits = pae_inter_logits + pae_intra_logits  # (b * m, t, t, n_pae)
+        else:
+            pae_logits = self.to_pae_logits(z)  # (b * m, t, t, n_pae)
+
+        if self.use_separate_heads:
+            pde_intra_logits = self.to_pde_intra_logits(z + z.transpose(1, 2))  # (b * m, t, t, n_pde)
+            pde_intra_logits = pde_intra_logits * is_same_chain.float().unsqueeze(-1)  # (b * m, t, t, n_pde)
+
+            pde_inter_logits = self.to_pde_inter_logits(z + z.transpose(1, 2))  # (b * m, t, t, n_pde)
+            pde_inter_logits = pde_inter_logits * is_different_chain.float().unsqueeze(-1)  # (b * m, t, t, n_pde)
+
+            pde_logits = pde_inter_logits + pde_intra_logits  # (b * m, t, t, n_pde)
+        else:
+            pde_logits = self.to_pde_logits(z + z.transpose(1, 2))  # (b * m, t, t, n_pde)
+        resolved_logits = self.to_resolved_logits(s)  # (b * m, t, c_resolved); c_resolved is 2 or 2 * n_slot
+        plddt_logits = self.to_plddt_logits(s)  # (b * m, t, c_plddt); c_plddt is n_plddt or n_slot * n_plddt
+
+        ligand_weight = 20
+        non_interface_weight = 1
+        interface_weight = 10
+
         token_type = feats["mol_type"]  # (b, t)
         token_type = token_type.repeat_interleave(multiplicity, 0)  # (b * m, t)
-        is_ligand_token = (token_type == const.chain_type_ids["NONPOLYMER"]).float()  # (b * m, t)
-
-        if self.token_level_confidence:
-            plddt = compute_aggregated_metric(plddt_logits)  # (b * m, n_item); n_item is t for token confidence, a for atom confidence
-            token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, t)
-            complex_plddt = (plddt * token_pad_mask).sum(dim=-1) / token_pad_mask.sum(dim=-1)  # (b * m,)
-
-            is_contact = (d < 8).float()  # (b * m, t, t)
-            is_different_chain = (
-                feats["asym_id"].unsqueeze(-1) != feats["asym_id"].unsqueeze(-2)
-            ).float()  # (b, t, t)
-            is_different_chain = is_different_chain.repeat_interleave(multiplicity, 0)  # (b * m, t, t)
-            token_interface_mask = torch.max(
-                is_contact * is_different_chain * (1 - is_ligand_token).unsqueeze(-1),
-                dim=-1,
-            ).values  # (b * m, t)
-            token_non_interface_mask = (1 - token_interface_mask) * (1 - is_ligand_token)  # (b * m, t)
-            iplddt_weight = (
-                is_ligand_token * ligand_weight
-                + token_interface_mask * interface_weight
-                + token_non_interface_mask * non_interface_weight
-            )  # (b * m, t)
-            complex_iplddt = (plddt * token_pad_mask * iplddt_weight).sum(dim=-1) / torch.sum(
-                token_pad_mask * iplddt_weight, dim=-1
-            )  # (b * m,)
-
-        else:
-            # token to atom conversion for resolved logits
-            b, n, _ = resolved_logits.shape
-            resolved_logits = resolved_logits.reshape(b, n, self.max_num_atoms_per_token, 2)  # (b * m, t, n_slot, 2)
-            resolved_logits = _token_slot_logits_to_atom_logits(
-                resolved_logits,
-                feats["atom_to_token"],
-                feats["atom_pad_mask"],
-                multiplicity=multiplicity,
-            )  # (b * m, a, 2)
-            plddt_logits = plddt_logits.reshape(b, n, self.max_num_atoms_per_token, -1)  # (b * m, t, n_slot, n_plddt)
-            plddt_logits = _token_slot_logits_to_atom_logits(
-                plddt_logits,
-                feats["atom_to_token"],
-                feats["atom_pad_mask"],
-                multiplicity=multiplicity,
-            )  # (b * m, a, n_plddt)
-            atom_pad_mask = feats["atom_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, a)
-            plddt = compute_aggregated_metric(plddt_logits)  # (b * m, n_item); n_item is t for token confidence, a for atom confidence
-
-            complex_plddt = (plddt * atom_pad_mask).sum(dim=-1) / atom_pad_mask.sum(dim=-1)  # (b * m,)
-            atom_to_token = feats["atom_to_token"].float().repeat_interleave(
-                multiplicity,
-                0,
-            )  # (b * m, a, t)
-            chain_id_token = feats["asym_id"].float().repeat_interleave(
-                multiplicity,
-                0,
-            )  # (b * m, t)
-            atom_type = torch.bmm(
-                atom_to_token,
-                token_type.float().unsqueeze(-1),
-            ).squeeze(-1)  # (b * m, a)
-            is_ligand_atom = (atom_type == const.chain_type_ids["NONPOLYMER"]).float()  # (b * m, a)
-            d_atom = torch.cdist(x_pred, x_pred)  # (b * m, a, a)
-            is_contact = (d_atom < 8).float()  # (b * m, a, a)
-            chain_id_atom = torch.bmm(atom_to_token, chain_id_token.unsqueeze(-1)).squeeze(-1)  # (b * m, a)
-            is_different_chain = (
-                chain_id_atom.unsqueeze(-1) != chain_id_atom.unsqueeze(-2)
-            ).float()  # (b * m, a, a)
-
-            atom_interface_mask = torch.max(
-                is_contact * is_different_chain * (1 - is_ligand_atom).unsqueeze(-1),
-                dim=-1,
-            ).values  # (b * m, a)
-            atom_non_interface_mask = (1 - atom_interface_mask) * (1 - is_ligand_atom)  # (b * m, a)
-            iplddt_weight = (
-                is_ligand_atom * ligand_weight
-                + atom_interface_mask * interface_weight
-                + atom_non_interface_mask * non_interface_weight
-            )  # (b * m, a)
-
-            complex_iplddt = (plddt * atom_pad_mask * iplddt_weight).sum(dim=-1) / torch.sum(
-                atom_pad_mask * iplddt_weight,
-                dim=-1,
-            )  # (b * m,)
-
-        # Compute the gPDE and giPDE
-        pde = compute_aggregated_metric(pde_logits, end=32)  # (b * m, t, t)
-        pred_distogram_prob = nn.functional.softmax(
-            pred_distogram_logits, dim=-1
-        ).repeat_interleave(multiplicity, 0)  # (b * m, t, t, n_bin)
-        contacts = torch.zeros((1, 1, 1, 64), dtype=pred_distogram_prob.dtype).to(
-            pred_distogram_prob.device
-        )  # (1, 1, 1, 64)
-        contacts[:, :, :, :20] = 1.0  # (1, 1, 1, 20) selected bins; contacts retains (1, 1, 1, 64)
-        prob_contact = (pred_distogram_prob * contacts).sum(-1)  # (b * m, t, t)
-        token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, t)
-        token_pad_pair_mask = (
-            token_pad_mask.unsqueeze(-1)
-            * token_pad_mask.unsqueeze(-2)
-            * (1 - torch.eye(token_pad_mask.shape[1], device=token_pad_mask.device).unsqueeze(0))
-        )  # (b * m, t, t)
-        token_pair_mask = token_pad_pair_mask * prob_contact  # (b * m, t, t)
-        complex_pde_numerator = (pde * token_pair_mask).sum(dim=(1, 2))  # (b * m,)
-        complex_pde_denominator = token_pair_mask.sum(dim=(1, 2))  # (b * m,)
-        complex_pde = complex_pde_numerator / torch.where(
-            complex_pde_denominator > 0,
-            complex_pde_denominator,
-            torch.ones_like(complex_pde_denominator),
-        )  # (b * m,)
-        asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)  # (b * m, t)
-        token_interface_pair_mask = token_pair_mask * (
-            asym_id.unsqueeze(-1) != asym_id.unsqueeze(-2)
-        )  # (b * m, t, t)
-        complex_ipde = (pde * token_interface_pair_mask).sum(dim=(1, 2)) / (
-            token_interface_pair_mask.sum(dim=(1, 2)) + 1e-5
-        )  # (b * m,)
-        out_dict = dict(
-            pde_logits=pde_logits,
-            plddt_logits=plddt_logits,
-            resolved_logits=resolved_logits,
-            pde=pde,
-            plddt=plddt,
-            complex_plddt=complex_plddt,
-            complex_iplddt=complex_iplddt,
-            complex_pde=complex_pde,
-            complex_ipde=complex_ipde,
-        )
-        out_dict["pae_logits"] = pae_logits  # (b * m, t, t, n_pae)
-        out_dict["pae"] = compute_aggregated_metric(pae_logits, end=32)  # (b * m, t, t)
-
-        ptm, iptm, ligand_iptm, protein_iptm, pair_chains_iptm = compute_ptms(
-            pae_logits, x_pred, feats, multiplicity
-        )  # four (b * m,) scores and a mapping of chain-pair scores (b * m,)
-        out_dict["ptm"] = ptm  # (b * m,)
-        out_dict["iptm"] = iptm  # (b * m,)
-        out_dict["ligand_iptm"] = ligand_iptm  # (b * m,)
-        out_dict["protein_iptm"] = protein_iptm  # (b * m,)
-        out_dict["pair_chains_iptm"] = pair_chains_iptm  # each chain-pair tensor: (b * m,)
-
-        return out_dict
+        is_ligand_token = (token_type == const.chain_type_ids["NONPOLYMER"]).float()  # (b * m, t)
+
+        if self.token_level_confidence:
+            plddt = compute_aggregated_metric(plddt_logits)  # (b * m, n_item); n_item is t for token confidence, a for atom confidence
+            token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, t)
+            complex_plddt = (plddt * token_pad_mask).sum(dim=-1) / token_pad_mask.sum(dim=-1)  # (b * m,)
+
+            is_contact = (d < 8).float()  # (b * m, t, t)
+            is_different_chain = (
+                feats["asym_id"].unsqueeze(-1) != feats["asym_id"].unsqueeze(-2)
+            ).float()  # (b, t, t)
+            is_different_chain = is_different_chain.repeat_interleave(multiplicity, 0)  # (b * m, t, t)
+            token_interface_mask = torch.max(
+                is_contact * is_different_chain * (1 - is_ligand_token).unsqueeze(-1),
+                dim=-1,
+            ).values  # (b * m, t)
+            token_non_interface_mask = (1 - token_interface_mask) * (1 - is_ligand_token)  # (b * m, t)
+            iplddt_weight = (
+                is_ligand_token * ligand_weight
+                + token_interface_mask * interface_weight
+                + token_non_interface_mask * non_interface_weight
+            )  # (b * m, t)
+            complex_iplddt = (plddt * token_pad_mask * iplddt_weight).sum(dim=-1) / torch.sum(
+                token_pad_mask * iplddt_weight, dim=-1
+            )  # (b * m,)
+
+        else:
+            # token to atom conversion for resolved logits
+            b, n, _ = resolved_logits.shape
+            resolved_logits = resolved_logits.reshape(b, n, self.max_num_atoms_per_token, 2)  # (b * m, t, n_slot, 2)
+            resolved_logits = _token_slot_logits_to_atom_logits(
+                resolved_logits,
+                feats["atom_to_token"],
+                feats["atom_pad_mask"],
+                multiplicity=multiplicity,
+            )  # (b * m, a, 2)
+            plddt_logits = plddt_logits.reshape(b, n, self.max_num_atoms_per_token, -1)  # (b * m, t, n_slot, n_plddt)
+            plddt_logits = _token_slot_logits_to_atom_logits(
+                plddt_logits,
+                feats["atom_to_token"],
+                feats["atom_pad_mask"],
+                multiplicity=multiplicity,
+            )  # (b * m, a, n_plddt)
+            atom_pad_mask = feats["atom_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, a)
+            plddt = compute_aggregated_metric(plddt_logits)  # (b * m, n_item); n_item is t for token confidence, a for atom confidence
+
+            complex_plddt = (plddt * atom_pad_mask).sum(dim=-1) / atom_pad_mask.sum(dim=-1)  # (b * m,)
+            atom_to_token = feats["atom_to_token"].float().repeat_interleave(
+                multiplicity,
+                0,
+            )  # (b * m, a, t)
+            chain_id_token = feats["asym_id"].float().repeat_interleave(
+                multiplicity,
+                0,
+            )  # (b * m, t)
+            atom_type = torch.bmm(
+                atom_to_token,
+                token_type.float().unsqueeze(-1),
+            ).squeeze(-1)  # (b * m, a)
+            is_ligand_atom = (atom_type == const.chain_type_ids["NONPOLYMER"]).float()  # (b * m, a)
+            d_atom = torch.cdist(x_pred, x_pred)  # (b * m, a, a)
+            is_contact = (d_atom < 8).float()  # (b * m, a, a)
+            chain_id_atom = torch.bmm(atom_to_token, chain_id_token.unsqueeze(-1)).squeeze(-1)  # (b * m, a)
+            is_different_chain = (
+                chain_id_atom.unsqueeze(-1) != chain_id_atom.unsqueeze(-2)
+            ).float()  # (b * m, a, a)
+
+            atom_interface_mask = torch.max(
+                is_contact * is_different_chain * (1 - is_ligand_atom).unsqueeze(-1),
+                dim=-1,
+            ).values  # (b * m, a)
+            atom_non_interface_mask = (1 - atom_interface_mask) * (1 - is_ligand_atom)  # (b * m, a)
+            iplddt_weight = (
+                is_ligand_atom * ligand_weight
+                + atom_interface_mask * interface_weight
+                + atom_non_interface_mask * non_interface_weight
+            )  # (b * m, a)
+
+            complex_iplddt = (plddt * atom_pad_mask * iplddt_weight).sum(dim=-1) / torch.sum(
+                atom_pad_mask * iplddt_weight,
+                dim=-1,
+            )  # (b * m,)
+
+        # Compute the gPDE and giPDE
+        pde = compute_aggregated_metric(pde_logits, end=32)  # (b * m, t, t)
+        pred_distogram_prob = nn.functional.softmax(
+            pred_distogram_logits, dim=-1
+        ).repeat_interleave(multiplicity, 0)  # (b * m, t, t, n_bin)
+        contacts = torch.zeros((1, 1, 1, 64), dtype=pred_distogram_prob.dtype).to(
+            pred_distogram_prob.device
+        )  # (1, 1, 1, 64)
+        contacts[:, :, :, :20] = 1.0  # (1, 1, 1, 20) selected bins; contacts retains (1, 1, 1, 64)
+        prob_contact = (pred_distogram_prob * contacts).sum(-1)  # (b * m, t, t)
+        token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, t)
+        token_pad_pair_mask = (
+            token_pad_mask.unsqueeze(-1)
+            * token_pad_mask.unsqueeze(-2)
+            * (1 - torch.eye(token_pad_mask.shape[1], device=token_pad_mask.device).unsqueeze(0))
+        )  # (b * m, t, t)
+        token_pair_mask = token_pad_pair_mask * prob_contact  # (b * m, t, t)
+        complex_pde_numerator = (pde * token_pair_mask).sum(dim=(1, 2))  # (b * m,)
+        complex_pde_denominator = token_pair_mask.sum(dim=(1, 2))  # (b * m,)
+        complex_pde = complex_pde_numerator / torch.where(
+            complex_pde_denominator > 0,
+            complex_pde_denominator,
+            torch.ones_like(complex_pde_denominator),
+        )  # (b * m,)
+        asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)  # (b * m, t)
+        token_interface_pair_mask = token_pair_mask * (
+            asym_id.unsqueeze(-1) != asym_id.unsqueeze(-2)
+        )  # (b * m, t, t)
+        complex_ipde = (pde * token_interface_pair_mask).sum(dim=(1, 2)) / (
+            token_interface_pair_mask.sum(dim=(1, 2)) + 1e-5
+        )  # (b * m,)
+        out_dict = dict(
+            pde_logits=pde_logits,
+            plddt_logits=plddt_logits,
+            resolved_logits=resolved_logits,
+            pde=pde,
+            plddt=plddt,
+            complex_plddt=complex_plddt,
+            complex_iplddt=complex_iplddt,
+            complex_pde=complex_pde,
+            complex_ipde=complex_ipde,
+        )
+        out_dict["pae_logits"] = pae_logits  # (b * m, t, t, n_pae)
+        out_dict["pae"] = compute_aggregated_metric(pae_logits, end=32)  # (b * m, t, t)
+
+        ptm, iptm, ligand_iptm, protein_iptm, pair_chains_iptm = compute_ptms(
+            pae_logits, x_pred, feats, multiplicity
+        )  # four (b * m,) scores and a mapping of chain-pair scores (b * m,)
+        out_dict["ptm"] = ptm  # (b * m,)
+        out_dict["iptm"] = iptm  # (b * m,)
+        out_dict["ligand_iptm"] = ligand_iptm  # (b * m,)
+        out_dict["protein_iptm"] = protein_iptm  # (b * m,)
+        out_dict["pair_chains_iptm"] = pair_chains_iptm  # each chain-pair tensor: (b * m,)
+
+        return out_dict

fastplms/models/boltz/vb_modules_diffusion_conditioning.py

@@ -1,147 +1,147 @@
-"""Precompute pair and atom biases used by the Boltz2 diffusion stack."""
-
-from __future__ import annotations
-
-import torch
-
-from collections.abc import Callable
-from torch import nn
-
-from .vb_modules_encodersv2 import AtomEncoder, PairwiseConditioning
-
-
-def _bias_projections(
-    depth: int,
-    input_dim: int,
-    num_heads: int,
-) -> nn.ModuleList:
-    """Build one normalized, bias-free projection per transformer block."""
-
-    return nn.ModuleList(
-        [
-            nn.Sequential(
-                nn.LayerNorm(input_dim),
-                nn.Linear(input_dim, num_heads, bias=False),
-            )
-            for _ in range(depth)
-        ]
-    )
-
-
-def _concatenate_biases(
-    projections: nn.ModuleList,
-    pair_features: torch.Tensor,
-) -> torch.Tensor:
-    # 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):
-    """Prepare conditioned atom features and per-layer attention biases."""
-
-    def __init__(
-        self,
-        token_s: int,
-        token_z: int,
-        atom_s: int,
-        atom_z: int,
-        atoms_per_window_queries: int = 32,
-        atoms_per_window_keys: int = 128,
-        atom_encoder_depth: int = 3,
-        atom_encoder_heads: int = 4,
-        token_transformer_depth: int = 24,
-        token_transformer_heads: int = 8,
-        atom_decoder_depth: int = 3,
-        atom_decoder_heads: int = 4,
-        atom_feature_dim: int = 128,
-        conditioning_transition_layers: int = 2,
-        use_no_atom_char: bool = False,
-        use_atom_backbone_feat: bool = False,
-        use_residue_feats_atoms: bool = False,
-    ) -> None:
-        super().__init__()
-        self.pairwise_conditioner = PairwiseConditioning(
-            token_z=token_z,
-            dim_token_rel_pos_feats=token_z,
-            num_transitions=conditioning_transition_layers,
-        )
-        self.atom_encoder = AtomEncoder(
-            atom_s=atom_s,
-            atom_z=atom_z,
-            token_s=token_s,
-            token_z=token_z,
-            atoms_per_window_queries=atoms_per_window_queries,
-            atoms_per_window_keys=atoms_per_window_keys,
-            atom_feature_dim=atom_feature_dim,
-            structure_prediction=True,
-            use_no_atom_char=use_no_atom_char,
-            use_atom_backbone_feat=use_atom_backbone_feat,
-            use_residue_feats_atoms=use_residue_feats_atoms,
-        )
-        self.atom_enc_proj_z = _bias_projections(
-            atom_encoder_depth,
-            atom_z,
-            atom_encoder_heads,
-        )
-        self.atom_dec_proj_z = _bias_projections(
-            atom_decoder_depth,
-            atom_z,
-            atom_decoder_heads,
-        )
-        self.token_trans_proj_z = _bias_projections(
-            token_transformer_depth,
-            token_z,
-            token_transformer_heads,
-        )
-
-    def forward(
-        self,
-        s_trunk: torch.Tensor,
-        z_trunk: torch.Tensor,
-        relative_position_encoding: torch.Tensor,
-        feats: dict[str, torch.Tensor],
-    ) -> tuple[
-        torch.Tensor,
-        torch.Tensor,
-        Callable[[torch.Tensor], torch.Tensor],
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor,
-    ]:
-        """Return conditioned atom tensors and concatenated layer biases.
-
-        ``S`` has shape ``(b, n, d_s)`` and each ``Z`` tensor has shape
-        ``(b, n, n, d_z)``.  Biases are concatenated in transformer-block
-        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,
-        )  # 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,
-            to_keys,
-            atom_encoder_bias,
-            atom_decoder_bias,
-            token_transformer_bias,
-        )  # tensor shapes are traced above; to_keys is the atom-key gatherer
+"""Precompute pair and atom biases used by the Boltz2 diffusion stack."""
+
+from __future__ import annotations
+
+import torch
+
+from collections.abc import Callable
+from torch import nn
+
+from .vb_modules_encodersv2 import AtomEncoder, PairwiseConditioning
+
+
+def _bias_projections(
+    depth: int,
+    input_dim: int,
+    num_heads: int,
+) -> nn.ModuleList:
+    """Build one normalized, bias-free projection per transformer block."""
+
+    return nn.ModuleList(
+        [
+            nn.Sequential(
+                nn.LayerNorm(input_dim),
+                nn.Linear(input_dim, num_heads, bias=False),
+            )
+            for _ in range(depth)
+        ]
+    )
+
+
+def _concatenate_biases(
+    projections: nn.ModuleList,
+    pair_features: torch.Tensor,
+) -> torch.Tensor:
+    # 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):
+    """Prepare conditioned atom features and per-layer attention biases."""
+
+    def __init__(
+        self,
+        token_s: int,
+        token_z: int,
+        atom_s: int,
+        atom_z: int,
+        atoms_per_window_queries: int = 32,
+        atoms_per_window_keys: int = 128,
+        atom_encoder_depth: int = 3,
+        atom_encoder_heads: int = 4,
+        token_transformer_depth: int = 24,
+        token_transformer_heads: int = 8,
+        atom_decoder_depth: int = 3,
+        atom_decoder_heads: int = 4,
+        atom_feature_dim: int = 128,
+        conditioning_transition_layers: int = 2,
+        use_no_atom_char: bool = False,
+        use_atom_backbone_feat: bool = False,
+        use_residue_feats_atoms: bool = False,
+    ) -> None:
+        super().__init__()
+        self.pairwise_conditioner = PairwiseConditioning(
+            token_z=token_z,
+            dim_token_rel_pos_feats=token_z,
+            num_transitions=conditioning_transition_layers,
+        )
+        self.atom_encoder = AtomEncoder(
+            atom_s=atom_s,
+            atom_z=atom_z,
+            token_s=token_s,
+            token_z=token_z,
+            atoms_per_window_queries=atoms_per_window_queries,
+            atoms_per_window_keys=atoms_per_window_keys,
+            atom_feature_dim=atom_feature_dim,
+            structure_prediction=True,
+            use_no_atom_char=use_no_atom_char,
+            use_atom_backbone_feat=use_atom_backbone_feat,
+            use_residue_feats_atoms=use_residue_feats_atoms,
+        )
+        self.atom_enc_proj_z = _bias_projections(
+            atom_encoder_depth,
+            atom_z,
+            atom_encoder_heads,
+        )
+        self.atom_dec_proj_z = _bias_projections(
+            atom_decoder_depth,
+            atom_z,
+            atom_decoder_heads,
+        )
+        self.token_trans_proj_z = _bias_projections(
+            token_transformer_depth,
+            token_z,
+            token_transformer_heads,
+        )
+
+    def forward(
+        self,
+        s_trunk: torch.Tensor,
+        z_trunk: torch.Tensor,
+        relative_position_encoding: torch.Tensor,
+        feats: dict[str, torch.Tensor],
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        Callable[[torch.Tensor], torch.Tensor],
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        """Return conditioned atom tensors and concatenated layer biases.
+
+        ``S`` has shape ``(b, n, d_s)`` and each ``Z`` tensor has shape
+        ``(b, n, n, d_z)``.  Biases are concatenated in transformer-block
+        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,
+        )  # 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,
+            to_keys,
+            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_diffusionv2.py

@@ -1,675 +1,675 @@
-"""Denoise Boltz2 atom coordinates with optional structure-steering potentials."""
-
-# Based on https://github.com/lucidrains/alphafold3-pytorch.
-# MIT License, Copyright (c) 2024 Phil Wang.
-
-from __future__ import annotations
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-from math import sqrt
-from typing import Any
-from einops import rearrange
-from torch import nn
-from torch.nn import Module
-from tqdm.auto import tqdm
-
-from . import vb_const as const
-from . import vb_layers_initialize as init
-from .vb_loss_diffusionv2 import (
-    smooth_lddt_loss,
-    weighted_rigid_align,
-)
-from .vb_modules_encodersv2 import (
-    AtomAttentionDecoder,
-    AtomAttentionEncoder,
-    SingleConditioning,
-)
-from .vb_modules_transformersv2 import (
-    DiffusionTransformer,
-)
-from .vb_modules_utils import (
-    LinearNoBias,
-    center_random_augmentation,
-    compute_random_augmentation,
-    default,
-    log,
-)
-from .vb_potentials_potentials import get_potentials
-
-
-class DiffusionModule(Module):
-    """Predict coordinate updates from token and local atom attention."""
-
-    def __init__(
-        self,
-        token_s: int,
-        atom_s: int,
-        atoms_per_window_queries: int = 32,
-        atoms_per_window_keys: int = 128,
-        sigma_data: int = 16,
-        dim_fourier: int = 256,
-        atom_encoder_depth: int = 3,
-        atom_encoder_heads: int = 4,
-        token_transformer_depth: int = 24,
-        token_transformer_heads: int = 8,
-        atom_decoder_depth: int = 3,
-        atom_decoder_heads: int = 4,
-        conditioning_transition_layers: int = 2,
-        activation_checkpointing: bool = False,
-        transformer_post_ln: bool = False,
-    ) -> None:
-        super().__init__()
-
-        self.atoms_per_window_queries = atoms_per_window_queries
-        self.atoms_per_window_keys = atoms_per_window_keys
-        self.sigma_data = sigma_data
-        self.activation_checkpointing = activation_checkpointing
-
-        # Token conditioning shares trunk features across repeated diffusion samples.
-        self.single_conditioner = SingleConditioning(
-            sigma_data=sigma_data,
-            token_s=token_s,
-            dim_fourier=dim_fourier,
-            num_transitions=conditioning_transition_layers,
-        )
-
-        self.atom_attention_encoder = AtomAttentionEncoder(
-            atom_s=atom_s,
-            token_s=token_s,
-            atoms_per_window_queries=atoms_per_window_queries,
-            atoms_per_window_keys=atoms_per_window_keys,
-            atom_encoder_depth=atom_encoder_depth,
-            atom_encoder_heads=atom_encoder_heads,
-            structure_prediction=True,
-            activation_checkpointing=activation_checkpointing,
-            transformer_post_layer_norm=transformer_post_ln,
-        )
-
-        self.s_to_a_linear = nn.Sequential(
-            nn.LayerNorm(2 * token_s), LinearNoBias(2 * token_s, 2 * token_s)
-        )
-        init.final_init_(self.s_to_a_linear[1].weight)
-
-        self.token_transformer = DiffusionTransformer(
-            dim=2 * token_s,
-            dim_single_cond=2 * token_s,
-            depth=token_transformer_depth,
-            heads=token_transformer_heads,
-            activation_checkpointing=activation_checkpointing,
-        )
-
-        self.a_norm = nn.LayerNorm(2 * token_s)
-
-        self.atom_attention_decoder = AtomAttentionDecoder(
-            atom_s=atom_s,
-            token_s=token_s,
-            attn_window_queries=atoms_per_window_queries,
-            attn_window_keys=atoms_per_window_keys,
-            atom_decoder_depth=atom_decoder_depth,
-            atom_decoder_heads=atom_decoder_heads,
-            activation_checkpointing=activation_checkpointing,
-        )
-
-    def forward(
-        self,
-        s_inputs: torch.Tensor,  # (b, t, d_s)
-        s_trunk: torch.Tensor,  # (b, t, d_s)
-        r_noisy: torch.Tensor,  # (b * m, a, 3)
-        times: torch.Tensor,  # (b * m,)
-        feats: dict[str, torch.Tensor],
-        diffusion_conditioning: dict[str, Any],
-        multiplicity: int = 1,
-    ) -> torch.Tensor:
-        # b: base batch; m: multiplicity; t: tokens; a: atoms; d_s: token width.
-        if self.activation_checkpointing:
-            s, _normed_fourier = torch.utils.checkpoint.checkpoint(
-                self.single_conditioner,
-                times,
-                s_trunk.repeat_interleave(multiplicity, 0),
-                s_inputs.repeat_interleave(multiplicity, 0),
-                use_reentrant=False,
-            )  # s: (b * m, t, 2 * d_s); Fourier features: (b * m, d_fourier)
-        else:
-            s, _normed_fourier = self.single_conditioner(
-                times,
-                s_trunk.repeat_interleave(multiplicity, 0),
-                s_inputs.repeat_interleave(multiplicity, 0),
-            )  # s: (b * m, t, 2 * d_s); Fourier features: (b * m, d_fourier)
-
-        # Sequence-local Atom Attention and aggregation to coarse-grained tokens
-        a, q_skip, c_skip, to_keys = self.atom_attention_encoder(
-            feats=feats,
-            q=diffusion_conditioning["q"].float(),
-            c=diffusion_conditioning["c"].float(),
-            atom_enc_bias=diffusion_conditioning["atom_enc_bias"].float(),
-            to_keys=diffusion_conditioning["to_keys"],
-            r=r_noisy,
-            multiplicity=multiplicity,
-        )  # a: (b * m, t, 2 * d_s); q_skip/c_skip: (b * m, n_atom, d_a)
-
-        # Full self-attention on token level
-        a = a + self.s_to_a_linear(s)  # (b * m, t, 2 * d_s)
-
-        mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, t)
-        a = self.token_transformer(
-            a,
-            mask=mask.float(),
-            s=s,
-            bias=diffusion_conditioning[
-                "token_trans_bias"
-            ].float(),  # note z is not expanded with multiplicity until after bias is computed
-            multiplicity=multiplicity,
-        )  # (b * m, t, 2 * d_s)
-        a = self.a_norm(a)  # (b * m, t, 2 * d_s)
-
-        # Broadcast token activations to atoms and run Sequence-local Atom Attention
-        r_update = self.atom_attention_decoder(
-            a=a,
-            q=q_skip,
-            c=c_skip,
-            atom_dec_bias=diffusion_conditioning["atom_dec_bias"].float(),
-            feats=feats,
-            multiplicity=multiplicity,
-            to_keys=to_keys,
-        )  # (b * m, n_atom, 3)
-
-        return r_update  # (b * m, n_atom, 3)
-
-
-class AtomDiffusion(Module):
-    def __init__(
-        self,
-        score_model_args: dict[str, Any],
-        num_sampling_steps: int = 5,  # number of sampling steps
-        sigma_min: float = 0.0004,  # min noise level
-        sigma_max: float = 160.0,  # max noise level
-        sigma_data: float = 16.0,  # standard deviation of data distribution
-        rho: float = 7,  # controls the sampling schedule
-        # Mean of the log-normal training-noise distribution.
-        P_mean: float = -1.2,
-        # Standard deviation of the log-normal training-noise distribution.
-        P_std: float = 1.5,
-        gamma_0: float = 0.8,
-        gamma_min: float = 1.0,
-        noise_scale: float = 1.003,
-        step_scale: float = 1.5,
-        step_scale_random: list[float] | None = None,
-        coordinate_augmentation: bool = True,
-        coordinate_augmentation_inference: bool | None = None,
-        compile_score: bool = False,
-        alignment_reverse_diff: bool = False,
-        synchronize_sigmas: bool = False,
-    ) -> None:
-        super().__init__()
-        self.score_model = DiffusionModule(
-            **score_model_args,
-        )
-        if compile_score:
-            self.score_model = torch.compile(self.score_model, dynamic=False, fullgraph=False)
-
-        self.sigma_min = sigma_min
-        self.sigma_max = sigma_max
-        self.sigma_data = sigma_data
-        self.rho = rho
-        self.P_mean = P_mean
-        self.P_std = P_std
-        self.num_sampling_steps = num_sampling_steps
-        self.gamma_0 = gamma_0
-        self.gamma_min = gamma_min
-        self.noise_scale = noise_scale
-        self.step_scale = step_scale
-        self.step_scale_random = step_scale_random
-        self.coordinate_augmentation = coordinate_augmentation
-        self.coordinate_augmentation_inference = (
-            coordinate_augmentation_inference
-            if coordinate_augmentation_inference is not None
-            else coordinate_augmentation
-        )
-        self.alignment_reverse_diff = alignment_reverse_diff
-        self.synchronize_sigmas = synchronize_sigmas
-
-        self.token_s = score_model_args["token_s"]
-        self.register_buffer("zero", torch.tensor(0.0), persistent=False)
-
-    @property
-    def device(self) -> torch.device:
-        return next(self.score_model.parameters()).device
-
-    def c_skip(self, sigma: torch.Tensor) -> torch.Tensor:
-        # sigma: (...); coefficient preserves shape (...).
-        return (self.sigma_data**2) / (sigma**2 + self.sigma_data**2)
-
-    def c_out(self, sigma: torch.Tensor) -> torch.Tensor:
-        # sigma: (...); coefficient preserves shape (...).
-        return sigma * self.sigma_data / torch.sqrt(self.sigma_data**2 + sigma**2)
-
-    def c_in(self, sigma: torch.Tensor) -> torch.Tensor:
-        # sigma: (...); coefficient preserves shape (...).
-        return 1 / torch.sqrt(sigma**2 + self.sigma_data**2)
-
-    def c_noise(self, sigma: torch.Tensor) -> torch.Tensor:
-        # sigma: (...); log-noise condition preserves shape (...).
-        return log(sigma / self.sigma_data) * 0.25
-
-    def preconditioned_network_forward(
-        self,
-        noised_atom_coords: torch.Tensor,  #: Float['b m 3'],
-        sigma: torch.Tensor | float,  #: Float['b'] | Float[' '] | float,
-        network_condition_kwargs: dict[str, Any],
-    ) -> torch.Tensor:
-        # noised_atom_coords: (s, a, 3); sigma: float or (s,).
-        batch, device = noised_atom_coords.shape[0], noised_atom_coords.device
-
-        if isinstance(sigma, float):
-            sigma = torch.full((batch,), sigma, device=device)  # (s,) after a scalar input is expanded
-
-        padded_sigma = rearrange(sigma, "b -> b 1 1")  # (s, 1, 1)
-
-        r_update = self.score_model(
-            r_noisy=self.c_in(padded_sigma) * noised_atom_coords,
-            times=self.c_noise(sigma),
-            **network_condition_kwargs,
-        )  # (s, a, 3)
-
-        denoised_coords = (
-            self.c_skip(padded_sigma) * noised_atom_coords + self.c_out(padded_sigma) * r_update
-        )  # (s, a, 3)
-        return denoised_coords
-
-    def sample_schedule(self, num_sampling_steps: int | None = None) -> torch.Tensor:
-        # n_step is the requested sampling-step count; the returned schedule has n_step + 1 values.
-        num_sampling_steps = default(num_sampling_steps, self.num_sampling_steps)
-        inv_rho = 1 / self.rho
-
-        steps = torch.arange(num_sampling_steps, device=self.device, dtype=torch.float32)  # (n_step,)
-        sigmas = (
-            self.sigma_max**inv_rho
-            + steps / (num_sampling_steps - 1) * (self.sigma_min**inv_rho - self.sigma_max**inv_rho)
-        ) ** self.rho  # (n_sigma,); schedule or batch axis is determined by the caller
-
-        sigmas = sigmas * self.sigma_data  # (n_sigma,); schedule or batch axis is determined by the caller
-
-        sigmas = F.pad(sigmas, (0, 1), value=0.0)  # last step is sigma value of 0.
-        return sigmas
-
-    def sample(
-        self,
-        atom_mask: torch.Tensor,
-        num_sampling_steps: int | None = None,
-        multiplicity: int = 1,
-        max_parallel_samples: int | None = None,
-        steering_args: dict[str, Any] | None = None,
-        verbose: bool = False,
-        **network_condition_kwargs: Any,
-    ) -> dict[str, Any]:
-        # a is atom count; m is multiplicity, including steering particles.
-        # s is the current sample axis after batch repetition or particle resampling.
-        if steering_args is not None and (
-            steering_args["fk_steering"]
-            or steering_args["physical_guidance_update"]
-            or steering_args["contact_guidance_update"]
-        ):
-            potentials = get_potentials(steering_args, boltz2=True)
-
-        if steering_args["fk_steering"]:
-            multiplicity = multiplicity * steering_args["num_particles"]
-            energy_traj = torch.empty((multiplicity, 0), device=self.device)  # (m, n_evaluation)
-            resample_weights = torch.ones(multiplicity, device=self.device).reshape(
-                -1, steering_args["num_particles"]
-            )  # (n_group, n_particle)
-        if steering_args["physical_guidance_update"] or steering_args["contact_guidance_update"]:
-            scaled_guidance_update = torch.zeros(
-                (multiplicity, *atom_mask.shape[1:], 3),
-                dtype=torch.float32,
-                device=self.device,
-            )  # (s, a, 3)
-        if max_parallel_samples is None:
-            max_parallel_samples = multiplicity
-
-        num_sampling_steps = default(num_sampling_steps, self.num_sampling_steps)
-        atom_mask = atom_mask.repeat_interleave(multiplicity, 0)  # (s, a); s expands on repetition and can shrink on resampling
-
-        shape = (*atom_mask.shape, 3)
-
-        # Pair each schedule value with the next sigma and gamma.
-        sigmas = self.sample_schedule(num_sampling_steps)  # (n_sigma,); schedule or batch axis is determined by the caller
-        gammas = torch.where(sigmas > self.gamma_min, self.gamma_0, 0.0)  # (n_step + 1,)
-        sigmas_and_gammas = list(zip(sigmas[:-1], sigmas[1:], gammas[1:], strict=True))
-        if self.training and self.step_scale_random is not None:
-            step_scale = np.random.choice(self.step_scale_random)
-        else:
-            step_scale = self.step_scale
-
-        # atom position is noise at the beginning
-        init_sigma = sigmas[0]  # ()
-        atom_coords = init_sigma * torch.randn(shape, device=self.device)  # (s, a, 3)
-        token_repr = None  # None; the current sampler does not produce token features
-        atom_coords_denoised = None  # (s, a, 3) or None before the first denoising step
-
-        for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(
-            tqdm(
-                sigmas_and_gammas,
-                desc="Boltz2: Diffusion sampling",
-                unit="step",
-                disable=not verbose,
-                dynamic_ncols=True,
-                leave=False,
-            )
-        ):
-            random_R, random_tr = compute_random_augmentation(
-                multiplicity, device=atom_coords.device, dtype=atom_coords.dtype
-            )  # (m, 3, 3), (m, 1, 3)
-            atom_coords = atom_coords - atom_coords.mean(dim=-2, keepdims=True)  # (s, a, 3)
-            atom_coords = torch.einsum("bmd,bds->bms", atom_coords, random_R) + random_tr  # (s, a, 3)
-            if atom_coords_denoised is not None:
-                atom_coords_denoised -= atom_coords_denoised.mean(dim=-2, keepdims=True)  # (s, a, 3) or None before the first denoising step
-                atom_coords_denoised = (
-                    torch.einsum("bmd,bds->bms", atom_coords_denoised, random_R) + random_tr
-                )  # (s, a, 3) or None before the first denoising step
-            if (
-                steering_args["physical_guidance_update"]
-                or steering_args["contact_guidance_update"]
-            ) and scaled_guidance_update is not None:
-                scaled_guidance_update = torch.einsum(
-                    "bmd,bds->bms", scaled_guidance_update, random_R
-                )  # (s, a, 3)
-
-            sigma_tm, sigma_t, gamma = sigma_tm.item(), sigma_t.item(), gamma.item()  # three scalar tensors () become Python floats
-
-            t_hat = sigma_tm * (1 + gamma)
-            steering_t = 1.0 - (step_idx / num_sampling_steps)
-            noise_var = self.noise_scale**2 * (t_hat**2 - sigma_tm**2)
-            eps = sqrt(noise_var) * torch.randn(shape, device=self.device)  # (s, a, 3)
-            atom_coords_noisy = atom_coords + eps  # (s, a, 3)
-
-            with torch.no_grad():
-                atom_coords_denoised = torch.zeros_like(atom_coords_noisy)  # (s, a, 3) or None before the first denoising step
-                sample_ids = torch.arange(multiplicity).to(atom_coords_noisy.device)  # (m,)
-                sample_ids_chunks = sample_ids.chunk(multiplicity % max_parallel_samples + 1)  # tuple of index tensors, each (n_chunk,)
-
-                for sample_ids_chunk in sample_ids_chunks:
-                    atom_coords_denoised_chunk = self.preconditioned_network_forward(
-                        atom_coords_noisy[sample_ids_chunk],
-                        t_hat,
-                        network_condition_kwargs=dict(
-                            multiplicity=sample_ids_chunk.numel(),
-                            **network_condition_kwargs,
-                        ),
-                    )  # (n_chunk, a, 3)
-                    atom_coords_denoised[sample_ids_chunk] = atom_coords_denoised_chunk  # (n_chunk, a, 3)
-
-                if steering_args["fk_steering"] and (
-                    (step_idx % steering_args["fk_resampling_interval"] == 0 and noise_var > 0)
-                    or step_idx == num_sampling_steps - 1
-                ):
-                    # Compute energy of x_0 prediction
-                    energy = torch.zeros(multiplicity, device=self.device)  # (m,)
-                    for potential in potentials:
-                        parameters = potential.compute_parameters(steering_t)
-                        if parameters["resampling_weight"] > 0:
-                            component_energy = potential.compute(
-                                atom_coords_denoised,
-                                network_condition_kwargs["feats"],
-                                parameters,
-                            )  # (m,)
-                            energy += parameters["resampling_weight"] * component_energy  # (m,)
-                    energy_traj = torch.cat((energy_traj, energy.unsqueeze(1)), dim=1)  # (m, n_evaluation)
-
-                    # Compute log G values
-                    if step_idx == 0:
-                        log_G = -1 * energy  # (m,)
-                    else:
-                        log_G = energy_traj[:, -2] - energy_traj[:, -1]  # (m,)
-
-                    # Compute ll difference between guided and unguided transition distribution
-                    if (
-                        steering_args["physical_guidance_update"]
-                        or steering_args["contact_guidance_update"]
-                    ) and noise_var > 0:
-                        ll_difference = (eps**2 - (eps + scaled_guidance_update) ** 2).sum(
-                            dim=(-1, -2)
-                        ) / (2 * noise_var)  # (m,)
-                    else:
-                        ll_difference = torch.zeros_like(energy)  # (m,)
-
-                    # Compute resampling weights
-                    resample_weights = F.softmax(
-                        (ll_difference + steering_args["fk_lambda"] * log_G).reshape(
-                            -1, steering_args["num_particles"]
-                        ),
-                        dim=1,
-                    )  # (n_group, n_particle)
-
-                # Compute guidance update to x_0 prediction
-                if (
-                    steering_args["physical_guidance_update"]
-                    or steering_args["contact_guidance_update"]
-                ) and step_idx < num_sampling_steps - 1:
-                    guidance_update = torch.zeros_like(atom_coords_denoised)  # (s, a, 3)
-                    for guidance_step in range(steering_args["num_gd_steps"]):
-                        energy_gradient = torch.zeros_like(atom_coords_denoised)  # (s, a, 3)
-                        for potential in potentials:
-                            parameters = potential.compute_parameters(steering_t)
-                            if (
-                                parameters["guidance_weight"] > 0
-                                and (guidance_step) % parameters["guidance_interval"] == 0
-                            ):
-                                energy_gradient += parameters[
-                                    "guidance_weight"
-                                ] * potential.compute_gradient(
-                                    atom_coords_denoised + guidance_update,
-                                    network_condition_kwargs["feats"],
-                                    parameters,
-                                )  # (s, a, 3)
-                        guidance_update -= energy_gradient  # (s, a, 3)
-                    atom_coords_denoised += guidance_update  # (s, a, 3) or None before the first denoising step
-                    scaled_guidance_update = (
-                        guidance_update * -1 * self.step_scale * (sigma_t - t_hat) / t_hat
-                    )  # (s, a, 3)
-
-                if steering_args["fk_steering"] and (
-                    (step_idx % steering_args["fk_resampling_interval"] == 0 and noise_var > 0)
-                    or step_idx == num_sampling_steps - 1
-                ):
-                    resample_indices = (
-                        torch.multinomial(
-                            resample_weights,
-                            resample_weights.shape[1] if step_idx < num_sampling_steps - 1 else 1,
-                            replacement=True,
-                        )
-                        + resample_weights.shape[1]
-                        * torch.arange(
-                            resample_weights.shape[0], device=resample_weights.device
-                        ).unsqueeze(-1)
-                    ).flatten()  # (s_next,); final resampling keeps one particle per group
-
-                    atom_coords = atom_coords[resample_indices]  # (s, a, 3)
-                    atom_coords_noisy = atom_coords_noisy[resample_indices]  # (s, a, 3)
-                    atom_mask = atom_mask[resample_indices]  # (s, a); s expands on repetition and can shrink on resampling
-                    if atom_coords_denoised is not None:
-                        atom_coords_denoised = atom_coords_denoised[resample_indices]  # (s, a, 3) or None before the first denoising step
-                    energy_traj = energy_traj[resample_indices]  # (s_next, n_evaluation)
-                    if (
-                        steering_args["physical_guidance_update"]
-                        or steering_args["contact_guidance_update"]
-                    ):
-                        scaled_guidance_update = scaled_guidance_update[resample_indices]  # (s, a, 3)
-                    if token_repr is not None:
-                        token_repr = token_repr[resample_indices]  # (s_next, *feature_shape)
-
-            if self.alignment_reverse_diff:
-                with torch.autocast("cuda", enabled=False):
-                    atom_coords_noisy = weighted_rigid_align(
-                        atom_coords_noisy.float(),
-                        atom_coords_denoised.float(),
-                        atom_mask.float(),
-                        atom_mask.float(),
-                    )  # (s, a, 3)
-
-                atom_coords_noisy = atom_coords_noisy.to(atom_coords_denoised)  # (s, a, 3)
-
-            denoised_over_sigma = (atom_coords_noisy - atom_coords_denoised) / t_hat  # (s, a, 3)
-            atom_coords_next = (
-                atom_coords_noisy + step_scale * (sigma_t - t_hat) * denoised_over_sigma
-            )  # (s, a, 3)
-
-            atom_coords = atom_coords_next  # (s, a, 3)
-
-        return dict(sample_atom_coords=atom_coords, diff_token_repr=token_repr)
-
-    def loss_weight(self, sigma: torch.Tensor) -> torch.Tensor:
-        # sigma: (...); weight preserves shape (...).
-        return (sigma**2 + self.sigma_data**2) / ((sigma * self.sigma_data) ** 2)
-
-    def noise_distribution(self, batch_size: int) -> torch.Tensor:
-        # Return one noise scale per example, shape (batch_size,).
-        return (
-            self.sigma_data
-            * (self.P_mean + self.P_std * torch.randn((batch_size,), device=self.device)).exp()
-        )
-
-    def forward(
-        self,
-        s_inputs: torch.Tensor,
-        s_trunk: torch.Tensor,
-        feats: dict[str, torch.Tensor],
-        diffusion_conditioning: dict[str, Any],
-        multiplicity: int = 1,
-    ) -> dict[str, Any]:
-        # training diffusion step
-        batch_size = feats["coords"].shape[0] // multiplicity
-
-        if self.synchronize_sigmas:
-            sigmas = self.noise_distribution(batch_size).repeat_interleave(multiplicity, 0)  # (n_sigma,); schedule or batch axis is determined by the caller
-        else:
-            sigmas = self.noise_distribution(batch_size * multiplicity)  # (n_sigma,); schedule or batch axis is determined by the caller
-        padded_sigmas = rearrange(sigmas, "b -> b 1 1")  # (b * m, 1, 1)
-
-        atom_coords = feats["coords"]  # (s, a, 3)
-
-        atom_mask = feats["atom_pad_mask"]  # (s, a); s expands on repetition and can shrink on resampling
-        atom_mask = atom_mask.repeat_interleave(multiplicity, 0)  # (s, a); s expands on repetition and can shrink on resampling
-
-        atom_coords = center_random_augmentation(
-            atom_coords, atom_mask, augmentation=self.coordinate_augmentation
-        )  # (s, a, 3)
-
-        noise = torch.randn_like(atom_coords)  # (b * m, a, 3)
-        noised_atom_coords = atom_coords + padded_sigmas * noise  # (b * m, a, 3)
-
-        denoised_atom_coords = self.preconditioned_network_forward(
-            noised_atom_coords,
-            sigmas,
-            network_condition_kwargs={
-                "s_inputs": s_inputs,
-                "s_trunk": s_trunk,
-                "feats": feats,
-                "multiplicity": multiplicity,
-                "diffusion_conditioning": diffusion_conditioning,
-            },
-        )  # (b * m, a, 3)
-
-        return {
-            "denoised_atom_coords": denoised_atom_coords,
-            "sigmas": sigmas,
-            "aligned_true_atom_coords": atom_coords,
-        }
-
-    def compute_loss(
-        self,
-        feats: dict[str, torch.Tensor],
-        out_dict: dict[str, torch.Tensor],
-        add_smooth_lddt_loss: bool = True,
-        nucleotide_loss_weight: float = 5.0,
-        ligand_loss_weight: float = 10.0,
-        multiplicity: int = 1,
-        filter_by_plddt: float = 0.0,
-    ) -> dict[str, Any]:
-        # Coordinates: (b * m, a, 3); base feature masks: (b, a).
-        with torch.autocast("cuda", enabled=False):
-            denoised_atom_coords = out_dict["denoised_atom_coords"].float()  # (b * m, a, 3)
-            sigmas = out_dict["sigmas"].float()  # (n_sigma,); schedule or batch axis is determined by the caller
-
-            resolved_atom_mask_uni = feats["atom_resolved_mask"].float()  # (b, a)
-
-            if filter_by_plddt > 0:
-                plddt_mask = feats["plddt"] > filter_by_plddt  # (b, a)
-                resolved_atom_mask_uni = resolved_atom_mask_uni * plddt_mask.float()  # (b, a)
-
-            resolved_atom_mask = resolved_atom_mask_uni.repeat_interleave(multiplicity, 0)  # (b * m, a)
-
-            align_weights = denoised_atom_coords.new_ones(denoised_atom_coords.shape[:2])  # (b * m, a)
-            atom_type = (
-                torch.bmm(
-                    feats["atom_to_token"].float(),
-                    feats["mol_type"].unsqueeze(-1).float(),
-                )
-                .squeeze(-1)
-                .long()
-            )  # (b, a)
-            atom_type_mult = atom_type.repeat_interleave(multiplicity, 0)  # (b * m, a)
-
-            align_weights = (
-                align_weights
-                * (
-                    1
-                    + nucleotide_loss_weight
-                    * (
-                        torch.eq(atom_type_mult, const.chain_type_ids["DNA"]).float()
-                        + torch.eq(atom_type_mult, const.chain_type_ids["RNA"]).float()
-                    )
-                    + ligand_loss_weight
-                    * torch.eq(atom_type_mult, const.chain_type_ids["NONPOLYMER"]).float()
-                ).float()
-            )  # (b * m, a)
-
-            atom_coords = out_dict["aligned_true_atom_coords"].float()  # (s, a, 3)
-            atom_coords_aligned_ground_truth = weighted_rigid_align(
-                atom_coords.detach(),
-                denoised_atom_coords.detach(),
-                align_weights.detach(),
-                mask=feats["atom_resolved_mask"]
-                .float()
-                .repeat_interleave(multiplicity, 0)
-                .detach(),
-            )  # (b * m, a, 3)
-
-            atom_coords_aligned_ground_truth = atom_coords_aligned_ground_truth.to(
-                denoised_atom_coords
-            )  # (b * m, a, 3)
-
-            # weighted MSE loss of denoised atom positions
-            mse_loss = ((denoised_atom_coords - atom_coords_aligned_ground_truth) ** 2).sum(dim=-1)  # (b * m, a)
-            mse_loss = torch.sum(mse_loss * align_weights * resolved_atom_mask, dim=-1) / (
-                torch.sum(3 * align_weights * resolved_atom_mask, dim=-1) + 1e-5
-            )  # (b * m,)
-
-            # weight by sigma factor
-            loss_weights = self.loss_weight(sigmas)  # (b * m,)
-            mse_loss = (mse_loss * loss_weights).mean()  # ()
-
-            total_loss = mse_loss  # ()
-
-            # proposed auxiliary smooth lddt loss
-            lddt_loss = self.zero  # ()
-            if add_smooth_lddt_loss:
-                lddt_loss = smooth_lddt_loss(
-                    denoised_atom_coords,
-                    feats["coords"],
-                    torch.eq(atom_type, const.chain_type_ids["DNA"]).float()
-                    + torch.eq(atom_type, const.chain_type_ids["RNA"]).float(),
-                    coords_mask=resolved_atom_mask_uni,
-                    multiplicity=multiplicity,
-                )  # ()
-
-                total_loss = total_loss + lddt_loss  # ()
-
-            loss_breakdown = {
-                "mse_loss": mse_loss,
-                "smooth_lddt_loss": lddt_loss,
-            }
-
-        return {"loss": total_loss, "loss_breakdown": loss_breakdown}
+"""Denoise Boltz2 atom coordinates with optional structure-steering potentials."""
+
+# Based on https://github.com/lucidrains/alphafold3-pytorch.
+# MIT License, Copyright (c) 2024 Phil Wang.
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from math import sqrt
+from typing import Any
+from einops import rearrange
+from torch import nn
+from torch.nn import Module
+from tqdm.auto import tqdm
+
+from . import vb_const as const
+from . import vb_layers_initialize as init
+from .vb_loss_diffusionv2 import (
+    smooth_lddt_loss,
+    weighted_rigid_align,
+)
+from .vb_modules_encodersv2 import (
+    AtomAttentionDecoder,
+    AtomAttentionEncoder,
+    SingleConditioning,
+)
+from .vb_modules_transformersv2 import (
+    DiffusionTransformer,
+)
+from .vb_modules_utils import (
+    LinearNoBias,
+    center_random_augmentation,
+    compute_random_augmentation,
+    default,
+    log,
+)
+from .vb_potentials_potentials import get_potentials
+
+
+class DiffusionModule(Module):
+    """Predict coordinate updates from token and local atom attention."""
+
+    def __init__(
+        self,
+        token_s: int,
+        atom_s: int,
+        atoms_per_window_queries: int = 32,
+        atoms_per_window_keys: int = 128,
+        sigma_data: int = 16,
+        dim_fourier: int = 256,
+        atom_encoder_depth: int = 3,
+        atom_encoder_heads: int = 4,
+        token_transformer_depth: int = 24,
+        token_transformer_heads: int = 8,
+        atom_decoder_depth: int = 3,
+        atom_decoder_heads: int = 4,
+        conditioning_transition_layers: int = 2,
+        activation_checkpointing: bool = False,
+        transformer_post_ln: bool = False,
+    ) -> None:
+        super().__init__()
+
+        self.atoms_per_window_queries = atoms_per_window_queries
+        self.atoms_per_window_keys = atoms_per_window_keys
+        self.sigma_data = sigma_data
+        self.activation_checkpointing = activation_checkpointing
+
+        # Token conditioning shares trunk features across repeated diffusion samples.
+        self.single_conditioner = SingleConditioning(
+            sigma_data=sigma_data,
+            token_s=token_s,
+            dim_fourier=dim_fourier,
+            num_transitions=conditioning_transition_layers,
+        )
+
+        self.atom_attention_encoder = AtomAttentionEncoder(
+            atom_s=atom_s,
+            token_s=token_s,
+            atoms_per_window_queries=atoms_per_window_queries,
+            atoms_per_window_keys=atoms_per_window_keys,
+            atom_encoder_depth=atom_encoder_depth,
+            atom_encoder_heads=atom_encoder_heads,
+            structure_prediction=True,
+            activation_checkpointing=activation_checkpointing,
+            transformer_post_layer_norm=transformer_post_ln,
+        )
+
+        self.s_to_a_linear = nn.Sequential(
+            nn.LayerNorm(2 * token_s), LinearNoBias(2 * token_s, 2 * token_s)
+        )
+        init.final_init_(self.s_to_a_linear[1].weight)
+
+        self.token_transformer = DiffusionTransformer(
+            dim=2 * token_s,
+            dim_single_cond=2 * token_s,
+            depth=token_transformer_depth,
+            heads=token_transformer_heads,
+            activation_checkpointing=activation_checkpointing,
+        )
+
+        self.a_norm = nn.LayerNorm(2 * token_s)
+
+        self.atom_attention_decoder = AtomAttentionDecoder(
+            atom_s=atom_s,
+            token_s=token_s,
+            attn_window_queries=atoms_per_window_queries,
+            attn_window_keys=atoms_per_window_keys,
+            atom_decoder_depth=atom_decoder_depth,
+            atom_decoder_heads=atom_decoder_heads,
+            activation_checkpointing=activation_checkpointing,
+        )
+
+    def forward(
+        self,
+        s_inputs: torch.Tensor,  # (b, t, d_s)
+        s_trunk: torch.Tensor,  # (b, t, d_s)
+        r_noisy: torch.Tensor,  # (b * m, a, 3)
+        times: torch.Tensor,  # (b * m,)
+        feats: dict[str, torch.Tensor],
+        diffusion_conditioning: dict[str, Any],
+        multiplicity: int = 1,
+    ) -> torch.Tensor:
+        # b: base batch; m: multiplicity; t: tokens; a: atoms; d_s: token width.
+        if self.activation_checkpointing:
+            s, _normed_fourier = torch.utils.checkpoint.checkpoint(
+                self.single_conditioner,
+                times,
+                s_trunk.repeat_interleave(multiplicity, 0),
+                s_inputs.repeat_interleave(multiplicity, 0),
+                use_reentrant=False,
+            )  # s: (b * m, t, 2 * d_s); Fourier features: (b * m, d_fourier)
+        else:
+            s, _normed_fourier = self.single_conditioner(
+                times,
+                s_trunk.repeat_interleave(multiplicity, 0),
+                s_inputs.repeat_interleave(multiplicity, 0),
+            )  # s: (b * m, t, 2 * d_s); Fourier features: (b * m, d_fourier)
+
+        # Sequence-local Atom Attention and aggregation to coarse-grained tokens
+        a, q_skip, c_skip, to_keys = self.atom_attention_encoder(
+            feats=feats,
+            q=diffusion_conditioning["q"].float(),
+            c=diffusion_conditioning["c"].float(),
+            atom_enc_bias=diffusion_conditioning["atom_enc_bias"].float(),
+            to_keys=diffusion_conditioning["to_keys"],
+            r=r_noisy,
+            multiplicity=multiplicity,
+        )  # a: (b * m, t, 2 * d_s); q_skip/c_skip: (b * m, n_atom, d_a)
+
+        # Full self-attention on token level
+        a = a + self.s_to_a_linear(s)  # (b * m, t, 2 * d_s)
+
+        mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)  # (b * m, t)
+        a = self.token_transformer(
+            a,
+            mask=mask.float(),
+            s=s,
+            bias=diffusion_conditioning[
+                "token_trans_bias"
+            ].float(),  # note z is not expanded with multiplicity until after bias is computed
+            multiplicity=multiplicity,
+        )  # (b * m, t, 2 * d_s)
+        a = self.a_norm(a)  # (b * m, t, 2 * d_s)
+
+        # Broadcast token activations to atoms and run Sequence-local Atom Attention
+        r_update = self.atom_attention_decoder(
+            a=a,
+            q=q_skip,
+            c=c_skip,
+            atom_dec_bias=diffusion_conditioning["atom_dec_bias"].float(),
+            feats=feats,
+            multiplicity=multiplicity,
+            to_keys=to_keys,
+        )  # (b * m, n_atom, 3)
+
+        return r_update  # (b * m, n_atom, 3)
+
+
+class AtomDiffusion(Module):
+    def __init__(
+        self,
+        score_model_args: dict[str, Any],
+        num_sampling_steps: int = 5,  # number of sampling steps
+        sigma_min: float = 0.0004,  # min noise level
+        sigma_max: float = 160.0,  # max noise level
+        sigma_data: float = 16.0,  # standard deviation of data distribution
+        rho: float = 7,  # controls the sampling schedule
+        # Mean of the log-normal training-noise distribution.
+        P_mean: float = -1.2,
+        # Standard deviation of the log-normal training-noise distribution.
+        P_std: float = 1.5,
+        gamma_0: float = 0.8,
+        gamma_min: float = 1.0,
+        noise_scale: float = 1.003,
+        step_scale: float = 1.5,
+        step_scale_random: list[float] | None = None,
+        coordinate_augmentation: bool = True,
+        coordinate_augmentation_inference: bool | None = None,
+        compile_score: bool = False,
+        alignment_reverse_diff: bool = False,
+        synchronize_sigmas: bool = False,
+    ) -> None:
+        super().__init__()
+        self.score_model = DiffusionModule(
+            **score_model_args,
+        )
+        if compile_score:
+            self.score_model = torch.compile(self.score_model, dynamic=False, fullgraph=False)
+
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.sigma_data = sigma_data
+        self.rho = rho
+        self.P_mean = P_mean
+        self.P_std = P_std
+        self.num_sampling_steps = num_sampling_steps
+        self.gamma_0 = gamma_0
+        self.gamma_min = gamma_min
+        self.noise_scale = noise_scale
+        self.step_scale = step_scale
+        self.step_scale_random = step_scale_random
+        self.coordinate_augmentation = coordinate_augmentation
+        self.coordinate_augmentation_inference = (
+            coordinate_augmentation_inference
+            if coordinate_augmentation_inference is not None
+            else coordinate_augmentation
+        )
+        self.alignment_reverse_diff = alignment_reverse_diff
+        self.synchronize_sigmas = synchronize_sigmas
+
+        self.token_s = score_model_args["token_s"]
+        self.register_buffer("zero", torch.tensor(0.0), persistent=False)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.score_model.parameters()).device
+
+    def c_skip(self, sigma: torch.Tensor) -> torch.Tensor:
+        # sigma: (...); coefficient preserves shape (...).
+        return (self.sigma_data**2) / (sigma**2 + self.sigma_data**2)
+
+    def c_out(self, sigma: torch.Tensor) -> torch.Tensor:
+        # sigma: (...); coefficient preserves shape (...).
+        return sigma * self.sigma_data / torch.sqrt(self.sigma_data**2 + sigma**2)
+
+    def c_in(self, sigma: torch.Tensor) -> torch.Tensor:
+        # sigma: (...); coefficient preserves shape (...).
+        return 1 / torch.sqrt(sigma**2 + self.sigma_data**2)
+
+    def c_noise(self, sigma: torch.Tensor) -> torch.Tensor:
+        # sigma: (...); log-noise condition preserves shape (...).
+        return log(sigma / self.sigma_data) * 0.25
+
+    def preconditioned_network_forward(
+        self,
+        noised_atom_coords: torch.Tensor,  #: Float['b m 3'],
+        sigma: torch.Tensor | float,  #: Float['b'] | Float[' '] | float,
+        network_condition_kwargs: dict[str, Any],
+    ) -> torch.Tensor:
+        # noised_atom_coords: (s, a, 3); sigma: float or (s,).
+        batch, device = noised_atom_coords.shape[0], noised_atom_coords.device
+
+        if isinstance(sigma, float):
+            sigma = torch.full((batch,), sigma, device=device)  # (s,) after a scalar input is expanded
+
+        padded_sigma = rearrange(sigma, "b -> b 1 1")  # (s, 1, 1)
+
+        r_update = self.score_model(
+            r_noisy=self.c_in(padded_sigma) * noised_atom_coords,
+            times=self.c_noise(sigma),
+            **network_condition_kwargs,
+        )  # (s, a, 3)
+
+        denoised_coords = (
+            self.c_skip(padded_sigma) * noised_atom_coords + self.c_out(padded_sigma) * r_update
+        )  # (s, a, 3)
+        return denoised_coords
+
+    def sample_schedule(self, num_sampling_steps: int | None = None) -> torch.Tensor:
+        # n_step is the requested sampling-step count; the returned schedule has n_step + 1 values.
+        num_sampling_steps = default(num_sampling_steps, self.num_sampling_steps)
+        inv_rho = 1 / self.rho
+
+        steps = torch.arange(num_sampling_steps, device=self.device, dtype=torch.float32)  # (n_step,)
+        sigmas = (
+            self.sigma_max**inv_rho
+            + steps / (num_sampling_steps - 1) * (self.sigma_min**inv_rho - self.sigma_max**inv_rho)
+        ) ** self.rho  # (n_sigma,); schedule or batch axis is determined by the caller
+
+        sigmas = sigmas * self.sigma_data  # (n_sigma,); schedule or batch axis is determined by the caller
+
+        sigmas = F.pad(sigmas, (0, 1), value=0.0)  # last step is sigma value of 0.
+        return sigmas
+
+    def sample(
+        self,
+        atom_mask: torch.Tensor,
+        num_sampling_steps: int | None = None,
+        multiplicity: int = 1,
+        max_parallel_samples: int | None = None,
+        steering_args: dict[str, Any] | None = None,
+        verbose: bool = False,
+        **network_condition_kwargs: Any,
+    ) -> dict[str, Any]:
+        # a is atom count; m is multiplicity, including steering particles.
+        # s is the current sample axis after batch repetition or particle resampling.
+        if steering_args is not None and (
+            steering_args["fk_steering"]
+            or steering_args["physical_guidance_update"]
+            or steering_args["contact_guidance_update"]
+        ):
+            potentials = get_potentials(steering_args, boltz2=True)
+
+        if steering_args["fk_steering"]:
+            multiplicity = multiplicity * steering_args["num_particles"]
+            energy_traj = torch.empty((multiplicity, 0), device=self.device)  # (m, n_evaluation)
+            resample_weights = torch.ones(multiplicity, device=self.device).reshape(
+                -1, steering_args["num_particles"]
+            )  # (n_group, n_particle)
+        if steering_args["physical_guidance_update"] or steering_args["contact_guidance_update"]:
+            scaled_guidance_update = torch.zeros(
+                (multiplicity, *atom_mask.shape[1:], 3),
+                dtype=torch.float32,
+                device=self.device,
+            )  # (s, a, 3)
+        if max_parallel_samples is None:
+            max_parallel_samples = multiplicity
+
+        num_sampling_steps = default(num_sampling_steps, self.num_sampling_steps)
+        atom_mask = atom_mask.repeat_interleave(multiplicity, 0)  # (s, a); s expands on repetition and can shrink on resampling
+
+        shape = (*atom_mask.shape, 3)
+
+        # Pair each schedule value with the next sigma and gamma.
+        sigmas = self.sample_schedule(num_sampling_steps)  # (n_sigma,); schedule or batch axis is determined by the caller
+        gammas = torch.where(sigmas > self.gamma_min, self.gamma_0, 0.0)  # (n_step + 1,)
+        sigmas_and_gammas = list(zip(sigmas[:-1], sigmas[1:], gammas[1:], strict=True))
+        if self.training and self.step_scale_random is not None:
+            step_scale = np.random.choice(self.step_scale_random)
+        else:
+            step_scale = self.step_scale
+
+        # atom position is noise at the beginning
+        init_sigma = sigmas[0]  # ()
+        atom_coords = init_sigma * torch.randn(shape, device=self.device)  # (s, a, 3)
+        token_repr = None  # None; the current sampler does not produce token features
+        atom_coords_denoised = None  # (s, a, 3) or None before the first denoising step
+
+        for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(
+            tqdm(
+                sigmas_and_gammas,
+                desc="Boltz2: Diffusion sampling",
+                unit="step",
+                disable=not verbose,
+                dynamic_ncols=True,
+                leave=False,
+            )
+        ):
+            random_R, random_tr = compute_random_augmentation(
+                multiplicity, device=atom_coords.device, dtype=atom_coords.dtype
+            )  # (m, 3, 3), (m, 1, 3)
+            atom_coords = atom_coords - atom_coords.mean(dim=-2, keepdims=True)  # (s, a, 3)
+            atom_coords = torch.einsum("bmd,bds->bms", atom_coords, random_R) + random_tr  # (s, a, 3)
+            if atom_coords_denoised is not None:
+                atom_coords_denoised -= atom_coords_denoised.mean(dim=-2, keepdims=True)  # (s, a, 3) or None before the first denoising step
+                atom_coords_denoised = (
+                    torch.einsum("bmd,bds->bms", atom_coords_denoised, random_R) + random_tr
+                )  # (s, a, 3) or None before the first denoising step
+            if (
+                steering_args["physical_guidance_update"]
+                or steering_args["contact_guidance_update"]
+            ) and scaled_guidance_update is not None:
+                scaled_guidance_update = torch.einsum(
+                    "bmd,bds->bms", scaled_guidance_update, random_R
+                )  # (s, a, 3)
+
+            sigma_tm, sigma_t, gamma = sigma_tm.item(), sigma_t.item(), gamma.item()  # three scalar tensors () become Python floats
+
+            t_hat = sigma_tm * (1 + gamma)
+            steering_t = 1.0 - (step_idx / num_sampling_steps)
+            noise_var = self.noise_scale**2 * (t_hat**2 - sigma_tm**2)
+            eps = sqrt(noise_var) * torch.randn(shape, device=self.device)  # (s, a, 3)
+            atom_coords_noisy = atom_coords + eps  # (s, a, 3)
+
+            with torch.no_grad():
+                atom_coords_denoised = torch.zeros_like(atom_coords_noisy)  # (s, a, 3) or None before the first denoising step
+                sample_ids = torch.arange(multiplicity).to(atom_coords_noisy.device)  # (m,)
+                sample_ids_chunks = sample_ids.chunk(multiplicity % max_parallel_samples + 1)  # tuple of index tensors, each (n_chunk,)
+
+                for sample_ids_chunk in sample_ids_chunks:
+                    atom_coords_denoised_chunk = self.preconditioned_network_forward(
+                        atom_coords_noisy[sample_ids_chunk],
+                        t_hat,
+                        network_condition_kwargs=dict(
+                            multiplicity=sample_ids_chunk.numel(),
+                            **network_condition_kwargs,
+                        ),
+                    )  # (n_chunk, a, 3)
+                    atom_coords_denoised[sample_ids_chunk] = atom_coords_denoised_chunk  # (n_chunk, a, 3)
+
+                if steering_args["fk_steering"] and (
+                    (step_idx % steering_args["fk_resampling_interval"] == 0 and noise_var > 0)
+                    or step_idx == num_sampling_steps - 1
+                ):
+                    # Compute energy of x_0 prediction
+                    energy = torch.zeros(multiplicity, device=self.device)  # (m,)
+                    for potential in potentials:
+                        parameters = potential.compute_parameters(steering_t)
+                        if parameters["resampling_weight"] > 0:
+                            component_energy = potential.compute(
+                                atom_coords_denoised,
+                                network_condition_kwargs["feats"],
+                                parameters,
+                            )  # (m,)
+                            energy += parameters["resampling_weight"] * component_energy  # (m,)
+                    energy_traj = torch.cat((energy_traj, energy.unsqueeze(1)), dim=1)  # (m, n_evaluation)
+
+                    # Compute log G values
+                    if step_idx == 0:
+                        log_G = -1 * energy  # (m,)
+                    else:
+                        log_G = energy_traj[:, -2] - energy_traj[:, -1]  # (m,)
+
+                    # Compute ll difference between guided and unguided transition distribution
+                    if (
+                        steering_args["physical_guidance_update"]
+                        or steering_args["contact_guidance_update"]
+                    ) and noise_var > 0:
+                        ll_difference = (eps**2 - (eps + scaled_guidance_update) ** 2).sum(
+                            dim=(-1, -2)
+                        ) / (2 * noise_var)  # (m,)
+                    else:
+                        ll_difference = torch.zeros_like(energy)  # (m,)
+
+                    # Compute resampling weights
+                    resample_weights = F.softmax(
+                        (ll_difference + steering_args["fk_lambda"] * log_G).reshape(
+                            -1, steering_args["num_particles"]
+                        ),
+                        dim=1,
+                    )  # (n_group, n_particle)
+
+                # Compute guidance update to x_0 prediction
+                if (
+                    steering_args["physical_guidance_update"]
+                    or steering_args["contact_guidance_update"]
+                ) and step_idx < num_sampling_steps - 1:
+                    guidance_update = torch.zeros_like(atom_coords_denoised)  # (s, a, 3)
+                    for guidance_step in range(steering_args["num_gd_steps"]):
+                        energy_gradient = torch.zeros_like(atom_coords_denoised)  # (s, a, 3)
+                        for potential in potentials:
+                            parameters = potential.compute_parameters(steering_t)
+                            if (
+                                parameters["guidance_weight"] > 0
+                                and (guidance_step) % parameters["guidance_interval"] == 0
+                            ):
+                                energy_gradient += parameters[
+                                    "guidance_weight"
+                                ] * potential.compute_gradient(
+                                    atom_coords_denoised + guidance_update,
+                                    network_condition_kwargs["feats"],
+                                    parameters,
+                                )  # (s, a, 3)
+                        guidance_update -= energy_gradient  # (s, a, 3)
+                    atom_coords_denoised += guidance_update  # (s, a, 3) or None before the first denoising step
+                    scaled_guidance_update = (
+                        guidance_update * -1 * self.step_scale * (sigma_t - t_hat) / t_hat
+                    )  # (s, a, 3)
+
+                if steering_args["fk_steering"] and (
+                    (step_idx % steering_args["fk_resampling_interval"] == 0 and noise_var > 0)
+                    or step_idx == num_sampling_steps - 1
+                ):
+                    resample_indices = (
+                        torch.multinomial(
+                            resample_weights,
+                            resample_weights.shape[1] if step_idx < num_sampling_steps - 1 else 1,
+                            replacement=True,
+                        )
+                        + resample_weights.shape[1]
+                        * torch.arange(
+                            resample_weights.shape[0], device=resample_weights.device
+                        ).unsqueeze(-1)
+                    ).flatten()  # (s_next,); final resampling keeps one particle per group
+
+                    atom_coords = atom_coords[resample_indices]  # (s, a, 3)
+                    atom_coords_noisy = atom_coords_noisy[resample_indices]  # (s, a, 3)
+                    atom_mask = atom_mask[resample_indices]  # (s, a); s expands on repetition and can shrink on resampling
+                    if atom_coords_denoised is not None:
+                        atom_coords_denoised = atom_coords_denoised[resample_indices]  # (s, a, 3) or None before the first denoising step
+                    energy_traj = energy_traj[resample_indices]  # (s_next, n_evaluation)
+                    if (
+                        steering_args["physical_guidance_update"]
+                        or steering_args["contact_guidance_update"]
+                    ):
+                        scaled_guidance_update = scaled_guidance_update[resample_indices]  # (s, a, 3)
+                    if token_repr is not None:
+                        token_repr = token_repr[resample_indices]  # (s_next, *feature_shape)
+
+            if self.alignment_reverse_diff:
+                with torch.autocast("cuda", enabled=False):
+                    atom_coords_noisy = weighted_rigid_align(
+                        atom_coords_noisy.float(),
+                        atom_coords_denoised.float(),
+                        atom_mask.float(),
+                        atom_mask.float(),
+                    )  # (s, a, 3)
+
+                atom_coords_noisy = atom_coords_noisy.to(atom_coords_denoised)  # (s, a, 3)
+
+            denoised_over_sigma = (atom_coords_noisy - atom_coords_denoised) / t_hat  # (s, a, 3)
+            atom_coords_next = (
+                atom_coords_noisy + step_scale * (sigma_t - t_hat) * denoised_over_sigma
+            )  # (s, a, 3)
+
+            atom_coords = atom_coords_next  # (s, a, 3)
+
+        return dict(sample_atom_coords=atom_coords, diff_token_repr=token_repr)
+
+    def loss_weight(self, sigma: torch.Tensor) -> torch.Tensor:
+        # sigma: (...); weight preserves shape (...).
+        return (sigma**2 + self.sigma_data**2) / ((sigma * self.sigma_data) ** 2)
+
+    def noise_distribution(self, batch_size: int) -> torch.Tensor:
+        # Return one noise scale per example, shape (batch_size,).
+        return (
+            self.sigma_data
+            * (self.P_mean + self.P_std * torch.randn((batch_size,), device=self.device)).exp()
+        )
+
+    def forward(
+        self,
+        s_inputs: torch.Tensor,
+        s_trunk: torch.Tensor,
+        feats: dict[str, torch.Tensor],
+        diffusion_conditioning: dict[str, Any],
+        multiplicity: int = 1,
+    ) -> dict[str, Any]:
+        # training diffusion step
+        batch_size = feats["coords"].shape[0] // multiplicity
+
+        if self.synchronize_sigmas:
+            sigmas = self.noise_distribution(batch_size).repeat_interleave(multiplicity, 0)  # (n_sigma,); schedule or batch axis is determined by the caller
+        else:
+            sigmas = self.noise_distribution(batch_size * multiplicity)  # (n_sigma,); schedule or batch axis is determined by the caller
+        padded_sigmas = rearrange(sigmas, "b -> b 1 1")  # (b * m, 1, 1)
+
+        atom_coords = feats["coords"]  # (s, a, 3)
+
+        atom_mask = feats["atom_pad_mask"]  # (s, a); s expands on repetition and can shrink on resampling
+        atom_mask = atom_mask.repeat_interleave(multiplicity, 0)  # (s, a); s expands on repetition and can shrink on resampling
+
+        atom_coords = center_random_augmentation(
+            atom_coords, atom_mask, augmentation=self.coordinate_augmentation
+        )  # (s, a, 3)
+
+        noise = torch.randn_like(atom_coords)  # (b * m, a, 3)
+        noised_atom_coords = atom_coords + padded_sigmas * noise  # (b * m, a, 3)
+
+        denoised_atom_coords = self.preconditioned_network_forward(
+            noised_atom_coords,
+            sigmas,
+            network_condition_kwargs={
+                "s_inputs": s_inputs,
+                "s_trunk": s_trunk,
+                "feats": feats,
+                "multiplicity": multiplicity,
+                "diffusion_conditioning": diffusion_conditioning,
+            },
+        )  # (b * m, a, 3)
+
+        return {
+            "denoised_atom_coords": denoised_atom_coords,
+            "sigmas": sigmas,
+            "aligned_true_atom_coords": atom_coords,
+        }
+
+    def compute_loss(
+        self,
+        feats: dict[str, torch.Tensor],
+        out_dict: dict[str, torch.Tensor],
+        add_smooth_lddt_loss: bool = True,
+        nucleotide_loss_weight: float = 5.0,
+        ligand_loss_weight: float = 10.0,
+        multiplicity: int = 1,
+        filter_by_plddt: float = 0.0,
+    ) -> dict[str, Any]:
+        # Coordinates: (b * m, a, 3); base feature masks: (b, a).
+        with torch.autocast("cuda", enabled=False):
+            denoised_atom_coords = out_dict["denoised_atom_coords"].float()  # (b * m, a, 3)
+            sigmas = out_dict["sigmas"].float()  # (n_sigma,); schedule or batch axis is determined by the caller
+
+            resolved_atom_mask_uni = feats["atom_resolved_mask"].float()  # (b, a)
+
+            if filter_by_plddt > 0:
+                plddt_mask = feats["plddt"] > filter_by_plddt  # (b, a)
+                resolved_atom_mask_uni = resolved_atom_mask_uni * plddt_mask.float()  # (b, a)
+
+            resolved_atom_mask = resolved_atom_mask_uni.repeat_interleave(multiplicity, 0)  # (b * m, a)
+
+            align_weights = denoised_atom_coords.new_ones(denoised_atom_coords.shape[:2])  # (b * m, a)
+            atom_type = (
+                torch.bmm(
+                    feats["atom_to_token"].float(),
+                    feats["mol_type"].unsqueeze(-1).float(),
+                )
+                .squeeze(-1)
+                .long()
+            )  # (b, a)
+            atom_type_mult = atom_type.repeat_interleave(multiplicity, 0)  # (b * m, a)
+
+            align_weights = (
+                align_weights
+                * (
+                    1
+                    + nucleotide_loss_weight
+                    * (
+                        torch.eq(atom_type_mult, const.chain_type_ids["DNA"]).float()
+                        + torch.eq(atom_type_mult, const.chain_type_ids["RNA"]).float()
+                    )
+                    + ligand_loss_weight
+                    * torch.eq(atom_type_mult, const.chain_type_ids["NONPOLYMER"]).float()
+                ).float()
+            )  # (b * m, a)
+
+            atom_coords = out_dict["aligned_true_atom_coords"].float()  # (s, a, 3)
+            atom_coords_aligned_ground_truth = weighted_rigid_align(
+                atom_coords.detach(),
+                denoised_atom_coords.detach(),
+                align_weights.detach(),
+                mask=feats["atom_resolved_mask"]
+                .float()
+                .repeat_interleave(multiplicity, 0)
+                .detach(),
+            )  # (b * m, a, 3)
+
+            atom_coords_aligned_ground_truth = atom_coords_aligned_ground_truth.to(
+                denoised_atom_coords
+            )  # (b * m, a, 3)
+
+            # weighted MSE loss of denoised atom positions
+            mse_loss = ((denoised_atom_coords - atom_coords_aligned_ground_truth) ** 2).sum(dim=-1)  # (b * m, a)
+            mse_loss = torch.sum(mse_loss * align_weights * resolved_atom_mask, dim=-1) / (
+                torch.sum(3 * align_weights * resolved_atom_mask, dim=-1) + 1e-5
+            )  # (b * m,)
+
+            # weight by sigma factor
+            loss_weights = self.loss_weight(sigmas)  # (b * m,)
+            mse_loss = (mse_loss * loss_weights).mean()  # ()
+
+            total_loss = mse_loss  # ()
+
+            # proposed auxiliary smooth lddt loss
+            lddt_loss = self.zero  # ()
+            if add_smooth_lddt_loss:
+                lddt_loss = smooth_lddt_loss(
+                    denoised_atom_coords,
+                    feats["coords"],
+                    torch.eq(atom_type, const.chain_type_ids["DNA"]).float()
+                    + torch.eq(atom_type, const.chain_type_ids["RNA"]).float(),
+                    coords_mask=resolved_atom_mask_uni,
+                    multiplicity=multiplicity,
+                )  # ()
+
+                total_loss = total_loss + lddt_loss  # ()
+
+            loss_breakdown = {
+                "mse_loss": mse_loss,
+                "smooth_lddt_loss": lddt_loss,
+            }
+
+        return {"loss": total_loss, "loss_breakdown": loss_breakdown}

fastplms/models/boltz/vb_modules_encodersv2.py

@@ -1,653 +1,653 @@
-"""Token, pair, and atom encoders used by the local Boltz2 runtime.
-
-The classes retain checkpoint-facing submodule names while separating feature
-assembly, window indexing, conditioning, and atom-token aggregation into small
-mechanism-specific units.  No production import depends on the upstream Boltz
-package.
-
-The Fourier and atom-attention mechanisms derive from AlphaFold 3 community
-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
-from torch import nn
-from torch.nn.functional import one_hot
-
-from . import vb_layers_initialize as init
-from .vb_layers_transition import Transition
-from .vb_modules_transformersv2 import AtomTransformer
-from .vb_modules_utils import LinearNoBias
-
-
-def _transition_stack(
-    count: int,
-    dim: int,
-    hidden_dim: int,
-) -> nn.ModuleList:
-    return nn.ModuleList([Transition(dim=dim, hidden=hidden_dim) for _ in range(count)])
-
-
-class FourierEmbedding(nn.Module):
-    """Embed diffusion time with fixed random Fourier frequencies."""
-
-    def __init__(self, dim: int) -> None:
-        super().__init__()
-        self.proj = nn.Linear(1, dim)
-        nn.init.normal_(self.proj.weight, mean=0, std=1)
-        nn.init.normal_(self.proj.bias, mean=0, std=1)
-        self.proj.requires_grad_(False)
-
-    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))  # (b, d)
-        return torch.cos(2 * pi * random_phase)  # (b, d)
-
-
-def _pairwise_difference(values: torch.Tensor) -> torch.Tensor:
-    # values: (b, n)
-    return values[:, :, None] - values[:, None, :]  # (b, n, n)
-
-
-class RelativePositionEncoder(nn.Module):
-    """Encode residue, token, entity, and symmetry relationships."""
-
-    def __init__(
-        self,
-        token_z: int,
-        r_max: int = 32,
-        s_max: int = 2,
-        fix_sym_check: bool = False,
-        cyclic_pos_enc: bool = False,
-    ) -> None:
-        super().__init__()
-        self.r_max = r_max
-        self.s_max = s_max
-        input_dim = 4 * (r_max + 1) + 2 * (s_max + 1) + 1
-        self.linear_layer = LinearNoBias(input_dim, token_z)
-        self.fix_sym_check = fix_sym_check
-        self.cyclic_pos_enc = cyclic_pos_enc
-
-    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, :]
-        )  # (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,
-            )  # (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,
-        )  # (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,
-        )  # (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,
-        )  # (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,
-        )  # (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(
-            (
-                residue_features.float(),
-                token_features.float(),
-                same_entity.unsqueeze(-1).float(),
-                symmetry_features.float(),
-            ),
-            dim=-1,
-        )  # (b, n, n, d_rel)
-        return self.linear_layer(pair_features)  # (b, n, n, token_z)
-
-
-class SingleConditioning(nn.Module):
-    """Condition token features on trunk inputs and diffusion time."""
-
-    def __init__(
-        self,
-        sigma_data: float,
-        token_s: int = 384,
-        dim_fourier: int = 256,
-        num_transitions: int = 2,
-        transition_expansion_factor: int = 2,
-        eps: float = 1e-20,
-        disable_times: bool = False,
-    ) -> None:
-        super().__init__()
-        self.eps = eps
-        self.sigma_data = sigma_data
-        self.disable_times = disable_times
-        conditioning_dim = 2 * token_s
-        self.norm_single = nn.LayerNorm(conditioning_dim)
-        self.single_embed = nn.Linear(conditioning_dim, conditioning_dim)
-        if not disable_times:
-            self.fourier_embed = FourierEmbedding(dim_fourier)
-            self.norm_fourier = nn.LayerNorm(dim_fourier)
-            self.fourier_to_single = LinearNoBias(dim_fourier, conditioning_dim)
-        self.transitions = _transition_stack(
-            num_transitions,
-            conditioning_dim,
-            transition_expansion_factor * conditioning_dim,
-        )
-
-    def forward(
-        self,
-        times: torch.Tensor,
-        s_trunk: torch.Tensor,
-        s_inputs: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Return conditioned token features ``S: (b, n, 2 * token_s)``."""
-
-        # 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)  # (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  # (b, n, 2 * token_s)
-        return single, normalized_fourier  # (b, n, 2 * token_s), (b, d_fourier) or None
-
-
-class PairwiseConditioning(nn.Module):
-    """Fuse trunk pair features with relative-position features."""
-
-    def __init__(
-        self,
-        token_z: int,
-        dim_token_rel_pos_feats: int,
-        num_transitions: int = 2,
-        transition_expansion_factor: int = 2,
-    ) -> None:
-        super().__init__()
-        combined_dim = token_z + dim_token_rel_pos_feats
-        self.dim_pairwise_init_proj = nn.Sequential(
-            nn.LayerNorm(combined_dim),
-            LinearNoBias(combined_dim, token_z),
-        )
-        self.transitions = _transition_stack(
-            num_transitions,
-            token_z,
-            transition_expansion_factor * token_z,
-        )
-
-    def forward(
-        self,
-        z_trunk: torch.Tensor,
-        token_rel_pos_feats: torch.Tensor,
-    ) -> 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)
-        )  # (b, n, n, token_z)
-        for transition in self.transitions:
-            pair = transition(pair) + pair  # (b, n, n, token_z)
-        return pair  # (b, n, n, token_z)
-
-
-def get_indexing_matrix(
-    k: int,
-    w: int,
-    h: int,
-    device: torch.device,
-) -> torch.Tensor:
-    """Build the atom-window gather matrix used for local attention keys."""
-
-    for name, value in (("k", k), ("w", w), ("h", h)):
-        if isinstance(value, bool) or not isinstance(value, int):
-            raise TypeError(f"{name} must be an int, got {type(value).__name__}.")
-        if value <= 0:
-            raise ValueError(f"{name} must be positive, got {value}.")
-    if w % 2 != 0:
-        raise ValueError(f"w must be even, got {w}.")
-    half_window = w // 2
-    if h % half_window != 0:
-        raise ValueError(f"h must be divisible by w // 2 ({half_window}), got {h}.")
-    key_blocks = h // half_window
-    if key_blocks % 2 != 0:
-        raise ValueError(
-            f"h must contain an even number of half-window key blocks; received {key_blocks}."
-        )
-
-    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(
-    single: torch.Tensor,
-    indexing_matrix: torch.Tensor,
-    w: int,
-    h: int,
-) -> torch.Tensor:
-    """Gather a sequence tensor into overlapping local key windows."""
-
-    b, n, d = single.shape
-    k = n // w
-    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):
-    """Encode reference atoms and local atom-pair geometry."""
-
-    def __init__(
-        self,
-        atom_s: int,
-        atom_z: int,
-        token_s: int,
-        token_z: int,
-        atoms_per_window_queries: int,
-        atoms_per_window_keys: int,
-        atom_feature_dim: int,
-        structure_prediction: bool = True,
-        use_no_atom_char: bool = False,
-        use_atom_backbone_feat: bool = False,
-        use_residue_feats_atoms: bool = False,
-    ) -> None:
-        super().__init__()
-        self.embed_atom_features = nn.Linear(atom_feature_dim, atom_s)
-        self.embed_atompair_ref_pos = LinearNoBias(3, atom_z)
-        self.embed_atompair_ref_dist = LinearNoBias(1, atom_z)
-        self.embed_atompair_mask = LinearNoBias(1, atom_z)
-        self.atoms_per_window_queries = atoms_per_window_queries
-        self.atoms_per_window_keys = atoms_per_window_keys
-        self.use_no_atom_char = use_no_atom_char
-        self.use_atom_backbone_feat = use_atom_backbone_feat
-        self.use_residue_feats_atoms = use_residue_feats_atoms
-        self.structure_prediction = structure_prediction
-
-        if structure_prediction:
-            self.s_to_c_trans = nn.Sequential(
-                nn.LayerNorm(token_s),
-                LinearNoBias(token_s, atom_s),
-            )
-            init.final_init_(self.s_to_c_trans[1].weight)
-            self.z_to_p_trans = nn.Sequential(
-                nn.LayerNorm(token_z),
-                LinearNoBias(token_z, atom_z),
-            )
-            init.final_init_(self.z_to_p_trans[1].weight)
-
-        self.c_to_p_trans_k = nn.Sequential(nn.ReLU(), LinearNoBias(atom_s, atom_z))
-        init.final_init_(self.c_to_p_trans_k[1].weight)
-        self.c_to_p_trans_q = nn.Sequential(nn.ReLU(), LinearNoBias(atom_s, atom_z))
-        init.final_init_(self.c_to_p_trans_q[1].weight)
-        self.p_mlp = nn.Sequential(
-            nn.ReLU(),
-            LinearNoBias(atom_z, atom_z),
-            nn.ReLU(),
-            LinearNoBias(atom_z, atom_z),
-            nn.ReLU(),
-            LinearNoBias(atom_z, atom_z),
-        )
-        init.final_init_(self.p_mlp[5].weight)
-
-    def _assemble_atom_features(
-        self,
-        feats: dict[str, torch.Tensor],
-        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:
-            feature_parts.append(feats["atom_backbone_feat"])
-        if self.use_residue_feats_atoms:
-            residue_features = torch.cat(
-                (
-                    feats["res_type"],
-                    feats["modified"].unsqueeze(-1),
-                    one_hot(feats["mol_type"], num_classes=4).float(),
-                ),
-                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,
-        feats: dict[str, torch.Tensor],
-        s_trunk: torch.Tensor | None = None,
-        z: torch.Tensor | None = None,
-    ) -> tuple[
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor,
-        Callable[[torch.Tensor], torch.Tensor],
-    ]:
-        """Return atom queries, conditioned singles, pairs, and key gatherer."""
-
-        with torch.autocast("cuda", enabled=False):
-            b, n, _ = feats["ref_pos"].shape
-            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
-            )  # (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)  # (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)
-            )  # (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  # (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()  # (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,
-                )  # (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,
-                    k,
-                    1,
-                    h,
-                    conditioned_single.shape[-1],
-                )
-            )  # (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):
-    """Run local atom attention and aggregate atoms to tokens."""
-
-    def __init__(
-        self,
-        atom_s: int,
-        token_s: int,
-        atoms_per_window_queries: int,
-        atoms_per_window_keys: int,
-        atom_encoder_depth: int = 3,
-        atom_encoder_heads: int = 4,
-        structure_prediction: bool = True,
-        activation_checkpointing: bool = False,
-        transformer_post_layer_norm: bool = False,
-    ) -> None:
-        super().__init__()
-        self.structure_prediction = structure_prediction
-        if structure_prediction:
-            self.r_to_q_trans = LinearNoBias(3, atom_s)
-            init.final_init_(self.r_to_q_trans.weight)
-        self.atom_encoder = AtomTransformer(
-            dim=atom_s,
-            dim_single_cond=atom_s,
-            attn_window_queries=atoms_per_window_queries,
-            attn_window_keys=atoms_per_window_keys,
-            depth=atom_encoder_depth,
-            heads=atom_encoder_heads,
-            activation_checkpointing=activation_checkpointing,
-            post_layer_norm=transformer_post_layer_norm,
-        )
-        output_dim = 2 * token_s if structure_prediction else token_s
-        self.atom_to_token_trans = nn.Sequential(
-            LinearNoBias(atom_s, output_dim),
-            nn.ReLU(),
-        )
-
-    def forward(
-        self,
-        feats: dict[str, torch.Tensor],
-        q: torch.Tensor,
-        c: torch.Tensor,
-        atom_enc_bias: torch.Tensor,
-        to_keys: Callable[[torch.Tensor], torch.Tensor],
-        r: torch.Tensor | None = None,
-        multiplicity: int = 1,
-    ) -> tuple[
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor,
-        Callable[[torch.Tensor], torch.Tensor],
-    ]:
-        b, n, _ = feats["ref_pos"].shape
-        del b, n
-        # 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)  # (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,
-            c=c,
-            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()  # (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):
-    """Map token updates back to atoms and predict coordinate displacements."""
-
-    def __init__(
-        self,
-        atom_s: int,
-        token_s: int,
-        attn_window_queries: int,
-        attn_window_keys: int,
-        atom_decoder_depth: int = 3,
-        atom_decoder_heads: int = 4,
-        activation_checkpointing: bool = False,
-        transformer_post_layer_norm: bool = False,
-    ) -> None:
-        super().__init__()
-        self.a_to_q_trans = LinearNoBias(2 * token_s, atom_s)
-        init.final_init_(self.a_to_q_trans.weight)
-        self.atom_decoder = AtomTransformer(
-            dim=atom_s,
-            dim_single_cond=atom_s,
-            attn_window_queries=attn_window_queries,
-            attn_window_keys=attn_window_keys,
-            depth=atom_decoder_depth,
-            heads=atom_decoder_heads,
-            activation_checkpointing=activation_checkpointing,
-            post_layer_norm=transformer_post_layer_norm,
-        )
-        if transformer_post_layer_norm:
-            self.atom_feat_to_atom_pos_update = LinearNoBias(atom_s, 3)
-            init.final_init_(self.atom_feat_to_atom_pos_update.weight)
-        else:
-            self.atom_feat_to_atom_pos_update = nn.Sequential(
-                nn.LayerNorm(atom_s),
-                LinearNoBias(atom_s, 3),
-            )
-            init.final_init_(self.atom_feat_to_atom_pos_update[1].weight)
-
-    def forward(
-        self,
-        a: torch.Tensor,
-        q: torch.Tensor,
-        c: torch.Tensor,
-        atom_dec_bias: torch.Tensor,
-        feats: dict[str, torch.Tensor],
-        to_keys: Callable[[torch.Tensor], torch.Tensor],
-        multiplicity: int = 1,
-    ) -> torch.Tensor:
-        """Return atom-coordinate updates ``R_update: (b, a, 3)``."""
-
-        with torch.autocast("cuda", enabled=False):
-            # 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,
-            c=c,
-            bias=atom_dec_bias,
-            multiplicity=multiplicity,
-            to_keys=to_keys,
-        )  # (b * m, a, atom_s)
-        return self.atom_feat_to_atom_pos_update(q)  # (b * m, a, 3)
+"""Token, pair, and atom encoders used by the local Boltz2 runtime.
+
+The classes retain checkpoint-facing submodule names while separating feature
+assembly, window indexing, conditioning, and atom-token aggregation into small
+mechanism-specific units.  No production import depends on the upstream Boltz
+package.
+
+The Fourier and atom-attention mechanisms derive from AlphaFold 3 community
+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
+from torch import nn
+from torch.nn.functional import one_hot
+
+from . import vb_layers_initialize as init
+from .vb_layers_transition import Transition
+from .vb_modules_transformersv2 import AtomTransformer
+from .vb_modules_utils import LinearNoBias
+
+
+def _transition_stack(
+    count: int,
+    dim: int,
+    hidden_dim: int,
+) -> nn.ModuleList:
+    return nn.ModuleList([Transition(dim=dim, hidden=hidden_dim) for _ in range(count)])
+
+
+class FourierEmbedding(nn.Module):
+    """Embed diffusion time with fixed random Fourier frequencies."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(1, dim)
+        nn.init.normal_(self.proj.weight, mean=0, std=1)
+        nn.init.normal_(self.proj.bias, mean=0, std=1)
+        self.proj.requires_grad_(False)
+
+    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))  # (b, d)
+        return torch.cos(2 * pi * random_phase)  # (b, d)
+
+
+def _pairwise_difference(values: torch.Tensor) -> torch.Tensor:
+    # values: (b, n)
+    return values[:, :, None] - values[:, None, :]  # (b, n, n)
+
+
+class RelativePositionEncoder(nn.Module):
+    """Encode residue, token, entity, and symmetry relationships."""
+
+    def __init__(
+        self,
+        token_z: int,
+        r_max: int = 32,
+        s_max: int = 2,
+        fix_sym_check: bool = False,
+        cyclic_pos_enc: bool = False,
+    ) -> None:
+        super().__init__()
+        self.r_max = r_max
+        self.s_max = s_max
+        input_dim = 4 * (r_max + 1) + 2 * (s_max + 1) + 1
+        self.linear_layer = LinearNoBias(input_dim, token_z)
+        self.fix_sym_check = fix_sym_check
+        self.cyclic_pos_enc = cyclic_pos_enc
+
+    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, :]
+        )  # (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,
+            )  # (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,
+        )  # (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,
+        )  # (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,
+        )  # (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,
+        )  # (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(
+            (
+                residue_features.float(),
+                token_features.float(),
+                same_entity.unsqueeze(-1).float(),
+                symmetry_features.float(),
+            ),
+            dim=-1,
+        )  # (b, n, n, d_rel)
+        return self.linear_layer(pair_features)  # (b, n, n, token_z)
+
+
+class SingleConditioning(nn.Module):
+    """Condition token features on trunk inputs and diffusion time."""
+
+    def __init__(
+        self,
+        sigma_data: float,
+        token_s: int = 384,
+        dim_fourier: int = 256,
+        num_transitions: int = 2,
+        transition_expansion_factor: int = 2,
+        eps: float = 1e-20,
+        disable_times: bool = False,
+    ) -> None:
+        super().__init__()
+        self.eps = eps
+        self.sigma_data = sigma_data
+        self.disable_times = disable_times
+        conditioning_dim = 2 * token_s
+        self.norm_single = nn.LayerNorm(conditioning_dim)
+        self.single_embed = nn.Linear(conditioning_dim, conditioning_dim)
+        if not disable_times:
+            self.fourier_embed = FourierEmbedding(dim_fourier)
+            self.norm_fourier = nn.LayerNorm(dim_fourier)
+            self.fourier_to_single = LinearNoBias(dim_fourier, conditioning_dim)
+        self.transitions = _transition_stack(
+            num_transitions,
+            conditioning_dim,
+            transition_expansion_factor * conditioning_dim,
+        )
+
+    def forward(
+        self,
+        times: torch.Tensor,
+        s_trunk: torch.Tensor,
+        s_inputs: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Return conditioned token features ``S: (b, n, 2 * token_s)``."""
+
+        # 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)  # (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  # (b, n, 2 * token_s)
+        return single, normalized_fourier  # (b, n, 2 * token_s), (b, d_fourier) or None
+
+
+class PairwiseConditioning(nn.Module):
+    """Fuse trunk pair features with relative-position features."""
+
+    def __init__(
+        self,
+        token_z: int,
+        dim_token_rel_pos_feats: int,
+        num_transitions: int = 2,
+        transition_expansion_factor: int = 2,
+    ) -> None:
+        super().__init__()
+        combined_dim = token_z + dim_token_rel_pos_feats
+        self.dim_pairwise_init_proj = nn.Sequential(
+            nn.LayerNorm(combined_dim),
+            LinearNoBias(combined_dim, token_z),
+        )
+        self.transitions = _transition_stack(
+            num_transitions,
+            token_z,
+            transition_expansion_factor * token_z,
+        )
+
+    def forward(
+        self,
+        z_trunk: torch.Tensor,
+        token_rel_pos_feats: torch.Tensor,
+    ) -> 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)
+        )  # (b, n, n, token_z)
+        for transition in self.transitions:
+            pair = transition(pair) + pair  # (b, n, n, token_z)
+        return pair  # (b, n, n, token_z)
+
+
+def get_indexing_matrix(
+    k: int,
+    w: int,
+    h: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Build the atom-window gather matrix used for local attention keys."""
+
+    for name, value in (("k", k), ("w", w), ("h", h)):
+        if isinstance(value, bool) or not isinstance(value, int):
+            raise TypeError(f"{name} must be an int, got {type(value).__name__}.")
+        if value <= 0:
+            raise ValueError(f"{name} must be positive, got {value}.")
+    if w % 2 != 0:
+        raise ValueError(f"w must be even, got {w}.")
+    half_window = w // 2
+    if h % half_window != 0:
+        raise ValueError(f"h must be divisible by w // 2 ({half_window}), got {h}.")
+    key_blocks = h // half_window
+    if key_blocks % 2 != 0:
+        raise ValueError(
+            f"h must contain an even number of half-window key blocks; received {key_blocks}."
+        )
+
+    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(
+    single: torch.Tensor,
+    indexing_matrix: torch.Tensor,
+    w: int,
+    h: int,
+) -> torch.Tensor:
+    """Gather a sequence tensor into overlapping local key windows."""
+
+    b, n, d = single.shape
+    k = n // w
+    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):
+    """Encode reference atoms and local atom-pair geometry."""
+
+    def __init__(
+        self,
+        atom_s: int,
+        atom_z: int,
+        token_s: int,
+        token_z: int,
+        atoms_per_window_queries: int,
+        atoms_per_window_keys: int,
+        atom_feature_dim: int,
+        structure_prediction: bool = True,
+        use_no_atom_char: bool = False,
+        use_atom_backbone_feat: bool = False,
+        use_residue_feats_atoms: bool = False,
+    ) -> None:
+        super().__init__()
+        self.embed_atom_features = nn.Linear(atom_feature_dim, atom_s)
+        self.embed_atompair_ref_pos = LinearNoBias(3, atom_z)
+        self.embed_atompair_ref_dist = LinearNoBias(1, atom_z)
+        self.embed_atompair_mask = LinearNoBias(1, atom_z)
+        self.atoms_per_window_queries = atoms_per_window_queries
+        self.atoms_per_window_keys = atoms_per_window_keys
+        self.use_no_atom_char = use_no_atom_char
+        self.use_atom_backbone_feat = use_atom_backbone_feat
+        self.use_residue_feats_atoms = use_residue_feats_atoms
+        self.structure_prediction = structure_prediction
+
+        if structure_prediction:
+            self.s_to_c_trans = nn.Sequential(
+                nn.LayerNorm(token_s),
+                LinearNoBias(token_s, atom_s),
+            )
+            init.final_init_(self.s_to_c_trans[1].weight)
+            self.z_to_p_trans = nn.Sequential(
+                nn.LayerNorm(token_z),
+                LinearNoBias(token_z, atom_z),
+            )
+            init.final_init_(self.z_to_p_trans[1].weight)
+
+        self.c_to_p_trans_k = nn.Sequential(nn.ReLU(), LinearNoBias(atom_s, atom_z))
+        init.final_init_(self.c_to_p_trans_k[1].weight)
+        self.c_to_p_trans_q = nn.Sequential(nn.ReLU(), LinearNoBias(atom_s, atom_z))
+        init.final_init_(self.c_to_p_trans_q[1].weight)
+        self.p_mlp = nn.Sequential(
+            nn.ReLU(),
+            LinearNoBias(atom_z, atom_z),
+            nn.ReLU(),
+            LinearNoBias(atom_z, atom_z),
+            nn.ReLU(),
+            LinearNoBias(atom_z, atom_z),
+        )
+        init.final_init_(self.p_mlp[5].weight)
+
+    def _assemble_atom_features(
+        self,
+        feats: dict[str, torch.Tensor],
+        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:
+            feature_parts.append(feats["atom_backbone_feat"])
+        if self.use_residue_feats_atoms:
+            residue_features = torch.cat(
+                (
+                    feats["res_type"],
+                    feats["modified"].unsqueeze(-1),
+                    one_hot(feats["mol_type"], num_classes=4).float(),
+                ),
+                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,
+        feats: dict[str, torch.Tensor],
+        s_trunk: torch.Tensor | None = None,
+        z: torch.Tensor | None = None,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        Callable[[torch.Tensor], torch.Tensor],
+    ]:
+        """Return atom queries, conditioned singles, pairs, and key gatherer."""
+
+        with torch.autocast("cuda", enabled=False):
+            b, n, _ = feats["ref_pos"].shape
+            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
+            )  # (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)  # (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)
+            )  # (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  # (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()  # (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,
+                )  # (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,
+                    k,
+                    1,
+                    h,
+                    conditioned_single.shape[-1],
+                )
+            )  # (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):
+    """Run local atom attention and aggregate atoms to tokens."""
+
+    def __init__(
+        self,
+        atom_s: int,
+        token_s: int,
+        atoms_per_window_queries: int,
+        atoms_per_window_keys: int,
+        atom_encoder_depth: int = 3,
+        atom_encoder_heads: int = 4,
+        structure_prediction: bool = True,
+        activation_checkpointing: bool = False,
+        transformer_post_layer_norm: bool = False,
+    ) -> None:
+        super().__init__()
+        self.structure_prediction = structure_prediction
+        if structure_prediction:
+            self.r_to_q_trans = LinearNoBias(3, atom_s)
+            init.final_init_(self.r_to_q_trans.weight)
+        self.atom_encoder = AtomTransformer(
+            dim=atom_s,
+            dim_single_cond=atom_s,
+            attn_window_queries=atoms_per_window_queries,
+            attn_window_keys=atoms_per_window_keys,
+            depth=atom_encoder_depth,
+            heads=atom_encoder_heads,
+            activation_checkpointing=activation_checkpointing,
+            post_layer_norm=transformer_post_layer_norm,
+        )
+        output_dim = 2 * token_s if structure_prediction else token_s
+        self.atom_to_token_trans = nn.Sequential(
+            LinearNoBias(atom_s, output_dim),
+            nn.ReLU(),
+        )
+
+    def forward(
+        self,
+        feats: dict[str, torch.Tensor],
+        q: torch.Tensor,
+        c: torch.Tensor,
+        atom_enc_bias: torch.Tensor,
+        to_keys: Callable[[torch.Tensor], torch.Tensor],
+        r: torch.Tensor | None = None,
+        multiplicity: int = 1,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        Callable[[torch.Tensor], torch.Tensor],
+    ]:
+        b, n, _ = feats["ref_pos"].shape
+        del b, n
+        # 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)  # (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,
+            c=c,
+            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()  # (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):
+    """Map token updates back to atoms and predict coordinate displacements."""
+
+    def __init__(
+        self,
+        atom_s: int,
+        token_s: int,
+        attn_window_queries: int,
+        attn_window_keys: int,
+        atom_decoder_depth: int = 3,
+        atom_decoder_heads: int = 4,
+        activation_checkpointing: bool = False,
+        transformer_post_layer_norm: bool = False,
+    ) -> None:
+        super().__init__()
+        self.a_to_q_trans = LinearNoBias(2 * token_s, atom_s)
+        init.final_init_(self.a_to_q_trans.weight)
+        self.atom_decoder = AtomTransformer(
+            dim=atom_s,
+            dim_single_cond=atom_s,
+            attn_window_queries=attn_window_queries,
+            attn_window_keys=attn_window_keys,
+            depth=atom_decoder_depth,
+            heads=atom_decoder_heads,
+            activation_checkpointing=activation_checkpointing,
+            post_layer_norm=transformer_post_layer_norm,
+        )
+        if transformer_post_layer_norm:
+            self.atom_feat_to_atom_pos_update = LinearNoBias(atom_s, 3)
+            init.final_init_(self.atom_feat_to_atom_pos_update.weight)
+        else:
+            self.atom_feat_to_atom_pos_update = nn.Sequential(
+                nn.LayerNorm(atom_s),
+                LinearNoBias(atom_s, 3),
+            )
+            init.final_init_(self.atom_feat_to_atom_pos_update[1].weight)
+
+    def forward(
+        self,
+        a: torch.Tensor,
+        q: torch.Tensor,
+        c: torch.Tensor,
+        atom_dec_bias: torch.Tensor,
+        feats: dict[str, torch.Tensor],
+        to_keys: Callable[[torch.Tensor], torch.Tensor],
+        multiplicity: int = 1,
+    ) -> torch.Tensor:
+        """Return atom-coordinate updates ``R_update: (b, a, 3)``."""
+
+        with torch.autocast("cuda", enabled=False):
+            # 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,
+            c=c,
+            bias=atom_dec_bias,
+            multiplicity=multiplicity,
+            to_keys=to_keys,
+        )  # (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

@@ -1,276 +1,276 @@
-"""Conditioned transformer blocks used by Boltz2 diffusion."""
-
-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
-
-from .vb_layers_attentionv2 import AttentionPairBias
-from .vb_modules_utils import LinearNoBias, SwiGLU
-
-
-class AdaLN(nn.Module):
-    """Normalize activations and apply scale and shift from conditioning S."""
-
-    def __init__(self, dim: int, dim_single_cond: int) -> None:
-        super().__init__()
-        self.a_norm = nn.LayerNorm(dim, elementwise_affine=False, bias=False)
-        self.s_norm = nn.LayerNorm(dim_single_cond, bias=False)
-        self.s_scale = nn.Linear(dim_single_cond, dim)
-        self.s_bias = LinearNoBias(dim_single_cond, dim)
-
-    def forward(self, a: Tensor, s: Tensor) -> Tensor:
-        # 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)  # (..., d)
-
-
-class ConditionedTransitionBlock(nn.Module):
-    """Apply a conditioned gated transition without changing tensor shape."""
-
-    def __init__(
-        self,
-        dim_single: int,
-        dim_single_cond: int,
-        expansion_factor: float = 2,
-    ) -> None:
-        super().__init__()
-        self.adaln = AdaLN(dim_single, dim_single_cond)
-        inner_dim = int(dim_single * expansion_factor)
-        self.swish_gate = nn.Sequential(
-            LinearNoBias(dim_single, inner_dim * 2),
-            SwiGLU(),
-        )
-        self.a_to_b = LinearNoBias(dim_single, inner_dim)
-        self.b_to_a = LinearNoBias(inner_dim, dim_single)
-
-        projection = nn.Linear(dim_single_cond, dim_single)
-        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:
-        # 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))
-
-
-class DiffusionTransformer(nn.Module):
-    """Run conditioned transformer layers over token or atom states."""
-
-    def __init__(
-        self,
-        depth: int,
-        heads: int,
-        dim: int = 384,
-        dim_single_cond: int | None = None,
-        pair_bias_attn: bool = True,
-        activation_checkpointing: bool = False,
-        post_layer_norm: bool = False,
-    ) -> None:
-        super().__init__()
-        conditioning_dim = dim if dim_single_cond is None else dim_single_cond
-        self.activation_checkpointing = activation_checkpointing
-        self.pair_bias_attn = pair_bias_attn
-        self.layers = nn.ModuleList(
-            [
-                DiffusionTransformerLayer(
-                    heads,
-                    dim,
-                    conditioning_dim,
-                    post_layer_norm,
-                )
-                for _ in range(depth)
-            ]
-        )
-
-    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  # 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(
-        self,
-        a: Tensor,
-        s: Tensor,
-        bias: Tensor | None = None,
-        mask: Tensor | None = None,
-        to_keys: Callable[[Tensor], Tensor] | None = None,
-        multiplicity: int = 1,
-    ) -> Tensor:
-        """Transform A and preserve shape ``(b * m, l, d)``."""
-
-        # 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(  # (b * m, l_q, d)
-                    layer,
-                    output,
-                    s,
-                    layer_bias,
-                    mask,
-                    to_keys,
-                    multiplicity,
-                    use_reentrant=False,
-                )
-            else:
-                output = layer(  # (b * m, l_q, d)
-                    output,
-                    s,
-                    layer_bias,
-                    mask,
-                    to_keys,
-                    multiplicity,
-                )
-        return output  # (b * m, l_q, d)
-
-
-class DiffusionTransformerLayer(nn.Module):
-    """One adaptive-normalization attention and transition block."""
-
-    def __init__(
-        self,
-        heads: int,
-        dim: int = 384,
-        dim_single_cond: int | None = None,
-        post_layer_norm: bool = False,
-    ) -> None:
-        super().__init__()
-        conditioning_dim = dim if dim_single_cond is None else dim_single_cond
-        self.adaln = AdaLN(dim, conditioning_dim)
-        self.pair_bias_attn = AttentionPairBias(
-            c_s=dim,
-            num_heads=heads,
-            compute_pair_bias=False,
-        )
-        self.output_projection_linear = nn.Linear(conditioning_dim, dim)
-        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(),
-        )
-        self.transition = ConditionedTransitionBlock(dim, conditioning_dim)
-        self.post_lnorm = nn.LayerNorm(dim) if post_layer_norm else nn.Identity()
-
-    def forward(
-        self,
-        a: Tensor,
-        s: Tensor,
-        bias: Tensor | None = None,
-        mask: Tensor | None = None,
-        to_keys: Callable[[Tensor], Tensor] | None = None,
-        multiplicity: int = 1,
-    ) -> 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)  # (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)  # (b * m, l_k, d)
-            key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)  # (b * m, l_k)
-
-        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  # (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):
-    """Apply the diffusion transformer over fixed atom-query windows."""
-
-    def __init__(
-        self,
-        attn_window_queries: int,
-        attn_window_keys: int,
-        **diffusion_transformer_kwargs: Any,
-    ) -> None:
-        super().__init__()
-        self.attn_window_queries = attn_window_queries
-        self.attn_window_keys = attn_window_keys
-        self.diffusion_transformer = DiffusionTransformer(
-            **diffusion_transformer_kwargs,
-        )
-
-    def forward(
-        self,
-        q: Tensor,
-        c: Tensor,
-        bias: Tensor,
-        to_keys: Callable[[Tensor], Tensor],
-        mask: Tensor,
-        multiplicity: int = 1,
-    ) -> Tensor:
-        """Transform atom tensor Q with shape ``(b, n_atoms, d)``."""
-
-        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(  # (b * m * k, w_q, w_k, n_layer * h)
-            pair_bias.shape[0] * window_count,
-            query_window,
-            key_window,
-            -1,
-        )
-
-        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(  # (b * m * k, w_q, d)
-            a=query_states,
-            s=conditioning,
-            bias=pair_bias,
-            mask=query_mask.float(),
-            multiplicity=1,
-            to_keys=windowed_keys,
-        )
-        # (b * m, n_atom, d)
-        return cast(Tensor, output.view(batch_size, window_count * query_window, width))
+"""Conditioned transformer blocks used by Boltz2 diffusion."""
+
+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
+
+from .vb_layers_attentionv2 import AttentionPairBias
+from .vb_modules_utils import LinearNoBias, SwiGLU
+
+
+class AdaLN(nn.Module):
+    """Normalize activations and apply scale and shift from conditioning S."""
+
+    def __init__(self, dim: int, dim_single_cond: int) -> None:
+        super().__init__()
+        self.a_norm = nn.LayerNorm(dim, elementwise_affine=False, bias=False)
+        self.s_norm = nn.LayerNorm(dim_single_cond, bias=False)
+        self.s_scale = nn.Linear(dim_single_cond, dim)
+        self.s_bias = LinearNoBias(dim_single_cond, dim)
+
+    def forward(self, a: Tensor, s: Tensor) -> Tensor:
+        # 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)  # (..., d)
+
+
+class ConditionedTransitionBlock(nn.Module):
+    """Apply a conditioned gated transition without changing tensor shape."""
+
+    def __init__(
+        self,
+        dim_single: int,
+        dim_single_cond: int,
+        expansion_factor: float = 2,
+    ) -> None:
+        super().__init__()
+        self.adaln = AdaLN(dim_single, dim_single_cond)
+        inner_dim = int(dim_single * expansion_factor)
+        self.swish_gate = nn.Sequential(
+            LinearNoBias(dim_single, inner_dim * 2),
+            SwiGLU(),
+        )
+        self.a_to_b = LinearNoBias(dim_single, inner_dim)
+        self.b_to_a = LinearNoBias(inner_dim, dim_single)
+
+        projection = nn.Linear(dim_single_cond, dim_single)
+        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:
+        # 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))
+
+
+class DiffusionTransformer(nn.Module):
+    """Run conditioned transformer layers over token or atom states."""
+
+    def __init__(
+        self,
+        depth: int,
+        heads: int,
+        dim: int = 384,
+        dim_single_cond: int | None = None,
+        pair_bias_attn: bool = True,
+        activation_checkpointing: bool = False,
+        post_layer_norm: bool = False,
+    ) -> None:
+        super().__init__()
+        conditioning_dim = dim if dim_single_cond is None else dim_single_cond
+        self.activation_checkpointing = activation_checkpointing
+        self.pair_bias_attn = pair_bias_attn
+        self.layers = nn.ModuleList(
+            [
+                DiffusionTransformerLayer(
+                    heads,
+                    dim,
+                    conditioning_dim,
+                    post_layer_norm,
+                )
+                for _ in range(depth)
+            ]
+        )
+
+    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  # 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(
+        self,
+        a: Tensor,
+        s: Tensor,
+        bias: Tensor | None = None,
+        mask: Tensor | None = None,
+        to_keys: Callable[[Tensor], Tensor] | None = None,
+        multiplicity: int = 1,
+    ) -> Tensor:
+        """Transform A and preserve shape ``(b * m, l, d)``."""
+
+        # 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(  # (b * m, l_q, d)
+                    layer,
+                    output,
+                    s,
+                    layer_bias,
+                    mask,
+                    to_keys,
+                    multiplicity,
+                    use_reentrant=False,
+                )
+            else:
+                output = layer(  # (b * m, l_q, d)
+                    output,
+                    s,
+                    layer_bias,
+                    mask,
+                    to_keys,
+                    multiplicity,
+                )
+        return output  # (b * m, l_q, d)
+
+
+class DiffusionTransformerLayer(nn.Module):
+    """One adaptive-normalization attention and transition block."""
+
+    def __init__(
+        self,
+        heads: int,
+        dim: int = 384,
+        dim_single_cond: int | None = None,
+        post_layer_norm: bool = False,
+    ) -> None:
+        super().__init__()
+        conditioning_dim = dim if dim_single_cond is None else dim_single_cond
+        self.adaln = AdaLN(dim, conditioning_dim)
+        self.pair_bias_attn = AttentionPairBias(
+            c_s=dim,
+            num_heads=heads,
+            compute_pair_bias=False,
+        )
+        self.output_projection_linear = nn.Linear(conditioning_dim, dim)
+        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(),
+        )
+        self.transition = ConditionedTransitionBlock(dim, conditioning_dim)
+        self.post_lnorm = nn.LayerNorm(dim) if post_layer_norm else nn.Identity()
+
+    def forward(
+        self,
+        a: Tensor,
+        s: Tensor,
+        bias: Tensor | None = None,
+        mask: Tensor | None = None,
+        to_keys: Callable[[Tensor], Tensor] | None = None,
+        multiplicity: int = 1,
+    ) -> 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)  # (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)  # (b * m, l_k, d)
+            key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)  # (b * m, l_k)
+
+        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  # (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):
+    """Apply the diffusion transformer over fixed atom-query windows."""
+
+    def __init__(
+        self,
+        attn_window_queries: int,
+        attn_window_keys: int,
+        **diffusion_transformer_kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self.attn_window_queries = attn_window_queries
+        self.attn_window_keys = attn_window_keys
+        self.diffusion_transformer = DiffusionTransformer(
+            **diffusion_transformer_kwargs,
+        )
+
+    def forward(
+        self,
+        q: Tensor,
+        c: Tensor,
+        bias: Tensor,
+        to_keys: Callable[[Tensor], Tensor],
+        mask: Tensor,
+        multiplicity: int = 1,
+    ) -> Tensor:
+        """Transform atom tensor Q with shape ``(b, n_atoms, d)``."""
+
+        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(  # (b * m * k, w_q, w_k, n_layer * h)
+            pair_bias.shape[0] * window_count,
+            query_window,
+            key_window,
+            -1,
+        )
+
+        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(  # (b * m * k, w_q, d)
+            a=query_states,
+            s=conditioning,
+            bias=pair_bias,
+            mask=query_mask.float(),
+            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

@@ -1,624 +1,624 @@
-"""Input, template, MSA, and prediction-head modules for the Boltz2 trunk."""
-
-from __future__ import annotations
-
-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
-
-from . import vb_const as const
-from .vb_layers_dropout import get_dropout_mask
-from .vb_layers_outer_product_mean import OuterProductMean
-from .vb_layers_pair_averaging import PairWeightedAveraging
-from .vb_layers_pairformer import PairformerNoSeqLayer, PairformerNoSeqModule
-from .vb_layers_transition import Transition
-from .vb_modules_encodersv2 import (
-    AtomAttentionEncoder,
-    AtomEncoder,
-    FourierEmbedding,
-)
-
-
-class ContactConditioning(nn.Module):
-    """Encode selected, unselected, and distance-threshold contact constraints."""
-
-    def __init__(self, token_z: int, cutoff_min: float, cutoff_max: float) -> None:
-        super().__init__()
-        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))  # (d_z,)
-        self.encoding_unselected = nn.Parameter(torch.zeros(token_z))  # (d_z,)
-        self.cutoff_min = cutoff_min
-        self.cutoff_max = cutoff_max
-
-    def forward(self, feats: dict[str, Tensor]) -> Tensor:
-        """Return contact tensor C with shape ``(b, l, l, d_z)``."""
-
-        if const.contact_conditioning_info["UNSPECIFIED"] != 0:
-            raise ValueError("UNSPECIFIED contact conditioning must use channel zero")
-        if const.contact_conditioning_info["UNSELECTED"] != 1:
-            raise ValueError("UNSELECTED contact conditioning must use channel one")
-
-        # 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(  # (b, l, l, c_contact - 1 + d_z)
-            [categories[..., 2:], normalized.unsqueeze(-1), fourier],
-            dim=-1,
-        )
-        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):
-    """Combine atom, residue, profile, and optional experimental features."""
-
-    def __init__(
-        self,
-        atom_s: int,
-        atom_z: int,
-        token_s: int,
-        token_z: int,
-        atoms_per_window_queries: int,
-        atoms_per_window_keys: int,
-        atom_feature_dim: int,
-        atom_encoder_depth: int,
-        atom_encoder_heads: int,
-        activation_checkpointing: bool = False,
-        add_method_conditioning: bool = False,
-        add_modified_flag: bool = False,
-        add_cyclic_flag: bool = False,
-        add_mol_type_feat: bool = False,
-        use_no_atom_char: bool = False,
-        use_atom_backbone_feat: bool = False,
-        use_residue_feats_atoms: bool = False,
-    ) -> None:
-        super().__init__()
-        self.token_s = token_s
-        self.add_method_conditioning = add_method_conditioning
-        self.add_modified_flag = add_modified_flag
-        self.add_cyclic_flag = add_cyclic_flag
-        self.add_mol_type_feat = add_mol_type_feat
-        self.atom_encoder = AtomEncoder(
-            atom_s=atom_s,
-            atom_z=atom_z,
-            token_s=token_s,
-            token_z=token_z,
-            atoms_per_window_queries=atoms_per_window_queries,
-            atoms_per_window_keys=atoms_per_window_keys,
-            atom_feature_dim=atom_feature_dim,
-            structure_prediction=False,
-            use_no_atom_char=use_no_atom_char,
-            use_atom_backbone_feat=use_atom_backbone_feat,
-            use_residue_feats_atoms=use_residue_feats_atoms,
-        )
-        self.atom_enc_proj_z = nn.Sequential(
-            nn.LayerNorm(atom_z),
-            nn.Linear(atom_z, atom_encoder_depth * atom_encoder_heads, bias=False),
-        )
-        self.atom_attention_encoder = AtomAttentionEncoder(
-            atom_s=atom_s,
-            token_s=token_s,
-            atoms_per_window_queries=atoms_per_window_queries,
-            atoms_per_window_keys=atoms_per_window_keys,
-            atom_encoder_depth=atom_encoder_depth,
-            atom_encoder_heads=atom_encoder_heads,
-            structure_prediction=False,
-            activation_checkpointing=activation_checkpointing,
-        )
-        self.res_type_encoding = nn.Linear(const.num_tokens, token_s, bias=False)
-        self.msa_profile_encoding = nn.Linear(const.num_tokens + 1, token_s, bias=False)
-
-        if add_method_conditioning:
-            self.method_conditioning_init = nn.Embedding(const.num_method_types, token_s)
-            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)  # (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)  # (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)  # (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)``."""
-
-        # 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}"]  # (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)  # (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,
-        )  # (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)  # (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)  # (b, l, d_s)
-
-
-class _TemplateBase(nn.Module):
-    def __init__(
-        self,
-        token_z: int,
-        template_dim: int,
-        template_blocks: int,
-        dropout: float,
-        pairwise_head_width: int,
-        pairwise_num_heads: int,
-        post_layer_norm: bool,
-        activation_checkpointing: bool,
-        min_dist: float,
-        max_dist: float,
-        num_bins: int,
-    ) -> None:
-        super().__init__()
-        self.min_dist = min_dist
-        self.max_dist = max_dist
-        self.num_bins = num_bins
-        self.relu = nn.ReLU()
-        self.z_norm = nn.LayerNorm(token_z)
-        self.v_norm = nn.LayerNorm(template_dim)
-        self.z_proj = nn.Linear(token_z, template_dim, bias=False)
-        feature_width = const.num_tokens * 2 + num_bins + 5
-        self.a_proj = nn.Linear(feature_width, template_dim, bias=False)
-        self.u_proj = nn.Linear(template_dim, token_z, bias=False)
-        self.pairformer = PairformerNoSeqModule(
-            template_dim,
-            num_blocks=template_blocks,
-            dropout=dropout,
-            pairwise_head_width=pairwise_head_width,
-            pairwise_num_heads=pairwise_num_heads,
-            post_layer_norm=post_layer_norm,
-            activation_checkpointing=activation_checkpointing,
-        )
-
-    def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
-        raise NotImplementedError
-
-    def _template_features(
-        self,
-        feats: dict[str, Tensor],
-        template_pair_mask: Tensor,
-    ) -> Tensor:
-        # 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(  # (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()  # (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)  # (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(  # (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(  # (b, t, l, l, n_token_type)
-                -1,
-                -1,
-                -1,
-                residue_type.size(2),
-                -1,
-            )
-            residue_j = residue_type[:, :, None, :].expand(  # (b, t, l, l, n_token_type)
-                -1,
-                -1,
-                residue_type.size(2),
-                -1,
-                -1,
-            )
-            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,
-        z: Tensor,
-        feats: dict[str, Tensor],
-        pair_mask: Tensor,
-        use_kernels: bool = False,
-    ) -> Tensor:
-        """Aggregate template pair tensor V into trunk update U."""
-
-        # 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()  # (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)  # (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  # (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(  # (b * t, l, l, d_t)
-            template_states,
-            expanded_mask,
-            use_kernels=use_kernels,
-        )
-        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]  # (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)))  # (b, l, l, d_z)
-
-
-class TemplateModule(_TemplateBase):
-    """Aggregate templates while restricting features to the same chain."""
-
-    def __init__(
-        self,
-        token_z: int,
-        template_dim: int,
-        template_blocks: int,
-        dropout: float = 0.25,
-        pairwise_head_width: int = 32,
-        pairwise_num_heads: int = 4,
-        post_layer_norm: bool = False,
-        activation_checkpointing: bool = False,
-        min_dist: float = 3.25,
-        max_dist: float = 50.75,
-        num_bins: int = 38,
-        **kwargs: Any,
-    ) -> None:
-        del kwargs
-        super().__init__(
-            token_z,
-            template_dim,
-            template_blocks,
-            dropout,
-            pairwise_head_width,
-            pairwise_num_heads,
-            post_layer_norm,
-            activation_checkpointing,
-            min_dist,
-            max_dist,
-            num_bins,
-        )
-
-    def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
-        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):
-    """Aggregate templates under per-template visibility groups."""
-
-    def __init__(
-        self,
-        token_z: int,
-        template_dim: int,
-        template_blocks: int,
-        dropout: float = 0.25,
-        pairwise_head_width: int = 32,
-        pairwise_num_heads: int = 4,
-        post_layer_norm: bool = False,
-        activation_checkpointing: bool = False,
-        min_dist: float = 3.25,
-        max_dist: float = 50.75,
-        num_bins: int = 38,
-        **kwargs: Any,
-    ) -> None:
-        del kwargs
-        super().__init__(
-            token_z,
-            template_dim,
-            template_blocks,
-            dropout,
-            pairwise_head_width,
-            pairwise_num_heads,
-            post_layer_norm,
-            activation_checkpointing,
-            min_dist,
-            max_dist,
-            num_bins,
-        )
-
-    def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
-        del count
-        visibility = feats["visibility_ids"]  # (b, t, l)
-        return (visibility[..., :, None] == visibility[..., None, :]).float()  # (b, t, l, l)
-
-
-class MSAModule(nn.Module):
-    """Embed and update an MSA before returning its accumulated pair update."""
-
-    def __init__(
-        self,
-        msa_s: int,
-        token_z: int,
-        token_s: int,
-        msa_blocks: int,
-        msa_dropout: float,
-        z_dropout: float,
-        pairwise_head_width: int = 32,
-        pairwise_num_heads: int = 4,
-        activation_checkpointing: bool = False,
-        use_paired_feature: bool = True,
-        subsample_msa: bool = False,
-        num_subsampled_msa: int = 1024,
-        **kwargs: Any,
-    ) -> None:
-        del kwargs
-        super().__init__()
-        self.msa_blocks = msa_blocks
-        self.msa_dropout = msa_dropout
-        self.z_dropout = z_dropout
-        self.use_paired_feature = use_paired_feature
-        self.activation_checkpointing = activation_checkpointing
-        self.subsample_msa = subsample_msa
-        self.num_subsampled_msa = num_subsampled_msa
-        self.s_proj = nn.Linear(token_s, msa_s, bias=False)
-        input_width = const.num_tokens + 2 + int(use_paired_feature)
-        self.msa_proj = nn.Linear(input_width, msa_s, bias=False)
-        self.layers = nn.ModuleList(
-            [
-                MSALayer(
-                    msa_s,
-                    token_z,
-                    msa_dropout,
-                    z_dropout,
-                    pairwise_head_width,
-                    pairwise_num_heads,
-                )
-                for _ in range(msa_blocks)
-            ]
-        )
-
-    @staticmethod
-    def _chunk_configuration(
-        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:
-            return True, 64, 32, 4, 128
-        return False, None, None, None, 512
-
-    def forward(
-        self,
-        z: Tensor,
-        emb: Tensor,
-        feats: dict[str, Tensor],
-        use_kernels: bool = False,
-    ) -> Tensor:
-        """Return updated pair tensor Z after every MSA block."""
-
-        chunking = self._chunk_configuration(z, self.training)
-        # 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()  # (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),  # (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))  # (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]  # (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,
-                pair_mask,
-                msa_mask,
-                *chunking,
-                use_kernels,
-            )
-            if self.activation_checkpointing and self.training:
-                pair_states, msa_states = checkpoint(  # (b, l, l, d_z), (b, s, l, d_m)
-                    layer,
-                    *arguments,
-                )
-            else:
-                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):
-    """Exchange information between MSA tensor M and pair tensor Z."""
-
-    def __init__(
-        self,
-        msa_s: int,
-        token_z: int,
-        msa_dropout: float,
-        z_dropout: float,
-        pairwise_head_width: int = 32,
-        pairwise_num_heads: int = 4,
-    ) -> None:
-        super().__init__()
-        self.msa_dropout = msa_dropout
-        self.msa_transition = Transition(msa_s, msa_s * 4)
-        self.pair_weighted_averaging = PairWeightedAveraging(
-            c_m=msa_s,
-            c_z=token_z,
-            c_h=32,
-            num_heads=8,
-        )
-        self.pairformer_layer = PairformerNoSeqLayer(
-            token_z=token_z,
-            dropout=z_dropout,
-            pairwise_head_width=pairwise_head_width,
-            pairwise_num_heads=pairwise_num_heads,
-        )
-        self.outer_product_mean = OuterProductMean(msa_s, 32, token_z)
-
-    def forward(
-        self,
-        z: Tensor,
-        m: Tensor,
-        token_mask: Tensor,
-        msa_mask: Tensor,
-        chunk_heads_pwa: bool = False,
-        chunk_size_transition_z: int | None = None,
-        chunk_size_transition_msa: int | None = None,
-        chunk_size_outer_product: int | None = None,
-        chunk_size_tri_attn: int | None = None,
-        use_kernels: bool = False,
-    ) -> tuple[Tensor, Tensor]:
-        """Return updated Z and M tensors."""
-
-        del chunk_size_transition_z
-        # 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(  # (b, s, l, d_m)
-            msa_states,
-            chunk_size_transition_msa,
-        )
-        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(  # (b, l, l, d_z)
-            pair_states,
-            token_mask,
-            chunk_size_tri_attn,
-            use_kernels=use_kernels,
-        )
-        return pair_states, msa_states  # (b, l, l, d_z), (b, s, l, d_m)
-
-
-class BFactorModule(nn.Module):
-    """Predict a per-token B-factor histogram."""
-
-    def __init__(self, token_s: int, num_bins: int) -> None:
-        super().__init__()
-        self.bfactor = nn.Linear(token_s, num_bins)
-        self.num_bins = num_bins
-
-    def forward(self, s: Tensor) -> Tensor:
-        # s: (..., d_s).
-        return cast(Tensor, self.bfactor(s))  # (..., n_bin)
-
-
-class DistogramModule(nn.Module):
-    """Predict symmetric residue-pair distance histograms."""
-
-    def __init__(self, token_z: int, num_bins: int, num_distograms: int = 1) -> None:
-        super().__init__()
-        self.distogram = nn.Linear(token_z, num_distograms * num_bins)
-        self.num_distograms = num_distograms
-        self.num_bins = num_bins
-
-    def forward(self, z: Tensor) -> Tensor:
-        # 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(
-                symmetric.shape[0],
-                symmetric.shape[1],
-                symmetric.shape[2],
-                self.num_distograms,
-                self.num_bins,
-            ),
-        )  # (b, l, l, n_distogram, n_bin)
+"""Input, template, MSA, and prediction-head modules for the Boltz2 trunk."""
+
+from __future__ import annotations
+
+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
+
+from . import vb_const as const
+from .vb_layers_dropout import get_dropout_mask
+from .vb_layers_outer_product_mean import OuterProductMean
+from .vb_layers_pair_averaging import PairWeightedAveraging
+from .vb_layers_pairformer import PairformerNoSeqLayer, PairformerNoSeqModule
+from .vb_layers_transition import Transition
+from .vb_modules_encodersv2 import (
+    AtomAttentionEncoder,
+    AtomEncoder,
+    FourierEmbedding,
+)
+
+
+class ContactConditioning(nn.Module):
+    """Encode selected, unselected, and distance-threshold contact constraints."""
+
+    def __init__(self, token_z: int, cutoff_min: float, cutoff_max: float) -> None:
+        super().__init__()
+        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))  # (d_z,)
+        self.encoding_unselected = nn.Parameter(torch.zeros(token_z))  # (d_z,)
+        self.cutoff_min = cutoff_min
+        self.cutoff_max = cutoff_max
+
+    def forward(self, feats: dict[str, Tensor]) -> Tensor:
+        """Return contact tensor C with shape ``(b, l, l, d_z)``."""
+
+        if const.contact_conditioning_info["UNSPECIFIED"] != 0:
+            raise ValueError("UNSPECIFIED contact conditioning must use channel zero")
+        if const.contact_conditioning_info["UNSELECTED"] != 1:
+            raise ValueError("UNSELECTED contact conditioning must use channel one")
+
+        # 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(  # (b, l, l, c_contact - 1 + d_z)
+            [categories[..., 2:], normalized.unsqueeze(-1), fourier],
+            dim=-1,
+        )
+        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):
+    """Combine atom, residue, profile, and optional experimental features."""
+
+    def __init__(
+        self,
+        atom_s: int,
+        atom_z: int,
+        token_s: int,
+        token_z: int,
+        atoms_per_window_queries: int,
+        atoms_per_window_keys: int,
+        atom_feature_dim: int,
+        atom_encoder_depth: int,
+        atom_encoder_heads: int,
+        activation_checkpointing: bool = False,
+        add_method_conditioning: bool = False,
+        add_modified_flag: bool = False,
+        add_cyclic_flag: bool = False,
+        add_mol_type_feat: bool = False,
+        use_no_atom_char: bool = False,
+        use_atom_backbone_feat: bool = False,
+        use_residue_feats_atoms: bool = False,
+    ) -> None:
+        super().__init__()
+        self.token_s = token_s
+        self.add_method_conditioning = add_method_conditioning
+        self.add_modified_flag = add_modified_flag
+        self.add_cyclic_flag = add_cyclic_flag
+        self.add_mol_type_feat = add_mol_type_feat
+        self.atom_encoder = AtomEncoder(
+            atom_s=atom_s,
+            atom_z=atom_z,
+            token_s=token_s,
+            token_z=token_z,
+            atoms_per_window_queries=atoms_per_window_queries,
+            atoms_per_window_keys=atoms_per_window_keys,
+            atom_feature_dim=atom_feature_dim,
+            structure_prediction=False,
+            use_no_atom_char=use_no_atom_char,
+            use_atom_backbone_feat=use_atom_backbone_feat,
+            use_residue_feats_atoms=use_residue_feats_atoms,
+        )
+        self.atom_enc_proj_z = nn.Sequential(
+            nn.LayerNorm(atom_z),
+            nn.Linear(atom_z, atom_encoder_depth * atom_encoder_heads, bias=False),
+        )
+        self.atom_attention_encoder = AtomAttentionEncoder(
+            atom_s=atom_s,
+            token_s=token_s,
+            atoms_per_window_queries=atoms_per_window_queries,
+            atoms_per_window_keys=atoms_per_window_keys,
+            atom_encoder_depth=atom_encoder_depth,
+            atom_encoder_heads=atom_encoder_heads,
+            structure_prediction=False,
+            activation_checkpointing=activation_checkpointing,
+        )
+        self.res_type_encoding = nn.Linear(const.num_tokens, token_s, bias=False)
+        self.msa_profile_encoding = nn.Linear(const.num_tokens + 1, token_s, bias=False)
+
+        if add_method_conditioning:
+            self.method_conditioning_init = nn.Embedding(const.num_method_types, token_s)
+            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)  # (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)  # (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)  # (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)``."""
+
+        # 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}"]  # (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)  # (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,
+        )  # (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)  # (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)  # (b, l, d_s)
+
+
+class _TemplateBase(nn.Module):
+    def __init__(
+        self,
+        token_z: int,
+        template_dim: int,
+        template_blocks: int,
+        dropout: float,
+        pairwise_head_width: int,
+        pairwise_num_heads: int,
+        post_layer_norm: bool,
+        activation_checkpointing: bool,
+        min_dist: float,
+        max_dist: float,
+        num_bins: int,
+    ) -> None:
+        super().__init__()
+        self.min_dist = min_dist
+        self.max_dist = max_dist
+        self.num_bins = num_bins
+        self.relu = nn.ReLU()
+        self.z_norm = nn.LayerNorm(token_z)
+        self.v_norm = nn.LayerNorm(template_dim)
+        self.z_proj = nn.Linear(token_z, template_dim, bias=False)
+        feature_width = const.num_tokens * 2 + num_bins + 5
+        self.a_proj = nn.Linear(feature_width, template_dim, bias=False)
+        self.u_proj = nn.Linear(template_dim, token_z, bias=False)
+        self.pairformer = PairformerNoSeqModule(
+            template_dim,
+            num_blocks=template_blocks,
+            dropout=dropout,
+            pairwise_head_width=pairwise_head_width,
+            pairwise_num_heads=pairwise_num_heads,
+            post_layer_norm=post_layer_norm,
+            activation_checkpointing=activation_checkpointing,
+        )
+
+    def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
+        raise NotImplementedError
+
+    def _template_features(
+        self,
+        feats: dict[str, Tensor],
+        template_pair_mask: Tensor,
+    ) -> Tensor:
+        # 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(  # (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()  # (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)  # (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(  # (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(  # (b, t, l, l, n_token_type)
+                -1,
+                -1,
+                -1,
+                residue_type.size(2),
+                -1,
+            )
+            residue_j = residue_type[:, :, None, :].expand(  # (b, t, l, l, n_token_type)
+                -1,
+                -1,
+                residue_type.size(2),
+                -1,
+                -1,
+            )
+            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,
+        z: Tensor,
+        feats: dict[str, Tensor],
+        pair_mask: Tensor,
+        use_kernels: bool = False,
+    ) -> Tensor:
+        """Aggregate template pair tensor V into trunk update U."""
+
+        # 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()  # (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)  # (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  # (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(  # (b * t, l, l, d_t)
+            template_states,
+            expanded_mask,
+            use_kernels=use_kernels,
+        )
+        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]  # (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)))  # (b, l, l, d_z)
+
+
+class TemplateModule(_TemplateBase):
+    """Aggregate templates while restricting features to the same chain."""
+
+    def __init__(
+        self,
+        token_z: int,
+        template_dim: int,
+        template_blocks: int,
+        dropout: float = 0.25,
+        pairwise_head_width: int = 32,
+        pairwise_num_heads: int = 4,
+        post_layer_norm: bool = False,
+        activation_checkpointing: bool = False,
+        min_dist: float = 3.25,
+        max_dist: float = 50.75,
+        num_bins: int = 38,
+        **kwargs: Any,
+    ) -> None:
+        del kwargs
+        super().__init__(
+            token_z,
+            template_dim,
+            template_blocks,
+            dropout,
+            pairwise_head_width,
+            pairwise_num_heads,
+            post_layer_norm,
+            activation_checkpointing,
+            min_dist,
+            max_dist,
+            num_bins,
+        )
+
+    def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
+        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):
+    """Aggregate templates under per-template visibility groups."""
+
+    def __init__(
+        self,
+        token_z: int,
+        template_dim: int,
+        template_blocks: int,
+        dropout: float = 0.25,
+        pairwise_head_width: int = 32,
+        pairwise_num_heads: int = 4,
+        post_layer_norm: bool = False,
+        activation_checkpointing: bool = False,
+        min_dist: float = 3.25,
+        max_dist: float = 50.75,
+        num_bins: int = 38,
+        **kwargs: Any,
+    ) -> None:
+        del kwargs
+        super().__init__(
+            token_z,
+            template_dim,
+            template_blocks,
+            dropout,
+            pairwise_head_width,
+            pairwise_num_heads,
+            post_layer_norm,
+            activation_checkpointing,
+            min_dist,
+            max_dist,
+            num_bins,
+        )
+
+    def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
+        del count
+        visibility = feats["visibility_ids"]  # (b, t, l)
+        return (visibility[..., :, None] == visibility[..., None, :]).float()  # (b, t, l, l)
+
+
+class MSAModule(nn.Module):
+    """Embed and update an MSA before returning its accumulated pair update."""
+
+    def __init__(
+        self,
+        msa_s: int,
+        token_z: int,
+        token_s: int,
+        msa_blocks: int,
+        msa_dropout: float,
+        z_dropout: float,
+        pairwise_head_width: int = 32,
+        pairwise_num_heads: int = 4,
+        activation_checkpointing: bool = False,
+        use_paired_feature: bool = True,
+        subsample_msa: bool = False,
+        num_subsampled_msa: int = 1024,
+        **kwargs: Any,
+    ) -> None:
+        del kwargs
+        super().__init__()
+        self.msa_blocks = msa_blocks
+        self.msa_dropout = msa_dropout
+        self.z_dropout = z_dropout
+        self.use_paired_feature = use_paired_feature
+        self.activation_checkpointing = activation_checkpointing
+        self.subsample_msa = subsample_msa
+        self.num_subsampled_msa = num_subsampled_msa
+        self.s_proj = nn.Linear(token_s, msa_s, bias=False)
+        input_width = const.num_tokens + 2 + int(use_paired_feature)
+        self.msa_proj = nn.Linear(input_width, msa_s, bias=False)
+        self.layers = nn.ModuleList(
+            [
+                MSALayer(
+                    msa_s,
+                    token_z,
+                    msa_dropout,
+                    z_dropout,
+                    pairwise_head_width,
+                    pairwise_num_heads,
+                )
+                for _ in range(msa_blocks)
+            ]
+        )
+
+    @staticmethod
+    def _chunk_configuration(
+        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:
+            return True, 64, 32, 4, 128
+        return False, None, None, None, 512
+
+    def forward(
+        self,
+        z: Tensor,
+        emb: Tensor,
+        feats: dict[str, Tensor],
+        use_kernels: bool = False,
+    ) -> Tensor:
+        """Return updated pair tensor Z after every MSA block."""
+
+        chunking = self._chunk_configuration(z, self.training)
+        # 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()  # (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),  # (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))  # (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]  # (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,
+                pair_mask,
+                msa_mask,
+                *chunking,
+                use_kernels,
+            )
+            if self.activation_checkpointing and self.training:
+                pair_states, msa_states = checkpoint(  # (b, l, l, d_z), (b, s, l, d_m)
+                    layer,
+                    *arguments,
+                )
+            else:
+                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):
+    """Exchange information between MSA tensor M and pair tensor Z."""
+
+    def __init__(
+        self,
+        msa_s: int,
+        token_z: int,
+        msa_dropout: float,
+        z_dropout: float,
+        pairwise_head_width: int = 32,
+        pairwise_num_heads: int = 4,
+    ) -> None:
+        super().__init__()
+        self.msa_dropout = msa_dropout
+        self.msa_transition = Transition(msa_s, msa_s * 4)
+        self.pair_weighted_averaging = PairWeightedAveraging(
+            c_m=msa_s,
+            c_z=token_z,
+            c_h=32,
+            num_heads=8,
+        )
+        self.pairformer_layer = PairformerNoSeqLayer(
+            token_z=token_z,
+            dropout=z_dropout,
+            pairwise_head_width=pairwise_head_width,
+            pairwise_num_heads=pairwise_num_heads,
+        )
+        self.outer_product_mean = OuterProductMean(msa_s, 32, token_z)
+
+    def forward(
+        self,
+        z: Tensor,
+        m: Tensor,
+        token_mask: Tensor,
+        msa_mask: Tensor,
+        chunk_heads_pwa: bool = False,
+        chunk_size_transition_z: int | None = None,
+        chunk_size_transition_msa: int | None = None,
+        chunk_size_outer_product: int | None = None,
+        chunk_size_tri_attn: int | None = None,
+        use_kernels: bool = False,
+    ) -> tuple[Tensor, Tensor]:
+        """Return updated Z and M tensors."""
+
+        del chunk_size_transition_z
+        # 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(  # (b, s, l, d_m)
+            msa_states,
+            chunk_size_transition_msa,
+        )
+        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(  # (b, l, l, d_z)
+            pair_states,
+            token_mask,
+            chunk_size_tri_attn,
+            use_kernels=use_kernels,
+        )
+        return pair_states, msa_states  # (b, l, l, d_z), (b, s, l, d_m)
+
+
+class BFactorModule(nn.Module):
+    """Predict a per-token B-factor histogram."""
+
+    def __init__(self, token_s: int, num_bins: int) -> None:
+        super().__init__()
+        self.bfactor = nn.Linear(token_s, num_bins)
+        self.num_bins = num_bins
+
+    def forward(self, s: Tensor) -> Tensor:
+        # s: (..., d_s).
+        return cast(Tensor, self.bfactor(s))  # (..., n_bin)
+
+
+class DistogramModule(nn.Module):
+    """Predict symmetric residue-pair distance histograms."""
+
+    def __init__(self, token_z: int, num_bins: int, num_distograms: int = 1) -> None:
+        super().__init__()
+        self.distogram = nn.Linear(token_z, num_distograms * num_bins)
+        self.num_distograms = num_distograms
+        self.num_bins = num_bins
+
+    def forward(self, z: Tensor) -> Tensor:
+        # 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(
+                symmetric.shape[0],
+                symmetric.shape[1],
+                symmetric.shape[2],
+                self.num_distograms,
+                self.num_bins,
+            ),
+        )  # (b, l, l, n_distogram, n_bin)

fastplms/models/boltz/vb_modules_utils.py

@@ -1,283 +1,283 @@
-"""Small geometry, augmentation, activation, and EMA utilities for Boltz2.
-
-The quaternion behavior follows PyTorch3D's public rotation convention. License
-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
-from torch import Tensor, nn
-from torch.nn import functional as F
-
-
-LinearNoBias = partial(nn.Linear, bias=False)
-
-
-def exists(value: object) -> bool:
-    """Return whether a value is not ``None``."""
-
-    return value is not None
-
-
-def default(value: Any, fallback: Any) -> Any:
-    """Return ``value`` unless it is ``None``."""
-
-    return fallback if value is None else value
-
-
-def log(tensor: Tensor, eps: float = 1e-20) -> Tensor:
-    """Compute a finite logarithm by applying a scalar lower bound."""
-
-    # 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:
-        # 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:
-    # 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,
-    )
-
-
-def center(atom_coords: Tensor, atom_mask: Tensor) -> Tensor:
-    """Center coordinate tensor X with shape ``(b, n_atoms, 3)``."""
-
-    # 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."""
-
-    # 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)``."""
-
-    # 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),
-        scale * (i_axis * j_axis + k_axis * real),
-        1 - scale * (i_axis * i_axis + k_axis * k_axis),
-        scale * (j_axis * k_axis - i_axis * real),
-        scale * (i_axis * k_axis - j_axis * real),
-        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))
-
-
-def random_quaternions(
-    n: int,
-    dtype: torch.dtype | None = None,
-    device: torch.device | str | None = None,
-) -> Tensor:
-    """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)  # (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(
-    n: int,
-    dtype: torch.dtype | None = None,
-    device: torch.device | str | None = None,
-) -> Tensor:
-    """Draw rotation tensor R with shape ``(n, 3, 3)``."""
-
-    # (n, 3, 3)
-    return quaternion_to_matrix(random_quaternions(n, dtype=dtype, device=device))
-
-
-def compute_random_augmentation(
-    multiplicity: int,
-    s_trans: float = 1.0,
-    device: torch.device | str | None = None,
-    dtype: torch.dtype = torch.float32,
-) -> tuple[Tensor, Tensor]:
-    """Draw independent rotations R and translations T for a replicated batch."""
-
-    rotations = random_rotations(multiplicity, dtype=dtype, device=device)  # (m, 3, 3)
-    translations = (  # (m, 1, 3)
-        torch.randn(
-            (multiplicity, 1, 3),
-            dtype=dtype,
-            device=device,
-        )
-        * s_trans
-    )
-    return rotations, translations  # (m, 3, 3), (m, 1, 3)
-
-
-def randomly_rotate(
-    coords: Tensor,
-    return_second_coords: bool = False,
-    second_coords: Tensor | None = None,
-) -> Tensor | tuple[Tensor, Tensor | None]:
-    """Rotate X and optionally Y using the same sampled rotation tensor R."""
-
-    # 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  # (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  # (b, n, 3), (b, n_second, 3) or None
-
-
-def center_random_augmentation(
-    atom_coords: Tensor,
-    atom_mask: Tensor,
-    s_trans: float = 1.0,
-    augmentation: bool = True,
-    centering: bool = True,
-    return_second_coords: bool = False,
-    second_coords: Tensor | None = None,
-) -> Tensor | tuple[Tensor, Tensor | None]:
-    """Center and rigidly augment coordinate tensors X and optional Y."""
-
-    # 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)  # (b, 1, 3)
-        primary = primary - centroid  # (b, n_atom, 3)
-        if secondary is not None:
-            secondary = secondary - centroid  # (b, n_second, 3)
-
-    if augmentation:
-        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  # (b, 1, 3)
-        primary = primary + translation  # (b, n_atom, 3)
-        if secondary is not None:
-            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
-
-
-class ExponentialMovingAverage:
-    """Maintain detached exponential moving averages of trainable parameters."""
-
-    def __init__(
-        self,
-        parameters: Iterable[nn.Parameter],
-        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 = [  # 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:
-        # 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
-        if self.num_updates is not None:
-            self.num_updates += 1
-            decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))
-        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))  # (...) 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:
-        # 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)  # (...) 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:
-        # 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)  # (...) 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,
-            "shadow_params": self.shadow_params,
-        }
-
-    def load_state_dict(
-        self,
-        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]
+"""Small geometry, augmentation, activation, and EMA utilities for Boltz2.
+
+The quaternion behavior follows PyTorch3D's public rotation convention. License
+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
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+LinearNoBias = partial(nn.Linear, bias=False)
+
+
+def exists(value: object) -> bool:
+    """Return whether a value is not ``None``."""
+
+    return value is not None
+
+
+def default(value: Any, fallback: Any) -> Any:
+    """Return ``value`` unless it is ``None``."""
+
+    return fallback if value is None else value
+
+
+def log(tensor: Tensor, eps: float = 1e-20) -> Tensor:
+    """Compute a finite logarithm by applying a scalar lower bound."""
+
+    # 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:
+        # 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:
+    # 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,
+    )
+
+
+def center(atom_coords: Tensor, atom_mask: Tensor) -> Tensor:
+    """Center coordinate tensor X with shape ``(b, n_atoms, 3)``."""
+
+    # 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."""
+
+    # 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)``."""
+
+    # 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),
+        scale * (i_axis * j_axis + k_axis * real),
+        1 - scale * (i_axis * i_axis + k_axis * k_axis),
+        scale * (j_axis * k_axis - i_axis * real),
+        scale * (i_axis * k_axis - j_axis * real),
+        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))
+
+
+def random_quaternions(
+    n: int,
+    dtype: torch.dtype | None = None,
+    device: torch.device | str | None = None,
+) -> Tensor:
+    """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)  # (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(
+    n: int,
+    dtype: torch.dtype | None = None,
+    device: torch.device | str | None = None,
+) -> Tensor:
+    """Draw rotation tensor R with shape ``(n, 3, 3)``."""
+
+    # (n, 3, 3)
+    return quaternion_to_matrix(random_quaternions(n, dtype=dtype, device=device))
+
+
+def compute_random_augmentation(
+    multiplicity: int,
+    s_trans: float = 1.0,
+    device: torch.device | str | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[Tensor, Tensor]:
+    """Draw independent rotations R and translations T for a replicated batch."""
+
+    rotations = random_rotations(multiplicity, dtype=dtype, device=device)  # (m, 3, 3)
+    translations = (  # (m, 1, 3)
+        torch.randn(
+            (multiplicity, 1, 3),
+            dtype=dtype,
+            device=device,
+        )
+        * s_trans
+    )
+    return rotations, translations  # (m, 3, 3), (m, 1, 3)
+
+
+def randomly_rotate(
+    coords: Tensor,
+    return_second_coords: bool = False,
+    second_coords: Tensor | None = None,
+) -> Tensor | tuple[Tensor, Tensor | None]:
+    """Rotate X and optionally Y using the same sampled rotation tensor R."""
+
+    # 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  # (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  # (b, n, 3), (b, n_second, 3) or None
+
+
+def center_random_augmentation(
+    atom_coords: Tensor,
+    atom_mask: Tensor,
+    s_trans: float = 1.0,
+    augmentation: bool = True,
+    centering: bool = True,
+    return_second_coords: bool = False,
+    second_coords: Tensor | None = None,
+) -> Tensor | tuple[Tensor, Tensor | None]:
+    """Center and rigidly augment coordinate tensors X and optional Y."""
+
+    # 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)  # (b, 1, 3)
+        primary = primary - centroid  # (b, n_atom, 3)
+        if secondary is not None:
+            secondary = secondary - centroid  # (b, n_second, 3)
+
+    if augmentation:
+        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  # (b, 1, 3)
+        primary = primary + translation  # (b, n_atom, 3)
+        if secondary is not None:
+            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
+
+
+class ExponentialMovingAverage:
+    """Maintain detached exponential moving averages of trainable parameters."""
+
+    def __init__(
+        self,
+        parameters: Iterable[nn.Parameter],
+        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 = [  # 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:
+        # 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
+        if self.num_updates is not None:
+            self.num_updates += 1
+            decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))
+        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))  # (...) 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:
+        # 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)  # (...) 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:
+        # 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)  # (...) 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,
+            "shadow_params": self.shadow_params,
+        }
+
+    def load_state_dict(
+        self,
+        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

@@ -1,896 +1,896 @@
-"""Differentiable structure-steering potentials for Boltz2 diffusion.
-
-Each potential follows the same mechanism: derive constraint indices and
-bounds from features, evaluate a geometric variable, map that variable to an
-energy, and optionally scatter its analytic derivative back to atom
-coordinates.  This module is maintained independently from the pinned parity
-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
-
-from . import vb_const as const
-from .vb_loss_diffusionv2 import weighted_rigid_align
-from .vb_potentials_schedules import (
-    ExponentialInterpolation,
-    ParameterSchedule,
-    PiecewiseStepFunction,
-)
-
-
-Parameter = ParameterSchedule | float | int | bool
-ParameterMap = dict[str, Parameter]
-
-
-@dataclass(frozen=True, slots=True)
-class _PreparedPotential:
-    coordinates: torch.Tensor  # (..., a, 3)
-    index: torch.Tensor  # (q, n)
-    function_args: tuple[Any, ...]
-    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(
-    coordinates: torch.Tensor,
-    com_index: torch.Tensor,
-    atom_pad_mask: torch.Tensor,
-) -> torch.Tensor:
-    # 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(
-    energy: torch.Tensor,
-    union_index: torch.Tensor,
-    union_lambda: float,
-) -> torch.Tensor:
-    # 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,
-    ).scatter_reduce(
-        -1,
-        union_index.expand_as(unnormalized),
-        unnormalized,
-        "sum",
-    )  # (..., n_union)
-    weights = unnormalized / partition[..., union_index]  # (..., n)
-    weights[partition[..., union_index] == 0] = 0  # (..., n)
-    return weights  # (..., n)
-
-
-def _scatter_constraint_gradients(
-    coefficients: torch.Tensor,
-    variable_gradient: torch.Tensor,
-    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)  # (..., q * n, 3)
-    if product.dim() > 3:
-        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):
-    """Base contract for energy and analytic-gradient steering potentials."""
-
-    def __init__(self, parameters: ParameterMap | None = None) -> None:
-        self.parameters = parameters
-
-    def _prepare(
-        self,
-        coordinates: torch.Tensor,
-        feats: dict[str, torch.Tensor],
-        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: (q, n)
-        com_index = atom_pad_mask = None
-        if com_args is not None:
-            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, :]  # (..., n, 3)
-
-        negation_mask = union_index = None
-        if operator_args is not None:
-            negation_mask, union_index = operator_args  # each (n,)
-        return _PreparedPotential(
-            coordinates=coordinates,
-            index=index,
-            function_args=args,
-            com_index=com_index,
-            atom_pad_mask=atom_pad_mask,
-            ref_coords=ref_coords,
-            ref_mask=ref_mask,
-            ref_token_index=ref_token_index,
-            negation_mask=negation_mask,
-            union_index=union_index,
-        )
-
-    def compute(
-        self,
-        coords: torch.Tensor,
-        feats: dict[str, torch.Tensor],
-        parameters: dict[str, Any],
-    ) -> 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)  # (...)
-        prepared = self._prepare(coords, feats, parameters, computed_args)
-        value = self.compute_variable(
-            prepared.coordinates,
-            prepared.index,
-            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"],
-            )  # (..., n)
-            return (energy * weights).sum(dim=-1)  # (...)
-        return energy.sum(dim=tuple(range(1, energy.dim())))  # (b,)
-
-    def compute_gradient(
-        self,
-        coords: torch.Tensor,
-        feats: dict[str, torch.Tensor],
-        parameters: dict[str, Any],
-    ) -> 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)  # (b, a, 3)
-        prepared = self._prepare(coords, feats, parameters, computed_args)
-        value, variable_gradient = self.compute_variable(
-            prepared.coordinates,
-            prepared.index,
-            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,
-            ).scatter_reduce(
-                -1,
-                prepared.union_index.expand_as(energy),
-                energy * weights,
-                "sum",
-            )  # (..., n_union)
-            coefficients = (
-                energy_derivative
-                * weights
-                * (
-                    1
-                    + parameters["union_lambda"]
-                    * (energy - union_energy[..., prepared.union_index])
-                )
-            )  # (..., n)
-        else:
-            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, :]  # (b, a, 3)
-        elif prepared.ref_token_index is not None:
-            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``."""
-
-        if self.parameters is None:
-            return None
-        return {
-            name: parameter.compute(t) if isinstance(parameter, ParameterSchedule) else parameter
-            for name, parameter in self.parameters.items()
-        }
-
-    @abstractmethod
-    def compute_function(
-        self,
-        value: torch.Tensor,
-        *args: Any,
-        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
-    def compute_variable(
-        self,
-        coords: torch.Tensor,
-        index: torch.Tensor,
-        ref_coords: torch.Tensor | None = None,
-        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
-    def compute_args(
-        self,
-        feats: dict[str, torch.Tensor],
-        parameters: dict[str, Any],
-    ) -> tuple[Any, ...]:
-        raise NotImplementedError
-
-    def get_reference_coords(
-        self,
-        feats: dict[str, torch.Tensor],
-        parameters: dict[str, Any],
-    ) -> tuple[None, None]:
-        del feats, parameters
-        return None, None
-
-
-class FlatBottomPotential(Potential):
-    """Linear penalty outside an allowed lower-to-upper interval."""
-
-    def compute_function(
-        self,
-        value: torch.Tensor,
-        k: torch.Tensor,
-        lower_bounds: torch.Tensor | None,
-        upper_bounds: torch.Tensor | None,
-        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):
-                raise TypeError(
-                    f"negation_mask must be a boolean tensor, got {type(negation_mask).__name__}."
-                )
-            if negation_mask.dtype != torch.bool:
-                raise TypeError(
-                    f"negation_mask must be a boolean tensor, got {negation_mask.dtype}."
-                )
-            if negation_mask.device != value.device:
-                raise ValueError(
-                    "negation_mask must be on the same device as value; "
-                    f"got {negation_mask.device} and {value.device}."
-                )
-            try:
-                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)  # (..., 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  # (..., 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  # (..., 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):
-    """Measure atom displacement after weighted rigid alignment."""
-
-    def compute_variable(
-        self,
-        coords: torch.Tensor,
-        index: torch.Tensor,
-        ref_coords: torch.Tensor,
-        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,
-        )  # (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  # (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  # (b, n_ref, n); (b, 1, n_ref, n, 3)
-
-
-class DistancePotential(Potential):
-    """Measure Euclidean distances for indexed atom pairs."""
-
-    def compute_variable(
-        self,
-        coords: torch.Tensor,
-        index: torch.Tensor,
-        ref_coords: torch.Tensor | None = None,
-        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],
-        )  # (..., n, 3)
-        distance = torch.linalg.norm(displacement, dim=-1)  # (..., n)
-        if not compute_gradient:
-            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)
-
-
-class DihedralPotential(Potential):
-    """Measure signed torsion angles for indexed atom quartets."""
-
-    def compute_variable(
-        self,
-        coords: torch.Tensor,
-        index: torch.Tensor,
-        ref_coords: torch.Tensor | None = None,
-        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, 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)  # (..., 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  # (..., n)
-
-        projection_i = (
-            (r_ij.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
-        ).unsqueeze(-1)  # (..., n, 1)
-        projection_l = (
-            (r_kl.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
-        ).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)
-
-
-class AbsDihedralPotential(DihedralPotential):
-    """Measure the unsigned magnitude of indexed torsion angles."""
-
-    def compute_variable(
-        self,
-        coords: torch.Tensor,
-        index: torch.Tensor,
-        ref_coords: torch.Tensor | None = None,
-        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))  # (..., n)
-        angle, gradient = super().compute_variable(
-            coords,
-            index,
-            compute_gradient=True,
-        )  # (..., 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(),
-            feats["asym_id"].unsqueeze(-1).float(),
-        )
-        .squeeze(-1)
-        .long()
-    )[0]  # (a,)
-
-
-class PoseBustersPotential(FlatBottomPotential, DistancePotential):
-    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]) -> tuple[Any, ...]:
-        # 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  # (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]) -> tuple[Any, ...]:
-        # 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]) -> tuple[Any, ...]:
-        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,
-        )  # (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]  # (2, n_connection)
-        connected_matrix = torch.eye(
-            num_chains,
-            device=atom_chain_id.device,
-            dtype=torch.bool,
-        )  # (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]) -> tuple[Any, ...]:
-        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),
-            (atom_chain_id, atom_pad_mask),
-            None,
-            None,
-        )
-
-
-def _oriented_bounds(
-    orientations: torch.Tensor,
-    positive_lower: float,
-    negative_upper: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    # 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]) -> tuple[Any, ...]:
-        # 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]) -> tuple[Any, ...]:
-        # 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]) -> tuple[Any, ...]:
-        # 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  # (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),
-            None,
-            None,
-            None,
-        )
-
-
-class TemplateReferencePotential(FlatBottomPotential, ReferencePotential):
-    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]) -> tuple[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_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()  # (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]  # (1, a, 1)
-        ref_atom_index = (
-            torch.bmm(
-                feats["token_to_rep_atom"].float(),
-                atom_indices,
-            )
-            .squeeze(-1)
-            .long()[0]  # (n,)
-        )
-        ref_token_index = (
-            torch.bmm(
-                feats["atom_to_token"].float(),
-                feats["token_index"].unsqueeze(-1).float(),
-            )
-            .squeeze(-1)
-            .long()[0]  # (a,)
-        )
-
-        index = torch.arange(
-            template_mask.shape[-1],
-            dtype=torch.long,
-            device=template_mask.device,
-        )[None]  # (1, n)
-        upper = torch.full(
-            template_mask.shape,
-            float("inf"),
-            device=index.device,
-            dtype=torch.float32,
-        )  # (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]]  # (n_ref, n)
-        # Returns index (1, n), function args (n_ref, n), and reference mappings.
-        return (
-            index,
-            (torch.ones_like(upper), None, upper),
-            None,
-            (ref_coords, ref_mask, ref_atom_index, ref_token_index),
-            None,
-        )
-
-
-class ContactPotentital(FlatBottomPotential, DistancePotential):
-    """Contact-union potential retaining the historical checkpoint name."""
-
-    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]) -> tuple[Any, ...]:
-        del parameters
-        # 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),
-            None,
-            None,
-            (
-                feats["contact_negation_mask"][0],
-                feats["contact_union_index"][0],
-            ),
-        )
-
-
-ContactPotential = ContactPotentital
-
-
-def get_potentials(
-    steering_args: dict[str, bool],
-    boltz2: bool = False,
-) -> list[Potential]:
-    """Build the ordered potential stack for requested steering modes."""
-
-    use_fk = steering_args["fk_steering"]
-    use_physical = steering_args["physical_guidance_update"]
-    use_contacts = steering_args.get("contact_guidance_update", False)
-    potentials: list[Potential] = []
-    if use_fk or use_physical:
-        potentials.extend(
-            (
-                SymmetricChainCOMPotential(
-                    {
-                        "guidance_interval": 4,
-                        "guidance_weight": 0.5 if use_physical else 0.0,
-                        "resampling_weight": 0.5,
-                        "buffer": ExponentialInterpolation(1.0, 5.0, -2.0),
-                    }
-                ),
-                VDWOverlapPotential(
-                    {
-                        "guidance_interval": 5,
-                        "guidance_weight": PiecewiseStepFunction(
-                            [0.4],
-                            [0.125, 0.0],
-                        )
-                        if use_physical
-                        else 0.0,
-                        "resampling_weight": PiecewiseStepFunction(
-                            [0.6],
-                            [0.01, 0.0],
-                        ),
-                        "buffer": 0.225,
-                    }
-                ),
-                ConnectionsPotential(
-                    {
-                        "guidance_interval": 1,
-                        "guidance_weight": 0.15 if use_physical else 0.0,
-                        "resampling_weight": 1.0,
-                        "buffer": 2.0,
-                    }
-                ),
-                PoseBustersPotential(
-                    {
-                        "guidance_interval": 1,
-                        "guidance_weight": 0.01 if use_physical else 0.0,
-                        "resampling_weight": 0.1,
-                        "bond_buffer": 0.125,
-                        "angle_buffer": 0.125,
-                        "clash_buffer": 0.10,
-                    }
-                ),
-                ChiralAtomPotential(
-                    {
-                        "guidance_interval": 1,
-                        "guidance_weight": 0.1 if use_physical else 0.0,
-                        "resampling_weight": 1.0,
-                        "buffer": 0.52360,
-                    }
-                ),
-                StereoBondPotential(
-                    {
-                        "guidance_interval": 1,
-                        "guidance_weight": 0.05 if use_physical else 0.0,
-                        "resampling_weight": 1.0,
-                        "buffer": 0.52360,
-                    }
-                ),
-                PlanarBondPotential(
-                    {
-                        "guidance_interval": 1,
-                        "guidance_weight": 0.05 if use_physical else 0.0,
-                        "resampling_weight": 1.0,
-                        "buffer": 0.26180,
-                    }
-                ),
-            )
-        )
-    if boltz2 and (use_fk or use_contacts):
-        potentials.extend(
-            (
-                ContactPotentital(
-                    {
-                        "guidance_interval": 4,
-                        "guidance_weight": PiecewiseStepFunction(
-                            [0.25, 0.75],
-                            [0.0, 0.5, 1.0],
-                        )
-                        if use_contacts
-                        else 0.0,
-                        "resampling_weight": 1.0,
-                        "union_lambda": ExponentialInterpolation(8.0, 0.0, -2.0),
-                    }
-                ),
-                TemplateReferencePotential(
-                    {
-                        "guidance_interval": 2,
-                        "guidance_weight": 0.1 if use_contacts else 0.0,
-                        "resampling_weight": 1.0,
-                    }
-                ),
-            )
-        )
-    return potentials
+"""Differentiable structure-steering potentials for Boltz2 diffusion.
+
+Each potential follows the same mechanism: derive constraint indices and
+bounds from features, evaluate a geometric variable, map that variable to an
+energy, and optionally scatter its analytic derivative back to atom
+coordinates.  This module is maintained independently from the pinned parity
+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
+
+from . import vb_const as const
+from .vb_loss_diffusionv2 import weighted_rigid_align
+from .vb_potentials_schedules import (
+    ExponentialInterpolation,
+    ParameterSchedule,
+    PiecewiseStepFunction,
+)
+
+
+Parameter = ParameterSchedule | float | int | bool
+ParameterMap = dict[str, Parameter]
+
+
+@dataclass(frozen=True, slots=True)
+class _PreparedPotential:
+    coordinates: torch.Tensor  # (..., a, 3)
+    index: torch.Tensor  # (q, n)
+    function_args: tuple[Any, ...]
+    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(
+    coordinates: torch.Tensor,
+    com_index: torch.Tensor,
+    atom_pad_mask: torch.Tensor,
+) -> torch.Tensor:
+    # 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(
+    energy: torch.Tensor,
+    union_index: torch.Tensor,
+    union_lambda: float,
+) -> torch.Tensor:
+    # 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,
+    ).scatter_reduce(
+        -1,
+        union_index.expand_as(unnormalized),
+        unnormalized,
+        "sum",
+    )  # (..., n_union)
+    weights = unnormalized / partition[..., union_index]  # (..., n)
+    weights[partition[..., union_index] == 0] = 0  # (..., n)
+    return weights  # (..., n)
+
+
+def _scatter_constraint_gradients(
+    coefficients: torch.Tensor,
+    variable_gradient: torch.Tensor,
+    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)  # (..., q * n, 3)
+    if product.dim() > 3:
+        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):
+    """Base contract for energy and analytic-gradient steering potentials."""
+
+    def __init__(self, parameters: ParameterMap | None = None) -> None:
+        self.parameters = parameters
+
+    def _prepare(
+        self,
+        coordinates: torch.Tensor,
+        feats: dict[str, torch.Tensor],
+        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: (q, n)
+        com_index = atom_pad_mask = None
+        if com_args is not None:
+            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, :]  # (..., n, 3)
+
+        negation_mask = union_index = None
+        if operator_args is not None:
+            negation_mask, union_index = operator_args  # each (n,)
+        return _PreparedPotential(
+            coordinates=coordinates,
+            index=index,
+            function_args=args,
+            com_index=com_index,
+            atom_pad_mask=atom_pad_mask,
+            ref_coords=ref_coords,
+            ref_mask=ref_mask,
+            ref_token_index=ref_token_index,
+            negation_mask=negation_mask,
+            union_index=union_index,
+        )
+
+    def compute(
+        self,
+        coords: torch.Tensor,
+        feats: dict[str, torch.Tensor],
+        parameters: dict[str, Any],
+    ) -> 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)  # (...)
+        prepared = self._prepare(coords, feats, parameters, computed_args)
+        value = self.compute_variable(
+            prepared.coordinates,
+            prepared.index,
+            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"],
+            )  # (..., n)
+            return (energy * weights).sum(dim=-1)  # (...)
+        return energy.sum(dim=tuple(range(1, energy.dim())))  # (b,)
+
+    def compute_gradient(
+        self,
+        coords: torch.Tensor,
+        feats: dict[str, torch.Tensor],
+        parameters: dict[str, Any],
+    ) -> 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)  # (b, a, 3)
+        prepared = self._prepare(coords, feats, parameters, computed_args)
+        value, variable_gradient = self.compute_variable(
+            prepared.coordinates,
+            prepared.index,
+            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,
+            ).scatter_reduce(
+                -1,
+                prepared.union_index.expand_as(energy),
+                energy * weights,
+                "sum",
+            )  # (..., n_union)
+            coefficients = (
+                energy_derivative
+                * weights
+                * (
+                    1
+                    + parameters["union_lambda"]
+                    * (energy - union_energy[..., prepared.union_index])
+                )
+            )  # (..., n)
+        else:
+            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, :]  # (b, a, 3)
+        elif prepared.ref_token_index is not None:
+            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``."""
+
+        if self.parameters is None:
+            return None
+        return {
+            name: parameter.compute(t) if isinstance(parameter, ParameterSchedule) else parameter
+            for name, parameter in self.parameters.items()
+        }
+
+    @abstractmethod
+    def compute_function(
+        self,
+        value: torch.Tensor,
+        *args: Any,
+        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
+    def compute_variable(
+        self,
+        coords: torch.Tensor,
+        index: torch.Tensor,
+        ref_coords: torch.Tensor | None = None,
+        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
+    def compute_args(
+        self,
+        feats: dict[str, torch.Tensor],
+        parameters: dict[str, Any],
+    ) -> tuple[Any, ...]:
+        raise NotImplementedError
+
+    def get_reference_coords(
+        self,
+        feats: dict[str, torch.Tensor],
+        parameters: dict[str, Any],
+    ) -> tuple[None, None]:
+        del feats, parameters
+        return None, None
+
+
+class FlatBottomPotential(Potential):
+    """Linear penalty outside an allowed lower-to-upper interval."""
+
+    def compute_function(
+        self,
+        value: torch.Tensor,
+        k: torch.Tensor,
+        lower_bounds: torch.Tensor | None,
+        upper_bounds: torch.Tensor | None,
+        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):
+                raise TypeError(
+                    f"negation_mask must be a boolean tensor, got {type(negation_mask).__name__}."
+                )
+            if negation_mask.dtype != torch.bool:
+                raise TypeError(
+                    f"negation_mask must be a boolean tensor, got {negation_mask.dtype}."
+                )
+            if negation_mask.device != value.device:
+                raise ValueError(
+                    "negation_mask must be on the same device as value; "
+                    f"got {negation_mask.device} and {value.device}."
+                )
+            try:
+                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)  # (..., 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  # (..., 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  # (..., 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):
+    """Measure atom displacement after weighted rigid alignment."""
+
+    def compute_variable(
+        self,
+        coords: torch.Tensor,
+        index: torch.Tensor,
+        ref_coords: torch.Tensor,
+        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,
+        )  # (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  # (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  # (b, n_ref, n); (b, 1, n_ref, n, 3)
+
+
+class DistancePotential(Potential):
+    """Measure Euclidean distances for indexed atom pairs."""
+
+    def compute_variable(
+        self,
+        coords: torch.Tensor,
+        index: torch.Tensor,
+        ref_coords: torch.Tensor | None = None,
+        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],
+        )  # (..., n, 3)
+        distance = torch.linalg.norm(displacement, dim=-1)  # (..., n)
+        if not compute_gradient:
+            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)
+
+
+class DihedralPotential(Potential):
+    """Measure signed torsion angles for indexed atom quartets."""
+
+    def compute_variable(
+        self,
+        coords: torch.Tensor,
+        index: torch.Tensor,
+        ref_coords: torch.Tensor | None = None,
+        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, 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)  # (..., 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  # (..., n)
+
+        projection_i = (
+            (r_ij.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
+        ).unsqueeze(-1)  # (..., n, 1)
+        projection_l = (
+            (r_kl.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
+        ).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)
+
+
+class AbsDihedralPotential(DihedralPotential):
+    """Measure the unsigned magnitude of indexed torsion angles."""
+
+    def compute_variable(
+        self,
+        coords: torch.Tensor,
+        index: torch.Tensor,
+        ref_coords: torch.Tensor | None = None,
+        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))  # (..., n)
+        angle, gradient = super().compute_variable(
+            coords,
+            index,
+            compute_gradient=True,
+        )  # (..., 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(),
+            feats["asym_id"].unsqueeze(-1).float(),
+        )
+        .squeeze(-1)
+        .long()
+    )[0]  # (a,)
+
+
+class PoseBustersPotential(FlatBottomPotential, DistancePotential):
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]) -> tuple[Any, ...]:
+        # 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  # (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]) -> tuple[Any, ...]:
+        # 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]) -> tuple[Any, ...]:
+        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,
+        )  # (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]  # (2, n_connection)
+        connected_matrix = torch.eye(
+            num_chains,
+            device=atom_chain_id.device,
+            dtype=torch.bool,
+        )  # (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]) -> tuple[Any, ...]:
+        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),
+            (atom_chain_id, atom_pad_mask),
+            None,
+            None,
+        )
+
+
+def _oriented_bounds(
+    orientations: torch.Tensor,
+    positive_lower: float,
+    negative_upper: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # 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]) -> tuple[Any, ...]:
+        # 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]) -> tuple[Any, ...]:
+        # 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]) -> tuple[Any, ...]:
+        # 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  # (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),
+            None,
+            None,
+            None,
+        )
+
+
+class TemplateReferencePotential(FlatBottomPotential, ReferencePotential):
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]) -> tuple[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_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()  # (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]  # (1, a, 1)
+        ref_atom_index = (
+            torch.bmm(
+                feats["token_to_rep_atom"].float(),
+                atom_indices,
+            )
+            .squeeze(-1)
+            .long()[0]  # (n,)
+        )
+        ref_token_index = (
+            torch.bmm(
+                feats["atom_to_token"].float(),
+                feats["token_index"].unsqueeze(-1).float(),
+            )
+            .squeeze(-1)
+            .long()[0]  # (a,)
+        )
+
+        index = torch.arange(
+            template_mask.shape[-1],
+            dtype=torch.long,
+            device=template_mask.device,
+        )[None]  # (1, n)
+        upper = torch.full(
+            template_mask.shape,
+            float("inf"),
+            device=index.device,
+            dtype=torch.float32,
+        )  # (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]]  # (n_ref, n)
+        # Returns index (1, n), function args (n_ref, n), and reference mappings.
+        return (
+            index,
+            (torch.ones_like(upper), None, upper),
+            None,
+            (ref_coords, ref_mask, ref_atom_index, ref_token_index),
+            None,
+        )
+
+
+class ContactPotentital(FlatBottomPotential, DistancePotential):
+    """Contact-union potential retaining the historical checkpoint name."""
+
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]) -> tuple[Any, ...]:
+        del parameters
+        # 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),
+            None,
+            None,
+            (
+                feats["contact_negation_mask"][0],
+                feats["contact_union_index"][0],
+            ),
+        )
+
+
+ContactPotential = ContactPotentital
+
+
+def get_potentials(
+    steering_args: dict[str, bool],
+    boltz2: bool = False,
+) -> list[Potential]:
+    """Build the ordered potential stack for requested steering modes."""
+
+    use_fk = steering_args["fk_steering"]
+    use_physical = steering_args["physical_guidance_update"]
+    use_contacts = steering_args.get("contact_guidance_update", False)
+    potentials: list[Potential] = []
+    if use_fk or use_physical:
+        potentials.extend(
+            (
+                SymmetricChainCOMPotential(
+                    {
+                        "guidance_interval": 4,
+                        "guidance_weight": 0.5 if use_physical else 0.0,
+                        "resampling_weight": 0.5,
+                        "buffer": ExponentialInterpolation(1.0, 5.0, -2.0),
+                    }
+                ),
+                VDWOverlapPotential(
+                    {
+                        "guidance_interval": 5,
+                        "guidance_weight": PiecewiseStepFunction(
+                            [0.4],
+                            [0.125, 0.0],
+                        )
+                        if use_physical
+                        else 0.0,
+                        "resampling_weight": PiecewiseStepFunction(
+                            [0.6],
+                            [0.01, 0.0],
+                        ),
+                        "buffer": 0.225,
+                    }
+                ),
+                ConnectionsPotential(
+                    {
+                        "guidance_interval": 1,
+                        "guidance_weight": 0.15 if use_physical else 0.0,
+                        "resampling_weight": 1.0,
+                        "buffer": 2.0,
+                    }
+                ),
+                PoseBustersPotential(
+                    {
+                        "guidance_interval": 1,
+                        "guidance_weight": 0.01 if use_physical else 0.0,
+                        "resampling_weight": 0.1,
+                        "bond_buffer": 0.125,
+                        "angle_buffer": 0.125,
+                        "clash_buffer": 0.10,
+                    }
+                ),
+                ChiralAtomPotential(
+                    {
+                        "guidance_interval": 1,
+                        "guidance_weight": 0.1 if use_physical else 0.0,
+                        "resampling_weight": 1.0,
+                        "buffer": 0.52360,
+                    }
+                ),
+                StereoBondPotential(
+                    {
+                        "guidance_interval": 1,
+                        "guidance_weight": 0.05 if use_physical else 0.0,
+                        "resampling_weight": 1.0,
+                        "buffer": 0.52360,
+                    }
+                ),
+                PlanarBondPotential(
+                    {
+                        "guidance_interval": 1,
+                        "guidance_weight": 0.05 if use_physical else 0.0,
+                        "resampling_weight": 1.0,
+                        "buffer": 0.26180,
+                    }
+                ),
+            )
+        )
+    if boltz2 and (use_fk or use_contacts):
+        potentials.extend(
+            (
+                ContactPotentital(
+                    {
+                        "guidance_interval": 4,
+                        "guidance_weight": PiecewiseStepFunction(
+                            [0.25, 0.75],
+                            [0.0, 0.5, 1.0],
+                        )
+                        if use_contacts
+                        else 0.0,
+                        "resampling_weight": 1.0,
+                        "union_lambda": ExponentialInterpolation(8.0, 0.0, -2.0),
+                    }
+                ),
+                TemplateReferencePotential(
+                    {
+                        "guidance_interval": 2,
+                        "guidance_weight": 0.1 if use_contacts else 0.0,
+                        "resampling_weight": 1.0,
+                    }
+                ),
+            )
+        )
+    return potentials

fastplms/models/boltz/vb_potentials_schedules.py

@@ -1,68 +1,68 @@
-"""Scalar schedules for structure-steering potentials."""
-
-from __future__ import annotations
-
-import math
-
-from abc import ABC, abstractmethod
-from collections.abc import Sequence
-
-
-class ParameterSchedule(ABC):
-    """Map normalized diffusion time ``t`` to a potential parameter."""
-
-    @abstractmethod
-    def compute(self, t: float) -> float:
-        """Evaluate the schedule at ``t``."""
-
-
-class ExponentialInterpolation(ParameterSchedule):
-    """Interpolate from ``start`` to ``end`` with exponential curvature."""
-
-    def __init__(self, start: float, end: float, alpha: float) -> None:
-        self.start = start
-        self.end = end
-        self.alpha = alpha
-
-    def compute(self, t: float) -> float:
-        span = self.end - self.start
-        if self.alpha == 0:
-            return self.start + span * t
-        numerator = math.exp(self.alpha * t) - 1
-        denominator = math.exp(self.alpha) - 1
-        return self.start + span * numerator / denominator
-
-
-class PiecewiseStepFunction(ParameterSchedule):
-    """Select values separated by strict upper thresholds.
-
-    A time exactly equal to a threshold remains in the lower interval.  This
-    boundary convention is part of the steering-input contract.
-    """
-
-    def __init__(
-        self,
-        thresholds: Sequence[float],
-        values: Sequence[float],
-    ) -> None:
-        self.thresholds = tuple(thresholds)
-        self.values = tuple(values)
-        if not self.thresholds:
-            raise ValueError("PiecewiseStepFunction requires at least one threshold.")
-        if len(self.values) != len(self.thresholds) + 1:
-            raise ValueError(
-                "PiecewiseStepFunction requires exactly one more value than threshold; "
-                f"received {len(self.values)} values and {len(self.thresholds)} thresholds."
-            )
-        if any(
-            current >= following
-            for current, following in zip(self.thresholds, self.thresholds[1:], strict=False)
-        ):
-            raise ValueError("PiecewiseStepFunction thresholds must be strictly increasing.")
-
-    def compute(self, t: float) -> float:
-        interval = next(
-            (index for index, threshold in enumerate(self.thresholds) if t <= threshold),
-            len(self.thresholds),
-        )
-        return self.values[interval]
+"""Scalar schedules for structure-steering potentials."""
+
+from __future__ import annotations
+
+import math
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+
+
+class ParameterSchedule(ABC):
+    """Map normalized diffusion time ``t`` to a potential parameter."""
+
+    @abstractmethod
+    def compute(self, t: float) -> float:
+        """Evaluate the schedule at ``t``."""
+
+
+class ExponentialInterpolation(ParameterSchedule):
+    """Interpolate from ``start`` to ``end`` with exponential curvature."""
+
+    def __init__(self, start: float, end: float, alpha: float) -> None:
+        self.start = start
+        self.end = end
+        self.alpha = alpha
+
+    def compute(self, t: float) -> float:
+        span = self.end - self.start
+        if self.alpha == 0:
+            return self.start + span * t
+        numerator = math.exp(self.alpha * t) - 1
+        denominator = math.exp(self.alpha) - 1
+        return self.start + span * numerator / denominator
+
+
+class PiecewiseStepFunction(ParameterSchedule):
+    """Select values separated by strict upper thresholds.
+
+    A time exactly equal to a threshold remains in the lower interval.  This
+    boundary convention is part of the steering-input contract.
+    """
+
+    def __init__(
+        self,
+        thresholds: Sequence[float],
+        values: Sequence[float],
+    ) -> None:
+        self.thresholds = tuple(thresholds)
+        self.values = tuple(values)
+        if not self.thresholds:
+            raise ValueError("PiecewiseStepFunction requires at least one threshold.")
+        if len(self.values) != len(self.thresholds) + 1:
+            raise ValueError(
+                "PiecewiseStepFunction requires exactly one more value than threshold; "
+                f"received {len(self.values)} values and {len(self.thresholds)} thresholds."
+            )
+        if any(
+            current >= following
+            for current, following in zip(self.thresholds, self.thresholds[1:], strict=False)
+        ):
+            raise ValueError("PiecewiseStepFunction thresholds must be strictly increasing.")
+
+    def compute(self, t: float) -> float:
+        interval = next(
+            (index for index, threshold in enumerate(self.thresholds) if t <= threshold),
+            len(self.thresholds),
+        )
+        return self.values[interval]

fastplms/models/boltz/vb_tri_attn_attention.py

@@ -1,151 +1,151 @@
-"""Triangular row and column attention for Boltz2 pair states."""
-
-from __future__ import annotations
-
-import torch
-
-from typing import cast
-from torch import Tensor, nn
-
-from .vb_tri_attn_primitives import Attention, LayerNorm, Linear
-from .vb_tri_attn_utils import chunk_layer, permute_final_dims
-
-
-class TriangleAttention(nn.Module):
-    """Apply pair-biased attention along one axis of pair tensor X."""
-
-    def __init__(
-        self,
-        c_in: int,
-        c_hidden: int,
-        no_heads: int,
-        starting: bool = True,
-        inf: float = 1e9,
-    ) -> None:
-        super().__init__()
-        self.c_in = c_in
-        self.c_hidden = c_hidden
-        self.no_heads = no_heads
-        self.starting = starting
-        self.inf = inf
-        self.layer_norm = LayerNorm(c_in)
-        self.linear = Linear(c_in, no_heads, bias=False, init="normal")
-        self.mha = Attention(c_in, c_in, c_in, c_hidden, no_heads)
-
-    @torch.jit.ignore  # type: ignore[untyped-decorator]
-    def _chunk(
-        self,
-        x: Tensor,
-        tri_bias: Tensor,
-        mask_bias: Tensor,
-        mask: Tensor,
-        chunk_size: int,
-        use_kernels: bool = False,
-    ) -> 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))
-
-        inputs = {
-            "q_x": x,
-            "kv_x": x,
-            "tri_bias": tri_bias,
-            "mask_bias": mask_bias,
-            "mask": mask,
-        }  # tensor values retain the input shapes above
-        return cast(
-            Tensor,
-            chunk_layer(
-                attention_call,
-                inputs,
-                chunk_size=chunk_size,
-                no_batch_dims=len(x.shape[:-2]),
-                _out=None,
-            ),
-        )
-
-    def _run_attention(
-        self,
-        states: Tensor,
-        triangle_bias: Tensor,
-        mask_bias: Tensor,
-        expanded_mask: Tensor,
-        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,
-                self._chunk(
-                    states,
-                    triangle_bias,
-                    mask_bias,
-                    expanded_mask,
-                    chunk_size,
-                    use_kernels=False,
-                ),
-            )
-        return cast(
-            Tensor,
-            self.mha(
-                states,
-                states,
-                triangle_bias,
-                mask_bias,
-                expanded_mask,
-                use_kernels=use_kernels,
-            ),
-        )
-
-    def forward(
-        self,
-        x: Tensor,
-        mask: Tensor | None = None,
-        chunk_size: int | None = None,
-        use_kernels: bool = False,
-    ) -> 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  # (..., l, l)
-        pair_states = x  # (..., l, l, d)
-        if not self.starting:
-            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,
-            mask_bias,
-            expanded_mask,
-            chunk_size,
-            use_kernels,
-        )  # (..., l, l, d)
-        return (
-            output if self.starting else output.transpose(-2, -3)
-        )  # (..., l, l, d)
-
-
-TriangleAttentionStartingNode = TriangleAttention
-
-
-class TriangleAttentionEndingNode(TriangleAttention):
-    """Apply triangular attention around each pair's ending node."""
-
-    def __init__(
-        self,
-        c_in: int,
-        c_hidden: int,
-        no_heads: int,
-        inf: float = 1e9,
-    ) -> None:
-        super().__init__(c_in, c_hidden, no_heads, starting=False, inf=inf)
+"""Triangular row and column attention for Boltz2 pair states."""
+
+from __future__ import annotations
+
+import torch
+
+from typing import cast
+from torch import Tensor, nn
+
+from .vb_tri_attn_primitives import Attention, LayerNorm, Linear
+from .vb_tri_attn_utils import chunk_layer, permute_final_dims
+
+
+class TriangleAttention(nn.Module):
+    """Apply pair-biased attention along one axis of pair tensor X."""
+
+    def __init__(
+        self,
+        c_in: int,
+        c_hidden: int,
+        no_heads: int,
+        starting: bool = True,
+        inf: float = 1e9,
+    ) -> None:
+        super().__init__()
+        self.c_in = c_in
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.starting = starting
+        self.inf = inf
+        self.layer_norm = LayerNorm(c_in)
+        self.linear = Linear(c_in, no_heads, bias=False, init="normal")
+        self.mha = Attention(c_in, c_in, c_in, c_hidden, no_heads)
+
+    @torch.jit.ignore  # type: ignore[untyped-decorator]
+    def _chunk(
+        self,
+        x: Tensor,
+        tri_bias: Tensor,
+        mask_bias: Tensor,
+        mask: Tensor,
+        chunk_size: int,
+        use_kernels: bool = False,
+    ) -> 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))
+
+        inputs = {
+            "q_x": x,
+            "kv_x": x,
+            "tri_bias": tri_bias,
+            "mask_bias": mask_bias,
+            "mask": mask,
+        }  # tensor values retain the input shapes above
+        return cast(
+            Tensor,
+            chunk_layer(
+                attention_call,
+                inputs,
+                chunk_size=chunk_size,
+                no_batch_dims=len(x.shape[:-2]),
+                _out=None,
+            ),
+        )
+
+    def _run_attention(
+        self,
+        states: Tensor,
+        triangle_bias: Tensor,
+        mask_bias: Tensor,
+        expanded_mask: Tensor,
+        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,
+                self._chunk(
+                    states,
+                    triangle_bias,
+                    mask_bias,
+                    expanded_mask,
+                    chunk_size,
+                    use_kernels=False,
+                ),
+            )
+        return cast(
+            Tensor,
+            self.mha(
+                states,
+                states,
+                triangle_bias,
+                mask_bias,
+                expanded_mask,
+                use_kernels=use_kernels,
+            ),
+        )
+
+    def forward(
+        self,
+        x: Tensor,
+        mask: Tensor | None = None,
+        chunk_size: int | None = None,
+        use_kernels: bool = False,
+    ) -> 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  # (..., l, l)
+        pair_states = x  # (..., l, l, d)
+        if not self.starting:
+            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,
+            mask_bias,
+            expanded_mask,
+            chunk_size,
+            use_kernels,
+        )  # (..., l, l, d)
+        return (
+            output if self.starting else output.transpose(-2, -3)
+        )  # (..., l, l, d)
+
+
+TriangleAttentionStartingNode = TriangleAttention
+
+
+class TriangleAttentionEndingNode(TriangleAttention):
+    """Apply triangular attention around each pair's ending node."""
+
+    def __init__(
+        self,
+        c_in: int,
+        c_hidden: int,
+        no_heads: int,
+        inf: float = 1e9,
+    ) -> None:
+        super().__init__(c_in, c_hidden, no_heads, starting=False, inf=inf)

fastplms/models/boltz/vb_tri_attn_primitives.py

@@ -1,266 +1,266 @@
-"""Attention primitives for the Boltz pair stack.
-
-The module retains the converted checkpoint names while providing a compact
-implementation around PyTorch operations.  Optional triangle kernels are
-loaded only at the call site, so importing FastPLMs does not initialize CUDA.
-
-Portions of the numerical contract follow OpenFold and DeepMind AlphaFold
-under Apache-2.0.  See ``THIRD_PARTY_NOTICES.md`` for provenance.
-"""
-
-from __future__ import annotations
-
-import importlib
-import math
-import torch
-
-from collections.abc import Callable, Sequence
-from importlib.util import find_spec
-from torch import nn
-
-from . import vb_layers_initialize as initialize
-from .vb_tri_attn_utils import flatten_final_dims, permute_final_dims
-
-
-def _initialize_weight(
-    weight: torch.Tensor,
-    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_,
-        "glorot": initialize.glorot_uniform_init_,
-        "normal": initialize.normal_init_,
-        "final": initialize.final_init_,
-        "gating": initialize.gating_init_,
-    }
-    try:
-        initializers[method](weight)
-    except KeyError as error:
-        raise ValueError(f"unknown linear initialization {method!r}") from error
-    if method == "gating" and bias is not None:
-        bias.fill_(1.0)
-
-
-class Linear(nn.Linear):
-    """Linear projection with the initializers used by converted checkpoints."""
-
-    def __init__(
-        self,
-        in_dim: int,
-        out_dim: int,
-        bias: bool = True,
-        init: str = "default",
-        init_fn: Callable[[torch.Tensor, torch.Tensor | None], None] | None = None,
-        precision: torch.dtype | None = None,
-    ) -> None:
-        super().__init__(in_dim, out_dim, bias=bias)
-        with torch.no_grad():
-            if self.bias is not None:
-                self.bias.zero_()
-            if init_fn is None:
-                _initialize_weight(self.weight, self.bias, init)
-            else:
-                init_fn(self.weight, self.bias)
-        self.precision = precision
-
-    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)
-                )  # (d_out,) or None
-                projected = nn.functional.linear(
-                    input.to(self.precision),
-                    self.weight.to(self.precision),
-                    bias,
-                )  # (..., 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)
-                )  # (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):
-    """Layer normalization that avoids an autocast upcast for BF16 inputs."""
-
-    def __init__(self, c_in: int, eps: float = 1e-5) -> None:
-        super().__init__()
-        self.c_in = (c_in,)
-        self.eps = eps
-        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(
-                    x,
-                    self.c_in,
-                    self.weight.to(x.dtype),
-                    self.bias.to(x.dtype),
-                    self.eps,
-                )
-        return nn.functional.layer_norm(x, self.c_in, self.weight, self.bias, self.eps)
-
-
-@torch.jit.ignore
-def softmax_no_cast(tensor: torch.Tensor, dim: int = -1) -> torch.Tensor:
-    """Apply softmax without promoting BF16 inputs under CUDA autocast."""
-
-    # tensor: (...); softmax preserves shape.
-    if tensor.dtype is torch.bfloat16:
-        with torch.autocast("cuda", enabled=False):
-            return nn.functional.softmax(tensor, dim=dim)
-    return nn.functional.softmax(tensor, dim=dim)
-
-
-def _attention(
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    biases: Sequence[torch.Tensor],
-) -> torch.Tensor:
-    """Compute biased scaled attention after query scaling."""
-
-    # 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  # (..., 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
-def kernel_triangular_attn(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    tri_bias: torch.Tensor,
-    mask: torch.Tensor,
-    scale: float,
-) -> 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
-    ):
-        raise RuntimeError(
-            "Boltz2 use_kernels=True requires cuequivariance_torch and the CUDA 13 "
-            "cuequivariance_ops_torch runtime from the 'structure,cueq' extras."
-        )
-    cueq = importlib.import_module("cuequivariance_torch")
-    return cueq.triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale)
-
-
-class Attention(nn.Module):
-    """Multi-head pair attention with optional query-dependent output gates."""
-
-    def __init__(
-        self,
-        c_q: int,
-        c_k: int,
-        c_v: int,
-        c_hidden: int,
-        no_heads: int,
-        gating: bool = True,
-    ) -> None:
-        super().__init__()
-        self.c_q = c_q
-        self.c_k = c_k
-        self.c_v = c_v
-        self.c_hidden = c_hidden
-        self.no_heads = no_heads
-        self.gating = gating
-
-        projected_dim = c_hidden * no_heads
-        self.linear_q = Linear(c_q, projected_dim, bias=False, init="glorot")
-        self.linear_k = Linear(c_k, projected_dim, bias=False, init="glorot")
-        self.linear_v = Linear(c_v, projected_dim, bias=False, init="glorot")
-        self.linear_o = Linear(projected_dim, c_q, bias=False, init="final")
-        self.linear_g = Linear(c_q, projected_dim, bias=False, init="gating") if gating else None
-        self.sigmoid = nn.Sigmoid()
-
-    def _prep_qkv(
-        self,
-        q_x: torch.Tensor,
-        kv_x: torch.Tensor,
-        apply_scale: bool = True,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Project Q, K, and V to ``(..., h, n, d)`` tensors."""
-
-        def split_heads(X: torch.Tensor) -> torch.Tensor:
-            # 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)  # (..., 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))  # (..., 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,
-        q_x: torch.Tensor,
-        kv_x: torch.Tensor,
-        tri_bias: torch.Tensor,
-        mask_bias: torch.Tensor,
-        mask: torch.Tensor,
-        use_kernels: bool = False,
-    ) -> 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: (..., h, n_q, d_h); k/v: (..., h, n_k, d_h)
-        if use_kernels:
-            output = kernel_triangular_attn(
-                q,
-                k,
-                v,
-                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)
-            )  # (..., h, n_q, d_h)
-        return self._wrap_up(
-            output.transpose(-2, -3), q_x
-        )  # (..., n_q, c_q)
+"""Attention primitives for the Boltz pair stack.
+
+The module retains the converted checkpoint names while providing a compact
+implementation around PyTorch operations.  Optional triangle kernels are
+loaded only at the call site, so importing FastPLMs does not initialize CUDA.
+
+Portions of the numerical contract follow OpenFold and DeepMind AlphaFold
+under Apache-2.0.  See ``THIRD_PARTY_NOTICES.md`` for provenance.
+"""
+
+from __future__ import annotations
+
+import importlib
+import math
+import torch
+
+from collections.abc import Callable, Sequence
+from importlib.util import find_spec
+from torch import nn
+
+from . import vb_layers_initialize as initialize
+from .vb_tri_attn_utils import flatten_final_dims, permute_final_dims
+
+
+def _initialize_weight(
+    weight: torch.Tensor,
+    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_,
+        "glorot": initialize.glorot_uniform_init_,
+        "normal": initialize.normal_init_,
+        "final": initialize.final_init_,
+        "gating": initialize.gating_init_,
+    }
+    try:
+        initializers[method](weight)
+    except KeyError as error:
+        raise ValueError(f"unknown linear initialization {method!r}") from error
+    if method == "gating" and bias is not None:
+        bias.fill_(1.0)
+
+
+class Linear(nn.Linear):
+    """Linear projection with the initializers used by converted checkpoints."""
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        bias: bool = True,
+        init: str = "default",
+        init_fn: Callable[[torch.Tensor, torch.Tensor | None], None] | None = None,
+        precision: torch.dtype | None = None,
+    ) -> None:
+        super().__init__(in_dim, out_dim, bias=bias)
+        with torch.no_grad():
+            if self.bias is not None:
+                self.bias.zero_()
+            if init_fn is None:
+                _initialize_weight(self.weight, self.bias, init)
+            else:
+                init_fn(self.weight, self.bias)
+        self.precision = precision
+
+    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)
+                )  # (d_out,) or None
+                projected = nn.functional.linear(
+                    input.to(self.precision),
+                    self.weight.to(self.precision),
+                    bias,
+                )  # (..., 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)
+                )  # (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):
+    """Layer normalization that avoids an autocast upcast for BF16 inputs."""
+
+    def __init__(self, c_in: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.c_in = (c_in,)
+        self.eps = eps
+        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(
+                    x,
+                    self.c_in,
+                    self.weight.to(x.dtype),
+                    self.bias.to(x.dtype),
+                    self.eps,
+                )
+        return nn.functional.layer_norm(x, self.c_in, self.weight, self.bias, self.eps)
+
+
+@torch.jit.ignore
+def softmax_no_cast(tensor: torch.Tensor, dim: int = -1) -> torch.Tensor:
+    """Apply softmax without promoting BF16 inputs under CUDA autocast."""
+
+    # tensor: (...); softmax preserves shape.
+    if tensor.dtype is torch.bfloat16:
+        with torch.autocast("cuda", enabled=False):
+            return nn.functional.softmax(tensor, dim=dim)
+    return nn.functional.softmax(tensor, dim=dim)
+
+
+def _attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    biases: Sequence[torch.Tensor],
+) -> torch.Tensor:
+    """Compute biased scaled attention after query scaling."""
+
+    # 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  # (..., 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
+def kernel_triangular_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    tri_bias: torch.Tensor,
+    mask: torch.Tensor,
+    scale: float,
+) -> 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
+    ):
+        raise RuntimeError(
+            "Boltz2 use_kernels=True requires cuequivariance_torch and the CUDA 13 "
+            "cuequivariance_ops_torch runtime from the 'structure,cueq' extras."
+        )
+    cueq = importlib.import_module("cuequivariance_torch")
+    return cueq.triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale)
+
+
+class Attention(nn.Module):
+    """Multi-head pair attention with optional query-dependent output gates."""
+
+    def __init__(
+        self,
+        c_q: int,
+        c_k: int,
+        c_v: int,
+        c_hidden: int,
+        no_heads: int,
+        gating: bool = True,
+    ) -> None:
+        super().__init__()
+        self.c_q = c_q
+        self.c_k = c_k
+        self.c_v = c_v
+        self.c_hidden = c_hidden
+        self.no_heads = no_heads
+        self.gating = gating
+
+        projected_dim = c_hidden * no_heads
+        self.linear_q = Linear(c_q, projected_dim, bias=False, init="glorot")
+        self.linear_k = Linear(c_k, projected_dim, bias=False, init="glorot")
+        self.linear_v = Linear(c_v, projected_dim, bias=False, init="glorot")
+        self.linear_o = Linear(projected_dim, c_q, bias=False, init="final")
+        self.linear_g = Linear(c_q, projected_dim, bias=False, init="gating") if gating else None
+        self.sigmoid = nn.Sigmoid()
+
+    def _prep_qkv(
+        self,
+        q_x: torch.Tensor,
+        kv_x: torch.Tensor,
+        apply_scale: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Project Q, K, and V to ``(..., h, n, d)`` tensors."""
+
+        def split_heads(X: torch.Tensor) -> torch.Tensor:
+            # 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)  # (..., 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))  # (..., 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,
+        q_x: torch.Tensor,
+        kv_x: torch.Tensor,
+        tri_bias: torch.Tensor,
+        mask_bias: torch.Tensor,
+        mask: torch.Tensor,
+        use_kernels: bool = False,
+    ) -> 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: (..., h, n_q, d_h); k/v: (..., h, n_k, d_h)
+        if use_kernels:
+            output = kernel_triangular_attn(
+                q,
+                k,
+                v,
+                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)
+            )  # (..., h, n_q, d_h)
+        return self._wrap_up(
+            output.transpose(-2, -3), q_x
+        )  # (..., n_q, c_q)

fastplms/models/boltz/vb_tri_attn_utils.py

@@ -1,341 +1,341 @@
-"""Tensor-tree and memory-bounded execution helpers for pair attention.
-
-These utilities implement the small subset of OpenFold-style chunking needed
-by the local Boltz pair stack.  The implementation is independent, but keeps
-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
-
-
-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  # 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:
-    """Permute only the final ``len(inds)`` axes of ``tensor``."""
-
-    # tensor: (..., d_0, ..., d_n); inds reorders the named final axes.
-    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)  # (..., d_0, ..., d_n)
-
-
-def is_fp16_enabled() -> bool:
-    """Return whether CUDA autocast currently targets ``float16``."""
-
-    return torch.is_autocast_enabled() and torch.get_autocast_dtype("cuda") == torch.float16
-
-
-def dict_map(
-    fn: Callable[[Any], Any],
-    dic: dict[Any, Any],
-    leaf_type: type | tuple[type, ...],
-) -> dict[Any, Any]:
-    """Apply ``fn`` to leaves in a nested dictionary tree."""
-
-    return {key: tree_map(fn, value, leaf_type) for key, value in dic.items()}
-
-
-def tree_map(
-    fn: Callable[[Any], Any],
-    tree: Any,
-    leaf_type: type | tuple[type, ...],
-) -> Any:
-    """Map a function over dict, list, and tuple containers."""
-
-    if isinstance(tree, leaf_type):
-        return fn(tree)
-    if isinstance(tree, dict):
-        return dict_map(fn, tree, leaf_type)
-    if isinstance(tree, list):
-        return [tree_map(fn, item, leaf_type) for item in tree]
-    if isinstance(tree, tuple):
-        return tuple(tree_map(fn, item, leaf_type) for item in tree)
-    raise ValueError(f"tree type {type(tree)!r} is not supported")
-
-
-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."""
-
-    # tensor: (..., d_0, ..., d_n)
-    return tensor.reshape(*tensor.shape[:-no_dims], -1)  # (..., prod(final dims))
-
-
-def _fetch_dims(tree: Any) -> list[torch.Size]:
-    """Collect tensor shapes from a supported tree."""
-
-    if isinstance(tree, torch.Tensor):
-        return [tree.shape]
-    if isinstance(tree, dict):
-        children = tree.values()
-    elif isinstance(tree, (list, tuple)):
-        children = tree
-    else:
-        raise ValueError(f"tree type {type(tree)!r} is not supported")
-
-    shapes: list[torch.Size] = []
-    for child in children:
-        shapes.extend(_fetch_dims(child))
-    return shapes
-
-
-@torch.jit.ignore
-def _flat_idx_to_idx(flat_idx: int, dims: tuple[int, ...]) -> tuple[int, ...]:
-    """Convert a row-major flat index into an index tuple."""
-
-    coordinates = [0] * len(dims)
-    remainder = flat_idx
-    for axis in range(len(dims) - 1, -1, -1):
-        remainder, coordinates[axis] = divmod(remainder, dims[axis])
-    return tuple(coordinates)
-
-
-def _ravel_index(index: Sequence[int], dims: Sequence[int]) -> int:
-    """Convert a row-major index tuple to a flat index."""
-
-    flat = 0
-    for coordinate, size in zip(index, dims, strict=True):
-        flat = flat * size + coordinate
-    return flat
-
-
-def _cover_flat_interval(
-    start: int,
-    stop: int,
-    dims: tuple[int, ...],
-) -> list[tuple[slice, ...]]:
-    """Cover ``[start, stop)`` with ordered contiguous tensor slices."""
-
-    if not dims:
-        return [tuple()]
-    if len(dims) == 1:
-        return [(slice(start, stop),)]
-
-    child_size = prod(dims[1:])
-    first_child = start // child_size
-    last_child = (stop - 1) // child_size
-    if first_child == last_child:
-        tail = _cover_flat_interval(
-            start % child_size,
-            (stop - 1) % child_size + 1,
-            dims[1:],
-        )
-        prefix = slice(first_child, first_child + 1)
-        return [(prefix, *item) for item in tail]
-
-    slices: list[tuple[slice, ...]] = []
-    start_offset = start % child_size
-    first_full_child = first_child
-    if start_offset:
-        prefix = slice(first_child, first_child + 1)
-        slices.extend(
-            (prefix, *item) for item in _cover_flat_interval(start_offset, child_size, dims[1:])
-        )
-        first_full_child += 1
-
-    stop_offset = stop % child_size
-    full_stop = last_child if stop_offset else last_child + 1
-    if first_full_child < full_stop:
-        slices.append((slice(first_full_child, full_stop),))
-
-    if stop_offset:
-        prefix = slice(last_child, last_child + 1)
-        slices.extend((prefix, *item) for item in _cover_flat_interval(0, stop_offset, dims[1:]))
-    return slices
-
-
-@torch.jit.ignore
-def _get_minimal_slice_set(
-    start: Sequence[int],
-    end: Sequence[int],
-    dims: Sequence[int],
-    start_edges: Sequence[bool] | None = None,
-    end_edges: Sequence[bool] | None = None,
-) -> list[tuple[slice, ...]]:
-    """Return ordered slices covering the inclusive row-major interval.
-
-    ``start_edges`` and ``end_edges`` remain accepted for compatibility.  The
-    interval decomposition derives the same information directly.
-    """
-
-    del start_edges, end_edges
-    shape = tuple(dims)
-    if not shape:
-        return [tuple()]
-    flat_start = _ravel_index(start, shape)
-    flat_stop = _ravel_index(end, shape) + 1
-    return _cover_flat_interval(flat_start, flat_stop, shape)
-
-
-@torch.jit.ignore
-def _chunk_slice(
-    tensor: torch.Tensor,
-    flat_start: int,
-    flat_end: int,
-    no_batch_dims: int,
-) -> torch.Tensor:
-    """Slice a flattened batch interval without flattening the full tensor."""
-
-    # tensor: (*batch_shape, *feature_shape).
-    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)
-    ]  # 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]
-    )  # (flat_end - flat_start, *feature_shape)
-
-
-def _prepare_input(
-    tensor: torch.Tensor,
-    batch_shape: tuple[int, ...],
-    no_batch_dims: int,
-    *,
-    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)  # (*batch_shape, *feature_shape)
-    if any(size != 1 for size in tensor.shape[:no_batch_dims]):
-        tensor = tensor.expand(
-            batch_shape + feature_shape
-        )  # (*batch_shape, *feature_shape)
-    return tensor.reshape(-1, *feature_shape)  # (flat_batch, *feature_shape)
-
-
-def _write_chunk(
-    destination: Any,
-    source: Any,
-    start: int,
-    stop: int,
-    *,
-    add_into_out: bool,
-) -> None:
-    if isinstance(source, dict):
-        for key, value in source.items():
-            _write_chunk(
-                destination[key],
-                value,
-                start,
-                stop,
-                add_into_out=add_into_out,
-            )
-        return
-    if isinstance(source, (tuple, list)):
-        for destination_item, source_item in zip(destination, source, strict=True):
-            _write_chunk(
-                destination_item,
-                source_item,
-                start,
-                stop,
-                add_into_out=add_into_out,
-            )
-        return
-    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  # (stop - start, *feature_shape)
-    else:
-        destination[start:stop] = source  # (stop - start, *feature_shape)
-
-
-def chunk_layer(
-    layer: Callable[..., Any],
-    inputs: dict[str, Any],
-    chunk_size: int,
-    no_batch_dims: int,
-    low_mem: bool = False,
-    _out: Any = None,
-    _add_into_out: bool = False,
-) -> Any:
-    """Apply ``layer`` to flattened batch chunks and reassemble its output."""
-
-    if not inputs:
-        raise ValueError("at least one input is required")
-    if chunk_size <= 0:
-        raise ValueError("chunk_size must be positive")
-
-    leading_shapes = [shape[:no_batch_dims] for shape in _fetch_dims(inputs)]
-    batch_shape = tuple(max(sizes) for sizes in zip(*leading_shapes, strict=True))
-    prepare = partial(
-        _prepare_input,
-        batch_shape=batch_shape,
-        no_batch_dims=no_batch_dims,
-        low_mem=low_mem,
-    )
-    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):
-        stop = min(flat_batch_size, start + chunk_size)
-        if low_mem:
-            select = partial(
-                _chunk_slice,
-                flat_start=start,
-                flat_end=stop,
-                no_batch_dims=len(batch_shape),
-            )
-        else:
-
-            def select(
-                tensor: torch.Tensor,
-                start: int = start,
-                stop: int = stop,
-            ) -> torch.Tensor:
-                return tensor if tensor.shape[0] == 1 else tensor[start:stop]
-
-        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,
-            start,
-            stop,
-            add_into_out=_add_into_out,
-        )
-
-    return tensor_tree_map(
-        lambda tensor: tensor.reshape(batch_shape + tuple(tensor.shape[1:])),
-        output,
-    )  # each tensor: (*batch_shape, *output_feature_shape)
+"""Tensor-tree and memory-bounded execution helpers for pair attention.
+
+These utilities implement the small subset of OpenFold-style chunking needed
+by the local Boltz pair stack.  The implementation is independent, but keeps
+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
+
+
+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  # 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:
+    """Permute only the final ``len(inds)`` axes of ``tensor``."""
+
+    # tensor: (..., d_0, ..., d_n); inds reorders the named final axes.
+    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)  # (..., d_0, ..., d_n)
+
+
+def is_fp16_enabled() -> bool:
+    """Return whether CUDA autocast currently targets ``float16``."""
+
+    return torch.is_autocast_enabled() and torch.get_autocast_dtype("cuda") == torch.float16
+
+
+def dict_map(
+    fn: Callable[[Any], Any],
+    dic: dict[Any, Any],
+    leaf_type: type | tuple[type, ...],
+) -> dict[Any, Any]:
+    """Apply ``fn`` to leaves in a nested dictionary tree."""
+
+    return {key: tree_map(fn, value, leaf_type) for key, value in dic.items()}
+
+
+def tree_map(
+    fn: Callable[[Any], Any],
+    tree: Any,
+    leaf_type: type | tuple[type, ...],
+) -> Any:
+    """Map a function over dict, list, and tuple containers."""
+
+    if isinstance(tree, leaf_type):
+        return fn(tree)
+    if isinstance(tree, dict):
+        return dict_map(fn, tree, leaf_type)
+    if isinstance(tree, list):
+        return [tree_map(fn, item, leaf_type) for item in tree]
+    if isinstance(tree, tuple):
+        return tuple(tree_map(fn, item, leaf_type) for item in tree)
+    raise ValueError(f"tree type {type(tree)!r} is not supported")
+
+
+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."""
+
+    # tensor: (..., d_0, ..., d_n)
+    return tensor.reshape(*tensor.shape[:-no_dims], -1)  # (..., prod(final dims))
+
+
+def _fetch_dims(tree: Any) -> list[torch.Size]:
+    """Collect tensor shapes from a supported tree."""
+
+    if isinstance(tree, torch.Tensor):
+        return [tree.shape]
+    if isinstance(tree, dict):
+        children = tree.values()
+    elif isinstance(tree, (list, tuple)):
+        children = tree
+    else:
+        raise ValueError(f"tree type {type(tree)!r} is not supported")
+
+    shapes: list[torch.Size] = []
+    for child in children:
+        shapes.extend(_fetch_dims(child))
+    return shapes
+
+
+@torch.jit.ignore
+def _flat_idx_to_idx(flat_idx: int, dims: tuple[int, ...]) -> tuple[int, ...]:
+    """Convert a row-major flat index into an index tuple."""
+
+    coordinates = [0] * len(dims)
+    remainder = flat_idx
+    for axis in range(len(dims) - 1, -1, -1):
+        remainder, coordinates[axis] = divmod(remainder, dims[axis])
+    return tuple(coordinates)
+
+
+def _ravel_index(index: Sequence[int], dims: Sequence[int]) -> int:
+    """Convert a row-major index tuple to a flat index."""
+
+    flat = 0
+    for coordinate, size in zip(index, dims, strict=True):
+        flat = flat * size + coordinate
+    return flat
+
+
+def _cover_flat_interval(
+    start: int,
+    stop: int,
+    dims: tuple[int, ...],
+) -> list[tuple[slice, ...]]:
+    """Cover ``[start, stop)`` with ordered contiguous tensor slices."""
+
+    if not dims:
+        return [tuple()]
+    if len(dims) == 1:
+        return [(slice(start, stop),)]
+
+    child_size = prod(dims[1:])
+    first_child = start // child_size
+    last_child = (stop - 1) // child_size
+    if first_child == last_child:
+        tail = _cover_flat_interval(
+            start % child_size,
+            (stop - 1) % child_size + 1,
+            dims[1:],
+        )
+        prefix = slice(first_child, first_child + 1)
+        return [(prefix, *item) for item in tail]
+
+    slices: list[tuple[slice, ...]] = []
+    start_offset = start % child_size
+    first_full_child = first_child
+    if start_offset:
+        prefix = slice(first_child, first_child + 1)
+        slices.extend(
+            (prefix, *item) for item in _cover_flat_interval(start_offset, child_size, dims[1:])
+        )
+        first_full_child += 1
+
+    stop_offset = stop % child_size
+    full_stop = last_child if stop_offset else last_child + 1
+    if first_full_child < full_stop:
+        slices.append((slice(first_full_child, full_stop),))
+
+    if stop_offset:
+        prefix = slice(last_child, last_child + 1)
+        slices.extend((prefix, *item) for item in _cover_flat_interval(0, stop_offset, dims[1:]))
+    return slices
+
+
+@torch.jit.ignore
+def _get_minimal_slice_set(
+    start: Sequence[int],
+    end: Sequence[int],
+    dims: Sequence[int],
+    start_edges: Sequence[bool] | None = None,
+    end_edges: Sequence[bool] | None = None,
+) -> list[tuple[slice, ...]]:
+    """Return ordered slices covering the inclusive row-major interval.
+
+    ``start_edges`` and ``end_edges`` remain accepted for compatibility.  The
+    interval decomposition derives the same information directly.
+    """
+
+    del start_edges, end_edges
+    shape = tuple(dims)
+    if not shape:
+        return [tuple()]
+    flat_start = _ravel_index(start, shape)
+    flat_stop = _ravel_index(end, shape) + 1
+    return _cover_flat_interval(flat_start, flat_stop, shape)
+
+
+@torch.jit.ignore
+def _chunk_slice(
+    tensor: torch.Tensor,
+    flat_start: int,
+    flat_end: int,
+    no_batch_dims: int,
+) -> torch.Tensor:
+    """Slice a flattened batch interval without flattening the full tensor."""
+
+    # tensor: (*batch_shape, *feature_shape).
+    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)
+    ]  # 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]
+    )  # (flat_end - flat_start, *feature_shape)
+
+
+def _prepare_input(
+    tensor: torch.Tensor,
+    batch_shape: tuple[int, ...],
+    no_batch_dims: int,
+    *,
+    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)  # (*batch_shape, *feature_shape)
+    if any(size != 1 for size in tensor.shape[:no_batch_dims]):
+        tensor = tensor.expand(
+            batch_shape + feature_shape
+        )  # (*batch_shape, *feature_shape)
+    return tensor.reshape(-1, *feature_shape)  # (flat_batch, *feature_shape)
+
+
+def _write_chunk(
+    destination: Any,
+    source: Any,
+    start: int,
+    stop: int,
+    *,
+    add_into_out: bool,
+) -> None:
+    if isinstance(source, dict):
+        for key, value in source.items():
+            _write_chunk(
+                destination[key],
+                value,
+                start,
+                stop,
+                add_into_out=add_into_out,
+            )
+        return
+    if isinstance(source, (tuple, list)):
+        for destination_item, source_item in zip(destination, source, strict=True):
+            _write_chunk(
+                destination_item,
+                source_item,
+                start,
+                stop,
+                add_into_out=add_into_out,
+            )
+        return
+    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  # (stop - start, *feature_shape)
+    else:
+        destination[start:stop] = source  # (stop - start, *feature_shape)
+
+
+def chunk_layer(
+    layer: Callable[..., Any],
+    inputs: dict[str, Any],
+    chunk_size: int,
+    no_batch_dims: int,
+    low_mem: bool = False,
+    _out: Any = None,
+    _add_into_out: bool = False,
+) -> Any:
+    """Apply ``layer`` to flattened batch chunks and reassemble its output."""
+
+    if not inputs:
+        raise ValueError("at least one input is required")
+    if chunk_size <= 0:
+        raise ValueError("chunk_size must be positive")
+
+    leading_shapes = [shape[:no_batch_dims] for shape in _fetch_dims(inputs)]
+    batch_shape = tuple(max(sizes) for sizes in zip(*leading_shapes, strict=True))
+    prepare = partial(
+        _prepare_input,
+        batch_shape=batch_shape,
+        no_batch_dims=no_batch_dims,
+        low_mem=low_mem,
+    )
+    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):
+        stop = min(flat_batch_size, start + chunk_size)
+        if low_mem:
+            select = partial(
+                _chunk_slice,
+                flat_start=start,
+                flat_end=stop,
+                no_batch_dims=len(batch_shape),
+            )
+        else:
+
+            def select(
+                tensor: torch.Tensor,
+                start: int = start,
+                stop: int = stop,
+            ) -> torch.Tensor:
+                return tensor if tensor.shape[0] == 1 else tensor[start:stop]
+
+        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,
+            start,
+            stop,
+            add_into_out=_add_into_out,
+        )
+
+    return tensor_tree_map(
+        lambda tensor: tensor.reshape(batch_shape + tuple(tensor.shape[1:])),
+        output,
+    )  # each tensor: (*batch_shape, *output_feature_shape)

fastplms/registry.py

@@ -87,6 +87,8 @@ _ROOT_FIELDS = frozenset(
         "families",
         "models",
         "runtime_assets",
+        "golden_artifacts",
+        "sparse_autoencoders",
     }
 )
 _UPSTREAM_FIELDS = frozenset(
@@ -94,6 +96,7 @@ _UPSTREAM_FIELDS = frozenset(
         "id",
         "path",
         "url",
+        "fetch_url",
         "revision",
         "license",
         "license_files",
@@ -177,6 +180,17 @@ _RUNTIME_ASSET_FIELDS = frozenset(
         "offline_behavior",
     }
 )
+_GOLDEN_ARTIFACT_FIELDS = frozenset(
+    {
+        "id",
+        "repository",
+        "revision",
+        "path",
+        "sha256",
+        "size",
+        "offline_behavior",
+    }
+)
 
 
 class RegistryError(ValueError):
@@ -262,6 +276,20 @@ class CheckpointSource:
         return MappingProxyType({item.path: item for item in self.files})
 
 
+@dataclass(frozen=True, slots=True)
+class SparseAutoencoderSpec:
+    """Pinned released SAE and its input contract, not its unknown training revision."""
+
+    id: str
+    base_model: str
+    checkpoint: CheckpointSource
+    layer: int
+    input_width: int
+    k: int
+    codebook_dim: int
+    input_kind: Literal["hidden_state"] = "hidden_state"
+
+
 @dataclass(frozen=True, slots=True)
 class OracleAsset:
     """Hash-pinned external file required by a native parity oracle."""
@@ -297,6 +325,19 @@ class OfficialGolden:
     tensors: FileDigest
 
 
+@dataclass(frozen=True, slots=True)
+class GoldenArtifact:
+    """Durable Hub copy of one model's official golden tensors at an immutable revision."""
+
+    model_id: str
+    repository: str
+    revision: str
+    path: str
+    sha256: str
+    size: int
+    offline_behavior: str
+
+
 @dataclass(frozen=True, slots=True)
 class UpstreamSource:
     """Pinned official implementation used as a parity oracle."""
@@ -309,6 +350,13 @@ class UpstreamSource:
     license_files: tuple[str, ...]
     license_digests: tuple[FileDigest, ...] = ()
     distribution_files: tuple[FileDigest, ...] = ()
+    # Where the pinned revision is fetched when `url`, the source of record, no longer serves it.
+    fetch_url: str = ""
+
+    @property
+    def clone_url(self) -> str:
+        """The URL that `.gitmodules` and source archives use."""
+        return self.fetch_url or self.url
 
 
 @dataclass(frozen=True, slots=True)
@@ -478,6 +526,8 @@ class ModelRegistry(Mapping[str, ModelSpec]):
         runtime_assets: Mapping[str, RuntimeAsset] = MappingProxyType({}),
         attention_kernels: Mapping[str, AttentionKernelSpec] = MappingProxyType({}),
         legal_files: tuple[FileDigest, ...] = (),
+        golden_artifacts: Mapping[str, GoldenArtifact] = MappingProxyType({}),
+        sparse_autoencoders: Mapping[str, SparseAutoencoderSpec] = MappingProxyType({}),
     ) -> None:
         self.schema_version = schema_version
         self.upstreams = MappingProxyType(dict(upstreams))
@@ -486,6 +536,8 @@ class ModelRegistry(Mapping[str, ModelSpec]):
         self._models = MappingProxyType(dict(models))
         self.runtime_assets = MappingProxyType(dict(runtime_assets))
         self.legal_files = legal_files
+        self.golden_artifacts = MappingProxyType(dict(golden_artifacts))
+        self.sparse_autoencoders = MappingProxyType(dict(sparse_autoencoders))
 
     def __getitem__(self, key: str) -> ModelSpec:
         return self._models[key]
@@ -501,6 +553,22 @@ class ModelRegistry(Mapping[str, ModelSpec]):
             raise KeyError(family_id)
         return tuple(model for model in self._models.values() if model.family.id == family_id)
 
+    def sae_for_base(
+        self, base_model: str, *, layer: int, k: int, codebook_dim: int
+    ) -> SparseAutoencoderSpec:
+        """Resolve an exact registered SAE selection without a Hub lookup or fallback."""
+        matches = tuple(
+            spec for spec in self.sparse_autoencoders.values()
+            if (spec.base_model, spec.layer, spec.k, spec.codebook_dim)
+            == (base_model, layer, k, codebook_dim)
+        )
+        if len(matches) != 1:
+            raise RegistryError(
+                f"Expected one pinned SAE for {base_model}, layer={layer}, k={k}, "
+                f"codebook_dim={codebook_dim}; found {len(matches)}."
+            )
+        return matches[0]
+
     def supported_attention_dtypes(
         self,
         family_id: str,
@@ -988,8 +1056,11 @@ def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
             raise RegistryError(f"Duplicate upstream path: {path!r}")
         paths.add(path)
         url = _require_str(value, "url", context)
-        if not url.startswith("https://github.com/") or not url.endswith(".git"):
-            raise RegistryError(f"{context}.url must be an HTTPS GitHub clone URL.")
+        fetch_url = _optional_str(value, "fetch_url", context) or ""
+        for field, candidate in (("url", url), ("fetch_url", fetch_url)):
+            is_github = candidate.startswith("https://github.com/")
+            if candidate and not (is_github and candidate.endswith(".git")):
+                raise RegistryError(f"{context}.{field} must be an HTTPS GitHub clone URL.")
         license_files = _require_str_list(value, "license_files", context)
         license_digests = _require_digest_list(value, "license_digests", context)
         if tuple(item.path for item in license_digests) != license_files:
@@ -1026,6 +1097,7 @@ def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
             license_files=license_files,
             license_digests=license_digests,
             distribution_files=distribution_files,
+            fetch_url=fetch_url,
         )
     return upstreams
 
@@ -1288,6 +1360,67 @@ def _parse_runtime_assets(
     return runtime_assets
 
 
+def _parse_golden_artifacts(
+    raw: object,
+    models: Mapping[str, ModelSpec],
+) -> dict[str, GoldenArtifact]:
+    """Parse the optional Hub copies of official golden tensors, keyed by model ID.
+
+    A copy must be byte-identical to the tensors its model's official_golden pins, so its
+    SHA-256 must equal that digest and the existing integrity check covers both files.
+    """
+    if raw is None:
+        return {}
+    if not isinstance(raw, list):
+        raise RegistryError("golden_artifacts must be an array of tables.")
+    golden_artifacts: dict[str, GoldenArtifact] = {}
+    for index, value in enumerate(raw):
+        context = f"golden_artifacts[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(value, _GOLDEN_ARTIFACT_FIELDS, context)
+        model_id = _require_str(value, "id", context)
+        model = models.get(model_id)
+        if model is None:
+            raise RegistryError(f"{context}.id references unknown model {model_id!r}.")
+        if model.official_golden is None:
+            raise RegistryError(f"{context}.id names {model_id!r}, which has no official_golden.")
+        if model_id in golden_artifacts:
+            raise RegistryError(f"Duplicate golden artifact for model {model_id!r}.")
+        repository = _require_str(value, "repository", context)
+        if _REPOSITORY_ID_RE.fullmatch(repository) is None:
+            raise RegistryError(f"{context}.repository must be a Hugging Face repository ID.")
+        revision = _require_str(value, "revision", context)
+        _validate_revision(revision, f"{context}.revision")
+        path = _require_str(value, "path", context)
+        expected_path = f"goldens/{model_id}.safetensors"
+        if path != expected_path:
+            raise RegistryError(f"{context}.path must be {expected_path!r}.")
+        sha256 = _require_str(value, "sha256", context)
+        if sha256 != model.official_golden.tensors.digest:
+            raise RegistryError(
+                f"{context}.sha256 must equal the official_golden tensors digest of {model_id!r}."
+            )
+        size = value.get("size")
+        if isinstance(size, bool) or not isinstance(size, int) or size <= 0:
+            raise RegistryError(f"{context}.size must be a positive byte count.")
+        offline_behavior = _require_str(value, "offline_behavior", context)
+        if offline_behavior not in _ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS:
+            raise RegistryError(
+                f"{context}.offline_behavior is unsupported: {offline_behavior!r}."
+            )
+        golden_artifacts[model_id] = GoldenArtifact(
+            model_id=model_id,
+            repository=repository,
+            revision=revision,
+            path=path,
+            sha256=sha256,
+            size=size,
+            offline_behavior=offline_behavior,
+        )
+    return golden_artifacts
+
+
 def _parse_confidence_adaptation(
     table: Mapping[str, Any], context: str, model_id: str | None = None
 ) -> ConfidenceAdaptation | None:
@@ -1667,6 +1800,57 @@ def _validate_registry(
             )
 
 
+def _parse_sparse_autoencoders(
+    raw: object, models: Mapping[str, ModelSpec]
+) -> dict[str, SparseAutoencoderSpec]:
+    """Validate optional SAE records independently of the base-model mapping."""
+    if raw is None:
+        return {}
+    if not isinstance(raw, list):
+        raise RegistryError("sparse_autoencoders must be an array of tables.")
+    allowed = frozenset({
+        "id", "base_model", "layer", "input_width", "k", "codebook_dim", "input_kind",
+        "checkpoint_repo", "checkpoint_revision", "checkpoint_files",
+    })
+    records: dict[str, SparseAutoencoderSpec] = {}
+    selections: set[tuple[str, int, int, int]] = set()
+    for index, table in enumerate(raw):
+        context = f"sparse_autoencoders[{index}]"
+        if not isinstance(table, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(table, allowed, context)
+        identifier = _require_str(table, "id", context)
+        if _IDENTIFIER_RE.fullmatch(identifier) is None or identifier in records:
+            raise RegistryError(f"{context} has an invalid or duplicate SAE id: {identifier!r}.")
+        base = _require_str(table, "base_model", context)
+        if base not in models or models[base].family.id != "esm_plusplus":
+            raise RegistryError(f"{context}.base_model must name a registered ESMC base.")
+        numbers: dict[str, int] = {}
+        for field in ("layer", "input_width", "k", "codebook_dim"):
+            value = table.get(field)
+            if type(value) is not int or value < (0 if field == "layer" else 1):
+                raise RegistryError(f"{context}.{field} must be a valid integer dimension.")
+            numbers[field] = value
+        if numbers["k"] > numbers["codebook_dim"]:
+            raise RegistryError(f"{context}.k exceeds codebook_dim.")
+        if _require_str(table, "input_kind", context) != "hidden_state":
+            raise RegistryError(f"{context} requires hidden_state input, not residual updates.")
+        checkpoint = _parse_checkpoint(table, "checkpoint", context)
+        expected_files = {"config.json", f"layer_{numbers['layer']}.safetensors"}
+        if set(checkpoint.file_map) != expected_files or any(
+            item.algorithm != "sha256" for item in checkpoint.files
+        ):
+            raise RegistryError(f"{context} must SHA-256 pin exactly {sorted(expected_files)}.")
+        selection = (base, numbers["layer"], numbers["k"], numbers["codebook_dim"])
+        if selection in selections:
+            raise RegistryError(f"{context} duplicates an SAE selection: {selection}.")
+        selections.add(selection)
+        records[identifier] = SparseAutoencoderSpec(
+            id=identifier, base_model=base, checkpoint=checkpoint, **numbers
+        )
+    return records
+
+
 def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry:
     try:
         manifest = tomllib.loads(raw_bytes.decode("utf-8"))
@@ -1685,6 +1869,8 @@ def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry:
     runtime_assets = _parse_runtime_assets(manifest.get("runtime_assets"), families)
     models = _parse_models(manifest.get("models"), families)
     _validate_registry(upstreams, attention_kernels, families, models)
+    golden_artifacts = _parse_golden_artifacts(manifest.get("golden_artifacts"), models)
+    sparse_autoencoders = _parse_sparse_autoencoders(manifest.get("sparse_autoencoders"), models)
     return ModelRegistry(
         schema_version=1,
         upstreams=upstreams,
@@ -1693,6 +1879,8 @@ def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry:
         models=models,
         runtime_assets=runtime_assets,
         legal_files=legal_files,
+        golden_artifacts=golden_artifacts,
+        sparse_autoencoders=sparse_autoencoders,
     )
 
 
@@ -1729,6 +1917,7 @@ __all__ = [
     "CheckpointSource",
     "FileDigest",
     "GenerationContract",
+    "GoldenArtifact",
     "ModelFamily",
     "ModelRegistry",
     "ModelSpec",
@@ -1737,6 +1926,7 @@ __all__ = [
     "RuntimeAsset",
     "RuntimeAssetTrustKind",
     "RuntimeExtra",
+    "SparseAutoencoderSpec",
     "TestTier",
     "UpstreamSource",
     "VramTier",

modeling_fastplms.py

@@ -13,7 +13,7 @@ from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "492b2c57aeef44795a403c858cf27184fde20e2b231a5e8278dfa5c2f8999aed":
+if RUNTIME_HASH != "bc88b7fe9bf0dd2e7b5b3337e94c4d316bf38106418cdccdd1e27882f91a1659":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []

requirements.txt

@@ -1,20 +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
+torch>=2.14
+transformers>=5.17
+huggingface-hub>=1.32
+tokenizers>=0.23
+safetensors>=0.8
+numpy>=2.5
+einops>=0.8
+tqdm>=4.70
+accelerate>=1.15
+biopython>=1.88
+biotite>=1.7
+brotli>=1.2
+msgpack>=1.2
+msgpack-numpy>=0.4.8
+omegaconf>=2.3
+rdkit>=2026.3
+scipy>=1.18
+zstandard>=0.25