Update FastPLMs runtime and model cards

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

README.md

@@ -16,16 +16,35 @@ Accepted inputs are amino-acid sequences tokenized to residue IDs.
 Supported Transformers entry points are `AutoConfig`, `AutoModel`,
 `AutoModelForMaskedLM`.
 
+## Capabilities
+
+| Feature | Status |
+| --- | --- |
+| Sequence classification | Unavailable: no advertised AutoClass |
+| Token classification | Unavailable: no advertised AutoClass |
+| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model |
+| Embeddings | Supported: shared ordered embedding API |
+| Test-time training | Supported: low-rank masked-residue adaptation |
+| Attention variants | Special: SDPA fidelity path; alternate backends have explicit bands |
+| Compliance | Declared: exact release evidence is required |
+
+A supported interface is not a pretrained downstream predictor. Classification
+heads start untrained, and declared compliance metadata is not a claim that an
+arbitrary local build passed its release gate.
+
 ## Install and platform requirements
 
-Install the current FastPLMs package:
+Install the direct dependencies published with this model:
 
 ```bash
-python -m pip install \
-  "fastplms @ git+https://github.com/Synthyra/FastPLMs.git"
+python -m pip install -r \
+  "https://huggingface.co/Synthyra/ESMplusplus_large/resolve/main/requirements.txt"
 ```
 
-Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. Eager, SDPA, and Flex use the core install. FlashAttention requires the `flash` extra, compatible CUDA hardware, and BF16 execution. The Hub quick start below requires network
+The FastPLMs implementation itself is embedded in the model repository and loaded
+by Transformers through `trust_remote_code=True`.
+
+Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network
 access on first download. For an air-gapped run, first build the manifest-pinned
 local artifact and use the offline form shown in the example.
 
@@ -38,23 +57,23 @@ model_id = "Synthyra/ESMplusplus_large"
 model = AutoModel.from_pretrained(
     model_id,
     trust_remote_code=True,
+    attn_implementation="sdpa",
 ).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/ESMplusplus_large` path, then pass `local_files_only=True`.
-
-Leave attention unspecified for the Transformers default. Supported explicit
-choices are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`,
-`flash_attention_3`.
-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 parameters loaded directly in BF16.
+For offline validation, replace `model_id` with the manifest-built
+`dist/hub/ESMplusplus_large` path and pass `local_files_only=True`.
+
+## Attention and compliance
+
+The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`,
+`flash_attention_3`. An unavailable requested backend raises instead of
+silently switching implementations.
+`output_attentions=True` may use the documented, one-call eager fallback solely
+to materialize attention tensors; the configured backend remains unchanged.
+
+This family declares the `compliance` tier. Release evidence binds the exact
+checkpoint, backend, dtype, hardware, inputs, and reference revision.
 
 ## Tokenization and forward inference
 
@@ -84,43 +103,78 @@ print(output.last_hidden_state.shape)
 
 ## Dataset embeddings
 
-The shared embedding API accepts sequences, `(id, sequence)` pairs,
-`EmbeddingInput` records, insertion-ordered `{id: sequence}` mappings, or a
-FASTA path. Results preserve order and duplicate identifiers:
+The shared embedding mixin preserves input order and biological-position
+masking. It accepts sequences, identified records, mappings, or a FASTA path:
 
 ```python
-result = model.embed_dataset(
+pooled = model.embed_dataset(
     ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"],
     batch_size=2,
     pooling=("mean", "std"),
 )
+residues = model.embed_dataset(
+    ["MSTNPKPQRKTKRNT"],
+    full_embeddings=True,
+)
+print(pooled[0].tensor.shape)   # (2 * d,)
+print(residues[0].tensor.shape) # (l, d)
+```
+
+Set `output` and `format="safetensors"` or `"sqlite"` for transactional,
+bounded-memory persistence. Resume verifies input order, model state, tokenizer
+policy, backend, dtype, and pooling configuration before appending.
+
+## PEFT fine-tuning
+
+Install the direct training dependencies, then attach LoRA to the loaded checkpoint:
+
+```bash
+python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
+```
 
-for record in result:
-    print(record.id, record.sequence, record.tensor.shape)
+```python
+from peft import LoraConfig, get_peft_model
+
+peft_model = get_peft_model(
+    model,
+    LoraConfig(
+        r=8,
+        lora_alpha=16,
+        target_modules="all-linear",
+    ),
+)
 ```
 
-Set `full_embeddings=True` for one residue tensor with shape `(l, d)` per
-sequence. Set `output` to a directory for bounded-memory, transactional
-safetensors with ordered-prefix resume, or choose `format="sqlite"` for
-batch-level database commits and exact resume. Pooling excludes boundary,
-padding, and other non-biological positions.
+This checkpoint has no advertised classifier. Supply the task-specific
+objective and preserve any new head through `modules_to_save`.
+All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and
+can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a
+support boundary. Record the target modules, base revision, data identity, and
+trainable parameter scope.
+
+## Test-time training
 
-For a long FASTA run, stream completed batches into SQLite:
+TTT samples masked views of one protein and updates only injected low-rank
+adapters. Base checkpoint weights remain frozen:
 
 ```python
-persisted = model.embed_dataset(
-    "proteins.fasta",
-    batch_size=64,
-    pooling=("mean",),
-    output="protein-embeddings.sqlite",
-    format="sqlite",
-    resume=True,
+from transformers import AutoModelForMaskedLM
+
+ttt_model = AutoModelForMaskedLM.from_pretrained(
+    "Synthyra/ESMplusplus_large",
+    trust_remote_code=True,
+)
+metrics = ttt_model.ttt(
+    seq="MSTNPKPQRKTKRNT",
+    ttt_config={"steps": 3, "batch_size": 1, "seed": 7},
 )
+ttt_model.save_pretrained("adapted", safe_serialization=True)
+ttt_model.ttt_reset()
+print(metrics)
 ```
 
-Resume verifies the input order, model state, tokenizer policy, backend, dtype,
-and pooling configuration. It never appends incompatible records to an
-existing run.
+Persisted adapters retain their deterministic reset state. TTT adds latency
+and memory, can worsen an output, and does not establish biological function.
 
 ## ESMC behavior
 
@@ -133,14 +187,9 @@ Those deviations produce diagnostic warnings rather than strict parity
 failures; dispatch integrity, masks, finite outputs, shapes, and catastrophic
 biological disagreement remain hard gates.
 
-The current locked GH200/aarch64 release image executes eager, SDPA, and Flex.
-It has no validated FlashAttention 2 kernel and no locked aarch64 artifact for
-the manifest-pinned FlashAttention 3 kernel. Both Flash backends remain
-supported, non-experimental interfaces, but requests fail closed without
-dispatch on this image. Prior real FlashAttention 2 execution was captured in
-separate workstation JUnit, but that immutable report and environment
-attestation are not bundled in this repository. It is not reused as a current
-ESMC release distribution or numerical claim.
+The current GH200/aarch64 release environment validates eager, SDPA, and Flex.
+Flash requests fail closed because compatible locked kernels are unavailable
+on this platform.
 
 When `sequence_id` is supplied, it is authoritative for ESMC attention grouping
 and padding, and `attention_mask` is ignored. Values greater than or equal to
@@ -149,48 +198,18 @@ use `attention_mask` as the padding contract.
 
 | Backend | Support | Measurement status |
 | --- | --- | --- |
-| `sdpa` | Recommended fidelity path | Pending complete validated 30-record frozen-head GH200/aarch64 set |
-| `eager` | Supported | Pending complete validated 30-record frozen-head GH200/aarch64 set |
-| `flash_attention_2` | Supported | Current GH200/aarch64 lock: unavailable; pending structured schema-v3 availability record |
-| `flex_attention` | Supported, numerically divergent | Pending complete validated 30-record frozen-head GH200/aarch64 set |
-| `flash_attention_3` | Supported, numerically divergent | Current GH200/aarch64 lock: unavailable; pending structured schema-v3 availability record |
-
-No threshold, report from another checkpoint, or result from another
-accelerator is substituted for a measurement. A release set contains all
-30 model/backend/panel records from one exact GH200 device and aarch64 runtime:
-18 eager/SDPA/Flex measurements include relative L2, Q99.9, residue
-cosine, pooled cosine, top-1, and Jensen-Shannon distributions; 12
-FlashAttention 2/3 records explicitly attest locked-platform unavailability.
-
-## Locked oracle package compatibility exception
-
-The frozen oracle lock permits exactly one nonzero `pip check` diagnostic:
-`nvidia-cusparselt-cu13 0.8.1 is not supported on this platform`. It applies only to
-`nvidia-cusparselt-cu13==0.8.1` on
-`NVIDIA GH200 480GB` / `linux` /
-`aarch64`. The vendor filename tag is
-`py3-none-manylinux2014_aarch64`, while the wheel metadata declares
-`py3-none-manylinux2014_sbsa`. The exact wheel is
-`nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl` with SHA-256 `4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f`.
-FastPLMs accepts this vendor metadata mismatch only after the lock, installed
-inventory, wheel bytes, metadata tag, and target identity all match. The wheel
-is not rewritten (`validated-vendor-metadata-exception-no-wheel-rewrite`). Any additional diagnostic or
-identity drift fails closed.
-
-
-No number is inferred from a threshold or copied from a different checkpoint.
-Before release, replace each pending cell only through the strict 30-record
-ingestion gate tied to this exact Hub revision, runtime revision, BF16 dtype,
-the current exact GH200/aarch64 accelerator and container identity, and both
-locked sequence panels. The set contains 18 eager/SDPA/Flex measurement records
-and 12 structured Flash locked-platform unavailable records. H100 and H200
-remain Hopper-class deployment examples, but they are not current ESMC release-confirmation evidence.
-Results are never
-carried across device, platform, source, lock, or image identities.
-Diagnostic JSON is written under `artifacts/diagnostics/esmc/`. Catastrophe
-guardrails (relative L2 `0.25`, Q99.9 `0.50`, residue cosine `0.90`, pooled
-cosine `0.95`, top-1 `0.80`, Jensen-Shannon `0.05`) detect corruption and do
-not constitute parity or quality claims.
+| `sdpa` | Recommended fidelity path | Pending release measurement |
+| `eager` | Supported | Pending release measurement |
+| `flash_attention_2` | Supported | Unavailable on current GH200/aarch64 lock |
+| `flex_attention` | Supported, numerically divergent | Pending release measurement |
+| `flash_attention_3` | Supported, numerically divergent | Unavailable on current GH200/aarch64 lock |
+
+Detailed backend measurements, release guardrails, and the GH200 package
+compatibility exception are maintained in the
+[attention backend guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/attention_backends.md)
+and
+[release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md).
+
 
 ## Runtime contract
 
@@ -201,7 +220,7 @@ not constitute parity or quality claims.
 - Precision policies: `default`
 - BF16 execution: `static_parameters`
 - Generation contract: `not_applicable`
-- Optional dependency group: `core`
+- Artifact dependency set: `core`
 - Weight publication allowed: `true`
 - Weight license status: `resolved`
 - Redistributable: `true`
@@ -212,28 +231,22 @@ not constitute parity or quality claims.
 - FastPLMs weights: `Synthyra/ESMplusplus_large`
 - Runtime revision: recorded separately in the built artifact and published commit
 - Source-tree and runtime-bundle SHA-256: recorded in `provenance.json`
-- Generator/schema version and complete/runtime-only attestations: recorded in `provenance.json`
 - Official checkpoint: `biohub/ESMC-600M`
 - Artifact source: `fast`
 - State transform: `esmc_to_fastplms_v1`
-- BF16 execution: `static_parameters`
 - Pinned upstreams: `biohub-esm`, `biohub-transformers`
-- Reference container: `reference-biohub-esm`
 - Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark`
 - Unresolved required file identities: `0`
 
-The local artifact records exact file identities, conversion details, source
-revisions, and legal texts in `provenance.json`. A nonzero unresolved count is a
-release blocker.
+`provenance.json` records exact file identities, conversion, source revisions,
+legal texts, schema, and attestations. A nonzero unresolved count blocks release.
 
 ## Validation boundary
 
-For tiers declared by the manifest, the release contract compares applicable
-semantic configuration, tokenizer behavior, state keys, shapes, dtypes,
-values, aliases, and representative inference with the pinned official
-implementation. This metadata does not by itself claim that a particular build
-passed, that one backend is faster, or that an output has biological or
-therapeutic validity.
+Declared tiers compare applicable configuration, tokenizer behavior, state,
+and representative inference with the pinned reference. Metadata alone does
+not claim a build passed, a backend is faster, or an output is biologically
+valid.
 
 ## License
 

config.json

@@ -13,11 +13,11 @@
   "fastplms_checkpoint_repo_id": "Synthyra/ESMplusplus_large",
   "fastplms_checkpoint_revision": "f813401638b3fddab09748aec1ad2bf537aa4208",
   "fastplms_model_id": "esmc_large",
-  "fastplms_release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "fastplms_runtime_bundle_sha256": "b01a730b69d0e25723da00aed2f25c2ede185e260e0ab1f7f29cc62934ab17d1",
-  "fastplms_runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_source_tree_sha256": "a95686f9e9cb31ebf4a3ced01c54787c6496336528fba0f696d407b053ebabc6",
+  "fastplms_release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "fastplms_runtime_bundle_sha256": "d73f544a7623d12eff2ba407e2d7e86417b7d731e36a2931515930eae7d7e8e5",
+  "fastplms_runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_source_tree_sha256": "bed9c7b1e1b0500dcd370e92e1f83b47494f9167d37b16e248f9d17e0135f822",
   "fastplms_weights_revision": "f813401638b3fddab09748aec1ad2bf537aa4208",
   "hidden_size": 1152,
   "initializer_range": 0.02,

fastplms/__init__.py

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

fastplms/attention/__init__.py

@@ -32,6 +32,7 @@ from .interfaces import (
     validate_transformers_attention_interfaces,
 )
 
+
 __all__ = [
     "FASTPLMS_ATTENTION_FUNCTIONS",
     "FASTPLMS_ATTENTION_MASKS",

fastplms/attention/_core.py

@@ -8,17 +8,17 @@ FastPLMs never downloads or compiles code.
 from __future__ import annotations
 
 import warnings
+import torch
 from collections import OrderedDict
 from collections.abc import Callable
 from enum import Enum
 from threading import RLock
-
-import torch
 from einops import rearrange
 from torch.nn import functional as F
 
 from ._kernel_lock import load_locked_kernel
 
+
 try:
     from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
 except ImportError:
@@ -117,12 +117,12 @@ def _get_flex_block_mask(
         raise RuntimeError(
             "'flex_attention' was requested, but torch.create_block_mask is unavailable."
         )
-    pattern = mask_pattern.detach().to(device=device).contiguous()
+    pattern = mask_pattern.detach().to(device=device).contiguous()  # mask_pattern.shape
     # One device-to-host transfer is required for an exact cache identity. Use
     # the contiguous buffer directly instead of materializing one Python int
     # per byte, which is prohibitively expensive for long batched sequences.
-    host_pattern = pattern.to(device="cpu").contiguous()
-    pattern_bytes = host_pattern.view(torch.uint8).numpy().tobytes(order="C")
+    host_pattern = pattern.to(device="cpu").contiguous()  # mask_pattern.shape
+    pattern_bytes = host_pattern.view(torch.uint8).numpy().tobytes(order="C")  # bytes
     cache_key = (
         str(device),
         None if dtype is None else str(dtype),
@@ -203,6 +203,7 @@ def _validate_kernels_flash_dtype(
 ) -> torch.dtype:
     """Reject dtypes outside the immutable kernel manifest before dispatch."""
 
+    # query_states, key_states, value_states: (b, l, h, d) or (t, h, d)
     tensor_dtypes = {query_states.dtype, key_states.dtype, value_states.dtype}
     if len(tensor_dtypes) != 1:
         observed = ", ".join(sorted(str(dtype) for dtype in tensor_dtypes))
@@ -243,6 +244,7 @@ def _validate_kernels_flash_device(
 ) -> torch.device:
     """Require Q, K, and V on one CUDA device before loading a kernel."""
 
+    # query_states, key_states, value_states: (b, l, h, d) or (t, h, d)
     devices = (query_states.device, key_states.device, value_states.device)
     if len(set(devices)) != 1:
         observed = ", ".join(str(device) for device in devices)
@@ -284,9 +286,10 @@ def _kernels_flash_forward(
     Failing to override when Q is pre-scaled applies the scale twice and breaks
     parity with eager attention and SDPA.
     """
+    # query_states, key_states, value_states: (b, l, h, d)
     flash_kernel, flash_kernel_variant = _ensure_flash_kernels_loaded(implementation)
     if flash_kernel_variant == "flash_attn2":
-        output = flash_kernel.flash_attn_func(
+        output = flash_kernel.flash_attn_func(  # (b, l, h, d) or tuple with that first
             q=query_states,
             k=key_states,
             v=value_states,
@@ -294,9 +297,9 @@ def _kernels_flash_forward(
             softmax_scale=softmax_scale,
             causal=causal,
         )
-        return output[0] if isinstance(output, tuple) else output
+        return output[0] if isinstance(output, tuple) else output  # (b, l, h, d)
     if flash_kernel_variant == "flash_attn3":
-        output = flash_kernel.flash_attn_func(
+        output = flash_kernel.flash_attn_func(  # (b, l, h, d) or tuple with that first
             q=query_states,
             k=key_states,
             v=value_states,
@@ -304,8 +307,8 @@ def _kernels_flash_forward(
             causal=causal,
         )
         if isinstance(output, tuple):
-            return output[0]
-        return output
+            return output[0]  # (b, l, h, d)
+        return output  # (b, l, h, d)
     raise RuntimeError(f"Unsupported FlashAttention kernel variant: {flash_kernel_variant}")
 
 
@@ -326,9 +329,11 @@ def _kernels_flash_varlen_forward(
     See `_kernels_flash_forward` docstring for why `softmax_scale=1.0` must be
     passed when Q has been pre-scaled by the caller.
     """
+    # query_states, key_states, value_states: (t, h, d)
+    # cu_seqlens_q, cu_seqlens_k: (b + 1,)
     flash_kernel, flash_kernel_variant = _ensure_flash_kernels_loaded(implementation)
     if flash_kernel_variant == "flash_attn2":
-        output = flash_kernel.flash_attn_varlen_func(
+        output = flash_kernel.flash_attn_varlen_func(  # (t, h, d) or tuple with that first
             q=query_states,
             k=key_states,
             v=value_states,
@@ -340,9 +345,9 @@ def _kernels_flash_varlen_forward(
             softmax_scale=softmax_scale,
             causal=causal,
         )
-        return output[0] if isinstance(output, tuple) else output
+        return output[0] if isinstance(output, tuple) else output  # (t, h, d)
     if flash_kernel_variant == "flash_attn3":
-        output = flash_kernel.flash_attn_varlen_func(
+        output = flash_kernel.flash_attn_varlen_func(  # (t, h, d) or tuple with that first
             q=query_states,
             k=key_states,
             v=value_states,
@@ -354,8 +359,8 @@ def _kernels_flash_varlen_forward(
             causal=causal,
         )
         if isinstance(output, tuple):
-            return output[0]
-        return output
+            return output[0]  # (t, h, d)
+        return output  # (t, h, d)
     raise RuntimeError(f"Unsupported FlashAttention kernel variant: {flash_kernel_variant}")
 
 
@@ -364,6 +369,7 @@ def _kernels_flash_varlen_forward(
 class IndexFirstAxis(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input, indices) -> torch.Tensor:
+        # input: (n, ...); indices: (m,)
         ctx.save_for_backward(indices)
         if input.ndim < 2:
             raise ValueError(
@@ -377,12 +383,13 @@ class IndexFirstAxis(torch.autograd.Function):
             )
         ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
         second_dim = other_shape.numel()
-        return torch.gather(
+        return torch.gather(  # (m, ...)
             rearrange(input, "b ... -> b (...)"), 0, indices.unsqueeze(1).expand(-1, second_dim)
         ).reshape(-1, *other_shape)
 
     @staticmethod
     def backward(ctx, grad_output) -> tuple[torch.Tensor, None]:
+        # grad_output: (m, ...)
         (indices,) = ctx.saved_tensors
         if grad_output.ndim < 2:
             raise RuntimeError(
@@ -390,19 +397,20 @@ class IndexFirstAxis(torch.autograd.Function):
                 "two dimensions."
             )
         other_shape = grad_output.shape[1:]
-        grad_output = rearrange(grad_output, "b ... -> b (...)")
-        grad_input = torch.zeros(
+        grad_output = rearrange(grad_output, "b ... -> b (...)")  # (m, product(...))
+        grad_input = torch.zeros(  # (n, product(...))
             [ctx.first_axis_dim, grad_output.shape[1]],
             device=grad_output.device,
             dtype=grad_output.dtype,
         )
         grad_input.scatter_(0, indices.unsqueeze(1).expand(-1, grad_output.shape[1]), grad_output)
-        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None  # (n, ...), None
 
 
 class IndexPutFirstAxis(torch.autograd.Function):
     @staticmethod
     def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor:
+        # values: (m, ...); indices: (m,)
         ctx.save_for_backward(indices)
         if indices.ndim != 1:
             raise ValueError(
@@ -414,16 +422,17 @@ class IndexPutFirstAxis(torch.autograd.Function):
                 "index_put_first_axis values must have at least two dimensions; "
                 f"received shape {tuple(values.shape)}."
             )
-        output = torch.zeros(
+        output = torch.zeros(  # (n, ...)
             first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
         )
         output[indices] = values
-        return output
+        return output  # (n, ...)
 
     @staticmethod
     def backward(ctx, grad_output) -> tuple[torch.Tensor, None, None]:
+        # grad_output: (n, ...)
         (indices,) = ctx.saved_tensors
-        return grad_output[indices], None, None
+        return grad_output[indices], None, None  # (m, ...), None, None
 
 
 index_first_axis = IndexFirstAxis.apply
@@ -433,8 +442,9 @@ index_put_first_axis = IndexPutFirstAxis.apply
 def pad_input(
     hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int
 ) -> torch.Tensor:
-    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
-    return rearrange(output, "(b s) ... -> b s ...", b=batch)
+    # hidden_states: (t, ...); indices: (t,)
+    output = index_put_first_axis(hidden_states, indices, batch * seqlen)  # (b * l, ...)
+    return rearrange(output, "(b s) ... -> b s ...", b=batch)  # (b, l, ...)
 
 
 def _unpad_input(
@@ -450,18 +460,19 @@ def _unpad_input(
     tuple[torch.Tensor, torch.Tensor],
     tuple[int, int],
 ]:
+    # query_layer, key_layer, value_layer: (b, l, h, d); attention_mask_2d: (b, l)
     batch_size, seq_len, num_heads, head_dim = query_layer.shape
-    seqlens = attention_mask_2d.sum(dim=1).int()
-    cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0))
+    seqlens = attention_mask_2d.sum(dim=1).int()  # (b,)
+    cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0))  # (b + 1,)
     max_seqlen = int(seqlens.max().item())
-    indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten()
-    query_layer = index_first_axis(
+    indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten()  # (t,)
+    query_layer = index_first_axis(  # (t, h, d)
         query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices
     )
-    key_layer = index_first_axis(
+    key_layer = index_first_axis(  # (t, h, d)
         key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices
     )
-    value_layer = index_first_axis(
+    value_layer = index_first_axis(  # (t, h, d)
         value_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices
     )
     return (
@@ -482,6 +493,7 @@ def _validate_flash_padding_mask(
 ) -> torch.Tensor:
     """Validate the self-attention padding mask used by the varlen kernels."""
 
+    # query_states, key_states, value_states: (b, l, h, d); attention_mask_2d: (b, l)
     if attention_mask_2d.ndim != 2:
         raise ValueError("FlashAttention padding masks must have shape (batch, sequence_length).")
     expected_shape = query_states.shape[:2]
@@ -497,7 +509,7 @@ def _validate_flash_padding_mask(
         )
     if attention_mask_2d.device != query_states.device:
         raise ValueError("FlashAttention padding mask and Q, K, and V must be on the same device.")
-    return attention_mask_2d.to(dtype=torch.bool)
+    return attention_mask_2d.to(dtype=torch.bool)  # (b, l)
 
 
 def kernels_flash_attention_func(
@@ -521,6 +533,8 @@ def kernels_flash_attention_func(
     `softmax_scale=1.0`. Otherwise the flash kernel applies its default scale
     again, yielding an effective `1/head_dim` scale that drifts across layers.
     """
+    # query_states, key_states, value_states: (b, l, h, d)
+    # attention_mask_2d: (b, l) or None
     _validate_kernels_flash_device(
         query_states,
         key_states,
@@ -534,11 +548,11 @@ def kernels_flash_attention_func(
         implementation,
     )
     if query_states.dtype != runtime_dtype:
-        query_states = query_states.to(dtype=runtime_dtype)
-        key_states = key_states.to(dtype=runtime_dtype)
-        value_states = value_states.to(dtype=runtime_dtype)
+        query_states = query_states.to(dtype=runtime_dtype)  # (b, l, h, d)
+        key_states = key_states.to(dtype=runtime_dtype)  # (b, l, h, d)
+        value_states = value_states.to(dtype=runtime_dtype)  # (b, l, h, d)
     if attention_mask_2d is not None:
-        attention_mask_2d = _validate_flash_padding_mask(
+        attention_mask_2d = _validate_flash_padding_mask(  # (b, l)
             query_states,
             key_states,
             value_states,
@@ -554,8 +568,13 @@ def kernels_flash_attention_func(
             indices_q,
             (cu_seqlens_q, cu_seqlens_k),
             (max_seqlen_q, max_seqlen_k),
-        ) = _unpad_input(query_states, key_states, value_states, attention_mask_2d)
-        attn_output_unpad = _kernels_flash_varlen_forward(
+        ) = _unpad_input(  # (t, h, d), (t, h, d), (t, h, d), (t,), (b + 1,), scalars
+            query_states,
+            key_states,
+            value_states,
+            attention_mask_2d,
+        )
+        attn_output_unpad = _kernels_flash_varlen_forward(  # (t, h, d)
             query_states=query_states,
             key_states=key_states,
             value_states=value_states,
@@ -567,10 +586,10 @@ def kernels_flash_attention_func(
             softmax_scale=softmax_scale,
             implementation=implementation,
         )
-        output = pad_input(attn_output_unpad, indices_q, batch_size, q_len)
-        return output.masked_fill(~attention_mask_2d[:, :, None, None], 0)
+        output = pad_input(attn_output_unpad, indices_q, batch_size, q_len)  # (b, l, h, d)
+        return output.masked_fill(~attention_mask_2d[:, :, None, None], 0)  # (b, l, h, d)
     else:
-        return _kernels_flash_forward(
+        return _kernels_flash_forward(  # (b, l, h, d)
             query_states=query_states,
             key_states=key_states,
             value_states=value_states,
@@ -706,6 +725,7 @@ def get_attention_mask(
 
     Returns (attention_mask_2d, attention_mask_4d, flex_block_mask).
     """
+    # attention_mask: (b, l) or None
     if attention_mask is None:
         return None, None, None
 
@@ -720,14 +740,14 @@ def get_attention_mask(
             "attention_mask shape must match the input batch and sequence dimensions; "
             f"expected {expected_shape}, received {tuple(attention_mask.shape)}."
         )
-    attention_mask_2d = attention_mask.to(device=device, dtype=torch.bool)
+    attention_mask_2d = attention_mask.to(device=device, dtype=torch.bool)  # (b, l)
     if not bool(attention_mask_2d.any(dim=1).all()):
         raise ValueError("attention_mask must keep at least one valid key per batch row.")
 
     effective_backend = resolve_attention_backend(effective_backend)
 
     if effective_backend.is_flash:
-        return attention_mask_2d, None, None
+        return attention_mask_2d, None, None  # (b, l), None, None
 
     if effective_backend == AttentionBackend.FLEX_ATTENTION:
         if create_block_mask is None:
@@ -751,12 +771,12 @@ def get_attention_mask(
             mask_semantics=mask_semantics,
             mask_mod=mask_mod,
         )
-        return attention_mask_2d, None, flex_block_mask
+        return attention_mask_2d, None, flex_block_mask  # (b, l), None, BlockMask
 
     # SDPA/manual masks only keys. Padding queries still attend to real keys, so
     # their outputs stay finite instead of softmaxing over all -inf scores.
-    attention_mask_4d = attention_mask_2d[:, None, None, :]
-    return attention_mask_2d, attention_mask_4d, None
+    attention_mask_4d = attention_mask_2d[:, None, None, :]  # (b, 1, 1, l)
+    return attention_mask_2d, attention_mask_4d, None  # (b, l), (b, 1, 1, l), None
 
 
 def bool_to_additive_mask(
@@ -770,10 +790,11 @@ def bool_to_additive_mask(
     That silently drops the mask. Always allocate a float tensor first, then fill it.
     This helper is the sanctioned way to build an SDPA additive mask from a bool validity mask.
     """
+    # bool_mask: (...)
     if bool_mask.dtype != torch.bool:
         raise TypeError(
             f"bool_to_additive_mask requires a bool tensor, got dtype={bool_mask.dtype}"
         )
-    additive = torch.zeros_like(bool_mask, dtype=dtype)
+    additive = torch.zeros_like(bool_mask, dtype=dtype)  # (...)
     additive.masked_fill_(bool_mask.logical_not(), float("-inf"))
-    return additive
+    return additive  # (...)

fastplms/attention/_kernel_lock.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import importlib.metadata
 import json
 import os
 from pathlib import Path
@@ -15,50 +14,35 @@ def require_kernels_package() -> None:
         import kernels  # noqa: F401
     except ImportError as error:
         raise RuntimeError(
-            "Precompiled FlashAttention requires the FastPLMs 'flash' extra."
+            "Precompiled FlashAttention requires requirements/features/flash.in."
         ) from error
 
 
 def _kernel_lock_path() -> Path:
-    """Return the lock from an artifact, checkout, or installed distribution."""
+    """Return the kernel lock from a Hub artifact or source checkout."""
     source_path = Path(__file__).resolve()
-    candidates = [
+    for candidate in (
         source_path.parents[1] / "kernels.lock",
         source_path.parents[3] / "kernels.lock",
-    ]
-    try:
-        import fastplms
-
-        candidates.extend(Path(root) / "kernels.lock" for root in fastplms.__path__)
-    except (ImportError, AttributeError):
-        pass
-    for candidate in candidates:
+    ):
         if candidate.is_file():
             return candidate
 
-    try:
-        distribution = importlib.metadata.distribution("fastplms")
-    except importlib.metadata.PackageNotFoundError as error:
-        raise RuntimeError("FastPLMs was installed without kernels.lock.") from error
-    for relative in distribution.files or ():
-        if relative.name != "kernels.lock":
-            continue
-        candidate = Path(distribution.locate_file(relative))
-        if candidate.is_file():
-            return candidate
-    raise RuntimeError("The installed FastPLMs distribution does not contain kernels.lock.")
+    raise RuntimeError(
+        "kernels.lock is missing from the Hugging Face artifact or source checkout."
+    )
 
 
 def _locked_entry(lock_path: Path, repository: str) -> dict[str, Any]:
     try:
-        data = json.loads(lock_path.read_text(encoding="utf-8"))
+        lock_entries = json.loads(lock_path.read_text(encoding="utf-8"))
     except (OSError, json.JSONDecodeError) as error:
-        raise RuntimeError(f"Unable to read the packaged kernel lock: {lock_path}") from error
-    if not isinstance(data, list):
+        raise RuntimeError(f"Unable to read the kernel lock: {lock_path}") from error
+    if not isinstance(lock_entries, list):
         raise RuntimeError("kernels.lock must contain a JSON list.")
-    if any(not isinstance(entry, dict) for entry in data):
+    if any(not isinstance(entry, dict) for entry in lock_entries):
         raise RuntimeError("Every kernels.lock entry must be a JSON object.")
-    matches = [entry for entry in data if entry.get("repo_id") == repository]
+    matches = [entry for entry in lock_entries if entry.get("repo_id") == repository]
     if len(matches) != 1:
         raise RuntimeError(
             f"kernels.lock must contain exactly one entry for {repository!r}; found {len(matches)}."
@@ -125,11 +109,11 @@ def _load_offline_locked_kernel(
         from kernels.variants import get_variants_local, resolve_variants
     except ImportError as error:
         raise RuntimeError(
-            "Precompiled FlashAttention requires the FastPLMs 'flash' extra."
+            "Precompiled FlashAttention requires requirements/features/flash.in."
         ) from error
 
-    parsed = get_variants_local(build_root)
-    parsed_names = {variant.variant_str for variant in parsed}
+    cached_variants = get_variants_local(build_root)
+    parsed_names = {variant.variant_str for variant in cached_variants}
     invalid = sorted(set(cached_names).difference(parsed_names))
     if invalid:
         raise RuntimeError(
@@ -137,7 +121,7 @@ def _load_offline_locked_kernel(
             f"{', '.join(invalid)}"
         )
 
-    compatible, _ = resolve_variants(parsed)
+    compatible, _ = resolve_variants(cached_variants)
     if len(compatible) != 1:
         names = ", ".join(variant.variant_str for variant in compatible) or "none"
         raise RuntimeError(
@@ -165,7 +149,7 @@ def load_locked_kernel(repository: str, revision: str) -> object:
         from kernels.lockfile import KernelLock
     except ImportError as error:
         raise RuntimeError(
-            "Precompiled FlashAttention requires the FastPLMs 'flash' extra."
+            "Precompiled FlashAttention requires requirements/features/flash.in."
         ) from error
 
     lock_path = _kernel_lock_path()

fastplms/attention/interfaces.py

@@ -2,11 +2,10 @@
 
 from __future__ import annotations
 
+import torch
 from collections.abc import Mapping
 from functools import partial
 from typing import Any
-
-import torch
 from transformers import AttentionInterface, AttentionMaskInterface
 
 from ._core import (
@@ -36,6 +35,7 @@ def _kernels_attention_forward(
     FastPLMs kernel adapter uses the latter layout internally.
     """
 
+    # query, key, value: (b, h, l, d); attention_mask: (b, l) or None
     dropout = float(kwargs.get("dropout", 0.0) or 0.0)
     if module.training and dropout:
         raise RuntimeError(
@@ -45,15 +45,15 @@ def _kernels_attention_forward(
     causal = bool(kwargs.get("is_causal", getattr(module, "is_causal", False)))
     softmax_scale = kwargs.get("scaling")
     output = kernels_flash_attention_func(
-        query_states=query.transpose(1, 2).contiguous(),
-        key_states=key.transpose(1, 2).contiguous(),
-        value_states=value.transpose(1, 2).contiguous(),
+        query_states=query.transpose(1, 2).contiguous(),  # (b, l, h, d)
+        key_states=key.transpose(1, 2).contiguous(),  # (b, l, h, d)
+        value_states=value.transpose(1, 2).contiguous(),  # (b, l, h, d)
         attention_mask_2d=attention_mask,
         causal=causal,
         softmax_scale=softmax_scale,
         implementation=implementation,
-    )
-    return output, None
+    )  # (b, l, h, d)
+    return output, None  # (b, l, h, d), None
 
 
 # Keep FastPLMs' kernels-only adapters local to this registry instance.

fastplms/embeddings/__init__.py

@@ -33,6 +33,7 @@ from .types import (
     TensorValue,
 )
 
+
 __all__ = [
     "DEFAULT_SHARD_SIZE",
     "POOLING_NAMES",

fastplms/embeddings/pooling.py

@@ -3,15 +3,16 @@
 from __future__ import annotations
 
 import math
-from collections.abc import Sequence
-
 import torch
+from collections.abc import Sequence
 from torch import Tensor
 
+
 POOLING_NAMES = frozenset({"mean", "max", "norm", "median", "std", "var", "cls", "parti"})
 
 
 def _validate_inputs(X: Tensor, M: Tensor) -> Tensor:
+    # X: (b, l, d); M: (b, l)
     if not isinstance(X, Tensor) or not isinstance(M, Tensor):
         raise TypeError("X and M must be tensors.")
     if X.ndim != 3:
@@ -24,12 +25,12 @@ def _validate_inputs(X: Tensor, M: Tensor) -> Tensor:
         raise TypeError("M must be a boolean or binary numeric residue mask.")
     if not bool(torch.isfinite(M).all()) or not bool(((M == 0) | (M == 1)).all()):
         raise ValueError("M must contain only finite binary mask values.")
-    M = M.to(device=X.device, dtype=torch.bool)
+    M = M.to(device=X.device, dtype=torch.bool)  # (b, l)
     if not bool(M.any(dim=1).all()):
         raise ValueError("Every sample must contain at least one biological residue.")
     if not bool((torch.isfinite(X) | ~M.unsqueeze(-1)).all()):
         raise ValueError("Biological residue embeddings produced non-finite output.")
-    return M
+    return M  # (b, l)
 
 
 def _pooled_attention(attentions: Tensor | Sequence[Tensor], *, batch_size: int) -> Tensor:
@@ -43,27 +44,27 @@ def _pooled_attention(attentions: Tensor | Sequence[Tensor], *, batch_size: int)
     if isinstance(attentions, Sequence):
         if not attentions:
             raise ValueError("parti received an empty attention sequence.")
-        # Each A_i has shape (b, h, l, l).
-        A = torch.stack(tuple(attentions), dim=1)
+        # Each A_i: (b, h, l, l).
+        A = torch.stack(tuple(attentions), dim=1)  # (b, n, h, l, l)
     else:
-        A = attentions
+        A = attentions  # (b, ..., l, l)
 
     if A.ndim == 5:
         if A.shape[0] != batch_size and A.shape[1] == batch_size:
-            A = A.transpose(0, 1)
+            A = A.transpose(0, 1)  # (b, n, h, l, l)
         if A.shape[0] != batch_size:
             raise ValueError("Five-dimensional attentions must use (b, n, h, l, l).")
-        A = A.flatten(1, 2).amax(dim=1)
+        A = A.flatten(1, 2).amax(dim=1)  # (b, l, l)
     elif A.ndim == 4:
         if A.shape[0] != batch_size:
             raise ValueError("Four-dimensional attentions must use (b, h, l, l).")
-        A = A.amax(dim=1)
+        A = A.amax(dim=1)  # (b, l, l)
     elif A.ndim == 3:
         if A.shape[0] != batch_size:
             raise ValueError("Three-dimensional attentions must use (b, l, l).")
     else:
         raise ValueError("Attentions must have shape (b, l, l), (b, h, l, l), or (b, n, h, l, l).")
-    return A
+    return A  # (b, l, l)
 
 
 def pagerank_weights(
@@ -79,6 +80,7 @@ def pagerank_weights(
     probabilities; dangling rows transition uniformly.
     """
 
+    # A: (l, l)
     if not isinstance(A, Tensor):
         raise TypeError("A must be a tensor.")
     if A.ndim != 2 or A.shape[0] != A.shape[1]:
@@ -103,19 +105,23 @@ def pagerank_weights(
     if not bool(torch.isfinite(A).all()):
         raise ValueError("A must contain only finite attention values.")
     work_dtype = torch.float64 if A.dtype == torch.float64 else torch.float32
-    P = A.detach().to(dtype=work_dtype).clamp_min(0)
-    row_sum = P.sum(dim=-1, keepdim=True)
-    uniform = torch.full_like(P, 1.0 / length)
-    P = torch.where(row_sum > 0, P / row_sum.clamp_min(torch.finfo(work_dtype).tiny), uniform)
-    p = torch.full((length,), 1.0 / length, device=P.device, dtype=work_dtype)
+    P = A.detach().to(dtype=work_dtype).clamp_min(0)  # (l, l)
+    row_sum = P.sum(dim=-1, keepdim=True)  # (l, 1)
+    uniform = torch.full_like(P, 1.0 / length)  # (l, l)
+    P = torch.where(  # (l, l)
+        row_sum > 0,
+        P / row_sum.clamp_min(torch.finfo(work_dtype).tiny),
+        uniform,
+    )
+    p = torch.full((length,), 1.0 / length, device=P.device, dtype=work_dtype)  # (l,)
     teleport = (1.0 - damping) / length
     for _ in range(max_iterations):
-        p_next = teleport + damping * (P.transpose(0, 1) @ p)
+        p_next = teleport + damping * (P.transpose(0, 1) @ p)  # (l,)
         if torch.linalg.vector_norm(p_next - p, ord=1) <= tolerance:
-            p = p_next
+            p = p_next  # (l,)
             break
-        p = p_next
-    return p / p.sum()
+        p = p_next  # (l,)
+    return p / p.sum()  # (l,)
 
 
 class Pooler:
@@ -157,28 +163,29 @@ class Pooler:
         attentions: Tensor | Sequence[Tensor] | None = None,
         attention_backend: str | None = None,
     ) -> Tensor:
-        M = _validate_inputs(X, residue_mask)
-        M_expanded = M.unsqueeze(-1)
-        count = M_expanded.sum(dim=1).clamp_min(1)
-        X_residues = X.masked_fill(~M_expanded, 0)
+        # X: (b, l, d); residue_mask: (b, l)
+        M = _validate_inputs(X, residue_mask)  # (b, l)
+        M_expanded = M.unsqueeze(-1)  # (b, l, 1)
+        count = M_expanded.sum(dim=1).clamp_min(1)  # (b, 1)
+        X_residues = X.masked_fill(~M_expanded, 0)  # (b, l, d)
         outputs: list[Tensor] = []
 
         for name in self.names:
             if name == "mean":
-                Y = X_residues.sum(dim=1) / count
+                Y = X_residues.sum(dim=1) / count  # (b, d)
             elif name == "max":
-                Y = X.masked_fill(~M_expanded, -torch.inf).max(dim=1).values
+                Y = X.masked_fill(~M_expanded, -torch.inf).max(dim=1).values  # (b, d)
             elif name == "norm":
-                Y = torch.linalg.vector_norm(X_residues, ord=2, dim=1)
+                Y = torch.linalg.vector_norm(X_residues, ord=2, dim=1)  # (b, d)
             elif name == "median":
-                Y = X.masked_fill(~M_expanded, torch.nan).nanmedian(dim=1).values
+                Y = X.masked_fill(~M_expanded, torch.nan).nanmedian(dim=1).values  # (b, d)
             elif name in {"var", "std"}:
-                mean = X_residues.sum(dim=1, keepdim=True) / count.unsqueeze(1)
-                centered = (X - mean).masked_fill(~M_expanded, 0)
-                variance = (centered**2).sum(dim=1) / count
-                Y = variance.sqrt() if name == "std" else variance
+                mean = X_residues.sum(dim=1, keepdim=True) / count.unsqueeze(1)  # (b, 1, d)
+                centered = (X - mean).masked_fill(~M_expanded, 0)  # (b, l, d)
+                variance = (centered**2).sum(dim=1) / count  # (b, d)
+                Y = variance.sqrt() if name == "std" else variance  # (b, d)
             elif name == "cls":
-                Y = X[:, 0]
+                Y = X[:, 0]  # (b, d)
             else:
                 if attention_backend != "eager":
                     raise ValueError(
@@ -189,14 +196,15 @@ class Pooler:
                     raise ValueError("parti requires model attention matrices.")
                 if int(M.sum(dim=1).max().item()) > 2048:
                     raise ValueError("parti supports at most 2,048 biological residues.")
-                A = _pooled_attention(attentions, batch_size=X.shape[0]).to(X.device)
+                A = _pooled_attention(attentions, batch_size=X.shape[0]).to(X.device)  # (b, l, l)
                 pooled: list[Tensor] = []
                 for X_i, M_i, A_i in zip(X, M, A, strict=True):
-                    indices = M_i.nonzero(as_tuple=True)[0]
-                    A_residue = A_i.index_select(0, indices).index_select(1, indices)
-                    w = pagerank_weights(A_residue).to(dtype=X.dtype)
-                    pooled.append(w @ X_i.index_select(0, indices))
-                Y = torch.stack(pooled)
+                    # X_i: (l, d); M_i: (l,); A_i: (l, l)
+                    indices = M_i.nonzero(as_tuple=True)[0]  # (r,)
+                    A_residue = A_i.index_select(0, indices).index_select(1, indices)  # (r, r)
+                    w = pagerank_weights(A_residue).to(dtype=X.dtype)  # (r,)
+                    pooled.append(w @ X_i.index_select(0, indices))  # (d,)
+                Y = torch.stack(pooled)  # (b, d)
             if not bool(torch.isfinite(Y).all()):
                 raise ValueError(
                     f"Pooling operation {name!r} produced non-finite output from "
@@ -204,7 +212,7 @@ class Pooler:
                 )
             outputs.append(Y)
 
-        return torch.cat(outputs, dim=-1)
+        return torch.cat(outputs, dim=-1)  # (b, len(self.names) * d)
 
 
 __all__ = ["POOLING_NAMES", "Pooler", "pagerank_weights"]

fastplms/embeddings/runner.py

@@ -7,12 +7,11 @@ import json
 import platform
 import sqlite3
 import tempfile
+import torch
 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Any, overload
-
-import torch
 from torch import Tensor
 
 from .pooling import Pooler
@@ -35,6 +34,7 @@ from .types import (
     LazyTensorReference,
 )
 
+
 _MAX_PARTI_RESIDUES = 2_048
 _RUN_FINGERPRINT_SCHEMA_VERSION = 3
 _MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2
@@ -45,6 +45,7 @@ _SUPPORTED_STORAGE_FORMATS = frozenset({"safetensors", "sqlite"})
 def _validate_parti_length(M: Tensor) -> None:
     """Reject an oversized attention graph before model inference."""
 
+    # M: (b, l)
     n_residues = int(M.to(dtype=torch.int64).sum(dim=1).max().item())
     if n_residues > _MAX_PARTI_RESIDUES:
         raise ValueError(f"parti supports at most {_MAX_PARTI_RESIDUES:,} biological residues.")
@@ -58,16 +59,17 @@ def select_hidden_state_embeddings(
     store_all_hidden_states: bool = False,
 ) -> Tensor:
     """Select one hidden state or stack every state without changing values."""
+    # last_hidden_state and each hidden_states entry: (b, l, d)
     if store_all_hidden_states:
         if not hidden_states:
             raise ValueError("store_all_hidden_states requires model hidden states.")
         # H has shape (b, n, l, d), where n follows the model's output order.
-        return torch.stack(hidden_states, dim=1)
+        return torch.stack(hidden_states, dim=1)  # (b, n, l, d)
     if hidden_state_index == -1:
-        return last_hidden_state
+        return last_hidden_state  # (b, l, d)
     if not hidden_states:
         raise ValueError("hidden_state_index requires model hidden states.")
-    return hidden_states[hidden_state_index]
+    return hidden_states[hidden_state_index]  # (b, l, d)
 
 
 def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]:
@@ -508,12 +510,17 @@ def _biological_residue_mask(
 ) -> Tensor:
     """Remove padding and tokenizer-declared special tokens from M."""
 
-    M = attention_mask.to(dtype=torch.bool)
+    # input_ids, attention_mask: (b, l)
+    M = attention_mask.to(dtype=torch.bool)  # (b, l)
     special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ()))
     if special_ids:
-        specials = torch.tensor(special_ids, device=input_ids.device, dtype=input_ids.dtype)
-        M = M & ~torch.isin(input_ids, specials)
-    return M
+        specials = torch.tensor(  # (n_special,)
+            special_ids,
+            device=input_ids.device,
+            dtype=input_ids.dtype,
+        )
+        M = M & ~torch.isin(input_ids, specials)  # (b, l)
+    return M  # (b, l)
 
 
 def _generic_embedding_batch(
@@ -535,20 +542,23 @@ def _generic_embedding_batch(
         output = model._embed(sequences, return_attention_mask=True, **model_kwargs)
         if not isinstance(output, tuple) or len(output) != 2:
             raise TypeError("E1 _embed must return (X, residue_mask).")
-        X, M = output
+        X, M = output  # (b, l, d), (b, l)
         preparer = getattr(model, "prep_tokens", None)
         if preparer is not None and hasattr(preparer, "get_batch_kwargs"):
             prepared = preparer.get_batch_kwargs(sequences, device=X.device)
-            input_ids = prepared["input_ids"]
-            boundary_ids = preparer.boundary_token_ids.to(
+            input_ids = prepared["input_ids"]  # (b, l)
+            boundary_ids = preparer.boundary_token_ids.to(  # (n_boundary,)
                 device=input_ids.device, dtype=input_ids.dtype
             )
             # E1 wraps each raw sequence in BOS, context-label, terminal-label,
             # and EOS tokens. Only amino-acid rows are biological residues.
-            M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids)
+            M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids)  # (b, l)
         if need_attentions:
             raise ValueError("parti is not available for tokenizer-free E1 embedding.")
-        return EmbeddingBatch(X=X, residue_mask=M.to(dtype=torch.bool))
+        return EmbeddingBatch(  # X: (b, l, d); residue_mask: (b, l)
+            X=X,
+            residue_mask=M.to(dtype=torch.bool),
+        )
     if tokenizer is None:
         raise ValueError("A tokenizer is required for this model's embedding path.")
 
@@ -572,14 +582,17 @@ def _generic_embedding_batch(
     else:
         encoded = tokenizer(sequences, **tokenize_kwargs)
     device = _model_device(model)
-    input_ids = encoded["input_ids"].to(device)
-    attention_mask = encoded.get("attention_mask", input_ids.new_ones(input_ids.shape)).to(device)
-    M = _biological_residue_mask(input_ids, attention_mask, tokenizer)
+    input_ids = encoded["input_ids"].to(device)  # (b, l)
+    attention_mask = encoded.get(  # (b, l)
+        "attention_mask",
+        input_ids.new_ones(input_ids.shape),
+    ).to(device)
+    M = _biological_residue_mask(input_ids, attention_mask, tokenizer)  # (b, l)
     if need_attentions:
         # Validate l before either the backbone or its quadratic attention graph
         # is materialized. M has shape (b, l).
         _validate_parti_length(M)
-    X = model._embed(input_ids, attention_mask, **model_kwargs)
+    X = model._embed(input_ids, attention_mask, **model_kwargs)  # (b, l, d)
     attentions = None
     if need_attentions:
         output = model(
@@ -588,10 +601,14 @@ def _generic_embedding_batch(
             output_attentions=True,
             return_dict=True,
         )
-        attentions = getattr(output, "attentions", None)
+        attentions = getattr(output, "attentions", None)  # each: (b, h, l, l)
         if attentions is None:
             raise ValueError("The model did not return attentions required by parti.")
-    return EmbeddingBatch(X=X, residue_mask=M, attentions=attentions)
+    return EmbeddingBatch(  # X: (b, l, d); M: (b, l)
+        X=X,
+        residue_mask=M,
+        attentions=attentions,
+    )
 
 
 def _first_metadata_value(*values: Any) -> Any:
@@ -638,6 +655,7 @@ def _model_identity_metadata(model: Any) -> dict[str, Any]:
 def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]:
     """Yield X in logical row-major order without materializing a full copy."""
 
+    # X: (...)
     if X.numel() == 0:
         return
     if X.ndim == 0:
@@ -649,7 +667,7 @@ def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]:
     if trailing_elements <= max_elements:
         rows_per_chunk = max(1, max_elements // trailing_elements)
         for start in range(0, X.shape[0], rows_per_chunk):
-            yield X[start : start + rows_per_chunk]
+            yield X[start : start + rows_per_chunk]  # (chunk_rows, ...)
         return
     for row in X:
         yield from _bounded_tensor_chunks(row, max_elements)
@@ -685,7 +703,7 @@ def _model_state_sha256(model: Any) -> str:
         digest.update(header)
         max_elements = max(1, _MODEL_STATE_HASH_CHUNK_BYTES // value.element_size())
         for chunk in _bounded_tensor_chunks(value.detach(), max_elements):
-            cpu_chunk = chunk.to(device="cpu").contiguous()
+            cpu_chunk = chunk.to(device="cpu").contiguous()  # chunk.shape
             digest.update(cpu_chunk.reshape(-1).view(torch.uint8).numpy().tobytes())
     return digest.hexdigest()
 
@@ -1306,22 +1324,24 @@ def embed_dataset(
                             normalized_decoder_inputs[position] for position in batch_positions
                         ]
                     if decoder_input_ids is not None:
-                        indices = torch.tensor(
+                        # decoder_input_ids: (n_records, l_decoder)
+                        indices = torch.tensor(  # (b,)
                             batch_positions,
                             device=decoder_input_ids.device,
                             dtype=torch.long,
                         )
-                        batch_model_kwargs["decoder_input_ids"] = decoder_input_ids.index_select(
-                            0, indices
+                        batch_model_kwargs["decoder_input_ids"] = (  # (b, l_decoder)
+                            decoder_input_ids.index_select(0, indices)
                         )
                     if decoder_attention_mask is not None:
-                        indices = torch.tensor(
+                        # decoder_attention_mask: (n_records, l_decoder)
+                        indices = torch.tensor(  # (b,)
                             batch_positions,
                             device=decoder_attention_mask.device,
                             dtype=torch.long,
                         )
                         batch_model_kwargs["decoder_attention_mask"] = (
-                            decoder_attention_mask.index_select(0, indices)
+                            decoder_attention_mask.index_select(0, indices)  # (b, l_decoder)
                         )
                 custom_batch = _embedding_batch_fn or getattr(model, "_embedding_batch", None)
                 if custom_batch is not None:
@@ -1348,8 +1368,8 @@ def embed_dataset(
                         need_attentions=need_attentions,
                         model_kwargs=batch_model_kwargs,
                     )
-                X = batch.X
-                raw_mask = batch.residue_mask
+                X = batch.X  # (b, l, d) or (b, n_states, l, d)
+                raw_mask = batch.residue_mask  # (b, l)
                 if not isinstance(X, Tensor) or not isinstance(raw_mask, Tensor):
                     raise TypeError("Embedding batches must provide Tensor X and residue_mask.")
                 if X.is_meta or raw_mask.is_meta:
@@ -1360,7 +1380,7 @@ def embed_dataset(
                     raise ValueError("Embedding residue_mask must contain finite binary values.")
                 if not bool(((raw_mask == 0) | (raw_mask == 1)).all()):
                     raise ValueError("Embedding residue_mask must contain finite binary values.")
-                M = raw_mask.to(device=X.device, dtype=torch.bool)
+                M = raw_mask.to(device=X.device, dtype=torch.bool)  # (b, l)
                 valid_X_shape = (
                     X.ndim == 3
                     and X.shape[0] == len(batch_records)
@@ -1384,7 +1404,7 @@ def embed_dataset(
                     )
                 if not bool(M.any(dim=1).all()):
                     raise ValueError("Every embedding sample must contain a biological residue.")
-                finite_selected = (
+                finite_selected = (  # X.shape
                     torch.isfinite(X) | ~M.unsqueeze(-1)
                     if X.ndim == 3
                     else torch.isfinite(X) | ~M[:, None, :, None]
@@ -1395,28 +1415,32 @@ def embed_dataset(
                     # Validate the biological graph only after mask integrity is established.
                     _validate_parti_length(M)
                 if dtype is not None:
-                    X = X.to(dtype=dtype)
+                    X = X.to(dtype=dtype)  # unchanged shape
 
                 if full_embeddings:
                     if X.ndim == 4:
                         values = [
-                            X_i[:, M_i, :].detach().cpu() for X_i, M_i in zip(X, M, strict=True)
+                            X_i[:, M_i, :].detach().cpu()  # (n_states, r_i, d)
+                            for X_i, M_i in zip(X, M, strict=True)
                         ]
                     else:
-                        values = [X_i[M_i].detach().cpu() for X_i, M_i in zip(X, M, strict=True)]
+                        values = [
+                            X_i[M_i].detach().cpu()  # (r_i, d)
+                            for X_i, M_i in zip(X, M, strict=True)
+                        ]
                 else:
                     if pooler is None:
                         raise RuntimeError(
                             "Pooled embedding output was requested without an initialized pooler."
                         )
-                    Y = pooler(
+                    Y = pooler(  # (b, n_poolers * d)
                         X,
                         M,
                         attentions=batch.attentions,
                         attention_backend=attention_backend,
                     )
                     pool_slices = pooler.output_slices(X.shape[-1])
-                    values = list(Y.detach().cpu().unbind(0))
+                    values = list(Y.detach().cpu().unbind(0))  # each: (n_poolers * d,)
                 for position, record, value in zip(
                     batch_positions, batch_records, values, strict=True
                 ):

fastplms/embeddings/storage.py

@@ -7,14 +7,13 @@ import io
 import json
 import sqlite3
 import struct
+import numpy as np
+import torch
 from bisect import bisect_right
 from collections.abc import Iterable, Iterator, Sequence
 from pathlib import Path
 from typing import Any, cast, overload
 from uuid import uuid4
-
-import numpy as np
-import torch
 from torch import Tensor
 
 from .types import (
@@ -23,6 +22,7 @@ from .types import (
     LazyTensorReference,
 )
 
+
 _DTYPE_NAMES: dict[torch.dtype, str] = {
     torch.float16: "float16",
     torch.bfloat16: "bfloat16",
@@ -80,22 +80,24 @@ def _persistent_metadata(
 def _tensor_bytes(X: Tensor) -> bytes:
     """Return the exact contiguous byte representation of X."""
 
-    X = X.detach().cpu().contiguous()
+    # X: (...)
+    X = X.detach().cpu().contiguous()  # (...)
     return X.view(torch.uint8).numpy().tobytes()
 
 
 def _bounded_tensor_chunks(X: Tensor, max_bytes: int) -> Iterator[Tensor]:
     """Yield row-major CPU chunks without materializing one full byte string."""
 
-    flattened = X.detach().to(device="cpu").reshape(-1)
+    # X: (...)
+    flattened = X.detach().to(device="cpu").reshape(-1)  # (n,)
     if flattened.numel() == 0:
         return
     chunk_elements = max(1, max_bytes // flattened.element_size())
     for start in range(0, flattened.numel(), chunk_elements):
-        chunk = flattened[start : start + chunk_elements]
+        chunk = flattened[start : start + chunk_elements]  # (n_chunk,)
         if chunk.stride(0) != 1:
-            chunk = chunk.clone(memory_format=torch.contiguous_format)
-        yield chunk
+            chunk = chunk.clone(memory_format=torch.contiguous_format)  # (n_chunk,)
+        yield chunk  # (n_chunk,)
 
 
 def _tensor_hash_chunks(X: Tensor) -> Iterator[bytes]:
@@ -136,9 +138,9 @@ def _decode_tensor(dtype_name: str, shape_json: str, data: bytes) -> Tensor:
         raise ValueError(f"Unsupported stored dtype {dtype_name!r}.") from error
     shape = tuple(json.loads(shape_json))
     # uint8 is used only as a byte-level carrier, preserving BF16 bits exactly.
-    byte_array = np.frombuffer(data, dtype=np.uint8).copy()
-    X = torch.from_numpy(byte_array).view(dtype)
-    return X.reshape(shape).clone()
+    byte_array = np.frombuffer(data, dtype=np.uint8).copy()  # (n_bytes,)
+    X = torch.from_numpy(byte_array).view(dtype)  # (n_elements,)
+    return X.reshape(shape).clone()  # shape
 
 
 def _index_path(path: str | Path) -> Path:
@@ -651,7 +653,7 @@ class SafetensorsStreamWriter:
 
         for record in records:
             position = self._record_count + len(self._pending)
-            tensor = record.load_tensor().detach().cpu().contiguous()
+            tensor = record.load_tensor().detach().cpu().contiguous()  # (...)
             if tensor.dtype not in _DTYPE_NAMES:
                 raise TypeError(f"Unsupported tensor dtype {tensor.dtype}.")
             nbytes = tensor.numel() * tensor.element_size()
@@ -949,7 +951,7 @@ def save_sqlite_result(result: EmbeddingResult, path: str | Path) -> EmbeddingRe
             (run_id, metadata_json),
         )
         for position, record in enumerate(result):
-            X = record.load_tensor().detach().cpu().contiguous()
+            X = record.load_tensor().detach().cpu().contiguous()  # (...)
             dtype_name, shape_json, data = _encode_tensor(X)
             digest = tensor_sha256(X)
             connection.execute(
@@ -1053,7 +1055,7 @@ def append_sqlite_records(
             )
         for offset, record in enumerate(records):
             position = start_position + offset
-            X = record.load_tensor().detach().cpu().contiguous()
+            X = record.load_tensor().detach().cpu().contiguous()  # (...)
             dtype_name, shape_json, data = _encode_tensor(X)
             digest = tensor_sha256(X)
             connection.execute(
@@ -1414,6 +1416,7 @@ def load_legacy_pth(path: str | Path, *, allow_unsafe_pickle: bool = False) -> E
         raise ValueError("A legacy .pth embedding file must contain a mapping.")
     records: list[EmbeddingRecord] = []
     for position, (sequence, X) in enumerate(payload.items()):
+        # X: (...)
         if not isinstance(sequence, str) or not isinstance(X, Tensor):
             raise ValueError("Legacy embedding mappings must use str keys and Tensor values.")
         records.append(EmbeddingRecord(str(position), sequence, X.detach().cpu()))
@@ -1450,8 +1453,10 @@ def _decode_legacy_sqlite_blob(
         expected = int(np.prod(shape, dtype=np.int64)) * numpy_dtype.itemsize
         if len(data) - offset != expected:
             raise ValueError("Legacy compact embedding payload length does not match shape.")
-        array = np.frombuffer(data, dtype=numpy_dtype, offset=offset).copy().reshape(shape)
-        return torch.from_numpy(array).to(dtype=target_dtype)
+        array = (  # shape
+            np.frombuffer(data, dtype=numpy_dtype, offset=offset).copy().reshape(shape)
+        )
+        return torch.from_numpy(array).to(dtype=target_dtype)  # shape
 
     try:
         loaded = torch.load(io.BytesIO(data), map_location="cpu", weights_only=True)
@@ -1470,11 +1475,13 @@ def _decode_legacy_sqlite_blob(
                 raise ValueError(
                     "Legacy raw FP32 payload length does not match fallback_shape."
                 ) from safe_error
-            array = np.frombuffer(data, dtype=np.float32).copy().reshape(fallback_shape)
-            return torch.from_numpy(array)
+            array = np.frombuffer(data, dtype=np.float32).copy().reshape(  # fallback_shape
+                fallback_shape
+            )
+            return torch.from_numpy(array)  # fallback_shape
     if not isinstance(loaded, Tensor):
         raise ValueError("Legacy serialized embedding payload must contain one tensor.")
-    return loaded.detach().cpu()
+    return loaded.detach().cpu()  # (...)
 
 
 def convert_legacy_sqlite(

fastplms/embeddings/types.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 from collections.abc import Callable, Iterator, Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import Any, Literal, overload
-
 from torch import Tensor
 
 
@@ -39,7 +38,7 @@ class LazyTensorReference:
 
         if not isinstance(verify, bool):
             raise TypeError("verify must be a boolean.")
-        X = self._loader()
+        X = self._loader()  # self.shape
         if not isinstance(X, Tensor):
             raise TypeError(f"Stored tensor loader for {self.key!r} must return a Tensor.")
         if tuple(X.shape) != self.shape:
@@ -57,7 +56,7 @@ class LazyTensorReference:
             digest = tensor_sha256(X)
             if digest != self.sha256:
                 raise ValueError(f"Stored tensor {self.key!r} failed SHA-256 verification.")
-        return X
+        return X  # self.shape
 
 
 TensorValue = Tensor | LazyTensorReference
@@ -85,8 +84,8 @@ class EmbeddingRecord:
         if not isinstance(verify, bool):
             raise TypeError("verify must be a boolean.")
         if isinstance(self.tensor, LazyTensorReference):
-            return self.tensor.load(verify=verify)
-        return self.tensor
+            return self.tensor.load(verify=verify)  # (...)
+        return self.tensor  # (...)
 
 
 class EmbeddingResult(Sequence[EmbeddingRecord]):

fastplms/models/__init__.py

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

fastplms/models/esm_plusplus/modeling_esm_plusplus.py

@@ -3,13 +3,12 @@
 from __future__ import annotations
 
 import math
-from dataclasses import dataclass
-from functools import partial
-from typing import ClassVar
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from dataclasses import dataclass
+from functools import partial
+from typing import ClassVar
 from einops import rearrange
 from tokenizers import Tokenizer
 from tokenizers.models import BPE
@@ -22,6 +21,7 @@ from transformers.modeling_outputs import (
     TokenClassifierOutput,
 )
 
+
 try:
     from fastplms.attention import (
         AttentionBackend,
@@ -176,7 +176,7 @@ class RotaryEmbedding(torch.nn.Module):
         scaling_factor: float = 1.0,
         pos_idx_in_fp32: bool = True,
         device: torch.device | None = None,
-    ):
+    ) -> None:
         super().__init__()
         self.dim, self.base = dim, float(base)
         self.interleaved, self.scale_base = interleaved, scale_base
@@ -192,7 +192,7 @@ class RotaryEmbedding(torch.nn.Module):
         self._cos_k_cached: torch.Tensor | None = None
         self._sin_k_cached: torch.Tensor | None = None
 
-    def reset_parameters(self, device: torch.device | str | None = None):
+    def reset_parameters(self, device: torch.device | str | None = None) -> None:
         """Rebuild the non-persistent frequency buffers on ``device``."""
         if device is not None:
             buffer_device = torch.device(device)
@@ -250,14 +250,14 @@ class RotaryEmbedding(torch.nn.Module):
         device: torch.device | None,
     ) -> torch.Tensor:
         position_dtype = torch.float32 if self.pos_idx_in_fp32 else self.inv_freq.dtype
-        positions = torch.arange(token_count, device=device, dtype=position_dtype)
+        positions = torch.arange(token_count, device=device, dtype=position_dtype)  # (l,)
         positions.div_(self.scaling_factor)
         frequencies = (
             self.inv_freq.to(torch.float32)
             if self.pos_idx_in_fp32 and self.inv_freq.dtype != torch.float32
             else self.inv_freq
         )
-        return torch.outer(positions, frequencies)
+        return torch.outer(positions, frequencies)  # (l, d / 2)
 
     def _update_cos_sin_cache(
         self, seqlen: int, device: torch.device | None = None, dtype: torch.dtype | None = None
@@ -267,9 +267,9 @@ class RotaryEmbedding(torch.nn.Module):
             return
 
         self._seq_len_cached = seqlen
-        angles = self._rotary_angles(seqlen, device)
-        cos_angles = torch.cos(angles)
-        sin_angles = torch.sin(angles)
+        angles = self._rotary_angles(seqlen, device)  # (l, d / 2)
+        cos_angles = torch.cos(angles)  # (l, d / 2)
+        sin_angles = torch.sin(angles)  # (l, d / 2)
         if self.scale is None:
             self._cos_cached = cos_angles.to(dtype)
             self._sin_cached = sin_angles.to(dtype)
@@ -314,7 +314,6 @@ class RotaryEmbedding(torch.nn.Module):
         )
 
 
-### Feedforward Network Components
 def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int:
     """Compute corrected dimension for SwiGLU."""
     return int(((expansion_ratio * d_model) + 255) // 256 * 256)
@@ -323,7 +322,7 @@ def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int:
 class SwiGLU(nn.Module):
     """SwiGLU activation function."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -341,7 +340,6 @@ def swiglu_ln_ffn(d_model: int, expansion_ratio: float) -> nn.Sequential:
     )
 
 
-### Attention
 class MultiHeadAttention(nn.Module):
     """Multi-head attention with rotary embeddings and configurable backend.
 
@@ -356,7 +354,7 @@ class MultiHeadAttention(nn.Module):
         d_model: int,
         n_heads: int,
         attn_backend: str = "sdpa",
-    ):
+    ) -> None:
         super().__init__()
         self.d_model = d_model
         self.n_heads = n_heads
@@ -373,8 +371,9 @@ class MultiHeadAttention(nn.Module):
         self.rotary = RotaryEmbedding(d_model // n_heads)
 
     def _apply_rotary(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        q = q.unflatten(-1, (self.n_heads, self.d_head))
-        k = k.unflatten(-1, (self.n_heads, self.d_head))
+        # q, k: (b, l, d)
+        q = q.unflatten(-1, (self.n_heads, self.d_head))  # (b, l, h, d_h)
+        k = k.unflatten(-1, (self.n_heads, self.d_head))  # (b, l, h, d_h)
         q, k = self.rotary(q, k)
         q = q.flatten(-2, -1)
         k = k.flatten(-2, -1)
@@ -389,7 +388,8 @@ class MultiHeadAttention(nn.Module):
         output_attentions: bool = False,
         output_s_max: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
-        qkv = self.layernorm_qkv(x)
+        # x: (b, l, d)
+        qkv = self.layernorm_qkv(x)  # (b, l, 3 * d)
         query_sequence, key_sequence, value_sequence = torch.chunk(qkv, 3, dim=-1)
         query_sequence, key_sequence = (
             self.q_ln(query_sequence).to(query_sequence.dtype),
@@ -398,7 +398,7 @@ class MultiHeadAttention(nn.Module):
         query_sequence, key_sequence = self._apply_rotary(query_sequence, key_sequence)
         query_heads, key_heads, value_heads = map(
             self.reshaper, (query_sequence, key_sequence, value_sequence)
-        )
+        )  # each (b, h, l, d_h)
 
         attn_output, attn_weights, s_max = self._attn(
             query_heads,
@@ -461,8 +461,8 @@ class MultiHeadAttention(nn.Module):
     def _compute_s_max(
         self, query_heads: torch.Tensor, key_heads: torch.Tensor
     ) -> list[torch.Tensor]:
-        q_norm = torch.linalg.vector_norm(query_heads, dim=-1)
-        k_norm = torch.linalg.vector_norm(key_heads, dim=-1)
+        q_norm = torch.linalg.vector_norm(query_heads, dim=-1)  # (b, h, l)
+        k_norm = torch.linalg.vector_norm(key_heads, dim=-1)  # (b, h, l)
         s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(
             dim=0
         ).values * self.scale
@@ -476,12 +476,15 @@ class MultiHeadAttention(nn.Module):
         attention_mask_4d: torch.Tensor | None = None,
         output_s_max: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]:
-        attn_weights = torch.matmul(query_heads, key_heads.transpose(-2, -1)) * self.scale
+        # query_heads, key_heads, value_heads: (b, h, l, d_h)
+        attn_weights = (
+            torch.matmul(query_heads, key_heads.transpose(-2, -1)) * self.scale
+        )  # (b, h, l, l)
         if attention_mask_4d is not None:
             attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf"))
         attn_weights = F.softmax(attn_weights, dim=-1)
-        context_heads = torch.matmul(attn_weights, value_heads)
-        attn_output = rearrange(context_heads, "b h s d -> b s (h d)")
+        context_heads = torch.matmul(attn_weights, value_heads)  # (b, h, l, d_h)
+        attn_output = rearrange(context_heads, "b h s d -> b s (h d)")  # (b, l, d)
         s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None
         return attn_output, attn_weights, s_max
 
@@ -548,7 +551,6 @@ class MultiHeadAttention(nn.Module):
         return rearrange(context_heads, "b h s d -> b s (h d)"), None
 
 
-### Regression Head
 def RegressionHead(d_model: int, output_dim: int, hidden_dim: int | None = None) -> nn.Module:
     """Create a regression head with optional hidden dimension.
 
@@ -566,7 +568,6 @@ def RegressionHead(d_model: int, output_dim: int, hidden_dim: int | None = None)
     )
 
 
-### Transformer Block
 class UnifiedTransformerBlock(nn.Module):
     """Transformer block with attention and feedforward layers."""
 
@@ -578,7 +579,7 @@ class UnifiedTransformerBlock(nn.Module):
         expansion_ratio: float = 8 / 3,
         dropout: float = 0.0,
         attn_backend: str = "sdpa",
-    ):
+    ) -> None:
         super().__init__()
         self.attn = MultiHeadAttention(d_model=d_model, n_heads=n_heads, attn_backend=attn_backend)
         self.ffn = swiglu_ln_ffn(d_model, expansion_ratio)
@@ -607,7 +608,6 @@ class UnifiedTransformerBlock(nn.Module):
         return x, attn_weights, s_max
 
 
-### Model Outputs
 @dataclass
 class TransformerOutput(ModelOutput):
     """Output type for transformer encoder."""
@@ -640,7 +640,6 @@ class ESMplusplusTokenClassifierOutput(TokenClassifierOutput):
     s_max: tuple[list[torch.Tensor], ...] | None = None
 
 
-### Transformer Stack
 class TransformerStack(nn.Module):
     """Stack of transformer blocks."""
 
@@ -651,7 +650,7 @@ class TransformerStack(nn.Module):
         n_layers: int,
         dropout: float = 0.0,
         attn_backend: str = "sdpa",
-    ):
+    ) -> None:
         super().__init__()
         self.attention_backend = resolve_attention_backend(attn_backend)
         self.blocks = nn.ModuleList(
@@ -690,6 +689,7 @@ class TransformerStack(nn.Module):
         output_s_max: bool | None = False,
         esmfold2_hidden_states: bool = False,
     ) -> TransformerOutput:
+        # x: (b, l, d); attention_mask, sequence_id: (b, l)
         hidden_states = () if output_hidden_states else None
         attentions = () if output_attentions else None
         full_s_max = () if output_s_max else None
@@ -971,7 +971,7 @@ class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin):
 
     config_class = ESMplusplusConfig
 
-    def __init__(self, config: ESMplusplusConfig, **kwargs):
+    def __init__(self, config: ESMplusplusConfig, **kwargs) -> None:
         PreTrainedESMplusplusModel.__init__(self, config, **kwargs)
         self.config = config
         self.vocab_size = config.vocab_size
@@ -1089,7 +1089,7 @@ class ESMplusplusForMaskedLM(
 
     config_class = ESMplusplusConfig
 
-    def __init__(self, config: ESMplusplusConfig, **kwargs):
+    def __init__(self, config: ESMplusplusConfig, **kwargs) -> None:
         PreTrainedESMplusplusModel.__init__(self, config, **kwargs)
         self.config = config
         self.vocab_size = config.vocab_size
@@ -1215,7 +1215,7 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixi
     Extends the base ESM++ model with a classification head.
     """
 
-    def __init__(self, config: ESMplusplusConfig, **kwargs):
+    def __init__(self, config: ESMplusplusConfig, **kwargs) -> None:
         pooling_types = kwargs.pop("pooling_types", None)
         if pooling_types is None:
             pooling_types = config.classifier_pooling_types or ["mean", "var"]
@@ -1355,7 +1355,7 @@ class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM, EmbeddingMixin):
     Extends the base ESM++ model with a token classification head.
     """
 
-    def __init__(self, config: ESMplusplusConfig, **kwargs):
+    def __init__(self, config: ESMplusplusConfig, **kwargs) -> None:
         ESMplusplusForMaskedLM.__init__(self, config, **kwargs)
         self.config = config
         self.num_labels = config.num_labels
@@ -1526,7 +1526,7 @@ class EsmSequenceTokenizer(PreTrainedTokenizerFast):
             **kwargs,
         )
 
-    # These are a footgun, we never use the `bos` token anywhere so we're just overriding it here.
+    # ESMC does not use BOS, so expose the sequence-start token through the HF BOS fields.
     @property
     def bos_token(self):
         return self.cls_token

fastplms/models/ttt.py

@@ -3,12 +3,13 @@ from __future__ import annotations
 import contextlib
 import math
 import numbers
-import typing as T
-from dataclasses import asdict, dataclass, fields
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from collections.abc import Iterator, Mapping
+from dataclasses import asdict, dataclass, fields
+from typing import Any
+
 
 _STANDARD_AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
 _TTT_SERIALIZATION_VERSION = 1
@@ -42,7 +43,7 @@ class TTTConfig:
         self.verify()
 
     @classmethod
-    def from_kwargs(cls, **kwargs: T.Any) -> TTTConfig:
+    def from_kwargs(cls, **kwargs: Any) -> TTTConfig:
         valid_names = {field.name for field in fields(cls)}
         unknown_names = set(kwargs) - valid_names
         if unknown_names:
@@ -53,7 +54,7 @@ class TTTConfig:
             kwargs["lora_target_modules"] = tuple(kwargs["lora_target_modules"])
         return cls(**kwargs)
 
-    def merged(self, overrides: T.Mapping[str, T.Any] | TTTConfig | None) -> TTTConfig:
+    def merged(self, overrides: Mapping[str, Any] | TTTConfig | None) -> TTTConfig:
         if overrides is None:
             return self
         if isinstance(overrides, TTTConfig):
@@ -65,7 +66,7 @@ class TTTConfig:
             values[name] = value
         return TTTConfig(**values)
 
-    def to_dict(self) -> dict[str, T.Any]:
+    def to_dict(self) -> dict[str, Any]:
         return asdict(self)
 
     def verify(self) -> None:
@@ -224,9 +225,12 @@ class LoraInjectedLinear(nn.Module):
         return self.linear._parameters["bias"]
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        base = self.linear(x)
-        delta = self.lora_up(self.lora_down(x.to(dtype=torch.float32))) * self.scale
-        return base + delta.to(dtype=base.dtype)
+        # x: (..., d_in)
+        base = self.linear(x)  # (..., d_out)
+        delta = (  # (..., d_out)
+            self.lora_up(self.lora_down(x.to(dtype=torch.float32))) * self.scale
+        )
+        return base + delta.to(dtype=base.dtype)  # (..., d_out)
 
     def reset_lora_parameters(self) -> None:
         with torch.no_grad():
@@ -235,7 +239,7 @@ class LoraInjectedLinear(nn.Module):
 
 
 class FastPLMTestTimeTrainingMixin:
-    def init_ttt(self, ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None) -> None:
+    def init_ttt(self, ttt_config: TTTConfig | Mapping[str, Any] | None = None) -> None:
         base_config = self.__dict__.get("_ttt_cfg")
         if base_config is None:
             base_config = TTTConfig()
@@ -245,7 +249,7 @@ class FastPLMTestTimeTrainingMixin:
         serialized = getattr(getattr(self, "config", None), "fastplms_ttt", None)
         serialized_initialized = False
         if serialized is not None:
-            if not isinstance(serialized, T.Mapping):
+            if not isinstance(serialized, Mapping):
                 raise ValueError("config.fastplms_ttt must be a mapping.")
             version = serialized.get("version")
             if version != _TTT_SERIALIZATION_VERSION:
@@ -254,7 +258,7 @@ class FastPLMTestTimeTrainingMixin:
                     f"{version!r}; expected {_TTT_SERIALIZATION_VERSION}."
                 )
             serialized_config = serialized.get("config")
-            if not isinstance(serialized_config, T.Mapping):
+            if not isinstance(serialized_config, Mapping):
                 raise ValueError("Serialized FastPLMs TTT state is missing its config mapping.")
             configured = TTTConfig.from_kwargs(**dict(serialized_config))
             initialized_value = serialized.get("initialized", False)
@@ -285,15 +289,15 @@ class FastPLMTestTimeTrainingMixin:
         self,
         seq: str | list[str] | None = None,
         input_ids: torch.Tensor | None = None,
-        **kwargs: T.Any,
+        **kwargs: Any,
     ) -> torch.Tensor | dict[str, torch.Tensor]:
         del kwargs
         if input_ids is not None:
-            return input_ids
+            return input_ids  # (b, l)
         if seq is None:
             raise ValueError("Pass either seq or input_ids for TTT.")
         tokenized = self.tokenizer(seq, return_tensors="pt", padding=True)
-        return tokenized["input_ids"]
+        return tokenized["input_ids"]  # (b, l)
 
     def _ttt_mask_token(self) -> int:
         return int(self.tokenizer.mask_token_id)
@@ -302,6 +306,7 @@ class FastPLMTestTimeTrainingMixin:
         return int(self.tokenizer.pad_token_id)
 
     def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+        # input_ids: (b, l)
         tokenizer = self.tokenizer
         special_ids = set(tokenizer.all_special_ids)
         vocab_size = int(self.config.vocab_size)
@@ -309,13 +314,13 @@ class FastPLMTestTimeTrainingMixin:
         if unknown_id is not None:
             special_ids.add(int(unknown_id))
 
-        vocab: T.Mapping[str, T.Any] = {}
+        vocab: Mapping[str, Any] = {}
         get_vocab = getattr(tokenizer, "get_vocab", None)
         if callable(get_vocab):
             vocab = get_vocab()
-        elif isinstance(getattr(tokenizer, "vocab", None), T.Mapping):
+        elif isinstance(getattr(tokenizer, "vocab", None), Mapping):
             vocab = tokenizer.vocab
-        elif isinstance(getattr(tokenizer, "_token_to_id", None), T.Mapping):
+        elif isinstance(getattr(tokenizer, "_token_to_id", None), Mapping):
             vocab = tokenizer._token_to_id
 
         ids: list[int] = []
@@ -334,20 +339,20 @@ class FastPLMTestTimeTrainingMixin:
                 "TTT could not resolve any canonical amino-acid token IDs from the tokenizer; "
                 "refusing to sample arbitrary or reserved vocabulary entries."
             )
-        return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype)
+        return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype)  # (c_aa,)
 
     def _ttt_predict_logits(
         self,
         batch: torch.Tensor | dict[str, torch.Tensor],
-        **kwargs: T.Any,
+        **kwargs: Any,
     ) -> torch.Tensor:
         del kwargs
         if isinstance(batch, dict):
             output = self(**batch)
-            return output.logits
-        attention_mask = batch.ne(self._ttt_padding_token())
+            return output.logits  # (b, l, c)
+        attention_mask = batch.ne(self._ttt_padding_token())  # (b, l)
         output = self(input_ids=batch, attention_mask=attention_mask)
-        return output.logits
+        return output.logits  # (b, l, c)
 
     def _ttt_eval_step(
         self,
@@ -355,8 +360,8 @@ class FastPLMTestTimeTrainingMixin:
         loss: float,
         seq: str | list[str] | None = None,
         input_ids: torch.Tensor | None = None,
-        **kwargs: T.Any,
-    ) -> tuple[dict[str, T.Any], float | None]:
+        **kwargs: Any,
+    ) -> tuple[dict[str, Any], float | None]:
         del step, loss, seq, input_ids, kwargs
         return {}, None
 
@@ -443,8 +448,8 @@ class FastPLMTestTimeTrainingMixin:
         for module in self._ttt_lora_modules():
             snapshot.append(
                 {
-                    "lora_down.weight": module.lora_down.weight.detach().clone(),
-                    "lora_up.weight": module.lora_up.weight.detach().clone(),
+                    "lora_down.weight": module.lora_down.weight.detach().clone(),  # (r, d_in)
+                    "lora_up.weight": module.lora_up.weight.detach().clone(),  # (d_out, r)
                 }
             )
         if not snapshot:
@@ -473,14 +478,14 @@ class FastPLMTestTimeTrainingMixin:
         for module in self._ttt_lora_modules():
             module.reset_lora_parameters()
 
-    def _ttt_serialized_contract(self) -> dict[str, T.Any]:
+    def _ttt_serialized_contract(self) -> dict[str, Any]:
         return {
             "version": _TTT_SERIALIZATION_VERSION,
             "initialized": bool(self._ttt_initialized),
             "config": self.ttt_config.to_dict(),
         }
 
-    def save_pretrained(self, save_directory: T.Any, *args: T.Any, **kwargs: T.Any) -> T.Any:
+    def save_pretrained(self, save_directory: Any, *args: Any, **kwargs: Any) -> Any:
         """Save initialized adapters, their reset baseline, and the TTT config.
 
         Adapter injection changes the module tree, so the serialized config must
@@ -523,16 +528,16 @@ class FastPLMTestTimeTrainingMixin:
         device: torch.device,
     ) -> torch.Tensor | dict[str, torch.Tensor]:
         if isinstance(batch, dict):
-            return {name: tensor.to(device) for name, tensor in batch.items()}
-        return batch.to(device)
+            return {name: tensor.to(device) for name, tensor in batch.items()}  # unchanged shapes
+        return batch.to(device)  # unchanged shape
 
     def _ttt_input_ids_from_batch(
         self,
         batch: torch.Tensor | dict[str, torch.Tensor],
     ) -> torch.Tensor:
         if isinstance(batch, dict):
-            return batch["input_ids"]
-        return batch
+            return batch["input_ids"]  # (b, l)
+        return batch  # (b, l)
 
     def _ttt_set_input_ids(
         self,
@@ -541,13 +546,14 @@ class FastPLMTestTimeTrainingMixin:
     ) -> torch.Tensor | dict[str, torch.Tensor]:
         if isinstance(batch, dict):
             updated = dict(batch)
-            updated["input_ids"] = input_ids
+            updated["input_ids"] = input_ids  # (b, l)
             return updated
-        return input_ids
+        return input_ids  # (b, l)
 
     def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
-        residue_ids = self._ttt_replacement_tokens(input_ids)
-        return torch.isin(input_ids, residue_ids)
+        # input_ids: (b, l)
+        residue_ids = self._ttt_replacement_tokens(input_ids)  # (c_aa,)
+        return torch.isin(input_ids, residue_ids)  # (b, l)
 
     def _ttt_validate_tokenized_batch(
         self,
@@ -570,12 +576,14 @@ class FastPLMTestTimeTrainingMixin:
                     "DPLM2 TTT could not resolve the structure-token boundary safely."
                 )
             pad_token = self._ttt_padding_token()
-            generic_aa_special_ids = torch.tensor(
+            generic_aa_special_ids = torch.tensor(  # (4,)
                 [int(self.config.vocab_size) + offset for offset in range(4)],
                 device=input_ids.device,
                 dtype=input_ids.dtype,
             )
-            is_structure = input_ids.ge(int(struct_boundary)) & input_ids.ne(pad_token)
+            is_structure = input_ids.ge(int(struct_boundary)) & input_ids.ne(  # (b, l)
+                pad_token
+            )
             is_structure &= ~torch.isin(input_ids, generic_aa_special_ids)
             if bool(is_structure.any()):
                 raise ValueError(
@@ -584,8 +592,11 @@ class FastPLMTestTimeTrainingMixin:
                 )
 
             if isinstance(batch, dict) and "type_ids" in batch:
-                type_ids = batch["type_ids"]
-                attention_mask = batch.get("attention_mask", input_ids.ne(pad_token)).bool()
+                type_ids = batch["type_ids"]  # (b, l)
+                attention_mask = batch.get(  # (b, l)
+                    "attention_mask",
+                    input_ids.ne(pad_token),
+                ).bool()
                 if bool(((type_ids == int(self.config.struct_type)) & attention_mask).any()):
                     raise ValueError(
                         "DPLM2 TTT currently supports amino-acid-only inputs; structure "
@@ -607,13 +618,15 @@ class FastPLMTestTimeTrainingMixin:
         cfg = self.ttt_config
         if input_ids.shape[1] <= cfg.crop_size:
             return batch
-        position_has_residue = self._ttt_non_special_mask(input_ids).any(dim=0).to(torch.int64)
-        prefix = F.pad(position_has_residue.cumsum(dim=0), (1, 0))
-        window_counts = prefix[cfg.crop_size :] - prefix[: -cfg.crop_size]
-        valid_starts = torch.where(window_counts > 0)[0]
+        position_has_residue = (  # (l,)
+            self._ttt_non_special_mask(input_ids).any(dim=0).to(torch.int64)
+        )
+        prefix = F.pad(position_has_residue.cumsum(dim=0), (1, 0))  # (l + 1,)
+        window_counts = prefix[cfg.crop_size :] - prefix[: -cfg.crop_size]  # (l-crop+1,)
+        valid_starts = torch.where(window_counts > 0)[0]  # (n_valid,)
         if valid_starts.numel() == 0:
             raise ValueError("TTT could not find a crop containing a biological residue token.")
-        selected = torch.randint(
+        selected = torch.randint(  # (1,)
             valid_starts.numel(),
             (1,),
             generator=generator,
@@ -625,11 +638,11 @@ class FastPLMTestTimeTrainingMixin:
             cropped = {}
             for name, tensor in batch.items():
                 if tensor.ndim >= 2 and tensor.shape[1] == input_ids.shape[1]:
-                    cropped[name] = tensor[:, start:end]
+                    cropped[name] = tensor[:, start:end]  # (b, crop_size, ...)
                 else:
                     cropped[name] = tensor
             return cropped
-        return input_ids[:, start:end]
+        return input_ids[:, start:end]  # (b, crop_size)
 
     def _ttt_sample_batch(
         self,
@@ -638,69 +651,69 @@ class FastPLMTestTimeTrainingMixin:
     ) -> tuple[torch.Tensor | dict[str, torch.Tensor], torch.Tensor]:
         cfg = self.ttt_config
         batch = self._ttt_sample_crop(tokenized, generator)
-        input_ids = self._ttt_input_ids_from_batch(batch)
-        row_has_residue = self._ttt_non_special_mask(input_ids).any(dim=1)
-        eligible_rows = torch.where(row_has_residue)[0]
+        input_ids = self._ttt_input_ids_from_batch(batch)  # (b, l)
+        row_has_residue = self._ttt_non_special_mask(input_ids).any(dim=1)  # (b,)
+        eligible_rows = torch.where(row_has_residue)[0]  # (n_eligible,)
         if eligible_rows.numel() == 0:
             raise ValueError(
                 "TTT sampled batch contains no trainable biological residue tokens."
             )
-        sampled_row_indices = torch.randint(
+        sampled_row_indices = torch.randint(  # (b_sample,)
             eligible_rows.numel(),
             (cfg.batch_size,),
             generator=generator,
             device=input_ids.device,
         )
-        rows = eligible_rows[sampled_row_indices]
+        rows = eligible_rows[sampled_row_indices]  # (b_sample,)
         if isinstance(batch, dict):
             sampled: torch.Tensor | dict[str, torch.Tensor] = {}
             for name, tensor in batch.items():
                 if tensor.ndim >= 1 and tensor.shape[0] == input_ids.shape[0]:
-                    sampled[name] = tensor.index_select(0, rows)
+                    sampled[name] = tensor.index_select(0, rows)  # (b_sample, ...)
                 else:
                     sampled[name] = tensor
         else:
-            sampled = input_ids.index_select(0, rows)
+            sampled = input_ids.index_select(0, rows)  # (b_sample, l)
 
-        sampled_ids = self._ttt_input_ids_from_batch(sampled)
-        labels = sampled_ids.clone()
-        non_special = self._ttt_non_special_mask(sampled_ids)
-        label_mask = torch.zeros_like(non_special)
+        sampled_ids = self._ttt_input_ids_from_batch(sampled)  # (b_sample, l)
+        labels = sampled_ids.clone()  # (b_sample, l)
+        non_special = self._ttt_non_special_mask(sampled_ids)  # (b_sample, l)
+        label_mask = torch.zeros_like(non_special)  # (b_sample, l)
         for row_idx in range(sampled_ids.shape[0]):
-            candidate_positions = torch.where(non_special[row_idx])[0]
+            candidate_positions = torch.where(non_special[row_idx])[0]  # (n_candidates,)
             if candidate_positions.numel() == 0:
                 continue
             num_mask = max(1, round(candidate_positions.numel() * cfg.mask_ratio))
-            order = torch.randperm(
+            order = torch.randperm(  # (n_candidates,)
                 candidate_positions.numel(),
                 generator=generator,
                 device=sampled_ids.device,
             )
-            chosen = candidate_positions[order[:num_mask]]
+            chosen = candidate_positions[order[:num_mask]]  # (n_mask,)
             label_mask[row_idx, chosen] = True
-        labels = labels.masked_fill(~label_mask, -100)
+        labels = labels.masked_fill(~label_mask, -100)  # (b_sample, l)
 
-        masked_ids = sampled_ids.clone()
-        chosen_positions = torch.where(label_mask)
+        masked_ids = sampled_ids.clone()  # (b_sample, l)
+        chosen_positions = torch.where(label_mask)  # two (n_chosen,) tensors
         if chosen_positions[0].numel() > 0:
-            random_values = torch.rand(
+            random_values = torch.rand(  # (n_chosen,)
                 chosen_positions[0].shape,
                 generator=generator,
                 device=sampled_ids.device,
             )
-            leave = random_values < cfg.bert_leave_prob
-            replace = (random_values >= cfg.bert_leave_prob) & (
+            leave = random_values < cfg.bert_leave_prob  # (n_chosen,)
+            replace = (random_values >= cfg.bert_leave_prob) & (  # (n_chosen,)
                 random_values < cfg.bert_leave_prob + cfg.bert_replace_prob
             )
-            mask = ~(leave | replace)
+            mask = ~(leave | replace)  # (n_chosen,)
             if mask.any():
                 masked_ids[
                     chosen_positions[0][mask],
                     chosen_positions[1][mask],
                 ] = self._ttt_mask_token()
             if replace.any():
-                replacement_tokens = self._ttt_replacement_tokens(sampled_ids)
-                replacement_idx = torch.randint(
+                replacement_tokens = self._ttt_replacement_tokens(sampled_ids)  # (c_aa,)
+                replacement_idx = torch.randint(  # (n_replace,)
                     replacement_tokens.shape[0],
                     (int(replace.sum().item()),),
                     generator=generator,
@@ -711,10 +724,10 @@ class FastPLMTestTimeTrainingMixin:
                     chosen_positions[1][replace],
                 ] = replacement_tokens[replacement_idx]
 
-        return self._ttt_set_input_ids(sampled, masked_ids), labels
+        return self._ttt_set_input_ids(sampled, masked_ids), labels  # batch, (b_sample, l)
 
     @contextlib.contextmanager
-    def _ttt_seed_scope(self, seed: int | None) -> T.Iterator[None]:
+    def _ttt_seed_scope(self, seed: int | None) -> Iterator[None]:
         if seed is None:
             yield
             return
@@ -736,9 +749,9 @@ class FastPLMTestTimeTrainingMixin:
         self,
         seq: str | list[str] | None = None,
         input_ids: torch.Tensor | None = None,
-        ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None,
-        **kwargs: T.Any,
-    ) -> dict[str, T.Any]:
+        ttt_config: TTTConfig | Mapping[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         if ttt_config is not None:
             if "_ttt_initialized" in self.__dict__ and self._ttt_initialized:
                 next_cfg = self.ttt_config.merged(ttt_config)
@@ -788,7 +801,7 @@ class FastPLMTestTimeTrainingMixin:
         module_modes = {module: module.training for module in self.modules()}
         requires_grad = {param: param.requires_grad for param in self.parameters()}
         losses: list[float] = []
-        step_metrics: list[dict[str, T.Any]] = []
+        step_metrics: list[dict[str, Any]] = []
         best_state: list[dict[str, torch.Tensor]] | None = None
         best_metric: float | None = None
         best_step = 0
@@ -805,14 +818,17 @@ class FastPLMTestTimeTrainingMixin:
                 optimizer.zero_grad(set_to_none=True)
                 total_micro_steps = cfg.steps * cfg.ags
                 for micro_step in range(total_micro_steps):
-                    batch, labels = self._ttt_sample_batch(tokenized, generator)
+                    batch, labels = self._ttt_sample_batch(  # batch, (b_sample, l)
+                        tokenized,
+                        generator,
+                    )
                     if not bool(labels.ne(-100).any()):
                         raise RuntimeError(
                             "TTT produced an all-ignored label batch; refusing a NaN update."
                         )
-                    logits = self._ttt_predict_logits(batch, **kwargs)
-                    labels = labels.to(device=logits.device)
-                    loss = F.cross_entropy(
+                    logits = self._ttt_predict_logits(batch, **kwargs)  # (b_sample, l, c)
+                    labels = labels.to(device=logits.device)  # (b_sample, l)
+                    loss = F.cross_entropy(  # ()
                         logits.reshape(-1, logits.shape[-1]),
                         labels.reshape(-1),
                         ignore_index=-100,

fastplms/registry.py

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

fastplms/runtime.py

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

modeling_fastplms.py

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

requirements.txt

@@ -0,0 +1,11 @@
+# Direct runtime dependencies for Synthyra/ESMplusplus_large.
+# FastPLMs source is embedded in this model repository.
+torch>=2.13,<2.14
+transformers>=5.13,<5.14
+huggingface-hub>=0.34,<2
+tokenizers>=0.22,<0.23
+safetensors>=0.5,<1
+numpy>=1.26,<3
+einops>=0.8,<1
+tqdm>=4.67,<5
+kernels>=0.15,<0.16

runtime-attestation.json

@@ -4,42 +4,43 @@
     "LICENSES/biohub-esm/LICENSE.md": "sha256:b63df9ca1dd96b3b21eec226b51b236d0bd152ac20eafc43aad46bf832b48d8a",
     "LICENSES/biohub-esm/THIRD_PARTY_NOTICE.md": "sha256:5bff8515ba4e0f53abdc43714c180b79c5b606160497d98de741a369cb9b6a23",
     "LICENSES/biohub-transformers/LICENSE": "sha256:77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049",
-    "README.md": "sha256:926c6c203528e78aeecc467fa325738d17378b9e1c715424869d1fdb802a04d3",
+    "README.md": "sha256:2115ee81a16a6c30e29ada82eaffdee331bf20360a5613850bdbfcc36046b077",
     "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
-    "config.json": "sha256:84629d9de640bb5c4c04e4830ce617aeeed7536e88635e0e9cbef5b49e6234e2",
-    "fastplms/__init__.py": "sha256:4fb3196022ca8ec699d59d09bdbc5f0184195552b773698ab9b061fe3cd7df12",
-    "fastplms/attention/__init__.py": "sha256:f60b9fecfb4bcb37a4e7c26dc2f752b9035f9cbad627b4a84213f3a92ec88f7d",
-    "fastplms/attention/_core.py": "sha256:8f7ec5b65bd8b6c6fa4951d50d1c0e499abf03ae00914794b51fc410201e3e33",
-    "fastplms/attention/_kernel_lock.py": "sha256:85d8521a2af5f94fad3948af3814db0c866c414ee4d43df797b9bd6f980e947b",
-    "fastplms/attention/interfaces.py": "sha256:1c6f06a8e411e0f9bf6d230522205c93ae46ea58864006fbb892aa05e5ca5749",
-    "fastplms/embeddings/__init__.py": "sha256:47ff8cdf682d44037dd9edab133e2e60675d60786bd5cf0bffd1998f31985555",
-    "fastplms/embeddings/pooling.py": "sha256:a140266ed6b1cc344c8507edc5c6c4f2dce464c3db70ba4b16c7ac2ba2fad96e",
-    "fastplms/embeddings/runner.py": "sha256:23ee4727a918d6d331f7a0f89b823d149f1a791f0c5586e3496d7b6eb2ce97e0",
-    "fastplms/embeddings/storage.py": "sha256:3fbe2bab75092e5a4cadf4d27e4752181d597469a65a55db085ceef808ed418e",
-    "fastplms/embeddings/types.py": "sha256:119718a20989d1ae5a60fabc0f5e98bdc172c5163b04db3d4554ac3956b30e52",
+    "config.json": "sha256:b0e9e94cc3b7d42bae5b7d1296eadfb7ba88b3e7ff87f5ea2f682764c58a51d3",
+    "fastplms/__init__.py": "sha256:8503e02debf24abc6d33c1913ff3eb9f245e63f7e62017e220d072c7393fdf2c",
+    "fastplms/attention/__init__.py": "sha256:ab3c0b6156968f418f3c97ea664959cb978f4137717b4362213bcfe25ab5e3d6",
+    "fastplms/attention/_core.py": "sha256:09d0f71a5ea2ad2138c9441f6a2f72bbf8e72729be41d8db737b163f36ce47da",
+    "fastplms/attention/_kernel_lock.py": "sha256:56455eb57cee9af08438eda52b775c9cff3cf994cadb9b116ce2ee10fd5f1801",
+    "fastplms/attention/interfaces.py": "sha256:6562a342961022fc2848d7004bd01fddaf8be8e08e782d0856e5c30a74232623",
+    "fastplms/embeddings/__init__.py": "sha256:cf0648de905a00ad16c0aece328eeb05f25739a4cf89b47b915e0f2b19c3ce26",
+    "fastplms/embeddings/pooling.py": "sha256:9bed8c466fce053333de9e25b04d1ab304f2954f403ffa0f5a3c749e8277b3ce",
+    "fastplms/embeddings/runner.py": "sha256:a7839763ad641a63d0e2671dcf34aa7b2a8114198950b64254188521b27453a8",
+    "fastplms/embeddings/storage.py": "sha256:98cf725d26766f32300f105345129ca6cf02ebd10628e746a7e71e107bc911a6",
+    "fastplms/embeddings/types.py": "sha256:c58f8c6a37a11a937c3f77b4e1617d5f213923d98ecaffddcdd283bd9dc42fe7",
     "fastplms/kernels.lock": "sha256:8e274e6f8f25fe272d95bb42849f672ebb897a67511e82be2b86c2a4b4962cdd",
     "fastplms/models.toml": "sha256:1b0d92911222f31b435bfee80251c9abe20b85392d65f8af8d285d1ccfbbc3c5",
-    "fastplms/models/__init__.py": "sha256:5e48c2cb3877aa6f42f3b5411d53b16bba2e32827bbde634f47f174c5cb36f86",
+    "fastplms/models/__init__.py": "sha256:d3dd084f5e8fafcf2d0fe4a26e4ab9284d3f9414b1f9a5bd86e308a406447c43",
     "fastplms/models/esm_plusplus/__init__.py": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
-    "fastplms/models/esm_plusplus/modeling_esm_plusplus.py": "sha256:7fc8047100a2d0ca230bc687126c4ccd3f46c0e51cefd39413e72e3842e14012",
-    "fastplms/models/ttt.py": "sha256:a0df4e98b02120d423e3c7ca9b866a8d0e3748b9076042a0102a838a11aed046",
-    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
-    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
-    "fastplms_bundle.py": "sha256:4e1c903d663ac29c7d42c6bfdf0369023eac9ad7b5075801f8efa4e65289ed41",
-    "modeling_fastplms.py": "sha256:b451c8179707dd905beee4e68cc7a4f675003e23941cae09e3f4716f3f600e06",
+    "fastplms/models/esm_plusplus/modeling_esm_plusplus.py": "sha256:ea812cc86c8cc5b3a41d82fdcbf56f999dd5599a181b1804b4c0f19b0f309557",
+    "fastplms/models/ttt.py": "sha256:f1d7f9298a930f4c40f8e321ffa08b09d1c5b8492bca623bdc7bfd03ec85774d",
+    "fastplms/registry.py": "sha256:cb37799dae5a4c0c780089ab2cca4cc56300c1c5f515b275892aa4e1c4b97eb6",
+    "fastplms/runtime.py": "sha256:9521c37fcd168bf9a4137277fb73672403ccf30d70d4845aa879fb22f42ddeb6",
+    "fastplms_bundle.py": "sha256:cd4e553471f54a7e7faf0fcc210dd75b8cbf990098116c040db3ff42fd947c43",
+    "modeling_fastplms.py": "sha256:486bbf6733e9e044c4d78cf2523f19ecd902b438dd37ee0e13356c01d5d7c761",
+    "requirements.txt": "sha256:32c675d7545024f2e9abae380ffd25690ec6e072f5d2cdb0cea9e25c57359fdd",
     "special_tokens_map.json": "sha256:0b7245ec86c8c3aeaf61523ba70dfa79be137e6283f127bd651adc30b4f15c74",
     "tokenizer.json": "sha256:8d3447b278176e65fb3ef0224472927bf5fee3be46ea2bd77fad0111423cee1f",
     "tokenizer_config.json": "sha256:e8d8e40c9f92b334f0272e80bb65ed4043cb9836523cbae899e9859e8cbb8833"
   },
   "model_id": "esmc_large",
   "redistributable": true,
-  "release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "runtime_bundle_sha256": "b01a730b69d0e25723da00aed2f25c2ede185e260e0ab1f7f29cc62934ab17d1",
-  "runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
+  "release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "runtime_bundle_sha256": "d73f544a7623d12eff2ba407e2d7e86417b7d731e36a2931515930eae7d7e8e5",
+  "runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
   "schema_version": 2,
   "scope": "runtime-only",
-  "source_tree_sha256": "a95686f9e9cb31ebf4a3ced01c54787c6496336528fba0f696d407b053ebabc6",
+  "source_tree_sha256": "bed9c7b1e1b0500dcd370e92e1f83b47494f9167d37b16e248f9d17e0135f822",
   "weights": {
     "repo_id": "Synthyra/ESMplusplus_large",
     "revision": "f813401638b3fddab09748aec1ad2bf537aa4208"