Update FastPLMs runtime files

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

LICENSES/FastPLMs-Apache-2.0.txt

@@ -0,0 +1,203 @@
+PLEASE NOTE THE APACHE LICENSE ONLY APPLIES TO THE CODE IN THE FastPLMs GITHUB AND ASSOCIATED HUGGINGFACE REPOSITORIES, NOT NECESSARILY THE MODEL WEIGHTS. THOSE LICENSES CAN BE FOUND HERE https://github.com/Synthyra/FastPLMs/tree/main/LICENSES
+
+                                Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

LICENSES/boltz/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Jeremy Wohlwend, Gabriele Corso, Saro Passaro
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,164 +1,143 @@
----
-library_name: transformers
-tags: []
----
-
-# NOTE
-The GitHub with the implementation and requirements can be found [here](https://github.com/Synthyra/FastPLMs.git).
-
-# Boltz2 AutoModel (Inference-only)
-This is a barebones Huggingface `AutoModel` compatible implementation of Boltz2 focused on fast inference workflows.
-
-The implementation is located in `fastplms/boltz/` and exposes:
-- `Boltz2Config`
-- `Boltz2Model`
-- `predict_structure(amino_acid_sequence, ...)`
-- `save_as_cif(structure_output, output_path, ...)`
-
-## Design goals
-- Inference-only (no training hooks, no Lightning trainer usage).
-- Lightweight runtime around `torch` + `transformers` (plus `numpy`).
-- AutoModel remote-code compatibility via `trust_remote_code=True`.
-- Confidence outputs included in prediction outputs (`plddt`, `ptm`, `iptm`, and derived confidence score when available).
-
-## Runtime note
-This implementation is self-contained inside `fastplms/boltz/` and does not require
-the original cloned `boltz` package at runtime.
-
-## Use with transformers
-
-### Load from an exported directory
-```python
-import torch
-from transformers import AutoModel
-
-model = AutoModel.from_pretrained(
-    "Synthyra/Boltz2",
-    trust_remote_code=True,
-    dtype=torch.float32,
-).eval()
-```
-
-### Predict structure from sequence
-```python
-out = model.predict_structure(
-    amino_acid_sequence="MSTNPKPQRKTKRNTNRRPQDVKFPGG",
-    recycling_steps=3,
-    num_sampling_steps=200,
-    diffusion_samples=1,
-)
-
-print(out.sample_atom_coords.shape)
-print(None if out.plddt is None else out.plddt.shape)
-```
-
-### Save CIF
-```python
-model.save_as_cif(out, "prediction.cif")
-```
-
-## Convert Boltz checkpoint to HF export
-Use:
-
-```bash
-py -m fastplms.boltz.get_weights --checkpoint_path fastplms/boltz/weights/boltz2_conf.ckpt --output_dir boltz2_automodel_export
-```
-
-The export directory contains:
-- `config.json`
-- `pytorch_model.bin`
-- `modeling_boltz2.py`
-- `minimal_featurizer.py`
-- `minimal_structures.py`
-- `cif_writer.py`
-- `vb_*.py` (self-contained vendored Boltz2 inference modules/constants)
-
-## Output object fields
-`predict_structure(...)` returns `Boltz2StructureOutput` with:
-- `sample_atom_coords`
-- `atom_pad_mask`
-- `plddt`
-- `complex_plddt`
-- `ptm`
-- `iptm`
-- `confidence_score` (derived when available)
-- `raw_output`
-
-## Limitations
-- Current featurization path is protein-only and minimal.
-- This implementation is meant for practical inference and export workflows, not full Boltz training parity.
-- Test-time training is not supported for Boltz2 in FastPLMs. TTT is currently limited to sequence PLMs plus ESMFold and ESMFold2 PLM backbones.
-
-## Docker-first compliance testing
-
-Build the container at repo root:
-
-```bash
-docker build -t fastplms-test -f Dockerfile .
-```
-
-Launch a test shell:
-
-```bash
-docker run --rm --gpus all -it -v ${PWD}:/workspace fastplms-test bash
-```
-
-Inside the container, run Boltz2 compliance against pip `boltz`:
-
-```bash
-python -m testing.run_boltz2_compliance --device cuda --dtype float32 --seed 42 --num-sequences 3 --recycling-steps 3 --num-sampling-steps 200 --diffusion-samples 1 --pass-coord-metric aligned --write-cif-artifacts
-```
-
-Artifacts are written to `testing/results/<timestamp>/boltz2_compliance/` by default:
-- `metrics.json`
-- `metrics.csv`
-- `summary.txt`
-- `structures/seq_<idx>/ours_seq<idx>.cif`
-- `structures/seq_<idx>/ref_seq<idx>.cif`
-
-Coordinate metrics now include both raw and rigid-aligned variants:
-- `coord_mae`, `coord_rmse`, `coord_max_abs` (raw frame-dependent deltas)
-- `coord_mae_aligned`, `coord_rmse_aligned`, `coord_max_abs_aligned` (Kabsch aligned)
-- `pairwise_dist_mae` (frame-invariant pairwise-distance delta)
-
-Pass/fail uses `--pass-coord-metric aligned` by default. Set `--pass-coord-metric raw` to use the raw coordinate thresholds.
-
-## Citations
-
-```bibtex
-@misc{FastPLMs,
-  author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
-  title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.},
-  year={2024},
-  url={https://huggingface.co/Synthyra/ESMplusplus_small},
-  DOI={10.57967/hf/3726},
-  publisher={Hugging Face}
-}
-```
-
-```bibtex
-@article{passaro2025boltz2,
-  title={Boltz-2: Exploring the Frontiers of Biomolecular Prediction},
-  author={Passaro, Saro and Corso, Gabriele and Wohlwend, Jeremy and Reveiz, Mateo and Bordes, Florian and Wicky, Basile and Dayan, Peter and Jing, Bowen},
-  journal={bioRxiv},
-  year={2025}
-}
-```
-
-```bibtex
-@article{wohlwend2024boltz1,
-  title={Boltz-1: Democratizing Biomolecular Interaction Modeling},
-  author={Wohlwend, Jeremy and Corso, Gabriele and Passaro, Saro and Reveiz, Mateo and Leidal, Ken and Swanson, Wojtek and Kher, Gilmer and Lember, Tommi and Jaakkola, Tommi},
-  journal={bioRxiv},
-  year={2024}
-}
-```
-
-```bibtex
-@inproceedings{paszke2019pytorch,
-  title={PyTorch: An Imperative Style, High-Performance Deep Learning Library},
-  author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith},
-  booktitle={Advances in Neural Information Processing Systems 32},
-  year={2019}
-}
-```
+---
+library_name: transformers
+license: "mit"
+tags:
+  - protein-language-model
+  - fastplms
+---
+
+<!-- Generated from src/fastplms/models.toml. Do not edit. -->
+
+# Synthyra/Boltz2
+
+This checkpoint packages the FastPLMs `Boltz2` implementation.
+
+Accepted inputs are raw amino-acid sequences through the convenience API, or
+prepared model features.
+Supported Transformers entry points are `AutoConfig`, `AutoModel`.
+
+## Install and platform requirements
+
+Install FastPLMs from the exact source revision paired with this model card:
+
+```bash
+python -m pip install \
+  "fastplms[structure] @ git+https://github.com/Synthyra/FastPLMs.git@1b9ce023f1e06571cf3e6324be0610ffa53e0a4a"
+```
+
+Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. Structure inference requires the `structure` extra and a CUDA device for the published execution contract. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network
+access on first download. For an air-gapped run, first build the manifest-pinned
+local artifact and use the offline form shown in the example.
+
+## Quick start
+
+```python
+from transformers import AutoModel
+
+model_id = "Synthyra/Boltz2"
+model = AutoModel.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+).eval()
+```
+
+This example uses the published Hub repository. For offline validation, build
+the manifest-pinned artifact and replace `model_id` with its local
+`dist/hub/Boltz2` path, then pass `local_files_only=True`.
+
+Leave attention unspecified for the Transformers default. Supported explicit
+choices are `eager`.
+Pass the selected name through `attn_implementation`.
+When an optimized backend cannot return full attention tensors,
+`output_attentions=True` emits one explicit runtime warning and uses a correctly
+masked eager implementation for that call only. The warning identifies the
+configured backend, effective backend, and reason. Configuration and later
+calls are unchanged.
+For BF16 execution, this family uses FP32 parameters with CUDA BF16 autocast.
+
+## Protein structure prediction
+
+The high-level helper prepares a protein-only input, runs the declared Boltz2
+inference core, and returns coordinates and confidence fields:
+
+```python
+import torch
+
+model = model.cuda().eval()
+output = model.predict_structure(
+    amino_acid_sequence="MSTNPKPQRKTKRNTNRRPQDVKFPGG",
+    recycling_steps=3,
+    num_sampling_steps=50,
+    diffusion_samples=1,
+    seed=7,
+)
+model.save_as_cif(output, "prediction.cif")
+
+print(output.sample_atom_coords.shape)
+print(output.plddt, output.ptm, output.iptm)
+```
+
+The validation boundary below describes the currently supported inference
+subset and its provisional status. The helper scopes and restores Python,
+NumPy, CPU Torch, and CUDA RNG state. Parameters and prepared features remain
+FP32; supported CUDA inference executes inside BF16 autocast.
+
+## Notes and limitations
+
+Boltz2 is provisional in FastPLMs 1.0. Exact configuration, the declared
+inference-core state, feature preparation, and seeded execution remain tested,
+but native-environment BF16 end-to-end inference currently exceeds the fixed
+numerical-equivalence limits. FastPLMs therefore does not claim official
+inference equivalence for this checkpoint yet. Work on that numerical gap
+continues independently of the ESM++ and ESMFold2 release gates.
+
+## Runtime contract
+
+- Public input: Raw amino-acid sequences through the convenience API, or prepared model features
+- Advertised AutoClasses: `AutoConfig`, `AutoModel`
+- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`
+- Attention implementations: `eager`
+- Precision policies: `default`
+- BF16 execution: `fp32_parameters_autocast`
+- Generation contract: `not_applicable`
+- Optional dependency group: `structure`
+- Weight publication allowed: `true`
+- Weight license status: `resolved`
+- Redistributable: `true`
+- Complete weight publication required: `false`
+
+## Provenance
+
+- FastPLMs weights: `Synthyra/Boltz2@3b148fc5efea109c065ec82ba8683d024de7134e`
+- Runtime revision: `1b9ce023f1e06571cf3e6324be0610ffa53e0a4a`
+- Runtime source-tree SHA-256: `21579947908f914e1b38ff1e69f9df9d22b7798899090355a53691a937b024e5`
+- Runtime bundle SHA-256: `297af2aeeb817b6e3a343bf73dbc2eb3b3e29a2ca88ccba8af1aeb2e12243ffd`
+- Generator/schema version and complete/runtime-only attestations: recorded in `provenance.json`
+- Official checkpoint: `boltz-community/boltz-2@6fdef46d763fee7fbb83ca5501ccceff43b85607`
+- Artifact source: `fast`
+- State transform: `boltz2_inference_core_v1`
+- BF16 execution: `fp32_parameters_autocast`
+- Pinned upstreams: `boltz`
+- Reference container: `reference-boltz2`
+- Release tiers: `structure`, `artifact`, `benchmark`
+- Unresolved required file identities: `0`
+
+The local artifact records exact file identities, conversion provenance, source
+revisions, and legal texts in `provenance.json`. A nonzero unresolved count is a
+release blocker.
+
+## Validation boundary
+
+For tiers declared by the manifest, the release contract compares applicable
+semantic configuration, tokenizer behavior, state keys, shapes, dtypes,
+values, aliases, and representative inference with the pinned official
+implementation. This metadata does not by itself claim that a particular build
+passed, that one backend is faster, or that an output has biological or
+therapeutic validity.
+
+## License
+
+Checkpoint terms: MIT. The Hub model-card identifier is
+`mit`. Applicable source licenses, notices, attribution,
+and conversion records are distributed with the local artifact. Review them
+before use.

THIRD_PARTY_NOTICES.md

@@ -0,0 +1,99 @@
+# Third-party notices
+
+FastPLMs implements interfaces and checkpoint mappings for independently
+released protein models. The pinned repositories under `vendor/upstream/` are
+parity oracles. Production code does not import them, and runtime images do not
+contain them.
+
+This notice is informational and is not legal advice. A checkpoint license can
+differ from the license covering its source implementation. The typed inventory
+in `src/fastplms/models.toml` and the verbatim files under `LICENSES/` are the
+distribution record.
+
+## ANKH
+
+The pinned ANKH implementation and the mirrored ANKH checkpoints are identified
+as CC BY-NC-SA 4.0. FastPLMs displays those terms but does not enforce them in
+software. Users are responsible for determining whether their use and
+redistribution comply. The complete text is in `LICENSES/ankh/LICENSE.md`.
+
+## Profluent-E1
+
+Profluent identifies its E1 model code as Apache-2.0. The E1 weights and full
+release are subject to the Profluent-E1 Clickthrough License Agreement and the
+incorporated attribution requirements. Any E1 distribution must retain all of
+the following files:
+
+- `LICENSES/e1/LICENSE`, the Profluent-E1 agreement
+- `LICENSES/e1/ATTRIBUTION`, the attribution guidelines
+- `LICENSES/e1/NOTICE`, the required notice
+- `LICENSES/e1/Apache-2.0.txt`, the code license
+- `LICENSES/e1/BSD-3-Clause.txt`, covering the FlashAttention-derived padding
+  utility identified by the official E1 source
+- `LICENSES/e1/MODIFICATIONS.md`, the FastPLMs modified-file notice
+
+The exact text `Profluent-E1` must remain prominently displayed in E1
+documentation and at each launch of an executable E1 workflow, as required by
+the upstream attribution guidelines. Certain commercial outputs, including
+specified pharmaceutical and target-related outputs, can require the separate
+`Built with Profluent-E1` statement described in `ATTRIBUTION`.
+
+## DPLM
+
+The pinned ByteDance DPLM repository is Apache-2.0. Its
+[README](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/README.md#overview)
+explicitly defines the repository release as including pretrained DPLM1 and
+DPLM2 weights, and the same revision carries the complete
+[Apache-2.0 license](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/LICENSE).
+FastPLMs records both checkpoint families as Apache-2.0 and distributes the
+verbatim license plus `LICENSES/dplm/PROVENANCE.md`. Converted weights retain
+those terms and remain subject to the ordinary artifact and publication gates.
+
+## Biohub
+
+The pinned Biohub ESM implementation is MIT and includes a separate
+`THIRD_PARTY_NOTICE.md`; both files are distributed under
+`LICENSES/biohub-esm/`. The pinned Biohub Transformers fork is Apache-2.0, with
+its complete text under `LICENSES/biohub-transformers/`.
+
+## Boltz
+
+The pinned Boltz source is MIT. The verbatim notice is in
+`LICENSES/boltz/LICENSE`.
+
+## Meta ESM and OpenFold
+
+The pinned Meta ESM source is MIT. The pinned OpenFold source is Apache-2.0.
+Their verbatim texts and revision-specific provenance notices are under
+`LICENSES/fair-esm/` and `LICENSES/openfold/`.
+
+The native H100 ESMFold reference image applies the tracked
+`docker/constraints/openfold-sm90.patch` to the copied OpenFold `setup.py`.
+This build-only change restricts the CUDA extension to `sm90` and selects the
+C++17 standard required by the reference PyTorch version. It leaves the pinned
+submodule, extension source, model classes, checkpoint data, and public API
+unchanged. The complete modified-file record is in
+`LICENSES/openfold/MODIFICATIONS.md`.
+
+The isolated reference image also includes Apache-2.0 PyTorch Lightning,
+TorchMetrics, Lightning Utilities, and NVIDIA DLLogger. Their exact versions or
+revision are pinned in `docker/constraints/esmfold.txt`; OpenFold imports them
+eagerly, and FastPLMs production code does not depend on them. DLLogger's exact
+source identity and installed-license handling are recorded in
+`LICENSES/dllogger/PROVENANCE.md`.
+
+## ProteinTTT
+
+The optional test-time training workflow is validated against the pinned
+ProteinTTT repository under its MIT license. Its verbatim license and
+revision-specific provenance are under `LICENSES/protein-ttt/`.
+
+## Conversion and packaging record
+
+For every supported family, `src/fastplms/models.toml` records an immutable
+official checkpoint revision, an immutable FastPLMs checkpoint revision, file
+digests, a named state transformation, and a mechanism-level conversion record.
+Generated artifacts reproduce that record in `provenance.json`. A release or
+artifact build must fail when a required file identity, legal text, attribution
+notice, modified-file notice, upstream revision, or conversion record is absent
+or differs from its manifest digest.

config.json

@@ -3,8 +3,8 @@
     "Boltz2Model"
   ],
   "auto_map": {
-    "AutoConfig": "modeling_boltz2.Boltz2Config",
-    "AutoModel": "modeling_boltz2.Boltz2Model"
+    "AutoConfig": "modeling_fastplms.Boltz2Config",
+    "AutoModel": "modeling_fastplms.Boltz2Model"
   },
   "core_kwargs": {
     "alpha_pae": 1,
@@ -133,6 +133,16 @@
   "default_recycling_steps": 3,
   "default_sampling_steps": 200,
   "dtype": "float32",
+  "fastplms_checkpoint_hash": "19fdd55cb04969d3e63f010a4655b73e52024bc88616f34bd0b5b38d74d303ca",
+  "fastplms_checkpoint_repo_id": "Synthyra/Boltz2",
+  "fastplms_checkpoint_revision": "3b148fc5efea109c065ec82ba8683d024de7134e",
+  "fastplms_model_id": "boltz2",
+  "fastplms_release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "fastplms_runtime_bundle_sha256": "297af2aeeb817b6e3a343bf73dbc2eb3b3e29a2ca88ccba8af1aeb2e12243ffd",
+  "fastplms_runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_source_tree_sha256": "21579947908f914e1b38ff1e69f9df9d22b7798899090355a53691a937b024e5",
+  "fastplms_weights_revision": "3b148fc5efea109c065ec82ba8683d024de7134e",
   "model_type": "boltz2_automodel",
   "num_bins": 64,
   "transformers_version": "4.57.6"

fastplms/__init__.py

@@ -0,0 +1,48 @@
+"""FastPLMs public package interface.
+
+The module uses lazy exports so importing :mod:`fastplms` does not initialize
+Torch, download checkpoints, construct tokenizers, or compile kernels.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Any
+
+__version__ = "1.0.0"
+
+_LAZY_EXPORTS = {
+    "CheckpointSource": ("fastplms.registry", "CheckpointSource"),
+    "EmbeddingInput": ("fastplms.embeddings", "EmbeddingInput"),
+    "EmbeddingRecord": ("fastplms.embeddings", "EmbeddingRecord"),
+    "EmbeddingResult": ("fastplms.embeddings", "EmbeddingResult"),
+    "FileDigest": ("fastplms.registry", "FileDigest"),
+    "ModelFamily": ("fastplms.registry", "ModelFamily"),
+    "ModelRegistry": ("fastplms.registry", "ModelRegistry"),
+    "ModelSpec": ("fastplms.registry", "ModelSpec"),
+    "OracleAsset": ("fastplms.registry", "OracleAsset"),
+    "RegistryError": ("fastplms.registry", "RegistryError"),
+    "RuntimeProfile": ("fastplms.runtime", "RuntimeProfile"),
+    "UpstreamSource": ("fastplms.registry", "UpstreamSource"),
+    "embed_dataset": ("fastplms.embeddings", "embed_dataset"),
+    "get_model_registry": ("fastplms.registry", "get_model_registry"),
+    "get_model_spec": ("fastplms.registry", "get_model_spec"),
+    "load_model_registry": ("fastplms.registry", "load_model_registry"),
+    "runtime_profile": ("fastplms.runtime", "runtime_profile"),
+}
+
+__all__ = ["__version__", *_LAZY_EXPORTS]
+
+
+def __getattr__(name: str) -> Any:
+    try:
+        module_name, attribute_name = _LAZY_EXPORTS[name]
+    except KeyError as error:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from error
+    value = getattr(import_module(module_name), attribute_name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()).union(__all__))

fastplms/models.toml

@@ -0,0 +1,1223 @@
+schema_version = 1
+legal_files = [
+  "LICENSE=sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
+  "THIRD_PARTY_NOTICES.md=sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
+]
+
+[[attention_kernels]]
+implementation = "flash_attention_2"
+repository = "kernels-community/flash-attn2"
+revision = "db6b51744f0cd7061386442c09df890fc6d9f47e"
+version = 2
+expected_variant = "flash_attn2"
+dtypes = ["bfloat16"]
+
+[[attention_kernels]]
+implementation = "flash_attention_3"
+repository = "kernels-community/flash-attn3"
+revision = "43f0bd269777115d94ff826e0d113ce9c1c9087b"
+version = 1
+expected_variant = "flash_attn3"
+dtypes = ["bfloat16"]
+
+[[runtime_assets]]
+id = "esmfold2_ccd"
+repository = "biohub/ESMFold2"
+revision = "1ebf0e3481a5184eb6171d40615c79e384b48796"
+path = "ccd.pkl"
+sha256 = "9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5"
+size = 417306584
+consumer_family = "esmfold2"
+trust_kind = "hash_pinned_pickle"
+license = "MIT"
+offline_behavior = "requires_cached_verified_file"
+
+[[upstreams]]
+id = "ankh"
+path = "vendor/upstream/ankh"
+url = "https://github.com/agemagician/Ankh.git"
+revision = "02b4e25ce5389b9e771c9df6e546c62af1216f8e"
+license = "CC-BY-NC-SA-4.0"
+license_files = ["LICENSE.md"]
+license_digests = ["LICENSE.md=sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088"]
+distribution_files = ["LICENSE.md=sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088"]
+
+[[upstreams]]
+id = "biohub-esm"
+path = "vendor/upstream/biohub-esm"
+url = "https://github.com/Biohub/esm.git"
+revision = "82ee35553d39169d678f784c8d3f8712ffd7d2c4"
+license = "MIT"
+license_files = ["LICENSE.md", "THIRD_PARTY_NOTICE.md"]
+license_digests = [
+  "LICENSE.md=sha256:b63df9ca1dd96b3b21eec226b51b236d0bd152ac20eafc43aad46bf832b48d8a",
+  "THIRD_PARTY_NOTICE.md=sha256:5bff8515ba4e0f53abdc43714c180b79c5b606160497d98de741a369cb9b6a23",
+]
+distribution_files = [
+  "LICENSE.md=sha256:b63df9ca1dd96b3b21eec226b51b236d0bd152ac20eafc43aad46bf832b48d8a",
+  "THIRD_PARTY_NOTICE.md=sha256:5bff8515ba4e0f53abdc43714c180b79c5b606160497d98de741a369cb9b6a23",
+]
+
+[[upstreams]]
+id = "biohub-transformers"
+path = "vendor/upstream/biohub-transformers"
+url = "https://github.com/Biohub/transformers.git"
+revision = "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf"
+license = "Apache-2.0"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049"]
+distribution_files = ["LICENSE=sha256:77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049"]
+
+[[upstreams]]
+id = "boltz"
+path = "vendor/upstream/boltz"
+url = "https://github.com/jwohlwend/boltz.git"
+revision = "b1ebfc46ecf57f5414e0d1a6f9027bbb122c53bc"
+license = "MIT"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:f0667fd5e66c51e1ba8ddaa0249c6d7225b30037e02c45782d8f2c2943ac2617"]
+distribution_files = ["LICENSE=sha256:f0667fd5e66c51e1ba8ddaa0249c6d7225b30037e02c45782d8f2c2943ac2617"]
+
+[[upstreams]]
+id = "dplm"
+path = "vendor/upstream/dplm"
+url = "https://github.com/bytedance/dplm.git"
+revision = "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d"
+license = "Apache-2.0"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"]
+distribution_files = [
+  "LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+  "PROVENANCE.md=sha256:a659f74be9073cf1ad2d2f7071531ca56959b421f111152cf4c41184ace5970e",
+]
+
+[[upstreams]]
+id = "e1"
+path = "vendor/upstream/e1"
+url = "https://github.com/Profluent-AI/E1.git"
+revision = "bfd2620a602248499f3d2583d85a7ecddf0b6e02"
+license = "Apache-2.0 AND Profluent-E1-Agreement"
+license_files = ["LICENSE", "ATTRIBUTION", "NOTICE"]
+license_digests = [
+  "LICENSE=sha256:8ef1dd556091544db3044164a8015424a3dcb3450fb3765a81b88463551bbe81",
+  "ATTRIBUTION=sha256:deb22b250f6491b649eda5c63e080dd56486b8d2736cea6a52ef875436214367",
+  "NOTICE=sha256:6de9db0320b4ee82f665c0951d8fd4cd53701a659c9dbce9bc3e3ea6afc4c6b3",
+]
+distribution_files = [
+  "LICENSE=sha256:8ef1dd556091544db3044164a8015424a3dcb3450fb3765a81b88463551bbe81",
+  "ATTRIBUTION=sha256:deb22b250f6491b649eda5c63e080dd56486b8d2736cea6a52ef875436214367",
+  "NOTICE=sha256:6de9db0320b4ee82f665c0951d8fd4cd53701a659c9dbce9bc3e3ea6afc4c6b3",
+  "Apache-2.0.txt=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+  "BSD-3-Clause.txt=sha256:36e1987f2f17db7f8ad36cd7a37dbb7aeaaf0ab68b97ab4b9d3556f3a7a76ae8",
+  "MODIFICATIONS.md=sha256:2506f47c0f5475af8e8ff2cff13eb8b79e8e25a08a054cdd617bf336536750ca",
+]
+
+[[upstreams]]
+id = "fair-esm"
+path = "vendor/upstream/fair-esm"
+url = "https://github.com/facebookresearch/esm.git"
+revision = "2b369911bb5b4b0dda914521b9475cad1656b2ac"
+license = "MIT"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93"]
+distribution_files = [
+  "LICENSE=sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93",
+  "PROVENANCE.md=sha256:950adb94daf15e646ddf226dacfe2a8e77801aa0793e439a9a3490a48eb666e7",
+]
+
+[[upstreams]]
+id = "openfold"
+path = "vendor/upstream/openfold"
+url = "https://github.com/aqlaboratory/openfold.git"
+revision = "4b41059694619831a7db195b7e0988fc4ff3a307"
+license = "Apache-2.0"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"]
+distribution_files = [
+  "LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+  "MODIFICATIONS.md=sha256:fd6f0aa1086a0c996cf967b326d18e965660cda0ad5c7f36a3474a8490720da3",
+  "PROVENANCE.md=sha256:48c903db43a217a3126afaefbac60b7ddac7efda2dfcc0cbff0bffc7d6c30081",
+]
+
+[[upstreams]]
+id = "protein-ttt"
+path = "vendor/upstream/protein-ttt"
+url = "https://github.com/anton-bushuiev/ProteinTTT.git"
+revision = "fde2817cd84b936167cc76ccabf31e5c0fe49962"
+license = "MIT"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df"]
+distribution_files = [
+  "LICENSE=sha256:bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df",
+  "PROVENANCE.md=sha256:dc641c37353c2efd50ccbdb316ca4aae495ec02c1563e0e15bac92f75fc482e5",
+]
+
+[families.esm2]
+architecture = "ESM2"
+upstreams = ["fair-esm"]
+tokenizer_mode = "tokenizer"
+public_input = "Amino-acid sequences tokenized to residue IDs"
+extra = "core"
+reference_container = "reference-esm2"
+reference_adapter = "tests.parity.support.reference_adapters.esm2"
+attention = ["eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "esm2_hf_to_fastplms_v1"
+conversion_provenance = "Input: the pinned official ESM2 state dictionary. Transformation: apply the deterministic esm2_hf_to_fastplms_v1 key map while preserving tensor values and materializing the tied input/output embedding values as independent tensors. Output: the pinned Synthyra FastPLMs checkpoint. Validation: release parity compares exact keys and values after the declared non-aliasing transform, tokenizer behavior, and inference. Limitation: any numerical rewrite requires a new transform identifier and exact conversion test."
+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"]
+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]
+architecture = "ESMC"
+upstreams = ["biohub-esm", "biohub-transformers"]
+tokenizer_mode = "tokenizer"
+public_input = "Amino-acid sequences tokenized to residue IDs"
+extra = "core"
+reference_container = "reference-biohub-esm"
+reference_adapter = "tests.parity.support.reference_adapters.esm_plusplus"
+attention = ["eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "static_parameters"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "esmc_to_fastplms_v1"
+conversion_provenance = "Input: the pinned Biohub ESMC checkpoint. Transformation: apply the deterministic esmc_to_fastplms_v1 parameter map into the FastPLMs ESMC modules. Output: the pinned Synthyra ESMplusplus checkpoint. Validation: release parity compares keys, shapes, dtypes, values, aliases, and live inference. Limitation: runtime attention and precision selection are not serialized weight transforms."
+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"]
+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" }
+
+[families.esm3]
+architecture = "ESM3"
+upstreams = ["biohub-esm", "biohub-transformers"]
+tokenizer_mode = "tokenizer"
+public_input = "Sequence, structure, and function tracks prepared through the multimodal helpers"
+extra = "core"
+reference_container = "reference-biohub-esm"
+reference_adapter = "tests.parity.support.reference_adapters.esm3"
+attention = ["eager", "sdpa", "flex_attention"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "large-sequence"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "esm3_to_fastplms_v1"
+conversion_provenance = "Input: the pinned Biohub ESM3 checkpoint. Transformation: apply the deterministic esm3_to_fastplms_v1 parameter map for the supported sequence and multimodal modules and expand BF16 checkpoint tensors to FP32 storage. Output: the pinned Synthyra ESM3 checkpoint. Validation: release parity compares exact state identity after the declared map and live feature behavior. Limitation: unsupported upstream modalities may not be inferred from this record."
+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"]
+auto_map = { AutoConfig = "fastplms.models.esm3.modeling_esm3.FastESM3Config", AutoModel = "fastplms.models.esm3.modeling_esm3.FastESM3Model" }
+
+[families.e1]
+architecture = "E1"
+upstreams = ["e1"]
+tokenizer_mode = "sequence"
+public_input = "Raw amino-acid sequences prepared by the native E1 adapter"
+extra = "core"
+reference_container = "reference-e1"
+reference_adapter = "tests.parity.support.reference_adapters.e1"
+attention = ["sdpa", "flex_attention"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "static_parameters"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "Profluent-E1-Agreement"
+hub_license = "other"
+hub_license_name = "Profluent-E1 Clickthrough License Agreement"
+hub_license_link = "https://github.com/Profluent-AI/E1/blob/bfd2620a602248499f3d2583d85a7ecddf0b6e02/LICENSE"
+weights_publication_allowed = true
+state_transform = "e1_to_fastplms_v1"
+conversion_provenance = "Input: the pinned Profluent-E1 checkpoint and tokenizer-free sequence contract. Transformation: apply e1_to_fastplms_v1 to the FastPLMs encoder and official task heads, storing floating tensors in BF16. Output: the pinned Synthyra Profluent-E1 checkpoint. Validation: release parity covers state identity after the declared cast, sequence and RAG preparation, aliases, and inference. Limitation: the FastPLMs scoring extension is not represented as an official E1 head."
+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"]
+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]
+architecture = "DPLM"
+upstreams = ["dplm"]
+tokenizer_mode = "tokenizer"
+public_input = "Amino-acid sequences tokenized to masked or partially masked residue IDs"
+extra = "core"
+reference_container = "reference-dplm"
+reference_adapter = "tests.parity.support.reference_adapters.dplm"
+attention = ["eager", "sdpa", "flex_attention", "flash_attention_3"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "Apache-2.0"
+hub_license = "apache-2.0"
+weights_publication_allowed = true
+state_transform = "dplm_to_fastplms_v1"
+conversion_provenance = "Input: the pinned official DPLM1 checkpoint. Transformation: apply dplm_to_fastplms_v1, omitting the unused absolute-position table for rotary checkpoints and materializing the tied input/output embedding values as independent tensors. Output: the pinned Synthyra DPLM checkpoint. Validation: release parity compares exact state identity after the declared transform, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/PROVENANCE.md. Limitation: redistribution remains subject to Apache-2.0 and the pinned provenance record; no broader rights are inferred."
+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"]
+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]
+architecture = "DPLM2"
+upstreams = ["dplm"]
+tokenizer_mode = "tokenizer"
+public_input = "Tokenized amino-acid and structure tracks with explicit modality boundaries"
+extra = "core"
+reference_container = "reference-dplm"
+reference_adapter = "tests.parity.support.reference_adapters.dplm2"
+attention = ["sdpa"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "Apache-2.0"
+hub_license = "apache-2.0"
+weights_publication_allowed = true
+state_transform = "dplm2_to_fastplms_v1"
+conversion_provenance = "Input: the pinned official DPLM2 checkpoint. Transformation: apply dplm2_to_fastplms_v1, retaining the independent language-model head and trained encoder contact head while omitting the unused absolute-position table for rotary checkpoints. Output: the pinned Synthyra DPLM2 checkpoint. Validation: release parity compares exact keys and values after the declared omission, non-aliasing, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/PROVENANCE.md. Limitation: no head exception is permitted by this record, and redistribution remains subject to Apache-2.0."
+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"]
+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"
+
+[families.ankh]
+architecture = "ANKH"
+upstreams = ["ankh"]
+tokenizer_mode = "tokenizer"
+public_input = "Amino-acid sequences tokenized for encoder or sequence-to-sequence use"
+extra = "core"
+reference_container = "reference-ankh"
+reference_adapter = "tests.parity.support.reference_adapters.ankh"
+attention = ["eager", "sdpa"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "static_parameters"
+precisions = ["default"]
+vram_tier = "large-sequence"
+checkpoint_license = "CC-BY-NC-SA-4.0"
+hub_license = "cc-by-nc-sa-4.0"
+weights_publication_allowed = true
+state_transform = "ankh_t5_to_fastplms_v1"
+conversion_provenance = "Input: the pinned official ANKH T5 checkpoint. Transformation: apply ankh_t5_to_fastplms_v1 to the official encoder and sequence-to-sequence heads. Output: the pinned Synthyra ANKH checkpoint. Validation: release parity compares exact mapped state, tokenizer behavior, official heads, and inference. Limitation: the separately named FastPLMs masked-language-model extension is not an official ANKH head."
+representative = "ankh_base"
+documentation = "docs/models.md#ankh"
+test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
+requires_complete_weight_publication = true
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "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]
+architecture = "Boltz2"
+upstreams = ["boltz"]
+tokenizer_mode = "structure"
+public_input = "Raw amino-acid sequences through the convenience API, or prepared model features"
+extra = "structure"
+reference_container = "reference-boltz2"
+reference_adapter = "tests.parity.support.reference_adapters.boltz"
+attention = ["eager"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "structure"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "boltz2_inference_core_v1"
+conversion_provenance = "Input: the pinned official Boltz2 checkpoint. Transformation: select and map the supported Boltz2 inference-core parameters with boltz2_inference_core_v1. Output: the pinned Synthyra Boltz2 checkpoint. Validation: release parity covers state identity for the declared subset, feature preparation, seeded inference, and structure outputs. Limitation: this record does not claim support for undeclared upstream training components."
+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"]
+auto_map = { AutoConfig = "fastplms.models.boltz.modeling_boltz2.Boltz2Config", AutoModel = "fastplms.models.boltz.modeling_boltz2.Boltz2Model" }
+
+[families.esmfold]
+architecture = "ESMFold"
+upstreams = ["fair-esm", "openfold"]
+tokenizer_mode = "structure"
+public_input = "Raw amino-acid sequences through folding helpers, or prepared residue tensors"
+extra = "structure"
+reference_container = "reference-esmfold"
+reference_adapter = "tests.parity.support.reference_adapters.esmfold"
+attention = ["eager", "sdpa", "flex_attention"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "structure"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "esmfold_meta_to_fastplms_v1"
+conversion_provenance = "Input: the pinned native Meta ESMFold checkpoint plus its pinned ESM2 backbone. Transformation: apply esmfold_meta_to_fastplms_v1 to map native ESM2 names into the structure-only FastPLMs backbone, retain folding tensors, omit five deterministically reconstructed geometry buffers, omit the folding-unused ESM2 masked-LM and contact-regression heads, and remove the obsolete random FastPLMs TTT head from earlier mirrors. Output: canonical FP32 FastPLMs ESMFold state with an explicit CUDA BF16-autocast execution path. Validation: release parity compares exact mapped keys, shapes, dtypes, values, aliases, semantic configuration, FP32 and BF16-compute seeded inference, and structure metrics with pLDDT normalized to (0, 1). Limitation: ESMFold TTT is rejected because the official checkpoint contains no trained masked-language-model head."
+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/esmfold"]
+auto_map = { AutoConfig = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig", AutoModel = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding" }
+
+[families.esmfold2]
+architecture = "ESMFold2"
+upstreams = ["biohub-esm", "biohub-transformers", "protein-ttt"]
+backbone_model = "esmc_6b"
+tokenizer_mode = "structure"
+public_input = "Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors"
+extra = "structure"
+reference_container = "reference-esmfold2"
+reference_adapter = "tests.parity.support.reference_adapters.esmfold2"
+attention = ["eager", "sdpa", "flex_attention"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["auto", "fp32", "bf16", "fp8"]
+experimental_precisions = ["fp8"]
+vram_tier = "structure-6b"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "identity"
+conversion_provenance = "Input: each pinned Biohub ESMFold2 checkpoint and its separately pinned ESMC checkpoint. Transformation: apply identity to preserve the folding checkpoint exactly, load its parameters in FP32 for CUDA BF16-autocast execution, retain canonical BF16 ESMC weights, and optionally rebuild exactly 80 ESMC attention output projections as transient Transformer Engine linears. Output: the corresponding pinned Synthyra ESMFold2 checkpoint plus its declared ESMC precision policy. Validation: release parity covers exact canonical state, learned projection, prepared features, and seeded BF16 folding; experimental FP8 validation covers strict unavailable-device behavior, all four variants, and three BF16-to-FP8 reload cycles on the standard variant. Limitation: only the four manifest-listed ESMFold2 variants are supported; FP8 is experimental, applies only to inference-time ESMC execution, and requires direct CUDA loading with Transformer Engine availability."
+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/esmfold2", "models/esm_plusplus", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model" }
+
+[[models]]
+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" }
+fast_repo = "Synthyra/ESM2-8M"
+fast_revision = "185ecbd45665d050a8dae326d91886d330c5f9d0"
+fast_files = [
+  "config.json=git-sha1:46d0a7b517f59123c6ebc6d1011585731cbab259",
+  "model.safetensors=sha256:c824e6ded5fb71c72bc5ac05300699947819023cb26cdaf6897665e6b2645e1b",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t6_8M_UR50D"
+official_revision = "c731040fcd8d73dceaa04b0a8e6329b345b0f5df"
+official_files = [
+  "config.json=git-sha1:c2c6e65a87d9d20d47699ae236d605b80c741dd3",
+  "model.safetensors=sha256:24c5fa474c48f3b754b86efe752d5f189d2bcd88190fa2270fc92b2ef3034189",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t6_8M_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t6_8M_UR50D.pt"
+sha256 = "46f002a9870c9bdecd0ea887acb1f9a38a6b561e8f8bf8a6990b679b9d31b928"
+size = 30099493
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t6_8M_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t6_8M_UR50D-contact-regression.pt"
+sha256 = "8f7a4557d57713b97ba0e484303007efb7230d25299c0ac47a0a1b12a87bbb9d"
+size = 1511
+
+[[models]]
+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" }
+fast_repo = "Synthyra/ESM2-35M"
+fast_revision = "37ab9f56b41e365b3bd9e25d6fefe9150fd910f0"
+fast_files = [
+  "config.json=git-sha1:4d428c9934572f39e2a00db162249971f37c88e4",
+  "model.safetensors=sha256:21d95ab6bb9aa91bfec87eff11da61a657b732f2df279cbddbae6a7f1f0bba9c",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t12_35M_UR50D"
+official_revision = "6fbf070e65b0b7291e7bbcd451118c216cff79d8"
+official_files = [
+  "config.json=git-sha1:3f64131bb610ed1ce482c4b5421fc358c785278f",
+  "model.safetensors=sha256:e35647818e0e064351d4531ed480d225a002567b4b2b93ad3a9246d753150fc0",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t12_35M_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t12_35M_UR50D.pt"
+sha256 = "7f21e80e61d16a71735163ef555d3009afb0c98da74c48e29df08606973cc55e"
+size = 134095705
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t12_35M_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t12_35M_UR50D-contact-regression.pt"
+sha256 = "16641e05d830d0ce863dd152dbb8c2f3ddfa3c3ec2a66080152c8abad01d8585"
+size = 1959
+
+[[models]]
+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" }
+fast_repo = "Synthyra/ESM2-150M"
+fast_revision = "979e0880dfc9e0c0080839b83d9d2dc05b92786a"
+fast_files = [
+  "config.json=git-sha1:efeae2af182b7d34dc35740a45f157661e7acdf4",
+  "model.safetensors=sha256:d1f7c60f98c31af328381519a750972b6a31b13b97aa7cca2e71b5ae1b3f8f53",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t30_150M_UR50D"
+official_revision = "a695f6045e2e32885fa60af20c13cb35398ce30c"
+official_files = [
+  "config.json=git-sha1:52e04179e6fbad6663a94ea5cc44f09d764c5cd4",
+  "model.safetensors=sha256:c3f1da8aea53bddd32c246c86168c23b9fd72341fb9db9a94436f855f5053566",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t30_150M_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t30_150M_UR50D.pt"
+sha256 = "881c7176cf198ef8dec26a3c375d40eb58d0c33df95c22562ca6cc6d3f812c62"
+size = 592774773
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t30_150M_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t30_150M_UR50D-contact-regression.pt"
+sha256 = "6a604b96722ed052eef8a094ad90b275ba2e987d406315dbed0bdc6b3c4238a7"
+size = 3431
+
+[[models]]
+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" }
+fast_repo = "Synthyra/ESM2-650M"
+fast_revision = "ca0718a5d52b80d5c60dd76860e55e061a95fb0a"
+fast_files = [
+  "config.json=git-sha1:88f6bd240680b29c3244df8292246048401f5caf",
+  "model.safetensors=sha256:a15142e94ecf36f0edde9b37796f591e609ebe1694ca411e93640f0ee384994a",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t33_650M_UR50D"
+official_revision = "08e4846e537177426273712802403f7ba8261b6c"
+official_files = [
+  "config.json=git-sha1:a956a25d277f30bd870d3760b9a116f19ead885e",
+  "model.safetensors=sha256:a08adabb949fa67ad3c14b509d04fd60368b35007b0095e3358f81200c4f4db0",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t33_650M_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt"
+sha256 = "ea9d0522b335a8778dea6535a65301f10208dece28cd5865482b0b1fc446168c"
+size = 2604537549
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t33_650M_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t33_650M_UR50D-contact-regression.pt"
+sha256 = "8ffe6edbd4173dc8d45c2cd5cb27d43aad77ec26b4c768200c58ae1f96693575"
+size = 3687
+
+[[models]]
+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" }
+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"
+fast_files = [
+  "config.json=git-sha1:94944ad6cabaa40a3ce1cbe6699cf464fdc1b2c0",
+  "model-00001-of-00003.safetensors=sha256:04b57854545c23779b562ee2ae22f10021ba0f4d586ba0ad482ee6eda187d562",
+  "model-00002-of-00003.safetensors=sha256:34954aaa05bc91635776ba6672946da5822626753d80db97b38c0538e9525102",
+  "model-00003-of-00003.safetensors=sha256:a6b3a55b9e3b2e1778de34c665c3dd17bdfdf6da9d6d5c97730c57168709ccae",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t36_3B_UR50D"
+official_revision = "476b639933c8baad5ad09a60ac1a87f987b656fc"
+official_files = [
+  "config.json=git-sha1:69e7563923f87d2d7439bfb83e5a19b44b46d71b",
+  "pytorch_model-00001-of-00002.bin=sha256:0f971f11c449d21422aa982b791619c10351972992c735f4c3cd43fe09790412",
+  "pytorch_model-00002-of-00002.bin=sha256:7560b46fc383c691fb74b915b7d4bcef40d3df181447f16ba4b298845e308d0c",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t36_3B_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt"
+sha256 = "7de8b4082ba15891959ab368b77ce3886697af1efb16d3c9e9e7b0c5d3f07500"
+size = 5678116398
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t36_3B_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t36_3B_UR50D-contact-regression.pt"
+sha256 = "4da500eab246481dc9c8c95bc7b1d02f2803d761c380b0e95186d4a07d0fc84e"
+size = 6759
+
+[[models]]
+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" }
+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"
+fast_files = [
+  "config.json=git-sha1:df2f44187157b0cc371c48c887b77b1783679201",
+  "model.safetensors=sha256:d099223765bc4f1ae8d6c7e18561ce41df1d54073fdc5327ef0a229235a8f52a",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
+  "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
+]
+official_repo = "biohub/ESMC-300M"
+official_revision = "a59b831785f907e96e6a246b1d142bfb76df31ee"
+official_files = [
+  "config.json=git-sha1:9a49eacf4e65c39f74381f0f0d240e3b89ef43d7",
+  "model.safetensors=sha256:0772d8fe64bb25e14fe6f23b80e3c9a7d215d0da3c6cba5bd356d7c0e0bb22cc",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c",
+  "tokenizer_config.json=git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61",
+]
+
+[[models]]
+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" }
+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"
+fast_files = [
+  "config.json=git-sha1:5736371902fe5d04e2859be30ac7dbd31b271b25",
+  "model.safetensors=sha256:4aff3f8c5de68c4d3e3824eb2c478e4a47355d3f849f3c745e5c8a5ee6cff851",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
+  "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
+]
+official_repo = "biohub/ESMC-600M"
+official_revision = "a7e82012c83126b9eedb055fea9fa84b6c02f094"
+official_files = [
+  "config.json=git-sha1:71c8241dc28a5fb636248267a0927c0242b264c1",
+  "model.safetensors=sha256:e4232c30fd35fe2f57051ec88a703996ac94520580b4b836894207a3d45d9ff8",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c",
+  "tokenizer_config.json=git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61",
+]
+
+[[models]]
+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" }
+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"
+fast_files = [
+  "config.json=git-sha1:e740cbcf211f2511c70c25a1ff6017a757ba7a69",
+  "model-00001-of-00006.safetensors=sha256:d30d18703453019f2d2d050866309888720c28eebc9a10307d1ddf3799e85a65",
+  "model-00002-of-00006.safetensors=sha256:b3d85378ab5023f4160a96e9c8cbd4cc6f78a771a83c856e88d48112f555bc13",
+  "model-00003-of-00006.safetensors=sha256:52595519b59349c5c6e373e6f5ca4a3d48ea6dde345f7e61e24766df5fab0e5b",
+  "model-00004-of-00006.safetensors=sha256:e46c6113c89c6f3e9b072c1bef02d763a625c37bcd8f9da2ed9363891c9a0758",
+  "model-00005-of-00006.safetensors=sha256:6d92cb2bf9791de644de2ae86f8523d802ac3b4aaabfff0716ab6c2b97f6fb14",
+  "model-00006-of-00006.safetensors=sha256:5fc1a8632490bb34162823c35d0d591337b9e4195b22cc0560741397a6e9d0b3",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
+  "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
+]
+official_repo = "biohub/ESMC-6B"
+official_revision = "45b0fa5d7fb06faefbd5e3b89bdcef35d564e79a"
+official_files = [
+  "config.json=git-sha1:19f5fb09e4f630fb5b748a497183c22a87ec5102",
+  "model-00001-of-00006.safetensors=sha256:bd90149ff223e6ac1a0cac6147a5ae0df20d3a21df4f65356a1f19cd14f4aa8a",
+  "model-00002-of-00006.safetensors=sha256:f75e2144d8269fe2eb4b3e0823fb089b94f176d8024153e85b8fb573a42294fa",
+  "model-00003-of-00006.safetensors=sha256:f699f01ecc9691d9c6470492765fe54b8b5d2e9f277c139e89427433ffdfe0b2",
+  "model-00004-of-00006.safetensors=sha256:46add1b7be098bbfdc3073884851ba3057f1b33ea23a158b650a37007dabd13d",
+  "model-00005-of-00006.safetensors=sha256:1e1cb62f060a34e18f54a31a76683ef888b8cec59e73315f5b31d25d45a1f88c",
+  "model-00006-of-00006.safetensors=sha256:56c73e13ae96e777ce65eee99364056069ef93b646470f352f83c5f1037b1b18",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c",
+  "tokenizer_config.json=git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61",
+]
+
+[[models]]
+id = "esm3_small"
+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" }
+fast_repo = "Synthyra/ESM3_small"
+fast_revision = "7ddb5a740f9e5f93933eb6410c0ee8684bc63ec1"
+fast_files = [
+  "config.json=git-sha1:60526e2fdd8af9d4fba17f323775458ef5a1a1f9",
+  "model-00001-of-00002.safetensors=sha256:a4c9b736c4c59d51180e966005a164859b47d5cd36e1f8ecdea619fbd34a0e92",
+  "model-00002-of-00002.safetensors=sha256:bea60e4e91b03bb00b6cedd29b07606b8543f0869fb74454af7b26e216d80d2b",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
+  "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
+]
+official_repo = "biohub/esm3-sm-open-v1"
+official_revision = "47f0545b2b6daf26a93439a3cd610f4f7f3d5478"
+official_files = [
+  "config.json=git-sha1:0967ef424bce6791893e9a57bb952f80fd536e93",
+  "data/weights/esm3_function_decoder_v0.pth=sha256:f76d074efcaccfe21365a4fa96f212dadd66798e1e49d809ab7ffbe025d227c9",
+  "data/weights/esm3_sm_open_v1.pth=sha256:5ead5a135c658068db6a4f1b933e72d6110992c4668822e1c0e2dcc53e38acd9",
+  "data/weights/esm3_structure_decoder_v0.pth=sha256:3b726258a44274792b40ce7ea307e10c5da09936368a4ffa2970264d909da65b",
+  "data/weights/esm3_structure_encoder_v0.pth=sha256:467acbaee703ba3ccde6e75241a912a316952e5ff071355f85c1d33c68704f40",
+]
+
+[[models]]
+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" }
+fast_repo = "Synthyra/Profluent-E1-150M"
+fast_revision = "7c5f3bbf697226a2e0900db7a100f9201774a907"
+fast_files = [
+  "config.json=git-sha1:562ef21e722ca708064fc3d54d25b731d4ac8171",
+  "model.safetensors=sha256:d779ed3a4e23799aafc932dc09c9963428d10aa7075999b5f8851b39c76b67f6",
+]
+official_repo = "Profluent-Bio/E1-150m"
+official_revision = "c4dbfe827e4aa6ed7f95eaef50dc1e084f4d77dc"
+official_files = [
+  "config.json=git-sha1:485e649199b46fe6ee7456bebf7aae9b3d4baeab",
+  "model.safetensors=sha256:ba2656339005e6598642836acfdafde480fecc7e145ce0058eb54adf572c3484",
+]
+
+[[models]]
+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" }
+fast_repo = "Synthyra/Profluent-E1-300M"
+fast_revision = "5ef52c0ad2ae2578f40622696b763523810e8e26"
+fast_files = [
+  "config.json=git-sha1:f5c91498b76a3e3282a0d716d87738abb1a1b6c1",
+  "model.safetensors=sha256:9271c4176a8a2e0905a0bb769570ba1c2978fb999a87da92db4cf2b041224864",
+]
+official_repo = "Profluent-Bio/E1-300m"
+official_revision = "5a2871c587eadbcc9237bc686ea45e5b4d28dfb3"
+official_files = [
+  "config.json=git-sha1:918cb09e6e96d4719ed85951f38c693360f9cdb8",
+  "model.safetensors=sha256:31e09a2542f45b04e6ce4adafb3b657f21e2d56d12bf68fd2266b1576a80bc9b",
+]
+
+[[models]]
+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" }
+fast_repo = "Synthyra/Profluent-E1-600M"
+fast_revision = "6c8bf0ec83b0e0178677c528b101efffd0677742"
+fast_files = [
+  "config.json=git-sha1:1d35c0b35b473259875fd29ee80167487a0d6afe",
+  "model.safetensors=sha256:793483b1b3411eab73fe5214b94d1424ca0545992dfac6889cfc0186af472363",
+]
+official_repo = "Profluent-Bio/E1-600m"
+official_revision = "52d959fb87a609d15cf223a485127b29ed5c382a"
+official_files = [
+  "config.json=git-sha1:8a0a439ed4201462bc01189c9f8b43523b257b5c",
+  "model.safetensors=sha256:cfc108d4b98baaa62932331b40be265eae39dc382595bc3cde4a5ab55db1bf7a",
+]
+
+[[models]]
+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" }
+fast_repo = "Synthyra/DPLM-150M"
+fast_revision = "90ba742754151a774f3b7ed580170d0a76b3e69d"
+fast_files = [
+  "config.json=git-sha1:117ac2c1222152ef378abaad1f605e18c4a18ab0",
+  "model.safetensors=sha256:8bac5ac767ceb8deb511b272d32883f811768d56cb25e920cea94ba9b979ca14",
+  "special_tokens_map.json=git-sha1:ef5f0f7d7baf4947564eafcf79972d272cd80a15",
+  "tokenizer_config.json=git-sha1:80100348e3f2b8ab05b59f3352ea7631685083cd",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "airkingbd/dplm_150m"
+official_revision = "49b7125a5d28c6418fcc2f3c4fe799352ac1488b"
+official_files = [
+  "config.json=git-sha1:4910cb02f1840e9ac577026f601829604af58c74",
+  "pytorch_model.bin=sha256:ea4eaa99536b60ed76f945f71a1a5e604f08447ec3def5104a93ca6001a59961",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models]]
+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" }
+fast_repo = "Synthyra/DPLM-650M"
+fast_revision = "05dc16d97c5c028aed924c9ed681cee4ab609760"
+fast_files = [
+  "config.json=git-sha1:3537150eb87b213a676d5840548625e220b60e8b",
+  "model.safetensors=sha256:e27a47b8ec1c078b3fccb36542210e20f0380c88828db2ca9acf3d8a25048bd8",
+  "special_tokens_map.json=git-sha1:ef5f0f7d7baf4947564eafcf79972d272cd80a15",
+  "tokenizer_config.json=git-sha1:80100348e3f2b8ab05b59f3352ea7631685083cd",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "airkingbd/dplm_650m"
+official_revision = "7a7e651baa667d094aba05e9dc1cf52a3332110a"
+official_files = [
+  "config.json=git-sha1:625574d625a4178ca6966e9545fee56026c0b634",
+  "pytorch_model.bin=sha256:db4e54343a89e7600f41c3aacbc593db1b0caee82ec28cab25ff2ae090eba39c",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models]]
+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" }
+fast_repo = "Synthyra/DPLM-3B"
+fast_revision = "7d764dd3d70ecf1ac0e64693de64a0064aacac65"
+fast_files = [
+  "config.json=git-sha1:7f5baf9426be06760c86882948b0f4af2e681e22",
+  "model-00001-of-00003.safetensors=sha256:37b54855d087ef3e7d883464ae9d5ea3127ec15a16c6323d91ad16a6b98305c9",
+  "model-00002-of-00003.safetensors=sha256:042604fefb05ea8c360a48416ce7ba662a4f90b176b4baf646c5c1814c35e6e8",
+  "model-00003-of-00003.safetensors=sha256:b9ae04012665163c3fc9781dd04fcd69738ac20c07e615e98fc4483fd2c4de45",
+  "special_tokens_map.json=git-sha1:ef5f0f7d7baf4947564eafcf79972d272cd80a15",
+  "tokenizer_config.json=git-sha1:80100348e3f2b8ab05b59f3352ea7631685083cd",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "airkingbd/dplm_3b"
+official_revision = "53849d4a7fe944ae0b9cf2bbc0d2cc0054795b51"
+official_files = [
+  "config.json=git-sha1:f6206456e8c2f22ebe1d37fce3b5d50fd8073e68",
+  "pytorch_model-00001-of-00004.bin=sha256:0bcb86a115fe744ed686756db143f78851304e855e2f83cec58681c6080ced5f",
+  "pytorch_model-00002-of-00004.bin=sha256:daf3324f3be949e7dd1c3c84b28da7fec5151b1890cb0904e73427266856a06f",
+  "pytorch_model-00003-of-00004.bin=sha256:dbbeb7924a21059854f994931e23590b054aa000b10370a71c052c4aa36e9246",
+  "pytorch_model-00004-of-00004.bin=sha256:21c01740d091487db43446489d8a893dea1fcc6f2e1c1991ece13945f7ab4e07",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models]]
+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" }
+artifact_source = "official"
+canonical_state_sha256 = "82e1751f59052b8de72b082517557db47947e8d9b4ac2f11278369e6c0cbf001"
+fast_repo = "Synthyra/DPLM2-150M"
+fast_revision = "182745b8dc5661f898481a4fa60a7af9d53385c4"
+fast_files = [
+  "config.json=git-sha1:07905a2e4327d27d073cd0390f140aec2976125a",
+  "model.safetensors=sha256:0a7751b3113027b1d9c966a5bda2d6ab831855de7aaa047b911731665a7c3cc6",
+  "special_tokens_map.json=git-sha1:e6378d20e897b8806734e65fd3ef9cf42a17631b",
+  "tokenizer_config.json=git-sha1:f2090783e3368b7323aa877e2b740e09f0862259",
+  "vocab.txt=git-sha1:9706a4277a5c39dc9b4ec7b283e8eb130ceaa7f2",
+]
+official_repo = "airkingbd/dplm2_150m"
+official_revision = "3451d984d06497f835ed49634bd68c9dfb54d730"
+official_files = [
+  "config.json=git-sha1:20f1e55c64fdc4d1d30f7b1df64b6167fa23dc7c",
+  "pytorch_model.bin=sha256:be7f5cf9e421f59fcc437e63ce1c7391099a314a4e9a4f10b8688785fa581238",
+  "special_tokens_map.json=git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a",
+  "tokenizer_config.json=git-sha1:fc8c21760dcff173955afb106859e5f015d4f757",
+  "vocab.txt=git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37",
+]
+
+[[models]]
+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" }
+artifact_source = "official"
+canonical_state_sha256 = "cba76b6602d2258de9fffff953b608d93cb8ef4a9e89b0bbd27e160c81e78bb4"
+fast_repo = "Synthyra/DPLM2-650M"
+fast_revision = "b9d8527a9473a54954fa2764f590b9ea1b435bb2"
+fast_files = [
+  "config.json=git-sha1:3e079579b214d48a09db57f2c60be6a1acea5baf",
+  "model.safetensors=sha256:92db08c7dbfd6c5e03fbfeaea3f36b09640ee794dcf5ea8d550527869a9f1d63",
+  "special_tokens_map.json=git-sha1:e6378d20e897b8806734e65fd3ef9cf42a17631b",
+  "tokenizer_config.json=git-sha1:f2090783e3368b7323aa877e2b740e09f0862259",
+  "vocab.txt=git-sha1:9706a4277a5c39dc9b4ec7b283e8eb130ceaa7f2",
+]
+official_repo = "airkingbd/dplm2_650m"
+official_revision = "0bc69b644976c6680ab7e26669854d1979e8876e"
+official_files = [
+  "config.json=git-sha1:4cce8d9dc212cdace0e20e89169790bcf199c158",
+  "pytorch_model.bin=sha256:8d6e08cc05e4858064a714013c74cc88c9caa2cc8b12c34605a3c24bcd877cfb",
+  "special_tokens_map.json=git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a",
+  "tokenizer_config.json=git-sha1:fc8c21760dcff173955afb106859e5f015d4f757",
+  "vocab.txt=git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37",
+]
+
+[[models]]
+id = "dplm2_3b"
+family = "dplm2"
+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" }
+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"
+fast_repo = "Synthyra/DPLM2-3B"
+fast_revision = "2a63babe8848abf5233d31bd55891dff8285fc50"
+fast_files = [
+  "config.json=git-sha1:5932b1d501fed28b84614e0d2c1ecc4e89f10d6e",
+  "model-00001-of-00003.safetensors=sha256:2ff393f6e8df1568ce075d50de69ff4e5e9d9886e5ec47e43d6c24df23459be3",
+  "model-00002-of-00003.safetensors=sha256:feb3cea852c2aa849cc30783a984a97f0d076990ade6606cda5e38bf2a5a9621",
+  "model-00003-of-00003.safetensors=sha256:9be363ddb98436af20901981ffbed2f1097377424987f6c1baad27d512b62e71",
+  "special_tokens_map.json=git-sha1:e6378d20e897b8806734e65fd3ef9cf42a17631b",
+  "tokenizer_config.json=git-sha1:f2090783e3368b7323aa877e2b740e09f0862259",
+  "vocab.txt=git-sha1:9706a4277a5c39dc9b4ec7b283e8eb130ceaa7f2",
+]
+official_repo = "airkingbd/dplm2_3b"
+official_revision = "9e77567926f98d1b997ea9131a8eeb035b9bf827"
+official_files = [
+  "config.json=git-sha1:22d51ce44cd6da8d819e0d00566987bb51d74753",
+  "pytorch_model-00001-of-00004.bin=sha256:d8c641eae6bf891581ec64d543169891b093e296f5679ac75c695bcf596b4211",
+  "pytorch_model-00002-of-00004.bin=sha256:6478ad86ec5fef3d1d26580493af2d8666009d3ff884f3f88548080c8bbf94b5",
+  "pytorch_model-00003-of-00004.bin=sha256:dde8f88dac4a6355488c2fb433ee12cd69f1169950566624fba43684d4d99dc6",
+  "pytorch_model-00004-of-00004.bin=sha256:17ec0145152bc10e4dd3b4c2edff337979f6b99ee7c7bfd6cf4e6dbd7262d079",
+  "special_tokens_map.json=git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a",
+  "tokenizer_config.json=git-sha1:fc8c21760dcff173955afb106859e5f015d4f757",
+  "vocab.txt=git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37",
+]
+
+[[models]]
+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" }
+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"
+fast_repo = "Synthyra/ANKH_base"
+fast_revision = "7ec329aae8e3e174bf22a1eb9e0e9fcc12b53092"
+fast_files = [
+  "config.json=git-sha1:7e1cbce6d08f9bb64eee4410899b1c6b4054f418",
+  "model.safetensors=sha256:b0d3473cac1bda90e39cde54f2abe86da1fc84f872c833ca3415672776dccb95",
+  "special_tokens_map.json=git-sha1:a2d8d626c31389a935e197fb94072e2414a6e7d1",
+  "tokenizer.json=git-sha1:0734d752d12d0f46ac96467fbceb1c4bfbeee0be",
+  "tokenizer_config.json=git-sha1:db0b80de72d3b16242b9eda74ed4663e39c65bcf",
+]
+official_repo = "ElnaggarLab/ankh-base"
+official_revision = "d99cb6b966530dfc2ae96bc69d9255c2a07308b0"
+official_files = [
+  "config.json=git-sha1:abd44a36b5469e9a7cb019e4059b5ac1392d8422",
+  "pytorch_model.bin=sha256:9b2a886374f0ff4a893f4e7a989deed76bb2458c8998bd5202ea8e97d92ddcc3",
+  "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791",
+  "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a",
+  "tokenizer_config.json=git-sha1:a8a872ae3441e7cc85ce19210dff1e4c5d2d7bd0",
+]
+
+[[models]]
+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" }
+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"
+fast_repo = "Synthyra/ANKH_large"
+fast_revision = "3be3df34140f49dc4e65bd1f247e3ce819e7fc59"
+fast_files = [
+  "config.json=git-sha1:272509deedb527e5c2c95b0c269194a44148fdcc",
+  "model.safetensors=sha256:e70b8f9755ac6bfe95d18359060ae9fe38fac63b12a89a886c83349d1adbaa53",
+  "special_tokens_map.json=git-sha1:a2d8d626c31389a935e197fb94072e2414a6e7d1",
+  "tokenizer.json=git-sha1:0734d752d12d0f46ac96467fbceb1c4bfbeee0be",
+  "tokenizer_config.json=git-sha1:2bcaff2567826f5f51188b00600d2c6e7bcea56e",
+]
+official_repo = "ElnaggarLab/ankh-large"
+official_revision = "74b371dbfa3ee0a05d32ae74df0c2e0b82d6b9a6"
+official_files = [
+  "config.json=git-sha1:1abf33e52ee3d6be67d780ec57d32ac2b27b5306",
+  "pytorch_model.bin=sha256:517b6e8b279dedcb477af240b35c46bd6eb3307723eb281e60d4b2c8a87b889b",
+  "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791",
+  "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a",
+  "tokenizer_config.json=git-sha1:d7fe02ba6f2b18d9ccfa19ac129c9fdc9ec24d09",
+]
+
+[[models]]
+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" }
+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"
+fast_repo = "Synthyra/ANKH2_large"
+fast_revision = "392de5ed52bbfd73b45f545e378aaebcff096d0e"
+fast_files = [
+  "config.json=git-sha1:66b6adc7215743a98a3229958bbd1c9c42b6108b",
+  "model.safetensors=sha256:be8e6242388d93b51cd9719a0e32cfc17a2e804786570c795ba332197eccb915",
+  "special_tokens_map.json=git-sha1:a2d8d626c31389a935e197fb94072e2414a6e7d1",
+  "tokenizer.json=git-sha1:0734d752d12d0f46ac96467fbceb1c4bfbeee0be",
+  "tokenizer_config.json=git-sha1:db0b80de72d3b16242b9eda74ed4663e39c65bcf",
+]
+official_repo = "ElnaggarLab/ankh2-ext2"
+official_revision = "aa9b9fa72288c47d9f618ce80c011e24b54e17a8"
+official_files = [
+  "config.json=git-sha1:9286bed4ecbc4f7113024919d16ec9719b0c0748",
+  "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a",
+  "pytorch_model.bin=sha256:2df583f28f111276ee22a7b76007f4297e9a69766d60bccd9c8d7169c06ac606",
+  "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791",
+  "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a",
+  "tokenizer_config.json=git-sha1:854e5db75dae8b1e9dd39c5bae80dae5508b3e25",
+]
+
+[[models]]
+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" }
+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"
+fast_repo = "Synthyra/ANKH3_large"
+fast_revision = "53600f175f328f986f43e55ca8ceb14935d337a4"
+fast_files = [
+  "config.json=git-sha1:432b09625d44a2eeab679fddb7495d42b560b7f9",
+  "model.safetensors=sha256:9f50f58cf5b3a537a0a41aa918695c3a26d7985dd0b2266642d6f86324c9e7a1",
+  "special_tokens_map.json=git-sha1:1fc3a4d6d4282e5201cd7c30d5c0a6a8bfa04f82",
+  "tokenizer.json=git-sha1:3d14291df2d6db3a183c5c4fe133afb330cc44cf",
+  "tokenizer_config.json=git-sha1:2005fec00a7ae9a49e248a1ecefbbd81c56674d6",
+]
+official_repo = "ElnaggarLab/ankh3-large"
+official_revision = "2be091622e8a393f0ef21735070084123c874b6e"
+official_files = [
+  "config.json=git-sha1:f5278f77d158cdd8a173df888e3ed365e84a80a3",
+  "generation_config.json=git-sha1:5767cc0cacebfd06884eb27ae1c796d3ca829fd2",
+  "pytorch_model.bin=sha256:26321a345e07a25b21c6c41b651c4db91b420892e52c0dcbc55bd7a8f510f95b",
+  "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4",
+  "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0",
+  "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca",
+  "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21",
+]
+
+[[models]]
+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" }
+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"
+fast_repo = "Synthyra/ANKH3_xl"
+fast_revision = "3cbf2c22c4f7d67bf0bfcbdcd500f41723e91d29"
+fast_files = [
+  "config.json=git-sha1:23f6d78ddcb3a031b88f876eaaf04c2fafaea46f",
+  "model-00001-of-00003.safetensors=sha256:39bd8f75cf98a67cf04055399f9fc401198f6fc2896b112aba9fd9ec9df52ab9",
+  "model-00002-of-00003.safetensors=sha256:9ff73233b39d2c200abb78e66b320c014ec61431bd6e1af36fb188a3cfa24c34",
+  "model-00003-of-00003.safetensors=sha256:c13125c02dbcd7f07bd412e9e085f2bca6624d2f1f45fedc95fb777f53161cbe",
+  "special_tokens_map.json=git-sha1:1fc3a4d6d4282e5201cd7c30d5c0a6a8bfa04f82",
+  "tokenizer.json=git-sha1:3d14291df2d6db3a183c5c4fe133afb330cc44cf",
+  "tokenizer_config.json=git-sha1:2005fec00a7ae9a49e248a1ecefbbd81c56674d6",
+]
+official_repo = "ElnaggarLab/ankh3-xl"
+official_revision = "e00113df5c95ef71df7ea3f5a73d56bd00e473a4"
+official_files = [
+  "config.json=git-sha1:f8997040e8913df75fd2eebe71a2a8eb750ed0d0",
+  "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a",
+  "pytorch_model-00001-of-00003.bin=sha256:2c9793cbee16697cd4149debe07d3a27143e280f6e970fa46042aae820fea981",
+  "pytorch_model-00002-of-00003.bin=sha256:31c5a860e414513c829ae52affb0970d7cef2c0545df2d6e1338b6806ab7174b",
+  "pytorch_model-00003-of-00003.bin=sha256:055a853bdd3623db95a637935aa299427e837cd8ea69fc04708b0262508bec75",
+  "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4",
+  "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0",
+  "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca",
+  "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21",
+]
+
+[[models]]
+id = "boltz2"
+family = "boltz2"
+size_category = "structure"
+generation_contract = "not_applicable"
+notes = "Boltz2 is provisional in FastPLMs 1.0. Exact configuration, the declared inference-core state, feature preparation, and seeded execution remain tested, but native-environment BF16 end-to-end inference currently exceeds the fixed numerical-equivalence limits. FastPLMs therefore does not claim official inference equivalence for this checkpoint yet. Work on that numerical gap continues independently of the ESM++ and ESMFold2 release gates."
+fast_repo = "Synthyra/Boltz2"
+fast_revision = "3b148fc5efea109c065ec82ba8683d024de7134e"
+fast_files = [
+  "config.json=git-sha1:8682ccb12e177e73bc7a351ff7e3af484bfb6fac",
+  "model.safetensors=sha256:5c863fd200a1613a0e311071e2ad73ab350635e3fd336e6822cf45c52cb960e5",
+]
+official_repo = "boltz-community/boltz-2"
+official_revision = "6fdef46d763fee7fbb83ca5501ccceff43b85607"
+official_files = [
+  "boltz2_conf.ckpt=sha256:090e82ac8c92f5e943fa1b39e7410a44027bea7243c0bbb3caa67a77fc1428e1",
+  "mols.tar=sha256:39e076d96dbec6b4e86982bbda16f3a53a2a60c9bdc17828d88f6f9a0c7d1fd7",
+]
+
+[[models]]
+id = "esmfold"
+family = "esmfold"
+size_category = "structure"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esmfold.json=sha256:380b9a96168410717d1f698feaabb826b1606444cbdeec86c2ea06d9ffe8f186", tensors = "tests/goldens/esmfold.safetensors=sha256:873b1b325a43d8e0f35f355c8914a2a9fe611cc48763875e9e6a22e09ec9ebcb" }
+fast_repo = "Synthyra/FastESMFold"
+fast_revision = "b88c8cb50d19b2cf7ab4fee4b0a61f5e02da7823"
+fast_files = [
+  "config.json=git-sha1:18e0091dcbf6140bf68924d53c4c8917b9cd90b1",
+  "model-00001-of-00003.safetensors=sha256:36fab9e5c96d409b2a34a8b4f1273acac8c07f119c32c4fcfa7d47bbcd55b83c",
+  "model-00002-of-00003.safetensors=sha256:34954aaa05bc91635776ba6672946da5822626753d80db97b38c0538e9525102",
+  "model-00003-of-00003.safetensors=sha256:2f1178cda0e6cff3b1e158e1acc59c83e3f4fc46e246388a5127bc56b8d9c4f2",
+  "special_tokens_map.json=git-sha1:53cd95604a28eb7e23da763c8da23f5006ab2179",
+  "tokenizer_config.json=git-sha1:10213f69b51b4b38876a29271b8f908e853a5800",
+  "vocab.txt=git-sha1:eee0a1fc93c82568f78f086550fbd7c591cf423a",
+]
+official_repo = "facebook/esmfold_v1"
+official_revision = "75a3841ee059df2bf4d56688166c8fb459ddd97a"
+official_files = [
+  "config.json=git-sha1:1232d0aee4be551021d8e70e66ed2b062df917bf",
+  "pytorch_model.bin=sha256:2ee07356b125d1e3e57503c204111fd7323347fc4735d41d3caac57c2a78e116",
+  "special_tokens_map.json=git-sha1:121c8d54f8ea66cdf678f48b3cb37c05b4de5c0d",
+  "tokenizer_config.json=git-sha1:aad24fba9f1bad2d74ed79d414ddcd60e6b0f812",
+  "vocab.txt=git-sha1:9abfdf5472c0ed970648b683b86ab131256b3e42",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esmfold_3B_v1.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt"
+sha256 = "e9a52579027e77d2d2e0a18218e755821f395730e86624cab9413dc117f5ca62"
+size = 2771653574
+
+[[models]]
+id = "esmfold2"
+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" }
+fast_repo = "Synthyra/ESMFold2"
+fast_revision = "cd5a0927cec585a778d983b99a8db23d2e9b281e"
+fast_files = [
+  "config.json=git-sha1:67e81ff571f393f0b630cd5a22398bd84979c030",
+  "model.safetensors=sha256:138fd4350d6892b81ce6be7ff9bf5a93ae9d4d3751f46a27438a3f9f0dcefa0e",
+]
+official_repo = "biohub/ESMFold2"
+official_revision = "1ebf0e3481a5184eb6171d40615c79e384b48796"
+official_files = [
+  "config.json=git-sha1:0300c084b990b2bd600efd9f538aa5de27109fea",
+  "model.safetensors=sha256:138fd4350d6892b81ce6be7ff9bf5a93ae9d4d3751f46a27438a3f9f0dcefa0e",
+]
+
+[[models]]
+id = "esmfold2_fast"
+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" }
+fast_repo = "Synthyra/ESMFold2-Fast"
+fast_revision = "407875bfcaa42552bfcb25acd67ee1888b790170"
+fast_files = [
+  "config.json=git-sha1:62ccca15a416a5dcbd02cd6ce161f432c7b4de58",
+  "model.safetensors=sha256:60ca19f2898188beba92944365f7b909efd9c99212f5018af75cc47cd9a6184a",
+]
+official_repo = "biohub/ESMFold2-Fast"
+official_revision = "b28d8ace5e05e61e5bec1e6820cfd3e221819d12"
+official_files = [
+  "config.json=git-sha1:c0ca526090fa7f8342ee4666d56e7fe3a4b8cbb2",
+  "model.safetensors=sha256:60ca19f2898188beba92944365f7b909efd9c99212f5018af75cc47cd9a6184a",
+]
+
+[[models]]
+id = "esmfold2_experimental_cutoff2025"
+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" }
+fast_repo = "Synthyra/ESMFold2-Experimental-Cutoff2025"
+fast_revision = "632ff4a9e68f1de78ee956a613267bdcdb5b354d"
+fast_files = [
+  "config.json=git-sha1:41119745d38bc5503a0212ad923e75211dec565f",
+  "model.safetensors=sha256:01358c317428d38535e3db513cab177336fc0f7fab0d84002e64b7741d5181b3",
+]
+official_repo = "biohub/ESMFold2-Experimental-Cutoff2025"
+official_revision = "56f94f5c1069ecde17512c96928850518340d287"
+official_files = [
+  "config.json=git-sha1:79ed0dc0f867b8f09bfa004d6f77397c2ab9b38d",
+  "model.safetensors=sha256:01358c317428d38535e3db513cab177336fc0f7fab0d84002e64b7741d5181b3",
+]
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel" }
+
+[[models]]
+id = "esmfold2_experimental_fast_cutoff2025"
+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" }
+fast_repo = "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025"
+fast_revision = "8f022c2514a6c32692aaca078a8391d6bc6c4bac"
+fast_files = [
+  "config.json=git-sha1:b9d39e941050179ca51faaed58cbbd77778c1143",
+  "model.safetensors=sha256:4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f",
+]
+official_repo = "biohub/ESMFold2-Experimental-Fast-Cutoff2025"
+official_revision = "74b88548bf19688b8727432db0d698cb2e1d8783"
+official_files = [
+  "config.json=git-sha1:0333d68ddb12ed2f066741dcb801142f466c0a2c",
+  "model.safetensors=sha256:4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f",
+]
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel" }

fastplms/models/__init__.py

@@ -0,0 +1,10 @@
+"""Lazy model-family namespace for FastPLMs.
+
+Model classes are resolved through Transformers AutoClasses and the typed
+registry. Importing this package therefore does not load checkpoints, create
+tokenizers, compile kernels, or initialize an accelerator runtime.
+"""
+
+from __future__ import annotations
+
+__all__: tuple[str, ...] = ()

fastplms/models/boltz/__init__.py

@@ -0,0 +1,13 @@
+from fastplms.models.boltz.modeling_boltz2 import (
+    Boltz2Config,
+    Boltz2Model,
+    Boltz2ModelOutput,
+    Boltz2StructureOutput,
+)
+
+__all__ = [
+    "Boltz2Config",
+    "Boltz2Model",
+    "Boltz2ModelOutput",
+    "Boltz2StructureOutput",
+]

fastplms/models/boltz/_pair_attention.py

@@ -0,0 +1,46 @@
+"""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)``."""
+
+    batch_size, sequence_length, width = states.shape
+    if width % num_heads:
+        raise ValueError(f"width {width} is not divisible by {num_heads} heads")
+    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)``."""
+
+    head_dim = query_states.shape[-1]
+    with torch.autocast("cuda", enabled=False):
+        # S is the FP32 attention-score tensor with shape (b, h, l_q, l_k).
+        scores = torch.einsum(
+            "bihd,bjhd->bhij",
+            query_states.float(),
+            key_states.float(),
+        )
+        scores = scores / math.sqrt(head_dim) + pair_bias.float()
+        scores = scores + (1 - key_mask[:, None, None].float()) * -mask_value
+        probabilities = scores.softmax(dim=-1)
+        output = torch.einsum(
+            "bhij,bjhd->bihd",
+            probabilities,
+            value_states.float(),
+        )
+    return output.to(value_states.dtype)

fastplms/models/boltz/cif_writer.py

@@ -0,0 +1,141 @@
+from pathlib import Path
+
+import numpy as np
+import torch
+
+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:
+    if plddt is None:
+        return np.ones((num_atoms,), dtype=np.float32) * 100.0
+
+    values = plddt.detach().cpu()
+    if values.ndim == 1:
+        values = values.unsqueeze(0)
+    if values.ndim != 2:
+        raise ValueError("Expected pLDDT tensor S with shape (n_samples, n_items).")
+    if not 0 <= sample_index < values.shape[0]:
+        raise IndexError("sample_index out of range for pLDDT.")
+
+    selected = values[sample_index]
+    if selected.shape[0] == num_atoms:
+        return (selected.numpy() * 100.0).astype(np.float32)
+
+    num_residues = max(atom_to_residue) + 1
+    if selected.shape[0] == num_residues:
+        expanded = np.zeros((num_atoms,), dtype=np.float32)
+        selected_np = selected.numpy()
+        for atom_idx, residue_idx in enumerate(atom_to_residue):
+            expanded[atom_idx] = selected_np[residue_idx] * 100.0
+        return expanded
+
+    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:
+    coords = atom_coords.detach().cpu()
+    if coords.ndim == 2:
+        coords = coords.unsqueeze(0)
+    if coords.ndim != 3 or coords.shape[-1] != 3:
+        raise ValueError(
+            "Expected coordinate tensor X with shape (n_samples, n_atoms, 3)."
+        )
+    if not 0 <= sample_index < coords.shape[0]:
+        raise IndexError("sample_index out of range.")
+    selected_coords_tensor = coords[sample_index]
+    all_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor))
+    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()
+
+    mask = atom_mask.detach().cpu()
+    if mask.ndim == 2:
+        mask = mask[0]
+    if mask.ndim != 1:
+        raise ValueError("Expected atom mask M with shape (n_atoms,).")
+    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]))
+    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,
+    )
+    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()
+        b_factor = float(b_iso[idx])
+
+        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

@@ -0,0 +1,786 @@
+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)
+    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)
+    # 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,
+    )
+
+
+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] = []
+    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)
+            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)
+    squared_norm = (quaternion * quaternion).sum(dim=1)
+    norm = torch.sqrt(squared_norm)
+    signed_norm = torch.where(quaternion[:, 0] < 0, -norm, norm)
+    quaternion = quaternion / signed_norm[:, None]
+
+    real, i, j, k = torch.unbind(quaternion, dim=-1)
+    two_s = 2.0 / (quaternion * quaternion).sum(dim=-1)
+    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,
+    )
+    return rotation.reshape(1, 3, 3)[0]
+
+
+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.
+    """
+    result = atom_positions.clone()
+    residue_index_tensor = torch.tensor(atom_residue_index, dtype=torch.long)
+    for residue_idx in range(max(num_residues - 1, 0)):
+        residue_mask = residue_index_tensor == residue_idx
+        if not torch.any(residue_mask):
+            raise RuntimeError(f"Residue index {residue_idx} has no atoms.")
+        residue_coords = result[residue_mask][None]
+        resolved_mask = torch.ones(
+            residue_coords.shape[:2], dtype=torch.bool, device=residue_coords.device
+        )
+        residue_center = torch.sum(
+            residue_coords * resolved_mask[:, :, None], dim=1, keepdim=True
+        ) / torch.sum(resolved_mask[:, :, None], dim=1, keepdim=True)
+        residue_coords = residue_coords - residue_center
+        rotation = _random_rotation_matrix()[None]
+        residue_coords = torch.einsum("bmd,bds->bms", residue_coords, rotation)
+        residue_coords = residue_coords + torch.randn_like(residue_coords[:, 0:1, :])
+        result[residue_mask] = residue_coords[0]
+    return result
+
+
+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.")
+
+    atom_positions = torch.tensor(np.asarray(atom_positions_np), dtype=torch.float32)
+    atom_positions = _center_and_augment_atoms_per_residue(
+        atom_positions=atom_positions,
+        atom_residue_index=template.atom_residue_index,
+        num_residues=num_tokens,
+    )
+
+    token_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
+    residue_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
+    asym_id = torch.zeros((1, num_tokens), dtype=torch.long)
+    entity_id = torch.zeros((1, num_tokens), dtype=torch.long)
+    sym_id = torch.zeros((1, num_tokens), dtype=torch.long)
+    mol_type = torch.full(
+        (1, num_tokens),
+        fill_value=const.chain_type_ids["PROTEIN"],
+        dtype=torch.long,
+    )
+
+    res_type_ids = torch.tensor(residue_token_ids, dtype=torch.long)
+    res_type = one_hot(res_type_ids, num_classes=const.num_tokens).unsqueeze(0)
+
+    # token_bonds encodes explicit covalent cross-links from structure bonds,
+    # NOT backbone peptide bonds (those are implicit via residue_index + asym_id).
+    # For a standard single-chain protein without cross-links, this is all zeros.
+    # This matches the official Boltz2 featurizer (featurizerv2.py lines 696-705).
+    token_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.float32)
+    type_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.long)
+    token_bonds = token_bonds.unsqueeze(0).unsqueeze(-1)
+    type_bonds = type_bonds.unsqueeze(0)
+
+    token_pad_mask = torch.ones((1, num_tokens), dtype=torch.float32)
+    token_resolved_mask = torch.ones((1, num_tokens), dtype=torch.float32)
+    token_disto_mask = torch.ones((1, num_tokens), dtype=torch.float32)
+
+    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,
+    )
+    contact_conditioning = one_hot(
+        contact_ids,
+        num_classes=num_contact_classes,
+    ).unsqueeze(0)
+    contact_threshold = torch.zeros((1, num_tokens, num_tokens), dtype=torch.float32)
+
+    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,
+    )
+    modified = torch.zeros((1, num_tokens), dtype=torch.long)
+    cyclic_period = torch.zeros((1, num_tokens), dtype=torch.float32)
+    affinity_token_mask = torch.zeros((1, num_tokens), dtype=torch.float32)
+
+    ref_pos = atom_positions.unsqueeze(0)
+    atom_pad_mask = torch.ones((1, num_atoms), dtype=torch.float32)
+    atom_resolved_mask = torch.ones((1, num_atoms), dtype=torch.bool)
+
+    atom_name_codes = torch.stack(
+        [_atom_name_to_codes(atom_name) for atom_name in template.atom_names],
+        dim=0,
+    )
+    ref_atom_name_chars = one_hot(atom_name_codes, num_classes=64).unsqueeze(0)
+
+    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)
+
+    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)
+    ref_chirality = torch.tensor(ref_chirality_values, dtype=torch.long).unsqueeze(0)
+    ref_space_uid = torch.tensor(template.atom_residue_index, dtype=torch.long).unsqueeze(0)
+
+    atom_to_token = one_hot(
+        torch.tensor(template.atom_residue_index, dtype=torch.long),
+        num_classes=num_tokens,
+    ).unsqueeze(0)
+    token_to_rep_atom = one_hot(
+        torch.tensor(residue_disto_atom_idx, dtype=torch.long),
+        num_classes=num_atoms,
+    ).unsqueeze(0)
+    token_to_center_atom = one_hot(
+        torch.tensor(residue_center_atom_idx, dtype=torch.long),
+        num_classes=num_atoms,
+    ).unsqueeze(0)
+    r_set_to_rep_atom = token_to_center_atom.clone()
+
+    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)
+
+    # X contains no observed coordinates for sequence-only inference.
+    coords = torch.zeros((1, 1, num_atoms, 3), dtype=torch.float32)
+    disto_coords_ensemble = torch.zeros(
+        (1, 1, num_tokens, 3),
+        dtype=torch.float32,
+    )
+
+    bfactor = torch.zeros((1, num_atoms), dtype=torch.float32)
+    atom_plddt = torch.ones((1, num_atoms), dtype=torch.float32)
+
+    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)
+        atom_pad_mask = torch.nn.functional.pad(atom_pad_mask, (0, pad_atoms), value=0.0)
+        atom_resolved_mask = torch.nn.functional.pad(
+            atom_resolved_mask,
+            (0, pad_atoms),
+            value=0.0,
+        )
+        ref_atom_name_chars = torch.nn.functional.pad(
+            ref_atom_name_chars,
+            (0, 0, 0, 0, 0, pad_atoms),
+            value=0.0,
+        )
+        ref_element = torch.nn.functional.pad(ref_element, (0, 0, 0, pad_atoms), value=0.0)
+        ref_charge = torch.nn.functional.pad(ref_charge, (0, pad_atoms), value=0.0)
+        ref_chirality = torch.nn.functional.pad(ref_chirality, (0, pad_atoms), value=0)
+        atom_backbone_feat = torch.nn.functional.pad(
+            atom_backbone_feat,
+            (0, 0, 0, pad_atoms),
+            value=0.0,
+        )
+        ref_space_uid = torch.nn.functional.pad(ref_space_uid, (0, pad_atoms), value=0)
+        coords = torch.nn.functional.pad(coords, (0, 0, 0, pad_atoms), value=0.0)
+        atom_to_token = torch.nn.functional.pad(atom_to_token, (0, 0, 0, pad_atoms), value=0.0)
+        token_to_rep_atom = torch.nn.functional.pad(
+            token_to_rep_atom,
+            (0, pad_atoms),
+            value=0.0,
+        )
+        token_to_center_atom = torch.nn.functional.pad(
+            token_to_center_atom,
+            (0, pad_atoms),
+            value=0.0,
+        )
+        r_set_to_rep_atom = torch.nn.functional.pad(
+            r_set_to_rep_atom,
+            (0, pad_atoms),
+            value=0.0,
+        )
+        bfactor = torch.nn.functional.pad(bfactor, (0, pad_atoms), value=0.0)
+        atom_plddt = torch.nn.functional.pad(atom_plddt, (0, pad_atoms), value=0.0)
+
+    frames_idx = torch.tensor(
+        residue_frame_atom_idx_flat,
+        dtype=torch.long,
+    ).reshape(num_tokens, 3)
+    frames_idx = frames_idx.unsqueeze(0).unsqueeze(1)
+    frame_resolved_mask = torch.zeros((1, 1, num_tokens), dtype=torch.bool)
+
+    msa = torch.tensor(residue_token_ids, dtype=torch.long).unsqueeze(0).unsqueeze(0)
+    msa_paired = torch.ones((1, 1, num_tokens), dtype=torch.float32)
+    deletion_value = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    has_deletion = torch.zeros((1, 1, num_tokens), dtype=torch.bool)
+    msa_mask = torch.ones((1, 1, num_tokens), dtype=torch.long)
+    deletion_mean = torch.zeros((1, num_tokens), dtype=torch.float32)
+    profile = (
+        one_hot(
+            torch.tensor(residue_token_ids, dtype=torch.long),
+            num_classes=const.num_tokens,
+        )
+        .float()
+        .unsqueeze(0)
+    )
+
+    template_restype = one_hot(
+        torch.zeros((1, 1, num_tokens), dtype=torch.long),
+        num_classes=const.num_tokens,
+    )
+    template_frame_rot = torch.zeros((1, 1, num_tokens, 3, 3), dtype=torch.float32)
+    template_frame_t = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
+    template_cb = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
+    template_ca = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
+    template_mask_cb = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    template_mask_frame = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    template_mask = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    query_to_template = torch.zeros((1, 1, num_tokens), dtype=torch.long)
+    visibility_ids = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+
+    disto_target = torch.zeros(
+        (1, num_tokens, num_tokens, 1, num_bins),
+        dtype=torch.float32,
+    )
+    disto_target[..., 0] = 1.0
+    disto_center = torch.zeros((1, num_tokens, 3), dtype=torch.float32)
+
+    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

fastplms/models/boltz/minimal_structures.py

@@ -0,0 +1,19 @@
+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

@@ -0,0 +1,993 @@
+import copy
+import inspect
+import random
+from collections.abc import Iterator, Mapping, Sequence
+from contextlib import contextmanager, nullcontext
+from dataclasses import dataclass
+from typing import Any, ClassVar
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch import Tensor
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+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]:
+    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()
+    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
+    torch_cuda: list[Tensor] | None
+
+
+@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
+    hidden_states: tuple[torch.Tensor, ...] | None = None
+    attentions: tuple[torch.Tensor, ...] | None = None
+    pdistogram: torch.Tensor | None = None
+    s: torch.Tensor | None = None
+    z: torch.Tensor | None = None
+    sample_atom_coords: torch.Tensor | None = None
+    diff_token_repr: torch.Tensor | None = None
+    s_conf: torch.Tensor | None = None
+    z_conf: torch.Tensor | None = None
+    pde_logits: torch.Tensor | None = None
+    plddt_logits: torch.Tensor | None = None
+    resolved_logits: torch.Tensor | None = None
+    pde: torch.Tensor | None = None
+    plddt: torch.Tensor | None = None
+    complex_plddt: torch.Tensor | None = None
+    complex_iplddt: torch.Tensor | None = None
+    complex_pde: torch.Tensor | None = None
+    complex_ipde: torch.Tensor | None = None
+    pae_logits: torch.Tensor | None = None
+    pae: torch.Tensor | None = None
+    ptm: torch.Tensor | None = None
+    iptm: torch.Tensor | None = None
+    ligand_iptm: torch.Tensor | None = None
+    protein_iptm: torch.Tensor | None = None
+    pair_chains_iptm: Any = None
+
+
+@dataclass
+class Boltz2StructureOutput(ModelOutput):
+    sample_atom_coords: torch.Tensor | None = None
+    atom_pad_mask: torch.Tensor | None = None
+    plddt: torch.Tensor | None = None
+    confidence_score: torch.Tensor | None = None
+    complex_plddt: torch.Tensor | None = None
+    iptm: torch.Tensor | None = None
+    ptm: torch.Tensor | None = None
+    sequence: str | None = None
+    structure_template: ProteinStructureTemplate | None = None
+    raw_output: dict[str, torch.Tensor] | None = None
+    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,
+    ) -> dict[str, Tensor]:
+        s_inputs = self.input_embedder(feats)
+        s_init = self.s_init(s_inputs)
+
+        z_init = self.z_init_1(s_inputs)[:, :, None] + self.z_init_2(s_inputs)[:, None, :]
+        relative_position_encoding = self.rel_pos(feats)
+        z_init = z_init + relative_position_encoding
+        z_init = z_init + self.token_bonds(feats["token_bonds"].float())
+        if self.bond_type_feature:
+            z_init = z_init + self.token_bonds_type(feats["type_bonds"].long())
+        z_init = z_init + self.contact_conditioning(feats)
+
+        s = torch.zeros_like(s_init)
+        z = torch.zeros_like(z_init)
+        mask = feats["token_pad_mask"].float()
+        pair_mask = mask[:, :, None] * mask[:, None, :]
+
+        if self.run_trunk_and_structure:
+            for _ in range(recycling_steps + 1):
+                s = s_init + self.s_recycle(self.s_norm(s))
+                z = z_init + self.z_recycle(self.z_norm(z))
+                z = z + self.msa_module(
+                    z,
+                    s_inputs,
+                    feats,
+                    use_kernels=self.use_kernels,
+                )
+                s, z = self.pairformer_module(
+                    s,
+                    z,
+                    mask=mask,
+                    pair_mask=pair_mask,
+                    use_kernels=self.use_kernels,
+                )
+
+        pdistogram = self.distogram_module(z)
+        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,
+                )
+            )
+            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,
+                )
+            output.update(struct_out)
+
+        if self.confidence_prediction:
+            if self.skip_run_structure:
+                x_pred = feats["coords"].repeat_interleave(diffusion_samples, 0)
+            else:
+                if "sample_atom_coords" not in output:
+                    raise RuntimeError("Structure sampling did not produce sample_atom_coords.")
+                x_pred = output["sample_atom_coords"]
+
+            if detach_confidence:
+                s_inputs_c = s_inputs.detach()
+                s_c = s.detach()
+                z_c = z.detach()
+                x_pred_c = x_pred.detach()
+                pdist_c = output["pdistogram"][:, :, :, 0].detach()
+            else:
+                s_inputs_c = s_inputs
+                s_c = s
+                z_c = z
+                x_pred_c = x_pred
+                pdist_c = output["pdistogram"][:, :, :, 0]
+
+            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,
+                )
+            )
+
+        return output
+
+
+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]:
+        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
+
+    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,
+        )
+        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"]
+        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)
+        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,
+    ) -> Boltz2ModelOutput | tuple[Any, ...]:
+        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,
+        )
+        token_state = raw_output.get("s")
+        pair_state = raw_output.get("z")
+        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]:
+        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)
+                else:
+                    moved[key] = value.to(device=self.device)
+            else:
+                moved[key] = value
+        return moved
+
+    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,
+    ) -> 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):
+            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
+                ),
+            )
+            feats = self._to_model_device(feats, float_dtype=torch.float32)
+            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,
+                )
+
+        sample_atom_coords_value = output.get("sample_atom_coords")
+        if not isinstance(sample_atom_coords_value, torch.Tensor):
+            raise RuntimeError("Boltz2 structure sampling did not return coordinate tensors.")
+        sample_atom_coords = sample_atom_coords_value.detach().cpu()
+        non_finite_mask = torch.logical_not(torch.isfinite(sample_atom_coords))
+        if torch.any(non_finite_mask):
+            raise RuntimeError(
+                "sample_atom_coords contains non-finite values. "
+                f"Non-finite count: {int(non_finite_mask.sum().item())}"
+            )
+        atom_pad_mask = feats["atom_pad_mask"][0].detach().cpu()
+        plddt = output["plddt"].detach().cpu() if "plddt" in output else None
+        complex_plddt = (
+            output["complex_plddt"].detach().cpu() if "complex_plddt" in output else None
+        )
+        iptm = output["iptm"].detach().cpu() if "iptm" in output else None
+        ptm = output["ptm"].detach().cpu() if "ptm" in output else None
+
+        confidence_score = None
+        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
+            else:
+                confidence_score = (4 * complex_plddt + iptm) / 5
+
+        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

@@ -0,0 +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",
+    ]
+]

fastplms/models/boltz/vb_layers_attention.py

@@ -0,0 +1,98 @@
+"""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:
+        if model_cache is not None and "z" in model_cache:
+            return model_cache["z"]
+        pair_bias = self.proj_z(pair_states)
+        if model_cache is not None:
+            model_cache["z"] = pair_bias
+        return pair_bias
+
+    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)``."""
+
+        sequence_states = self.norm_s(s) if self.initial_norm else s
+        if to_keys is not None:
+            key_states = to_keys(sequence_states)
+            key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)
+        else:
+            key_states = sequence_states if k_in is None else k_in
+            key_mask = mask
+
+        query = reshape_heads(self.proj_q(sequence_states), self.num_heads)
+        key = reshape_heads(self.proj_k(key_states), self.num_heads)
+        value = reshape_heads(self.proj_v(key_states), self.num_heads)
+        pair_bias = self._resolve_pair_bias(z, model_cache).repeat_interleave(
+            multiplicity,
+            dim=0,
+        )
+        attended = pair_biased_attention(
+            query,
+            key,
+            value,
+            pair_bias,
+            key_mask,
+            self.inf,
+        )
+        attended = attended.reshape(s.shape[0], -1, self.c_s)
+        gate = self.proj_g(sequence_states).sigmoid()
+        return self.proj_o(gate * attended)

fastplms/models/boltz/vb_layers_attentionv2.py

@@ -0,0 +1,73 @@
+"""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)``."""
+
+        query = reshape_heads(self.proj_q(s), self.num_heads)
+        key = reshape_heads(self.proj_k(k_in), self.num_heads)
+        value = reshape_heads(self.proj_v(k_in), self.num_heads)
+        pair_bias = self.proj_z(z).repeat_interleave(multiplicity, dim=0)
+        attended = pair_biased_attention(
+            query,
+            key,
+            value,
+            pair_bias,
+            mask,
+            self.inf,
+        )
+        attended = attended.reshape(s.shape[0], -1, self.c_s)
+        gate = self.proj_g(s).sigmoid()
+        return self.proj_o(gate * attended)

fastplms/models/boltz/vb_layers_confidence_utils.py

@@ -0,0 +1,226 @@
+"""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."""
+
+    norm1 = torch.norm(v1, dim=1, keepdim=True)
+    norm2 = torch.norm(v2, dim=1, keepdim=True)
+    unit1 = v1 / (norm1 + 1e-6)
+    unit2 = v2 / (norm2 + 1e-6)
+    separated = torch.abs(torch.sum(unit1 * unit2, dim=1)) < 0.9063
+    return separated & (norm1.reshape(-1) > 1e-2) & (norm2.reshape(-1) > 1e-2)
+
+
+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)
+
+
+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:
+    token_chain_ids = feats["asym_id"]
+    for batch_index, sample_coordinates in enumerate(coordinates):
+        token_offset = 0
+        atom_offset = 0
+        for chain_id in torch.unique(token_chain_ids[batch_index]):
+            token_mask = (token_chain_ids[batch_index] == chain_id) * feats["token_pad_mask"][
+                batch_index
+            ]
+            atom_mask = (atom_chain_ids[batch_index] == chain_id) * feats["atom_pad_mask"][
+                batch_index
+            ]
+            token_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()
+                chain_coordinates = sample_coordinates[:, chain_atom_mask]
+                differences = chain_coordinates[:, None, :, :] - chain_coordinates[:, :, None, :]
+                distances = differences.square().sum(dim=-1) ** 0.5
+                valid = chain_resolved[chain_atom_mask]
+                invalid_pairs = 1 - (valid[None, :] * valid[:, None]).to(torch.float32)
+                invalid_pairs[invalid_pairs == 1] = torch.inf
+                nearest = torch.sort(distances + invalid_pairs, dim=2).indices
+                frames = (
+                    torch.cat(
+                        (nearest[:, :, 1:2], nearest[:, :, 0:1], nearest[:, :, 2:3]),
+                        dim=2,
+                    )
+                    + atom_offset
+                )
+                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)
+    expanded_batch, _, _ = pred_atom_coords.shape
+    base_batch = expanded_batch // multiplicity
+    coordinates = pred_atom_coords.reshape(base_batch, multiplicity, -1, 3)
+    frame_indices = frames_idx_true.clone().repeat_interleave(multiplicity, 0)
+    frame_indices = frame_indices.reshape(base_batch, multiplicity, -1, 3)
+    _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]
+    sample_indices = torch.arange(multiplicity, device=frame_indices.device)[None, :, None, None]
+    frame_coordinates = coordinates[
+        batch_indices,
+        sample_indices,
+        frame_indices,
+    ].reshape(-1, 3, 3)
+    valid_frames = compute_collinear_mask(
+        frame_coordinates[:, 1] - frame_coordinates[:, 0],
+        frame_coordinates[:, 1] - frame_coordinates[:, 2],
+    ).reshape(base_batch, multiplicity, -1)
+    return frame_indices, valid_frames * feats["token_pad_mask"][:, None, :]
+
+
+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,
+    )
+    probabilities = nn.functional.softmax(logits, dim=-1)
+    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."""
+
+    distance_scale = 1.24 * (torch.clip(n_res, min=19) - 15) ** (1 / 3) - 1.8
+    return 1 / (1 + (d / distance_scale) ** 2)
+
+
+def _maximum_masked_tm(
+    expected_tm: torch.Tensor,
+    pair_mask: torch.Tensor,
+) -> torch.Tensor:
+    per_anchor = torch.sum(expected_tm * pair_mask, dim=-1) / (torch.sum(pair_mask, dim=-1) + 1e-5)
+    return torch.max(per_anchor, dim=1).values
+
+
+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."""
+
+    _, frame_mask = compute_frame_pred(
+        x_preds,
+        feats["frames_idx"],
+        feats,
+        multiplicity,
+        inference=True,
+    )
+    token_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
+    valid_anchor = frame_mask.reshape(-1, frame_mask.shape[-1])
+    base_pair_mask = valid_anchor[:, :, None] * token_mask[:, None, :] * token_mask[:, :, None]
+    asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)
+    interface_mask = base_pair_mask * (asym_id[:, None, :] != asym_id[:, :, None])
+
+    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)
+    n_res = token_mask.sum(dim=-1, keepdim=True)
+    tm_values = tm_function(pae_centers, n_res).unsqueeze(1).unsqueeze(2)
+    expected_tm = torch.sum(nn.functional.softmax(logits, dim=-1) * tm_values, dim=-1)
+
+    ptm = _maximum_masked_tm(expected_tm, base_pair_mask)
+    iptm = _maximum_masked_tm(expected_tm, interface_mask)
+
+    token_type = feats["mol_type"].repeat_interleave(multiplicity, 0)
+    ligand = (token_type == const.chain_type_ids["NONPOLYMER"]).float()
+    protein = (token_type == const.chain_type_ids["PROTEIN"]).float()
+    ligand_protein = (
+        ligand[:, :, None] * protein[:, None, :] + protein[:, :, None] * ligand[:, None, :]
+    )
+    protein_protein = protein[:, :, None] * protein[:, None, :]
+    ligand_iptm = _maximum_masked_tm(expected_tm, interface_mask * ligand_protein)
+    protein_iptm = _maximum_masked_tm(expected_tm, interface_mask * protein_protein)
+
+    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

@@ -0,0 +1,40 @@
+"""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
+    return keep * (1.0 / (1.0 - probability))

fastplms/models/boltz/vb_layers_initialize.py

@@ -0,0 +1,98 @@
+"""Torch-only parameter initialization for the Boltz2 runtime."""
+
+from __future__ import annotations
+
+import math
+from typing import Literal
+
+import torch
+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."""
+
+    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."""
+
+    trunc_normal_init_(weights)
+
+
+def he_normal_init_(weights: Tensor) -> None:
+    """Initialize W using twice the fan-in variance."""
+
+    trunc_normal_init_(weights, scale=2.0)
+
+
+def glorot_uniform_init_(weights: Tensor) -> None:
+    """Initialize W with Xavier uniform values."""
+
+    torch.nn.init.xavier_uniform_(weights, gain=1.0)
+
+
+def _fill_(tensor: Tensor, value: float) -> None:
+    with torch.no_grad():
+        tensor.fill_(value)
+
+
+def final_init_(weights: Tensor) -> None:
+    """Zero the final projection W."""
+
+    _fill_(weights, 0.0)
+
+
+def gating_init_(weights: Tensor) -> None:
+    """Zero the gating projection W."""
+
+    _fill_(weights, 0.0)
+
+
+def bias_init_zero_(bias: Tensor) -> None:
+    """Set the bias tensor to zero."""
+
+    _fill_(bias, 0.0)
+
+
+def bias_init_one_(bias: Tensor) -> None:
+    """Set the bias tensor to one."""
+
+    _fill_(bias, 1.0)
+
+
+def normal_init_(weights: Tensor) -> None:
+    """Initialize W with linear Kaiming-normal values."""
+
+    torch.nn.init.kaiming_normal_(weights, nonlinearity="linear")
+
+
+def ipa_point_weights_init_(weights: Tensor) -> None:
+    """Initialize W so applying softplus yields one."""
+
+    _fill_(weights, 0.541324854612918)

fastplms/models/boltz/vb_layers_outer_product_mean.py

@@ -0,0 +1,87 @@
+"""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."""
+
+        counts: Tensor | None = None
+        for start in range(0, mask.shape[1], 64):
+            mask_slice = mask[:, start : start + 64]
+            contribution = (mask_slice[:, :, None, :] * mask_slice[:, :, :, None]).sum(dim=1)
+            counts = contribution if counts is None else counts + contribution
+        if counts is None:
+            raise ValueError("MSA depth must be positive")
+        return counts.clamp(min=1)
+
+    @staticmethod
+    def _outer_product(left: Tensor, right: Tensor, counts: Tensor) -> Tensor:
+        outer = torch.einsum("bsic,bsjd->bijcd", left, right)
+        return outer.reshape(*outer.shape[:3], -1) / counts
+
+    def _chunked_projection(
+        self,
+        left: Tensor,
+        right: Tensor,
+        counts: Tensor,
+        chunk_size: int,
+        target: Tensor,
+    ) -> Tensor:
+        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)
+            weight = self.proj_o.weight[
+                :,
+                start * self.c_hidden : stop * self.c_hidden,
+            ]
+            contribution = outer.to(target) @ weight.T
+            output = contribution if output is None else output + contribution
+        if output is None:
+            raise ValueError("c_hidden must be positive")
+        return output + self.proj_o.bias
+
+    def forward(self, m: Tensor, mask: Tensor, chunk_size: int | None = None) -> Tensor:
+        """Return pair tensor Z with shape ``(b, l, l, d_z)``."""
+
+        expanded_mask = mask.unsqueeze(-1).to(m)
+        normalized = self.norm(m)
+        left = self.proj_a(normalized) * expanded_mask
+        right = self.proj_b(normalized) * expanded_mask
+
+        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)
+            return self._chunked_projection(
+                left,
+                right,
+                counts,
+                chunk_size,
+                normalized,
+            )
+
+        pair_mask = expanded_mask[:, :, None, :] * expanded_mask[:, :, :, None]
+        counts = pair_mask.sum(dim=1).clamp(min=1)
+        outer = self._outer_product(left.float(), right.float(), counts)
+        return self.proj_o(outer.to(normalized))

fastplms/models/boltz/vb_layers_pair_averaging.py

@@ -0,0 +1,91 @@
+"""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:
+        logits = self.proj_z(pair_states).permute(0, 3, 1, 2)
+        logits = logits + (1 - mask[:, None]) * -self.inf
+        return torch.softmax(logits, dim=-1)
+
+    def _all_heads(self, msa_states: Tensor, pair_states: Tensor, mask: Tensor) -> Tensor:
+        values = self.proj_m(msa_states).reshape(
+            *msa_states.shape[:3],
+            self.num_heads,
+            self.c_h,
+        )
+        values = values.permute(0, 3, 1, 2, 4)
+        weights = self._attention_weights(pair_states, mask)
+        gate = self.proj_g(msa_states).sigmoid()
+        attended = torch.einsum("bhij,bhsjd->bhsid", weights, values)
+        attended = attended.permute(0, 2, 3, 1, 4)
+        attended = attended.reshape(*attended.shape[:3], self.num_heads * self.c_h)
+        return self.proj_o(gate * attended)
+
+    def _one_head(self, msa_states: Tensor, pair_states: Tensor, mask: Tensor, head: int) -> Tensor:
+        start = head * self.c_h
+        stop = start + self.c_h
+        value = msa_states @ self.proj_m.weight[start:stop].T
+        value = value.reshape(*value.shape[:3], 1, self.c_h).permute(0, 3, 1, 2, 4)
+        logits = pair_states @ self.proj_z.weight[head : head + 1].T
+        logits = logits.permute(0, 3, 1, 2)
+        weights = torch.softmax(logits + (1 - mask[:, None]) * -self.inf, dim=-1)
+        gate = (msa_states @ self.proj_g.weight[start:stop].T).sigmoid()
+        attended = torch.einsum("bhij,bhsjd->bhsid", weights, value)
+        attended = attended.permute(0, 2, 3, 1, 4).reshape(
+            *msa_states.shape[:3],
+            self.c_h,
+        )
+        return (gate * attended) @ self.proj_o.weight[:, start:stop].T
+
+    def forward(
+        self,
+        m: Tensor,
+        z: Tensor,
+        mask: Tensor,
+        chunk_heads: bool = True,
+    ) -> Tensor:
+        """Return updated M with shape ``(b, s, l, d_m)``."""
+
+        msa_states = self.norm_m(m)
+        pair_states = self.norm_z(z)
+        if not chunk_heads or self.training:
+            return self._all_heads(msa_states, pair_states, mask)
+
+        output: Tensor | None = None
+        for head in range(self.num_heads):
+            contribution = self._one_head(msa_states, pair_states, mask, head)
+            output = contribution if output is None else output + contribution
+        if output is None:
+            raise ValueError("num_heads must be positive")
+        return output

fastplms/models/boltz/vb_layers_pairformer.py

@@ -0,0 +1,337 @@
+"""Pairformer layers for joint or pair-only Boltz2 representations."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import torch
+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."""
+
+    multiplication_kernel = use_kernels or use_cuequiv_mul
+    attention_kernel = use_kernels or use_cuequiv_attn
+    output = pair_states
+    for update in (modules.tri_mul_out, modules.tri_mul_in):
+        dropout = get_dropout_mask(modules.dropout, output, modules.training)
+        output = output + dropout * update(
+            output,
+            mask=pair_mask,
+            use_kernels=multiplication_kernel,
+        )
+
+    dropout = get_dropout_mask(modules.dropout, output, modules.training)
+    output = output + dropout * modules.tri_att_start(
+        output,
+        mask=pair_mask,
+        chunk_size=chunk_size,
+        use_kernels=attention_kernel,
+    )
+    dropout = get_dropout_mask(
+        modules.dropout,
+        output,
+        modules.training,
+        columnwise=True,
+    )
+    output = output + dropout * modules.tri_att_end(
+        output,
+        mask=pair_mask,
+        chunk_size=chunk_size,
+        use_kernels=attention_kernel,
+    )
+    return cast(Tensor, output + modules.transition_z(output))
+
+
+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,
+        )
+        with torch.autocast("cuda", enabled=False):
+            normalized = self.pre_norm_s(s.float())
+            sequence_output = s.float() + self.attention(
+                s=normalized,
+                z=pair_output.float(),
+                mask=mask.float(),
+                k_in=normalized,
+            )
+            sequence_output = sequence_output + self.transition_s(sequence_output)
+            sequence_output = cast(Tensor, self.s_post_norm(sequence_output))
+        return sequence_output, pair_output
+
+
+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."""
+
+        chunk_size = _triangle_chunk_size(z, self.training)
+        sequence_output, pair_output = s, 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,
+                )
+            else:
+                sequence_output, pair_output = layer(
+                    sequence_output,
+                    pair_output,
+                    mask,
+                    pair_mask,
+                    chunk_size,
+                    use_kernels,
+                )
+        return sequence_output, pair_output
+
+
+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."""
+
+        chunk_size = _triangle_chunk_size(z, self.training)
+        output = z
+        for layer in self.layers:
+            if self.activation_checkpointing:
+                output = checkpoint(
+                    layer,
+                    output,
+                    pair_mask,
+                    chunk_size,
+                    use_kernels,
+                    use_reentrant=False,
+                )
+            else:
+                output = layer(output, pair_mask, chunk_size, use_kernels)
+        return output

fastplms/models/boltz/vb_layers_transition.py

@@ -0,0 +1,69 @@
+"""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."""
+
+        gate = self.silu(torch.matmul(normalized, self.fc1.weight[start:stop].T))
+        value = torch.matmul(normalized, self.fc2.weight[start:stop].T)
+        return torch.matmul(gate * value, self.fc3.weight[:, start:stop].T)
+
+    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)
+        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)
+        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),
+            )
+            output = contribution if output is None else output + contribution
+        if output is None:  # hidden is normally positive; keep malformed configs explicit.
+            raise ValueError("hidden must be positive")
+        return output

fastplms/models/boltz/vb_layers_triangular_mult.py

@@ -0,0 +1,125 @@
+"""Incoming and outgoing triangular multiplicative pair updates."""
+
+from __future__ import annotations
+
+import importlib
+from importlib.util import find_spec
+from typing import Literal
+
+import torch
+from torch import Tensor, nn
+
+from . import vb_layers_initialize as init
+
+TriangleDirection = Literal["incoming", "outgoing"]
+
+
+@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."""
+
+    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:
+        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,
+        )
+
+    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)
+        projected = self.p_in(normalized) * self.g_in(normalized).sigmoid()
+        projected = projected * mask.unsqueeze(-1)
+        left, right = torch.chunk(projected.float(), 2, dim=-1)
+        combined = torch.einsum(self.equation, left, right)
+        return self.p_out(self.norm_out(combined)) * self.g_out(normalized).sigmoid()
+
+
+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

@@ -0,0 +1,178 @@
+"""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:
+    return (coordinates * weights).sum(dim=-2, keepdim=True) / weights.sum(
+        dim=-2,
+        keepdim=True,
+    )
+
+
+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)
+
+    true_centroid = _weighted_centroid(true_coords, point_weights)
+    pred_centroid = _weighted_centroid(pred_coords, point_weights)
+    true_centered = true_coords - true_centroid
+    pred_centered = pred_coords - pred_centroid
+
+    covariance = einsum(
+        point_weights * pred_centered,
+        true_centered,
+        "... n i, ... n j -> ... i j",
+    )
+    original_dtype = covariance.dtype
+    covariance_fp32 = covariance.to(torch.float32)
+    left_vectors, singular_values, right_vectors_h = torch.linalg.svd(
+        covariance_fp32,
+        driver="gesvd" if covariance_fp32.is_cuda else None,
+    )
+    right_vectors = right_vectors_h.mH
+    _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)
+    orientation = torch.eye(
+        coordinate_dim,
+        dtype=covariance_fp32.dtype,
+        device=covariance.device,
+    )[None].repeat(*batch_shape, 1, 1)
+    orientation[..., -1, -1] = torch.det(preliminary_rotation)
+    rotation = einsum(
+        left_vectors,
+        orientation,
+        right_vectors,
+        "... i j, ... j k, ... l k -> ... i l",
+    ).to(original_dtype)
+
+    aligned = einsum(true_centered, rotation, "... n i, ... j i -> ... n j") + pred_centroid
+    aligned.detach_()
+    return aligned
+
+
+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:
+    true_distances = torch.cdist(true_coords, true_coords)
+    nucleotide_rows = is_nucleotide.bool().unsqueeze(-1).expand_as(true_distances)
+    pair_mask = torch.where(
+        nucleotide_rows,
+        true_distances < nucleic_acid_cutoff,
+        true_distances < other_cutoff,
+    )
+    pair_mask &= ~torch.eye(
+        pred_coords.shape[0],
+        dtype=torch.bool,
+        device=pred_coords.device,
+    )
+    coordinate_rows = coords_mask.bool()
+    pair_mask &= coordinate_rows.unsqueeze(-1)
+    pair_mask &= coordinate_rows.unsqueeze(-2)
+
+    pair_indices = pair_mask.nonzero()
+    true_pair_distances = true_distances[pair_indices[:, 0], pair_indices[:, 1]]
+    pred_pair_distances = functional.pairwise_distance(
+        pred_coords[pair_indices[:, 0]],
+        pred_coords[pair_indices[:, 1]],
+    )
+    distance_error = torch.abs(true_pair_distances - pred_pair_distances)
+    smooth_agreement = (
+        sum(torch.sigmoid(threshold - distance_error) for threshold in (0.5, 1.0, 2.0, 4.0)) / 4.0
+    )
+    return smooth_agreement.sum() / (pair_indices.shape[0] + 1e-5)
+
+
+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])
+    ]
+    return 1.0 - torch.stack(agreements).mean(dim=0)

fastplms/models/boltz/vb_modules_confidencev2.py

@@ -0,0 +1,542 @@
+import torch
+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.
+    """
+
+    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()
+    assignments = atom_to_token.bool() & valid_atoms.unsqueeze(-1)
+    token_index = assignments.to(dtype=torch.int64).argmax(dim=-1)
+    # 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
+    slot_index = (cumulative_slots * assignments).sum(dim=-1)
+
+    slots_per_token = token_logits.shape[2]
+    flattened_index = token_index * slots_per_token + slot_index
+    flattened_index = flattened_index.masked_fill(~valid_atoms, 0)
+    flattened_index = flattened_index.repeat_interleave(multiplicity, dim=0)
+    expanded_atom_mask = valid_atoms.repeat_interleave(multiplicity, dim=0)
+
+    flattened_logits = token_logits.flatten(1, 2)
+    gather_index = flattened_index.unsqueeze(-1).expand(
+        -1,
+        atom_count,
+        flattened_logits.shape[-1],
+    )
+    atom_logits = torch.gather(flattened_logits, dim=1, index=gather_index)
+    return atom_logits * expanded_atom_mask.unsqueeze(-1).to(dtype=atom_logits.dtype)
+
+
+class ConfidenceModule(nn.Module):
+    """Algorithm 31"""
+
+    def __init__(
+        self,
+        token_s,
+        token_z,
+        pairformer_args: dict,
+        num_dist_bins=64,
+        token_level_confidence=True,
+        max_dist=22,
+        add_s_to_z_prod=False,
+        add_s_input_to_s=False,
+        add_z_input_to_z=False,
+        maximum_bond_distance=0,
+        bond_type_feature=False,
+        confidence_args: dict | None = None,
+        compile_pairformer=False,
+        fix_sym_check=False,
+        cyclic_pos_enc=False,
+        return_latent_feats=False,
+        conditioning_cutoff_min=None,
+        conditioning_cutoff_max=None,
+        **kwargs,
+    ):
+        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)
+        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,  # Float['b n ts']
+        s,  # Float['b n ts']
+        z,  # Float['b n n tz']
+        x_pred,  # Float['bm m 3']
+        feats,
+        pred_distogram_logits,
+        multiplicity=1,
+        run_sequentially=False,
+        use_kernels: bool = False,
+    ):
+        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
+            elif x_pred.shape[0] == batch_size * multiplicity:
+                sample_coordinates = x_pred.reshape(
+                    batch_size,
+                    multiplicity,
+                    *x_pred.shape[1:],
+                )
+            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]
+                    out_dict[key] = torch.stack(values, dim=1).flatten(0, 1)
+                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
+                            ]
+                            chains_iptm[chain_idx2] = torch.stack(
+                                values,
+                                dim=1,
+                            ).flatten(0, 1)
+                        pair_chains_iptm[chain_idx1] = chains_iptm
+                    out_dict[key] = pair_chains_iptm
+            return out_dict
+
+        s_inputs = self.s_inputs_norm(s_inputs)
+        if not self.no_update_s:
+            s = self.s_norm(s)
+
+        if self.add_s_input_to_s:
+            s = s + self.s_input_to_s(s_inputs)
+
+        z = self.z_norm(z)
+
+        if self.add_z_input_to_z:
+            relative_position_encoding = self.rel_pos(feats)
+            z = z + relative_position_encoding
+            z = z + self.token_bonds(feats["token_bonds"].float())
+            if self.bond_type_feature:
+                z = z + self.token_bonds_type(feats["type_bonds"].long())
+            z = z + self.contact_conditioning(feats)
+
+        s = s.repeat_interleave(multiplicity, 0)
+
+        z = (
+            z
+            + self.s_to_z(s_inputs)[:, :, None, :]
+            + self.s_to_z_transpose(s_inputs)[:, None, :, :]
+        )
+        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, :, :]
+            )
+
+        z = z.repeat_interleave(multiplicity, 0)
+        s_inputs = s_inputs.repeat_interleave(multiplicity, 0)
+
+        token_to_rep_atom = feats["token_to_rep_atom"]
+        token_to_rep_atom = token_to_rep_atom.repeat_interleave(multiplicity, 0)
+        if len(x_pred.shape) == 4:
+            b, multiplicity, n, _ = x_pred.shape
+            x_pred = x_pred.reshape(b * multiplicity, n, -1)
+        else:
+            _, n, _ = x_pred.shape
+        x_pred_repr = torch.bmm(token_to_rep_atom.float(), x_pred)
+        d = torch.cdist(x_pred_repr, x_pred_repr)
+        distogram = (d.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()
+        distogram = self.dist_bin_pairwise_embed(distogram)
+        z = z + distogram
+
+        mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
+        pair_mask = mask[:, :, None] * mask[:, None, :]
+
+        s_t, z_t = self.pairformer_stack(
+            s, z, mask=mask, pair_mask=pair_mask, use_kernels=use_kernels
+        )
+
+        # AF3 has residual connections, we remove them
+        s = s_t
+        z = z_t
+
+        out_dict = {}
+
+        if self.return_latent_feats:
+            out_dict["s_conf"] = s
+            out_dict["z_conf"] = z
+
+        # confidence heads
+        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,
+        token_z,
+        num_plddt_bins=50,
+        num_pde_bins=64,
+        num_pae_bins=64,
+        token_level_confidence=True,
+        use_separate_heads: bool = False,
+        **kwargs,
+    ):
+        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,  # Float['b n ts']
+        z,  # Float['b n n tz']
+        x_pred,  # Float['bm m 3']
+        d,
+        feats,
+        pred_distogram_logits,
+        multiplicity=1,
+    ):
+        if self.use_separate_heads:
+            asym_id_token = feats["asym_id"]
+            is_same_chain = asym_id_token.unsqueeze(-1) == asym_id_token.unsqueeze(-2)
+            is_different_chain = ~is_same_chain
+
+        if self.use_separate_heads:
+            pae_intra_logits = self.to_pae_intra_logits(z)
+            pae_intra_logits = pae_intra_logits * is_same_chain.float().unsqueeze(-1)
+
+            pae_inter_logits = self.to_pae_inter_logits(z)
+            pae_inter_logits = pae_inter_logits * is_different_chain.float().unsqueeze(-1)
+
+            pae_logits = pae_inter_logits + pae_intra_logits
+        else:
+            pae_logits = self.to_pae_logits(z)
+
+        if self.use_separate_heads:
+            pde_intra_logits = self.to_pde_intra_logits(z + z.transpose(1, 2))
+            pde_intra_logits = pde_intra_logits * is_same_chain.float().unsqueeze(-1)
+
+            pde_inter_logits = self.to_pde_inter_logits(z + z.transpose(1, 2))
+            pde_inter_logits = pde_inter_logits * is_different_chain.float().unsqueeze(-1)
+
+            pde_logits = pde_inter_logits + pde_intra_logits
+        else:
+            pde_logits = self.to_pde_logits(z + z.transpose(1, 2))
+        resolved_logits = self.to_resolved_logits(s)
+        plddt_logits = self.to_plddt_logits(s)
+
+        ligand_weight = 20
+        non_interface_weight = 1
+        interface_weight = 10
+
+        token_type = feats["mol_type"]
+        token_type = token_type.repeat_interleave(multiplicity, 0)
+        is_ligand_token = (token_type == const.chain_type_ids["NONPOLYMER"]).float()
+
+        if self.token_level_confidence:
+            plddt = compute_aggregated_metric(plddt_logits)
+            token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
+            complex_plddt = (plddt * token_pad_mask).sum(dim=-1) / token_pad_mask.sum(dim=-1)
+
+            is_contact = (d < 8).float()
+            is_different_chain = (
+                feats["asym_id"].unsqueeze(-1) != feats["asym_id"].unsqueeze(-2)
+            ).float()
+            is_different_chain = is_different_chain.repeat_interleave(multiplicity, 0)
+            token_interface_mask = torch.max(
+                is_contact * is_different_chain * (1 - is_ligand_token).unsqueeze(-1),
+                dim=-1,
+            ).values
+            token_non_interface_mask = (1 - token_interface_mask) * (1 - is_ligand_token)
+            iplddt_weight = (
+                is_ligand_token * ligand_weight
+                + token_interface_mask * interface_weight
+                + token_non_interface_mask * non_interface_weight
+            )
+            complex_iplddt = (plddt * token_pad_mask * iplddt_weight).sum(dim=-1) / torch.sum(
+                token_pad_mask * iplddt_weight, dim=-1
+            )
+
+        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)
+            resolved_logits = _token_slot_logits_to_atom_logits(
+                resolved_logits,
+                feats["atom_to_token"],
+                feats["atom_pad_mask"],
+                multiplicity=multiplicity,
+            )
+            plddt_logits = plddt_logits.reshape(b, n, self.max_num_atoms_per_token, -1)
+            plddt_logits = _token_slot_logits_to_atom_logits(
+                plddt_logits,
+                feats["atom_to_token"],
+                feats["atom_pad_mask"],
+                multiplicity=multiplicity,
+            )
+            atom_pad_mask = feats["atom_pad_mask"].repeat_interleave(multiplicity, 0)
+            plddt = compute_aggregated_metric(plddt_logits)
+
+            complex_plddt = (plddt * atom_pad_mask).sum(dim=-1) / atom_pad_mask.sum(dim=-1)
+            atom_to_token = feats["atom_to_token"].float().repeat_interleave(
+                multiplicity,
+                0,
+            )
+            chain_id_token = feats["asym_id"].float().repeat_interleave(
+                multiplicity,
+                0,
+            )
+            atom_type = torch.bmm(
+                atom_to_token,
+                token_type.float().unsqueeze(-1),
+            ).squeeze(-1)
+            is_ligand_atom = (atom_type == const.chain_type_ids["NONPOLYMER"]).float()
+            d_atom = torch.cdist(x_pred, x_pred)
+            is_contact = (d_atom < 8).float()
+            chain_id_atom = torch.bmm(atom_to_token, chain_id_token.unsqueeze(-1)).squeeze(-1)
+            is_different_chain = (
+                chain_id_atom.unsqueeze(-1) != chain_id_atom.unsqueeze(-2)
+            ).float()
+
+            atom_interface_mask = torch.max(
+                is_contact * is_different_chain * (1 - is_ligand_atom).unsqueeze(-1),
+                dim=-1,
+            ).values
+            atom_non_interface_mask = (1 - atom_interface_mask) * (1 - is_ligand_atom)
+            iplddt_weight = (
+                is_ligand_atom * ligand_weight
+                + atom_interface_mask * interface_weight
+                + atom_non_interface_mask * non_interface_weight
+            )
+
+            complex_iplddt = (plddt * atom_pad_mask * iplddt_weight).sum(dim=-1) / torch.sum(
+                atom_pad_mask * iplddt_weight,
+                dim=-1,
+            )
+
+        # Compute the gPDE and giPDE
+        pde = compute_aggregated_metric(pde_logits, end=32)
+        pred_distogram_prob = nn.functional.softmax(
+            pred_distogram_logits, dim=-1
+        ).repeat_interleave(multiplicity, 0)
+        contacts = torch.zeros((1, 1, 1, 64), dtype=pred_distogram_prob.dtype).to(
+            pred_distogram_prob.device
+        )
+        contacts[:, :, :, :20] = 1.0
+        prob_contact = (pred_distogram_prob * contacts).sum(-1)
+        token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
+        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))
+        )
+        token_pair_mask = token_pad_pair_mask * prob_contact
+        complex_pde_numerator = (pde * token_pair_mask).sum(dim=(1, 2))
+        complex_pde_denominator = token_pair_mask.sum(dim=(1, 2))
+        complex_pde = complex_pde_numerator / torch.where(
+            complex_pde_denominator > 0,
+            complex_pde_denominator,
+            torch.ones_like(complex_pde_denominator),
+        )
+        asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0)
+        token_interface_pair_mask = token_pair_mask * (
+            asym_id.unsqueeze(-1) != asym_id.unsqueeze(-2)
+        )
+        complex_ipde = (pde * token_interface_pair_mask).sum(dim=(1, 2)) / (
+            token_interface_pair_mask.sum(dim=(1, 2)) + 1e-5
+        )
+        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
+        out_dict["pae"] = compute_aggregated_metric(pae_logits, end=32)
+
+        ptm, iptm, ligand_iptm, protein_iptm, pair_chains_iptm = compute_ptms(
+            pae_logits, x_pred, feats, multiplicity
+        )
+        out_dict["ptm"] = ptm
+        out_dict["iptm"] = iptm
+        out_dict["ligand_iptm"] = ligand_iptm
+        out_dict["protein_iptm"] = protein_iptm
+        out_dict["pair_chains_iptm"] = pair_chains_iptm
+
+        return out_dict

fastplms/models/boltz/vb_modules_diffusion_conditioning.py

@@ -0,0 +1,137 @@
+"""Precompute pair and atom biases used by the Boltz2 diffusion stack."""
+
+from __future__ import annotations
+
+import torch
+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:
+    return torch.cat([projection(pair_features) for projection in projections], dim=-1)
+
+
+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,
+        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.
+        """
+
+        z_conditioned = self.pairwise_conditioner(
+            z_trunk,
+            relative_position_encoding,
+        )
+        q, c, p, to_keys = self.atom_encoder(
+            feats=feats,
+            s_trunk=s_trunk,
+            z=z_conditioned,
+        )
+        atom_encoder_bias = _concatenate_biases(self.atom_enc_proj_z, p)
+        atom_decoder_bias = _concatenate_biases(self.atom_dec_proj_z, p)
+        token_transformer_bias = _concatenate_biases(
+            self.token_trans_proj_z,
+            z_conditioned,
+        )
+        return (
+            q,
+            c,
+            to_keys,
+            atom_encoder_bias,
+            atom_decoder_bias,
+            token_transformer_bias,
+        )

fastplms/models/boltz/vb_modules_diffusionv2.py

@@ -0,0 +1,654 @@
+# Based on https://github.com/lucidrains/alphafold3-pytorch.
+# MIT License, Copyright (c) 2024 Phil Wang.
+
+from __future__ import annotations
+
+from math import sqrt
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+from torch.nn import Module
+
+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):
+    """Diffusion module"""
+
+    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
+
+        # conditioning
+        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,
+            # post_layer_norm=transformer_post_ln,
+        )
+
+        self.a_norm = nn.LayerNorm(2 * token_s)  # if not transformer_post_ln else nn.Identity()
+
+        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,
+            # transformer_post_layer_norm=transformer_post_ln,
+        )
+
+    def forward(
+        self,
+        s_inputs,  # Float['b n ts']
+        s_trunk,  # Float['b n ts']
+        r_noisy,  # Float['bm m 3']
+        times,  # Float['bm 1 1']
+        feats,
+        diffusion_conditioning,
+        multiplicity=1,
+    ):
+        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,
+            )
+        else:
+            s, _normed_fourier = self.single_conditioner(
+                times,
+                s_trunk.repeat_interleave(multiplicity, 0),
+                s_inputs.repeat_interleave(multiplicity, 0),
+            )
+
+        # 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,  # Float['b m 3'],
+            multiplicity=multiplicity,
+        )
+
+        # Full self-attention on token level
+        a = a + self.s_to_a_linear(s)
+
+        mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0)
+        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,
+        )
+        a = self.a_norm(a)
+
+        # 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,
+        )
+
+        return r_update
+
+
+class AtomDiffusion(Module):
+    def __init__(
+        self,
+        score_model_args,
+        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 | None = None,
+        coordinate_augmentation: bool = True,
+        coordinate_augmentation_inference=None,
+        compile_score: bool = False,
+        alignment_reverse_diff: bool = False,
+        synchronize_sigmas: bool = False,
+    ):
+        super().__init__()
+        self.score_model = DiffusionModule(
+            **score_model_args,
+        )
+        if compile_score:
+            self.score_model = torch.compile(self.score_model, dynamic=False, fullgraph=False)
+
+        # parameters
+        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):
+        return next(self.score_model.parameters()).device
+
+    def c_skip(self, sigma):
+        return (self.sigma_data**2) / (sigma**2 + self.sigma_data**2)
+
+    def c_out(self, sigma):
+        return sigma * self.sigma_data / torch.sqrt(self.sigma_data**2 + sigma**2)
+
+    def c_in(self, sigma):
+        return 1 / torch.sqrt(sigma**2 + self.sigma_data**2)
+
+    def c_noise(self, sigma):
+        return log(sigma / self.sigma_data) * 0.25
+
+    def preconditioned_network_forward(
+        self,
+        noised_atom_coords,  #: Float['b m 3'],
+        sigma,  #: Float['b'] | Float[' '] | float,
+        network_condition_kwargs: dict,
+    ):
+        batch, device = noised_atom_coords.shape[0], noised_atom_coords.device
+
+        if isinstance(sigma, float):
+            sigma = torch.full((batch,), sigma, device=device)
+
+        padded_sigma = rearrange(sigma, "b -> b 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,
+        )
+
+        denoised_coords = (
+            self.c_skip(padded_sigma) * noised_atom_coords + self.c_out(padded_sigma) * r_update
+        )
+        return denoised_coords
+
+    def sample_schedule(self, num_sampling_steps=None):
+        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)
+        sigmas = (
+            self.sigma_max**inv_rho
+            + steps / (num_sampling_steps - 1) * (self.sigma_min**inv_rho - self.sigma_max**inv_rho)
+        ) ** self.rho
+
+        sigmas = sigmas * self.sigma_data
+
+        sigmas = F.pad(sigmas, (0, 1), value=0.0)  # last step is sigma value of 0.
+        return sigmas
+
+    def sample(
+        self,
+        atom_mask,
+        num_sampling_steps=None,
+        multiplicity=1,
+        max_parallel_samples=None,
+        steering_args=None,
+        **network_condition_kwargs,
+    ):
+        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)
+            resample_weights = torch.ones(multiplicity, device=self.device).reshape(
+                -1, steering_args["num_particles"]
+            )
+        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,
+            )
+        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)
+
+        shape = (*atom_mask.shape, 3)
+
+        # Pair each schedule value with the next sigma and gamma.
+        sigmas = self.sample_schedule(num_sampling_steps)
+        gammas = torch.where(sigmas > self.gamma_min, self.gamma_0, 0.0)
+        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)
+        token_repr = None
+        atom_coords_denoised = None
+
+        # gradually denoise
+        for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(sigmas_and_gammas):
+            random_R, random_tr = compute_random_augmentation(
+                multiplicity, device=atom_coords.device, dtype=atom_coords.dtype
+            )
+            atom_coords = atom_coords - atom_coords.mean(dim=-2, keepdims=True)
+            atom_coords = torch.einsum("bmd,bds->bms", atom_coords, random_R) + random_tr
+            if atom_coords_denoised is not None:
+                atom_coords_denoised -= atom_coords_denoised.mean(dim=-2, keepdims=True)
+                atom_coords_denoised = (
+                    torch.einsum("bmd,bds->bms", atom_coords_denoised, random_R) + random_tr
+                )
+            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
+                )
+
+            sigma_tm, sigma_t, gamma = sigma_tm.item(), sigma_t.item(), gamma.item()
+
+            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)
+            atom_coords_noisy = atom_coords + eps
+
+            with torch.no_grad():
+                atom_coords_denoised = torch.zeros_like(atom_coords_noisy)
+                sample_ids = torch.arange(multiplicity).to(atom_coords_noisy.device)
+                sample_ids_chunks = sample_ids.chunk(multiplicity % max_parallel_samples + 1)
+
+                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,
+                        ),
+                    )
+                    atom_coords_denoised[sample_ids_chunk] = atom_coords_denoised_chunk
+
+                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)
+                    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,
+                            )
+                            energy += parameters["resampling_weight"] * component_energy
+                    energy_traj = torch.cat((energy_traj, energy.unsqueeze(1)), dim=1)
+
+                    # Compute log G values
+                    if step_idx == 0:
+                        log_G = -1 * energy
+                    else:
+                        log_G = energy_traj[:, -2] - energy_traj[:, -1]
+
+                    # 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)
+                    else:
+                        ll_difference = torch.zeros_like(energy)
+
+                    # Compute resampling weights
+                    resample_weights = F.softmax(
+                        (ll_difference + steering_args["fk_lambda"] * log_G).reshape(
+                            -1, steering_args["num_particles"]
+                        ),
+                        dim=1,
+                    )
+
+                # 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)
+                    for guidance_step in range(steering_args["num_gd_steps"]):
+                        energy_gradient = torch.zeros_like(atom_coords_denoised)
+                        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,
+                                )
+                        guidance_update -= energy_gradient
+                    atom_coords_denoised += guidance_update
+                    scaled_guidance_update = (
+                        guidance_update * -1 * self.step_scale * (sigma_t - t_hat) / t_hat
+                    )
+
+                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()
+
+                    atom_coords = atom_coords[resample_indices]
+                    atom_coords_noisy = atom_coords_noisy[resample_indices]
+                    atom_mask = atom_mask[resample_indices]
+                    if atom_coords_denoised is not None:
+                        atom_coords_denoised = atom_coords_denoised[resample_indices]
+                    energy_traj = energy_traj[resample_indices]
+                    if (
+                        steering_args["physical_guidance_update"]
+                        or steering_args["contact_guidance_update"]
+                    ):
+                        scaled_guidance_update = scaled_guidance_update[resample_indices]
+                    if token_repr is not None:
+                        token_repr = token_repr[resample_indices]
+
+            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(),
+                    )
+
+                atom_coords_noisy = atom_coords_noisy.to(atom_coords_denoised)
+
+            denoised_over_sigma = (atom_coords_noisy - atom_coords_denoised) / t_hat
+            atom_coords_next = (
+                atom_coords_noisy + step_scale * (sigma_t - t_hat) * denoised_over_sigma
+            )
+
+            atom_coords = atom_coords_next
+
+        return dict(sample_atom_coords=atom_coords, diff_token_repr=token_repr)
+
+    def loss_weight(self, sigma):
+        return (sigma**2 + self.sigma_data**2) / ((sigma * self.sigma_data) ** 2)
+
+    def noise_distribution(self, 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,
+        s_trunk,
+        feats,
+        diffusion_conditioning,
+        multiplicity=1,
+    ):
+        # training diffusion step
+        batch_size = feats["coords"].shape[0] // multiplicity
+
+        if self.synchronize_sigmas:
+            sigmas = self.noise_distribution(batch_size).repeat_interleave(multiplicity, 0)
+        else:
+            sigmas = self.noise_distribution(batch_size * multiplicity)
+        padded_sigmas = rearrange(sigmas, "b -> b 1 1")
+
+        atom_coords = feats["coords"]
+
+        atom_mask = feats["atom_pad_mask"]
+        atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
+
+        atom_coords = center_random_augmentation(
+            atom_coords, atom_mask, augmentation=self.coordinate_augmentation
+        )
+
+        noise = torch.randn_like(atom_coords)
+        noised_atom_coords = atom_coords + padded_sigmas * noise
+
+        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,
+            },
+        )
+
+        return {
+            "denoised_atom_coords": denoised_atom_coords,
+            "sigmas": sigmas,
+            "aligned_true_atom_coords": atom_coords,
+        }
+
+    def compute_loss(
+        self,
+        feats,
+        out_dict,
+        add_smooth_lddt_loss=True,
+        nucleotide_loss_weight=5.0,
+        ligand_loss_weight=10.0,
+        multiplicity=1,
+        filter_by_plddt=0.0,
+    ):
+        with torch.autocast("cuda", enabled=False):
+            denoised_atom_coords = out_dict["denoised_atom_coords"].float()
+            sigmas = out_dict["sigmas"].float()
+
+            resolved_atom_mask_uni = feats["atom_resolved_mask"].float()
+
+            if filter_by_plddt > 0:
+                plddt_mask = feats["plddt"] > filter_by_plddt
+                resolved_atom_mask_uni = resolved_atom_mask_uni * plddt_mask.float()
+
+            resolved_atom_mask = resolved_atom_mask_uni.repeat_interleave(multiplicity, 0)
+
+            align_weights = denoised_atom_coords.new_ones(denoised_atom_coords.shape[:2])
+            atom_type = (
+                torch.bmm(
+                    feats["atom_to_token"].float(),
+                    feats["mol_type"].unsqueeze(-1).float(),
+                )
+                .squeeze(-1)
+                .long()
+            )
+            atom_type_mult = atom_type.repeat_interleave(multiplicity, 0)
+
+            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()
+            )
+
+            atom_coords = out_dict["aligned_true_atom_coords"].float()
+            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(),
+            )
+
+            # Cast back
+            atom_coords_aligned_ground_truth = atom_coords_aligned_ground_truth.to(
+                denoised_atom_coords
+            )
+
+            # weighted MSE loss of denoised atom positions
+            mse_loss = ((denoised_atom_coords - atom_coords_aligned_ground_truth) ** 2).sum(dim=-1)
+            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
+            )
+
+            # weight by sigma factor
+            loss_weights = self.loss_weight(sigmas)
+            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

@@ -0,0 +1,577 @@
+"""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
+
+from collections.abc import Callable
+from functools import partial
+from math import pi
+
+import torch
+from torch import nn
+from torch.nn.functional import one_hot
+
+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))
+        return torch.cos(2 * pi * random_phase)
+
+
+def _pairwise_difference(values: torch.Tensor) -> torch.Tensor:
+    return values[:, :, None] - values[:, None, :]
+
+
+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, :]
+        same_residue = feats["residue_index"][:, :, None] == feats["residue_index"][:, None, :]
+        same_entity = feats["entity_id"][:, :, None] == feats["entity_id"][:, None, :]
+
+        residue_offset = _pairwise_difference(feats["residue_index"])
+        if self.cyclic_pos_enc and torch.any(feats["cyclic_period"] > 0):
+            period = torch.where(
+                feats["cyclic_period"] > 0,
+                feats["cyclic_period"],
+                torch.zeros_like(feats["cyclic_period"]) + 10000,
+            )
+            residue_offset = (residue_offset - period * torch.round(residue_offset / period)).long()
+        residue_offset = torch.clip(
+            residue_offset + self.r_max,
+            0,
+            2 * self.r_max,
+        )
+        residue_offset = torch.where(
+            same_chain,
+            residue_offset,
+            torch.zeros_like(residue_offset) + 2 * self.r_max + 1,
+        )
+        residue_features = one_hot(residue_offset, 2 * self.r_max + 2)
+
+        token_offset = torch.clip(
+            _pairwise_difference(feats["token_index"]) + self.r_max,
+            0,
+            2 * self.r_max,
+        )
+        token_offset = torch.where(
+            same_chain & same_residue,
+            token_offset,
+            torch.zeros_like(token_offset) + 2 * self.r_max + 1,
+        )
+        token_features = one_hot(token_offset, 2 * self.r_max + 2)
+
+        symmetry_offset = torch.clip(
+            _pairwise_difference(feats["sym_id"]) + self.s_max,
+            0,
+            2 * self.s_max,
+        )
+        invalid_symmetry = ~same_entity if self.fix_sym_check else same_chain
+        symmetry_offset = torch.where(
+            invalid_symmetry,
+            torch.zeros_like(symmetry_offset) + 2 * self.s_max + 1,
+            symmetry_offset,
+        )
+        symmetry_features = one_hot(symmetry_offset, 2 * self.s_max + 2)
+
+        pair_features = torch.cat(
+            (
+                residue_features.float(),
+                token_features.float(),
+                same_entity.unsqueeze(-1).float(),
+                symmetry_features.float(),
+            ),
+            dim=-1,
+        )
+        return self.linear_layer(pair_features)
+
+
+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)``."""
+
+        single = torch.cat((s_trunk, s_inputs), dim=-1)
+        single = self.single_embed(self.norm_single(single))
+        normalized_fourier = None
+        if not self.disable_times:
+            fourier = self.fourier_embed(times)
+            normalized_fourier = self.norm_fourier(fourier)
+            time_condition = self.fourier_to_single(normalized_fourier)
+            single = time_condition[:, None, :] + single
+        for transition in self.transitions:
+            single = transition(single) + single
+        return single, normalized_fourier
+
+
+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))
+        for transition in self.transitions:
+            pair = transition(pair) + pair
+        return pair
+
+
+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)
+    relative_blocks = ((positions.unsqueeze(0) - positions.unsqueeze(1)) + key_blocks // 2).clamp(
+        min=0, max=key_blocks + 1
+    )
+    relative_blocks = relative_blocks.view(k, 2, 2 * k)[:, 0, :]
+    selectors = one_hot(relative_blocks, num_classes=key_blocks + 2)[..., 1:-1].transpose(1, 0)
+    return selectors.reshape(2 * k, key_blocks * k).float()
+
+
+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)
+    gathered = torch.einsum("b j i d, j k -> b k i d", half_windows, indexing_matrix)
+    return gathered.reshape(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_parts = [
+            feats["ref_pos"],
+            feats["ref_charge"].unsqueeze(-1),
+            feats["ref_element"],
+        ]
+        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,
+            )
+            feature_parts.append(torch.bmm(feats["atom_to_token"].float(), residue_features))
+        return torch.cat(feature_parts, dim=-1)
+
+    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()
+            atom_positions = feats["ref_pos"]
+            atom_space = feats["ref_space_uid"]
+            atom_features = self._assemble_atom_features(feats, b, n)
+            conditioned_single = self.embed_atom_features(atom_features)
+
+            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)
+            to_keys = partial(single_to_keys, indexing_matrix=indexing, w=w, h=h)
+
+            position_queries = atom_positions.view(b, k, w, 1, 3)
+            position_keys = to_keys(atom_positions).view(b, k, 1, h, 3)
+            displacement = position_keys - position_queries
+            inverse_squared_distance = 1 / (
+                1 + torch.sum(displacement * displacement, dim=-1, keepdim=True)
+            )
+
+            mask_queries = atom_mask.view(b, k, w, 1)
+            mask_keys = to_keys(atom_mask.unsqueeze(-1).float()).view(b, k, 1, h).bool()
+            space_queries = atom_space.view(b, k, w, 1)
+            space_keys = to_keys(atom_space.unsqueeze(-1).float()).view(b, k, 1, h).long()
+            valid_pair = (
+                (mask_queries & mask_keys & (space_queries == space_keys)).float().unsqueeze(-1)
+            )
+
+            pair = self.embed_atompair_ref_pos(displacement) * valid_pair
+            pair = pair + self.embed_atompair_ref_dist(inverse_squared_distance) * valid_pair
+            pair = pair + self.embed_atompair_mask(valid_pair) * valid_pair
+            query = conditioned_single
+
+            if self.structure_prediction:
+                if s_trunk is None or z is None:
+                    raise ValueError("structure prediction requires S and Z trunk tensors")
+                atom_to_token = feats["atom_to_token"].float()
+                single_update = self.s_to_c_trans(s_trunk.float())
+                single_update = torch.bmm(atom_to_token, single_update)
+                conditioned_single = conditioned_single + single_update.to(conditioned_single)
+
+                token_queries = atom_to_token.view(b, k, w, atom_to_token.shape[-1])
+                token_keys = to_keys(atom_to_token)
+                pair_update = self.z_to_p_trans(z.float())
+                pair_update = torch.einsum(
+                    "bijd,bwki,bwlj->bwkld",
+                    pair_update,
+                    token_queries,
+                    token_keys,
+                )
+                pair = pair + pair_update.to(pair)
+
+            pair = pair + self.c_to_p_trans_q(
+                conditioned_single.view(b, k, w, 1, conditioned_single.shape[-1])
+            )
+            pair = pair + self.c_to_p_trans_k(
+                to_keys(conditioned_single).view(
+                    b,
+                    k,
+                    1,
+                    h,
+                    conditioned_single.shape[-1],
+                )
+            )
+            pair = pair + self.p_mlp(pair)
+        return query, conditioned_single, pair, to_keys
+
+
+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]:
+        b, n, _ = feats["ref_pos"].shape
+        del b, n
+        atom_mask = feats["atom_pad_mask"].bool()
+        if self.structure_prediction:
+            if r is None:
+                raise ValueError("structure prediction requires atom coordinates R")
+            q = q.repeat_interleave(multiplicity, 0)
+            q = q + self.r_to_q_trans(r)
+        c = c.repeat_interleave(multiplicity, 0)
+        atom_mask = atom_mask.repeat_interleave(multiplicity, 0)
+        q = self.atom_encoder(
+            q=q,
+            mask=atom_mask,
+            c=c,
+            bias=atom_enc_bias,
+            multiplicity=multiplicity,
+            to_keys=to_keys,
+        )
+
+        with torch.autocast("cuda", enabled=False):
+            atom_update = self.atom_to_token_trans(q).float()
+            atom_to_token = feats["atom_to_token"].float()
+            atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)
+            atom_to_token_mean = atom_to_token / (atom_to_token.sum(dim=1, keepdim=True) + 1e-6)
+            token_update = torch.bmm(atom_to_token_mean.transpose(1, 2), atom_update)
+        return token_update.to(q), q, c, to_keys
+
+
+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):
+            atom_to_token = feats["atom_to_token"].float()
+            atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)
+            token_update = self.a_to_q_trans(a.float())
+            atom_update = torch.bmm(atom_to_token, token_update)
+        q = q + atom_update.to(q)
+        atom_mask = feats["atom_pad_mask"].repeat_interleave(multiplicity, 0)
+        q = self.atom_decoder(
+            q=q,
+            mask=atom_mask,
+            c=c,
+            bias=atom_dec_bias,
+            multiplicity=multiplicity,
+            to_keys=to_keys,
+        )
+        return self.atom_feat_to_atom_pos_update(q)

fastplms/models/boltz/vb_modules_transformersv2.py

@@ -0,0 +1,254 @@
+"""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:
+        normalized = self.a_norm(a)
+        conditioning = self.s_norm(s)
+        output = self.s_scale(conditioning).sigmoid() * normalized + self.s_bias(conditioning)
+        return cast(Tensor, output)
+
+
+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)
+        nn.init.constant_(projection.bias, -2.0)
+        self.output_projection = nn.Sequential(projection, nn.Sigmoid())
+
+    def forward(self, a: Tensor, s: Tensor) -> Tensor:
+        normalized = self.adaln(a, s)
+        hidden = self.swish_gate(normalized) * self.a_to_b(normalized)
+        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:
+        if not self.pair_bias_attn:
+            return None
+        if bias is None:
+            raise ValueError("pair bias is required when pair_bias_attn=True")
+        batch_size, query_length, key_length, width = bias.shape
+        depth = len(self.layers)
+        if depth == 0 or width % depth:
+            raise ValueError("pair-bias width must be divisible by transformer depth")
+        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)``."""
+
+        layer_biases = self._split_pair_bias(bias)
+        output = a
+        for index, layer in enumerate(self.layers):
+            layer_bias = None if layer_biases is None else layer_biases[..., index, :]
+            if self.activation_checkpointing:
+                output = checkpoint(
+                    layer,
+                    output,
+                    s,
+                    layer_bias,
+                    mask,
+                    to_keys,
+                    multiplicity,
+                    use_reentrant=False,
+                )
+            else:
+                output = layer(
+                    output,
+                    s,
+                    layer_bias,
+                    mask,
+                    to_keys,
+                    multiplicity,
+                )
+        return output
+
+
+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)
+        nn.init.constant_(self.output_projection_linear.bias, -2.0)
+        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."""
+
+        if bias is None or mask is None:
+            raise ValueError("diffusion attention requires pair bias and a key mask")
+        normalized = self.adaln(a, s)
+        key_states = normalized
+        key_mask = mask
+        if to_keys is not None:
+            key_states = to_keys(normalized)
+            key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1)
+
+        attended = self.pair_bias_attn(
+            s=normalized,
+            z=bias,
+            mask=key_mask,
+            multiplicity=multiplicity,
+            k_in=key_states,
+        )
+        output = a + self.output_projection(s) * attended
+        output = output + self.transition(output, s)
+        return cast(Tensor, self.post_lnorm(output))
+
+
+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
+        key_window = self.attn_window_keys
+        batch_size, atom_count, width = q.shape
+        window_count = atom_count // query_window
+        query_states = q.view(batch_size * window_count, query_window, -1)
+        conditioning = c.view(batch_size * window_count, query_window, -1)
+        query_mask = mask.view(batch_size * window_count, query_window)
+        pair_bias = bias.repeat_interleave(multiplicity, dim=0)
+        pair_bias = pair_bias.view(
+            pair_bias.shape[0] * window_count,
+            query_window,
+            key_window,
+            -1,
+        )
+
+        def windowed_keys(states: Tensor) -> Tensor:
+            merged = states.view(batch_size, window_count * query_window, -1)
+            return to_keys(merged).view(batch_size * window_count, key_window, -1)
+
+        output = self.diffusion_transformer(
+            a=query_states,
+            s=conditioning,
+            bias=pair_bias,
+            mask=query_mask.float(),
+            multiplicity=1,
+            to_keys=windowed_keys,
+        )
+        return cast(Tensor, output.view(batch_size, window_count * query_window, width))

fastplms/models/boltz/vb_modules_trunkv2.py

@@ -0,0 +1,595 @@
+"""Input, template, MSA, and prediction-head modules for the Boltz2 trunk."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import torch
+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))
+        self.encoding_unselected = nn.Parameter(torch.zeros(token_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")
+
+        categories = feats["contact_conditioning"]
+        threshold = feats["contact_threshold"]
+        normalized = (threshold - self.cutoff_min) / (self.cutoff_max - self.cutoff_min)
+        fourier = self.fourier_embedding(normalized.flatten()).reshape((*normalized.shape, -1))
+        selected_features = torch.cat(
+            [categories[..., 2:], normalized.unsqueeze(-1), fourier],
+            dim=-1,
+        )
+        selected = self.encoder(selected_features)
+        special = categories[..., :2]
+        return cast(
+            Tensor,
+            selected * (1 - special.sum(dim=-1, keepdim=True))
+            + self.encoding_unspecified * special[..., 0:1]
+            + self.encoding_unselected * special[..., 1:2],
+        )
+
+
+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)
+        if add_modified_flag:
+            self.modified_conditioning_init = nn.Embedding(2, token_s)
+            self.modified_conditioning_init.weight.data.fill_(0)
+        if add_cyclic_flag:
+            self.cyclic_conditioning_init = nn.Linear(1, token_s, bias=False)
+            self.cyclic_conditioning_init.weight.data.fill_(0)
+        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)
+
+    def forward(self, feats: dict[str, Tensor], affinity: bool = False) -> Tensor:
+        """Return embedded sequence tensor S with shape ``(b, l, d_s)``."""
+
+        residue_type = feats["res_type"].float()
+        suffix = "_affinity" if affinity else ""
+        profile = feats[f"profile{suffix}"]
+        deletion_mean = feats[f"deletion_mean{suffix}"].unsqueeze(-1)
+        atom_queries, atom_conditioning, atom_pairs, to_keys = self.atom_encoder(feats)
+        atom_bias = self.atom_enc_proj_z(atom_pairs)
+        atom_output, _, _, _ = self.atom_attention_encoder(
+            feats=feats,
+            q=atom_queries,
+            c=atom_conditioning,
+            atom_enc_bias=atom_bias,
+            to_keys=to_keys,
+        )
+        output = (
+            atom_output
+            + self.res_type_encoding(residue_type)
+            + self.msa_profile_encoding(torch.cat([profile, deletion_mean], dim=-1))
+        )
+        if self.add_method_conditioning:
+            output = output + self.method_conditioning_init(feats["method_feature"])
+        if self.add_modified_flag:
+            output = output + self.modified_conditioning_init(feats["modified"])
+        if self.add_cyclic_flag:
+            cyclic = feats["cyclic_period"].clamp(max=1.0).unsqueeze(-1)
+            output = output + self.cyclic_conditioning_init(cyclic)
+        if self.add_mol_type_feat:
+            output = output + self.mol_type_conditioning_init(feats["mol_type"])
+        return cast(Tensor, output)
+
+
+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 = feats["template_restype"]
+        cb_mask = feats["template_mask_cb"]
+        frame_mask = feats["template_mask_frame"]
+        cb_pair_mask = (cb_mask[..., :, None] * cb_mask[..., None, :]).unsqueeze(-1)
+        frame_pair_mask = (frame_mask[..., :, None] * frame_mask[..., None, :]).unsqueeze(-1)
+        with torch.autocast(device_type="cuda", enabled=False):
+            cb_distances = torch.cdist(feats["template_cb"], feats["template_cb"])
+            boundaries = torch.linspace(
+                self.min_dist,
+                self.max_dist,
+                self.num_bins - 1,
+            ).to(cb_distances.device)
+            bins = (cb_distances[..., None] > boundaries).sum(dim=-1).long()
+            distogram = F.one_hot(bins, num_classes=self.num_bins)
+
+            rotations = feats["template_frame_rot"].unsqueeze(2).transpose(-1, -2)
+            translations = feats["template_frame_t"].unsqueeze(2).unsqueeze(-1)
+            ca_coordinates = feats["template_ca"].unsqueeze(3).unsqueeze(-1)
+            vectors = torch.matmul(rotations, ca_coordinates - translations)
+            norms = torch.norm(vectors, dim=-1, keepdim=True)
+            unit_vectors = torch.where(
+                norms > 0,
+                vectors / norms,
+                torch.zeros_like(vectors),
+            ).squeeze(-1)
+            pair_features = torch.cat(
+                [distogram, cb_pair_mask, unit_vectors, frame_pair_mask],
+                dim=-1,
+            )
+            pair_features = pair_features * template_pair_mask.unsqueeze(-1)
+            residue_i = residue_type[:, :, :, None].expand(
+                -1,
+                -1,
+                -1,
+                residue_type.size(2),
+                -1,
+            )
+            residue_j = residue_type[:, :, None, :].expand(
+                -1,
+                -1,
+                residue_type.size(2),
+                -1,
+                -1,
+            )
+            return cast(
+                Tensor,
+                self.a_proj(torch.cat([pair_features, residue_i, residue_j], dim=-1)),
+            )
+
+    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."""
+
+        residue_type = feats["template_restype"]
+        batch_size, template_count = residue_type.shape[:2]
+        template_present = feats["template_mask"].any(dim=2).float()
+        present_count = template_present.sum(dim=1).clamp(min=1)
+        features = self._template_features(
+            feats,
+            self._template_pair_mask(feats, template_count),
+        )
+        expanded_mask = pair_mask[:, None].expand(-1, template_count, -1, -1)
+        expanded_mask = expanded_mask.reshape(
+            batch_size * template_count,
+            *expanded_mask.shape[2:],
+        )
+        template_states = self.z_proj(self.z_norm(z[:, None])) + features
+        template_states = template_states.view(
+            batch_size * template_count,
+            *template_states.shape[2:],
+        )
+        template_states = template_states + self.pairformer(
+            template_states,
+            expanded_mask,
+            use_kernels=use_kernels,
+        )
+        template_states = self.v_norm(template_states).view(
+            batch_size,
+            template_count,
+            *template_states.shape[1:],
+        )
+        weights = template_present[:, :, None, None, None]
+        aggregate = (template_states * weights).sum(dim=1)
+        aggregate = aggregate / present_count[:, None, None, None].to(template_states)
+        return cast(Tensor, self.u_proj(self.relu(aggregate)))
+
+
+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"]
+        same_chain = (asym_id[:, :, None] == asym_id[:, None, :]).float()
+        return same_chain[:, None].expand(-1, count, -1, -1)
+
+
+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"]
+        return (visibility[..., :, None] == visibility[..., None, :]).float()
+
+
+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]:
+        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)
+        msa = feats["msa"]
+        if msa.dtype in (torch.long, torch.int32, torch.int64):
+            msa = F.one_hot(msa, num_classes=const.num_tokens).float()
+        msa_mask = feats["msa_mask"]
+        components = [
+            msa,
+            feats["has_deletion"].unsqueeze(-1),
+            feats["deletion_value"].unsqueeze(-1),
+        ]
+        if self.use_paired_feature:
+            components.append(feats["msa_paired"].unsqueeze(-1))
+        msa_input = torch.cat(components, dim=-1)
+        if self.subsample_msa:
+            indices = torch.randperm(msa.shape[1])[: self.num_subsampled_msa]
+            msa_input = msa_input[:, indices]
+            msa_mask = msa_mask[:, indices]
+
+        msa_states = self.msa_proj(msa_input) + self.s_proj(emb).unsqueeze(1)
+        token_mask = feats["token_pad_mask"].float()
+        pair_mask = token_mask[:, :, None] * token_mask[:, None, :]
+        pair_states = z
+        for layer in self.layers:
+            arguments = (
+                pair_states,
+                msa_states,
+                pair_mask,
+                msa_mask,
+                *chunking,
+                use_kernels,
+            )
+            if self.activation_checkpointing and self.training:
+                pair_states, msa_states = checkpoint(
+                    layer,
+                    *arguments,
+                )
+            else:
+                pair_states, msa_states = layer(*arguments)
+        return pair_states
+
+
+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
+        dropout = get_dropout_mask(self.msa_dropout, m, self.training)
+        msa_states = m + dropout * self.pair_weighted_averaging(
+            m,
+            z,
+            token_mask,
+            chunk_heads_pwa,
+        )
+        msa_states = msa_states + self.msa_transition(
+            msa_states,
+            chunk_size_transition_msa,
+        )
+        pair_states = z + self.outer_product_mean(
+            msa_states,
+            msa_mask,
+            chunk_size_outer_product,
+        )
+        pair_states = self.pairformer_layer(
+            pair_states,
+            token_mask,
+            chunk_size_tri_attn,
+            use_kernels=use_kernels,
+        )
+        return pair_states, msa_states
+
+
+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:
+        return cast(Tensor, self.bfactor(s))
+
+
+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:
+        symmetric = z + z.transpose(1, 2)
+        logits = self.distogram(symmetric)
+        return cast(
+            Tensor,
+            logits.reshape(
+                symmetric.shape[0],
+                symmetric.shape[1],
+                symmetric.shape[2],
+                self.num_distograms,
+                self.num_bins,
+            ),
+        )

fastplms/models/boltz/vb_modules_utils.py

@@ -0,0 +1,257 @@
+"""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
+
+from collections.abc import Iterable, Mapping, Sequence
+from functools import partial
+from typing import Any
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+LinearNoBias = partial(nn.Linear, bias=False)
+
+
+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."""
+
+    return torch.log(tensor.clamp(min=eps))
+
+
+class SwiGLU(nn.Module):
+    """Split X in half and apply a SiLU-gated linear unit."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        values, gates = x.chunk(2, dim=-1)
+        return F.silu(gates) * values
+
+
+def _masked_center(coordinates: Tensor, mask: Tensor) -> Tensor:
+    weights = mask[:, :, None]
+    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)``."""
+
+    return atom_coords - _masked_center(atom_coords, atom_mask)
+
+
+def _copysign(magnitudes: Tensor, signs: Tensor) -> Tensor:
+    """Apply the elementwise sign of S to magnitude tensor M."""
+
+    return torch.where((magnitudes < 0) != (signs < 0), -magnitudes, magnitudes)
+
+
+def quaternion_to_matrix(quaternions: Tensor) -> Tensor:
+    """Convert real-first quaternion tensor Q from ``(..., 4)`` to ``(..., 3, 3)``."""
+
+    real, i_axis, j_axis, k_axis = torch.unbind(quaternions, dim=-1)
+    scale = 2.0 / (quaternions * quaternions).sum(dim=-1)
+    entries = (
+        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),
+    )
+    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)
+    squared_norm = (samples * samples).sum(dim=1)
+    signed_norm = _copysign(torch.sqrt(squared_norm), samples[:, 0])
+    return samples / signed_norm[:, None]
+
+
+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)``."""
+
+    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)
+    translations = (
+        torch.randn(
+            (multiplicity, 1, 3),
+            dtype=dtype,
+            device=device,
+        )
+        * s_trans
+    )
+    return rotations, translations
+
+
+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."""
+
+    rotations = random_rotations(len(coords), coords.dtype, coords.device)
+    rotated = torch.einsum("bmd,bds->bms", coords, rotations)
+    if not return_second_coords:
+        return rotated
+    rotated_second = (
+        None if second_coords is None else torch.einsum("bmd,bds->bms", second_coords, rotations)
+    )
+    return rotated, rotated_second
+
+
+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."""
+
+    primary = atom_coords
+    secondary = second_coords
+    if centering:
+        centroid = _masked_center(primary, atom_mask)
+        primary = primary - centroid
+        if secondary is not None:
+            secondary = secondary - centroid
+
+    if augmentation:
+        primary, secondary = randomly_rotate(
+            primary,
+            return_second_coords=True,
+            second_coords=secondary,
+        )
+        translation = torch.randn_like(primary[:, 0:1, :]) * s_trans
+        primary = primary + translation
+        if secondary is not None:
+            secondary = secondary + translation
+
+    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:
+        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 = [
+            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]:
+        return [parameter for parameter in parameters if parameter.requires_grad]
+
+    def update(self, parameters: Iterable[nn.Parameter]) -> None:
+        trainable = self._trainable(parameters)
+        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))
+
+    def compatible(self, parameters: Sequence[Tensor]) -> bool:
+        """Return whether parameter count and tensor shapes match the EMA state."""
+
+        return len(self.shadow_params) == len(parameters) and all(
+            shadow.shape == parameter.shape
+            for shadow, parameter in zip(self.shadow_params, parameters, strict=True)
+        )
+
+    def copy_to(self, parameters: Iterable[nn.Parameter]) -> None:
+        trainable = self._trainable(parameters)
+        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)
+
+    def store(self, parameters: Iterable[nn.Parameter]) -> None:
+        self.collected_params = [parameter.clone() for parameter in parameters]
+
+    def restore(self, parameters: Iterable[nn.Parameter]) -> None:
+        current = list(parameters)
+        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)
+
+    def state_dict(self) -> dict[str, Any]:
+        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:
+        self.decay = float(state_dict["decay"])
+        updates = state_dict["num_updates"]
+        self.num_updates = None if updates is None else int(updates)
+        self.shadow_params = [tensor.to(device) for tensor in state_dict["shadow_params"]]
+
+    def to(self, device: torch.device | str) -> None:
+        self.shadow_params = [tensor.to(device) for tensor in self.shadow_params]

fastplms/models/boltz/vb_potentials_potentials.py

@@ -0,0 +1,845 @@
+"""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
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from . import vb_const as const
+from .vb_loss_diffusionv2 import weighted_rigid_align
+from .vb_potentials_schedules import (
+    ExponentialInterpolation,
+    ParameterSchedule,
+    PiecewiseStepFunction,
+)
+
+Parameter = ParameterSchedule | float | int | bool
+ParameterMap = dict[str, Parameter]
+
+
+@dataclass(frozen=True, slots=True)
+class _PreparedPotential:
+    coordinates: torch.Tensor
+    index: torch.Tensor
+    function_args: tuple[Any, ...]
+    com_index: torch.Tensor | None
+    atom_pad_mask: torch.Tensor | None
+    ref_coords: torch.Tensor | None
+    ref_mask: torch.Tensor | None
+    ref_token_index: torch.Tensor | None
+    negation_mask: torch.Tensor | None
+    union_index: torch.Tensor | None
+
+
+def _reduce_atoms_to_centers(
+    coordinates: torch.Tensor,
+    com_index: torch.Tensor,
+    atom_pad_mask: torch.Tensor,
+) -> torch.Tensor:
+    unpadded_index = com_index[atom_pad_mask]
+    unpadded_coordinates = coordinates[..., atom_pad_mask, :]
+    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",
+    )
+
+
+def _union_weights(
+    energy: torch.Tensor,
+    union_index: torch.Tensor,
+    union_lambda: float,
+) -> torch.Tensor:
+    unnormalized = torch.exp(-union_lambda * energy)
+    partition = torch.zeros(
+        (*energy.shape[:-1], union_index.max() + 1),
+        device=union_index.device,
+    ).scatter_reduce(
+        -1,
+        union_index.expand_as(unnormalized),
+        unnormalized,
+        "sum",
+    )
+    weights = unnormalized / partition[..., union_index]
+    weights[partition[..., union_index] == 0] = 0
+    return weights
+
+
+def _scatter_constraint_gradients(
+    coefficients: torch.Tensor,
+    variable_gradient: torch.Tensor,
+    index: torch.Tensor,
+    coordinates: torch.Tensor,
+) -> torch.Tensor:
+    product = coefficients.tile(variable_gradient.shape[-3]).unsqueeze(
+        -1
+    ) * variable_gradient.flatten(start_dim=-3, end_dim=-2)
+    if product.dim() > 3:
+        product = product.sum(dim=list(range(1, product.dim() - 2)))
+    scatter_index = (
+        index.flatten(start_dim=0, end_dim=1).unsqueeze(-1).expand((*coordinates.shape[:-2], -1, 3))
+    )
+    return torch.zeros_like(coordinates).scatter_reduce(
+        -2,
+        scatter_index,
+        product,
+        "sum",
+    )
+
+
+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:
+        if computed_args is None:
+            computed_args = self.compute_args(feats, parameters)
+        index, args, com_args, ref_args, operator_args = computed_args
+        com_index = atom_pad_mask = None
+        if com_args is not None:
+            com_index, atom_pad_mask = com_args
+            coordinates = _reduce_atoms_to_centers(
+                coordinates,
+                com_index,
+                atom_pad_mask,
+            )
+
+        ref_coords = ref_mask = ref_token_index = None
+        if ref_args is not None:
+            ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args
+            coordinates = coordinates[..., ref_atom_index, :]
+
+        negation_mask = union_index = None
+        if operator_args is not None:
+            negation_mask, union_index = operator_args
+        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."""
+
+        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,
+        )
+        energy = self.compute_function(
+            value,
+            *prepared.function_args,
+            negation_mask=prepared.negation_mask,
+            compute_derivative=False,
+        )
+        if prepared.union_index is not None:
+            weights = _union_weights(
+                energy,
+                prepared.union_index,
+                parameters["union_lambda"],
+            )
+            return (energy * weights).sum(dim=-1)
+        return energy.sum(dim=tuple(range(1, energy.dim())))
+
+    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."""
+
+        computed_args = self.compute_args(feats, parameters)
+        if computed_args[0].shape[1] == 0:
+            return torch.zeros_like(coords)
+        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,
+        )
+        energy, energy_derivative = self.compute_function(
+            value,
+            *prepared.function_args,
+            negation_mask=prepared.negation_mask,
+            compute_derivative=True,
+        )
+        if prepared.union_index is not None:
+            weights = _union_weights(
+                energy,
+                prepared.union_index,
+                parameters["union_lambda"],
+            )
+            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",
+            )
+            coefficients = (
+                energy_derivative
+                * weights
+                * (
+                    1
+                    + parameters["union_lambda"]
+                    * (energy - union_energy[..., prepared.union_index])
+                )
+            )
+        else:
+            coefficients = energy_derivative
+
+        atom_gradient = _scatter_constraint_gradients(
+            coefficients,
+            variable_gradient,
+            prepared.index,
+            prepared.coordinates,
+        )
+        if prepared.com_index is not None:
+            atom_gradient = atom_gradient[..., prepared.com_index, :]
+        elif prepared.ref_token_index is not None:
+            atom_gradient = atom_gradient[..., prepared.ref_token_index, :]
+        return atom_gradient
+
+    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]:
+        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]:
+        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]:
+        lower = (
+            (torch.full_like(value, float("-inf")) if lower_bounds is None else lower_bounds)
+            .expand_as(value)
+            .clone()
+        )
+        upper = (
+            (torch.full_like(value, float("inf")) if upper_bounds is None else upper_bounds)
+            .expand_as(value)
+            .clone()
+        )
+
+        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)
+            except RuntimeError as error:
+                raise ValueError(
+                    "negation_mask must be broadcastable to value shape "
+                    f"{tuple(value.shape)}, got {tuple(negation_mask.shape)}."
+                ) from error
+            unbounded_below = torch.isneginf(lower)
+            unbounded_above = torch.isposinf(upper)
+            valid_negation = unbounded_below | unbounded_above | expanded_negation_mask
+            if not bool(torch.all(valid_negation).item()):
+                raise ValueError(
+                    "negation_mask may be false only where at least one bound is infinite."
+                )
+            select_upper = ~unbounded_above & ~expanded_negation_mask
+            lower[select_upper] = upper[select_upper]
+            upper[select_upper] = float("inf")
+            select_lower = ~unbounded_below & ~expanded_negation_mask
+            upper[select_lower] = lower[select_lower]
+            lower[select_lower] = float("-inf")
+
+        below = value < lower
+        above = value > upper
+        energy = torch.zeros_like(value)
+        energy[below] = (k * (lower - value))[below]
+        energy[above] = (k * (value - upper))[above]
+        if not compute_derivative:
+            return energy
+        derivative = torch.zeros_like(value)
+        derivative[below] = -k.expand_as(below)[below]
+        derivative[above] = k.expand_as(above)[above]
+        return energy, derivative
+
+
+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]:
+        aligned_reference = weighted_rigid_align(
+            ref_coords.float(),
+            coords[:, index].float(),
+            ref_mask,
+            ref_mask,
+        )
+        displacement = coords[:, index] - aligned_reference
+        distance = torch.linalg.norm(displacement, dim=-1)
+        if not compute_gradient:
+            return distance
+        unit_displacement = displacement / distance.unsqueeze(-1)
+        gradient = (unit_displacement * ref_mask.unsqueeze(-1)).unsqueeze(1)
+        return distance, gradient
+
+
+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]:
+        del ref_coords, ref_mask
+        displacement = coords.index_select(-2, index[0]) - coords.index_select(
+            -2,
+            index[1],
+        )
+        distance = torch.linalg.norm(displacement, dim=-1)
+        if not compute_gradient:
+            return distance
+        unit_displacement = displacement / distance.unsqueeze(-1)
+        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]:
+        del ref_coords, ref_mask
+        r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1])
+        r_kj = coords.index_select(-2, index[2]) - coords.index_select(-2, index[1])
+        r_kl = coords.index_select(-2, index[2]) - coords.index_select(-2, index[3])
+        n_ijk = torch.cross(r_ij, r_kj, dim=-1)
+        n_jkl = torch.cross(r_kj, r_kl, dim=-1)
+        r_kj_norm = torch.linalg.norm(r_kj, dim=-1)
+        n_ijk_norm = torch.linalg.norm(n_ijk, dim=-1)
+        n_jkl_norm = torch.linalg.norm(n_jkl, dim=-1)
+
+        orientation = torch.sign(
+            r_kj.unsqueeze(-2) @ torch.cross(n_ijk, n_jkl, dim=-1).unsqueeze(-1)
+        ).squeeze(-1, -2)
+        cosine = (n_ijk.unsqueeze(-2) @ n_jkl.unsqueeze(-1)).squeeze(-1, -2) / (
+            n_ijk_norm * n_jkl_norm
+        )
+        angle = orientation * torch.arccos(torch.clamp(cosine, -1 + 1e-8, 1 - 1e-8))
+        if not compute_gradient:
+            return angle
+
+        projection_i = (
+            (r_ij.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
+        ).unsqueeze(-1)
+        projection_l = (
+            (r_kl.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2)
+        ).unsqueeze(-1)
+        grad_i = n_ijk * (r_kj_norm / n_ijk_norm**2).unsqueeze(-1)
+        grad_l = -n_jkl * (r_kj_norm / n_jkl_norm**2).unsqueeze(-1)
+        grad_j = (projection_i - 1) * grad_i - projection_l * grad_l
+        grad_k = (projection_l - 1) * grad_l - projection_i * grad_i
+        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]:
+        del ref_coords, ref_mask
+        if not compute_gradient:
+            return torch.abs(super().compute_variable(coords, index))
+        angle, gradient = super().compute_variable(
+            coords,
+            index,
+            compute_gradient=True,
+        )
+        gradient[(angle < 0)[..., None, :, None].expand_as(gradient)] *= -1
+        return torch.abs(angle), gradient
+
+
+def _element_radii(feats: dict[str, torch.Tensor]) -> torch.Tensor:
+    element_radii = torch.zeros(
+        const.num_elements,
+        dtype=torch.float32,
+        device=feats["ref_element"].device,
+    )
+    element_radii[1:119] = torch.tensor(
+        const.vdw_radii,
+        dtype=torch.float32,
+        device=element_radii.device,
+    )
+    return (feats["ref_element"].float() @ element_radii.unsqueeze(-1)).squeeze(-1)[0]
+
+
+def _atom_chain_ids(feats: dict[str, torch.Tensor]) -> torch.Tensor:
+    return (
+        torch.bmm(
+            feats["atom_to_token"].float(),
+            feats["asym_id"].unsqueeze(-1).float(),
+        )
+        .squeeze(-1)
+        .long()
+    )[0]
+
+
+class PoseBustersPotential(FlatBottomPotential, DistancePotential):
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
+        pair_index = feats["rdkit_bounds_index"][0]
+        lower = feats["rdkit_lower_bounds"][0].clone()
+        upper = feats["rdkit_upper_bounds"][0].clone()
+        bond = feats["rdkit_bounds_bond_mask"][0]
+        angle = feats["rdkit_bounds_angle_mask"][0]
+        lower[bond * ~angle] *= 1.0 - parameters["bond_buffer"]
+        upper[bond * ~angle] *= 1.0 + parameters["bond_buffer"]
+        lower[~bond * angle] *= 1.0 - parameters["angle_buffer"]
+        upper[~bond * angle] *= 1.0 + parameters["angle_buffer"]
+        shared_buffer = min(parameters["bond_buffer"], parameters["angle_buffer"])
+        lower[bond * angle] *= 1.0 - shared_buffer
+        upper[bond * angle] *= 1.0 + shared_buffer
+        lower[~bond * ~angle] *= 1.0 - parameters["clash_buffer"]
+        upper[~bond * ~angle] = float("inf")
+
+        atom_radii = _element_radii(feats)
+        bond_cutoff = 0.35 + atom_radii[pair_index].mean(dim=0)
+        lower[~bond] = torch.max(lower[~bond], bond_cutoff[~bond])
+        upper[bond] = torch.min(upper[bond], bond_cutoff[bond])
+        return pair_index, (torch.ones_like(lower), lower, upper), None, None, None
+
+
+class ConnectionsPotential(FlatBottomPotential, DistancePotential):
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
+        pair_index = feats["connected_atom_index"][0]
+        upper = torch.full(
+            (pair_index.shape[1],),
+            parameters["buffer"],
+            device=pair_index.device,
+        )
+        return pair_index, (torch.ones_like(upper), None, upper), None, None, None
+
+
+class VDWOverlapPotential(FlatBottomPotential, DistancePotential):
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
+        atom_chain_id = _atom_chain_ids(feats)
+        atom_pad_mask = feats["atom_pad_mask"][0].bool()
+        chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask])
+        nonion_atom = (chain_sizes > 1)[atom_chain_id]
+        atom_radii = _element_radii(feats)
+        pair_index = torch.triu_indices(
+            atom_chain_id.shape[0],
+            atom_chain_id.shape[0],
+            1,
+            device=atom_chain_id.device,
+        )
+        pair_pad_mask = atom_pad_mask[pair_index].all(dim=0)
+        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]
+        connected_matrix = torch.eye(
+            num_chains,
+            device=atom_chain_id.device,
+            dtype=torch.bool,
+        )
+        connected_matrix[connected[0], connected[1]] = True
+        connected_matrix[connected[1], connected[0]] = True
+        connected_pair = connected_matrix[
+            atom_chain_id[pair_index[0]],
+            atom_chain_id[pair_index[1]],
+        ]
+        pair_index = pair_index[:, pair_pad_mask * pair_ion_mask * ~connected_pair]
+        lower = atom_radii[pair_index].sum(dim=0) * (1.0 - parameters["buffer"])
+        return pair_index, (torch.ones_like(lower), lower, None), None, None, None
+
+
+class SymmetricChainCOMPotential(FlatBottomPotential, DistancePotential):
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
+        atom_chain_id = _atom_chain_ids(feats)
+        atom_pad_mask = feats["atom_pad_mask"][0].bool()
+        nonion_chain = torch.bincount(atom_chain_id[atom_pad_mask]) > 1
+        pair_index = feats["symmetric_chain_index"][0]
+        pair_index = pair_index[
+            :,
+            nonion_chain[pair_index[0]] * nonion_chain[pair_index[1]],
+        ]
+        lower = torch.full(
+            (pair_index.shape[1],),
+            parameters["buffer"],
+            dtype=torch.float32,
+            device=pair_index.device,
+        )
+        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]:
+    lower = torch.zeros(orientations.shape, device=orientations.device)
+    upper = torch.zeros(orientations.shape, device=orientations.device)
+    lower[orientations] = positive_lower
+    upper[orientations] = float("inf")
+    lower[~orientations] = float("-inf")
+    upper[~orientations] = negative_upper
+    return lower, upper
+
+
+class StereoBondPotential(FlatBottomPotential, AbsDihedralPotential):
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
+        index = feats["stereo_bond_index"][0]
+        orientation = feats["stereo_bond_orientations"][0].bool()
+        lower, upper = _oriented_bounds(
+            orientation,
+            torch.pi - parameters["buffer"],
+            parameters["buffer"],
+        )
+        return index, (torch.ones_like(lower), lower, upper), None, None, None
+
+
+class ChiralAtomPotential(FlatBottomPotential, DihedralPotential):
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
+        index = feats["chiral_atom_index"][0]
+        orientation = feats["chiral_atom_orientations"][0].bool()
+        lower, upper = _oriented_bounds(
+            orientation,
+            parameters["buffer"],
+            -parameters["buffer"],
+        )
+        return index, (torch.ones_like(lower), lower, upper), None, None, None
+
+
+class PlanarBondPotential(FlatBottomPotential, AbsDihedralPotential):
+    def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]):
+        bond_index = feats["planar_bond_index"][0].T
+        improper_pattern = torch.tensor(
+            [[1, 2, 3, 0], [4, 5, 0, 3]],
+            device=bond_index.device,
+        ).T
+        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,
+        )
+        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]):
+        del parameters
+        if "template_mask_cb" not in feats or "template_force" not in feats:
+            return torch.empty((1, 0)), None, None, None, None
+        template_mask = feats["template_mask_cb"][feats["template_force"]]
+        if template_mask.shape[0] == 0:
+            return torch.empty((1, 0)), None, None, None, None
+
+        ref_coords = feats["template_cb"][feats["template_force"]].clone()
+        ref_mask = feats["template_mask_cb"][feats["template_force"]].clone()
+        atom_indices = torch.arange(
+            feats["atom_pad_mask"].shape[1],
+            device=feats["atom_pad_mask"].device,
+            dtype=torch.float32,
+        )[None, :, None]
+        ref_atom_index = (
+            torch.bmm(
+                feats["token_to_rep_atom"].float(),
+                atom_indices,
+            )
+            .squeeze(-1)
+            .long()[0]
+        )
+        ref_token_index = (
+            torch.bmm(
+                feats["atom_to_token"].float(),
+                feats["token_index"].unsqueeze(-1).float(),
+            )
+            .squeeze(-1)
+            .long()[0]
+        )
+
+        index = torch.arange(
+            template_mask.shape[-1],
+            dtype=torch.long,
+            device=template_mask.device,
+        )[None]
+        upper = torch.full(
+            template_mask.shape,
+            float("inf"),
+            device=index.device,
+            dtype=torch.float32,
+        )
+        reference_indices = torch.argwhere(template_mask).T
+        upper[reference_indices.unbind()] = feats["template_force_threshold"][
+            feats["template_force"]
+        ][reference_indices[0]]
+        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]):
+        del parameters
+        index = feats["contact_pair_index"][0]
+        upper = feats["contact_thresholds"][0].clone()
+        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

@@ -0,0 +1,67 @@
+"""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

@@ -0,0 +1,145 @@
+"""Triangular row and column attention for Boltz2 pair states."""
+
+from __future__ import annotations
+
+from typing import cast
+
+import torch
+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."""
+
+        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,
+        }
+        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:
+        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
+        pair_states = x
+        if not self.starting:
+            pair_states = pair_states.transpose(-2, -3)
+            pair_mask = pair_mask.transpose(-1, -2)
+
+        normalized = self.layer_norm(pair_states)
+        expanded_mask = pair_mask[..., :, None, None, :]
+        mask_bias = self.inf * (expanded_mask - 1)
+        triangle_bias = permute_final_dims(self.linear(normalized), (2, 0, 1))
+        triangle_bias = triangle_bias.unsqueeze(-4)
+        output = self._run_attention(
+            normalized,
+            triangle_bias,
+            mask_bias,
+            expanded_mask,
+            chunk_size,
+            use_kernels,
+        )
+        return output if self.starting else output.transpose(-2, -3)
+
+
+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

@@ -0,0 +1,237 @@
+"""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
+from collections.abc import Callable, Sequence
+from importlib.util import find_spec
+
+import torch
+from torch import nn
+
+from . import vb_layers_initialize as initialize
+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:
+    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_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)
+                projected = nn.functional.linear(
+                    input.to(self.precision),
+                    self.weight.to(self.precision),
+                    bias,
+                )
+            return projected.to(input_dtype)
+        if input_dtype is torch.bfloat16:
+            with torch.autocast("cuda", enabled=False):
+                bias = None if self.bias is None else self.bias.to(input_dtype)
+                return nn.functional.linear(input, self.weight.to(input_dtype), bias)
+        return nn.functional.linear(input, self.weight, self.bias)
+
+
+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))
+        self.bias = nn.Parameter(torch.zeros(c_in))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        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."""
+
+    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."""
+
+    scores = torch.matmul(query, permute_final_dims(key, (1, 0)))
+    for bias in biases:
+        scores += bias
+    probabilities = softmax_no_cast(scores, dim=-1)
+    return torch.matmul(probabilities, value)
+
+
+@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."""
+
+    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:
+            return X.view(*X.shape[:-1], self.no_heads, -1).transpose(-2, -3)
+
+        q = split_heads(self.linear_q(q_x))
+        k = split_heads(self.linear_k(kv_x))
+        v = split_heads(self.linear_v(kv_x))
+        if apply_scale:
+            q /= math.sqrt(self.c_hidden)
+        return q, k, v
+
+    def _wrap_up(self, output: torch.Tensor, q_x: torch.Tensor) -> torch.Tensor:
+        """Apply the optional gate, merge heads, and project the update."""
+
+        if self.linear_g is not None:
+            gate = self.sigmoid(self.linear_g(q_x))
+            gate = gate.view(*gate.shape[:-1], self.no_heads, -1)
+            output = output * gate
+        return self.linear_o(flatten_final_dims(output, 2))
+
+    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)
+        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),
+            )
+        else:
+            output = _attention(q, k, v, (mask_bias, tri_bias))
+        return self._wrap_up(output.transpose(-2, -3), q_x)

fastplms/models/boltz/vb_tri_attn_utils.py

@@ -0,0 +1,324 @@
+"""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
+
+from collections.abc import Callable, Sequence
+from functools import partial
+from math import prod
+from typing import Any
+
+import torch
+
+
+def add(left: torch.Tensor, right: torch.Tensor, inplace: bool) -> torch.Tensor:
+    """Add ``right`` to ``left``, optionally reusing ``left`` storage."""
+
+    if inplace:
+        left += right
+        return left
+    return left + right
+
+
+def permute_final_dims(tensor: torch.Tensor, inds: Sequence[int]) -> torch.Tensor:
+    """Permute only the final ``len(inds)`` axes of ``tensor``."""
+
+    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)
+
+
+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."""
+
+    return tensor.reshape(*tensor.shape[:-no_dims], -1)
+
+
+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."""
+
+    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)]
+    feature_shape = tuple(tensor.shape[no_batch_dims:])
+    return torch.cat([piece.reshape(-1, *feature_shape) for piece in pieces])
+
+
+def _prepare_input(
+    tensor: torch.Tensor,
+    batch_shape: tuple[int, ...],
+    no_batch_dims: int,
+    *,
+    low_mem: bool,
+) -> torch.Tensor:
+    feature_shape = tuple(tensor.shape[no_batch_dims:])
+    if low_mem:
+        return tensor.expand(batch_shape + feature_shape)
+    if any(size != 1 for size in tensor.shape[:no_batch_dims]):
+        tensor = tensor.expand(batch_shape + feature_shape)
+    return tensor.reshape(-1, *feature_shape)
+
+
+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
+    else:
+        destination[start:stop] = source
+
+
+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)
+
+    output = None
+    if _out is not None:
+        output = tensor_tree_map(
+            lambda tensor: tensor.reshape(-1, *tensor.shape[no_batch_dims:]),
+            _out,
+        )
+
+    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)
+        output_chunk = layer(**input_chunk)
+
+        if output is None:
+            output = tensor_tree_map(
+                lambda tensor: tensor.new_zeros((flat_batch_size, *tensor.shape[1:])),
+                output_chunk,
+            )
+        _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,
+    )

fastplms/registry.py

@@ -0,0 +1,1486 @@
+"""Typed access to the FastPLMs model and provenance manifest.
+
+The registry is intentionally independent of Torch and Transformers. Tooling can
+therefore inspect supported checkpoints, licenses, and reference sources without
+initializing a model runtime or downloading any files.
+"""
+
+from __future__ import annotations
+
+import re
+import tomllib
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass
+from functools import lru_cache
+from importlib import resources
+from pathlib import Path, PurePosixPath, PureWindowsPath
+from types import MappingProxyType
+from typing import Any, Literal, cast
+from urllib.parse import urlparse
+
+_HEX_RE = re.compile(r"^[0-9a-f]+$")
+_IDENTIFIER_RE = re.compile(r"^[a-z0-9][a-z0-9_-]*$")
+_HUB_LICENSE_NAME_RE = re.compile(r"[^a-z0-9.]+")
+_WINDOWS_INVALID_PATH_CHARACTERS = frozenset('<>:"|?*')
+_WINDOWS_RESERVED_PATH_NAMES = frozenset(
+    {"AUX", "CON", "NUL", "PRN"}
+    | {f"COM{index}" for index in range(1, 10)}
+    | {f"LPT{index}" for index in range(1, 10)}
+)
+_REPOSITORY_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*$")
+_REFERENCE_CONTAINER_RE = re.compile(r"^reference-[a-z0-9]+(?:-[a-z0-9]+)*$")
+_REFERENCE_ADAPTER_RE = re.compile(
+    r"^tests\.parity\.support\.reference_adapters\.[a-z_][a-z0-9_]*$"
+)
+_DOCUMENTATION_FRAGMENT_RE = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$")
+_ALLOWED_ATTENTION = frozenset(
+    {"eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"}
+)
+_ALLOWED_DTYPES = frozenset({"float32", "bfloat16"})
+_ALLOWED_PRECISIONS = frozenset({"default", "auto", "fp32", "bf16", "fp8"})
+_ALLOWED_BF16_EXECUTIONS = frozenset({"static_parameters", "fp32_parameters_autocast"})
+HUB_LICENSE_IDENTIFIERS = frozenset({"mit", "apache-2.0", "cc-by-nc-sa-4.0", "other"})
+_ALLOWED_TOKENIZER_MODES = frozenset({"tokenizer", "sequence", "structure"})
+_ALLOWED_SIZE_CATEGORIES = frozenset({"small", "medium", "large", "xlarge", "structure"})
+RuntimeExtra = Literal["core", "structure"]
+TestTier = Literal["check", "compliance", "structure", "feature", "artifact", "benchmark"]
+VramTier = Literal["sequence", "large-sequence", "structure", "structure-6b"]
+GenerationContract = Literal["not_applicable", "required", "official_unavailable"]
+RuntimeAssetTrustKind = Literal["hash_pinned_pickle"]
+Bf16Execution = Literal["static_parameters", "fp32_parameters_autocast"]
+DtypeName = Literal["float32", "bfloat16"]
+_ALLOWED_EXTRAS = frozenset({"core", "structure"})
+_ALLOWED_TEST_TIERS = frozenset(
+    {"check", "compliance", "structure", "feature", "artifact", "benchmark"}
+)
+_ALLOWED_VRAM_TIERS = frozenset({"sequence", "large-sequence", "structure", "structure-6b"})
+_ALLOWED_GENERATION_CONTRACTS = frozenset({"not_applicable", "required", "official_unavailable"})
+_ALLOWED_RUNTIME_ASSET_TRUST_KINDS = frozenset({"hash_pinned_pickle"})
+_ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS = frozenset({"requires_cached_verified_file"})
+_ALLOWED_AUTO_CLASSES = frozenset(
+    {
+        "AutoConfig",
+        "AutoModel",
+        "AutoModelForMaskedLM",
+        "AutoModelForProteinFolding",
+        "AutoModelForSequenceClassification",
+        "AutoModelForSeq2SeqLM",
+        "AutoModelForTokenClassification",
+    }
+)
+_WEIGHT_SUFFIXES = (".bin", ".ckpt", ".pt", ".pth", ".safetensors")
+_ALLOWED_ORACLE_ASSET_ROLES = frozenset({"weights", "contact_regression"})
+_FAIR_ESM_ASSET_HOST = "dl.fbaipublicfiles.com"
+_ROOT_FIELDS = frozenset(
+    {
+        "schema_version",
+        "legal_files",
+        "attention_kernels",
+        "upstreams",
+        "families",
+        "models",
+        "runtime_assets",
+    }
+)
+_UPSTREAM_FIELDS = frozenset(
+    {
+        "id",
+        "path",
+        "url",
+        "revision",
+        "license",
+        "license_files",
+        "license_digests",
+        "distribution_files",
+    }
+)
+_FAMILY_FIELDS = frozenset(
+    {
+        "architecture",
+        "upstreams",
+        "tokenizer_mode",
+        "public_input",
+        "extra",
+        "reference_container",
+        "reference_adapter",
+        "attention",
+        "dtypes",
+        "bf16_execution",
+        "precisions",
+        "experimental_precisions",
+        "vram_tier",
+        "checkpoint_license",
+        "hub_license",
+        "hub_license_name",
+        "hub_license_link",
+        "state_transform",
+        "conversion_provenance",
+        "representative",
+        "documentation",
+        "test_tiers",
+        "runtime_paths",
+        "requires_complete_weight_publication",
+        "weights_publication_allowed",
+        "auto_map",
+        "tokenizer_class",
+        "backbone_model",
+    }
+)
+_MODEL_FIELDS = frozenset(
+    {
+        "id",
+        "family",
+        "size_category",
+        "generation_contract",
+        "fast_repo",
+        "fast_revision",
+        "fast_files",
+        "fast_unresolved_files",
+        "official_repo",
+        "official_revision",
+        "official_files",
+        "official_unresolved_files",
+        "oracle_assets",
+        "official_golden",
+        "artifact_source",
+        "canonical_state_sha256",
+        "tokenizer_source",
+        "auto_map",
+        "notes",
+        "msa_conditioning",
+    }
+)
+_RUNTIME_ASSET_FIELDS = frozenset(
+    {
+        "id",
+        "repository",
+        "revision",
+        "path",
+        "sha256",
+        "size",
+        "consumer_family",
+        "trust_kind",
+        "license",
+        "offline_behavior",
+    }
+)
+
+
+class RegistryError(ValueError):
+    """Raised when the model manifest is incomplete or internally inconsistent."""
+
+
+def _portable_relative_path(value: str, context: str) -> PurePosixPath:
+    """Return one normalized cross-platform relative path or fail closed."""
+
+    posix = PurePosixPath(value)
+    windows = PureWindowsPath(value)
+    unsafe_windows_part = any(
+        part.rstrip(" .") != part
+        or part.split(".", maxsplit=1)[0].upper() in _WINDOWS_RESERVED_PATH_NAMES
+        or any(
+            ord(character) < 32 or character in _WINDOWS_INVALID_PATH_CHARACTERS
+            for character in part
+        )
+        for part in posix.parts
+    )
+    if (
+        not value
+        or not posix.parts
+        or posix == PurePosixPath(".")
+        or posix.is_absolute()
+        or windows.is_absolute()
+        or windows.drive
+        or "\\" in value
+        or "." in posix.parts
+        or ".." in posix.parts
+        or value != posix.as_posix()
+        or any(
+            part.lower() in {".git", ".cache", "__pycache__"}
+            for part in posix.parts
+        )
+        or unsafe_windows_part
+    ):
+        raise RegistryError(f"{context} is not portable: {value!r}")
+    return posix
+
+
+@dataclass(frozen=True, slots=True)
+class FileDigest:
+    """Expected content identity for one pinned file."""
+
+    path: str
+    algorithm: str
+    digest: str
+
+    @classmethod
+    def parse(cls, value: str) -> FileDigest:
+        try:
+            path, encoded_digest = value.split("=", maxsplit=1)
+            algorithm, digest = encoded_digest.split(":", maxsplit=1)
+        except ValueError as error:
+            raise RegistryError("File digests must use '<path>=<algorithm>:<digest>'.") from error
+
+        _portable_relative_path(path, "Checkpoint file path")
+
+        expected_length = {"git-sha1": 40, "sha256": 64}.get(algorithm)
+        if expected_length is None:
+            raise RegistryError(f"Unsupported file digest algorithm: {algorithm!r}")
+        if len(digest) != expected_length or _HEX_RE.fullmatch(digest) is None:
+            raise RegistryError(f"Invalid {algorithm} digest for {path!r}: {digest!r}")
+        return cls(path=path, algorithm=algorithm, digest=digest)
+
+    @property
+    def encoded(self) -> str:
+        return f"{self.algorithm}:{self.digest}"
+
+
+@dataclass(frozen=True, slots=True)
+class CheckpointSource:
+    """One immutable Hugging Face repository snapshot."""
+
+    repo_id: str
+    revision: str
+    files: tuple[FileDigest, ...]
+    unresolved_files: tuple[str, ...] = ()
+
+    @property
+    def file_map(self) -> Mapping[str, FileDigest]:
+        return MappingProxyType({item.path: item for item in self.files})
+
+
+@dataclass(frozen=True, slots=True)
+class OracleAsset:
+    """Hash-pinned external file required by a native parity oracle."""
+
+    role: str
+    path: str
+    url: str
+    sha256: str
+    size: int
+
+
+@dataclass(frozen=True, slots=True)
+class RuntimeAsset:
+    """Immutable runtime data with an explicit deserialization trust boundary."""
+
+    id: str
+    repository: str
+    revision: str
+    path: str
+    sha256: str
+    size: int
+    consumer_family: str
+    trust_kind: RuntimeAssetTrustKind
+    license_expression: str
+    offline_behavior: str
+
+
+@dataclass(frozen=True, slots=True)
+class OfficialGolden:
+    """Hash-pinned official output bundle required by the check tier."""
+
+    metadata: FileDigest
+    tensors: FileDigest
+
+
+@dataclass(frozen=True, slots=True)
+class UpstreamSource:
+    """Pinned official implementation used as a parity oracle."""
+
+    id: str
+    path: str
+    url: str
+    revision: str
+    license_expression: str
+    license_files: tuple[str, ...]
+    license_digests: tuple[FileDigest, ...] = ()
+    distribution_files: tuple[FileDigest, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class AttentionKernelSpec:
+    """Immutable Hugging Face kernel used by one attention backend."""
+
+    implementation: str
+    repository: str
+    revision: str
+    version: int
+    expected_variant: str
+    dtypes: tuple[DtypeName, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class ModelFamily:
+    """Shared runtime and compliance contract for one architecture family."""
+
+    id: str
+    architecture: str
+    upstreams: tuple[str, ...]
+    tokenizer_mode: str
+    public_input: str
+    extra: RuntimeExtra
+    reference_container: str
+    reference_adapter: str
+    attention: tuple[str, ...]
+    dtypes: tuple[DtypeName, ...]
+    bf16_execution: Bf16Execution
+    precisions: tuple[str, ...]
+    vram_tier: VramTier
+    checkpoint_license: str
+    hub_license: str
+    state_transform: str
+    representative: str
+    documentation: str
+    test_tiers: tuple[TestTier, ...]
+    runtime_paths: tuple[str, ...]
+    auto_map_items: tuple[tuple[str, str], ...]
+    requires_complete_weight_publication: bool = False
+    weights_publication_allowed: bool = False
+    experimental_precisions: tuple[str, ...] = ()
+    tokenizer_class: str | None = None
+    hub_license_name: str | None = None
+    hub_license_link: str | None = None
+    conversion_provenance: str = ""
+    backbone_model: str | None = None
+
+    @property
+    def auto_map(self) -> Mapping[str, str]:
+        return MappingProxyType(dict(self.auto_map_items))
+
+    @property
+    def hub_license_metadata(self) -> Mapping[str, str]:
+        """Return valid Hugging Face model-card license fields."""
+
+        metadata = {"license": self.hub_license}
+        if self.hub_license_name is not None:
+            # Hugging Face validates custom license names as lowercase slugs,
+            # while the manifest retains the reader-facing display name used
+            # in generated prose.
+            metadata["license_name"] = _HUB_LICENSE_NAME_RE.sub(
+                "-",
+                self.hub_license_name.lower(),
+            ).strip("-.")
+        if self.hub_license_link is not None:
+            metadata["license_link"] = self.hub_license_link
+        return MappingProxyType(metadata)
+
+    @property
+    def stable_precisions(self) -> tuple[str, ...]:
+        """Return precision policies covered by the release contract."""
+
+        experimental = set(self.experimental_precisions)
+        return tuple(precision for precision in self.precisions if precision not in experimental)
+
+
+@dataclass(frozen=True, slots=True)
+class ModelSpec:
+    """Complete immutable source and runtime contract for one checkpoint."""
+
+    id: str
+    family: ModelFamily
+    fast: CheckpointSource
+    official: CheckpointSource
+    size_category: str
+    generation_contract: GenerationContract = "not_applicable"
+    oracle_assets: tuple[OracleAsset, ...] = ()
+    official_golden: OfficialGolden | None = None
+    artifact_source: str = "fast"
+    canonical_state_sha256: str | None = None
+    tokenizer_source_id: str | None = None
+    auto_map_items: tuple[tuple[str, str], ...] = ()
+    notes: str = ""
+    msa_conditioning: bool | None = None
+
+    @property
+    def is_deep_reference(self) -> bool:
+        return self.id == self.family.representative
+
+    @property
+    def auto_map(self) -> Mapping[str, str]:
+        if self.auto_map_items:
+            return MappingProxyType(dict(self.auto_map_items))
+        return self.family.auto_map
+
+    @property
+    def artifact_checkpoint(self) -> CheckpointSource:
+        """Return the checkpoint selected for local artifact construction."""
+
+        return self.fast if self.artifact_source == "fast" else self.official
+
+    @property
+    def oracle_asset_map(self) -> Mapping[str, OracleAsset]:
+        """Return native oracle assets keyed by their declared role."""
+
+        return MappingProxyType({asset.role: asset for asset in self.oracle_assets})
+
+
+class ModelRegistry(Mapping[str, ModelSpec]):
+    """Validated mapping of model IDs to typed model specifications."""
+
+    def __init__(
+        self,
+        *,
+        schema_version: int,
+        upstreams: Mapping[str, UpstreamSource],
+        families: Mapping[str, ModelFamily],
+        models: Mapping[str, ModelSpec],
+        runtime_assets: Mapping[str, RuntimeAsset] = MappingProxyType({}),
+        attention_kernels: Mapping[str, AttentionKernelSpec] = MappingProxyType({}),
+        legal_files: tuple[FileDigest, ...] = (),
+    ) -> None:
+        self.schema_version = schema_version
+        self.upstreams = MappingProxyType(dict(upstreams))
+        self.attention_kernels = MappingProxyType(dict(attention_kernels))
+        self.families = MappingProxyType(dict(families))
+        self._models = MappingProxyType(dict(models))
+        self.runtime_assets = MappingProxyType(dict(runtime_assets))
+        self.legal_files = legal_files
+
+    def __getitem__(self, key: str) -> ModelSpec:
+        return self._models[key]
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._models)
+
+    def __len__(self) -> int:
+        return len(self._models)
+
+    def by_family(self, family_id: str) -> tuple[ModelSpec, ...]:
+        if family_id not in self.families:
+            raise KeyError(family_id)
+        return tuple(model for model in self._models.values() if model.family.id == family_id)
+
+    def supported_attention_dtypes(
+        self,
+        family_id: str,
+        implementation: str,
+    ) -> tuple[DtypeName, ...]:
+        """Return manifest-supported dtypes for one family/backend pair."""
+
+        family = self.families[family_id]
+        if implementation not in family.attention:
+            raise KeyError(
+                f"Family {family_id!r} does not advertise attention backend "
+                f"{implementation!r}."
+            )
+        kernel = self.attention_kernels.get(implementation)
+        if kernel is None:
+            return family.dtypes
+        return tuple(dtype for dtype in family.dtypes if dtype in kernel.dtypes)
+
+    def require_resolved(self, model_id: str | None = None) -> None:
+        """Fail release validation when required file identities remain unresolved."""
+
+        selected = self._models.values() if model_id is None else (self._models[model_id],)
+        unresolved: list[str] = []
+        for model in selected:
+            for label, checkpoint in (("fast", model.fast), ("official", model.official)):
+                for path in checkpoint.unresolved_files:
+                    unresolved.append(f"{model.id}.{label}:{path}")
+        if unresolved:
+            detail = ", ".join(unresolved)
+            raise RegistryError(f"Release provenance is unresolved: {detail}")
+
+
+def _reject_unknown_fields(
+    table: Mapping[str, Any],
+    allowed: frozenset[str],
+    context: str,
+) -> None:
+    unknown = sorted(set(table).difference(allowed))
+    if unknown:
+        raise RegistryError(f"{context} contains unknown fields: {unknown}.")
+
+
+def _require_str(table: Mapping[str, Any], key: str, context: str) -> str:
+    value = table.get(key)
+    if not isinstance(value, str) or not value.strip():
+        raise RegistryError(f"{context}.{key} must be a non-empty string.")
+    return value
+
+
+def _require_enum(
+    table: Mapping[str, Any],
+    key: str,
+    context: str,
+    allowed: frozenset[str],
+) -> str:
+    value = _require_str(table, key, context)
+    if value not in allowed:
+        raise RegistryError(
+            f"{context}.{key} must be one of {sorted(allowed)}; received {value!r}."
+        )
+    return value
+
+
+def _parse_reference_container(table: Mapping[str, Any], context: str) -> str:
+    value = _require_str(table, "reference_container", context)
+    if _REFERENCE_CONTAINER_RE.fullmatch(value) is None:
+        raise RegistryError(
+            f"{context}.reference_container must be a portable 'reference-<name>' target."
+        )
+    return value
+
+
+def _parse_reference_adapter(table: Mapping[str, Any], context: str) -> str:
+    value = _require_str(table, "reference_adapter", context)
+    if _REFERENCE_ADAPTER_RE.fullmatch(value) is None:
+        raise RegistryError(
+            f"{context}.reference_adapter must name one module under "
+            "tests.parity.support.reference_adapters."
+        )
+    return value
+
+
+def _parse_documentation_path(table: Mapping[str, Any], context: str) -> str:
+    value = _require_str(table, "documentation", context)
+    if value.count("#") > 1 or "\\" in value:
+        raise RegistryError(f"{context}.documentation must be a portable documentation path.")
+    raw_path, separator, fragment = value.partition("#")
+    path = PurePosixPath(raw_path)
+    if (
+        path.is_absolute()
+        or ".." in path.parts
+        or len(path.parts) < 2
+        or path.parts[0] != "docs"
+        or path.suffix != ".md"
+        or path.as_posix() != raw_path
+    ):
+        raise RegistryError(
+            f"{context}.documentation must reference a normalized Markdown file under docs/."
+        )
+    if separator and _DOCUMENTATION_FRAGMENT_RE.fullmatch(fragment) is None:
+        raise RegistryError(f"{context}.documentation has an invalid heading fragment.")
+    return value
+
+
+def _require_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]:
+    value = table.get(key)
+    if not isinstance(value, list) or not value or any(not isinstance(item, str) for item in value):
+        raise RegistryError(f"{context}.{key} must be a non-empty string array.")
+    result = tuple(value)
+    if len(set(result)) != len(result):
+        raise RegistryError(f"{context}.{key} contains duplicate values.")
+    return result
+
+
+def _optional_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]:
+    value = table.get(key, [])
+    if not isinstance(value, list) or any(not isinstance(item, str) for item in value):
+        raise RegistryError(f"{context}.{key} must be a string array.")
+    result = tuple(value)
+    if len(set(result)) != len(result):
+        raise RegistryError(f"{context}.{key} contains duplicate values.")
+    return result
+
+
+def _optional_str(table: Mapping[str, Any], key: str, context: str) -> str | None:
+    value = table.get(key)
+    if value is None:
+        return None
+    if (
+        not isinstance(value, str)
+        or not value.strip()
+        or value != value.strip()
+        or "\n" in value
+        or "\r" in value
+    ):
+        raise RegistryError(f"{context}.{key} must be a non-empty single-line string.")
+    return value
+
+
+def _parse_hub_license(
+    table: Mapping[str, Any],
+    *,
+    checkpoint_license: str,
+    context: str,
+) -> tuple[str, str | None, str | None]:
+    expected_fields = {"hub_license", "hub_license_name", "hub_license_link"}
+    unknown_fields = sorted(
+        key for key in table if key.startswith("hub_") and key not in expected_fields
+    )
+    if unknown_fields:
+        raise RegistryError(f"{context} contains unsupported Hub license fields: {unknown_fields}.")
+    identifier = _require_str(table, "hub_license", context)
+    if identifier not in HUB_LICENSE_IDENTIFIERS:
+        raise RegistryError(
+            f"{context}.hub_license must be a supported Hugging Face license identifier."
+        )
+    expected_identifier: str | None = None
+    for prefix, candidate in (
+        ("MIT", "mit"),
+        ("Apache-2.0", "apache-2.0"),
+        ("CC-BY-NC-SA-4.0", "cc-by-nc-sa-4.0"),
+        ("Profluent-E1-Agreement", "other"),
+        ("Unresolved", "other"),
+    ):
+        if checkpoint_license.startswith(prefix):
+            expected_identifier = candidate
+            break
+    if expected_identifier is None:
+        raise RegistryError(
+            f"{context}.checkpoint_license has no declared Hugging Face identifier mapping."
+        )
+    if identifier != expected_identifier:
+        raise RegistryError(
+            f"{context}.hub_license must be {expected_identifier!r} for "
+            f"checkpoint terms {checkpoint_license!r}."
+        )
+
+    name = _optional_str(table, "hub_license_name", context)
+    link = _optional_str(table, "hub_license_link", context)
+    if identifier != "other":
+        if name is not None or link is not None:
+            raise RegistryError(
+                f"{context} may define hub_license_name and hub_license_link only "
+                "when hub_license='other'."
+            )
+        return identifier, None, None
+    if name is None or link is None:
+        raise RegistryError(
+            f"{context} must define hub_license_name and hub_license_link when hub_license='other'."
+        )
+    parsed_link = urlparse(link)
+    if (
+        parsed_link.scheme != "https"
+        or not parsed_link.netloc
+        or not parsed_link.path
+        or parsed_link.username is not None
+        or parsed_link.password is not None
+    ):
+        raise RegistryError(f"{context}.hub_license_link must be an absolute HTTPS URL.")
+    return identifier, name, link
+
+
+def _require_digest_list(
+    table: Mapping[str, Any], key: str, context: str
+) -> tuple[FileDigest, ...]:
+    encoded = _require_str_list(table, key, context)
+    result = tuple(FileDigest.parse(value) for value in encoded)
+    paths = [item.path for item in result]
+    if len(paths) != len(set(paths)):
+        raise RegistryError(f"{context}.{key} contains duplicate paths.")
+    return result
+
+
+def _validate_revision(revision: str, context: str) -> None:
+    if len(revision) != 40 or _HEX_RE.fullmatch(revision) is None:
+        raise RegistryError(f"{context} must be an immutable 40-character commit revision.")
+
+
+def _parse_checkpoint(table: Mapping[str, Any], prefix: str, context: str) -> CheckpointSource:
+    repo_id = _require_str(table, f"{prefix}_repo", context)
+    if _REPOSITORY_ID_RE.fullmatch(repo_id) is None:
+        raise RegistryError(f"{context}.{prefix}_repo must be a Hugging Face repository ID.")
+    revision = _require_str(table, f"{prefix}_revision", context)
+    _validate_revision(revision, f"{context}.{prefix}_revision")
+    encoded_files = _require_str_list(table, f"{prefix}_files", context)
+    files = tuple(FileDigest.parse(value) for value in encoded_files)
+    paths = [item.path for item in files]
+    if len(paths) != len(set(paths)):
+        raise RegistryError(f"{context}.{prefix}_files contains duplicate paths.")
+    if not any(item.path.endswith(_WEIGHT_SUFFIXES) for item in files):
+        raise RegistryError(f"{context}.{prefix}_files does not identify a weight file.")
+    unresolved_files = _optional_str_list(table, f"{prefix}_unresolved_files", context)
+    for unresolved_path in unresolved_files:
+        _portable_relative_path(unresolved_path, "Unresolved checkpoint path")
+        if unresolved_path in paths:
+            raise RegistryError(
+                f"{context}.{prefix} marks {unresolved_path!r} both resolved and unresolved."
+            )
+    return CheckpointSource(
+        repo_id=repo_id,
+        revision=revision,
+        files=files,
+        unresolved_files=unresolved_files,
+    )
+
+
+def _parse_oracle_assets(table: Mapping[str, Any], context: str) -> tuple[OracleAsset, ...]:
+    raw = table.get("oracle_assets", [])
+    if not isinstance(raw, list):
+        raise RegistryError(f"{context}.oracle_assets must be an array of tables.")
+    result: list[OracleAsset] = []
+    for index, value in enumerate(raw):
+        asset_context = f"{context}.oracle_assets[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{asset_context} must be a table.")
+        expected_fields = {"role", "path", "url", "sha256", "size"}
+        if set(value) != expected_fields:
+            raise RegistryError(f"{asset_context} must contain exactly {sorted(expected_fields)}.")
+        role = _require_str(value, "role", asset_context)
+        if role not in _ALLOWED_ORACLE_ASSET_ROLES:
+            raise RegistryError(f"Unsupported oracle asset role: {role!r}.")
+        path = _require_str(value, "path", asset_context)
+        try:
+            normalized_path = _portable_relative_path(path, "Oracle asset path")
+        except RegistryError as error:
+            raise RegistryError(f"Invalid oracle asset path: {path!r}.") from error
+        if normalized_path.suffix != ".pt":
+            raise RegistryError(f"Invalid oracle asset path: {path!r}.")
+        url = _require_str(value, "url", asset_context)
+        parsed_url = urlparse(url)
+        if (
+            parsed_url.scheme != "https"
+            or parsed_url.hostname != _FAIR_ESM_ASSET_HOST
+            or parsed_url.path != f"/fair-esm/{path}"
+            or parsed_url.params
+            or parsed_url.query
+            or parsed_url.fragment
+        ):
+            raise RegistryError(f"Invalid fair-esm oracle asset URL: {url!r}.")
+        sha256 = _require_str(value, "sha256", asset_context)
+        if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None:
+            raise RegistryError(f"Invalid oracle asset SHA-256 for {path!r}.")
+        size = value.get("size")
+        if isinstance(size, bool) or not isinstance(size, int) or size <= 0:
+            raise RegistryError(f"{asset_context}.size must be a positive byte count.")
+        result.append(
+            OracleAsset(
+                role=role,
+                path=path,
+                url=url,
+                sha256=sha256,
+                size=size,
+            )
+        )
+    roles = [asset.role for asset in result]
+    paths = [asset.path for asset in result]
+    urls = [asset.url for asset in result]
+    if (
+        len(roles) != len(set(roles))
+        or len(paths) != len(set(paths))
+        or len(urls) != len(set(urls))
+    ):
+        raise RegistryError(f"{context}.oracle_assets contains duplicate identities.")
+    return tuple(result)
+
+
+def _parse_official_golden(
+    table: Mapping[str, Any],
+    model_id: str,
+    context: str,
+) -> OfficialGolden | None:
+    raw = table.get("official_golden")
+    if raw is None:
+        return None
+    if not isinstance(raw, dict) or set(raw) != {"metadata", "tensors"}:
+        raise RegistryError(
+            f"{context}.official_golden must contain exactly 'metadata' and 'tensors'."
+        )
+    parsed: dict[str, FileDigest] = {}
+    for role in ("metadata", "tensors"):
+        value = raw[role]
+        if not isinstance(value, str):
+            raise RegistryError(f"{context}.official_golden.{role} must be a file digest.")
+        digest = FileDigest.parse(value)
+        if digest.algorithm != "sha256":
+            raise RegistryError(
+                f"{context}.official_golden.{role} must use an immutable SHA-256 digest."
+            )
+        expected = f"tests/goldens/{model_id}.{'json' if role == 'metadata' else 'safetensors'}"
+        if digest.path != expected:
+            raise RegistryError(f"{context}.official_golden.{role} must use path {expected!r}.")
+        parsed[role] = digest
+    return OfficialGolden(metadata=parsed["metadata"], tensors=parsed["tensors"])
+
+
+def _parse_attention_kernels(raw: object) -> dict[str, AttentionKernelSpec]:
+    if not isinstance(raw, list) or not raw:
+        raise RegistryError("The manifest must contain [[attention_kernels]] entries.")
+    result: dict[str, AttentionKernelSpec] = {}
+    expected_variants = {
+        "flash_attention_2": "flash_attn2",
+        "flash_attention_3": "flash_attn3",
+    }
+    for index, value in enumerate(raw):
+        context = f"attention_kernels[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        expected_fields = frozenset(
+            {
+                "implementation",
+                "repository",
+                "revision",
+                "version",
+                "expected_variant",
+                "dtypes",
+            }
+        )
+        _reject_unknown_fields(value, expected_fields, context)
+        implementation = _require_str(value, "implementation", context)
+        if implementation not in expected_variants:
+            raise RegistryError(f"Unsupported attention kernel {implementation!r}.")
+        if implementation in result:
+            raise RegistryError(f"Duplicate attention kernel {implementation!r}.")
+        repository = _require_str(value, "repository", context)
+        if _REPOSITORY_ID_RE.fullmatch(repository) is None:
+            raise RegistryError(f"Invalid attention-kernel repository {repository!r}.")
+        revision = _require_str(value, "revision", context)
+        _validate_revision(revision, f"{context}.revision")
+        kernel_version = value.get("version")
+        if (
+            isinstance(kernel_version, bool)
+            or not isinstance(kernel_version, int)
+            or kernel_version <= 0
+        ):
+            raise RegistryError(f"{context}.version must be a positive integer.")
+        expected_variant = _require_str(value, "expected_variant", context)
+        if expected_variant != expected_variants[implementation]:
+            raise RegistryError(
+                f"{context}.expected_variant must be {expected_variants[implementation]!r}."
+            )
+        dtypes = _require_str_list(value, "dtypes", context)
+        if not set(dtypes).issubset(_ALLOWED_DTYPES):
+            raise RegistryError(f"{context}.dtypes contains unsupported dtypes.")
+        result[implementation] = AttentionKernelSpec(
+            implementation=implementation,
+            repository=repository,
+            revision=revision,
+            version=kernel_version,
+            expected_variant=expected_variant,
+            dtypes=cast(tuple[DtypeName, ...], dtypes),
+        )
+    if set(result) != set(expected_variants):
+        raise RegistryError("The manifest must pin both FlashAttention kernel versions.")
+    return result
+
+
+def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
+    if not isinstance(raw, list) or not raw:
+        raise RegistryError("The manifest must contain at least one [[upstreams]] entry.")
+    result: dict[str, UpstreamSource] = {}
+    paths: set[str] = set()
+    for index, value in enumerate(raw):
+        context = f"upstreams[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(value, _UPSTREAM_FIELDS, context)
+        source_id = _require_str(value, "id", context)
+        if _IDENTIFIER_RE.fullmatch(source_id) is None:
+            raise RegistryError(f"Invalid upstream ID: {source_id!r}")
+        if source_id in result:
+            raise RegistryError(f"Duplicate upstream ID: {source_id!r}")
+        revision = _require_str(value, "revision", context)
+        _validate_revision(revision, f"{context}.revision")
+        path = _require_str(value, "path", context)
+        try:
+            normalized_path = _portable_relative_path(path, f"{context}.path")
+        except RegistryError as error:
+            raise RegistryError(
+                f"{context}.path must be a normalized directory directly under "
+                "'vendor/upstream/'."
+            ) from error
+        if (
+            normalized_path.parts[:2] != ("vendor", "upstream")
+            or len(normalized_path.parts) != 3
+        ):
+            raise RegistryError(
+                f"{context}.path must be a normalized directory directly under "
+                "'vendor/upstream/'."
+            )
+        if path in paths:
+            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.")
+        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:
+            raise RegistryError(
+                f"{context}.license_digests must cover license_files in the same order."
+            )
+        distribution_files = _require_digest_list(value, "distribution_files", context)
+        distribution_map = {item.path: item for item in distribution_files}
+        for canonical in license_digests:
+            distributed = distribution_map.get(canonical.path)
+            if distributed is None or distributed.encoded != canonical.encoded:
+                raise RegistryError(
+                    f"{context}.distribution_files must include an exact copy of "
+                    f"{canonical.path!r}."
+                )
+        if source_id == "e1":
+            required_e1 = {
+                "LICENSE",
+                "ATTRIBUTION",
+                "NOTICE",
+                "Apache-2.0.txt",
+                "BSD-3-Clause.txt",
+                "MODIFICATIONS.md",
+            }
+            missing_e1 = sorted(required_e1.difference(distribution_map))
+            if missing_e1:
+                raise RegistryError(f"{context} is missing E1 legal files: {missing_e1}")
+        result[source_id] = UpstreamSource(
+            id=source_id,
+            path=path,
+            url=url,
+            revision=revision,
+            license_expression=_require_str(value, "license", context),
+            license_files=license_files,
+            license_digests=license_digests,
+            distribution_files=distribution_files,
+        )
+    return result
+
+
+def _parse_families(
+    raw: object,
+    upstreams: Mapping[str, UpstreamSource],
+) -> dict[str, ModelFamily]:
+    if not isinstance(raw, dict) or not raw:
+        raise RegistryError("The manifest must contain [families.<id>] tables.")
+    result: dict[str, ModelFamily] = {}
+    for family_id, value in raw.items():
+        context = f"families.{family_id}"
+        if _IDENTIFIER_RE.fullmatch(family_id) is None or not isinstance(value, dict):
+            raise RegistryError(f"Invalid family table: {family_id!r}")
+        checkpoint_license = _require_str(value, "checkpoint_license", context)
+        hub_license, hub_license_name, hub_license_link = _parse_hub_license(
+            value,
+            checkpoint_license=checkpoint_license,
+            context=context,
+        )
+        _reject_unknown_fields(value, _FAMILY_FIELDS, context)
+        source_ids = _require_str_list(value, "upstreams", context)
+        unknown_sources = sorted(set(source_ids).difference(upstreams))
+        if unknown_sources:
+            raise RegistryError(f"{context} references unknown upstreams: {unknown_sources}")
+        tokenizer_mode = _require_str(value, "tokenizer_mode", context)
+        if tokenizer_mode not in _ALLOWED_TOKENIZER_MODES:
+            raise RegistryError(f"Unsupported tokenizer mode in {context}: {tokenizer_mode!r}")
+        public_input = _require_str(value, "public_input", context)
+        attention = _require_str_list(value, "attention", context)
+        if not set(attention).issubset(_ALLOWED_ATTENTION):
+            raise RegistryError(f"Unsupported attention implementation in {context}.")
+        dtypes = _require_str_list(value, "dtypes", context)
+        if not set(dtypes).issubset(_ALLOWED_DTYPES):
+            raise RegistryError(f"Unsupported dtype in {context}.")
+        bf16_execution = cast(
+            Bf16Execution,
+            _require_enum(
+                value,
+                "bf16_execution",
+                context,
+                _ALLOWED_BF16_EXECUTIONS,
+            ),
+        )
+        precisions = _require_str_list(value, "precisions", context)
+        if not set(precisions).issubset(_ALLOWED_PRECISIONS):
+            raise RegistryError(f"Unsupported precision policy in {context}.")
+        experimental_precisions = _optional_str_list(
+            value,
+            "experimental_precisions",
+            context,
+        )
+        unknown_experimental_precisions = sorted(
+            set(experimental_precisions).difference(precisions)
+        )
+        if unknown_experimental_precisions:
+            raise RegistryError(
+                f"{context}.experimental_precisions must be a subset of precisions; "
+                f"unknown values: {unknown_experimental_precisions}."
+            )
+        extra = cast(RuntimeExtra, _require_enum(value, "extra", context, _ALLOWED_EXTRAS))
+        vram_tier = cast(
+            VramTier,
+            _require_enum(value, "vram_tier", context, _ALLOWED_VRAM_TIERS),
+        )
+        test_tiers_raw = _require_str_list(value, "test_tiers", context)
+        unknown_test_tiers = sorted(set(test_tiers_raw).difference(_ALLOWED_TEST_TIERS))
+        if unknown_test_tiers:
+            raise RegistryError(
+                f"{context}.test_tiers contains unsupported tiers: {unknown_test_tiers}."
+            )
+        test_tiers = cast(tuple[TestTier, ...], test_tiers_raw)
+        reference_container = _parse_reference_container(value, context)
+        reference_adapter = _parse_reference_adapter(value, context)
+        documentation = _parse_documentation_path(value, context)
+        runtime_paths = _require_str_list(value, "runtime_paths", context)
+        if len(runtime_paths) != len(set(runtime_paths)):
+            raise RegistryError(f"{context}.runtime_paths must not contain duplicates.")
+        for runtime_path in runtime_paths:
+            try:
+                _portable_relative_path(runtime_path, f"{context}.runtime_paths entry")
+            except RegistryError as error:
+                raise RegistryError(
+                    f"Unsafe runtime path in {context}: {runtime_path!r}"
+                ) from error
+            if runtime_path.startswith("vendor/"):
+                raise RegistryError(f"Unsafe runtime path in {context}: {runtime_path!r}")
+        requires_complete_weight_publication = value.get(
+            "requires_complete_weight_publication",
+            False,
+        )
+        if not isinstance(requires_complete_weight_publication, bool):
+            raise RegistryError(
+                f"{context}.requires_complete_weight_publication must be a boolean."
+            )
+        if "weights_publication_allowed" not in value:
+            raise RegistryError(
+                f"{context}.weights_publication_allowed must be declared explicitly."
+            )
+        weights_publication_allowed = value["weights_publication_allowed"]
+        if not isinstance(weights_publication_allowed, bool):
+            raise RegistryError(f"{context}.weights_publication_allowed must be a boolean.")
+        raw_auto_map = value.get("auto_map")
+        if not isinstance(raw_auto_map, dict) or not raw_auto_map:
+            raise RegistryError(f"{context}.auto_map must be a non-empty table.")
+        auto_map: list[tuple[str, str]] = []
+        for auto_class, class_path in raw_auto_map.items():
+            if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str):
+                raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}")
+            if not class_path.startswith("fastplms.") or class_path.count(".") < 2:
+                raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}")
+            auto_map.append((auto_class, class_path))
+        tokenizer_class = value.get("tokenizer_class")
+        if tokenizer_class is not None:
+            if tokenizer_mode != "tokenizer":
+                raise RegistryError(
+                    f"{context}.tokenizer_class requires tokenizer_mode='tokenizer'."
+                )
+            if (
+                not isinstance(tokenizer_class, str)
+                or not tokenizer_class.startswith("fastplms.")
+                or tokenizer_class.count(".") < 2
+            ):
+                raise RegistryError(
+                    f"Invalid tokenizer class path in {context}: {tokenizer_class!r}"
+                )
+        backbone_model = value.get("backbone_model")
+        if backbone_model is not None and (
+            not isinstance(backbone_model, str)
+            or _IDENTIFIER_RE.fullmatch(backbone_model) is None
+        ):
+            raise RegistryError(
+                f"{context}.backbone_model must be a valid manifest model ID."
+            )
+        state_transform = _require_str(value, "state_transform", context)
+        conversion_provenance = _require_str(value, "conversion_provenance", context)
+        required_sections = ("Input:", "Transformation:", "Output:", "Validation:", "Limitation:")
+        missing_sections = [
+            section for section in required_sections if section not in conversion_provenance
+        ]
+        if missing_sections or state_transform not in conversion_provenance:
+            raise RegistryError(
+                f"{context}.conversion_provenance must identify {state_transform!r} and "
+                f"contain mechanism-first sections; missing {missing_sections}."
+            )
+        result[family_id] = ModelFamily(
+            id=family_id,
+            architecture=_require_str(value, "architecture", context),
+            upstreams=source_ids,
+            tokenizer_mode=tokenizer_mode,
+            public_input=public_input,
+            extra=extra,
+            reference_container=reference_container,
+            reference_adapter=reference_adapter,
+            attention=attention,
+            dtypes=cast(tuple[DtypeName, ...], dtypes),
+            bf16_execution=bf16_execution,
+            precisions=precisions,
+            experimental_precisions=experimental_precisions,
+            vram_tier=vram_tier,
+            checkpoint_license=checkpoint_license,
+            hub_license=hub_license,
+            state_transform=state_transform,
+            representative=_require_str(value, "representative", context),
+            documentation=documentation,
+            test_tiers=test_tiers,
+            runtime_paths=runtime_paths,
+            auto_map_items=tuple(auto_map),
+            requires_complete_weight_publication=requires_complete_weight_publication,
+            weights_publication_allowed=weights_publication_allowed,
+            tokenizer_class=tokenizer_class,
+            hub_license_name=hub_license_name,
+            hub_license_link=hub_license_link,
+            conversion_provenance=conversion_provenance,
+            backbone_model=backbone_model,
+        )
+    return result
+
+
+def _parse_runtime_assets(
+    raw: object,
+    families: Mapping[str, ModelFamily],
+) -> dict[str, RuntimeAsset]:
+    if not isinstance(raw, list) or not raw:
+        raise RegistryError("The manifest must contain at least one [[runtime_assets]] entry.")
+    result: dict[str, RuntimeAsset] = {}
+    identities: set[tuple[str, str, str]] = set()
+    for index, value in enumerate(raw):
+        context = f"runtime_assets[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(value, _RUNTIME_ASSET_FIELDS, context)
+        asset_id = _require_str(value, "id", context)
+        if _IDENTIFIER_RE.fullmatch(asset_id) is None:
+            raise RegistryError(f"Invalid runtime asset ID: {asset_id!r}")
+        if asset_id in result:
+            raise RegistryError(f"Duplicate runtime asset ID: {asset_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)
+        try:
+            normalized_path = _portable_relative_path(path, "Runtime asset path")
+        except RegistryError as error:
+            raise RegistryError(f"Runtime asset path is not portable: {path!r}") from error
+        sha256 = _require_str(value, "sha256", context)
+        if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None:
+            raise RegistryError(f"Invalid runtime asset SHA-256 for {path!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.")
+        consumer_family = _require_str(value, "consumer_family", context)
+        if consumer_family not in families:
+            raise RegistryError(
+                f"{context}.consumer_family references unknown family {consumer_family!r}."
+            )
+        trust_kind = cast(
+            RuntimeAssetTrustKind,
+            _require_enum(
+                value,
+                "trust_kind",
+                context,
+                _ALLOWED_RUNTIME_ASSET_TRUST_KINDS,
+            ),
+        )
+        license_expression = _require_str(value, "license", context)
+        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}."
+            )
+        if trust_kind == "hash_pinned_pickle" and normalized_path.suffix != ".pkl":
+            raise RegistryError(
+                f"{context}.path must end in '.pkl' for trust_kind='hash_pinned_pickle'."
+            )
+        identity = (repository, revision, path)
+        if identity in identities:
+            raise RegistryError(f"Duplicate runtime asset identity: {identity!r}")
+        identities.add(identity)
+        result[asset_id] = RuntimeAsset(
+            id=asset_id,
+            repository=repository,
+            revision=revision,
+            path=path,
+            sha256=sha256,
+            size=size,
+            consumer_family=consumer_family,
+            trust_kind=trust_kind,
+            license_expression=license_expression,
+            offline_behavior=offline_behavior,
+        )
+    return result
+
+
+def _parse_models(
+    raw: object,
+    families: Mapping[str, ModelFamily],
+) -> dict[str, ModelSpec]:
+    if not isinstance(raw, list) or not raw:
+        raise RegistryError("The manifest must contain at least one [[models]] entry.")
+    result: dict[str, ModelSpec] = {}
+    fast_repositories: set[str] = set()
+    for index, value in enumerate(raw):
+        context = f"models[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(value, _MODEL_FIELDS, context)
+        model_id = _require_str(value, "id", context)
+        if _IDENTIFIER_RE.fullmatch(model_id) is None:
+            raise RegistryError(f"Invalid model ID: {model_id!r}")
+        if model_id in result:
+            raise RegistryError(f"Duplicate model ID: {model_id!r}")
+        family_id = _require_str(value, "family", context)
+        if family_id not in families:
+            raise RegistryError(f"{context} references unknown family {family_id!r}.")
+        fast = _parse_checkpoint(value, "fast", context)
+        official = _parse_checkpoint(value, "official", context)
+        if fast.repo_id in fast_repositories:
+            raise RegistryError(f"Duplicate FastPLMs repository ID: {fast.repo_id!r}")
+        fast_repositories.add(fast.repo_id)
+        family = families[family_id]
+        oracle_assets = _parse_oracle_assets(value, context)
+        official_golden = _parse_official_golden(value, model_id, context)
+        size_category = _require_str(value, "size_category", context)
+        if size_category not in _ALLOWED_SIZE_CATEGORIES:
+            raise RegistryError(f"Unsupported size category in {context}: {size_category!r}")
+        generation_contract = cast(
+            GenerationContract,
+            _require_enum(
+                value,
+                "generation_contract",
+                context,
+                _ALLOWED_GENERATION_CONTRACTS,
+            ),
+        )
+        if family.tokenizer_mode == "structure" and size_category != "structure":
+            raise RegistryError(
+                f"Structure checkpoint {model_id!r} must use size_category='structure'."
+            )
+        artifact_source = value.get("artifact_source", "fast")
+        if artifact_source not in {"fast", "official"}:
+            raise RegistryError(f"{context}.artifact_source must be 'fast' or 'official'.")
+        canonical_state_sha256 = value.get("canonical_state_sha256")
+        if artifact_source == "official":
+            if (
+                not isinstance(canonical_state_sha256, str)
+                or len(canonical_state_sha256) != 64
+                or _HEX_RE.fullmatch(canonical_state_sha256) is None
+            ):
+                raise RegistryError(
+                    f"{context}.canonical_state_sha256 must be a SHA-256 commitment "
+                    "for an official-source artifact."
+                )
+        elif canonical_state_sha256 is not None:
+            raise RegistryError(
+                f"{context}.canonical_state_sha256 is restricted to official-source artifacts."
+            )
+        if family.tokenizer_mode == "tokenizer" and not any(
+            "tokenizer" in item.path or "vocab" in item.path for item in fast.files
+        ):
+            raise RegistryError(f"{context} does not pin a tokenizer asset.")
+        tokenizer_source_id = value.get("tokenizer_source")
+        if tokenizer_source_id is not None and (
+            family.tokenizer_mode != "tokenizer"
+            or not isinstance(tokenizer_source_id, str)
+            or _IDENTIFIER_RE.fullmatch(tokenizer_source_id) is None
+        ):
+            raise RegistryError(f"{context}.tokenizer_source is invalid.")
+        notes = value.get("notes", "")
+        if not isinstance(notes, str):
+            raise RegistryError(f"{context}.notes must be a string.")
+        msa_conditioning = value.get("msa_conditioning")
+        if family_id == "esmfold2":
+            if not isinstance(msa_conditioning, bool):
+                raise RegistryError(
+                    f"{context}.msa_conditioning must be an explicit boolean for "
+                    "ESMFold2 checkpoints."
+                )
+        elif "msa_conditioning" in value:
+            raise RegistryError(
+                f"{context}.msa_conditioning is only valid for ESMFold2 checkpoints."
+            )
+        raw_auto_map = value.get("auto_map")
+        auto_map: list[tuple[str, str]] = []
+        if raw_auto_map is not None:
+            if not isinstance(raw_auto_map, dict) or not raw_auto_map:
+                raise RegistryError(f"{context}.auto_map must be a non-empty table.")
+            for auto_class, class_path in raw_auto_map.items():
+                if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str):
+                    raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}")
+                if not class_path.startswith("fastplms.") or class_path.count(".") < 2:
+                    raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}")
+                auto_map.append((auto_class, class_path))
+        result[model_id] = ModelSpec(
+            id=model_id,
+            family=family,
+            fast=fast,
+            official=official,
+            size_category=size_category,
+            generation_contract=generation_contract,
+            oracle_assets=oracle_assets,
+            official_golden=official_golden,
+            artifact_source=artifact_source,
+            canonical_state_sha256=canonical_state_sha256,
+            tokenizer_source_id=tokenizer_source_id,
+            auto_map_items=tuple(auto_map),
+            notes=notes,
+            msa_conditioning=msa_conditioning,
+        )
+    return result
+
+
+def _validate_registry(
+    upstreams: Mapping[str, UpstreamSource],
+    attention_kernels: Mapping[str, AttentionKernelSpec],
+    families: Mapping[str, ModelFamily],
+    models: Mapping[str, ModelSpec],
+) -> None:
+    for spec in models.values():
+        if spec.tokenizer_source_id is None:
+            continue
+        source = models.get(spec.tokenizer_source_id)
+        if source is None:
+            raise RegistryError(
+                f"Model {spec.id!r} references unknown tokenizer source "
+                f"{spec.tokenizer_source_id!r}."
+            )
+        if not any(
+            PurePosixPath(item.path).name
+            in {
+                "added_tokens.json",
+                "merges.txt",
+                "sentencepiece.bpe.model",
+                "special_tokens_map.json",
+                "spiece.model",
+                "tokenizer.json",
+                "tokenizer_config.json",
+                "vocab.json",
+                "vocab.txt",
+            }
+            for item in source.official.files
+        ):
+            raise RegistryError(
+                f"Tokenizer source {source.id!r} has no official tokenizer assets."
+            )
+    expected_esmfold2 = {
+        "esmfold2": ("Synthyra/ESMFold2", "biohub/ESMFold2"),
+        "esmfold2_fast": ("Synthyra/ESMFold2-Fast", "biohub/ESMFold2-Fast"),
+        "esmfold2_experimental_cutoff2025": (
+            "Synthyra/ESMFold2-Experimental-Cutoff2025",
+            "biohub/ESMFold2-Experimental-Cutoff2025",
+        ),
+        "esmfold2_experimental_fast_cutoff2025": (
+            "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025",
+            "biohub/ESMFold2-Experimental-Fast-Cutoff2025",
+        ),
+    }
+    actual_esmfold2 = {
+        model.id: (model.fast.repo_id, model.official.repo_id)
+        for model in models.values()
+        if model.family.id == "esmfold2"
+    }
+    if actual_esmfold2 != expected_esmfold2:
+        raise RegistryError(
+            "ESMFold2 support must contain exactly the four approved model IDs and "
+            "official/Synthyra repositories."
+        )
+
+    golden_paths: list[str] = []
+    for model in models.values():
+        if model.official_golden is not None:
+            golden_paths.extend(
+                (
+                    model.official_golden.metadata.path,
+                    model.official_golden.tensors.path,
+                )
+            )
+    if len(golden_paths) != len(set(golden_paths)):
+        raise RegistryError("Official golden paths must be unique across model declarations.")
+    unused_upstreams = sorted(
+        set(upstreams).difference(
+            source for family in families.values() for source in family.upstreams
+        )
+    )
+    if unused_upstreams:
+        raise RegistryError(
+            f"Upstream sources are not connected to a model family: {unused_upstreams}"
+        )
+    advertised_flash = {
+        implementation
+        for family in families.values()
+        for implementation in family.attention
+        if implementation.startswith("flash_attention_")
+    }
+    missing_kernels = sorted(advertised_flash.difference(attention_kernels))
+    if missing_kernels:
+        raise RegistryError(
+            f"Advertised FlashAttention backends lack kernel specs: {missing_kernels}."
+        )
+    for family in families.values():
+        for implementation in family.attention:
+            kernel = attention_kernels.get(implementation)
+            if kernel is not None and not set(family.dtypes).intersection(kernel.dtypes):
+                raise RegistryError(
+                    f"Family {family.id!r} and attention kernel {implementation!r} "
+                    "have no supported dtype in common."
+                )
+        family_models = [model for model in models.values() if model.family.id == family.id]
+        if not family_models:
+            raise RegistryError(f"Family {family.id!r} has no checkpoints.")
+        representative = models.get(family.representative)
+        if representative is None or representative.family.id != family.id:
+            raise RegistryError(
+                f"Family {family.id!r} has invalid representative {family.representative!r}."
+            )
+        if family.backbone_model is not None and family.backbone_model not in models:
+            raise RegistryError(
+                f"Family {family.id!r} references unknown backbone model "
+                f"{family.backbone_model!r}."
+            )
+
+
+def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry:
+    try:
+        data = tomllib.loads(raw_bytes.decode("utf-8"))
+    except (UnicodeDecodeError, tomllib.TOMLDecodeError) as error:
+        raise RegistryError(f"Unable to parse model manifest: {error}") from error
+    _reject_unknown_fields(data, _ROOT_FIELDS, "manifest")
+    if data.get("schema_version") != 1:
+        raise RegistryError("Unsupported model manifest schema_version; expected 1.")
+    legal_files = _require_digest_list(data, "legal_files", "manifest")
+    required_legal_paths = {"LICENSE", "THIRD_PARTY_NOTICES.md"}
+    if {item.path for item in legal_files} != required_legal_paths:
+        raise RegistryError("manifest.legal_files must contain LICENSE and THIRD_PARTY_NOTICES.md.")
+    attention_kernels = _parse_attention_kernels(data.get("attention_kernels"))
+    upstreams = _parse_upstreams(data.get("upstreams"))
+    families = _parse_families(data.get("families"), upstreams)
+    runtime_assets = _parse_runtime_assets(data.get("runtime_assets"), families)
+    models = _parse_models(data.get("models"), families)
+    _validate_registry(upstreams, attention_kernels, families, models)
+    return ModelRegistry(
+        schema_version=1,
+        upstreams=upstreams,
+        attention_kernels=attention_kernels,
+        families=families,
+        models=models,
+        runtime_assets=runtime_assets,
+        legal_files=legal_files,
+    )
+
+
+def load_model_registry(path: str | Path | None = None) -> ModelRegistry:
+    """Load and validate a model manifest without importing model code."""
+
+    if path is None:
+        manifest = resources.files("fastplms").joinpath("models.toml")
+        return _load_manifest_bytes(manifest.read_bytes())
+    return _load_manifest_bytes(Path(path).read_bytes())
+
+
+@lru_cache(maxsize=1)
+def get_model_registry() -> ModelRegistry:
+    """Return the validated package registry, cached after its first read."""
+
+    return load_model_registry()
+
+
+def get_model_spec(model_id: str) -> ModelSpec:
+    """Return one model specification by its stable manifest ID."""
+
+    try:
+        return get_model_registry()[model_id]
+    except KeyError as error:
+        supported = ", ".join(get_model_registry())
+        raise KeyError(
+            f"Unknown FastPLMs model ID {model_id!r}. Supported IDs: {supported}"
+        ) from error
+
+
+__all__ = [
+    "HUB_LICENSE_IDENTIFIERS",
+    "CheckpointSource",
+    "FileDigest",
+    "GenerationContract",
+    "ModelFamily",
+    "ModelRegistry",
+    "ModelSpec",
+    "OracleAsset",
+    "RegistryError",
+    "RuntimeAsset",
+    "RuntimeAssetTrustKind",
+    "RuntimeExtra",
+    "TestTier",
+    "UpstreamSource",
+    "VramTier",
+    "get_model_registry",
+    "get_model_spec",
+    "load_model_registry",
+]

fastplms/runtime.py

@@ -0,0 +1,68 @@
+"""Explicit, reversible Torch runtime configuration.
+
+Importing FastPLMs does not change global Torch settings. Callers that want a
+runtime profile opt in with :func:`runtime_profile` and receive their previous
+settings back when the context exits.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+MatmulPrecision = Literal["highest", "high", "medium"]
+
+
+@dataclass(frozen=True, slots=True)
+class RuntimeProfile:
+    """Requested Torch settings for a bounded inference or training block."""
+
+    float32_matmul_precision: MatmulPrecision = "highest"
+    allow_tf32: bool | None = None
+
+
+@contextmanager
+def runtime_profile(profile: RuntimeProfile | None = None) -> Iterator[None]:
+    """Apply a Torch runtime profile and restore the previous global settings.
+
+    The default profile requests the highest float32 matrix-multiplication
+    precision and leaves TF32 policy unchanged. Torch is imported only when the
+    context is entered.
+    """
+
+    import torch
+
+    selected = profile or RuntimeProfile()
+    previous_matmul_precision = torch.get_float32_matmul_precision()
+    matmul_backend = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
+    cudnn_backend = getattr(torch.backends, "cudnn", None)
+    previous_matmul_tf32 = (
+        getattr(matmul_backend, "allow_tf32", None) if matmul_backend is not None else None
+    )
+    previous_cudnn_tf32 = (
+        getattr(cudnn_backend, "allow_tf32", None) if cudnn_backend is not None else None
+    )
+
+    torch.set_float32_matmul_precision(selected.float32_matmul_precision)
+    if selected.allow_tf32 is not None:
+        if matmul_backend is not None and hasattr(matmul_backend, "allow_tf32"):
+            matmul_backend.allow_tf32 = selected.allow_tf32
+        if cudnn_backend is not None and hasattr(cudnn_backend, "allow_tf32"):
+            cudnn_backend.allow_tf32 = selected.allow_tf32
+    try:
+        yield
+    finally:
+        torch.set_float32_matmul_precision(previous_matmul_precision)
+        if selected.allow_tf32 is not None:
+            if matmul_backend is not None and previous_matmul_tf32 is not None:
+                matmul_backend.allow_tf32 = previous_matmul_tf32
+            if cudnn_backend is not None and previous_cudnn_tf32 is not None:
+                cudnn_backend.allow_tf32 = previous_cudnn_tf32
+
+
+__all__ = ["MatmulPrecision", "RuntimeProfile", "runtime_profile"]

modeling_fastplms.py

@@ -0,0 +1,230 @@
+"""Generated bridge to the unchanged FastPLMs package sources."""
+
+import base64
+import hashlib
+import importlib
+import importlib.util
+import sys
+import tempfile
+from importlib.metadata import PackageNotFoundError, distribution
+from io import BytesIO
+from pathlib import Path
+from zipfile import ZIP_DEFLATED, ZipFile
+
+from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
+
+if RUNTIME_HASH != "297af2aeeb817b6e3a343bf73dbc2eb3b3e29a2ca88ccba8af1aeb2e12243ffd":
+    raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
+
+_RUNTIME_TEMPORARIES = []
+
+def _archive_runtime_hashes(payload):
+    result = {}
+    with ZipFile(BytesIO(payload)) as archive:
+        for member in archive.infolist():
+            name = member.filename
+            parts = Path(name).parts
+            if (
+                member.is_dir()
+                or "\\" in name
+                or not parts
+                or parts[0] != "fastplms"
+                or len(parts) < 2
+                or any(part in {"", ".", ".."} for part in parts)
+                or Path(name).suffix in {".pyc", ".pyo"}
+                or member.flag_bits & 0x1
+                or member.compress_type != ZIP_DEFLATED
+                or member.external_attr >> 16 != 0o100644
+            ):
+                raise RuntimeError("Embedded FastPLMs archive has an unsafe path.")
+            relative = Path(*parts[1:]).as_posix()
+            if relative in result:
+                raise RuntimeError("Embedded FastPLMs archive repeats a path.")
+            result[relative] = hashlib.sha256(archive.read(member)).hexdigest()
+    return result
+
+def _ensure_runtime():
+    payload = base64.b85decode("".join(RUNTIME_DATA))
+    if hashlib.sha256(payload).hexdigest() != RUNTIME_HASH:
+        raise RuntimeError("Embedded FastPLMs runtime hash mismatch.")
+    expected = _archive_runtime_hashes(payload)
+    temporary = tempfile.TemporaryDirectory(prefix="fastplms-artifact-runtime-")
+    try:
+        runtime_root = Path(temporary.name)
+        with ZipFile(BytesIO(payload)) as archive:
+            for member in archive.infolist():
+                target = runtime_root.joinpath(*Path(member.filename).parts)
+                target.parent.mkdir(parents=True, exist_ok=True)
+                with target.open("xb") as handle:
+                    handle.write(archive.read(member))
+        package_root = runtime_root / "fastplms"
+        if _runtime_file_hashes(package_root) != expected:
+            raise RuntimeError(
+                "Private FastPLMs runtime differs from the embedded archive."
+            )
+    except BaseException:
+        temporary.cleanup()
+        raise
+    _RUNTIME_TEMPORARIES.append(temporary)
+    return package_root
+
+def _runtime_file_hashes(package_root):
+    result = {}
+    for path in sorted(package_root.rglob("*")):
+        relative = path.relative_to(package_root)
+        if path.is_symlink():
+            raise RuntimeError("Private FastPLMs runtime contains a symlink.")
+        if path.is_dir():
+            continue
+        if path.suffix in {".pyc", ".pyo"}:
+            raise RuntimeError("Private FastPLMs runtime contains bytecode.")
+        if not path.is_file():
+            raise RuntimeError("Private FastPLMs runtime contains a non-file entry.")
+        result[relative.as_posix()] = hashlib.sha256(path.read_bytes()).hexdigest()
+    return result
+
+def _installed_runtime_digest(installed_root, relative):
+    candidate = installed_root / relative
+    if candidate.is_file():
+        return hashlib.sha256(candidate.read_bytes()).hexdigest()
+    if relative != "kernels.lock":
+        return None
+    try:
+        installed_distribution = distribution("fastplms")
+    except PackageNotFoundError:
+        return None
+    for entry in installed_distribution.files or ():
+        normalized = str(entry).replace("\\", "/")
+        if normalized.endswith(".dist-info/kernels.lock"):
+            lock_path = Path(installed_distribution.locate_file(entry))
+            if lock_path.is_file():
+                return hashlib.sha256(lock_path.read_bytes()).hexdigest()
+    return None
+
+def _extend_loaded_package_paths(package_root):
+    for name, module in list(sys.modules.items()):
+        if name != "fastplms" and not name.startswith("fastplms."):
+            continue
+        paths = getattr(module, "__path__", None)
+        if paths is None:
+            continue
+        relative = name.split(".")[1:]
+        candidate = package_root.joinpath(*relative)
+        candidate_text = str(candidate)
+        if candidate.is_dir() and candidate_text not in paths:
+            paths.append(candidate_text)
+
+def _merge_runtime(installed, package_root):
+    incoming = _runtime_file_hashes(package_root)
+    known = dict(getattr(installed, "__fastplms_artifact_runtime_files__", {}))
+    installed_root_text = getattr(
+        installed, "__fastplms_artifact_installed_root__", None
+    )
+    if not known:
+        installed_file = getattr(installed, "__file__", None)
+        if installed_file is None:
+            raise RuntimeError(
+                "The loaded fastplms package has no source path and cannot be verified "
+                "against the embedded artifact runtime."
+            )
+        installed_root = Path(installed_file).resolve().parent
+        for relative, digest in incoming.items():
+            if _installed_runtime_digest(installed_root, relative) != digest:
+                raise RuntimeError(
+                    "The installed FastPLMs runtime differs from this artifact at "
+                    f"{relative!r}. Install the artifact's matching FastPLMs release "
+                    "or use a separate Python process."
+                )
+        installed_root_text = str(installed_root)
+        installed.__fastplms_artifact_installed_root__ = installed_root_text
+    conflicts = sorted(
+        relative
+        for relative, digest in incoming.items()
+        if relative in known and known[relative] != digest
+    )
+    if conflicts:
+        raise RuntimeError(
+            "FastPLMs artifacts contain incompatible runtime sources at "
+            + ", ".join(repr(path) for path in conflicts[:5])
+            + ". Load incompatible releases in separate Python processes."
+        )
+    if installed_root_text is not None:
+        installed_root = Path(installed_root_text)
+        for relative, digest in incoming.items():
+            if relative in known:
+                continue
+            if _installed_runtime_digest(installed_root, relative) != digest:
+                raise RuntimeError(
+                    "The installed FastPLMs runtime differs from this artifact at "
+                    f"{relative!r}. Install the artifact's matching FastPLMs release "
+                    "or use a separate Python process."
+                )
+    known.update(incoming)
+    installed.__fastplms_artifact_runtime_files__ = known
+    roots = list(getattr(installed, "__fastplms_artifact_runtime_roots__", ()))
+    if str(package_root) not in roots:
+        roots.append(str(package_root))
+    installed.__fastplms_artifact_runtime_roots__ = tuple(roots)
+    temporaries = list(
+        getattr(installed, "__fastplms_artifact_runtime_temporaries__", ())
+    )
+    for temporary in _RUNTIME_TEMPORARIES:
+        if temporary not in temporaries:
+            temporaries.append(temporary)
+    installed.__fastplms_artifact_runtime_temporaries__ = tuple(temporaries)
+    hashes = set(getattr(installed, "__fastplms_artifact_runtime_hashes__", ()))
+    hashes.add(RUNTIME_HASH)
+    installed.__fastplms_artifact_runtime_hashes__ = frozenset(hashes)
+    _extend_loaded_package_paths(package_root)
+    return installed
+
+def _import_without_bytecode(module_name):
+    previous = sys.dont_write_bytecode
+    sys.dont_write_bytecode = True
+    try:
+        return importlib.import_module(module_name)
+    finally:
+        sys.dont_write_bytecode = previous
+
+def _install_runtime():
+    installed = sys.modules.get("fastplms")
+    hashes = getattr(installed, "__fastplms_artifact_runtime_hashes__", ())
+    if RUNTIME_HASH in hashes:
+        return installed
+    package_root = _ensure_runtime()
+    if installed is not None:
+        return _merge_runtime(installed, package_root)
+    spec = importlib.util.spec_from_file_location(
+        "fastplms",
+        package_root / "__init__.py",
+        submodule_search_locations=[str(package_root)],
+    )
+    if spec is None or spec.loader is None:
+        raise ImportError("Unable to load the embedded FastPLMs runtime.")
+    package = importlib.util.module_from_spec(spec)
+    package.__fastplms_artifact_runtime_hash__ = RUNTIME_HASH
+    package.__fastplms_artifact_runtime_hashes__ = frozenset({RUNTIME_HASH})
+    package.__fastplms_artifact_runtime_files__ = _runtime_file_hashes(package_root)
+    package.__fastplms_artifact_runtime_roots__ = (str(package_root),)
+    package.__fastplms_artifact_runtime_temporaries__ = tuple(
+        _RUNTIME_TEMPORARIES
+    )
+    sys.modules["fastplms"] = package
+    previous = sys.dont_write_bytecode
+    sys.dont_write_bytecode = True
+    try:
+        try:
+            spec.loader.exec_module(package)
+        except BaseException:
+            sys.modules.pop("fastplms", None)
+            raise
+    finally:
+        sys.dont_write_bytecode = previous
+    return package
+
+_install_runtime()
+_module_225 = _import_without_bytecode("fastplms.models.boltz.modeling_boltz2")
+Boltz2Config = _module_225.Boltz2Config
+Boltz2Config.__module__ = __name__
+Boltz2Model = _module_225.Boltz2Model
+Boltz2Model.__module__ = __name__

runtime-attestation.json

@@ -0,0 +1,60 @@
+{
+  "files": {
+    "LICENSES/FastPLMs-Apache-2.0.txt": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
+    "LICENSES/boltz/LICENSE": "sha256:f0667fd5e66c51e1ba8ddaa0249c6d7225b30037e02c45782d8f2c2943ac2617",
+    "README.md": "sha256:49eaec9931001fa9175755c386a28403a5215f3ef7e79f9331d9442246785614",
+    "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
+    "config.json": "sha256:f285511527e78e8800f52f948cf6b1779ac2444c84ee60f25c53ba098ab343af",
+    "fastplms/__init__.py": "sha256:4fb3196022ca8ec699d59d09bdbc5f0184195552b773698ab9b061fe3cd7df12",
+    "fastplms/models.toml": "sha256:05a8399f084a5babb5f0916cee7e564c4030767f3ff0230c4f46e539209847d1",
+    "fastplms/models/__init__.py": "sha256:5e48c2cb3877aa6f42f3b5411d53b16bba2e32827bbde634f47f174c5cb36f86",
+    "fastplms/models/boltz/__init__.py": "sha256:54f7ad44ddbb3fd160382a50bb5d5a8a5b86cb66978b7d72445194fc21caa673",
+    "fastplms/models/boltz/_pair_attention.py": "sha256:13d16831664ac7d0bfc859e2d4b97d523f1c2ae0fe4a0b11f65b9a3197dffccd",
+    "fastplms/models/boltz/cif_writer.py": "sha256:095eeab50ee8b97fe0160b4c1f9c6df2be1412a1de13a1d8ad00f9847032a7ab",
+    "fastplms/models/boltz/minimal_featurizer.py": "sha256:35e5f4f85d24d5ba8d952ee64f1ddb98e9ce5245f364d231da185ae8a3b28ce1",
+    "fastplms/models/boltz/minimal_structures.py": "sha256:df7f1ec765ddbb5a730edb9f304fe312a5d3c431ae564ca9b97e6537e0ebf476",
+    "fastplms/models/boltz/modeling_boltz2.py": "sha256:c0d00a63a23762853257a4a4c718f7a4ded358c900562ca260ed3596a113d934",
+    "fastplms/models/boltz/vb_const.py": "sha256:67d3cfdb7f0facdb217fe9220404d598ac7924a672f1a2720f7d2d82b23270df",
+    "fastplms/models/boltz/vb_layers_attention.py": "sha256:ad09126f467b9c73f764b6258ae3f2bce9f9752f641c0347052f12548abdfe4b",
+    "fastplms/models/boltz/vb_layers_attentionv2.py": "sha256:407b1ce628d53c3e0b6f801d989b0ae59bd87d9f9e5d1155e61c1e1ec5a78e0f",
+    "fastplms/models/boltz/vb_layers_confidence_utils.py": "sha256:61b024671ca2eaaa08eb6ddfd75f374f8db97ef6ef75ffcfc8c10347f4fd239a",
+    "fastplms/models/boltz/vb_layers_dropout.py": "sha256:61db64dc4e6be29a41bad2705f73cc489dcf3227cc6e8914c1ea9e6f45ba57a5",
+    "fastplms/models/boltz/vb_layers_initialize.py": "sha256:c1f5a5d56843438bd11639778eed06352e3d4995d2d12c967865d1221208048c",
+    "fastplms/models/boltz/vb_layers_outer_product_mean.py": "sha256:60478bea7d9d6d7f518ab1461c37b91e42d9fa0ab8cfcd777be88c1ddbfff07e",
+    "fastplms/models/boltz/vb_layers_pair_averaging.py": "sha256:e626502b549348ecf26077d0cb1f88f37970b6b9987a6780b4599c5a0b1c0ada",
+    "fastplms/models/boltz/vb_layers_pairformer.py": "sha256:fb056b76231366a1a93373ff2effddde887768dcfca1c9ebc46a567ba39eacb2",
+    "fastplms/models/boltz/vb_layers_transition.py": "sha256:62d306f33a600d08f8bb60a1a8cb9a00a253f81c05ef8c2af3f4fb7936ea8c18",
+    "fastplms/models/boltz/vb_layers_triangular_mult.py": "sha256:ace40becff882828ae3aa2b4f77e1a3d854d2adf7180216013996e38dd0e6fdd",
+    "fastplms/models/boltz/vb_loss_diffusionv2.py": "sha256:3778aea2ac541a3c173d9cc89b0d1b2cff63555571e0b3776b98c746daa0720c",
+    "fastplms/models/boltz/vb_modules_confidencev2.py": "sha256:a80a6f711e9ea021c45d89967fd37f7e668eea9ab209f9c70f2562d695a15041",
+    "fastplms/models/boltz/vb_modules_diffusion_conditioning.py": "sha256:a3e94f4493c056e54691b8289cbc35de3408b41bcbc7a75a8786eb1c273c1c31",
+    "fastplms/models/boltz/vb_modules_diffusionv2.py": "sha256:0077b40a72d9bf19904e0ab04fdb78b2c642ec6169b688732a1d221aa20a3637",
+    "fastplms/models/boltz/vb_modules_encodersv2.py": "sha256:9e78b92db67911795eb7691c7538edd12153ba30f4918ff8abf9874a99d02943",
+    "fastplms/models/boltz/vb_modules_transformersv2.py": "sha256:3eb5790030c333452beddeb3fe3583e65529a70f0e443d9d44304f7447de48f6",
+    "fastplms/models/boltz/vb_modules_trunkv2.py": "sha256:355cf735d57393167273b3be1d802ef495f69f8d20900e4cd2ce9b2df2147c25",
+    "fastplms/models/boltz/vb_modules_utils.py": "sha256:7a8f348cb8c703c6fef98312bea1beb179ceae5ae00eadc6962780a722171d05",
+    "fastplms/models/boltz/vb_potentials_potentials.py": "sha256:f332824228c5fabca51dd4695334e4a2822c6c36cca99dfc0c7f659eaf6062a3",
+    "fastplms/models/boltz/vb_potentials_schedules.py": "sha256:91997ee4b7a2a5ba1f40475c9e4fe6a58927f5ef7b51012743615b0f05e278d1",
+    "fastplms/models/boltz/vb_tri_attn_attention.py": "sha256:3c599f47575bfa762afd1b2fa70d6135ad66af72bfff8b703803d1f34f861080",
+    "fastplms/models/boltz/vb_tri_attn_primitives.py": "sha256:99b149ed2775b2d58baf4d5b4be414cfb2183234a0f46567ac6a37d9327909e7",
+    "fastplms/models/boltz/vb_tri_attn_utils.py": "sha256:df831665ee4bf485c5f95b3bc2a04dcb37aa002b03c42417036fc866e87622f1",
+    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
+    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
+    "fastplms_bundle.py": "sha256:f432c0047654b903f893db123760044e99b5a4caf176c62807a78c0d20c3fa33",
+    "modeling_fastplms.py": "sha256:29988e291b3f7a779d722a37d8acec6114c130fe389feae43fb6485e58a2cccf"
+  },
+  "model_id": "boltz2",
+  "redistributable": true,
+  "release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "runtime_bundle_sha256": "297af2aeeb817b6e3a343bf73dbc2eb3b3e29a2ca88ccba8af1aeb2e12243ffd",
+  "runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "schema_version": 2,
+  "scope": "runtime-only",
+  "source_tree_sha256": "21579947908f914e1b38ff1e69f9df9d22b7798899090355a53691a937b024e5",
+  "weights": {
+    "repo_id": "Synthyra/Boltz2",
+    "revision": "3b148fc5efea109c065ec82ba8683d024de7134e"
+  },
+  "weights_license_status": "resolved"
+}