slide2vec 6.3.3__tar.gz → 6.3.4__tar.gz

{slide2vec-6.3.3 → slide2vec-6.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 6.3.3
+Version: 6.3.4
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0
@@ -27,18 +27,18 @@ Requires-Dist: torchvision
 Requires-Dist: transformers
 Requires-Dist: wandb
 Requires-Dist: einops
-Requires-Dist: timm
+Requires-Dist: timm>=1.0.3
 Requires-Dist: huggingface_hub
 Provides-Extra: hoptimus
 Requires-Dist: torch>=2.0; extra == "hoptimus"
 Requires-Dist: torchvision>=0.15.0; extra == "hoptimus"
 Requires-Dist: xformers>=0.0.18; extra == "hoptimus"
 Provides-Extra: virchow
-Requires-Dist: timm>=0.9.11; extra == "virchow"
+Requires-Dist: timm>=1.0.3; extra == "virchow"
 Requires-Dist: torch>=2.0; extra == "virchow"
 Provides-Extra: uni
 Requires-Dist: torch>=2.0; extra == "uni"
-Requires-Dist: timm>=0.9.8; extra == "uni"
+Requires-Dist: timm>=1.0.3; extra == "uni"
 Requires-Dist: xformers>=0.0.18; extra == "uni"
 Provides-Extra: prism
 Requires-Dist: transformers~=4.53.0; extra == "prism"

{slide2vec-6.3.3 → slide2vec-6.3.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "slide2vec"
-version = "6.3.3"
+version = "6.3.4"
 description = "Embedding of whole slide images with Foundation Models"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -33,7 +33,7 @@ dependencies = [
     "transformers",
     "wandb",
     "einops",
-    "timm",
+    "timm>=1.0.3",
     "huggingface_hub",
 ]
 
@@ -51,12 +51,12 @@ hoptimus = [
     "xformers>=0.0.18",
 ]
 virchow = [
-    "timm>=0.9.11",
+    "timm>=1.0.3",
     "torch>=2.0",
 ]
 uni = [
     "torch>=2.0",
-    "timm>=0.9.8",
+    "timm>=1.0.3",
     "xformers>=0.0.18",
 ]
 prism = [
@@ -180,7 +180,7 @@ no_implicit_reexport = true
 max-line-length = 160
 
 [tool.bumpver]
-current_version = "6.3.3"
+current_version = "6.3.4"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit = false
 tag = false

{slide2vec-6.3.3 → slide2vec-6.3.4}/slide2vec/__init__.py

@@ -31,7 +31,7 @@ from slide2vec.encoders import (
 )
 
 
-__version__ = "6.3.3"
+__version__ = "6.3.4"
 
 __all__ = [
     "Model",

{slide2vec-6.3.3 → slide2vec-6.3.4}/slide2vec/api.py

@@ -26,6 +26,7 @@ from slide2vec.encoders.registry import (
     encoder_registry,
     list_encoder_provider_diagnostics,
     resolve_preprocessing_fields,
+    validate_feature_kind_capability,
 )
 from slide2vec.encoders.validation import validate_encoder_config
 from slide2vec.runtime.model_settings import (
@@ -457,8 +458,10 @@ class DenseOptions:
     #: Fractional window overlap in ``[0, 1)`` for the sliding path (ignored when
     #: ``window_size is None``).
     overlap: float = 0.0
-    #: ``"patch_features"`` (the patch-token grid) or ``"cls_attention"`` (CLS/register
-    #: self-attention grid).
+    #: ``"patch_features"`` (the patch-token grid after the backbone's final norm),
+    #: ``"patch_features_prenorm"`` (the same grid tapped *before* the final norm, timm's
+    #: ``features_only`` tap; timm ViT encoders only) or ``"cls_attention"``
+    #: (CLS/register self-attention grid).
     feature_kind: str = "patch_features"
     #: Transformer blocks whose CLS attention is read (``cls_attention`` only).
     attention_blocks: tuple[int, ...] = (-1,)
@@ -508,8 +511,10 @@ class DenseImageOptions:
     #: Fractional window overlap in ``[0, 1)`` for the sliding path (ignored when
     #: ``window_size is None``).
     overlap: float = 0.0
-    #: ``"patch_features"`` (the patch-token grid) or ``"cls_attention"`` (CLS/register
-    #: self-attention grid).
+    #: ``"patch_features"`` (the patch-token grid after the backbone's final norm),
+    #: ``"patch_features_prenorm"`` (the same grid tapped *before* the final norm, timm's
+    #: ``features_only`` tap; timm ViT encoders only) or ``"cls_attention"``
+    #: (CLS/register self-attention grid).
     feature_kind: str = "patch_features"
     #: Transformer blocks whose CLS attention is read (``cls_attention`` only).
     attention_blocks: tuple[int, ...] = (-1,)
@@ -1091,14 +1096,20 @@ class Model:
         parent stage and every torchrun rank — declares for itself.
 
         *dense* is a :class:`DenseOptions` (ROIs on a slide) or a :class:`DenseImageOptions`
-        (pre-cropped images); only the supervision geometry is read here, and the two state
-        it the same way.
+        (pre-cropped images); only the supervision geometry and the requested feature kind
+        are read here, and the two state them the same way.
+
+        The feature kind is checked against the registered encoder class at the same
+        point, so a tap this encoder does not implement (``patch_features_prenorm`` on a
+        non-timm backbone, say) fails before any backend is loaded and before resume
+        inspects or invalidates an existing artifact.
         """
         contract = EncoderInputContract.declared_dense(
             self.name,
             target_size_px=dense.target_size,
             window_size=None if dense.window_size is None else int(dense.window_size),
         )
+        validate_feature_kind_capability(self.name, feature_kind=str(dense.feature_kind))
         self._encoder_input = contract
         plan = contract.plan
         if emit_run_info and plan.requires_variable_model_input:

{slide2vec-6.3.3 → slide2vec-6.3.4}/slide2vec/encoders/base.py

@@ -462,6 +462,21 @@ class TileEncoder(Encoder):
             "tokens can be reshaped into a spatial grid."
         )
 
+    def encode_tiles_dense_prenorm(self, batch: Tensor) -> Tensor:
+        """Dense grid tapped before the backbone's final normalisation. (B, C, H, W) -> (B, d, h, w).
+
+        Same grid as :meth:`encode_tiles_dense` but from the last block's output,
+        *before* the final LayerNorm — the tap timm's ``features_only`` backbones
+        expose (``forward_intermediates(norm=False)``) and that kaiko-ai/eva's
+        segmentation decoders were trained on. Default: unsupported; the timm ViT
+        base overrides it.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support pre-norm dense feature "
+            "extraction (feature_kind='patch_features_prenorm'). The tap is "
+            "available for timm ViT tile encoders only."
+        )
+
     def encode_tiles_attention(
         self,
         batch: Tensor,
@@ -634,15 +649,8 @@ class TimmTileEncoder(TileEncoder):
         """Number of leading non-patch tokens (CLS + register tokens)."""
         return int(self._model.num_prefix_tokens)
 
-    def encode_tiles_dense(self, batch: Tensor) -> Tensor:
-        """Encode tiles into a dense spatial grid. (B, C, H, W) -> (B, d, h, w).
-
-        Runs the frozen backbone's ``forward_features`` and folds the patch-token
-        sequence back into its spatial grid (CLS/register tokens discarded). The
-        backbone must accept ``batch`` at its current spatial size (timm ViTs need
-        ``dynamic_img_size=True`` for sizes other than their native input), and
-        ``H, W`` must be divisible by the patch size.
-        """
+    def _check_dense_batch(self, batch: Tensor) -> tuple[int, int]:
+        """Validate a dense ``(B, C, H, W)`` batch; return ``(H, W)``."""
         if batch.ndim != 4:
             raise ValueError(
                 "encode_tiles_dense expects a (B, C, H, W) batch, got shape "
@@ -656,6 +664,19 @@ class TimmTileEncoder(TileEncoder):
                 f"divisible by the patch size: got {height}x{width}, patch "
                 f"{patch_h}x{patch_w}. Pad the tile up to a patch multiple first."
             )
+        return int(height), int(width)
+
+    def encode_tiles_dense(self, batch: Tensor) -> Tensor:
+        """Encode tiles into a dense spatial grid. (B, C, H, W) -> (B, d, h, w).
+
+        Runs the frozen backbone's ``forward_features`` and folds the patch-token
+        sequence back into its spatial grid (CLS/register tokens discarded). The
+        backbone must accept ``batch`` at its current spatial size (timm ViTs need
+        ``dynamic_img_size=True`` for sizes other than their native input), and
+        ``H, W`` must be divisible by the patch size.
+        """
+        height, width = self._check_dense_batch(batch)
+        patch_h, patch_w = self._dense_patch_size()
         tokens = self._model.forward_features(batch)
         return reshape_tokens_to_grid(
             tokens,
@@ -665,6 +686,25 @@ class TimmTileEncoder(TileEncoder):
             encoder_name=type(self).__name__,
         )
 
+    def encode_tiles_dense_prenorm(self, batch: Tensor) -> Tensor:
+        """Pre-norm dense grid via timm's ``forward_intermediates(norm=False)``.
+
+        ``indices=1`` selects the last block, ``output_fmt="NCHW"`` folds the patch
+        tokens into the grid (prefix tokens dropped), ``norm=False`` skips the final
+        LayerNorm — exactly timm's ``features_only=True, out_indices=1`` path.
+        """
+        self._check_dense_batch(batch)
+        forward_intermediates = getattr(self._model, "forward_intermediates", None)
+        if forward_intermediates is None:
+            raise NotImplementedError(
+                f"{type(self).__name__} backbone has no forward_intermediates; the "
+                "pre-norm dense tap needs a timm ViT-family model."
+            )
+        grids = forward_intermediates(
+            batch, indices=1, norm=False, intermediates_only=True, output_fmt="NCHW"
+        )
+        return grids[-1]
+
     def encode_tiles_attention(
         self,
         batch: Tensor,

{slide2vec-6.3.3 → slide2vec-6.3.4}/slide2vec/encoders/registry.py

@@ -149,6 +149,7 @@ class EncoderCapabilities:
     level: str
     pooled: bool
     dense: bool
+    dense_prenorm: bool
     attention: bool
     slide: bool
     patient: bool
@@ -164,6 +165,7 @@ def resolve_encoder_capabilities(encoder_name: str) -> EncoderCapabilities:
     _validate_encoder_capability_contract(encoder_name, encoder_cls, metadata)
     level = resolve_encoder_level(encoder_name, metadata)
     dense = all(_dense_class_contract(encoder_cls))
+    dense_prenorm = dense and _supports_dense_prenorm(encoder_cls)
     attention = _supports_attention(encoder_cls)
     tile_encoder = None
     tile_encoder_output_variant = None
@@ -189,6 +191,7 @@ def resolve_encoder_capabilities(encoder_name: str) -> EncoderCapabilities:
         level=level,
         pooled=issubclass(encoder_cls, TileEncoder),
         dense=dense,
+        dense_prenorm=dense_prenorm,
         attention=attention,
         slide=issubclass(encoder_cls, (SlideEncoder, PatientEncoder)),
         patient=issubclass(encoder_cls, PatientEncoder),
@@ -239,6 +242,43 @@ def _dense_class_contract(encoder_cls: type) -> tuple[bool, ...]:
     )
 
 
+def _supports_dense_prenorm(encoder_cls: type) -> bool:
+    return (
+        issubclass(encoder_cls, TileEncoder)
+        and getattr(encoder_cls, "encode_tiles_dense_prenorm")
+        is not TileEncoder.encode_tiles_dense_prenorm
+    )
+
+
+#: Dense feature kind -> the ``TileEncoder`` method that produces it.
+_FEATURE_KIND_METHODS = {
+    "patch_features": "encode_tiles_dense",
+    "patch_features_prenorm": "encode_tiles_dense_prenorm",
+    "cls_attention": "encode_tiles_attention",
+}
+
+
+def validate_feature_kind_capability(encoder_name: str, *, feature_kind: str) -> None:
+    """Reject a dense feature kind whose method the encoder class leaves unimplemented.
+
+    Resolved from the registered class, so it runs before any backend is loaded and before
+    any on-disk artifact is inspected or invalidated. Every dense route (persisted ROI,
+    persisted image, and live kit) goes through it via the dense contract declaration.
+    """
+    encoder_cls = encoder_registry.require(encoder_name)
+    method_name = _FEATURE_KIND_METHODS.get(str(feature_kind))
+    if method_name is None:
+        raise ValueError(
+            f"unsupported feature_kind {feature_kind!r}; expected one of "
+            + ", ".join(repr(kind) for kind in _FEATURE_KIND_METHODS)
+        )
+    if getattr(encoder_cls, method_name) is getattr(TileEncoder, method_name):
+        raise ValueError(
+            f"Encoder {encoder_name!r} does not support {feature_kind}; choose a "
+            "registered dense feature kind implemented by this encoder."
+        )
+
+
 def _supports_attention(encoder_cls: type) -> bool:
     return (
         issubclass(encoder_cls, TileEncoder)

{slide2vec-6.3.3 → slide2vec-6.3.4}/slide2vec/runtime/dense_encode.py

@@ -7,8 +7,7 @@ from typing import TYPE_CHECKING, Any, Callable, cast
 
 import torch
 
-from slide2vec.encoders.base import TileEncoder
-from slide2vec.encoders.registry import encoder_registry, resolve_patch_size
+from slide2vec.encoders.registry import resolve_patch_size
 from slide2vec.runtime.dense_regions import (
     DenseGridGeometry,
     DenseGridEncoder,
@@ -94,10 +93,6 @@ class _DenseEncodePlan:
             attention_blocks=dense.attention_blocks,
             attention_include_registers=dense.attention_include_registers,
         )
-        _validate_registered_feature_capability(
-            contract_plan.tile_encoder_name,
-            feature_kind=str(dense.feature_kind),
-        )
         if execution.precision is None:
             raise ValueError("Live dense encoder precision must be resolved before loading")
         output_precision = resolve_output_precision(
@@ -122,21 +117,6 @@ class _DenseEncodePlan:
         )
 
 
-def _validate_registered_feature_capability(
-    encoder_name: str, *, feature_kind: str
-) -> None:
-    """Reject a feature method left at ``TileEncoder``'s unsupported default."""
-    encoder_cls = encoder_registry.require(encoder_name)
-    method_name = {
-        "patch_features": "encode_tiles_dense",
-        "cls_attention": "encode_tiles_attention",
-    }[feature_kind]
-    if getattr(encoder_cls, method_name) is getattr(TileEncoder, method_name):
-        raise ValueError(
-            f"Encoder {encoder_name!r} does not support {feature_kind}; choose a "
-            "registered dense feature kind implemented by this encoder."
-        )
-
 
 @dataclass(frozen=True, kw_only=True)
 class _DenseItemPreprocessor:

{slide2vec-6.3.3 → slide2vec-6.3.4}/slide2vec/runtime/dense_regions.py

@@ -51,6 +51,10 @@ from slide2vec.runtime.slide_encode import slide_encode_autocast_ctx
 from slide2vec.runtime.tiling import resolve_slide_backend
 
 
+#: Dense feature kinds a request may ask for, each backed by one ``TileEncoder`` method.
+DENSE_FEATURE_KINDS: tuple[str, ...] = ("patch_features", "patch_features_prenorm", "cls_attention")
+
+
 def _resolve_output_dtype(output_dtype: "torch.dtype | None", precision: str) -> "torch.dtype":
     """Resolve the dtype emitted grids are materialized in.
 
@@ -152,10 +156,10 @@ def validate_dense_request_settings(
                 f"dimension; target_size={geometry.target_size} requires bottom/right "
                 f"pad={geometry.pad}. Choose a larger target or another pad_mode."
             )
-    if feature_kind not in {"patch_features", "cls_attention"}:
+    if feature_kind not in DENSE_FEATURE_KINDS:
         raise ValueError(
-            f"unsupported feature_kind {feature_kind!r}; "
-            "expected 'patch_features' or 'cls_attention'"
+            f"unsupported feature_kind {feature_kind!r}; expected one of "
+            + ", ".join(repr(kind) for kind in DENSE_FEATURE_KINDS)
         )
     if window_size is not None:
         if int(window_size) <= 0:
@@ -210,6 +214,8 @@ def _resolve_encode_fn(
 ) -> Callable[[torch.Tensor], torch.Tensor]:
     if feature_kind == "patch_features":
         return model.encode_tiles_dense
+    if feature_kind == "patch_features_prenorm":
+        return model.encode_tiles_dense_prenorm
     if feature_kind == "cls_attention":
         blocks = tuple(int(b) for b in attention_blocks)
         include_registers = bool(attention_include_registers)
@@ -221,7 +227,8 @@ def _resolve_encode_fn(
 
         return encode_fn
     raise ValueError(
-        f"unsupported feature_kind {feature_kind!r}; expected 'patch_features' or 'cls_attention'"
+        f"unsupported feature_kind {feature_kind!r}; expected one of "
+        + ", ".join(repr(kind) for kind in DENSE_FEATURE_KINDS)
     )
 
 

{slide2vec-6.3.3 → slide2vec-6.3.4}/slide2vec.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 6.3.3
+Version: 6.3.4
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0
@@ -27,18 +27,18 @@ Requires-Dist: torchvision
 Requires-Dist: transformers
 Requires-Dist: wandb
 Requires-Dist: einops
-Requires-Dist: timm
+Requires-Dist: timm>=1.0.3
 Requires-Dist: huggingface_hub
 Provides-Extra: hoptimus
 Requires-Dist: torch>=2.0; extra == "hoptimus"
 Requires-Dist: torchvision>=0.15.0; extra == "hoptimus"
 Requires-Dist: xformers>=0.0.18; extra == "hoptimus"
 Provides-Extra: virchow
-Requires-Dist: timm>=0.9.11; extra == "virchow"
+Requires-Dist: timm>=1.0.3; extra == "virchow"
 Requires-Dist: torch>=2.0; extra == "virchow"
 Provides-Extra: uni
 Requires-Dist: torch>=2.0; extra == "uni"
-Requires-Dist: timm>=0.9.8; extra == "uni"
+Requires-Dist: timm>=1.0.3; extra == "uni"
 Requires-Dist: xformers>=0.0.18; extra == "uni"
 Provides-Extra: prism
 Requires-Dist: transformers~=4.53.0; extra == "prism"

{slide2vec-6.3.3 → slide2vec-6.3.4}/slide2vec.egg-info/requires.txt

@@ -10,7 +10,7 @@ torchvision
 transformers
 wandb
 einops
-timm
+timm>=1.0.3
 huggingface_hub
 
 [docs]
@@ -96,11 +96,11 @@ transformers==4.46.0
 
 [uni]
 torch>=2.0
-timm>=0.9.8
+timm>=1.0.3
 xformers>=0.0.18
 
 [virchow]
-timm>=0.9.11
+timm>=1.0.3
 torch>=2.0
 
 [waiv]

{slide2vec-6.3.3 → slide2vec-6.3.4}/tests/test_dense_encode_kit.py

@@ -247,6 +247,21 @@ def test_prepare_dense_encoder_rejects_unsupported_attention_before_loading(
     assert load_attempted is False
 
 
+@pytest.mark.parametrize(("model_name", "target_size"), [("phikon", 224), ("hibou-b", 224)])
+def test_prepare_dense_encoder_rejects_unsupported_prenorm_before_loading(
+    model_name, target_size, monkeypatch
+):
+    model = Model(name=model_name, device="cpu")
+    monkeypatch.setattr(
+        model, "_load_backend", lambda: (_ for _ in ()).throw(AssertionError("must not load"))
+    )
+    with pytest.raises(ValueError, match="does not support patch_features_prenorm"):
+        model.prepare_dense_encoder(
+            dense=_dense_image(target_size=target_size, feature_kind="patch_features_prenorm"),
+            execution=ExecutionOptions(num_gpus=1, precision="fp32"),
+        )
+
+
 def test_prepare_dense_encoder_rejects_non_tile_models_before_loading(monkeypatch):
     model = Model(name="gigapath-slide", device="cpu")
     load_attempted = False

{slide2vec-6.3.3 → slide2vec-6.3.4}/tests/test_dense_image_stage.py

@@ -22,7 +22,7 @@ PIL = pytest.importorskip("PIL")
 
 from PIL import Image  # noqa: E402
 
-from slide2vec.api import DenseImageOptions, ExecutionOptions, ImageSpec  # noqa: E402
+from slide2vec.api import DenseImageOptions, ExecutionOptions, ImageSpec, Model  # noqa: E402
 from slide2vec.encoders.base import TimmTileEncoder  # noqa: E402
 from slide2vec.runtime import dense_image_stage  # noqa: E402
 from slide2vec.runtime.types import LoadedModel  # noqa: E402
@@ -1085,3 +1085,32 @@ def test_worker_encodes_only_its_rank_shard(tmp_path, monkeypatch):
     # World of 2 ranks: plan_contiguous_shards([a,b,c,d], 2) => rank 1 owns [c, d].
     dense_dir = tmp_path / "out" / "dense_image_embeddings"
     assert sorted(p.name for p in dense_dir.glob("*.pt")) == ["c.pt", "d.pt"]
+
+
+def test_embed_images_dense_rejects_unsupported_prenorm_before_loading_or_resume(
+    tmp_path, monkeypatch
+):
+    """An unsupported feature kind fails from registry metadata on the persisted image route.
+
+    Phikon is not a timm ViT, so ``patch_features_prenorm`` is rejected before the backend
+    loads and before resume inspects the output dir: an existing post-norm payload and its
+    sidecar are left untouched (the earlier behaviour deleted the sidecar first).
+    """
+    model = Model(name="phikon", device="cpu")
+    monkeypatch.setattr(
+        model, "_load_backend", lambda: (_ for _ in ()).throw(AssertionError("must not load"))
+    )
+    specs = _images(tmp_path, ["a"])
+    out_dir = tmp_path / "out"
+    existing = out_dir / "dense_image_embeddings"
+    existing.mkdir(parents=True)
+    (existing / "a.pt").write_bytes(b"payload")
+    (existing / "a.meta.json").write_text("{}")
+    with pytest.raises(ValueError, match="does not support patch_features_prenorm"):
+        model.embed_images_dense(
+            specs,
+            dense=_dense(feature_kind="patch_features_prenorm"),
+            execution=ExecutionOptions(num_gpus=1, precision="fp32", output_dir=out_dir),
+        )
+    assert (existing / "a.pt").read_bytes() == b"payload"
+    assert (existing / "a.meta.json").read_text() == "{}"

{slide2vec-6.3.3 → slide2vec-6.3.4}/tests/test_dense_regions.py

@@ -161,7 +161,7 @@ def _make_tiling_result(
     )
 
 
-@pytest.mark.parametrize("feature_kind", ["patch_features", "cls_attention"])
+@pytest.mark.parametrize("feature_kind", ["patch_features", "patch_features_prenorm", "cls_attention"])
 @pytest.mark.parametrize("window_size", [None, 32], ids=["whole", "window32"])
 def test_iter_regions_dense_yields_grid_per_coordinate_in_order(fake_backend, window_size, feature_kind):
     enc = _encoder()
@@ -235,6 +235,8 @@ def _reference_grid(
     batch = padded.unsqueeze(0)
     if feature_kind == "patch_features":
         encode_fn = enc.encode_tiles_dense
+    elif feature_kind == "patch_features_prenorm":
+        encode_fn = enc.encode_tiles_dense_prenorm
     else:
         encode_fn = enc.encode_tiles_attention
     with torch.inference_mode():
@@ -248,7 +250,7 @@ def _reference_grid(
     return out.detach().float().cpu().numpy()[0]
 
 
-@pytest.mark.parametrize("feature_kind", ["patch_features", "cls_attention"])
+@pytest.mark.parametrize("feature_kind", ["patch_features", "patch_features_prenorm", "cls_attention"])
 @pytest.mark.parametrize("window_size", [None, 32], ids=["whole", "window32"])
 def test_iter_regions_dense_matches_direct_encode(fake_backend, window_size, feature_kind):
     """Each yielded grid is byte-identical to a hand-rolled transform+pad+encode.
@@ -276,7 +278,7 @@ def test_iter_regions_dense_matches_direct_encode(fake_backend, window_size, fea
         np.testing.assert_array_equal(grid, ref)
 
 
-@pytest.mark.parametrize("feature_kind", ["patch_features", "cls_attention"])
+@pytest.mark.parametrize("feature_kind", ["patch_features", "patch_features_prenorm", "cls_attention"])
 def test_iter_regions_dense_area_resizes_when_read_differs_from_requested(fake_backend, feature_kind):
     """When ``read_tile_size_px != requested_tile_size_px`` the region is area-resized.
 
@@ -333,7 +335,7 @@ def test_iter_regions_dense_empty_coordinates_yields_nothing(fake_backend):
     assert fake_backend.open_count == 0  # nothing read -> slide never opened
 
 
-@pytest.mark.parametrize("feature_kind", ["patch_features", "cls_attention"])
+@pytest.mark.parametrize("feature_kind", ["patch_features", "patch_features_prenorm", "cls_attention"])
 @pytest.mark.parametrize("window_size", [None, 32], ids=["whole", "window32"])
 def test_iter_regions_dense_streams_one_batch_at_a_time(fake_backend, window_size, feature_kind):
     """Reads advance one batch at a time; first grids land before all coords are read.
@@ -369,7 +371,7 @@ def test_iter_regions_dense_streams_one_batch_at_a_time(fake_backend, window_siz
     assert len(backend.locations_read) == len(coords)  # total reads never exceed the coordinate count
 
 
-@pytest.mark.parametrize("feature_kind", ["patch_features", "cls_attention"])
+@pytest.mark.parametrize("feature_kind", ["patch_features", "patch_features_prenorm", "cls_attention"])
 def test_iter_regions_dense_is_batch_invariant(fake_backend, feature_kind):
     """Composition is irrelevant: only ``B`` matters, not how coords are grouped (adr/0002).
 
@@ -448,3 +450,64 @@ def test_iter_regions_dense_rejects_bfloat16_output_eagerly(fake_backend):
         )
     assert fake_backend.backend.locations_read == []
     assert fake_backend.open_count == 0
+
+
+@pytest.mark.parametrize("model_name", ["vit_tiny_patch16_224", "vit_small_patch14_reg4_dinov2"])
+def test_prenorm_dense_grid_is_the_last_block_output_before_the_final_norm(model_name):
+    """``encode_tiles_dense_prenorm`` is ``encode_tiles_dense`` minus the backbone's final norm.
+
+    Applying the backbone's own ``norm`` to the pre-norm grid reproduces the post-norm grid
+    exactly (timm's ``features_only`` tap), for a plain ViT and a register-token ViT.
+    """
+    enc = TimmTileEncoder(model_name, pretrained=False, num_classes=0, dynamic_img_size=True)
+    patch = enc.patch_size[0]
+    batch = torch.randn(2, 3, 4 * patch, 3 * patch, generator=torch.Generator().manual_seed(0))
+    with torch.inference_mode():
+        post = enc.encode_tiles_dense(batch)
+        pre = enc.encode_tiles_dense_prenorm(batch)
+        normed = enc._model.norm(pre.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+    assert tuple(pre.shape) == tuple(post.shape) == (2, enc.encode_dim, 4, 3)
+    assert not torch.allclose(pre, post)
+    assert torch.allclose(normed, post, atol=1e-5)
+
+
+def test_prenorm_dense_grid_has_literal_values_for_a_fixed_backbone():
+    """Literal oracle: with every weight zeroed, each block adds only its MLP output bias.
+
+    Penultimate block bias 5, last block bias 3, final LayerNorm bias 2: the pre-norm grid
+    must be the last block's output ``5 + 3 = 8`` everywhere (not the penultimate block's
+    5, and not the normed 2), and the post-norm grid must be the LayerNorm bias 2 (a
+    constant input has zero variance, so only the bias survives).
+    """
+    enc = _encoder()
+    with torch.no_grad():
+        for parameter in enc._model.parameters():
+            parameter.zero_()
+        enc._model.blocks[-2].mlp.fc2.bias.fill_(5.0)
+        enc._model.blocks[-1].mlp.fc2.bias.fill_(3.0)
+        enc._model.norm.bias.fill_(2.0)
+    batch = torch.ones(2, 3, 32, 48)
+    with torch.inference_mode():
+        pre = enc.encode_tiles_dense_prenorm(batch)
+        post = enc.encode_tiles_dense(batch)
+    assert tuple(pre.shape) == tuple(post.shape) == (2, enc.encode_dim, 2, 3)
+    torch.testing.assert_close(pre, torch.full_like(pre, 8.0))
+    torch.testing.assert_close(post, torch.full_like(post, 2.0))
+
+
+def test_prenorm_dense_rejects_batches_not_divisible_by_the_patch_size():
+    enc = _encoder()
+    with pytest.raises(ValueError, match="divisible by the patch size"):
+        enc.encode_tiles_dense_prenorm(torch.randn(1, 3, 40, 32))
+
+
+def test_validate_dense_request_settings_rejects_unknown_feature_kind():
+    from slide2vec.runtime.dense_regions import validate_dense_request_settings
+
+    geometry = compute_dense_geometry(target_size=32, patch_size=(16, 16))
+    with pytest.raises(ValueError, match="unsupported feature_kind 'tokens'"):
+        validate_dense_request_settings(
+            geometry, pad_mode="reflect", window_size=None,
+            overlap=0.0, feature_kind="tokens", attention_blocks=(-1,),
+            attention_include_registers=False,
+        )

{slide2vec-6.3.3 → slide2vec-6.3.4}/tests/test_dense_stage.py

@@ -533,3 +533,34 @@ def test_region_specs_round_trip_through_request(tmp_path):
     )
     assert request["slides"][1]["read_plan"] == specs[2].read_plan.to_dict()
     assert dense_stage.region_specs_from_request(request) == specs
+
+
+def test_embed_regions_dense_rejects_unsupported_prenorm_before_loading_or_resume(
+    tmp_path, monkeypatch
+):
+    """The persisted ROI route rejects an unimplemented feature kind from registry metadata.
+
+    Nothing is read, loaded or partitioned for resume: an existing artifact directory is
+    left exactly as it was.
+    """
+    from slide2vec.api import Model
+
+    model = Model(name="phikon", device="cpu")
+    monkeypatch.setattr(
+        model, "_load_backend", lambda: (_ for _ in ()).throw(AssertionError("must not load"))
+    )
+    monkeypatch.setattr(
+        dense_stage,
+        "partition_regions_by_resume",
+        lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("must not resume")),
+    )
+    out_dir = tmp_path / "out"
+    (out_dir / "dense_embeddings" / "s0").mkdir(parents=True)
+    (out_dir / "dense_embeddings" / "s0" / "0_0.meta.json").write_text("{}")
+    with pytest.raises(ValueError, match="does not support patch_features_prenorm"):
+        model.embed_regions_dense(
+            [_regions()],
+            dense=_dense(feature_kind="patch_features_prenorm"),
+            execution=ExecutionOptions(num_gpus=1, precision="fp32", output_dir=out_dir),
+        )
+    assert (out_dir / "dense_embeddings" / "s0" / "0_0.meta.json").read_text() == "{}"