Update FastPLMs runtime and model cards

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

README.md

@@ -17,16 +17,35 @@ Supported Transformers entry points are `AutoConfig`, `AutoModel`,
 `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`,
 `AutoModelForTokenClassification`.
 
+## Capabilities
+
+| Feature | Status |
+| --- | --- |
+| Sequence classification | Supported: base weights with an untrained task head |
+| Token classification | Supported: base weights with an untrained task head |
+| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` |
+| Embeddings | Supported: shared ordered embedding API |
+| Test-time training | Supported: low-rank masked-residue adaptation |
+| Attention variants | Supported: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` |
+| 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/ESM2-3B/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.
 
@@ -39,23 +58,23 @@ model_id = "Synthyra/ESM2-3B"
 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/ESM2-3B` 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 FP32 parameters with CUDA BF16 autocast.
+For offline validation, replace `model_id` with the manifest-built
+`dist/hub/ESM2-3B` 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
 
@@ -85,43 +104,119 @@ 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.
+
+## Downstream classification
+
+Both downstream AutoClasses reuse the checkpoint backbone and initialize a new,
+untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have
+shape `(b, l)` and use `-100` outside biological positions:
+
+```python
+import torch
+from transformers import AutoTokenizer
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoModelForTokenClassification,
+)
+
+model_id = "Synthyra/ESM2-3B"
+sequence_model = AutoModelForSequenceClassification.from_pretrained(
+    model_id, num_labels=2, trust_remote_code=True
+).eval()
+token_model = AutoModelForTokenClassification.from_pretrained(
+    model_id, num_labels=3, trust_remote_code=True
+).eval()
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"]
+batch = tokenizer(sequences, padding=True, return_tensors="pt")
+biological = batch["attention_mask"].bool()
+for special_id in tokenizer.all_special_ids:
+    biological &= batch["input_ids"].ne(special_id)
+
+sequence_labels = torch.zeros(len(sequences), dtype=torch.long)
+token_labels = torch.full_like(batch["input_ids"], -100)
+token_labels[biological] = 0
+
+with torch.inference_mode():
+    sequence_output = sequence_model(**batch, labels=sequence_labels)
+    token_output = token_model(**batch, labels=token_labels)
+print(sequence_output.logits.shape)  # (b, 2)
+print(token_output.logits.shape)     # (b, l, 3)
+```
+
+## 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, TaskType, get_peft_model
+
+peft_model = get_peft_model(
+    sequence_model,
+    LoraConfig(
+        task_type=TaskType.SEQ_CLS,
+        r=8,
+        lora_alpha=16,
+        target_modules="all-linear",
+        modules_to_save=["classifier"],
+    ),
+)
 ```
 
-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 advertises a classification head, so the separately trained
+`classifier` is saved with the adapter.
+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/ESM2-3B",
+    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.
 
 ## Masked language modeling and contacts
 
@@ -174,7 +269,7 @@ required.
 - Precision policies: `default`
 - BF16 execution: `fp32_parameters_autocast`
 - Generation contract: `not_applicable`
-- Optional dependency group: `core`
+- Artifact dependency set: `core`
 - Weight publication allowed: `true`
 - Weight license status: `resolved`
 - Redistributable: `true`
@@ -185,28 +280,22 @@ required.
 - FastPLMs weights: `Synthyra/ESM2-3B`
 - 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: `facebook/esm2_t36_3B_UR50D`
 - Artifact source: `fast`
 - State transform: `esm2_hf_to_fastplms_v1`
-- BF16 execution: `fp32_parameters_autocast`
 - Pinned upstreams: `fair-esm`
-- Reference container: `reference-esm2`
 - 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

@@ -19,11 +19,11 @@
   "fastplms_checkpoint_repo_id": "Synthyra/ESM2-3B",
   "fastplms_checkpoint_revision": "ff89d0180f414ab9c677219a25da79bf09185456",
   "fastplms_model_id": "esm2_3b",
-  "fastplms_release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "fastplms_runtime_bundle_sha256": "0e9cee71faeeb3dd6406d10d4107c209b10379ab0ee07d0115b13102b6525460",
-  "fastplms_runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_source_tree_sha256": "fb8779cf2f90db290a898b170b6c6d2ebe9835b89dc946cc99a94417cb586aa4",
+  "fastplms_release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "fastplms_runtime_bundle_sha256": "f22353ef386f607784889dc303488a2dc5e99b009a5e423ab3eff1f1eb03a9a4",
+  "fastplms_runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_source_tree_sha256": "818bf2bf5f573adc842b55f4f8feff926ddc25f904c5769ff5c382077d3a78a2",
   "fastplms_weights_revision": "ff89d0180f414ab9c677219a25da79bf09185456",
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.0,

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_rotary.py

@@ -15,8 +15,9 @@ from torch import nn
 def _rotate_half(tensor: torch.Tensor) -> torch.Tensor:
     """Rotate the final dimension of X by 90 degrees in paired subspaces."""
 
-    first, second = tensor.chunk(2, dim=-1)
-    return torch.cat((-second, first), dim=-1)
+    # tensor: (..., d)
+    first, second = tensor.chunk(2, dim=-1)  # (..., d / 2), (..., d / 2)
+    return torch.cat((-second, first), dim=-1)  # (..., d)
 
 
 def apply_rotary_pos_emb(
@@ -26,9 +27,10 @@ def apply_rotary_pos_emb(
 ) -> torch.Tensor:
     """Apply cached rotary factors to X with shape ``(b, h, l, d)``."""
 
-    cos = cos[:, :, : tensor.shape[-2], :]
-    sin = sin[:, :, : tensor.shape[-2], :]
-    return tensor * cos + _rotate_half(tensor) * sin
+    # tensor: (b, h, l, d); cos, sin: (1, 1, l_cache, d)
+    cos = cos[:, :, : tensor.shape[-2], :]  # (1, 1, l, d)
+    sin = sin[:, :, : tensor.shape[-2], :]  # (1, 1, l, d)
+    return tensor * cos + _rotate_half(tensor) * sin  # (b, h, l, d)
 
 
 class RotaryEmbedding(nn.Module):
@@ -38,7 +40,9 @@ class RotaryEmbedding(nn.Module):
 
     def __init__(self, dim: int) -> None:
         super().__init__()
-        frequencies = 1.0 / (10_000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
+        frequencies = 1.0 / (  # (d / 2,)
+            10_000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim)
+        )
         # Keep this persistent to preserve the historical checkpoint schema.
         self.register_buffer("inv_freq", frequencies)
         self._seq_len_cached: int | None = None
@@ -50,6 +54,7 @@ class RotaryEmbedding(nn.Module):
         tensor: torch.Tensor,
         seq_dimension: int = 2,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        # tensor: (..., l, d)
         seq_len = tensor.shape[seq_dimension]
         cache_stale = (
             self._cos_cached is None
@@ -59,23 +64,29 @@ class RotaryEmbedding(nn.Module):
         )
         if cache_stale:
             self._seq_len_cached = seq_len
-            positions = torch.arange(seq_len, device=tensor.device).type_as(self.inv_freq)
-            angles = torch.outer(positions, self.inv_freq)
-            angles = torch.cat((angles, angles), dim=-1).to(tensor.device)
-            self._cos_cached = angles.cos()[None, None, :, :]
-            self._sin_cached = angles.sin()[None, None, :, :]
+            positions = torch.arange(seq_len, device=tensor.device).type_as(  # (l,)
+                self.inv_freq
+            )
+            angles = torch.outer(positions, self.inv_freq)  # (l, d / 2)
+            angles = torch.cat((angles, angles), dim=-1).to(tensor.device)  # (l, d)
+            self._cos_cached = angles.cos()[None, None, :, :]  # (1, 1, l, d)
+            self._sin_cached = angles.sin()[None, None, :, :]  # (1, 1, l, d)
 
         assert self._cos_cached is not None
         assert self._sin_cached is not None
-        return self._cos_cached, self._sin_cached
+        return self._cos_cached, self._sin_cached  # (1, 1, l, d), (1, 1, l, d)
 
     def forward(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        cos, sin = self._update_cos_sin_tables(key, seq_dimension=-2)
+        # query, key: (b, h, l, d)
+        cos, sin = self._update_cos_sin_tables(  # (1, 1, l, d), (1, 1, l, d)
+            key,
+            seq_dimension=-2,
+        )
         return (
-            apply_rotary_pos_emb(query, cos, sin).to(dtype=query.dtype),
-            apply_rotary_pos_emb(key, cos, sin).to(dtype=key.dtype),
+            apply_rotary_pos_emb(query, cos, sin).to(dtype=query.dtype),  # (b, h, l, d)
+            apply_rotary_pos_emb(key, cos, sin).to(dtype=key.dtype),  # (b, h, l, d)
         )

fastplms/models/esm2/modeling_fastesm.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import Any, ClassVar
-
 import torch
 import torch.nn as nn
+from dataclasses import dataclass
+from typing import Any, ClassVar
 from einops import rearrange
 from torch.nn import functional as F
 from transformers import EsmTokenizer, PretrainedConfig, PreTrainedModel
@@ -27,6 +26,7 @@ from transformers.models.esm.modeling_esm import (
 
 from fastplms.models._esm_rotary import RotaryEmbedding
 
+
 try:
     from fastplms.attention import (
         AttentionBackend,
@@ -199,7 +199,7 @@ class FastEsmTokenizer(EsmTokenizer):
 
 
 class EsmSelfAttention(nn.Module):
-    def __init__(self, config, position_embedding_type: str | None = None):
+    def __init__(self, config, position_embedding_type: str | None = None) -> None:
         super().__init__()
         if config.hidden_size % config.num_attention_heads != 0:
             raise ValueError(
@@ -233,18 +233,22 @@ class EsmSelfAttention(nn.Module):
         output_attentions: bool = False,
         output_s_max: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
+        # hidden_states: (b, l, d); masks: (b, l) and (b, 1, 1, l)
         batch_size, seq_length = hidden_states.shape[:-1]
         hidden_shape = (batch_size, seq_length, -1, self.attention_head_size)
-        query_heads = self.query(hidden_states).view(hidden_shape).transpose(1, 2)
-        key_heads = self.key(hidden_states).view(hidden_shape).transpose(1, 2)
-        value_heads = self.value(hidden_states).view(hidden_shape).transpose(1, 2)
+        query_heads = self.query(hidden_states).view(hidden_shape).transpose(1, 2)  # (b, h, l, d_h)
+        key_heads = self.key(hidden_states).view(hidden_shape).transpose(1, 2)  # (b, h, l, d_h)
+        value_heads = self.value(hidden_states).view(hidden_shape).transpose(1, 2)  # (b, h, l, d_h)
 
-        query_heads = query_heads * self.scale
+        query_heads = query_heads * self.scale  # (b, h, l, d_h)
 
         if self.position_embedding_type == "rotary":
-            query_heads, key_heads = self.rotary_embeddings(query_heads, key_heads)
+            query_heads, key_heads = self.rotary_embeddings(  # both (b, h, l, d_h)
+                query_heads,
+                key_heads,
+            )
 
-        attn_output, attn_weights, s_max = self._attn(
+        attn_output, attn_weights, s_max = self._attn(  # (b, l, d), (b, h, l, l), heads
             query_heads,
             key_heads,
             value_heads,
@@ -254,7 +258,7 @@ class EsmSelfAttention(nn.Module):
             output_attentions=output_attentions,
             output_s_max=output_s_max,
         )
-        return attn_output, attn_weights, s_max
+        return attn_output, attn_weights, s_max  # (b, l, d), optional (b, h, l, l), heads
 
     def _attn(
         self,
@@ -313,10 +317,13 @@ class EsmSelfAttention(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)
-        s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(dim=0).values
-        return [s_max_bound[h] for h in range(self.num_attention_heads)]
+        # query_heads, key_heads: (b, h, l, d_h)
+        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 = (  # (h,)
+            q_norm.max(dim=-1).values * k_norm.max(dim=-1).values
+        ).max(dim=0).values
+        return [s_max_bound[h] for h in range(self.num_attention_heads)]  # h scalars
 
     def _manual_attn(
         self,
@@ -326,16 +333,24 @@ class EsmSelfAttention(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(-1, -2))
+        # query_heads, key_heads, value_heads: (b, h, l, d_h)
+        attn_weights = torch.matmul(query_heads, key_heads.transpose(-1, -2))  # (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)
+            attn_weights = attn_weights.masked_fill(  # (b, h, l, l)
+                attention_mask_4d.logical_not(),
+                float("-inf"),
+            )
+        attn_weights = F.softmax(attn_weights, dim=-1)  # (b, h, l, l)
         if self.dropout_prob > 0 and self.training:
-            attn_weights = F.dropout(attn_weights, p=self.dropout_prob, training=self.training)
-        context_heads = torch.matmul(attn_weights, value_heads)
-        attn_output = rearrange(context_heads, "b h s d -> b s (h d)")
+            attn_weights = F.dropout(  # (b, h, l, l)
+                attn_weights,
+                p=self.dropout_prob,
+                training=self.training,
+            )
+        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
+        return attn_output, attn_weights, s_max  # (b, l, d), (b, h, l, l), heads
 
     def _kernels_flash_attn(
         self,
@@ -344,14 +359,14 @@ class EsmSelfAttention(nn.Module):
         value_heads: torch.Tensor,
         attention_mask_2d: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, None]:
-        query_tokens = query_heads.transpose(1, 2).contiguous()
-        key_tokens = key_heads.transpose(1, 2).contiguous()
-        value_tokens = value_heads.transpose(1, 2).contiguous()
+        query_tokens = query_heads.transpose(1, 2).contiguous()  # (b, l, h, d_h)
+        key_tokens = key_heads.transpose(1, 2).contiguous()  # (b, l, h, d_h)
+        value_tokens = value_heads.transpose(1, 2).contiguous()  # (b, l, h, d_h)
         # Q has been pre-scaled by self.scale = 1/sqrt(head_dim) in forward().
         # Pass softmax_scale=1.0 to prevent the kernel from applying its default
         # 1/sqrt(head_dim) scale on top (which would yield effective scale
         # 1/head_dim and break parity vs sdpa).
-        attn_output = kernels_flash_attention_func(
+        attn_output = kernels_flash_attention_func(  # (b, l, h, d_h)
             query_states=query_tokens,
             key_states=key_tokens,
             value_states=value_tokens,
@@ -360,7 +375,7 @@ class EsmSelfAttention(nn.Module):
             softmax_scale=1.0,
             implementation=self.attn_backend.value,
         )
-        return rearrange(attn_output, "b s h d -> b s (h d)"), None
+        return rearrange(attn_output, "b s h d -> b s (h d)"), None  # (b, l, d), None
 
     def _flex_attn(
         self,
@@ -378,10 +393,10 @@ class EsmSelfAttention(nn.Module):
             shape=tuple(query_heads.shape),
             mask_semantics="padding",
         )
-        context_heads = fn(
+        context_heads = fn(  # (b, h, l, d_h)
             query_heads, key_heads, value_heads, block_mask=flex_block_mask, scale=1.0
         )
-        return rearrange(context_heads, "b h s d -> b s (h d)"), None
+        return rearrange(context_heads, "b h s d -> b s (h d)"), None  # (b, l, d), None
 
     def _sdpa_attn(
         self,
@@ -402,7 +417,7 @@ class EsmSelfAttention(nn.Module):
 
 
 class EsmAttention(nn.Module):
-    def __init__(self, config):
+    def __init__(self, config) -> None:
         super().__init__()
         self.self = EsmSelfAttention(config)
         self.output = EsmSelfOutput(config)
@@ -417,8 +432,9 @@ class EsmAttention(nn.Module):
         output_attentions: bool = False,
         output_s_max: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
-        hidden_states_ln = self.LayerNorm(hidden_states)
-        attn_output, attn_weights, s_max = self.self(
+        # hidden_states: (b, l, d)
+        hidden_states_ln = self.LayerNorm(hidden_states)  # (b, l, d)
+        attn_output, attn_weights, s_max = self.self(  # (b, l, d), optional (b, h, l, l), heads
             hidden_states_ln,
             attention_mask_2d=attention_mask_2d,
             attention_mask_4d=attention_mask_4d,
@@ -426,12 +442,12 @@ class EsmAttention(nn.Module):
             output_attentions=output_attentions,
             output_s_max=output_s_max,
         )
-        attention_output = self.output(attn_output, hidden_states)
-        return attention_output, attn_weights, s_max
+        attention_output = self.output(attn_output, hidden_states)  # (b, l, d)
+        return attention_output, attn_weights, s_max  # (b, l, d), optional weights, heads
 
 
 class EsmLayer(nn.Module):
-    def __init__(self, config):
+    def __init__(self, config) -> None:
         super().__init__()
         self.chunk_size_feed_forward = config.chunk_size_feed_forward
         self.seq_len_dim = 1
@@ -449,7 +465,7 @@ class EsmLayer(nn.Module):
         output_attentions: bool = False,
         output_s_max: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
-        attention_output, attn_weights, s_max = self.attention(
+        attention_output, attn_weights, s_max = self.attention(  # (b, l, d), weights, heads
             hidden_states,
             attention_mask_2d=attention_mask_2d,
             attention_mask_4d=attention_mask_4d,
@@ -457,18 +473,19 @@ class EsmLayer(nn.Module):
             output_attentions=output_attentions,
             output_s_max=output_s_max,
         )
-        layer_output = self.feed_forward_chunk(attention_output)
-        return layer_output, attn_weights, s_max
+        layer_output = self.feed_forward_chunk(attention_output)  # (b, l, d)
+        return layer_output, attn_weights, s_max  # (b, l, d), weights, heads
 
-    def feed_forward_chunk(self, attention_output):
-        attention_output_ln = self.LayerNorm(attention_output)
-        intermediate_output = self.intermediate(attention_output_ln)
-        layer_output = self.output(intermediate_output, attention_output)
-        return layer_output
+    def feed_forward_chunk(self, attention_output: torch.Tensor) -> torch.Tensor:
+        # attention_output: (b, l, d)
+        attention_output_ln = self.LayerNorm(attention_output)  # (b, l, d)
+        intermediate_output = self.intermediate(attention_output_ln)  # (b, l, d_ff)
+        layer_output = self.output(intermediate_output, attention_output)  # (b, l, d)
+        return layer_output  # (b, l, d)
 
 
 class EsmEncoder(nn.Module):
-    def __init__(self, config):
+    def __init__(self, config) -> None:
         super().__init__()
         self.config = config
         self.attention_backend = resolve_attention_backend(config.attn_backend)
@@ -484,6 +501,7 @@ class EsmEncoder(nn.Module):
         output_attentions: bool = False,
         output_s_max: bool = False,
     ) -> FastEsmEncoderOutput:
+        # hidden_states: (b, l, d); attention_mask: (b, l)
         all_hidden_states = () if output_hidden_states else None
         all_attentions = () if output_attentions else None
         full_s_max = () if output_s_max else None
@@ -507,6 +525,7 @@ class EsmEncoder(nn.Module):
                 all_hidden_states = (*all_hidden_states, hidden_states)
 
             if self.gradient_checkpointing and self.training:
+                # hidden_states: (b, l, d)
                 hidden_states, attn_weights, s_max = self._gradient_checkpointing_func(
                     layer_module.__call__,
                     hidden_states,
@@ -517,7 +536,7 @@ class EsmEncoder(nn.Module):
                     output_s_max,
                 )
             else:
-                hidden_states, attn_weights, s_max = layer_module(
+                hidden_states, attn_weights, s_max = layer_module(  # (b, l, d), weights, heads
                     hidden_states,
                     attention_mask_2d=attention_mask_2d,
                     attention_mask_4d=attention_mask_4d,
@@ -532,7 +551,7 @@ class EsmEncoder(nn.Module):
                 full_s_max = (*full_s_max, s_max)
 
         if self.emb_layer_norm_after:
-            hidden_states = self.emb_layer_norm_after(hidden_states)
+            hidden_states = self.emb_layer_norm_after(hidden_states)  # (b, l, d)
 
         if output_hidden_states:
             all_hidden_states = (*all_hidden_states, hidden_states)
@@ -756,7 +775,7 @@ class FAST_ESM_ENCODER(FastEsmPreTrainedModel, EmbeddingMixin):
 
 
 class FastEsmModel(FastEsmPreTrainedModel, EmbeddingMixin):
-    def __init__(self, config, add_pooling_layer: bool | None = None, **kwargs):
+    def __init__(self, config, add_pooling_layer: bool | None = None, **kwargs) -> None:
         FastEsmPreTrainedModel.__init__(self, config, **kwargs)
         self.config = config
         self.esm = FAST_ESM_ENCODER(config)
@@ -822,8 +841,8 @@ class FastEsmModel(FastEsmPreTrainedModel, EmbeddingMixin):
             output_s_max=output_s_max,
             return_dict=True,
         )
-        sequence_output = outputs.last_hidden_state
-        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+        sequence_output = outputs.last_hidden_state  # (b, l, d)
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None  # (b, d)
 
         result = FastEsmEncoderOutput(
             last_hidden_state=sequence_output,
@@ -836,7 +855,7 @@ class FastEsmModel(FastEsmPreTrainedModel, EmbeddingMixin):
 
 
 class FastEsmForMaskedLM(FastPLMTestTimeTrainingMixin, FastEsmPreTrainedModel, EmbeddingMixin):
-    def __init__(self, config, **kwargs):
+    def __init__(self, config, **kwargs) -> None:
         FastEsmPreTrainedModel.__init__(self, config, **kwargs)
         self.esm = FAST_ESM_ENCODER(config, add_pooling_layer=False)
         self.lm_head = EsmLMHead(config)
@@ -913,13 +932,13 @@ class FastEsmForMaskedLM(FastPLMTestTimeTrainingMixin, FastEsmPreTrainedModel, E
             output_s_max=output_s_max,
             return_dict=True,
         )
-        sequence_output = outputs.last_hidden_state
-        prediction_scores = self.lm_head(sequence_output)
+        sequence_output = outputs.last_hidden_state  # (b, l, d)
+        prediction_scores = self.lm_head(sequence_output)  # (b, l, c)
 
         loss = None
         if labels is not None:
-            labels = labels.to(prediction_scores.device)
-            loss = self.loss_fct(
+            labels = labels.to(prediction_scores.device)  # (b, l)
+            loss = self.loss_fct(  # ()
                 prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)
             )
 
@@ -935,7 +954,7 @@ class FastEsmForMaskedLM(FastPLMTestTimeTrainingMixin, FastEsmPreTrainedModel, E
 
 
 class FastEsmForSequenceClassification(FastEsmPreTrainedModel, EmbeddingMixin):
-    def __init__(self, config, **kwargs):
+    def __init__(self, config, **kwargs) -> None:
         FastEsmPreTrainedModel.__init__(self, config, **kwargs)
         self.num_labels = config.num_labels
         self.config = config
@@ -994,12 +1013,12 @@ class FastEsmForSequenceClassification(FastEsmPreTrainedModel, EmbeddingMixin):
             output_s_max=output_s_max,
             return_dict=True,
         )
-        sequence_output = outputs.last_hidden_state
-        logits = self.classifier(sequence_output)
+        sequence_output = outputs.last_hidden_state  # (b, l, d)
+        logits = self.classifier(sequence_output)  # (b, c)
 
         loss = None
         if labels is not None:
-            labels = labels.to(logits.device)
+            labels = labels.to(logits.device)  # (b,) or (b, c)
             if self.config.problem_type is None:
                 if self.num_labels == 1:
                     self.config.problem_type = "regression"
@@ -1012,13 +1031,13 @@ class FastEsmForSequenceClassification(FastEsmPreTrainedModel, EmbeddingMixin):
 
             if self.config.problem_type == "regression":
                 if self.num_labels == 1:
-                    loss = self.mse(logits.squeeze(), labels.squeeze())
+                    loss = self.mse(logits.squeeze(), labels.squeeze())  # ()
                 else:
-                    loss = self.mse(logits, labels)
+                    loss = self.mse(logits, labels)  # ()
             elif self.config.problem_type == "single_label_classification":
-                loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1))
+                loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1))  # ()
             elif self.config.problem_type == "multi_label_classification":
-                loss = self.bce(logits, labels)
+                loss = self.bce(logits, labels)  # ()
 
         result = EsmSequenceClassifierOutput(
             loss=loss,
@@ -1031,7 +1050,7 @@ class FastEsmForSequenceClassification(FastEsmPreTrainedModel, EmbeddingMixin):
 
 
 class FastEsmForTokenClassification(FastEsmPreTrainedModel, EmbeddingMixin):
-    def __init__(self, config, **kwargs):
+    def __init__(self, config, **kwargs) -> None:
         FastEsmPreTrainedModel.__init__(self, config, **kwargs)
         self.num_labels = config.num_labels
         self.esm = FAST_ESM_ENCODER(config, add_pooling_layer=False)
@@ -1088,14 +1107,14 @@ class FastEsmForTokenClassification(FastEsmPreTrainedModel, EmbeddingMixin):
             output_s_max=output_s_max,
             return_dict=True,
         )
-        sequence_output = outputs.last_hidden_state
-        sequence_output = self.dropout(sequence_output)
-        logits = self.classifier(sequence_output)
+        sequence_output = outputs.last_hidden_state  # (b, l, d)
+        sequence_output = self.dropout(sequence_output)  # (b, l, d)
+        logits = self.classifier(sequence_output)  # (b, l, c)
 
         loss = None
         if labels is not None:
-            labels = labels.to(logits.device)
-            loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            labels = labels.to(logits.device)  # (b, l)
+            loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))  # ()
 
         result = EsmTokenClassifierOutput(
             loss=loss,

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 != "0e9cee71faeeb3dd6406d10d4107c209b10379ab0ee07d0115b13102b6525460":
+if RUNTIME_HASH != "f22353ef386f607784889dc303488a2dc5e99b009a5e423ab3eff1f1eb03a9a4":
     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,14 +179,14 @@ def _install_runtime():
     return package
 
 _install_runtime()
-_module_225 = _import_without_bytecode("fastplms.models.esm2.modeling_fastesm")
-FastEsmConfig = _module_225.FastEsmConfig
+_module_181 = _import_without_bytecode("fastplms.models.esm2.modeling_fastesm")
+FastEsmConfig = _module_181.FastEsmConfig
 FastEsmConfig.__module__ = __name__
-FastEsmForMaskedLM = _module_225.FastEsmForMaskedLM
+FastEsmForMaskedLM = _module_181.FastEsmForMaskedLM
 FastEsmForMaskedLM.__module__ = __name__
-FastEsmForSequenceClassification = _module_225.FastEsmForSequenceClassification
+FastEsmForSequenceClassification = _module_181.FastEsmForSequenceClassification
 FastEsmForSequenceClassification.__module__ = __name__
-FastEsmForTokenClassification = _module_225.FastEsmForTokenClassification
+FastEsmForTokenClassification = _module_181.FastEsmForTokenClassification
 FastEsmForTokenClassification.__module__ = __name__
-FastEsmModel = _module_225.FastEsmModel
+FastEsmModel = _module_181.FastEsmModel
 FastEsmModel.__module__ = __name__

requirements.txt

@@ -0,0 +1,11 @@
+# Direct runtime dependencies for Synthyra/ESM2-3B.
+# 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

@@ -3,43 +3,44 @@
     "LICENSES/FastPLMs-Apache-2.0.txt": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
     "LICENSES/fair-esm/LICENSE": "sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93",
     "LICENSES/fair-esm/PROVENANCE.md": "sha256:950adb94daf15e646ddf226dacfe2a8e77801aa0793e439a9a3490a48eb666e7",
-    "README.md": "sha256:4fea7a1fc1b8e8f3120a0f4b96132a1e24d4ca4f79413d4d2cd97b972204ab3c",
+    "README.md": "sha256:1775e5fe2c71285b4c5c8b8af9adbe0e4694b90d3922069b52ae5852fa74482a",
     "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
-    "config.json": "sha256:481eb8d940168cefbcfbc6689d0bffbbf0404eff418ae392fc049420b5fbd9da",
-    "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:147c67ca43b3a952ab393892a953a5be4a1bc61c2d89713d280f0a09144ba730",
+    "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/_esm_rotary.py": "sha256:4dcee8af74adf2cfcabeeaadf4915f1da5fb6c0672d583bf3ee5478ac1290805",
+    "fastplms/models/__init__.py": "sha256:d3dd084f5e8fafcf2d0fe4a26e4ab9284d3f9414b1f9a5bd86e308a406447c43",
+    "fastplms/models/_esm_rotary.py": "sha256:8e041bc4a056373e07fbf598e6acb6075f889d1bce12057715369804befb37ec",
     "fastplms/models/esm2/__init__.py": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
-    "fastplms/models/esm2/modeling_fastesm.py": "sha256:034d3790790294354c5fb561a2c9fe2e4e3851f4b4962f41fc132bf87c465078",
-    "fastplms/models/ttt.py": "sha256:a0df4e98b02120d423e3c7ca9b866a8d0e3748b9076042a0102a838a11aed046",
-    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
-    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
-    "fastplms_bundle.py": "sha256:a86427af9f7da529c4259415bc6e083b65e62afaae44774ee27be8b7b16071f5",
-    "modeling_fastplms.py": "sha256:1a8fdc9190a50bcb69b8d61011d8669bb0bb5e27f1b7817ae08220df9e0996d9",
+    "fastplms/models/esm2/modeling_fastesm.py": "sha256:119462776260e6eea6409c7203491d9100bfa51418986e9354c2b12f3684d31b",
+    "fastplms/models/ttt.py": "sha256:f1d7f9298a930f4c40f8e321ffa08b09d1c5b8492bca623bdc7bfd03ec85774d",
+    "fastplms/registry.py": "sha256:cb37799dae5a4c0c780089ab2cca4cc56300c1c5f515b275892aa4e1c4b97eb6",
+    "fastplms/runtime.py": "sha256:9521c37fcd168bf9a4137277fb73672403ccf30d70d4845aa879fb22f42ddeb6",
+    "fastplms_bundle.py": "sha256:96be83e24e3ce995195d5a51e06d360e80ae1a60a5483494d2da74e0c7e601d6",
+    "modeling_fastplms.py": "sha256:ca2bed6fb4e70e8e24f271abed55283c8f179d6794c5d860a11f09a67e4260dc",
+    "requirements.txt": "sha256:a93eb8b0b6c7a4b3453d0af7465d72c1147d29a556be5182fe63a02194006767",
     "special_tokens_map.json": "sha256:3aedcd4211c0d43aec4e607ff60a63255f3174ead795e997350f09a5f8cd9ee1",
     "tokenizer_config.json": "sha256:7e9161ecdb548ec45a41cbc6b24aa4476fdd418461f491c4207baa99419a29ad",
     "vocab.txt": "sha256:0b82cc0a7c7cf9e567b1e5892d793285b9fbae822c964ca48696f7db44598e03"
   },
   "model_id": "esm2_3b",
   "redistributable": true,
-  "release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "runtime_bundle_sha256": "0e9cee71faeeb3dd6406d10d4107c209b10379ab0ee07d0115b13102b6525460",
-  "runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
+  "release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "runtime_bundle_sha256": "f22353ef386f607784889dc303488a2dc5e99b009a5e423ab3eff1f1eb03a9a4",
+  "runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
   "schema_version": 2,
   "scope": "runtime-only",
-  "source_tree_sha256": "fb8779cf2f90db290a898b170b6c6d2ebe9835b89dc946cc99a94417cb586aa4",
+  "source_tree_sha256": "818bf2bf5f573adc842b55f4f8feff926ddc25f904c5769ff5c382077d3a78a2",
   "weights": {
     "repo_id": "Synthyra/ESM2-3B",
     "revision": "ff89d0180f414ab9c677219a25da79bf09185456"