Update FastPLMs runtime and model cards

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

README.md

@@ -18,15 +18,34 @@ Supported Transformers entry points are `AutoConfig`, `AutoModel`,
 `AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`,
 `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: encoder or explicitly prepared decoder states |
+| Test-time training | Supported: low-rank masked-residue adaptation |
+| Attention variants | Supported: `eager`, `sdpa` |
+| 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/ANKH2_large/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.
@@ -40,22 +59,22 @@ model_id = "Synthyra/ANKH2_large"
 model = AutoModel.from_pretrained(
     model_id,
     trust_remote_code=True,
+    attn_implementation="sdpa",
 ).eval()
 ```
 
-This example uses the published Hub repository. For offline validation, build
-the manifest-pinned artifact and replace `model_id` with its local
-`dist/hub/ANKH2_large` path, then pass `local_files_only=True`.
+For offline validation, replace `model_id` with the manifest-built
+`dist/hub/ANKH2_large` path and pass `local_files_only=True`.
+
+## Attention and compliance
 
-Leave attention unspecified for the Transformers default. Supported explicit
-choices are `eager`, `sdpa`.
-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 `eager`, `sdpa`. An unavailable requested backend raises
+instead of silently switching implementations.
+`output_attentions=True` may use the documented, one-call eager fallback solely
+to materialize attention tensors; the configured backend remains unchanged.
+
+This family declares the `compliance` tier. Release evidence binds the exact
+checkpoint, backend, dtype, hardware, inputs, and reference revision.
 
 ## Tokenization and forward inference
 
@@ -86,8 +105,8 @@ print(output.last_hidden_state.shape)
 
 ## Dataset embeddings
 
-Dataset embeddings default to the encoder's final hidden state. Layer indices
-use the selected stack's native hidden-state order:
+Dataset embeddings default to the encoder final state. Select a native encoder
+layer directly:
 
 ```python
 encoder_result = model.embed_dataset(
@@ -96,20 +115,11 @@ encoder_result = model.embed_dataset(
     hidden_state_index=-1,
     full_embeddings=True,
 )
-all_encoder_layers = model.embed_dataset(
-    ["MSTNPKPQRKTKRNT"],
-    hidden_state_source="encoder",
-    store_all_hidden_states=True,
-    full_embeddings=True,
-)
+print(encoder_result[0].tensor.shape)  # (l, d)
 ```
 
 Decoder representations require `AutoModelForSeq2SeqLM` and exactly one
-explicit, aligned `decoder_inputs` sequence or `decoder_input_ids` tensor. ANKH
-does not infer a shifted source sequence because official tasks use prompts,
-sentinel tokens, or generated tokens that depend on the task. Protein inputs
-remain raw residue strings and sentinel prompts remain tight, as in
-`M<extra_id_0>`:
+aligned decoder input. ANKH does not invent a shifted target:
 
 ```python
 from transformers import AutoModelForSeq2SeqLM
@@ -125,12 +135,105 @@ decoder_result = seq2seq.embed_dataset(
     decoder_inputs=["M<extra_id_0>"],
     full_embeddings=True,
 )
+print(decoder_result[0].tensor.shape)  # (decoder_length, d)
+```
+
+Pooling excludes boundary, padding, sentinel, and other non-biological
+positions. Persisted results record the selected stack, layer, inputs, masks,
+and alignment policy.
+
+## Downstream classification
+
+Both downstream AutoClasses reuse the checkpoint backbone and initialize a new,
+untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have
+shape `(b, l)` and use `-100` outside biological positions:
+
+```python
+import torch
+from transformers import AutoTokenizer
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoModelForTokenClassification,
+)
+
+model_id = "Synthyra/ANKH2_large"
+sequence_model = AutoModelForSequenceClassification.from_pretrained(
+    model_id, num_labels=2, trust_remote_code=True
+).eval()
+token_model = AutoModelForTokenClassification.from_pretrained(
+    model_id, num_labels=3, trust_remote_code=True
+).eval()
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"]
+batch = tokenizer(sequences, padding=True, return_tensors="pt")
+biological = batch["attention_mask"].bool()
+for special_id in tokenizer.all_special_ids:
+    biological &= batch["input_ids"].ne(special_id)
+
+sequence_labels = torch.zeros(len(sequences), dtype=torch.long)
+token_labels = torch.full_like(batch["input_ids"], -100)
+token_labels[biological] = 0
+
+with torch.inference_mode():
+    sequence_output = sequence_model(**batch, labels=sequence_labels)
+    token_output = token_model(**batch, labels=token_labels)
+print(sequence_output.logits.shape)  # (b, 2)
+print(token_output.logits.shape)     # (b, l, 3)
+```
+
+## PEFT fine-tuning
+
+Install the direct training dependencies, then attach LoRA to the loaded checkpoint:
+
+```bash
+python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
+```
+
+```python
+from peft import LoraConfig, TaskType, get_peft_model
+
+peft_model = get_peft_model(
+    sequence_model,
+    LoraConfig(
+        task_type=TaskType.SEQ_CLS,
+        r=8,
+        lora_alpha=16,
+        target_modules="all-linear",
+        modules_to_save=["classifier"],
+    ),
+)
+```
+
+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/ANKH2_large",
+    trust_remote_code=True,
+)
+metrics = ttt_model.ttt(
+    seq="MSTNPKPQRKTKRNT",
+    ttt_config={"steps": 3, "batch_size": 1, "seed": 7},
+)
+ttt_model.save_pretrained("adapted", safe_serialization=True)
+ttt_model.ttt_reset()
+print(metrics)
 ```
 
-`decoder_attention_mask` is accepted only with `decoder_input_ids`. Decoder
-pooling excludes start, EOS, padding, sentinel, and other tokenizer-special
-positions. Persisted results record the selected stack and layer, decoder input
-and mask fingerprints, input-position alignment, and biological-mask policy.
+Persisted adapters retain their deterministic reset state. TTT adds latency
+and memory, can worsen an output, and does not establish biological function.
 
 ## Encoder and sequence-to-sequence use
 
@@ -177,7 +280,7 @@ masked-LM extension and is not an official ANKH head.
 - Precision policies: `default`
 - BF16 execution: `static_parameters`
 - Generation contract: `required`
-- Optional dependency group: `core`
+- Artifact dependency set: `core`
 - Weight publication allowed: `true`
 - Weight license status: `resolved`
 - Redistributable: `true`
@@ -188,30 +291,24 @@ masked-LM extension and is not an official ANKH head.
 - FastPLMs weights: `Synthyra/ANKH2_large`
 - Runtime revision: recorded separately in the built artifact and published commit
 - Source-tree and runtime-bundle SHA-256: recorded in `provenance.json`
-- Generator/schema version and complete/runtime-only attestations: recorded in `provenance.json`
 - Canonical transformed state SHA-256: `597c4fe2fa8711f11a25317905f1d62fa92905e55fdd5c0a79614cd9c9d2bca3`
 - Conversion equality attestation: recorded in `provenance.json`
 - Official checkpoint: `ElnaggarLab/ankh2-ext2`
 - Artifact source: `official`
 - State transform: `ankh_t5_to_fastplms_v1`
-- BF16 execution: `static_parameters`
 - Pinned upstreams: `ankh`
-- Reference container: `reference-ankh`
 - 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

@@ -22,11 +22,11 @@
   "fastplms_checkpoint_repo_id": "ElnaggarLab/ankh2-ext2",
   "fastplms_checkpoint_revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8",
   "fastplms_model_id": "ankh2_large",
-  "fastplms_release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "fastplms_runtime_bundle_sha256": "405e7dddef77117e55232e1112faf088db88bb9b13e212a63b3318c4bc0d5a03",
-  "fastplms_runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "fastplms_source_tree_sha256": "e377aef8902771bb8bdabcae9c8853ab50de04d85cacd3d75c6e83508f6d52ed",
+  "fastplms_release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "fastplms_runtime_bundle_sha256": "437c5f5dcc809678e99b8e36d962e78b7960170c81d0b6289d2baeef7920e40f",
+  "fastplms_runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "fastplms_source_tree_sha256": "44f108a32fffbe0e689434fbff109aea7d03a7085a550ed5771f2eac434f130d",
   "fastplms_weights_revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8",
   "feed_forward_proj": "gated-silu",
   "initializer_factor": 1.0,

fastplms/__init__.py

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

fastplms/attention/__init__.py

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

fastplms/attention/_core.py

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

fastplms/attention/_kernel_lock.py

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

fastplms/attention/interfaces.py

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

fastplms/embeddings/__init__.py

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

fastplms/embeddings/pooling.py

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

fastplms/embeddings/runner.py

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

fastplms/embeddings/storage.py

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

fastplms/embeddings/types.py

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

fastplms/models/__init__.py

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

fastplms/models/ankh/modeling_ankh.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
 import math
+import torch
+import torch.nn as nn
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from numbers import Real
 from typing import Any, ClassVar
-
-import torch
-import torch.nn as nn
 from tokenizers import pre_tokenizers
 from torch.nn import functional as F
 from transformers import (
@@ -23,6 +22,7 @@ from transformers.modeling_outputs import (
     TokenClassifierOutput,
 )
 
+
 try:
     from fastplms.attention import (
         AttentionBackend,
@@ -405,13 +405,14 @@ def _biological_token_mask(
 class AnkhRMSNorm(nn.Module):
     """T5-style RMS layer norm: scales without mean subtraction or bias."""
 
-    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
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        # hidden_states: (..., d)
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)  # (..., 1)
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         return self.weight * hidden_states.to(self.weight.dtype)
 
@@ -425,7 +426,7 @@ def _gelu_new(x: torch.Tensor) -> torch.Tensor:
 class AnkhGatedFFN(nn.Module):
     """T5-style gated feed-forward: activation(wi_0(x)) * wi_1(x) -> wo."""
 
-    def __init__(self, config: FastAnkhConfig):
+    def __init__(self, config: FastAnkhConfig) -> None:
         super().__init__()
         self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False)
         self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False)
@@ -434,6 +435,7 @@ class AnkhGatedFFN(nn.Module):
         self.dropout = nn.Dropout(config.dropout_rate)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # hidden_states: (b, l, d)
         hidden_states = self.act(self.wi_0(hidden_states)) * self.wi_1(hidden_states)
         return self.wo(self.dropout(hidden_states))
 
@@ -451,7 +453,11 @@ class AnkhSelfAttention(nn.Module):
     receive the precomputed bias through the forward call.
     """
 
-    def __init__(self, config: FastAnkhConfig, has_relative_attention_bias: bool = False):
+    def __init__(
+        self,
+        config: FastAnkhConfig,
+        has_relative_attention_bias: bool = False,
+    ) -> None:
         super().__init__()
         self.num_heads = config.num_heads
         self.d_kv = config.d_kv
@@ -515,8 +521,8 @@ class AnkhSelfAttention(nn.Module):
             num_buckets=self.relative_attention_num_buckets,
             max_distance=self.relative_attention_max_distance,
         )
-        values = self.relative_attention_bias(buckets)  # A has shape (q, k, h).
-        return values.permute(2, 0, 1).unsqueeze(0)  # A has shape (1, h, q, k).
+        values = self.relative_attention_bias(buckets)  # (q, k, h)
+        return values.permute(2, 0, 1).unsqueeze(0)  # (1, h, q, k)
 
     # ---- Forward ----
 
@@ -529,12 +535,13 @@ class AnkhSelfAttention(nn.Module):
         effective_backend: AttentionBackend | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
         """Returns (attn_output, attn_weights_or_none, position_bias)."""
+        # hidden_states: (b, l, d)
         batch_size, seq_length = hidden_states.shape[:2]
         hidden_shape = (batch_size, seq_length, self.num_heads, self.d_kv)
 
-        query_heads = self.q(hidden_states).view(hidden_shape).transpose(1, 2)
-        key_heads = self.k(hidden_states).view(hidden_shape).transpose(1, 2)
-        value_heads = self.v(hidden_states).view(hidden_shape).transpose(1, 2)
+        query_heads = self.q(hidden_states).view(hidden_shape).transpose(1, 2)  # (b, h, l, d_h)
+        key_heads = self.k(hidden_states).view(hidden_shape).transpose(1, 2)  # (b, h, l, d_h)
+        value_heads = self.v(hidden_states).view(hidden_shape).transpose(1, 2)  # (b, h, l, d_h)
 
         # The first layer computes the bias once; later layers reuse it.
         if position_bias is None and self.has_relative_attention_bias:
@@ -572,7 +579,7 @@ class AnkhSelfAttention(nn.Module):
         value_heads: torch.Tensor,
         position_bias: torch.Tensor | None,
     ) -> torch.Tensor:
-        # A is the additive position bias with shape (1, h, q, k), including padding.
+        # position_bias: (1, h, l, l), including padding
         # Never mutate torch.backends.cuda process-global reduction policy from
         # a model forward. Concurrent model requests must not change each
         # other's numerical behavior or restore a stale process setting.
@@ -583,7 +590,7 @@ class AnkhSelfAttention(nn.Module):
             attn_mask=position_bias,
             dropout_p=self.dropout_prob if self.training else 0.0,
             scale=self.scale,
-        )
+        )  # (b, h, l, d_h)
         return (
             context_heads.transpose(1, 2)
             .contiguous()
@@ -597,7 +604,10 @@ class AnkhSelfAttention(nn.Module):
         value_heads: torch.Tensor,
         position_bias: torch.Tensor | None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        attn_weights = torch.matmul(query_heads, key_heads.transpose(-1, -2)) * self.scale
+        # query_heads, key_heads, value_heads: (b, h, l, d_h)
+        attn_weights = (
+            torch.matmul(query_heads, key_heads.transpose(-1, -2)) * self.scale
+        )  # (b, h, l, l)
         if position_bias is not None:
             attn_weights = attn_weights + position_bias
         attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(attn_weights)
@@ -607,7 +617,7 @@ class AnkhSelfAttention(nn.Module):
                 p=self.dropout_prob,
                 training=self.training,
             )
-        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)
             .contiguous()
@@ -624,7 +634,11 @@ class AnkhSelfAttention(nn.Module):
 class AnkhSelfAttentionLayer(nn.Module):
     """Wraps AnkhSelfAttention + layer_norm to match T5Block.layer[0] key naming."""
 
-    def __init__(self, config: FastAnkhConfig, has_relative_attention_bias: bool = False):
+    def __init__(
+        self,
+        config: FastAnkhConfig,
+        has_relative_attention_bias: bool = False,
+    ) -> None:
         super().__init__()
         self.SelfAttention = AnkhSelfAttention(config, has_relative_attention_bias)
         self.layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon)
@@ -653,7 +667,7 @@ class AnkhSelfAttentionLayer(nn.Module):
 class AnkhFFLayer(nn.Module):
     """Wraps AnkhGatedFFN + layer_norm to match T5Block.layer[1] key naming."""
 
-    def __init__(self, config: FastAnkhConfig):
+    def __init__(self, config: FastAnkhConfig) -> None:
         super().__init__()
         self.DenseReluDense = AnkhGatedFFN(config)
         self.layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon)
@@ -668,7 +682,11 @@ class AnkhFFLayer(nn.Module):
 class AnkhBlock(nn.Module):
     """Single transformer block with T5-compatible .layer ModuleList naming."""
 
-    def __init__(self, config: FastAnkhConfig, has_relative_attention_bias: bool = False):
+    def __init__(
+        self,
+        config: FastAnkhConfig,
+        has_relative_attention_bias: bool = False,
+    ) -> None:
         super().__init__()
         self.layer = nn.ModuleList(
             [
@@ -782,7 +800,7 @@ class FAST_ANKH_ENCODER(AnkhPreTrainedModel, EmbeddingMixin):
     block.{i}.layer.1.DenseReluDense.*, final_layer_norm.*.
     """
 
-    def __init__(self, config: FastAnkhConfig, **kwargs):
+    def __init__(self, config: FastAnkhConfig, **kwargs) -> None:
         AnkhPreTrainedModel.__init__(self, config, **kwargs)
         self.config = config
 
@@ -965,7 +983,7 @@ class FastAnkhModel(AnkhPreTrainedModel, EmbeddingMixin):
         r"^lm_head\.",
     ]
 
-    def __init__(self, config: FastAnkhConfig, **kwargs):
+    def __init__(self, config: FastAnkhConfig, **kwargs) -> None:
         AnkhPreTrainedModel.__init__(self, config, **kwargs)
         self.config = config
         self.shared = nn.Embedding(config.vocab_size, config.d_model)
@@ -1029,7 +1047,7 @@ class FastAnkhForMaskedLMExtension(
     _tied_weights_keys: ClassVar[dict[str, str]] = {"encoder.embed_tokens.weight": "shared.weight"}
     _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"^decoder\."]
 
-    def __init__(self, config: FastAnkhConfig, **kwargs):
+    def __init__(self, config: FastAnkhConfig, **kwargs) -> None:
         # The historical Synthyra extension stores an independent output head.
         config.tie_word_embeddings = False
         AnkhPreTrainedModel.__init__(self, config, **kwargs)
@@ -1156,7 +1174,7 @@ class FastAnkhForConditionalGeneration(
     embedding_unsupported_pooling = ("cls",)
     _fastplms_attention_implementations = ("eager",)
 
-    def __init__(self, config: FastAnkhConfig, **kwargs):
+    def __init__(self, config: FastAnkhConfig, **kwargs) -> None:
         requested_backend = getattr(config, "_attn_implementation", None) or config.attn_backend
         if requested_backend not in (None, "eager"):
             raise ValueError(
@@ -1448,7 +1466,7 @@ class FastAnkhForSequenceClassification(AnkhPreTrainedModel, EmbeddingMixin):
         r"^lm_head\.",
     ]
 
-    def __init__(self, config: FastAnkhConfig, **kwargs):
+    def __init__(self, config: FastAnkhConfig, **kwargs) -> None:
         AnkhPreTrainedModel.__init__(self, config, **kwargs)
         self.num_labels = config.num_labels
         self.config = config
@@ -1555,7 +1573,7 @@ class FastAnkhForTokenClassification(AnkhPreTrainedModel, EmbeddingMixin):
         r"^lm_head\.",
     ]
 
-    def __init__(self, config: FastAnkhConfig, **kwargs):
+    def __init__(self, config: FastAnkhConfig, **kwargs) -> None:
         AnkhPreTrainedModel.__init__(self, config, **kwargs)
         self.num_labels = config.num_labels
         self.shared = nn.Embedding(config.vocab_size, config.d_model)

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 != "405e7dddef77117e55232e1112faf088db88bb9b13e212a63b3318c4bc0d5a03":
+if RUNTIME_HASH != "437c5f5dcc809678e99b8e36d962e78b7960170c81d0b6289d2baeef7920e40f":
     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,16 +179,16 @@ def _install_runtime():
     return package
 
 _install_runtime()
-_module_225 = _import_without_bytecode("fastplms.models.ankh.modeling_ankh")
-FastAnkhConfig = _module_225.FastAnkhConfig
+_module_181 = _import_without_bytecode("fastplms.models.ankh.modeling_ankh")
+FastAnkhConfig = _module_181.FastAnkhConfig
 FastAnkhConfig.__module__ = __name__
-FastAnkhForConditionalGeneration = _module_225.FastAnkhForConditionalGeneration
+FastAnkhForConditionalGeneration = _module_181.FastAnkhForConditionalGeneration
 FastAnkhForConditionalGeneration.__module__ = __name__
-FastAnkhForMaskedLMExtension = _module_225.FastAnkhForMaskedLMExtension
+FastAnkhForMaskedLMExtension = _module_181.FastAnkhForMaskedLMExtension
 FastAnkhForMaskedLMExtension.__module__ = __name__
-FastAnkhForSequenceClassification = _module_225.FastAnkhForSequenceClassification
+FastAnkhForSequenceClassification = _module_181.FastAnkhForSequenceClassification
 FastAnkhForSequenceClassification.__module__ = __name__
-FastAnkhForTokenClassification = _module_225.FastAnkhForTokenClassification
+FastAnkhForTokenClassification = _module_181.FastAnkhForTokenClassification
 FastAnkhForTokenClassification.__module__ = __name__
-FastAnkhModel = _module_225.FastAnkhModel
+FastAnkhModel = _module_181.FastAnkhModel
 FastAnkhModel.__module__ = __name__

requirements.txt

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

runtime-attestation.json

@@ -2,41 +2,42 @@
   "files": {
     "LICENSES/FastPLMs-Apache-2.0.txt": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
     "LICENSES/ankh/LICENSE.md": "sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088",
-    "README.md": "sha256:87671dd466292a782bedf8dfe95bc595f03450e253243f2da217b3898acfc182",
+    "README.md": "sha256:131315548dc250a4359361d16d9355080d1598c8bff6fc9902abe68c9e966cb2",
     "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
-    "config.json": "sha256:0ba3a92b852ff36f7658a466903ee970ea45a59c9c555fe87b119607bbc32ff2",
-    "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:c2e2327b8765473cbe9004169bbd860ff7dca916cfd66653c4a08c6c3a419bce",
+    "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/ankh/__init__.py": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
-    "fastplms/models/ankh/modeling_ankh.py": "sha256:b114c73e4aaedc567e3351b425ddddacf35661f09b3414cf59b6db0b712c82c3",
-    "fastplms/models/ttt.py": "sha256:a0df4e98b02120d423e3c7ca9b866a8d0e3748b9076042a0102a838a11aed046",
-    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
-    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
-    "fastplms_bundle.py": "sha256:a06395708d58c64035ff39e6e38280d6fcb733d223de3709a21f1a5c87cd97b7",
-    "modeling_fastplms.py": "sha256:ccf1ac016e0fdcaaed4376db29e5b29d843d776c48c977f1006edfaab51c8307",
+    "fastplms/models/ankh/modeling_ankh.py": "sha256:9ff235c9bd529a4c2ba63293376e6479161d364e7a4884c4063a0eb3ed1c5719",
+    "fastplms/models/ttt.py": "sha256:f1d7f9298a930f4c40f8e321ffa08b09d1c5b8492bca623bdc7bfd03ec85774d",
+    "fastplms/registry.py": "sha256:cb37799dae5a4c0c780089ab2cca4cc56300c1c5f515b275892aa4e1c4b97eb6",
+    "fastplms/runtime.py": "sha256:9521c37fcd168bf9a4137277fb73672403ccf30d70d4845aa879fb22f42ddeb6",
+    "fastplms_bundle.py": "sha256:af54f8786c1124b45e07ef857486e14163bb79266fa98fc2002a82296ce89fbb",
+    "modeling_fastplms.py": "sha256:37b0ea758dc1a50e7ca7225a175ca983e109efe4342759055c58520ef271f49f",
+    "requirements.txt": "sha256:3ae7e085e6ec8ed936981f14e58d1199b4a26ba79cf90cb5d01585946ad8a96a",
     "special_tokens_map.json": "sha256:c8995d2f8037fe3a8cfdef30475365e1c314c417880b75abaf8296e3c05d42d6",
     "tokenizer.json": "sha256:b4533f607d9fd665f2d9d94b0cf71870a6fc2fc2ae7cbd516d0e43a9efb406fd",
     "tokenizer_config.json": "sha256:dba1f9315ded007fc2e5bcb2ed4bfb6fff81c9becdc97b73a5713ba5ddc01afc"
   },
   "model_id": "ankh2_large",
   "redistributable": true,
-  "release_tool_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
-  "release_tool_sha256": "19b4247aba04a74de5069397518374ab2c3b9e0ce18ba9f36de34f0f4c5d1ceb",
-  "runtime_bundle_sha256": "405e7dddef77117e55232e1112faf088db88bb9b13e212a63b3318c4bc0d5a03",
-  "runtime_revision": "95691a87c781f052f65b6ab3ca99dfce4ebb59c7",
+  "release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
+  "runtime_bundle_sha256": "437c5f5dcc809678e99b8e36d962e78b7960170c81d0b6289d2baeef7920e40f",
+  "runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
   "schema_version": 2,
   "scope": "runtime-only",
-  "source_tree_sha256": "e377aef8902771bb8bdabcae9c8853ab50de04d85cacd3d75c6e83508f6d52ed",
+  "source_tree_sha256": "44f108a32fffbe0e689434fbff109aea7d03a7085a550ed5771f2eac434f130d",
   "weights": {
     "repo_id": "Synthyra/ANKH2_large",
     "revision": "729167c1980316ae61691338838447491926033f"