chore: import upstream snapshot with attribution

tests/models/backbone/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/models/backbone/test_backbone.py

@@ -0,0 +1,119 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for dual-projector backbone joiner routing."""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from rfdetr.models.backbone import Joiner
+from rfdetr.utilities.tensors import NestedTensor
+
+
+class _FakeBackbone(nn.Module):
+    """Backbone shim used to validate Joiner contract changes."""
+
+    def __init__(
+        self,
+        features: list[NestedTensor],
+        cross_attention_features: list[object] | None,
+    ) -> None:
+        super().__init__()
+        self._features = features
+        self._cross_attention_features = cross_attention_features
+
+    def forward(self, tensor: torch.Tensor | NestedTensor):
+        if isinstance(tensor, torch.Tensor):
+            feats = [f.tensors for f in self._features]
+            masks = [f.mask for f in self._features]
+            return feats, masks, self._cross_attention_features
+        return self._features, self._cross_attention_features
+
+
+class _FakePositionEncoding(nn.Module):
+    """Tiny callable that behaves like a position encoder."""
+
+    def forward(self, nested_tensor: NestedTensor | torch.Tensor, align_dim_orders: bool = False) -> torch.Tensor:
+        if isinstance(nested_tensor, NestedTensor):
+            base = nested_tensor.tensors
+        else:
+            base = nested_tensor
+        if base.dim() == 3:
+            base = base[:, None]
+        return torch.zeros((base.shape[0], 1, base.shape[-2], base.shape[-1]), dtype=base.dtype, device=base.device)
+
+
+def _feature(shape: tuple[int, ...], batch_size: int = 2) -> NestedTensor:
+    channels, height, width = shape
+    return NestedTensor(
+        tensors=torch.ones((batch_size, channels, height, width), dtype=torch.float32),
+        mask=torch.zeros((batch_size, height, width), dtype=torch.bool),
+    )
+
+
+def _input_tensor(batch_size: int = 2) -> tuple[NestedTensor, torch.Tensor]:
+    return (
+        NestedTensor(
+            tensors=torch.ones((batch_size, 3, 16, 16), dtype=torch.float32),
+            mask=torch.zeros((batch_size, 16, 16), dtype=torch.bool),
+        ),
+        torch.ones((batch_size, 3, 16, 16), dtype=torch.float32),
+    )
+
+
+def test_joiner_dual_projector_disabled_contract() -> None:
+    """Joiner should forward one feature stream and a ``None`` cross-attention stream when disabled."""
+    features = [_feature((256, 16, 16))]
+    joiner = Joiner(_FakeBackbone(features, None), _FakePositionEncoding())
+
+    input_tensor, image = _input_tensor()
+
+    _, _, cross_attention = joiner(input_tensor)
+    assert cross_attention is None
+    assert len(joiner(input_tensor)[0]) == 1
+
+    exported = joiner.forward_export(image)
+    assert exported[3] is None
+    assert len(exported[0]) == 1
+    assert exported[2][0].shape == (2, 16, 16)
+
+
+def test_joiner_dual_projector_enabled_contract() -> None:
+    """Joiner should forward cross-attention features in parallel with feature features when enabled."""
+    features = [_feature((256, 16, 16)), _feature((256, 8, 8))]
+    cross_attention_features = [_feature((256, 16, 16)), _feature((256, 8, 8))]
+    joiner = Joiner(_FakeBackbone(features, cross_attention_features), _FakePositionEncoding())
+
+    input_tensor, _ = _input_tensor()
+
+    feature_tensors, _, cross_attention = joiner(input_tensor)
+    assert len(feature_tensors) == len(cross_attention)
+    assert all(f.tensors.shape == c.tensors.shape for f, c in zip(feature_tensors, cross_attention))
+    assert all(f.mask is not None for f in cross_attention)
+
+
+def test_joiner_forward_export_contract() -> None:
+    """Exported joiner contracts should remain 4-tuples and preserve cross-attention stream arity."""
+    exported_features = [torch.ones(2, 256, 16, 16), torch.ones(2, 256, 8, 8)]
+    exported_masks = [torch.zeros(2, 16, 16, dtype=torch.bool), torch.zeros(2, 8, 8, dtype=torch.bool)]
+    export_backbone = _FakeBackbone(
+        [NestedTensor(t, mask) for t, mask in zip(exported_features, exported_masks)],
+        [torch.ones(2, 256, 16, 16), torch.ones(2, 256, 8, 8)],
+    )
+    joiner = Joiner(export_backbone, _FakePositionEncoding())
+
+    outputs = joiner.forward_export(torch.ones(2, 3, 16, 16))
+    feats_out, masks_out, poss, cross_attention = outputs
+
+    assert len(feats_out) == len(exported_features)
+    assert len(masks_out) == len(exported_masks)
+    assert feats_out[0].shape == exported_features[0].shape
+    assert masks_out[0].shape == exported_masks[0].shape
+    assert len(outputs) == 4
+    assert poss[0].shape == exported_features[0][:, :1, :, :].shape
+    assert isinstance(cross_attention, list)
+    assert all(isinstance(feature, torch.Tensor) for feature in cross_attention)

tests/models/backbone/test_backbone_export.py

@@ -0,0 +1,52 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for backbone export behavior."""
+
+import sys
+from types import ModuleType
+
+from rfdetr.models.backbone.backbone import Backbone
+
+
+class TestBackboneExport:
+    """Tests for ``Backbone.export``."""
+
+    def test_export_without_lora_encoder_skips_peft_import_and_warning(self, monkeypatch) -> None:
+        """Non-LoRA exports should not warn just because peft is unavailable."""
+        backbone = object.__new__(Backbone)
+        backbone.encoder = object()
+        warning_messages: list[str] = []
+
+        monkeypatch.delitem(sys.modules, "peft", raising=False)
+        monkeypatch.setattr("rfdetr.models.backbone.backbone.logger.warning", warning_messages.append)
+
+        backbone.export()
+
+        assert warning_messages == []
+
+    def test_export_replaces_peft_encoder_with_merged_encoder(self, monkeypatch) -> None:
+        """Export should replace PEFT wrapper with merged base encoder."""
+
+        class _MergedEncoder:
+            pass
+
+        class _FakePeftModel:
+            def __init__(self) -> None:
+                self._merged = _MergedEncoder()
+
+            def merge_and_unload(self):
+                return self._merged
+
+        peft_module = ModuleType("peft")
+        peft_module.PeftModel = _FakePeftModel
+        monkeypatch.setitem(sys.modules, "peft", peft_module)
+
+        backbone = object.__new__(Backbone)
+        backbone.encoder = _FakePeftModel()
+
+        backbone.export()
+
+        assert isinstance(backbone.encoder, _MergedEncoder)

tests/models/backbone/test_windowed_dino.py

@@ -0,0 +1,544 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+import pytest
+import torch
+
+from rfdetr.models.backbone.dinov2_with_windowed_attn import (
+    Dinov2WithRegistersAttention,
+    Dinov2WithRegistersSdpaAttention,
+    WindowedDinov2WithRegistersBackbone,
+    WindowedDinov2WithRegistersConfig,
+    WindowedDinov2WithRegistersEmbeddings,
+    WindowedDinov2WithRegistersModel,
+    _find_pruneable_heads_and_indices,
+    _get_aligned_output_features_output_indices,
+)
+
+
+def test_window_partition_forward_rectangular_preserves_shapes():
+    """Regression test for WindowedDinov2WithRegistersEmbeddings.forward with rectangular input.
+
+    Ensures window partitioning logic correctly handles H != W.
+    """
+    # Params: H_patches=6, W_patches=4, num_windows=2 -> 3x2 patches per window
+    batch_size, hidden_size, patch_size, num_windows = 1, 64, 16, 2
+    hp, wp, nr = 6, 4, 4
+    h, w = hp * patch_size, wp * patch_size
+
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=hidden_size,
+        patch_size=patch_size,
+        num_windows=num_windows,
+        image_size=h,  # square image_size for positional embeddings
+        num_register_tokens=nr,
+    )
+    model = WindowedDinov2WithRegistersEmbeddings(config)
+
+    # Input is rectangular
+    pixel_values = torch.randn(batch_size, 3, h, w)
+    result = model(pixel_values)
+
+    expected_batch = batch_size * (num_windows**2)
+    expected_seq_len = 1 + nr + (hp // num_windows) * (wp // num_windows)
+
+    assert result.shape == (expected_batch, expected_seq_len, hidden_size)
+
+
+# Before fix in PR #448 the reshape used num_h_patches_per_window in both the height
+# AND width dimension. This only fails when height and width produce different patch
+# counts, so all tests below use non-square images (hp != wp).
+
+
+@pytest.mark.parametrize(
+    "hp, wp, num_windows",
+    [
+        (4, 6, 2),  # wider than tall
+        (6, 4, 2),  # taller than wide
+        (6, 9, 3),  # 3-window grid, non-square
+        (8, 4, 2),  # 2:1 aspect ratio
+    ],
+)
+def test_window_partition_nonsquare_does_not_raise(hp, wp, num_windows):
+    """Before the fix, the reshape used num_h_patches_per_window for the width dimension, so the total element count
+    mismatched and PyTorch raised a RuntimeError for any non-square image.
+
+    The fix replaces that variable with num_w_patches_per_window, making the operation valid for all shapes.
+    """
+    hidden_size, patch_size, nr = 32, 16, 0
+    h, w = hp * patch_size, wp * patch_size
+
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=hidden_size,
+        patch_size=patch_size,
+        num_windows=num_windows,
+        image_size=max(h, w),
+        num_register_tokens=nr,
+    )
+    model = WindowedDinov2WithRegistersEmbeddings(config)
+    pixel_values = torch.randn(1, 3, h, w)
+
+    # This line would raise RuntimeError before the fix
+    result = model(pixel_values)
+
+    expected_batch = num_windows**2
+    expected_seq_len = 1 + (hp // num_windows) * (wp // num_windows)
+    assert result.shape == (expected_batch, expected_seq_len, hidden_size)
+
+
+def test_window_partition_correct_window_content():
+    """Verifies that after windowing each window contains the spatially correct patch tokens — not just that the shape
+    is right.
+
+    Layout with hp=4, wp=6, num_windows=2 (2x2 grid of windows):   Window (0,0): rows 0-1, cols 0-2   Window (0,1): rows
+    0-1, cols 3-5   Window (1,0): rows 2-3, cols 0-2   Window (1,1): rows 2-3, cols 3-5
+
+    Before the fix the reshape used num_h_patches_per_window for the width dim so it raised an error and never produced
+    window content at all.
+    """
+    hidden_size, patch_size, num_windows, nr = 1, 16, 2, 0
+    hp, wp = 4, 6
+    h, w = hp * patch_size, wp * patch_size
+    batch_size = 1
+
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=hidden_size,
+        patch_size=patch_size,
+        num_windows=num_windows,
+        image_size=max(h, w),
+        num_register_tokens=nr,
+        num_hidden_layers=1,
+        num_attention_heads=1,
+    )
+    model = WindowedDinov2WithRegistersEmbeddings(config)
+
+    # Disable position embeddings and cls token so we can track patch identity.
+    # Each patch gets a unique value equal to its flat index (row * wp + col).
+    with torch.no_grad():
+        model.position_embeddings.zero_()
+        model.cls_token.zero_()
+
+    # Build a synthetic patch embedding: patch at (row, col) has value row*wp+col.
+    # Shape after patch projection: (1, hp*wp, 1) — hidden_size=1 for simplicity.
+    patch_ids = torch.arange(hp * wp, dtype=torch.float).view(1, hp * wp, 1)
+
+    # Bypass the full forward pass and exercise the windowing logic directly.
+    pixel_tokens = patch_ids  # (1, 24, 1)
+    pixel_tokens_2d = pixel_tokens.view(batch_size, hp, wp, hidden_size)  # (1,4,6,1)
+
+    num_h_patches_per_window = hp // num_windows  # 2
+    num_w_patches_per_window = wp // num_windows  # 3
+
+    # --- correct reshape (the fix) ---
+    windowed = pixel_tokens_2d.reshape(
+        batch_size * num_windows,
+        num_h_patches_per_window,
+        num_windows,
+        num_w_patches_per_window,
+        hidden_size,
+    )
+    windowed = windowed.permute(0, 2, 1, 3, 4)
+    windowed = windowed.reshape(
+        batch_size * num_windows**2,
+        num_h_patches_per_window * num_w_patches_per_window,
+        hidden_size,
+    )
+
+    # Expected content for each of the 4 windows (6 patches each):
+    expected = torch.tensor(
+        [
+            # Window 0 (rows 0-1, cols 0-2): ids 0,1,2, 6,7,8
+            [[0.0], [1.0], [2.0], [6.0], [7.0], [8.0]],
+            # Window 1 (rows 0-1, cols 3-5): ids 3,4,5, 9,10,11
+            [[3.0], [4.0], [5.0], [9.0], [10.0], [11.0]],
+            # Window 2 (rows 2-3, cols 0-2): ids 12,13,14, 18,19,20
+            [[12.0], [13.0], [14.0], [18.0], [19.0], [20.0]],
+            # Window 3 (rows 2-3, cols 3-5): ids 15,16,17, 21,22,23
+            [[15.0], [16.0], [17.0], [21.0], [22.0], [23.0]],
+        ]
+    )
+    assert torch.equal(windowed, expected), f"Window content mismatch:\n{windowed}\n!=\n{expected}"
+
+
+def test_buggy_reshape_raises_for_nonsquare():
+    """Directly demonstrates what the pre-fix code did: using num_h_patches_per_window in the width position of the
+    reshape causes a RuntimeError when the element count is not divisible by the (wrong) shape.
+
+    With hidden_size=1 and hp=4, wp=6, num_windows=2 the total elements are 24 but the buggy target dims (2,2,2,2,-1)
+    require a non-integer last dimension, so PyTorch raises RuntimeError.
+    """
+    hp, wp = 4, 6  # non-square: width > height
+    num_windows = 2
+    hidden_size = 1  # chosen so total / buggy-fixed-dims is non-integer
+
+    num_h_patches_per_window = hp // num_windows  # 2
+    num_w_patches_per_window = wp // num_windows  # 3
+    batch_size = 1
+
+    # Simulate pixel_tokens_with_pos_embed after the .view() call
+    pixel_tokens_2d = torch.randn(batch_size, hp, wp, hidden_size)
+
+    # The correct reshape (post-fix) must succeed
+    pixel_tokens_2d.reshape(
+        batch_size * num_windows,
+        num_h_patches_per_window,
+        num_windows,
+        num_w_patches_per_window,  # correct
+        hidden_size,
+    )
+
+    # The buggy reshape (pre-fix) must raise RuntimeError:
+    # total elements = 1*4*6*1 = 24,  fixed-dims product = 2*2*2*2 = 16,  16 ∤ 24.
+    with pytest.raises(RuntimeError):
+        pixel_tokens_2d.reshape(
+            batch_size * num_windows,
+            num_h_patches_per_window,
+            num_windows,
+            num_h_patches_per_window,  # bug: height used for width
+            -1,
+        )
+
+
+def test_buggy_reshape_silent_corruption_for_nonsquare():
+    """When hidden_size happens to make the total element count divisible by the buggy target shape, PyTorch does NOT
+    raise — instead the last dimension is inflated, which silently corrupts the tensor layout.
+
+    Pre-fix with hp=4, wp=6, hidden_size=8, num_windows=2:   total elements = 1*4*6*8 = 192 buggy fixed dims = 2*2*2*2 =
+    16  →  last dim inferred as 192/16 = 12 (not 8)
+
+    The fix ensures the correct reshape always yields a last dim equal to hidden_size.
+    """
+    hp, wp = 4, 6
+    num_windows = 2
+    hidden_size = 8
+
+    num_h_patches_per_window = hp // num_windows  # 2
+    num_w_patches_per_window = wp // num_windows  # 3
+    batch_size = 1
+
+    pixel_tokens_2d = torch.randn(batch_size, hp, wp, hidden_size)
+
+    # Buggy reshape silently infers last dim = 12 (not 8)
+    buggy_out = pixel_tokens_2d.reshape(
+        batch_size * num_windows,
+        num_h_patches_per_window,
+        num_windows,
+        num_h_patches_per_window,  # bug
+        -1,
+    )
+    assert buggy_out.shape[-1] != hidden_size, "Buggy reshape should produce wrong last dim"
+
+    # Correct reshape always yields last dim == hidden_size
+    correct_out = pixel_tokens_2d.reshape(
+        batch_size * num_windows,
+        num_h_patches_per_window,
+        num_windows,
+        num_w_patches_per_window,  # fix
+        hidden_size,
+    )
+    assert correct_out.shape[-1] == hidden_size
+
+
+# ---------------------------------------------------------------------------
+# Tests for locally-copied utility functions (removed from transformers v5 public API)
+# ---------------------------------------------------------------------------
+
+
+class TestGetAlignedOutputFeaturesOutputIndices:
+    """Tests for the local copy of get_aligned_output_features_output_indices."""
+
+    def test_both_none_returns_last_stage(self):
+        stage_names = ["stage1", "stage2", "stage3"]
+        features, indices = _get_aligned_output_features_output_indices(None, None, stage_names)
+        assert features == ["stage3"]
+        assert indices == [2]
+
+    def test_only_out_features_derives_indices(self):
+        stage_names = ["stem", "layer1", "layer2", "layer3"]
+        features, indices = _get_aligned_output_features_output_indices(["layer1", "layer3"], None, stage_names)
+        assert features == ["layer1", "layer3"]
+        assert indices == [1, 3]
+
+    def test_only_out_indices_derives_features(self):
+        stage_names = ["stem", "layer1", "layer2", "layer3"]
+        features, indices = _get_aligned_output_features_output_indices(None, [0, 2], stage_names)
+        assert features == ["stem", "layer2"]
+        assert indices == [0, 2]
+
+    def test_both_provided_returns_as_is(self):
+        stage_names = ["stem", "layer1", "layer2"]
+        features, indices = _get_aligned_output_features_output_indices(["layer1"], [1], stage_names)
+        assert features == ["layer1"]
+        assert indices == [1]
+
+    def test_out_indices_converted_to_list(self):
+        """out_indices supplied as a tuple must be returned as a list."""
+        stage_names = ["stem", "layer1", "layer2"]
+        _, indices = _get_aligned_output_features_output_indices(None, (1, 2), stage_names)
+        assert isinstance(indices, list)
+        assert indices == [1, 2]
+
+
+class TestFindPruneableHeadsAndIndices:
+    """Tests for the local copy of find_pruneable_heads_and_indices."""
+
+    def test_no_pruning_returns_full_index(self):
+        heads, index = _find_pruneable_heads_and_indices(set(), n_heads=4, head_size=3, already_pruned_heads=set())
+        assert len(heads) == 0
+        assert len(index) == 12  # 4 * 3, nothing masked
+
+    @pytest.mark.parametrize(
+        "head_to_prune, expected_index",
+        [
+            pytest.param({0}, list(range(3, 12)), id="prune-first-head"),
+            pytest.param({3}, list(range(9)), id="prune-last-head"),
+        ],
+    )
+    def test_prune_single_head_removes_correct_rows(self, head_to_prune, expected_index):
+        # Head N masked → N*head_size indices removed; remaining = n_heads*head_size - head_size = 9
+        heads, index = _find_pruneable_heads_and_indices(
+            head_to_prune, n_heads=4, head_size=3, already_pruned_heads=set()
+        )
+        assert heads == head_to_prune
+        assert len(index) == 9
+        assert index.tolist() == expected_index
+
+    def test_already_pruned_head_adjusts_offset(self):
+        # Head 0 was already pruned. Now pruning head 1 (which is now effective head 0
+        # after offset adjustment) should remove 3 more indices from the effective mask.
+        heads, index = _find_pruneable_heads_and_indices({1}, n_heads=4, head_size=3, already_pruned_heads={0})
+        assert 1 in heads
+        assert len(index) == 9  # 4*3 - 3 pruned
+
+
+# ---------------------------------------------------------------------------
+# Smoke tests for WindowedDinov2WithRegistersBackbone
+# ---------------------------------------------------------------------------
+
+
+def _minimal_backbone_config(**kwargs) -> WindowedDinov2WithRegistersConfig:
+    """Return the smallest valid config for backbone instantiation tests."""
+    defaults = dict(
+        hidden_size=32,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        intermediate_size=64,
+        patch_size=16,
+        image_size=64,
+        num_register_tokens=0,
+        num_windows=1,
+    )
+    defaults.update(kwargs)
+    return WindowedDinov2WithRegistersConfig(**defaults)
+
+
+class TestWindowedDinov2WithRegistersBackbone:
+    """Smoke tests that guard against _init_transformers_backbone() API regressions."""
+
+    @pytest.mark.parametrize(
+        "attr",
+        [
+            pytest.param("stage_names", id="stage_names"),
+            pytest.param("out_features", id="out_features"),
+        ],
+    )
+    def test_instantiation_sets_list_attribute(self, attr):
+        config = _minimal_backbone_config()
+        backbone = WindowedDinov2WithRegistersBackbone(config)
+        assert hasattr(backbone, attr)
+        assert isinstance(getattr(backbone, attr), list)
+        assert len(getattr(backbone, attr)) > 0
+
+    def test_forward_returns_backbone_output(self):
+        config = _minimal_backbone_config()
+        backbone = WindowedDinov2WithRegistersBackbone(config)
+        backbone.eval()
+        pixel_values = torch.randn(1, 3, 64, 64)
+        with torch.no_grad():
+            output = backbone(pixel_values)
+        assert hasattr(output, "feature_maps")
+        assert len(output.feature_maps) == len(backbone.out_features)
+
+
+# ---------------------------------------------------------------------------
+# Test for output_attentions=True SDPA fallback path
+# ---------------------------------------------------------------------------
+
+
+class TestSdpaFallbackWithOutputAttentions:
+    """Guards the output_attentions behaviour in windowed attention."""
+
+    def test_output_attentions_true_raises(self):
+        """Windowed attention explicitly does not support output_attentions=True."""
+        config = _minimal_backbone_config()
+        model = WindowedDinov2WithRegistersModel(config)
+        model.eval()
+        pixel_values = torch.randn(1, 3, 64, 64)
+        with torch.no_grad():
+            with pytest.raises(AssertionError, match="output_attentions is not supported for windowed attention"):
+                model(pixel_values, output_attentions=True)
+
+
+class TestSetAttnImplementation:
+    """Tests for WindowedDinov2WithRegistersModel.set_attn_implementation."""
+
+    @pytest.mark.parametrize(
+        "switches, expected_impl, expected_cls",
+        [
+            pytest.param(["eager"], "eager", Dinov2WithRegistersAttention, id="sdpa-to-eager"),
+            pytest.param(["eager", "sdpa"], "sdpa", Dinov2WithRegistersSdpaAttention, id="roundtrip-back-to-sdpa"),
+        ],
+    )
+    def test_switch_updates_config_and_layers(self, switches, expected_impl, expected_cls):
+        """After each call in *switches*, config and all layer attention modules reflect the final impl."""
+        config = _minimal_backbone_config()
+        model = WindowedDinov2WithRegistersModel(config)
+
+        for impl in switches:
+            model.set_attn_implementation(impl)
+
+        assert model.config._attn_implementation == expected_impl
+        for layer in model.encoder.layer:
+            assert type(layer.attention) is expected_cls
+
+    def test_invalid_implementation_raises(self):
+        """Passing an unknown key raises ValueError with a clear message."""
+        config = _minimal_backbone_config()
+        model = WindowedDinov2WithRegistersModel(config)
+
+        with pytest.raises(ValueError, match="Unknown attn_implementation"):
+            model.set_attn_implementation("flash_attention_2")
+
+
+@pytest.mark.parametrize(
+    "h, w, num_windows, should_raise",
+    [
+        pytest.param(64, 64, 2, False, id="valid-square"),
+        pytest.param(64, 96, 2, False, id="valid-rectangular"),
+        pytest.param(32, 32, 1, False, id="num_windows-1-valid"),
+        pytest.param(33, 64, 2, True, id="h-not-divisible"),
+        pytest.param(64, 33, 2, True, id="w-not-divisible"),
+        pytest.param(33, 33, 2, True, id="both-not-divisible"),
+    ],
+)
+def test_forward_validates_spatial_dims(h: int, w: int, num_windows: int, should_raise: bool) -> None:
+    """WindowedDinov2WithRegistersEmbeddings raises ValueError for incompatible dims.
+
+    Both H and W must be divisible by patch_size * num_windows.  The check must survive Python's -O flag (assert would
+    be silently stripped).
+    """
+    patch_size = 16
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=32,
+        patch_size=patch_size,
+        num_windows=num_windows,
+        image_size=max(h, w),
+        num_register_tokens=0,
+    )
+    model = WindowedDinov2WithRegistersEmbeddings(config)
+    pixel_values = torch.randn(1, 3, h, w)
+    if should_raise:
+        with pytest.raises(ValueError, match="divisible"):
+            model(pixel_values)
+    else:
+        model(pixel_values)  # must not raise
+
+
+def _make_small_model() -> WindowedDinov2WithRegistersModel:
+    """Return the smallest valid WindowedDinov2WithRegistersModel for unit tests."""
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        image_size=32,
+        patch_size=16,
+        num_register_tokens=2,
+    )
+    return WindowedDinov2WithRegistersModel(config)
+
+
+class TestSetAttnImplementationPreservesWeights:
+    """set_attn_implementation must transfer trained weights to the new attention module.
+
+    Before the fix the method replaced each layer's attention with a freshly constructed (randomly initialised) module,
+    silently discarding all trained q/k/v/output weights.
+    """
+
+    @pytest.mark.parametrize(
+        "from_impl, to_impl",
+        [
+            pytest.param("sdpa", "eager", id="sdpa_to_eager"),
+            pytest.param("eager", "sdpa", id="eager_to_sdpa"),
+        ],
+    )
+    def test_query_weight_preserved_after_switch(self, from_impl: str, to_impl: str) -> None:
+        """After switching implementation the query weight tensor must be unchanged."""
+        model = _make_small_model()
+        model.set_attn_implementation(from_impl)
+
+        # Record the query weights of every layer before switching.
+        before = [layer.attention.attention.query.weight.clone() for layer in model.encoder.layer]
+
+        model.set_attn_implementation(to_impl)
+
+        after = [layer.attention.attention.query.weight for layer in model.encoder.layer]
+        for layer_idx, (w_before, w_after) in enumerate(zip(before, after)):
+            assert torch.equal(w_before, w_after), (
+                f"Layer {layer_idx}: query weight changed after set_attn_implementation({from_impl!r} → {to_impl!r})"
+            )
+
+    def test_key_and_value_weights_preserved(self) -> None:
+        """Key and value weights must also survive the implementation switch."""
+        model = _make_small_model()
+        model.set_attn_implementation("sdpa")
+
+        key_before = [layer.attention.attention.key.weight.clone() for layer in model.encoder.layer]
+        val_before = [layer.attention.attention.value.weight.clone() for layer in model.encoder.layer]
+
+        model.set_attn_implementation("eager")
+
+        for layer_idx, layer in enumerate(model.encoder.layer):
+            assert torch.equal(key_before[layer_idx], layer.attention.attention.key.weight), (
+                f"Layer {layer_idx}: key weight changed after implementation switch"
+            )
+            assert torch.equal(val_before[layer_idx], layer.attention.attention.value.weight), (
+                f"Layer {layer_idx}: value weight changed after implementation switch"
+            )
+
+    def test_output_dense_weight_preserved(self) -> None:
+        """The output projection (dense) weight must survive the implementation switch."""
+        model = _make_small_model()
+        model.set_attn_implementation("sdpa")
+
+        dense_before = [layer.attention.output.dense.weight.clone() for layer in model.encoder.layer]
+
+        model.set_attn_implementation("eager")
+
+        for layer_idx, layer in enumerate(model.encoder.layer):
+            assert torch.equal(dense_before[layer_idx], layer.attention.output.dense.weight), (
+                f"Layer {layer_idx}: output dense weight changed after implementation switch"
+            )
+
+    def test_config_updated_after_switch(self) -> None:
+        """config._attn_implementation must reflect the new implementation after switching."""
+        model = _make_small_model()
+        model.set_attn_implementation("eager")
+        assert model.config._attn_implementation == "eager"
+        model.set_attn_implementation("sdpa")
+        assert model.config._attn_implementation == "sdpa"
+
+    def test_attention_module_type_after_switch(self) -> None:
+        """After switching to eager, every layer must hold a non-SDPA attention class."""
+        model = _make_small_model()
+        model.set_attn_implementation("eager")
+        for layer in model.encoder.layer:
+            assert isinstance(layer.attention, Dinov2WithRegistersAttention)
+            assert not isinstance(layer.attention, Dinov2WithRegistersSdpaAttention)
+
+    def test_invalid_implementation_raises_value_error(self) -> None:
+        """An unknown implementation name must raise ValueError before touching any layer."""
+        model = _make_small_model()
+        with pytest.raises(ValueError, match="Unknown attn_implementation"):
+            model.set_attn_implementation("flash_attn")