Update FastPLMs runtime and model cards

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

README.md

@@ -19,15 +19,34 @@ 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 | Special: tokenizer-free raw-sequence preparation |
+| Test-time training | Supported: low-rank masked-residue adaptation |
+| Attention variants | Supported: `sdpa`, `flex_attention` |
+| 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/Profluent-E1-600M/resolve/main/requirements.txt"
 ```
 
+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 declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. 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.
@@ -41,62 +60,137 @@ model_id = "Synthyra/Profluent-E1-600M"
 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/Profluent-E1-600M` path, then pass `local_files_only=True`.
+For offline validation, replace `model_id` with the manifest-built
+`dist/hub/Profluent-E1-600M` path and pass `local_files_only=True`.
+
+## Attention and compliance
 
-Leave attention unspecified for the Transformers default. Supported explicit
-choices are `sdpa`, `flex_attention`.
-Pass the selected name through `attn_implementation`.
-When an optimized backend cannot return full attention tensors,
-`output_attentions=True` emits one explicit runtime warning and uses a correctly
-masked eager implementation for that call only. The warning identifies the
-configured backend, effective backend, and reason. Configuration and later
-calls are unchanged.
-For BF16 execution, this family uses parameters loaded directly in BF16.
+The quick start selects `sdpa` explicitly. Declared variants are `sdpa`, `flex_attention`. 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.
 
 ## 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 (
+    AutoModelForSequenceClassification,
+    AutoModelForTokenClassification,
+)
+
+model_id = "Synthyra/Profluent-E1-600M"
+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()
+sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"]
+batch = sequence_model.prep_tokens.get_batch_kwargs(
+    sequences,
+    device=sequence_model.device,
+)
+biological = batch["sequence_ids"].ne(-1)
+
+sequence_labels = torch.zeros(len(sequences), dtype=torch.long)
+token_labels = torch.full_like(batch["input_ids"], -100)
+token_labels[biological] = 0
 
-for record in result:
-    print(record.id, record.sequence, record.tensor.shape)
+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)
 ```
 
-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.
+## PEFT fine-tuning
+
+Install the direct training dependencies, then attach LoRA to the loaded checkpoint:
 
-For a long FASTA run, stream completed batches into SQLite:
+```bash
+python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
+```
 
 ```python
-persisted = model.embed_dataset(
-    "proteins.fasta",
-    batch_size=64,
-    pooling=("mean",),
-    output="protein-embeddings.sqlite",
-    format="sqlite",
-    resume=True,
+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"],
+    ),
 )
 ```
 
-Resume verifies the input order, model state, tokenizer policy, backend, dtype,
-and pooling configuration. It never appends incompatible records to an
-existing run.
+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
+
+TTT samples masked views of one protein and updates only injected low-rank
+adapters. Base checkpoint weights remain frozen:
+
+```python
+from transformers import AutoModelForMaskedLM
+
+ttt_model = AutoModelForMaskedLM.from_pretrained(
+    "Synthyra/Profluent-E1-600M",
+    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)
+```
+
+Persisted adapters retain their deterministic reset state. TTT adds latency
+and memory, can worsen an output, and does not establish biological function.
 
 ## Tokenizer-free E1 input
 
@@ -126,7 +220,7 @@ the attribution required by the upstream agreement.
 - Precision policies: `default`
 - BF16 execution: `static_parameters`
 - Generation contract: `not_applicable`
-- Optional dependency group: `core`
+- Artifact dependency set: `core`
 - Weight publication allowed: `true`
 - Weight license status: `resolved`
 - Redistributable: `true`
@@ -137,28 +231,22 @@ the attribution required by the upstream agreement.
 - FastPLMs weights: `Synthyra/Profluent-E1-600M`
 - 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: `Profluent-Bio/E1-600m`
 - Artifact source: `fast`
 - State transform: `e1_to_fastplms_v1`
-- BF16 execution: `static_parameters`
 - Pinned upstreams: `e1`
-- Reference container: `reference-e1`
 - 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

@@ -18,11 +18,11 @@
   "fastplms_checkpoint_repo_id": "Synthyra/Profluent-E1-600M",
   "fastplms_checkpoint_revision": "6c8bf0ec83b0e0178677c528b101efffd0677742",
   "fastplms_model_id": "e1_600m",
-  "fastplms_release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "fastplms_runtime_bundle_sha256": "5289065bf7679630038a2e794052e4364a152c4f103c36c80ab9b8544919e924",
-  "fastplms_runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_source_tree_sha256": "c583bf72e969d20248d4f2a01c81579fffa3002e2e19c4e90acebbbacf1faeab",
+  "fastplms_release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "fastplms_runtime_bundle_sha256": "a8dda20d6a4a55e6808c7c2f2c9477c113edaa7d6ec6ef57f4a898cd47b73d48",
+  "fastplms_runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_source_tree_sha256": "9e3cbab448ffde4fbbba79f17cf0321645584acd534c5f96c78b1017d9f3806b",
   "fastplms_weights_revision": "6c8bf0ec83b0e0178677c528b101efffd0677742",
   "gated_mlp": true,
   "global_attention_every_n_layers": 3,

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/e1/attention.py

@@ -3,9 +3,8 @@
 from __future__ import annotations
 
 import os
-from collections.abc import Callable
-
 import torch
+from collections.abc import Callable
 from torch.nn.attention.flex_attention import _create_sparse_block_from_block_mask
 
 from fastplms.attention import (
@@ -23,6 +22,7 @@ from fastplms.attention import (
 
 @torch.compiler.disable
 def create_block_causal_mask_optimized(sequence_ids: torch.Tensor) -> BlockMask:
+    # sequence_ids: (b, l)
     if create_block_mask is None:
         raise RuntimeError("Flex Attention block-mask creation is unavailable in this environment.")
     # Assumes sequence_ids is sorted in increasing order for each batch item, except for
@@ -42,6 +42,7 @@ def create_block_causal_mask_optimized(sequence_ids: torch.Tensor) -> BlockMask:
 
 @torch.compiler.disable
 def create_within_seq_block_mask(sequence_ids: torch.Tensor) -> BlockMask:
+    # sequence_ids: (b, l)
     if create_block_mask is None:
         raise RuntimeError("Flex Attention block-mask creation is unavailable in this environment.")
     def document_mask(b, h, q_idx, kv_idx):  # type: ignore[no-untyped-def]
@@ -58,17 +59,19 @@ def create_within_seq_block_mask(sequence_ids: torch.Tensor) -> BlockMask:
 
 
 def build_within_seq_mask_4d(sequence_ids: torch.Tensor) -> torch.Tensor:
-    not_pad = sequence_ids != -1
-    same_seq = sequence_ids.unsqueeze(-1) == sequence_ids.unsqueeze(-2)
-    valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2)
-    return (same_seq & valid).unsqueeze(1)
+    # sequence_ids: (b, l)
+    not_pad = sequence_ids != -1  # (b, l)
+    same_seq = sequence_ids.unsqueeze(-1) == sequence_ids.unsqueeze(-2)  # (b, l, l)
+    valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2)  # (b, l, l)
+    return (same_seq & valid).unsqueeze(1)  # (b, 1, l, l)
 
 
 def build_block_causal_mask_4d(sequence_ids: torch.Tensor) -> torch.Tensor:
-    not_pad = sequence_ids != -1
-    causal = sequence_ids.unsqueeze(-1) >= sequence_ids.unsqueeze(-2)
-    valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2)
-    return (causal & valid).unsqueeze(1)
+    # sequence_ids: (b, l)
+    not_pad = sequence_ids != -1  # (b, l)
+    causal = sequence_ids.unsqueeze(-1) >= sequence_ids.unsqueeze(-2)  # (b, l, l)
+    valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2)  # (b, l, l)
+    return (causal & valid).unsqueeze(1)  # (b, 1, l, l)
 
 
 def flex_attention_func(
@@ -84,9 +87,9 @@ def flex_attention_func(
         raise RuntimeError("Flex Attention is not available in this environment.")
     if score_mod is not None:
         raise NotImplementedError("E1 Flex Attention does not support score_mod.")
-    query_states = query_states.transpose(1, 2).contiguous()  # (bs, nh, seqlen, hs)
-    key_states = key_states.transpose(1, 2).contiguous()  # (bs, nkv, seqlen, hs)
-    value_states = value_states.transpose(1, 2).contiguous()  # (bs, nkv, seqlen, hs)
+    query_states = query_states.transpose(1, 2).contiguous()  # (b, h, l, d)
+    key_states = key_states.transpose(1, 2).contiguous()  # (b, h_kv, l, d)
+    value_states = value_states.transpose(1, 2).contiguous()  # (b, h_kv, l, d)
 
     fn = _get_flex_attention_fn(
         device=query_states.device,
@@ -97,7 +100,7 @@ def flex_attention_func(
     )
     if fn is None:
         raise RuntimeError("Flex Attention is not available in this environment.")
-    outputs = fn(
+    outputs = fn(  # (b, h, l, d)
         query_states,
         key_states,
         value_states,
@@ -106,19 +109,20 @@ def flex_attention_func(
         enable_gqa=query_states.shape[1] != key_states.shape[1],  # if nkv != nh
     )
 
-    outputs = outputs.transpose(1, 2)  # (bs, seqlen, nh, hs)
-    return outputs
+    outputs = outputs.transpose(1, 2)  # (b, l, h, d)
+    return outputs  # (b, l, h, d)
 
 
 def kernels_flash_attention_func(
-    query_states: torch.Tensor,  # (bs, seqlen, nh, hs)
-    key_states: torch.Tensor,  # (bs, seqlen, nkv, hs)
-    value_states: torch.Tensor,  # (bs, seqlen, nkv, hs)
+    query_states: torch.Tensor,  # (b, l_q, h, d)
+    key_states: torch.Tensor,  # (b, l_kv, h_kv, d)
+    value_states: torch.Tensor,  # (b, l_kv, h_kv, d)
     q_sequence_ids: torch.Tensor,
     k_sequence_ids: torch.Tensor,
     causal: bool = False,
     implementation: str = "flash_attention_3",
-) -> torch.Tensor:  # (bs, seqlen, nh, hs)
+) -> torch.Tensor:  # (b, l_q, h, d)
+    # q_sequence_ids: (b, l_q); k_sequence_ids: (b, l_kv)
     _ensure_flash_kernels_loaded(implementation)
 
     if not causal:
@@ -130,9 +134,15 @@ def kernels_flash_attention_func(
             indices_q,
             (cu_seqlens_q, cu_seqlens_k),
             (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
-        ) = _unpad_input(query_states, key_states, value_states, q_sequence_ids, k_sequence_ids)
+        ) = _unpad_input(  # Q: (t_q, h, d); K/V: (t_kv, h_kv, d)
+            query_states,
+            key_states,
+            value_states,
+            q_sequence_ids,
+            k_sequence_ids,
+        )
 
-        attn_output_unpad = _kernels_flash_varlen_forward(
+        attn_output_unpad = _kernels_flash_varlen_forward(  # (t_q, h, d)
             query_states,
             key_states,
             value_states,
@@ -143,14 +153,19 @@ def kernels_flash_attention_func(
             causal=False,
             implementation=implementation,
         )
-        attn_output = pad_input(attn_output_unpad, indices_q, batch_size, q_len)
+        attn_output = pad_input(  # (b, l_q, h, d)
+            attn_output_unpad,
+            indices_q,
+            batch_size,
+            q_len,
+        )
 
     else:
-        attn_output = _kernels_flash_forward(
+        attn_output = _kernels_flash_forward(  # (b, l_q, h, d)
             query_states, key_states, value_states, causal=True, implementation=implementation
         )
 
-    return attn_output
+    return attn_output  # (b, l_q, h, d)
 
 
 def block_min_max_seq_ids(
@@ -159,7 +174,8 @@ def block_min_max_seq_ids(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Map each physical attention block to its first and last sequence."""
 
-    total_tokens = sequence_lengths.sum()
+    # sequence_lengths: (n,)
+    total_tokens = sequence_lengths.sum()  # ()
     block_count = int(
         torch.div(
             total_tokens + block_size - 1,
@@ -167,22 +183,22 @@ def block_min_max_seq_ids(
             rounding_mode="floor",
         ).item()
     )
-    padded_tokens = block_count * block_size - total_tokens
-    lengths_with_tail = torch.cat(
+    padded_tokens = block_count * block_size - total_tokens  # ()
+    lengths_with_tail = torch.cat(  # (n + 1,)
         (sequence_lengths, padded_tokens.to(sequence_lengths).reshape(1)),
     )
-    sequence_ends = lengths_with_tail.to(torch.long).cumsum(dim=0)
-    block_starts = torch.arange(
+    sequence_ends = lengths_with_tail.to(torch.long).cumsum(dim=0)  # (n + 1,)
+    block_starts = torch.arange(  # (n_blocks,)
         start=0,
         end=block_count * block_size,
         step=block_size,
         dtype=torch.long,
         device=sequence_lengths.device,
     )
-    block_last_tokens = block_starts + block_size - 1
-    first_sequence = torch.searchsorted(sequence_ends, block_starts, right=True)
-    last_sequence = torch.searchsorted(sequence_ends, block_last_tokens, right=True)
-    return first_sequence, last_sequence
+    block_last_tokens = block_starts + block_size - 1  # (n_blocks,)
+    first_sequence = torch.searchsorted(sequence_ends, block_starts, right=True)  # (n_blocks,)
+    last_sequence = torch.searchsorted(sequence_ends, block_last_tokens, right=True)  # (n_blocks,)
+    return first_sequence, last_sequence  # (n_blocks,), (n_blocks,)
 
 
 def get_overlapping_blocks(
@@ -191,32 +207,42 @@ def get_overlapping_blocks(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Classify query/key block pairs as full, partial, or disjoint."""
 
-    q_first, q_last = block_min_max_seq_ids(q_lengths)
-    k_first, k_last = block_min_max_seq_ids(k_lengths)
-    intersection_start = torch.maximum(q_first[:, None], k_first[None, :])
-    intersection_end = torch.minimum(q_last[:, None], k_last[None, :])
-    intersects = intersection_start <= intersection_end
-    both_blocks_are_single_sequence = (q_first == q_last)[:, None] & (k_first == k_last)[None, :]
-    full_blocks = intersects & both_blocks_are_single_sequence
-    return full_blocks, intersects & ~both_blocks_are_single_sequence
+    # q_lengths: (n_q,); k_lengths: (n_k,)
+    q_first, q_last = block_min_max_seq_ids(q_lengths)  # (q_blocks,), (q_blocks,)
+    k_first, k_last = block_min_max_seq_ids(k_lengths)  # (k_blocks,), (k_blocks,)
+    intersection_start = torch.maximum(  # (q_blocks, k_blocks)
+        q_first[:, None],
+        k_first[None, :],
+    )
+    intersection_end = torch.minimum(  # (q_blocks, k_blocks)
+        q_last[:, None],
+        k_last[None, :],
+    )
+    intersects = intersection_start <= intersection_end  # (q_blocks, k_blocks)
+    both_blocks_are_single_sequence = (  # (q_blocks, k_blocks)
+        (q_first == q_last)[:, None] & (k_first == k_last)[None, :]
+    )
+    full_blocks = intersects & both_blocks_are_single_sequence  # (q_blocks, k_blocks)
+    return full_blocks, intersects & ~both_blocks_are_single_sequence  # both (q_blocks, k_blocks)
 
 
 def _document_ids(sequence_lengths: torch.Tensor) -> torch.Tensor:
-    sequence_numbers = torch.arange(
+    # sequence_lengths: (n,)
+    sequence_numbers = torch.arange(  # (n,)
         sequence_lengths.numel(),
         device=sequence_lengths.device,
         dtype=torch.long,
     )
-    return sequence_numbers.repeat_interleave(sequence_lengths.to(torch.long))
+    return sequence_numbers.repeat_interleave(sequence_lengths.to(torch.long))  # (t,)
 
 
 @torch.compiler.disable
 def direct_block_mask(q_lengths: torch.Tensor, k_lengths: torch.Tensor) -> BlockMask:
     """Build a packed-sequence mask from preclassified sparse blocks."""
 
-    full, partial = get_overlapping_blocks(q_lengths, k_lengths)
-    q_document = _document_ids(q_lengths)
-    k_document = _document_ids(k_lengths)
+    full, partial = get_overlapping_blocks(q_lengths, k_lengths)  # both (q_blocks, k_blocks)
+    q_document = _document_ids(q_lengths)  # (t_q,)
+    k_document = _document_ids(k_lengths)  # (t_k,)
 
     def same_document(
         _batch: torch.Tensor,
@@ -239,8 +265,8 @@ def direct_block_mask(q_lengths: torch.Tensor, k_lengths: torch.Tensor) -> Block
 def doc_id_mask(q_lengths: torch.Tensor, k_lengths: torch.Tensor) -> BlockMask:
     if create_block_mask is None:
         raise RuntimeError("Flex Attention block-mask creation is unavailable in this environment.")
-    q_document = _document_ids(q_lengths)
-    k_document = _document_ids(k_lengths)
+    q_document = _document_ids(q_lengths)  # (t_q,)
+    k_document = _document_ids(k_lengths)  # (t_k,)
 
     def same_document(
         _batch: torch.Tensor,
@@ -268,6 +294,8 @@ def varlen_flex_attention_func(
     q_sequence_ids: torch.Tensor,
     k_sequence_ids: torch.Tensor,
 ) -> torch.Tensor:
+    # query_states: (b, l_q, h, d); key_states, value_states: (b, l_kv, h_kv, d)
+    # q_sequence_ids: (b, l_q); k_sequence_ids: (b, l_kv)
     if flex_attention is None:
         raise RuntimeError("Flex Attention is not available in this environment.")
     batch_size, q_len = query_states.shape[0], query_states.shape[1]
@@ -278,14 +306,20 @@ def varlen_flex_attention_func(
         indices_q,
         (cu_seqlens_q, cu_seqlens_k),
         (_max_seqlen_in_batch_q, _max_seqlen_in_batch_k),
-    ) = _unpad_input(query_states, key_states, value_states, q_sequence_ids, k_sequence_ids)
+    ) = _unpad_input(  # Q: (t_q, h, d); K/V: (t_kv, h_kv, d)
+        query_states,
+        key_states,
+        value_states,
+        q_sequence_ids,
+        k_sequence_ids,
+    )
 
-    query_states = query_states.unsqueeze(0).transpose(1, 2).contiguous()
-    key_states = key_states.unsqueeze(0).transpose(1, 2).contiguous()
-    value_states = value_states.unsqueeze(0).transpose(1, 2).contiguous()
+    query_states = query_states.unsqueeze(0).transpose(1, 2).contiguous()  # (1, h, t_q, d)
+    key_states = key_states.unsqueeze(0).transpose(1, 2).contiguous()  # (1, h_kv, t_kv, d)
+    value_states = value_states.unsqueeze(0).transpose(1, 2).contiguous()  # (1, h_kv, t_kv, d)
 
-    seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
-    seqlens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1]
+    seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]  # (n,)
+    seqlens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1]  # (n,)
     block_mask = block_mask_creator(seqlens_q, seqlens_k)
 
     packed_lengths = (
@@ -302,7 +336,7 @@ def varlen_flex_attention_func(
     )
     if fn is None:
         raise RuntimeError("Flex Attention is not available in this environment.")
-    attn_output_unpad = fn(
+    attn_output_unpad = fn(  # (1, h, t_q, d)
         query_states,
         key_states,
         value_states,
@@ -310,39 +344,44 @@ def varlen_flex_attention_func(
         enable_gqa=query_states.shape[1] != key_states.shape[1],
     )
 
-    attn_output = pad_input(
+    attn_output = pad_input(  # (b, l_q, h, d)
         attn_output_unpad.transpose(1, 2).squeeze(0), indices_q, batch_size, q_len
     )
 
-    return attn_output
+    return attn_output  # (b, l_q, h, d)
 
 
 def _get_unpad_data(sequence_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
     """Return packed indices and run lengths for the non-padding sequence IDs."""
 
-    flat_ids = sequence_ids.reshape(-1)
-    non_pad_indices = torch.where(flat_ids.ne(-1))[0]
+    # sequence_ids: (b, l)
+    flat_ids = sequence_ids.reshape(-1)  # (b * l,)
+    non_pad_indices = torch.where(flat_ids.ne(-1))[0]  # (t,)
     if non_pad_indices.numel() == 0:
         raise ValueError("Packed attention requires at least one non-padding token.")
 
-    valid_ids = flat_ids.index_select(0, non_pad_indices)
-    row_ids = torch.div(
+    valid_ids = flat_ids.index_select(0, non_pad_indices)  # (t,)
+    row_ids = torch.div(  # (t,)
         non_pad_indices,
         sequence_ids.shape[1],
         rounding_mode="floor",
     )
-    run_starts = torch.ones_like(valid_ids, dtype=torch.bool)
+    run_starts = torch.ones_like(valid_ids, dtype=torch.bool)  # (t,)
     run_starts[1:] = (valid_ids[1:] != valid_ids[:-1]) | (row_ids[1:] != row_ids[:-1])
-    start_indices = torch.where(run_starts)[0]
-    end_indices = torch.cat((start_indices[1:], start_indices.new_tensor([valid_ids.numel()])))
-    sequence_lengths = end_indices - start_indices
-    cumulative_lengths = torch.cat(
+    start_indices = torch.where(run_starts)[0]  # (n,)
+    end_indices = torch.cat(  # (n,)
+        (start_indices[1:], start_indices.new_tensor([valid_ids.numel()]))
+    )
+    sequence_lengths = end_indices - start_indices  # (n,)
+    cumulative_lengths = torch.cat(  # (n + 1,)
         (
             torch.zeros(1, dtype=torch.int32, device=sequence_ids.device),
             sequence_lengths.cumsum(dim=0, dtype=torch.int32),
         ),
     )
-    return non_pad_indices, cumulative_lengths, int(sequence_lengths.max().item())
+    return non_pad_indices, cumulative_lengths, int(  # (t,), (n + 1,), scalar
+        sequence_lengths.max().item()
+    )
 
 
 def _unpad_input(
@@ -359,6 +398,8 @@ def _unpad_input(
     tuple[torch.Tensor, torch.Tensor],
     tuple[int, int],
 ]:
+    # query_layer: (b, l_q, h, d); key_layer, value_layer: (b, l_kv, h_kv, d)
+    # q_sequence_ids: (b, l_q); k_sequence_ids: (b, l_kv)
     for name, layer in (
         ("query_layer", query_layer),
         ("key_layer", key_layer),
@@ -402,12 +443,14 @@ def _unpad_input(
             f"{query_length} > {kv_seq_len}"
         )
 
-    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(k_sequence_ids)
+    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(  # (t_kv,), (n + 1,), scalar
+        k_sequence_ids
+    )
 
-    key_layer = index_first_axis(
+    key_layer = index_first_axis(  # (t_kv, h_kv, d)
         key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
     )
-    value_layer = index_first_axis(
+    value_layer = index_first_axis(  # (t_kv, h_kv, d)
         value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
     )
 
@@ -416,9 +459,10 @@ def _unpad_input(
         cu_seqlens_q = cu_seqlens_k
         max_seqlen_in_batch_q = max_seqlen_in_batch_k
     else:
+        # (t_q,), (n + 1,), scalar
         indices_q, cu_seqlens_q, max_seqlen_in_batch_q = _get_unpad_data(q_sequence_ids)
 
-    query_layer = index_first_axis(
+    query_layer = index_first_axis(  # (t_q, h, d)
         query_layer.reshape(batch_size * query_length, num_q_heads, head_dim), indices_q
     )
 

fastplms/models/e1/cache.py

@@ -2,9 +2,8 @@
 
 from __future__ import annotations
 
-from typing import Any
-
 import torch
+from typing import Any
 from transformers.modeling_outputs import ModelOutput
 from transformers.utils import logging
 
@@ -44,9 +43,9 @@ class DynamicCache:
             tuple[`torch.Tensor`, `torch.Tensor`]: Cached K and V, each with shape
             (b, l, h, d).
         """
-        # Lazy initialization
+        # key_states, value_states: (b, l_new, h, d)
         if len(self.key_cache) <= layer_idx:
-            # There may be skipped layers, fill them with empty lists
+            # Empty tensors preserve skipped layer indices until those layers receive state.
             for _ in range(len(self.key_cache), layer_idx):
                 self.key_cache.append(torch.tensor([]))
                 self.value_cache.append(torch.tensor([]))
@@ -60,12 +59,18 @@ class DynamicCache:
             self.key_cache[layer_idx] = key_states
             self.value_cache[layer_idx] = value_states
         else:
-            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=1)
-            self.value_cache[layer_idx] = torch.cat(
+            self.key_cache[layer_idx] = torch.cat(  # (b, l_cached + l_new, h, d)
+                [self.key_cache[layer_idx], key_states],
+                dim=1,
+            )
+            self.value_cache[layer_idx] = torch.cat(  # (b, l_cached + l_new, h, d)
                 [self.value_cache[layer_idx], value_states], dim=1
             )
 
-        return self.key_cache[layer_idx], self.value_cache[layer_idx]
+        return (  # (b, l_total, h, d), (b, l_total, h, d)
+            self.key_cache[layer_idx],
+            self.value_cache[layer_idx],
+        )
 
     def get_seq_length(self, layer_idx: int = 0) -> int:
         """Return the cached sequence length for one layer."""
@@ -95,19 +100,22 @@ class DynamicCache:
         """Repeat the cache `repeats` times in the batch dimension. Used in contrastive search."""
         for layer_idx in range(len(self.key_cache)):
             if self.key_cache[layer_idx].numel():
+                # (b * repeats, l, h, d)
                 self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(
                     repeats, dim=0
                 )
                 self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(
                     repeats, dim=0
-                )
+                )  # (b * repeats, l, h, d)
 
-    def batch_select_indices(self, indices: torch.Tensor) -> None:
+    def batch_select_indices(self, indices: torch.Tensor | list[int]) -> None:
         """Keep selected rows of the cache batch dimension."""
         for layer_idx in range(len(self.key_cache)):
             if self.key_cache[layer_idx].numel():
-                self.key_cache[layer_idx] = self.key_cache[layer_idx][indices, ...]
-                self.value_cache[layer_idx] = self.value_cache[layer_idx][indices, ...]
+                self.key_cache[layer_idx] = self.key_cache[layer_idx][indices, ...]  # (n, l, h, d)
+                self.value_cache[layer_idx] = self.value_cache[layer_idx][
+                    indices, ...
+                ]  # (n, l, h, d)
 
 
 class KVCache:
@@ -158,7 +166,7 @@ class KVCache:
             )
             return
 
-        batch_size = batch["input_ids"].shape[0]
+        batch_size = batch["input_ids"].shape[0]  # b
 
         unique_context = contexts[0]
         unique_context_len = context_lens[0]
@@ -171,10 +179,10 @@ class KVCache:
         past_key_values = self.cache_dict[unique_context]
         batch["past_key_values"] = past_key_values
 
-        # Remove context from the input fields
+        # A cached prefix leaves only query-suffix tokens for the model call.
         for field_name in self.tensor_input_field_names:
             if batch.get(field_name) is not None:
-                batch[field_name] = batch[field_name][:, unique_context_len:]
+                batch[field_name] = batch[field_name][:, unique_context_len:]  # (b, l_suffix, ...)
 
     def after_forward(self, batch: dict[str, Any], outputs: ModelOutput) -> None:
         contexts = batch.get("context")
@@ -209,18 +217,21 @@ class KVCache:
             self.cache_dict[unique_context] = past_key_values
             self.cache_queue.append(unique_context)
 
-            # Remove context from the input fields
+            # The first uncached call returns the full sequence; expose its query suffix.
             for field_name in self.tensor_input_field_names:
                 if field_name in batch and batch[field_name] is not None:
-                    batch[field_name] = batch[field_name][:, unique_context_len:]
+                    batch[field_name] = batch[field_name][
+                        :, unique_context_len:
+                    ]  # (b, l_suffix, ...)
 
-            # Remove context from the output fields
             for field_name in self.tensor_output_field_names:
                 if field_name in outputs and outputs[field_name] is not None:
-                    outputs[field_name] = outputs[field_name][:, unique_context_len:]
+                    outputs[field_name] = outputs[field_name][
+                        :, unique_context_len:
+                    ]  # (b, l_suffix, ...)
             if "hidden_states" in outputs and outputs["hidden_states"] is not None:
                 hidden_states = outputs["hidden_states"]
-                sliced_hidden_states = tuple(
+                sliced_hidden_states = tuple(  # each: (b, l_suffix, d)
                     hidden_state[:, unique_context_len:] for hidden_state in hidden_states
                 )
                 outputs["hidden_states"] = sliced_hidden_states

fastplms/models/e1/modeling_e1.py

@@ -3,21 +3,21 @@ from __future__ import annotations
 import hashlib
 import os
 import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
 from collections import defaultdict
 from contextvars import ContextVar
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any, ClassVar, TypedDict
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
 from tqdm.auto import tqdm
 from transformers import PretrainedConfig, PreTrainedModel
 from transformers.activations import ACT2FN
 from transformers.modeling_outputs import ModelOutput
 from transformers.utils import logging
 
+
 try:
     from fastplms.attention import (
         AttentionBackend,
@@ -282,7 +282,7 @@ class RotaryPositionalEmbedding(nn.Module):
         max_position_embeddings: int = 2048,
         base: int = 10000,
         device: torch.device | None = None,
-    ):
+    ) -> None:
         super().__init__()
 
         self.dim = dim
@@ -311,11 +311,11 @@ class RotaryPositionalEmbedding(nn.Module):
         self.max_seq_len_cached = seq_len
         inv_freq = self.base ** -(
             torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim
-        )
+        )  # (d / 2,)
         self.inv_freq = inv_freq
-        t = torch.arange(seq_len, device=device, dtype=torch.float32)
-        angles = torch.outer(t, inv_freq)
-        angles = torch.cat((angles, angles), dim=1)
+        t = torch.arange(seq_len, device=device, dtype=torch.float32)  # (l,)
+        angles = torch.outer(t, inv_freq)  # (l, d / 2)
+        angles = torch.cat((angles, angles), dim=1)  # (l, d)
         self.cos_cached = angles.cos()
         self.sin_cached = angles.sin()
 
@@ -326,18 +326,17 @@ class RotaryPositionalEmbedding(nn.Module):
         position_ids: torch.LongTensor,
         seq_len: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        # Q and K have shape (b, l, h, d).
+        # q, k: (b, l, h, d)
         device, dtype = q.device, q.dtype
         seq_len = position_ids.max().item() + 1 if seq_len is None else seq_len
 
         if seq_len > self.max_seq_len_cached:
             self._set_sin_cos_cache(seq_len=seq_len, device=device)
 
-        # Selecting by position gives C and S shape (b, l, d). Insert a head
-        # axis so they broadcast over Q and K with shape (b, l, h, d).
+        # Selecting by position inserts a head axis for broadcasting.
         idxs = position_ids.to(device)
-        cos = self.cos_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs]
-        sin = self.sin_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs]
+        cos = self.cos_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs]  # (b, l, 1, d)
+        sin = self.sin_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs]  # (b, l, 1, d)
 
         # Apply the real and imaginary parts of the rotary transform to Q and K.
         # Both halves reuse C and S, so rotate_half supplies the cross terms.
@@ -349,7 +348,7 @@ class RotaryPositionalEmbedding(nn.Module):
 class Attention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper."""
 
-    def __init__(self, config: E1Config, layer_idx: int):
+    def __init__(self, config: E1Config, layer_idx: int) -> None:
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
@@ -407,14 +406,21 @@ class Attention(nn.Module):
         past_key_value: DynamicCache | None = None,
         use_cache: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # hidden_states: (b, l, d); position_ids: (b, l)
         bsz, q_len, _ = hidden_states.size()
         query_states: torch.Tensor = self.q_proj(hidden_states)
         key_states: torch.Tensor = self.k_proj(hidden_states)
         val_states: torch.Tensor = self.v_proj(hidden_states)
 
-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)
-        key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim)
-        val_states = val_states.view(bsz, q_len, self.num_kv_heads, self.head_dim)
+        query_states = query_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        )  # (b, l, h, d_h)
+        key_states = key_states.view(
+            bsz, q_len, self.num_kv_heads, self.head_dim
+        )  # (b, l, h_kv, d_h)
+        val_states = val_states.view(
+            bsz, q_len, self.num_kv_heads, self.head_dim
+        )  # (b, l, h_kv, d_h)
 
         if self.clip_qkv is not None:
             query_states = query_states.clamp(-self.clip_qkv, self.clip_qkv)
@@ -456,6 +462,7 @@ class Attention(nn.Module):
         use_cache: bool = False,
         effective_backend: AttentionBackend | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None, DynamicCache | None, list[torch.Tensor] | None]:
+        # hidden_states: (b, l, d); position and sequence IDs: (b, l)
         is_cache_prefilled = (
             use_cache
             and past_key_value is not None
@@ -586,12 +593,12 @@ class Attention(nn.Module):
         query_states: torch.Tensor,  # Q has shape (b, l, h, d).
         key_states: torch.Tensor,  # K has shape (b, l, h_kv, d).
     ) -> list[torch.Tensor]:
-        query_heads = query_states.transpose(1, 2).contiguous()
-        key_heads = key_states.transpose(1, 2).contiguous()
+        query_heads = query_states.transpose(1, 2).contiguous()  # (b, h, l, d_h)
+        key_heads = key_states.transpose(1, 2).contiguous()  # (b, h_kv, l, d_h)
         key_heads = repeat_kv(key_heads, self.num_key_value_groups)
         scale = 1.0 / (self.head_dim**0.5)
-        q_norm = torch.linalg.vector_norm(query_heads, dim=-1)
-        k_norm = torch.linalg.vector_norm(key_heads, dim=-1)
+        q_norm = torch.linalg.vector_norm(query_heads, dim=-1)  # (b, h, l)
+        k_norm = torch.linalg.vector_norm(key_heads, dim=-1)  # (b, h, l)
         s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(
             dim=0
         ).values * scale
@@ -763,14 +770,14 @@ class Attention(nn.Module):
         else:
             attention_mask_4d = None
 
-        query_heads = query_states.transpose(1, 2).contiguous()
-        key_heads = key_states.transpose(1, 2).contiguous()
-        value_heads = val_states.transpose(1, 2).contiguous()
+        query_heads = query_states.transpose(1, 2).contiguous()  # (b, h, l, d_h)
+        key_heads = key_states.transpose(1, 2).contiguous()  # (b, h_kv, l, d_h)
+        value_heads = val_states.transpose(1, 2).contiguous()  # (b, h_kv, l, d_h)
         key_heads = repeat_kv(key_heads, self.num_key_value_groups)
         value_heads = repeat_kv(value_heads, self.num_key_value_groups)
         context_heads = F.scaled_dot_product_attention(
             query_heads, key_heads, value_heads, attn_mask=attention_mask_4d
-        )
+        )  # (b, h, l, d_h)
         attn_output = (
             context_heads.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous()
         )
@@ -807,13 +814,15 @@ class Attention(nn.Module):
         else:
             attention_mask_4d = None
 
-        query_heads = query_states.transpose(1, 2).contiguous()
-        key_heads = key_states.transpose(1, 2).contiguous()
-        value_heads = val_states.transpose(1, 2).contiguous()
+        query_heads = query_states.transpose(1, 2).contiguous()  # (b, h, l, d_h)
+        key_heads = key_states.transpose(1, 2).contiguous()  # (b, h_kv, l, d_h)
+        value_heads = val_states.transpose(1, 2).contiguous()  # (b, h_kv, l, d_h)
         key_heads = repeat_kv(key_heads, self.num_key_value_groups)
         value_heads = repeat_kv(value_heads, self.num_key_value_groups)
         scale = 1.0 / (self.head_dim**0.5)
-        attn_weights = torch.matmul(query_heads, key_heads.transpose(-2, -1)) * scale
+        attn_weights = (
+            torch.matmul(query_heads, key_heads.transpose(-2, -1)) * scale
+        )  # (b, h, l, l)
         if attention_mask_4d is not None:
             attention_mask_4d = attention_mask_4d.to(dtype=torch.bool)
             attn_weights = attn_weights.masked_fill(
@@ -823,7 +832,7 @@ class Attention(nn.Module):
         attn_weights = F.softmax(attn_weights, dim=-1)
         if attention_mask_4d is not None:
             attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), 0.0)
-        context_heads = torch.matmul(attn_weights, value_heads)
+        context_heads = torch.matmul(attn_weights, value_heads)  # (b, h, l, d_h)
         attn_output = (
             context_heads.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous()
         )
@@ -832,7 +841,7 @@ class Attention(nn.Module):
 
 
 class MLP(nn.Module):
-    def __init__(self, config: E1Config):
+    def __init__(self, config: E1Config) -> None:
         super().__init__()
         self.ffn_dim = config.intermediate_size
         self.hidden_dim = config.hidden_size
@@ -845,7 +854,7 @@ class MLP(nn.Module):
 
 
 class GLUMLP(nn.Module):
-    def __init__(self, config: E1Config):
+    def __init__(self, config: E1Config) -> None:
         super().__init__()
         self.ffn_dim = config.intermediate_size
         self.hidden_dim = config.hidden_size
@@ -861,7 +870,7 @@ class GLUMLP(nn.Module):
 
 
 class FFN(nn.Module):
-    def __init__(self, config: E1Config):
+    def __init__(self, config: E1Config) -> None:
         super().__init__()
         mlp_cls = GLUMLP if config.gated_mlp else MLP
         self.mlp = mlp_cls(config)
@@ -927,7 +936,7 @@ class E1TokenClassificationOutputWithPast(ModelOutput):
 
 
 class RMSNorm(nn.Module):
-    def __init__(self, hidden_size: int, eps: float = 1e-6):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
@@ -941,7 +950,7 @@ class RMSNorm(nn.Module):
 
 
 class NormAttentionNorm(nn.Module):
-    def __init__(self, config: E1Config, layer_idx: int):
+    def __init__(self, config: E1Config, layer_idx: int) -> None:
         super().__init__()
         self.self_attn = Attention(config, layer_idx)
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -988,7 +997,7 @@ class NormAttentionNorm(nn.Module):
 
 
 class DecoderLayer(nn.Module):
-    def __init__(self, config: E1Config, layer_idx: int):
+    def __init__(self, config: E1Config, layer_idx: int) -> None:
         super().__init__()
         self.initializer_range = config.initializer_range
         self.hidden_size = config.hidden_size
@@ -1179,7 +1188,7 @@ class FAST_E1_ENCODER(E1PreTrainedModel, EmbeddingMixin):
     config_class = E1Config
     _is_internal_encoder = True
 
-    def __init__(self, config: E1Config, **kwargs):
+    def __init__(self, config: E1Config, **kwargs) -> None:
         E1PreTrainedModel.__init__(self, config, **kwargs)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
@@ -1540,7 +1549,7 @@ class E1Model(E1PreTrainedModel, EmbeddingMixin):
     config: E1Config
     config_class = E1Config
 
-    def __init__(self, config: E1Config, **kwargs):
+    def __init__(self, config: E1Config, **kwargs) -> None:
         E1PreTrainedModel.__init__(self, config, **kwargs)
         self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs)
         self.post_init()
@@ -1589,7 +1598,7 @@ class E1ForMaskedLM(FastPLMTestTimeTrainingMixin, E1PreTrainedModel, EmbeddingMi
     config: E1Config
     config_class = E1Config
 
-    def __init__(self, config: E1Config, **kwargs):
+    def __init__(self, config: E1Config, **kwargs) -> None:
         E1PreTrainedModel.__init__(self, config, **kwargs)
         self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs)
         self.vocab_size = config.vocab_size
@@ -2092,7 +2101,7 @@ class E1ForSequenceClassification(E1PreTrainedModel, EmbeddingMixin):
     config: E1Config
     config_class = E1Config
 
-    def __init__(self, config: E1Config, **kwargs):
+    def __init__(self, config: E1Config, **kwargs) -> None:
         pooling_types = kwargs.pop("pooling_types", None)
         if pooling_types is None:
             pooling_types = ["mean", "var"]
@@ -2227,7 +2236,7 @@ class E1ForTokenClassification(E1PreTrainedModel, EmbeddingMixin):
     config: E1Config
     config_class = E1Config
 
-    def __init__(self, config: E1Config, **kwargs):
+    def __init__(self, config: E1Config, **kwargs) -> None:
         E1PreTrainedModel.__init__(self, config, **kwargs)
         self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs)
         self.vocab_size = config.vocab_size

fastplms/models/e1/preparation.py

@@ -4,12 +4,12 @@ from __future__ import annotations
 
 import itertools
 import os
-from dataclasses import dataclass
-
 import torch
+from dataclasses import dataclass
 from tokenizers import Tokenizer
 from torch.nn.utils.rnn import pad_sequence
 
+
 PAD_TOKEN_ID = 0
 BOS_TOKEN_ID = 1
 EOS_TOKEN_ID = 2
@@ -100,7 +100,7 @@ class E1BatchPreparer:
         revision: str | None = None,
         token: str | bool | None = None,
         preserve_context_labels: bool = False,
-    ):
+    ) -> None:
         self.tokenizer = tokenizer or get_tokenizer(
             tokenizer_source,
             local_files_only=local_files_only,
@@ -111,7 +111,7 @@ class E1BatchPreparer:
         self.data_prep_config = data_prep_config or DataPrepConfig()
         self.pad_token_id = self.tokenizer.token_to_id("<pad>")
         self.preserve_context_labels = preserve_context_labels
-        self.boundary_token_ids = torch.tensor(
+        self.boundary_token_ids = torch.tensor(  # (5,)
             [self.tokenizer.token_to_id(token) for token in ["<bos>", "<eos>", "1", "2", "<pad>"]]
         ).long()
         self.mask_token = "?"  # nosec
@@ -135,12 +135,11 @@ class E1BatchPreparer:
         device: torch.device | None = None,
         non_blocking: bool = False,
     ) -> dict[str, torch.Tensor | list[str] | list[int]]:
+        # Each sequence encoding contains aligned one-dimensional tensors of length l_i.
         device = torch.device("cpu") if device is None else device
         non_blocking = non_blocking and device.type == "cuda"
         padded_encodings = {}
-        # Note: We use -1 as the padding value for sequence and position ids because the 0 value
-        # is a valid value for sequence and position ids. -1 is then used to distinguish valid
-        # tokens from padding tokens, for example, when doing padding/unpadding for flash attention.
+        # Sequence and position ID zero is valid, so -1 unambiguously marks padding.
         for key, padding_value in {
             "input_ids": self.pad_token_id,
             "sequence_ids": -1,
@@ -148,7 +147,7 @@ class E1BatchPreparer:
             "global_position_ids": -1,
             "labels": self.pad_token_id,
         }.items():
-            padded_encodings[key] = pad_sequence(
+            padded_encodings[key] = pad_sequence(  # (b, l_max)
                 [enc[key] for enc in sequence_encodings],
                 batch_first=True,
                 padding_value=padding_value,
@@ -169,30 +168,33 @@ class E1BatchPreparer:
             )
 
         encodings = tuple(self.prepare_singleseq(item) for item in sequences)
-        token_counts = torch.tensor(
+        token_counts = torch.tensor(  # (n,)
             [encoding["input_ids"].numel() for encoding in encodings],
             dtype=torch.long,
         )
-        input_ids = torch.cat(tuple(encoding["input_ids"] for encoding in encodings))
-        labels = torch.cat(tuple(encoding["labels"] for encoding in encodings))
+        input_ids = torch.cat(tuple(encoding["input_ids"] for encoding in encodings))  # (t,)
+        labels = torch.cat(tuple(encoding["labels"] for encoding in encodings))  # (t,)
         positions = tuple(encoding["position_ids"] for encoding in encodings)
-        within_seq_position_ids = torch.cat(positions)
+        within_seq_position_ids = torch.cat(positions)  # (t,)
 
         # Offsets preserve gaps left by optional X-token removal.
-        position_spans = torch.tensor(
+        position_spans = torch.tensor(  # (n,)
             [int(position_ids[-1].item()) + 1 for position_ids in positions],
             dtype=torch.long,
         )
-        position_offsets = torch.cat(
+        position_offsets = torch.cat(  # (n,)
             (torch.zeros(1, dtype=torch.long), position_spans[:-1]),
         ).cumsum(dim=0)
-        global_position_ids = torch.cat(
+        global_position_ids = torch.cat(  # (t,)
             tuple(
                 position_ids + offset
                 for position_ids, offset in zip(positions, position_offsets, strict=True)
             )
         )
-        sequence_ids = torch.arange(len(encodings), dtype=torch.long).repeat_interleave(
+        sequence_ids = torch.arange(  # (t,)
+            len(encodings),
+            dtype=torch.long,
+        ).repeat_interleave(
             token_counts
         )
 
@@ -241,24 +243,30 @@ class E1BatchPreparer:
             )
 
         symbols = itertools.chain(("<bos>", "1"), sequence, ("2", "<eos>"))
-        tokens = torch.tensor(
+        tokens = torch.tensor(  # (l + 4,)
             [self.vocab[symbol] for symbol in symbols],
             dtype=torch.long,
         )
-        position_ids = torch.arange(tokens.numel(), dtype=torch.long)
+        position_ids = torch.arange(tokens.numel(), dtype=torch.long)  # (l + 4,)
 
         if self.data_prep_config.remove_X_tokens:
-            keep = tokens.ne(self.X_token_id)
-            tokens = tokens[keep]
-            position_ids = position_ids[keep]
-
-        return {"input_ids": tokens, "labels": tokens, "position_ids": position_ids}
+            keep = tokens.ne(self.X_token_id)  # (l + 4,)
+            tokens = tokens[keep]  # (t,)
+            position_ids = position_ids[keep]  # (t,)
+
+        return {  # Each tensor: (t,)
+            "input_ids": tokens,
+            "labels": tokens,
+            "position_ids": position_ids,
+        }
 
     def get_boundary_token_mask(self, tokens: torch.Tensor) -> torch.BoolTensor:
-        return torch.isin(tokens, self.boundary_token_ids.to(tokens.device))
+        # tokens: (...)
+        return torch.isin(tokens, self.boundary_token_ids.to(tokens.device))  # (...)
 
     def get_mask_positions_mask(self, tokens: torch.Tensor) -> torch.BoolTensor:
-        return tokens == self.mask_token_id
+        # tokens: (...)
+        return tokens == self.mask_token_id  # (...)
 
     def validate_sequence(self, sequence: str) -> bool:
         if not isinstance(sequence, str):

fastplms/models/e1/retrieval.py

@@ -19,6 +19,8 @@ import time
 import urllib.error
 import urllib.parse
 import urllib.request
+import numpy as np
+import torch
 from collections import defaultdict, namedtuple
 from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
@@ -26,18 +28,15 @@ from datetime import UTC, datetime
 from email.utils import parsedate_to_datetime
 from pathlib import Path, PurePosixPath, PureWindowsPath
 from typing import TYPE_CHECKING, Any, TypedDict
-
-import numpy as np
-import torch
 from tqdm.auto import tqdm
 from transformers import PreTrainedModel
 from transformers.utils import logging
 
 from fastplms.embeddings import Pooler
-
 from .cache import KVCache
 from .preparation import DataPrepConfig, E1BatchPreparer, get_context
 
+
 if TYPE_CHECKING:
     from .modeling_e1 import E1MaskedLMOutputWithPast
 
@@ -259,19 +258,22 @@ def convert_to_tensor(
             "MSA rows must have equal aligned lengths after removing insertions: "
             f"{sorted(lengths)}"
         )
-    array = np.vstack(
+    array = np.vstack(  # (n, l)
         [np.frombuffer(byte_sequence, dtype=np.uint8) for byte_sequence in byte_sequences]
     )
-    return torch.from_numpy(array).to(device)
+    return torch.from_numpy(array).to(device)  # (n, l)
 
 
 def get_num_neighbors(byte_seqs: torch.ByteTensor, sim_threshold: float = 0.8) -> list[int]:
+    # byte_seqs: (n, l)
     gap_token_id = np.frombuffer(b"-", np.uint8)[0].item()
-    seq_lens = (byte_seqs != gap_token_id).sum(dim=1)
+    seq_lens = (byte_seqs != gap_token_id).sum(dim=1)  # (n,)
     num_neighbors: list[int] = []
     for i in range(byte_seqs.shape[0]):
-        query_non_gaps = byte_seqs[i] != gap_token_id
-        seqs_sim = (byte_seqs[:, query_non_gaps] == byte_seqs[i, query_non_gaps]).sum(
+        query_non_gaps = byte_seqs[i] != gap_token_id  # (l,)
+        seqs_sim = (  # (n,)
+            byte_seqs[:, query_non_gaps] == byte_seqs[i, query_non_gaps]
+        ).sum(
             dim=1
         ) / seq_lens
         num_neighbors.append(int((seqs_sim >= sim_threshold).sum().item()))
@@ -279,7 +281,8 @@ def get_num_neighbors(byte_seqs: torch.ByteTensor, sim_threshold: float = 0.8) -
 
 
 def get_similarity_to_query(byte_seqs: torch.ByteTensor) -> torch.FloatTensor:
-    return (byte_seqs == byte_seqs[0, :]).sum(dim=1) / byte_seqs.shape[1]
+    # byte_seqs: (n, l)
+    return (byte_seqs == byte_seqs[0, :]).sum(dim=1) / byte_seqs.shape[1]  # (n,)
 
 
 def sample_context(
@@ -296,19 +299,19 @@ def sample_context(
     cache_num_neighbors_path: str | None = None,
 ) -> tuple[str, list[str]]:
     msa_sequences = parse_msa(msa_path)
-    msa_as_byte_tensor = convert_to_tensor(msa_sequences, device)
+    msa_as_byte_tensor = convert_to_tensor(msa_sequences, device)  # (n, l)
     if cache_num_neighbors_path is not None and os.path.exists(cache_num_neighbors_path):
-        num_neighbors = np.load(cache_num_neighbors_path)
+        num_neighbors = np.load(cache_num_neighbors_path)  # (n,)
     else:
-        num_neighbors = np.array(
+        num_neighbors = np.array(  # (n,)
             get_num_neighbors(msa_as_byte_tensor, neighbor_similarity_lower_bound)
         )
         if cache_num_neighbors_path is not None:
             np.save(cache_num_neighbors_path, num_neighbors)
 
-    sampling_weights = 1.0 / num_neighbors
-    query_similarity = get_similarity_to_query(msa_as_byte_tensor)
-    filtered_mask = (query_similarity <= max_query_similarity) & (
+    sampling_weights = 1.0 / num_neighbors  # (n,)
+    query_similarity = get_similarity_to_query(msa_as_byte_tensor)  # (n,)
+    filtered_mask = (query_similarity <= max_query_similarity) & (  # (n,)
         query_similarity >= min_query_similarity
     )
     if int(filtered_mask.sum()) < 1:
@@ -317,8 +320,12 @@ def sample_context(
             f"{min_query_similarity} <= query_similarity <= {max_query_similarity}."
         )
 
-    filtered_weights = np.where(filtered_mask.cpu().numpy(), sampling_weights, 0.0)
-    sampled_indices = np.random.default_rng(seed).choice(
+    filtered_weights = np.where(  # (n,)
+        filtered_mask.cpu().numpy(),
+        sampling_weights,
+        0.0,
+    )
+    sampled_indices = np.random.default_rng(seed).choice(  # (n_sampled,)
         len(filtered_weights),
         size=min(max_num_samples, int(filtered_mask.sum())),
         p=filtered_weights / filtered_weights.sum(),
@@ -380,7 +387,9 @@ def sample_multiple_contexts(
                 max_token_length=context_specification.max_token_length,
                 max_query_similarity=context_specification.max_query_similarity,
                 min_query_similarity=context_specification.min_query_similarity,
-                neighbor_similarity_lower_bound=context_specification.neighbor_similarity_lower_bound,
+                neighbor_similarity_lower_bound=(
+                    context_specification.neighbor_similarity_lower_bound
+                ),
                 use_full_sequences_in_context=use_full_sequences_in_context,
                 full_sequences_path=full_sequences_path,
                 seed=seed + i,
@@ -633,16 +642,17 @@ class ContextCache:
 
 
 def compute_ppll(logits: torch.Tensor, token_ids: torch.Tensor) -> float:
+    # logits: (l, c); token_ids: (l,)
     if token_ids.numel() == 0:
         raise ValueError("Cannot score an empty token sequence")
     if token_ids.device != logits.device:
-        token_ids = token_ids.to(logits.device)
+        token_ids = token_ids.to(logits.device)  # (l,)
     if logits.shape[0] != token_ids.shape[0]:
         raise ValueError(
             f"Logits length {logits.shape[0]} != token_ids length {token_ids.shape[0]}"
         )
-    probs = logits.softmax(dim=-1)
-    token_probs = probs.gather(dim=1, index=token_ids.unsqueeze(1)).squeeze(1)
+    probs = logits.softmax(dim=-1)  # (l, c)
+    token_probs = probs.gather(dim=1, index=token_ids.unsqueeze(1)).squeeze(1)  # (l,)
     return float(token_probs.mean().item())
 
 
@@ -716,12 +726,14 @@ class _E1ContextPredictor:
         self, sequences: list[str], sequence_metadata: list[dict[str, str | int]]
     ) -> list[E1Prediction]:
         outputs = self.predict_batch_padded(sequences)
-        outputs["logits"] = outputs["logits"].float()
-        outputs["embeddings"] = outputs["embeddings"].float()
+        outputs["logits"] = outputs["logits"].float()  # (b, l, c)
+        outputs["embeddings"] = outputs["embeddings"].float()  # (b, l, d)
 
-        token_mask = outputs["non_boundary_token_mask"] & outputs["last_sequence_mask"]
+        token_mask = (  # (b, l)
+            outputs["non_boundary_token_mask"] & outputs["last_sequence_mask"]
+        )
         if self.save_masked_positions_only:
-            token_mask = token_mask & outputs["mask_positions_mask"]
+            token_mask = token_mask & outputs["mask_positions_mask"]  # (b, l)
 
         predictions: list[E1Prediction] = []
         for i in range(len(sequences)):
@@ -729,17 +741,21 @@ class _E1ContextPredictor:
             if "context_id" in sequence_metadata[i]:
                 pred["context_id"] = sequence_metadata[i]["context_id"]
             if "logits" in self.fields_to_save:
-                pred["logits"] = outputs["logits"][i, token_mask[i]]
+                pred["logits"] = outputs["logits"][i, token_mask[i]]  # (r_i, c)
                 if not self.keep_predictions_in_gpu:
-                    pred["logits"] = pred["logits"].to("cpu")
+                    pred["logits"] = pred["logits"].to("cpu")  # (r_i, c)
             if "token_embeddings" in self.fields_to_save:
-                pred["token_embeddings"] = outputs["embeddings"][i, token_mask[i]]
+                pred["token_embeddings"] = outputs["embeddings"][i, token_mask[i]]  # (r_i, d)
                 if not self.keep_predictions_in_gpu:
-                    pred["token_embeddings"] = pred["token_embeddings"].to("cpu")
+                    pred["token_embeddings"] = pred["token_embeddings"].to("cpu")  # (r_i, d)
             if "mean_token_embeddings" in self.fields_to_save:
-                pred["mean_token_embeddings"] = outputs["embeddings"][i, token_mask[i]].mean(dim=0)
+                pred["mean_token_embeddings"] = outputs["embeddings"][
+                    i, token_mask[i]
+                ].mean(dim=0)  # (d,)
                 if not self.keep_predictions_in_gpu:
-                    pred["mean_token_embeddings"] = pred["mean_token_embeddings"].to("cpu")
+                    pred["mean_token_embeddings"] = pred["mean_token_embeddings"].to(  # (d,)
+                        "cpu"
+                    )
             predictions.append(pred)
         return predictions
 
@@ -767,12 +783,16 @@ class _E1ContextPredictor:
             if self.kv_cache is not None:
                 self.kv_cache.after_forward(batch, output)
 
-            padding_mask = batch["input_ids"] == self.batch_preparer.pad_token_id
-            last_sequence_mask = (
+            padding_mask = batch["input_ids"] == self.batch_preparer.pad_token_id  # (b, l)
+            last_sequence_mask = (  # (b, l)
                 batch["sequence_ids"] == batch["sequence_ids"].max(dim=1).values[:, None]
             )
-            boundary_token_mask = self.batch_preparer.get_boundary_token_mask(batch["input_ids"])
-            mask_positions_mask = self.batch_preparer.get_mask_positions_mask(batch["input_ids"])
+            boundary_token_mask = self.batch_preparer.get_boundary_token_mask(  # (b, l)
+                batch["input_ids"]
+            )
+            mask_positions_mask = self.batch_preparer.get_mask_positions_mask(  # (b, l)
+                batch["input_ids"]
+            )
             return {
                 "logits": output.logits,
                 "embeddings": output.last_hidden_state,
@@ -817,17 +837,18 @@ def _pool_hidden_states(
     pooling_types: list[str],
     device: torch.device,
 ) -> torch.Tensor:
+    # Each hidden_list entry: (l_i, d)
     pooler = Pooler(pooling_types)
     max_len = max(hidden.shape[0] for hidden in hidden_list)
     hidden_dim = hidden_list[0].shape[1]
     batch_size = len(hidden_list)
-    padded = torch.zeros(batch_size, max_len, hidden_dim, device=device)
-    attention_mask = torch.zeros(batch_size, max_len, device=device)
+    padded = torch.zeros(batch_size, max_len, hidden_dim, device=device)  # (b, l_max, d)
+    attention_mask = torch.zeros(batch_size, max_len, device=device)  # (b, l_max)
     for i, hidden in enumerate(hidden_list):
         seq_len = hidden.shape[0]
         padded[i, :seq_len] = hidden
         attention_mask[i, :seq_len] = 1.0
-    return pooler(padded, attention_mask)
+    return pooler(padded, attention_mask)  # (b, n_poolers * d)
 
 
 def _forward_for_embedding(

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 != "5289065bf7679630038a2e794052e4364a152c4f103c36c80ab9b8544919e924":
+if RUNTIME_HASH != "a8dda20d6a4a55e6808c7c2f2c9477c113edaa7d6ec6ef57f4a898cd47b73d48":
     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.e1.modeling_e1")
-E1Config = _module_225.E1Config
+_module_181 = _import_without_bytecode("fastplms.models.e1.modeling_e1")
+E1Config = _module_181.E1Config
 E1Config.__module__ = __name__
-E1ForMaskedLM = _module_225.E1ForMaskedLM
+E1ForMaskedLM = _module_181.E1ForMaskedLM
 E1ForMaskedLM.__module__ = __name__
-E1ForSequenceClassification = _module_225.E1ForSequenceClassification
+E1ForSequenceClassification = _module_181.E1ForSequenceClassification
 E1ForSequenceClassification.__module__ = __name__
-E1ForTokenClassification = _module_225.E1ForTokenClassification
+E1ForTokenClassification = _module_181.E1ForTokenClassification
 E1ForTokenClassification.__module__ = __name__
-E1Model = _module_225.E1Model
+E1Model = _module_181.E1Model
 E1Model.__module__ = __name__

requirements.txt

@@ -0,0 +1,10 @@
+# Direct runtime dependencies for Synthyra/Profluent-E1-600M.
+# 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

runtime-attestation.json

@@ -7,43 +7,44 @@
     "LICENSES/e1/LICENSE": "sha256:8ef1dd556091544db3044164a8015424a3dcb3450fb3765a81b88463551bbe81",
     "LICENSES/e1/MODIFICATIONS.md": "sha256:2506f47c0f5475af8e8ff2cff13eb8b79e8e25a08a054cdd617bf336536750ca",
     "LICENSES/e1/NOTICE": "sha256:6de9db0320b4ee82f665c0951d8fd4cd53701a659c9dbce9bc3e3ea6afc4c6b3",
-    "README.md": "sha256:ce6fa1b7d092ed995e6d53650d2257576a0a7b5da86c5eb0b9f92ca3e509eac0",
+    "README.md": "sha256:f4e31044d720a5b058539f147dd9b1d4d85d34f2a6929dbf4d3684a07974b6de",
     "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
-    "config.json": "sha256:7019ed218cf9d6c05a571a6322b114a43d40a7ca727aef4606c92af0ad63b5f8",
-    "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:a679591daa125bbdf6689af936ebfb72b90a06d7b83c0b64845d7d0ed7b4f78d",
+    "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/models.toml": "sha256:1b0d92911222f31b435bfee80251c9abe20b85392d65f8af8d285d1ccfbbc3c5",
-    "fastplms/models/__init__.py": "sha256:5e48c2cb3877aa6f42f3b5411d53b16bba2e32827bbde634f47f174c5cb36f86",
+    "fastplms/models/__init__.py": "sha256:d3dd084f5e8fafcf2d0fe4a26e4ab9284d3f9414b1f9a5bd86e308a406447c43",
     "fastplms/models/e1/__init__.py": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
-    "fastplms/models/e1/attention.py": "sha256:8449b2bafba9573dfcfe4ad907e1fce24c55a4f279f1f72cd52227e48dcecf67",
-    "fastplms/models/e1/cache.py": "sha256:b17a0a743505f567b387d1edb8550f028a8770db1ba56ad8ad2ff0d3fce19c46",
-    "fastplms/models/e1/modeling_e1.py": "sha256:96545247491395291060e2cbb7b6b082ed6073e0fbf12d4656e490010aae5ae4",
-    "fastplms/models/e1/preparation.py": "sha256:3c1cc96d36cbfa74dc0f847cc5f44c6edd527d8e223d0a564d4891a041761395",
-    "fastplms/models/e1/retrieval.py": "sha256:d077ef10100cf4f1482fbc7c8a944c241e61061a94374699156388d9e8e5c4db",
+    "fastplms/models/e1/attention.py": "sha256:114bff9ed29a9f98a91965a57b702413f4a0e84bdd4063e7047e03a49a61241a",
+    "fastplms/models/e1/cache.py": "sha256:b4d065b22b3be06f929eac451a0e5a4658ac3ad366fc6ee5e9ac195af821b434",
+    "fastplms/models/e1/modeling_e1.py": "sha256:bb1c963731b7277a67d444ea0df6ed340384a6fafe80a0042e943fb844ae26e0",
+    "fastplms/models/e1/preparation.py": "sha256:b207f85032a3fbac8494ec958ecdfa37de470f6996b8568da9374e2e7ee0598c",
+    "fastplms/models/e1/retrieval.py": "sha256:6f4ac4b9f21ea874d99b91ee1f366bc456af1c0060e678e3ee7cda4c2820cefd",
     "fastplms/models/e1/tokenizer.json": "sha256:65d4345dbe908d3deb554891dc01b94e215fb17cd9e0bd0535e062e256737415",
-    "fastplms/models/ttt.py": "sha256:a0df4e98b02120d423e3c7ca9b866a8d0e3748b9076042a0102a838a11aed046",
-    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
-    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
-    "fastplms_bundle.py": "sha256:0b53f8920820cd377c6a342390542caa7e886e27005d94b20fd7d8d9c2811cc4",
-    "modeling_fastplms.py": "sha256:11932687d6d095013137669f2d4faeda96420b4ad5db16a3f90eb93cebd26c9d"
+    "fastplms/models/ttt.py": "sha256:f1d7f9298a930f4c40f8e321ffa08b09d1c5b8492bca623bdc7bfd03ec85774d",
+    "fastplms/registry.py": "sha256:cb37799dae5a4c0c780089ab2cca4cc56300c1c5f515b275892aa4e1c4b97eb6",
+    "fastplms/runtime.py": "sha256:9521c37fcd168bf9a4137277fb73672403ccf30d70d4845aa879fb22f42ddeb6",
+    "fastplms_bundle.py": "sha256:c5015ec74cd70db2cd82ed30d675510757d26bcd18a184c73560f85a8b1e9f74",
+    "modeling_fastplms.py": "sha256:bb40b09e9d61bf567d0590cb22ae7a5f8f7fb74d3694fd87c42c862c7301ecf3",
+    "requirements.txt": "sha256:9dba69cb9fd6b1c083423aa17df29bb5f98d46e7189b8c5b2393a45501056ebf"
   },
   "model_id": "e1_600m",
   "redistributable": true,
-  "release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "runtime_bundle_sha256": "5289065bf7679630038a2e794052e4364a152c4f103c36c80ab9b8544919e924",
-  "runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
+  "release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "runtime_bundle_sha256": "a8dda20d6a4a55e6808c7c2f2c9477c113edaa7d6ec6ef57f4a898cd47b73d48",
+  "runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
   "schema_version": 2,
   "scope": "runtime-only",
-  "source_tree_sha256": "c583bf72e969d20248d4f2a01c81579fffa3002e2e19c4e90acebbbacf1faeab",
+  "source_tree_sha256": "9e3cbab448ffde4fbbba79f17cf0321645584acd534c5f96c78b1017d9f3806b",
   "weights": {
     "repo_id": "Synthyra/Profluent-E1-600M",
     "revision": "6c8bf0ec83b0e0178677c528b101efffd0677742"