chore: import upstream snapshot with attribution

tests/models/__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/__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")

tests/models/conftest.py

@@ -0,0 +1,24 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Shared fixtures for the models test suite."""
+
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def reset_torch_safe_globals():
+    """Reset torch serialization safe globals after each test.
+
+    Prevents cross-test state contamination caused by ``_safe_torch_load``'s Attempt 2 path, which calls
+    ``torch.serialization.add_safe_globals``. Without this reset, globals registered by one test bleed into subsequent
+    tests and can mask trust-gate failures.
+    """
+    yield
+    try:
+        torch.serialization.clear_safe_globals()
+    except AttributeError:
+        pass  # torch <2.4 does not have clear_safe_globals

tests/models/heads/__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/heads/test_detection_head.py

@@ -0,0 +1,274 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for _resize_linear(), LWDETR.reinitialize_detection_head(), and _aggregate_keypoint_class_logits().
+
+These tests guard against the out_features staleness bug where in-place .data mutation did not update
+nn.Linear.out_features, causing ONNX export to emit stale (pre-fine-tuning) class counts.
+
+Also covers the spurious "Keypoint class-logit boost has N classes but detection head has M" warning that fired when
+num_keypoints_per_class exactly covered all foreground classes (correct configuration) but the comparison was against
+class_embed.out_features which includes the background slot (+1).
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from torch import nn
+
+from rfdetr.models.lwdetr import LWDETR, _resize_linear
+
+
+def _make_minimal_lwdetr(num_classes: int = 91, two_stage: bool = False) -> LWDETR:
+    """Construct the smallest viable LWDETR without loading pretrained weights.
+
+    Uses a MagicMock backbone and transformer with hidden_dim=4 so the model can be constructed in milliseconds without
+    any network I/O.
+
+    Args:
+        num_classes: Initial number of output classes passed to LWDETR.
+        two_stage: Whether to enable two-stage mode (creates enc_out_class_embed).
+
+    Returns:
+        An LWDETR instance with hidden_dim=4, num_queries=2, group_detr=1.
+
+    Examples:
+        >>> model = _make_minimal_lwdetr(num_classes=91)
+        >>> isinstance(model, LWDETR)
+        True
+    """
+    hidden_dim = 4
+    backbone = MagicMock()
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer.decoder = MagicMock()
+    transformer.decoder.bbox_embed = None
+    return LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=2,
+        group_detr=1,
+        two_stage=two_stage,
+    )
+
+
+def _make_keypoint_lwdetr(num_classes: int, num_keypoints_per_class: list[int]) -> LWDETR:
+    """Construct a minimal keypoint-capable LWDETR with detection head resized to num_classes+1.
+
+    Mirrors what happens after loading a pretrained checkpoint and fine-tuning to num_classes
+    foreground categories: reinitialize_detection_head is called with num_classes+1 (includes
+    background), so class_embed.out_features == num_classes+1 in the returned model.
+
+    Args:
+        num_classes: Number of foreground detection classes.
+        num_keypoints_per_class: Keypoint count per foreground class.
+
+    Returns:
+        An LWDETR with use_grouppose_keypoints=True and class_embed.out_features==num_classes+1.
+
+    Examples:
+        >>> model = _make_keypoint_lwdetr(num_classes=2, num_keypoints_per_class=[17, 4])
+        >>> model.class_embed.out_features
+        3
+    """
+    hidden_dim = 4
+    backbone = MagicMock()
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer.decoder = MagicMock()
+    transformer.decoder.bbox_embed = None
+    transformer.decoder.num_keypoints_per_class = num_keypoints_per_class
+    transformer.decoder.keypoint_class_mask = torch.zeros(1, 1, dtype=torch.bool)
+    transformer.num_keypoints_per_class = num_keypoints_per_class
+    model = LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=2,
+        group_detr=1,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=num_keypoints_per_class,
+    )
+    # Simulate post-checkpoint-load state: detection head includes background slot.
+    model.reinitialize_detection_head(num_classes + 1)
+    return model
+
+
+def _keypoint_tensor(num_keypoints_per_class: list[int], batch: int = 1, seq: int = 1) -> torch.Tensor:
+    """Build a zero keypoint prediction tensor with the shape expected by _aggregate_keypoint_class_logits.
+
+    The second-to-last dimension must equal num_kp_classes * max_kp (padded layout).
+
+    Args:
+        num_keypoints_per_class: Keypoint schema for the model.
+        batch: Batch size dimension.
+        seq: Sequence (query) dimension.
+
+    Returns:
+        Zero tensor of shape (batch, seq, num_kp_classes * max_kp, 8).
+
+    Examples:
+        >>> t = _keypoint_tensor([17, 4])
+        >>> t.shape
+        torch.Size([1, 1, 34, 8])
+    """
+    num_kp_classes = len(num_keypoints_per_class)
+    max_kp = max(num_keypoints_per_class) if any(num_keypoints_per_class) else 1
+    total_padded = num_kp_classes * max_kp
+    return torch.zeros(batch, seq, total_padded, 8)
+
+
+class TestResizeLinear:
+    """Unit tests for _resize_linear() — verifies out_features, weight shape, and bias shape."""
+
+    def test_shrink_out_features(self) -> None:
+        """Shrink: out_features equals the requested smaller class count."""
+        result = _resize_linear(nn.Linear(256, 91), 8)
+        assert result.out_features == 8, f"Expected out_features=8, got {result.out_features}"
+        assert result.weight.shape == (8, 256), f"Expected weight (8, 256), got {result.weight.shape}"
+        assert result.bias is not None
+        assert result.bias.shape == (8,), f"Expected bias (8,), got {result.bias.shape}"
+
+    def test_expand_out_features(self) -> None:
+        """Expand: out_features equals the requested larger class count via tiling."""
+        result = _resize_linear(nn.Linear(256, 10), 25)
+        assert result.out_features == 25, f"Expected out_features=25, got {result.out_features}"
+        assert result.weight.shape == (25, 256), f"Expected weight (25, 256), got {result.weight.shape}"
+        assert result.bias is not None
+        assert result.bias.shape == (25,), f"Expected bias (25,), got {result.bias.shape}"
+
+    def test_same_size_preserves_values(self) -> None:
+        """Same size: shapes and weight/bias values are preserved exactly."""
+        linear = nn.Linear(256, 91)
+        result = _resize_linear(linear, 91)
+        assert result.out_features == 91
+        assert result.weight.shape == (91, 256)
+        assert result.bias is not None
+        assert result.bias.shape == (91,)
+        assert torch.allclose(result.weight.data, linear.weight.data)
+        assert torch.allclose(result.bias.data, linear.bias.data)
+
+    def test_no_bias_returns_no_bias(self) -> None:
+        """Bias=False input: returned module has bias=None and out_features is correct."""
+        linear = nn.Linear(256, 91, bias=False)
+        result = _resize_linear(linear, 8)
+        assert result.out_features == 8, f"Expected out_features=8, got {result.out_features}"
+        assert result.bias is None, "Expected bias=None for bias=False input"
+
+
+class TestReinitializeDetectionHead:
+    """Integration tests for LWDETR.reinitialize_detection_head().
+
+    Uses a minimal LWDETR (hidden_dim=4, no real backbone) to verify that out_features is updated on the replaced
+    nn.Linear modules — the core invariant required for correct ONNX export.
+    """
+
+    def test_updates_class_embed_out_features(self) -> None:
+        """class_embed.out_features must reflect num_classes after reinitialize.
+
+        The `num_outputs_including_background` argument represents the total number of classifier outputs (foreground
+        classes plus background).
+        """
+        num_outputs_including_background = 8
+        model = _make_minimal_lwdetr(num_classes=91)
+        model.reinitialize_detection_head(num_outputs_including_background)
+        assert model.class_embed.out_features == num_outputs_including_background, (
+            f"Expected class_embed.out_features={num_outputs_including_background}, "
+            f"got {model.class_embed.out_features}"
+        )
+        assert model.class_embed.weight.shape == (num_outputs_including_background, 4), (
+            f"Expected weight ({num_outputs_including_background}, 4), got {model.class_embed.weight.shape}"
+        )
+
+    def test_two_stage_updates_enc_out_class_embed(self) -> None:
+        """enc_out_class_embed entries must also have updated out_features in two-stage mode.
+
+        The `num_outputs_including_background` argument represents the total number of classifier outputs (foreground
+        classes plus background).
+        """
+        num_outputs_including_background = 8
+        model = _make_minimal_lwdetr(num_classes=91, two_stage=True)
+        model.reinitialize_detection_head(num_outputs_including_background)
+        enc_embeds = model.transformer.enc_out_class_embed
+        assert len(enc_embeds) > 0, "enc_out_class_embed should be non-empty in two-stage mode"
+        for i, embed in enumerate(enc_embeds):
+            assert embed.out_features == num_outputs_including_background, (
+                f"enc_out_class_embed[{i}].out_features={embed.out_features}, "
+                f"expected {num_outputs_including_background}"
+            )
+            assert embed.weight.shape == (num_outputs_including_background, 4), (
+                f"enc_out_class_embed[{i}].weight.shape={embed.weight.shape}, "
+                f"expected ({num_outputs_including_background}, 4)"
+            )
+
+
+class TestAggregateKeypointClassLogits:
+    """Regression tests for LWDETR._aggregate_keypoint_class_logits().
+
+    Guards against a spurious warning that fired when num_keypoints_per_class exactly covered all
+    foreground classes: class_embed.out_features includes background (+1), so len(schema)==num_classes
+    always satisfied schema_len < detection_num_classes, triggering the warning incorrectly.
+
+    Uses _kp_zero_pad_warned as a proxy for whether the warning fired — the rf-detr logger uses
+    propagate=False which prevents standard caplog capture.
+    """
+
+    @pytest.mark.parametrize(
+        "num_classes,num_keypoints_per_class",
+        [
+            pytest.param(1, [17], id="coco-person-1class"),
+            pytest.param(2, [17, 4], id="basketball-2class"),
+            pytest.param(3, [17, 4, 0], id="3class-schema-covers-all"),
+        ],
+    )
+    def test_no_warning_when_schema_covers_all_foreground_classes(
+        self,
+        num_classes: int,
+        num_keypoints_per_class: list[int],
+    ) -> None:
+        """No warning when num_keypoints_per_class covers exactly all foreground detection classes.
+
+        Regression: the comparison used class_embed.out_features (num_classes+1) instead of
+        num_classes, so a fully correct schema always triggered the spurious mismatch warning.
+        """
+        model = _make_keypoint_lwdetr(num_classes=num_classes, num_keypoints_per_class=num_keypoints_per_class)
+        fake_kp = _keypoint_tensor(num_keypoints_per_class)
+
+        model._aggregate_keypoint_class_logits(fake_kp)
+
+        assert not model._kp_zero_pad_warned, (
+            f"Spurious warning fired for schema={num_keypoints_per_class} with num_classes={num_classes}"
+        )
+
+    def test_warning_fires_when_schema_shorter_than_foreground_classes(self) -> None:
+        """Warning fires when schema covers fewer classes than foreground detection classes.
+
+        Scenario: 3 foreground classes but schema only covers 1 (e.g. only person has keypoints).
+        The two uncovered foreground classes receive zero boost — a real mismatch worth warning about.
+        """
+        model = _make_keypoint_lwdetr(num_classes=3, num_keypoints_per_class=[17])
+        fake_kp = _keypoint_tensor([17])
+
+        model._aggregate_keypoint_class_logits(fake_kp)
+
+        assert model._kp_zero_pad_warned, "Expected warning flag set for schema shorter than foreground class count"
+
+    def test_output_shape_matches_detection_head(self) -> None:
+        """Output shape is (batch, seq, detection_num_classes) regardless of schema length."""
+        num_classes = 2
+        schema = [17, 4]
+        model = _make_keypoint_lwdetr(num_classes=num_classes, num_keypoints_per_class=schema)
+        batch, seq = 2, 10
+        fake_kp = _keypoint_tensor(schema, batch=batch, seq=seq)
+
+        out = model._aggregate_keypoint_class_logits(fake_kp)
+
+        assert out.shape == (batch, seq, num_classes + 1), (
+            f"Expected shape {(batch, seq, num_classes + 1)}, got {out.shape}"
+        )

tests/models/heads/test_keypoint_heads.py

@@ -0,0 +1,278 @@
+# ------------------------------------------------------------------------
+# 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.heads import ConditionalQueryInitializer
+from rfdetr.models.heads.keypoints import (
+    compute_keypoint_matching_cost,
+    compute_l1_keypoint_loss,
+)
+
+
+def test_conditional_query_initializer_shape() -> None:
+    """Initializer output should have expected batch/query/out dimensions."""
+    initializer = ConditionalQueryInitializer(dim=32, num_queries=11, out_dim=16)
+    query_features = torch.randn(3, 32)
+    queries = initializer(query_features)
+
+    assert queries.shape == (3, 11, 16)
+
+
+def test_conditional_query_initializer_zero_adaln_identity() -> None:
+    """A zeroed AdaLN gate should make initializer return the unmodified learned queries."""
+    initializer = ConditionalQueryInitializer(dim=16, num_queries=5, out_dim=16)
+    query_features = torch.randn(4, 16)
+    output = initializer(query_features)
+    expected = initializer.queries.unsqueeze(0).expand_as(output)
+
+    assert torch.equal(output, expected)
+
+
+def test_compute_l1_keypoint_loss_smoke() -> None:
+    """Loss helper should emit four finite vectors with matching target batch shape."""
+    pred_keypoints = torch.randn(3, 17, 7)
+    target_keypoints = torch.rand(3, 17, 3)
+    target_keypoints[:, :, 2] = 2.0
+    target_classes = torch.tensor([0, 0, 0], dtype=torch.int64)
+    target_areas = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
+    losses = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=target_classes,
+        target_areas=target_areas,
+        num_keypoints_per_class=[17],
+    )
+
+    assert len(losses) == 4
+    for loss in losses:
+        assert loss.shape == (3,)
+        assert torch.isfinite(loss).all()
+
+
+def test_compute_l1_keypoint_loss_skips_visible_zero_area_nll_residuals() -> None:
+    """Visible keypoints on zero-area targets should not produce non-finite Gaussian NLL."""
+    pred_keypoints = torch.zeros(1, 17, 7)
+    target_keypoints = torch.rand(1, 17, 3)
+    target_keypoints[:, :, 2] = 2.0
+    losses = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([0.0], dtype=torch.float32),
+        num_keypoints_per_class=[17],
+    )
+
+    for loss in losses:
+        assert torch.isfinite(loss).all()
+
+
+def test_compute_l1_keypoint_loss_uses_raw_rflow_gaussian_nll() -> None:
+    """Perfect keypoints should use raw r-flow NLL without a floor shift."""
+    pred_keypoints = torch.zeros(1, 1, 7)
+    pred_keypoints[:, :, 4] = 0.3
+    pred_keypoints[:, :, 6] = -0.2
+    target_keypoints = torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32)
+
+    _, _, _, nll = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([1.0], dtype=torch.float32),
+        num_keypoints_per_class=[1],
+    )
+
+    torch.testing.assert_close(nll, torch.tensor([-0.1]), rtol=1e-4, atol=1e-6)
+
+
+def test_compute_l1_keypoint_loss_does_not_clamp_log_cholesky_nll() -> None:
+    """Large precision log-diagonals should remain raw to match r-flow."""
+    pred_keypoints = torch.zeros(1, 1, 7)
+    pred_keypoints[:, :, 4] = 25.0
+    target_keypoints = torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32)
+
+    _, _, _, nll = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([1.0], dtype=torch.float32),
+        num_keypoints_per_class=[1],
+    )
+
+    torch.testing.assert_close(nll, torch.tensor([-25.0]), rtol=1e-4, atol=1e-6)
+
+
+def test_compute_l1_keypoint_loss_raw_nll_gradients_match_reference_formula() -> None:
+    """The implemented NLL gradients should match the raw r-flow Gaussian formula."""
+    pred_keypoints = torch.tensor([[[0.2, -0.1, 0.0, 0.0, 0.3, 0.1, -0.2]]], requires_grad=True)
+    target_keypoints = torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32)
+    target_areas = torch.tensor([1.0], dtype=torch.float32)
+    _, _, _, nll = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=target_areas,
+        num_keypoints_per_class=[1],
+    )
+    nll.sum().backward()
+    grad = pred_keypoints.grad.detach().clone()
+
+    raw_pred_keypoints = pred_keypoints.detach().clone().requires_grad_(True)
+    dx = raw_pred_keypoints[:, :, 0] - target_keypoints[:, :, 0]
+    dy = raw_pred_keypoints[:, :, 1] - target_keypoints[:, :, 1]
+    log_l11 = raw_pred_keypoints[:, :, 4]
+    l21 = raw_pred_keypoints[:, :, 5]
+    log_l22 = raw_pred_keypoints[:, :, 6]
+    u0 = log_l11.exp() * dx + l21 * dy
+    u1 = log_l22.exp() * dy
+    raw_nll = 0.5 * (u0 * u0 + u1 * u1) / target_areas.unsqueeze(1) - (log_l11 + log_l22)
+    raw_nll.sum().backward()
+
+    torch.testing.assert_close(nll.detach(), raw_nll.detach().reshape(-1), rtol=1e-4, atol=1e-6)
+    torch.testing.assert_close(grad, raw_pred_keypoints.grad, rtol=1e-4, atol=1e-6)
+
+
+def test_compute_l1_keypoint_loss_rejects_missing_schema() -> None:
+    """Missing keypoint schema should fail before producing zero supervision."""
+    pred_keypoints = torch.randn(1, 17, 7)
+    target_keypoints = torch.rand(1, 17, 3)
+
+    with pytest.raises(ValueError, match="num_keypoints_per_class must be non-empty"):
+        compute_l1_keypoint_loss(
+            all_pred_keypoints=pred_keypoints,
+            target_keypoints=target_keypoints,
+            target_classes=torch.tensor([0], dtype=torch.int64),
+            target_areas=torch.tensor([1.0], dtype=torch.float32),
+            num_keypoints_per_class=[],
+        )
+
+
+def test_compute_keypoint_matching_cost_smoke() -> None:
+    """Matching-cost helper should return a four-term cost tensor for each target."""
+    all_pred_keypoints = torch.randn(2, 4, 17, 7)
+    target_keypoints = torch.rand(2, 17, 3)
+    target_keypoints[:, :, 2] = 2.0
+    target_classes = torch.tensor([0, 0], dtype=torch.int64)
+    target_areas = torch.tensor([1.0, 2.0], dtype=torch.float32)
+    cost_l1, cost_findable, cost_visible, cost_nll = compute_keypoint_matching_cost(
+        all_pred_keypoints=all_pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=target_classes,
+        target_areas=target_areas,
+        num_keypoints_per_class=[17],
+    )
+
+    assert cost_l1.shape == (2, 4, 2)
+    assert cost_findable.shape == (2, 4, 2)
+    assert cost_visible.shape == (2, 4, 2)
+    assert cost_nll.shape == (2, 4, 2)
+    assert torch.isfinite(cost_l1).all()
+    assert torch.isfinite(cost_findable).all()
+    assert torch.isfinite(cost_visible).all()
+    assert torch.isfinite(cost_nll).all()
+
+
+def test_compute_keypoint_matching_cost_skips_zero_area_nll_residuals() -> None:
+    """Zero-area targets should not produce non-finite keypoint matching costs."""
+    all_pred_keypoints = torch.zeros(1, 2, 17, 7)
+    target_keypoints = torch.rand(1, 17, 3)
+    target_keypoints[:, :, 2] = 2.0
+    costs = compute_keypoint_matching_cost(
+        all_pred_keypoints=all_pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([0.0], dtype=torch.float32),
+        num_keypoints_per_class=[17],
+    )
+
+    for cost in costs:
+        assert torch.isfinite(cost).all()
+
+
+def test_compute_keypoint_matching_cost_does_not_clamp_log_cholesky_nll() -> None:
+    """Matching NLL should use raw precision log-diagonals to match r-flow."""
+    all_pred_keypoints = torch.zeros(1, 1, 1, 7)
+    all_pred_keypoints[:, :, :, 4] = 25.0
+    target_keypoints = torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32)
+
+    _, _, _, cost_nll = compute_keypoint_matching_cost(
+        all_pred_keypoints=all_pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([1.0], dtype=torch.float32),
+        num_keypoints_per_class=[1],
+    )
+
+    torch.testing.assert_close(cost_nll, torch.tensor([[[-25.0]]]), rtol=1e-4, atol=1e-6)
+
+
+def test_compute_keypoint_matching_cost_rejects_missing_schema() -> None:
+    """Missing keypoint schema should fail before matcher costs become keypoint no-ops."""
+    all_pred_keypoints = torch.randn(1, 2, 17, 7)
+    target_keypoints = torch.rand(1, 17, 3)
+
+    with pytest.raises(ValueError, match="num_keypoints_per_class must be non-empty"):
+        compute_keypoint_matching_cost(
+            all_pred_keypoints=all_pred_keypoints,
+            target_keypoints=target_keypoints,
+            target_classes=torch.tensor([0], dtype=torch.int64),
+            target_areas=torch.tensor([1.0], dtype=torch.float32),
+            num_keypoints_per_class=[],
+        )
+
+
+class TestComputeKeypointMatchingCostSmoke:
+    """Group: compute_keypoint_matching_cost — shape and boundary checks."""
+
+    def test_n_targets_zero_returns_four_empty_cost_tensors(self) -> None:
+        """Empty target set should return four finite (B, Q, 0) cost tensors immediately."""
+        b, q = 2, 4
+        all_pred_keypoints = torch.randn(b, q, 17, 7)
+
+        cost_l1, cost_findable, cost_visible, cost_nll = compute_keypoint_matching_cost(
+            all_pred_keypoints=all_pred_keypoints,
+            target_keypoints=torch.empty(0, 17, 3),
+            target_classes=torch.empty(0, dtype=torch.int64),
+            target_areas=torch.empty(0),
+            num_keypoints_per_class=[17],
+        )
+
+        for cost, name in (
+            (cost_l1, "cost_l1"),
+            (cost_findable, "cost_findable"),
+            (cost_visible, "cost_visible"),
+            (cost_nll, "cost_nll"),
+        ):
+            assert cost.shape == (b, q, 0), f"{name}: expected shape ({b}, {q}, 0), got {cost.shape}"
+            assert torch.isfinite(cost).all(), f"{name}: expected all-finite tensor, got non-finite values"
+
+
+class TestComputeL1KeypointLossOobClass:
+    """Group: compute_l1_keypoint_loss — out-of-range class index handling."""
+
+    def test_class_index_out_of_range_returns_zero_losses_without_raising(self) -> None:
+        """Out-of-range class index should emit a warning and return zeros, not raise."""
+        pred_keypoints = torch.randn(1, 17, 7)
+        target_keypoints = torch.rand(1, 17, 3)
+        target_keypoints[:, :, 2] = 2.0
+        # class index 2 is out of range for num_keypoints_per_class=[17] (only class 0 defined)
+        result = compute_l1_keypoint_loss(
+            all_pred_keypoints=pred_keypoints,
+            target_keypoints=target_keypoints,
+            target_classes=torch.tensor([2], dtype=torch.int64),
+            target_areas=torch.tensor([1.0], dtype=torch.float32),
+            num_keypoints_per_class=[17],
+        )
+
+        assert len(result) == 4, f"Expected 4-tuple, got {len(result)} elements"
+        for i, loss in enumerate(result):
+            assert loss.shape == (1,), f"Loss[{i}]: expected shape (1,), got {loss.shape}"
+            torch.testing.assert_close(
+                loss,
+                torch.zeros(1),
+                msg=f"Loss[{i}]: expected all zeros for out-of-range class, got {loss}",
+            )

tests/models/heads/test_segmentation_head.py

@@ -0,0 +1,263 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for DepthwiseConvBlock and _DepthwiseConvWithoutCuDNN (segmentation head)."""
+
+from contextlib import contextmanager
+
+import pytest
+import torch
+
+from rfdetr.models.heads.segmentation import DepthwiseConvBlock
+
+
+@pytest.fixture(autouse=True)
+def _reset_random_seeds() -> None:
+    """Reset random seeds before each test for reproducibility."""
+    torch.manual_seed(42)
+
+
+@pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param("cpu", id="cpu"),
+        pytest.param(
+            "cuda",
+            id="gpu",
+            marks=[
+                pytest.mark.gpu,
+                pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason="CUDA is not available",
+                ),
+            ],
+        ),
+    ],
+)
+def test_depthwise_conv_block_forward(device: str) -> None:
+    """DepthwiseConvBlock forward pass produces correct output shape without error."""
+    block = DepthwiseConvBlock(dim=8).to(device)
+    x = torch.randn(1, 8, 4, 4, device=device)
+    y = block(x)
+    assert y.shape == x.shape
+
+
+def test_depthwise_conv_forward_disables_cudnn(monkeypatch) -> None:
+    """Depthwise conv should execute with cuDNN disabled during forward."""
+    block = DepthwiseConvBlock(dim=8)
+    enabled_calls: list[bool] = []
+    original_flags = torch.backends.cudnn.flags
+
+    @contextmanager
+    def _tracking_flags(*, enabled: bool):
+        enabled_calls.append(enabled)
+        with original_flags(enabled=enabled):
+            yield
+
+    monkeypatch.setattr(torch.backends.cudnn, "flags", _tracking_flags)
+
+    x = torch.randn(1, 8, 4, 4)
+    y = block(x)
+    assert y.shape == x.shape
+    assert enabled_calls, "torch.backends.cudnn.flags was never called"
+    assert all(not e for e in enabled_calls)
+
+
+def test_depthwise_conv_backward_disables_cudnn(monkeypatch) -> None:
+    """Backward pass must also run with cuDNN disabled (issue #731).
+
+    The previous fix (PR #728) only wrapped the forward pass in a context manager.  The backward kernels ran with cuDNN
+    re-enabled, causing RuntimeError on T4/P100 GPUs.
+    """
+    block = DepthwiseConvBlock(dim=8)
+    enabled_calls: list[bool] = []
+    original_flags = torch.backends.cudnn.flags
+
+    @contextmanager
+    def _tracking_flags(*, enabled: bool):
+        enabled_calls.append(enabled)
+        with original_flags(enabled=enabled):
+            yield
+
+    monkeypatch.setattr(torch.backends.cudnn, "flags", _tracking_flags)
+
+    x = torch.randn(1, 8, 4, 4, requires_grad=True)
+    y = block(x)
+    y.sum().backward()
+
+    assert x.grad is not None
+    assert x.grad.shape == x.shape
+    # cuDNN must be disabled for both forward and backward
+    assert len(enabled_calls) >= 2
+    assert all(not e for e in enabled_calls)
+
+
+@pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param("cpu", id="cpu"),
+        pytest.param(
+            "cuda",
+            id="gpu",
+            marks=[
+                pytest.mark.gpu,
+                pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason="CUDA is not available",
+                ),
+            ],
+        ),
+    ],
+)
+def test_depthwise_conv_backward_produces_correct_gradients(device: str) -> None:
+    """Backward pass through DepthwiseConvBlock produces valid gradients."""
+    block = DepthwiseConvBlock(dim=8).to(device)
+    x = torch.randn(1, 8, 4, 4, device=device, requires_grad=True)
+    y = block(x)
+    y.sum().backward()
+    assert x.grad is not None
+    assert x.grad.shape == x.shape
+    assert torch.isfinite(x.grad).all()
+
+
+def test_depthwise_conv_gradients_match_reference() -> None:
+    """Custom autograd Function gradients match nn.Conv2d gradients.
+
+    Verifies that _DepthwiseConvWithoutCuDNN produces the same gradients as a standard nn.Conv2d forward+backward (run
+    with cuDNN disabled globally).
+    """
+    torch.manual_seed(42)
+    dim = 8
+    block = DepthwiseConvBlock(dim=dim)
+
+    # Reference: run standard nn.Conv2d with cuDNN globally disabled
+    x_ref = torch.randn(1, dim, 4, 4, requires_grad=True)
+    with torch.backends.cudnn.flags(enabled=False):
+        y_ref = block.dwconv(x_ref)
+    y_ref.sum().backward()
+
+    x_ref_grad = x_ref.grad.clone()
+    weight_ref_grad = block.dwconv.weight.grad.clone()
+    bias_ref_grad = block.dwconv.bias.grad.clone()
+
+    # Our implementation via _depthwise_conv.  zero_grad() so that the second
+    # backward does not accumulate into weight.grad from the first run.
+    block.zero_grad()
+    x_test = x_ref.detach().clone().requires_grad_(True)
+    y_test = block._depthwise_conv(x_test)
+    y_test.sum().backward()
+
+    assert torch.allclose(y_ref, y_test, atol=1e-6)
+    assert torch.allclose(x_ref_grad, x_test.grad, atol=1e-6)
+    assert torch.allclose(weight_ref_grad, block.dwconv.weight.grad, atol=1e-6)
+    assert torch.allclose(bias_ref_grad, block.dwconv.bias.grad, atol=1e-6)
+
+
+def test_depthwise_conv_backward_fp16_grad_output() -> None:
+    """Backward must not crash when grad_output is fp16 (AMP 16-mixed on T4/P100).
+
+    On T4/P100, trainer resolves amp=True to '16-mixed'.  In that mode the backward receives fp16 grad_output while the
+    saved weight stays fp32. Without explicit dtype casting, conv2d_input raises:
+        RuntimeError: expected scalar type Half but found Float
+    """
+    dim = 8
+    block = DepthwiseConvBlock(dim=dim)
+    x = torch.randn(1, dim, 4, 4, requires_grad=True)
+
+    # Simulate 16-mixed backward: forward in fp32, grad_output arrives as fp16
+    y = block._depthwise_conv(x)
+    grad_output = torch.ones_like(y, dtype=torch.float16)
+    y.backward(grad_output)
+
+    assert x.grad is not None
+    assert x.grad.dtype == torch.float32
+    assert torch.isfinite(x.grad).all()
+
+
+def test_depthwise_conv_backward_bf16_activation_keeps_grads_fp32() -> None:
+    """grad_input and weight.grad must be fp32 when saved activation x is bf16 (issue #959).
+
+    Under bf16-mixed AMP, the activation x entering _DepthwiseConvWithoutCuDNN is bf16 while weight stays fp32.  The old
+    code cast grad_input back to x.dtype (bf16), propagating bf16 gradients to fp32 backbone parameters so that
+    param.grad.dtype became bf16.  Fused AdamW then crashed with 'params, grads, exp_avgs, and exp_avg_sqs must have
+    same dtype, device, and layout' (see issue #959).  The fix keeps grad_input in weight.dtype (fp32).
+
+    This test drives the backward directly with a bf16 saved activation to reproduce the dtype that is present at
+    training time without requiring a GPU.
+    """
+    import types
+
+    from rfdetr.models.heads.segmentation import _DepthwiseConvWithoutCuDNN
+
+    dim = 8
+    weight = torch.randn(dim, 1, 3, 3, requires_grad=True)  # fp32 parameter (never cast by AMP)
+    x_bf16 = torch.randn(1, dim, 4, 4, dtype=torch.bfloat16)  # bf16 activation (cast by AMP forward)
+    grad_output = torch.ones(1, dim, 4, 4, dtype=torch.bfloat16)  # bf16 grad (from bf16 backward)
+
+    # Build a minimal context mirroring what ctx would contain after the AMP forward pass.
+    ctx = types.SimpleNamespace()
+    ctx.saved_tensors = (x_bf16, weight)
+    ctx.has_bias = False
+    ctx.stride = (1, 1)
+    ctx.padding = (1, 1)
+    ctx.dilation = (1, 1)
+    ctx.groups = dim
+    ctx.needs_input_grad = [True, True, False, False, False, False, False]
+
+    grad_input, grad_weight, *_ = _DepthwiseConvWithoutCuDNN.backward(ctx, grad_output)
+
+    assert grad_input is not None, "grad_input should not be None when needs_input_grad[0] is True"
+    assert grad_input.dtype == torch.float32, (
+        f"grad_input is {grad_input.dtype} — bf16 grad_input propagates to fp32 backbone params "
+        "and crashes fused AdamW (issue #959)"
+    )
+    assert grad_weight is not None, "grad_weight should not be None when needs_input_grad[1] is True"
+    assert grad_weight.dtype == torch.float32, (
+        f"grad_weight is {grad_weight.dtype} — weight grad must stay fp32 to match param dtype"
+    )
+
+
+def test_depthwise_conv_no_cudnn_bias_none() -> None:
+    """_DepthwiseConvWithoutCuDNN forward and backward work correctly with bias=None.
+
+    Exercises the ctx.has_bias=False branch in forward and the grad_bias=None return in backward — never reached via
+    DepthwiseConvBlock (always has bias).
+    """
+    from rfdetr.models.heads.segmentation import _DepthwiseConvWithoutCuDNN
+
+    dim = 8
+    weight = torch.randn(dim, 1, 3, 3, requires_grad=True)
+    x = torch.randn(1, dim, 4, 4, requires_grad=True)
+    y = _DepthwiseConvWithoutCuDNN.apply(x, weight, None, (1, 1), (1, 1), (1, 1), dim)
+    y_ref = torch.nn.functional.conv2d(x.detach(), weight.detach(), None, stride=1, padding=1, dilation=1, groups=dim)
+    assert torch.allclose(y.detach(), y_ref, atol=1e-6)
+    y.sum().backward()
+    assert x.grad is not None
+    assert x.grad.shape == x.shape
+    assert torch.isfinite(x.grad).all()
+    assert weight.grad is not None
+    assert weight.grad.shape == weight.shape
+    assert torch.isfinite(weight.grad).all()
+
+
+@pytest.mark.parametrize(
+    "layer_scale_init_value", [pytest.param(0, id="no_gamma"), pytest.param(1e-6, id="with_gamma")]
+)
+def test_depthwise_conv_block_layer_scale(layer_scale_init_value: float) -> None:
+    """DepthwiseConvBlock with and without layer scaling produces valid output and gradients.
+
+    Exercises the gamma=None (layer_scale_init_value=0) and gamma!=None (layer_scale_init_value>0) branches in
+    DepthwiseConvBlock.forward().
+    """
+    block = DepthwiseConvBlock(dim=8, layer_scale_init_value=layer_scale_init_value)
+    x = torch.randn(1, 8, 4, 4, requires_grad=True)
+    y = block(x)
+    assert y.shape == x.shape
+    y.sum().backward()
+    assert x.grad is not None
+    assert torch.isfinite(x.grad).all()
+    if layer_scale_init_value > 0:
+        assert block.gamma is not None
+        assert block.gamma.grad is not None

tests/models/test_builder_characterization.py

@@ -0,0 +1,433 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Characterization tests for build_model() and build_criterion_and_postprocessors().
+
+These tests pin the current behavior of the legacy namespace-based builder functions. They serve as a safety net during
+the config-native builder refactoring: any change that alters these outputs is a regression.
+
+All tests in this file must pass against the CURRENT codebase.
+"""
+
+import pytest
+import torch
+
+from rfdetr._namespace import _namespace_from_configs
+from rfdetr.config import (
+    RFDETRBaseConfig,
+    RFDETRKeypointPreviewConfig,
+    RFDETRNanoConfig,
+    RFDETRSegNanoConfig,
+    SegmentationTrainConfig,
+    TrainConfig,
+)
+from rfdetr.models.criterion import SetCriterion
+from rfdetr.models.lwdetr import LWDETR, build_criterion_and_postprocessors, build_model
+from rfdetr.models.postprocess import PostProcess
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_ns(mc=None, tc=None):
+    """Build a namespace suitable for builder functions."""
+    mc = mc or RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+    tc = tc or TrainConfig(dataset_dir="/tmp")
+    return _namespace_from_configs(mc, tc)
+
+
+# ---------------------------------------------------------------------------
+# build_model characterization
+# ---------------------------------------------------------------------------
+
+
+class TestBuildModelCharacterization:
+    """Pin current build_model() behaviour for the standard code path."""
+
+    def test_returns_lwdetr_instance(self) -> None:
+        ns = _make_ns()
+        model = build_model(ns)
+        assert isinstance(model, LWDETR)
+
+    def test_num_classes_plus_one(self) -> None:
+        """build_model applies the +1 background class convention."""
+        mc = RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.class_embed.out_features == 6
+
+    def test_num_queries_forwarded(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.num_queries == mc.num_queries
+
+    @pytest.mark.parametrize(
+        "config_class, expected_queries",
+        [
+            pytest.param(RFDETRBaseConfig, 300, id="base"),
+            pytest.param(RFDETRNanoConfig, 300, id="nano"),
+            pytest.param(RFDETRSegNanoConfig, 100, id="seg_nano"),
+        ],
+    )
+    def test_num_queries_per_config_variant(self, config_class, expected_queries) -> None:
+        mc = config_class(pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.num_queries == expected_queries
+
+    def test_segmentation_head_none_for_detection(self) -> None:
+        ns = _make_ns()
+        model = build_model(ns)
+        assert model.segmentation_head is None
+
+    def test_segmentation_head_present_for_seg_config(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.segmentation_head is not None
+
+    def test_aux_loss_enabled_by_default(self) -> None:
+        ns = _make_ns()
+        model = build_model(ns)
+        assert model.aux_loss is True
+
+    def test_group_detr_forwarded(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.group_detr == mc.group_detr
+
+    def test_num_feature_levels_set_on_args(self) -> None:
+        """build_model mutates args.num_feature_levels = len(projector_scale)."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        build_model(ns)
+        assert ns.num_feature_levels == len(mc.projector_scale)
+
+    @pytest.mark.parametrize(
+        "config_class, expected_param_count_range",
+        [
+            pytest.param(RFDETRBaseConfig, (25_000_000, 40_000_000), id="base"),
+            pytest.param(RFDETRNanoConfig, (25_000_000, 40_000_000), id="nano"),
+        ],
+    )
+    def test_param_count_in_expected_range(self, config_class, expected_param_count_range) -> None:
+        """Sanity check that the model has a plausible number of parameters."""
+        mc = config_class(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        total = sum(p.numel() for p in model.parameters())
+        low, high = expected_param_count_range
+        assert low <= total <= high, f"Expected param count in [{low}, {high}], got {total}"
+
+    def test_encoder_only_returns_triple(self) -> None:
+        """When encoder_only=True, build_model returns (encoder, None, None)."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        ns.encoder_only = True
+        result = build_model(ns)
+        assert isinstance(result, tuple), f"Expected tuple, got {type(result)}"
+        assert len(result) == 3
+        encoder, second, third = result
+        assert second is None
+        assert third is None
+        assert encoder is not None
+
+    def test_backbone_only_returns_triple(self) -> None:
+        """When backbone_only=True, build_model returns (backbone, None, None)."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        ns.backbone_only = True
+        result = build_model(ns)
+        assert isinstance(result, tuple), f"Expected tuple, got {type(result)}"
+        assert len(result) == 3
+        backbone, second, third = result
+        assert second is None
+        assert third is None
+        assert backbone is not None
+
+
+# ---------------------------------------------------------------------------
+# build_criterion_and_postprocessors characterization
+# ---------------------------------------------------------------------------
+
+
+class TestBuildCriterionCharacterization:
+    """Pin current build_criterion_and_postprocessors() behaviour."""
+
+    def test_returns_criterion_and_postprocess(self) -> None:
+        ns = _make_ns()
+        criterion, postprocess = build_criterion_and_postprocessors(ns)
+        assert isinstance(criterion, SetCriterion)
+        assert isinstance(postprocess, PostProcess)
+
+    def test_detection_losses_list(self) -> None:
+        """Detection-only config has exactly ['labels', 'boxes', 'cardinality']."""
+        ns = _make_ns()
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.losses == ["labels", "boxes", "cardinality"]
+
+    def test_segmentation_losses_include_masks(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        ns = _make_ns(mc=mc, tc=tc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert "masks" in criterion.losses
+
+    def test_num_select_forwarded_to_postprocess(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        _, postprocess = build_criterion_and_postprocessors(ns)
+        assert postprocess.num_select == 100
+
+    def test_num_select_default_for_base(self) -> None:
+        ns = _make_ns()
+        _, postprocess = build_criterion_and_postprocessors(ns)
+        assert postprocess.num_select == 300
+
+    def test_weight_dict_contains_base_losses(self) -> None:
+        ns = _make_ns()
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert "loss_ce" in criterion.weight_dict
+        assert "loss_bbox" in criterion.weight_dict
+        assert "loss_giou" in criterion.weight_dict
+
+    def test_weight_dict_values_match_namespace(self) -> None:
+        ns = _make_ns()
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.weight_dict["loss_ce"] == ns.cls_loss_coef
+        assert criterion.weight_dict["loss_bbox"] == ns.bbox_loss_coef
+        assert criterion.weight_dict["loss_giou"] == ns.giou_loss_coef
+
+    def test_segmentation_weight_dict_contains_mask_losses(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        ns = _make_ns(mc=mc, tc=tc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert "loss_mask_ce" in criterion.weight_dict
+        assert "loss_mask_dice" in criterion.weight_dict
+
+    def test_aux_loss_expands_weight_dict(self) -> None:
+        """With aux_loss=True and 3 dec_layers, weight_dict has aux entries _0 and _1."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        assert ns.aux_loss is True
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        # dec_layers=3 -> 2 aux layers (0 and 1)
+        assert "loss_ce_0" in criterion.weight_dict
+        assert "loss_ce_1" in criterion.weight_dict
+
+    def test_two_stage_adds_enc_losses(self) -> None:
+        """With two_stage=True, weight_dict has '_enc' suffix entries."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        assert ns.two_stage is True
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert "loss_ce_enc" in criterion.weight_dict
+        assert "loss_bbox_enc" in criterion.weight_dict
+        assert "loss_giou_enc" in criterion.weight_dict
+
+    def test_criterion_num_classes_plus_one(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.num_classes == 6
+
+    def test_focal_alpha_forwarded(self) -> None:
+        ns = _make_ns()
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.focal_alpha == pytest.approx(0.25)
+
+    def test_group_detr_forwarded_to_criterion(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.group_detr == mc.group_detr
+
+    def test_segmentation_criterion_has_mask_point_sample_ratio(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        ns = _make_ns(mc=mc, tc=tc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.mask_point_sample_ratio == 16
+
+    def test_ia_bce_loss_forwarded(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.ia_bce_loss == mc.ia_bce_loss
+
+
+# ---------------------------------------------------------------------------
+# _build_model_context characterization
+# ---------------------------------------------------------------------------
+
+
+class TestBuildModelContextCharacterization:
+    """Pin current _build_model_context() behaviour.
+
+    _build_model_context is the inference-path factory used by RFDETR.get_model(). It has zero test coverage today.
+    """
+
+    def test_returns_model_context(self) -> None:
+        from rfdetr.detr import ModelContext, _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert isinstance(ctx, ModelContext)
+
+    def test_model_is_lwdetr(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert isinstance(ctx.model, LWDETR)
+
+    def test_postprocess_is_postprocess(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert isinstance(ctx.postprocess, PostProcess)
+
+    def test_resolution_from_config(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.resolution == mc.resolution
+
+    def test_device_from_config(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.device == torch.device("cpu")
+
+    def test_torch_device_cpu_from_config(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device=torch.device("cpu"))
+        ctx = _build_model_context(mc)
+        assert ctx.device == torch.device("cpu")
+
+    def test_class_names_none_without_pretrain(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.class_names is None
+
+    def test_num_select_on_postprocess(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.postprocess.num_select == 100
+
+    def test_keypoint_preview_postprocess_has_keypoint_schema(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.postprocess.num_keypoints_per_class == [17]
+
+    def test_args_namespace_attached(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert hasattr(ctx.args, "num_classes")
+        assert hasattr(ctx.args, "num_select")
+
+    def test_inference_model_initially_none(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.inference_model is None
+
+    def test_args_dataset_dir_does_not_leak_cwd(self) -> None:
+        """The serialized namespace must not embed the caller's realpathed CWD as dataset_dir."""
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.args.dataset_dir is None
+
+    def test_args_output_dir_does_not_leak_cwd(self) -> None:
+        """The serialized namespace must keep a relative output_dir, not the caller's realpathed CWD."""
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.args.output_dir == "output"
+
+
+# ---------------------------------------------------------------------------
+# RFDETRModelModule.__init__ characterization
+# ---------------------------------------------------------------------------
+
+
+class TestRFDETRModelModuleInitCharacterization:
+    """Pin RFDETRModelModule.__init__() structural outputs.
+
+    The existing test_module_model.py tests the init via mocked build_model and build_namespace. These tests exercise
+    the REAL init path (no mocks) to characterize what a freshly built module looks like.
+    """
+
+    def _make_module(self, mc=None, tc=None):
+        from rfdetr.training.module_model import RFDETRModelModule
+
+        mc = mc or RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        tc = tc or TrainConfig(dataset_dir="/tmp")
+        return RFDETRModelModule(mc, tc)
+
+    def test_model_attribute_is_lwdetr(self) -> None:
+        module = self._make_module()
+        # model could be wrapped by torch.compile, so check the underlying type
+        underlying = getattr(module.model, "_orig_mod", module.model)
+        assert isinstance(underlying, LWDETR)
+
+    def test_criterion_is_set_criterion(self) -> None:
+        module = self._make_module()
+        assert isinstance(module.criterion, SetCriterion)
+
+    def test_postprocess_is_postprocess(self) -> None:
+        module = self._make_module()
+        assert isinstance(module.postprocess, PostProcess)
+
+    def test_strict_loading_false(self) -> None:
+        """strict_loading=False allows partial state-dict loading."""
+        module = self._make_module()
+        assert module.strict_loading is False
+
+    def test_configs_stored(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        tc = TrainConfig(dataset_dir="/tmp")
+        module = self._make_module(mc=mc, tc=tc)
+        assert module.model_config is mc
+        assert module.train_config is tc
+
+    def test_criterion_num_classes_matches_model(self) -> None:
+        """Criterion and model must agree on num_classes (both use +1 convention)."""
+        mc = RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        module = self._make_module(mc=mc)
+        underlying = getattr(module.model, "_orig_mod", module.model)
+        assert module.criterion.num_classes == underlying.class_embed.out_features
+
+    def test_postprocess_num_select_matches_config(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        module = self._make_module(mc=mc, tc=tc)
+        assert module.postprocess.num_select == mc.num_select
+
+    def test_segmentation_criterion_with_seg_config(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        module = self._make_module(mc=mc, tc=tc)
+        assert "masks" in module.criterion.losses

tests/models/test_builders.py

@@ -0,0 +1,162 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Characterization tests for config-native builder functions.
+
+These tests validate build_model_from_config() and build_criterion_from_config() which accept Pydantic config objects
+directly instead of requiring a pre-built SimpleNamespace. If these functions cannot be imported, all tests skip via the
+module-level pytestmark.
+"""
+
+import pytest
+
+from rfdetr.config import (
+    RFDETRBaseConfig,
+    RFDETRSegNanoConfig,
+    SegmentationTrainConfig,
+    TrainConfig,
+)
+
+try:
+    from rfdetr.models import build_criterion_from_config, build_model_from_config
+
+    HAS_CONFIG_BUILDERS = True
+except ImportError:
+    HAS_CONFIG_BUILDERS = False
+
+pytestmark = pytest.mark.skipif(
+    not HAS_CONFIG_BUILDERS,
+    reason="config-native builder functions are not importable",
+)
+
+
+class TestBuildModelFromConfig:
+    """Tests for build_model_from_config(model_config, defaults=MODEL_DEFAULTS)."""
+
+    def test_returns_lwdetr_for_base_config(self) -> None:
+        """build_model_from_config with RFDETRBaseConfig returns an LWDETR instance."""
+        from rfdetr.models.lwdetr import LWDETR
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        model = build_model_from_config(mc)
+        assert isinstance(model, LWDETR), f"Expected LWDETR instance, got {type(model).__name__}"
+
+    def test_num_classes_correct(self) -> None:
+        """num_classes=5 in config should produce class_embed with out_features=6.
+
+        build_model adds +1 to num_classes (background class convention).
+        """
+        mc = RFDETRBaseConfig(num_classes=5)
+        model = build_model_from_config(mc)
+        assert model.class_embed.out_features == 6, (
+            f"Expected class_embed.out_features=6 (num_classes+1), got {model.class_embed.out_features}"
+        )
+
+    def test_parity_with_build_model_via_namespace(self) -> None:
+        """Parameter count must match between config-native and namespace paths."""
+        from rfdetr._namespace import _namespace_from_configs
+        from rfdetr.models.lwdetr import build_model
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+
+        model_config_native = build_model_from_config(mc, tc)
+        ns = _namespace_from_configs(mc, tc)
+        model_namespace = build_model(ns)
+
+        params_native = sum(p.numel() for p in model_config_native.parameters())
+        params_namespace = sum(p.numel() for p in model_namespace.parameters())
+        assert params_native == params_namespace, (
+            f"Parameter count mismatch: config-native={params_native}, namespace={params_namespace}"
+        )
+
+    def test_segmentation_head_created_when_true(self) -> None:
+        """RFDETRSegNanoConfig has segmentation_head=True; model must have it."""
+        mc = RFDETRSegNanoConfig()
+        model = build_model_from_config(mc)
+        assert model.segmentation_head is not None, "Expected segmentation_head to be created for RFDETRSegNanoConfig"
+
+    def test_drop_path_uses_train_config_value(self) -> None:
+        """Non-default TrainConfig.drop_path must reach the model builder path."""
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp", drop_path=0.2)
+        model = build_model_from_config(mc, tc)
+
+        layers = model._get_backbone_encoder_layers()
+        assert layers is not None
+        assert hasattr(layers[-1], "drop_path")
+        assert layers[-1].drop_path.drop_prob == pytest.approx(0.2)
+
+    def test_rejects_encoder_only_defaults(self) -> None:
+        """The config-native builder guarantees an LWDETR return value."""
+        from dataclasses import replace
+
+        from rfdetr.models import MODEL_DEFAULTS
+
+        mc = RFDETRBaseConfig(num_classes=80)
+
+        with pytest.raises(ValueError, match="encoder_only=False"):
+            build_model_from_config(mc, defaults=replace(MODEL_DEFAULTS, encoder_only=True))
+
+    def test_rejects_backbone_only_defaults(self) -> None:
+        """backbone_only=True in defaults must also raise ValueError."""
+        from dataclasses import replace
+
+        from rfdetr.models import MODEL_DEFAULTS
+
+        mc = RFDETRBaseConfig(num_classes=80)
+
+        with pytest.raises(ValueError, match="backbone_only=False"):
+            build_model_from_config(mc, defaults=replace(MODEL_DEFAULTS, backbone_only=True))
+
+    def test_none_train_config_uses_dummy(self) -> None:
+        """build_model_from_config with train_config=None must not raise."""
+        mc = RFDETRBaseConfig(num_classes=80)
+        model = build_model_from_config(mc, train_config=None)
+        assert model is not None, "Expected a model, got None"
+
+
+class TestBuildCriterionFromConfig:
+    """Tests for build_criterion_from_config(model_config, train_config, defaults)."""
+
+    def test_returns_tuple(self) -> None:
+        """build_criterion_from_config must return a 2-tuple (SetCriterion, PostProcess)."""
+        from rfdetr.models.criterion import SetCriterion
+        from rfdetr.models.postprocess import PostProcess
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+        result = build_criterion_from_config(mc, tc)
+        assert isinstance(result, tuple), f"Expected tuple, got {type(result).__name__}"
+        assert len(result) == 2, f"Expected 2-tuple, got {len(result)}-tuple"
+        criterion, postprocess = result
+        assert isinstance(criterion, SetCriterion), f"Expected SetCriterion, got {type(criterion).__name__}"
+        assert isinstance(postprocess, PostProcess), f"Expected PostProcess, got {type(postprocess).__name__}"
+
+    def test_num_select_postprocess(self) -> None:
+        """RFDETRSegNanoConfig has num_select=100; PostProcess must reflect it."""
+        mc = RFDETRSegNanoConfig()
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        _, postprocess = build_criterion_from_config(mc, tc)
+        assert postprocess.num_select == 100, f"Expected PostProcess.num_select=100, got {postprocess.num_select}"
+
+    def test_segmentation_losses_included(self) -> None:
+        """With segmentation config, 'masks' must be in criterion.losses."""
+        mc = RFDETRSegNanoConfig()
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        criterion, _ = build_criterion_from_config(mc, tc)
+        assert "masks" in criterion.losses, f"Expected 'masks' in criterion.losses, got {criterion.losses}"
+
+    def test_custom_defaults_focal_alpha_applied(self) -> None:
+        """Custom focal_alpha in ModelDefaults must reach SetCriterion."""
+        from dataclasses import replace
+
+        from rfdetr.models import MODEL_DEFAULTS
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+        custom_defaults = replace(MODEL_DEFAULTS, focal_alpha=0.5)
+        criterion, _ = build_criterion_from_config(mc, tc, defaults=custom_defaults)
+        assert criterion.focal_alpha == pytest.approx(0.5), f"Expected focal_alpha=0.5, got {criterion.focal_alpha}"

tests/models/test_config.py

@@ -0,0 +1,925 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import os
+import warnings
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from pydantic import ValidationError
+
+from rfdetr.config import (
+    KeypointTrainConfig,
+    ModelConfig,
+    PretrainWeightsCompatibilityWarning,
+    RFDETRBaseConfig,
+    RFDETRLargeConfig,
+    RFDETRMediumConfig,
+    RFDETRNanoConfig,
+    RFDETRSeg2XLargeConfig,
+    RFDETRSegLargeConfig,
+    RFDETRSegMediumConfig,
+    RFDETRSegNanoConfig,
+    RFDETRSegSmallConfig,
+    RFDETRSegXLargeConfig,
+    RFDETRSmallConfig,
+    SegmentationTrainConfig,
+    TrainConfig,
+    _detect_device,
+)
+
+
+@pytest.fixture
+def sample_model_config() -> dict[str, object]:
+    return {
+        "encoder": "dinov2_windowed_small",
+        "out_feature_indexes": [1, 2, 3],
+        "dec_layers": 3,
+        "projector_scale": ["P3"],
+        "hidden_dim": 256,
+        "patch_size": 14,
+        "num_windows": 2,
+        "sa_nheads": 8,
+        "ca_nheads": 8,
+        "dec_n_points": 4,
+        "resolution": 384,
+        "positional_encoding_size": 256,
+    }
+
+
+class TestModelConfigValidation:
+    def test_rejects_unknown_fields(self, sample_model_config) -> None:
+        sample_model_config["unknown"] = "value"
+
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'unknown'"):
+            ModelConfig(**sample_model_config)
+
+    def test_rejects_unknown_attribute_assignment(self, sample_model_config) -> None:
+        config = ModelConfig(**sample_model_config)
+
+        with pytest.raises(ValueError, match=r"Unknown attribute: 'unknown'\."):
+            setattr(config, "unknown", "value")
+
+    def test_accepts_indexed_cuda_device_string(self, sample_model_config) -> None:
+        config = ModelConfig(**sample_model_config, device="cuda:1")
+        assert config.device == "cuda:1"
+
+    def test_accepts_torch_device(self, sample_model_config) -> None:
+        config = ModelConfig(**sample_model_config, device=torch.device("cuda:2"))
+        assert config.device == "cuda:2"
+
+    def test_rejects_non_string_non_torch_device_with_validation_error(self, sample_model_config) -> None:
+        with pytest.raises(ValidationError, match="device must be a string or torch\\.device\\."):
+            ModelConfig(**sample_model_config, device=123)
+
+    def test_rejects_invalid_device_string(self, sample_model_config) -> None:
+        with pytest.raises(ValidationError, match="Invalid device specifier: 'notadevice'\\."):
+            ModelConfig(**sample_model_config, device="notadevice")
+
+    @pytest.mark.parametrize(
+        "encoder",
+        [
+            pytest.param("dinov2_windowed_small", id="windowed_small"),
+            pytest.param("dinov2_windowed_base", id="windowed_base"),
+            pytest.param("dinov2_registers_windowed_small", id="registers_windowed_small"),
+        ],
+    )
+    def test_accepts_valid_encoder(self, sample_model_config, encoder: str) -> None:
+        """ModelConfig accepts every value in the EncoderName Literal."""
+        config = ModelConfig(**{**sample_model_config, "encoder": encoder})
+        assert config.encoder == encoder
+
+    def test_rejects_invalid_encoder(self, sample_model_config) -> None:
+        """ModelConfig raises ValidationError for encoder strings outside the Literal."""
+        with pytest.raises(ValidationError):
+            ModelConfig(**{**sample_model_config, "encoder": "dinov2_invalid_backbone"})
+
+    def test_rejects_negative_postprocess_trace_alpha(self, sample_model_config) -> None:
+        """ModelConfig rejects negative uncertainty score-fusion exponents."""
+        with pytest.raises(ValidationError):
+            ModelConfig(**sample_model_config, postprocess_trace_alpha=-0.1)
+
+    def test_postprocess_trace_alpha_defaults_to_keypoint_fusion_value(self, sample_model_config) -> None:
+        """ModelConfig defaults to the keypoint uncertainty score-fusion exponent."""
+        config = ModelConfig(**sample_model_config)
+
+        assert config.postprocess_trace_alpha == 0.2
+
+    def test_pretrain_weights_absolute_path_realpath_normalised(self, tmp_path) -> None:
+        """An absolute pathlib.Path for pretrain_weights is stored as the realpath-normalised string."""
+        weights_path = tmp_path / "weights.pth"
+
+        config = RFDETRBaseConfig(pretrain_weights=weights_path)
+
+        assert config.pretrain_weights == os.path.realpath(os.fspath(weights_path))
+
+    @pytest.mark.parametrize(
+        "field",
+        [
+            pytest.param("dataset_dir", id="dataset_dir"),
+            pytest.param("output_dir", id="output_dir"),
+        ],
+    )
+    def test_train_dir_fields_accept_path(self, tmp_path, field: str) -> None:
+        """TrainConfig dataset/output dir fields accept pathlib.Path and store the realpath-normalised string."""
+        path = tmp_path / "artifact"
+        kwargs = {"dataset_dir": str(tmp_path)}
+        kwargs[field] = path
+
+        config = TrainConfig(**kwargs)
+
+        assert getattr(config, field) == os.path.realpath(os.fspath(path))
+
+    def test_accepts_bare_path_object_for_pretrain_weights(self) -> None:
+        """Bare pretrained weight Path values resolve the same as bare strings."""
+        path_config = RFDETRBaseConfig(pretrain_weights=Path("rf-detr-base.pth"))
+        string_config = RFDETRBaseConfig(pretrain_weights="rf-detr-base.pth")
+
+        assert path_config.pretrain_weights == string_config.pretrain_weights
+
+    @pytest.mark.parametrize(
+        "value",
+        [
+            # PTL trainer.fit(ckpt_path=...) sentinels — must pass through verbatim.
+            pytest.param("best", id="sentinel_best"),
+            pytest.param("last", id="sentinel_last"),
+            pytest.param("hpc", id="sentinel_hpc"),
+            pytest.param("registry:model-name", id="sentinel_registry"),
+            # A relative Path proves os.fspath coercion without realpath resolution.
+            pytest.param(Path("checkpoints/last.ckpt"), id="path_object"),
+        ],
+    )
+    def test_resume_coerced_via_fspath_without_realpath(self, value) -> None:
+        """``resume`` accepts pathlib.Path and is coerced to ``str`` via ``os.fspath`` only.
+
+        Unlike ``dataset_dir``/``output_dir``, ``resume`` is forwarded verbatim to PyTorch Lightning's
+        ``trainer.fit(ckpt_path=...)``, which also accepts sentinels such as ``"last"``. Realpath-normalising it would
+        rewrite those sentinels (and relative paths) into spurious absolute paths, so the value must be preserved.
+        """
+        config = TrainConfig(dataset_dir="/tmp", resume=value)
+
+        assert config.resume == os.fspath(value)
+
+
+class TestRFDETRBaseConfigEncoder:
+    """Encoder field validation on RFDETRBaseConfig (no fixture needed — has defaults)."""
+
+    def test_accepts_registers_windowed_small(self) -> None:
+        """RFDETRBaseConfig accepts the new dinov2_registers_windowed_small encoder."""
+        config = RFDETRBaseConfig(encoder="dinov2_registers_windowed_small", pretrain_weights=None)
+        assert config.encoder == "dinov2_registers_windowed_small"
+
+    def test_rejects_invalid_encoder(self) -> None:
+        """RFDETRBaseConfig raises ValidationError for unknown encoder strings."""
+        with pytest.raises(ValidationError):
+            RFDETRBaseConfig(encoder="not_a_real_encoder", pretrain_weights=None)
+
+
+class TestSegmentationTrainConfigNumSelect:
+    """Unit tests for SegmentationTrainConfig.num_select default and per-model values."""
+
+    def test_defaults_to_none(self) -> None:
+        config = SegmentationTrainConfig(dataset_dir="/tmp")
+        assert config.num_select is None
+
+    def test_explicit_value_is_accepted(self) -> None:
+        # Explicitly setting num_select on SegmentationTrainConfig is deprecated (Item #3).
+        with pytest.warns(DeprecationWarning, match="TrainConfig.num_select is deprecated"):
+            config = SegmentationTrainConfig(dataset_dir="/tmp", num_select=42)
+        assert config.num_select == 42
+
+    @pytest.mark.parametrize(
+        "config_class, expected_num_select",
+        [
+            (RFDETRSegNanoConfig, 100),
+            (RFDETRSegSmallConfig, 100),
+            (RFDETRSegMediumConfig, 200),
+            (RFDETRSegLargeConfig, 200),
+            (RFDETRSegXLargeConfig, 300),
+            (RFDETRSeg2XLargeConfig, 300),
+        ],
+    )
+    def test_model_config_has_variant_specific_num_select(self, config_class, expected_num_select) -> None:
+        assert config_class().num_select == expected_num_select
+
+
+class TestTrainConfigRejectsUnknownKwargs:
+    """TrainConfig must raise on unknown/typo'd kwargs instead of silently ignoring them (extra='forbid')."""
+
+    def test_typo_kwarg_raises_with_helpful_message(self, tmp_path) -> None:
+        """A typo'd kwarg (epoch instead of epochs) raises listing the unknown and available parameters."""
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'epoch'"):
+            TrainConfig(dataset_dir=str(tmp_path), output_dir=str(tmp_path), epoch=5)
+
+    def test_typo_error_lists_available_parameters(self, tmp_path) -> None:
+        """The rejection message includes the available parameter list so the typo is easy to fix."""
+        with pytest.raises(ValidationError, match=r"Available parameter\(s\):.*epochs"):
+            TrainConfig(dataset_dir=str(tmp_path), output_dir=str(tmp_path), epoch=5)
+
+    @pytest.mark.parametrize(
+        "config_class",
+        [
+            pytest.param(SegmentationTrainConfig, id="segmentation"),
+            pytest.param(KeypointTrainConfig, id="keypoint"),
+        ],
+    )
+    def test_subclasses_reject_unknown_kwargs(self, tmp_path, config_class) -> None:
+        """TrainConfig subclasses inherit the forbid behaviour."""
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'epoch'"):
+            config_class(dataset_dir=str(tmp_path), output_dir=str(tmp_path), epoch=5)
+
+    def test_get_train_config_raises_for_typo_kwarg(self, tmp_path) -> None:
+        """The public RFDETR.get_train_config path surfaces the typo instead of swallowing it."""
+        from types import SimpleNamespace
+
+        from rfdetr.detr import RFDETR
+
+        stub = SimpleNamespace(_train_config_class=TrainConfig)
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'epoch'"):
+            RFDETR.get_train_config(stub, dataset_dir=str(tmp_path), output_dir=str(tmp_path), epoch=5)
+
+
+class TestTrainConfigT42PromotedFields:
+    """T4-2: Promoted fields exist with correct defaults; device field is absent."""
+
+    def _tc(self, tmp_path, **kwargs):
+        defaults = dict(dataset_dir=str(tmp_path), output_dir=str(tmp_path), tensorboard=False)
+        defaults.update(kwargs)
+        return TrainConfig(**defaults)
+
+    # --- device field removed ---
+
+    def test_device_not_in_model_fields(self):
+        """Device must not appear in TrainConfig.model_fields (PTL auto-detects accelerator)."""
+        assert "device" not in TrainConfig.model_fields
+
+    def test_device_kwarg_rejected(self, tmp_path):
+        """Passing device= directly to TrainConfig raises (extra='forbid'); RFDETR.train() pops it beforehand."""
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'device'"):
+            self._tc(tmp_path, device="cpu")
+
+    # --- promoted fields: defaults ---
+
+    def test_clip_max_norm_default(self, tmp_path):
+        """clip_max_norm defaults to 0.1."""
+        assert self._tc(tmp_path).clip_max_norm == pytest.approx(0.1)
+
+    def test_seed_default_is_none(self, tmp_path):
+        """Seed defaults to None (no seeding)."""
+        assert self._tc(tmp_path).seed is None
+
+    def test_sync_bn_default_is_false(self, tmp_path):
+        """sync_bn defaults to False."""
+        assert self._tc(tmp_path).sync_bn is False
+
+    def test_fp16_eval_default_is_false(self, tmp_path):
+        """fp16_eval defaults to False."""
+        assert self._tc(tmp_path).fp16_eval is False
+
+    def test_lr_scheduler_default_is_step(self, tmp_path):
+        """lr_scheduler defaults to 'step'."""
+        assert self._tc(tmp_path).lr_scheduler == "step"
+
+    def test_lr_min_factor_default(self, tmp_path):
+        """lr_min_factor defaults to 0.0."""
+        assert self._tc(tmp_path).lr_min_factor == pytest.approx(0.0)
+
+    def test_dont_save_weights_default_is_false(self, tmp_path):
+        """dont_save_weights defaults to False."""
+        assert self._tc(tmp_path).dont_save_weights is False
+
+    def test_run_test_default_is_false(self, tmp_path):
+        """run_test defaults to False to avoid extra full-dataset test passes."""
+        assert self._tc(tmp_path).run_test is False
+
+    def test_eval_interval_default_is_one(self, tmp_path):
+        """eval_interval defaults to 1 (evaluate each epoch)."""
+        assert self._tc(tmp_path).eval_interval == 1
+
+    def test_skip_best_epochs_default_is_zero(self, tmp_path):
+        """skip_best_epochs defaults to 0 for backward compatibility."""
+        assert self._tc(tmp_path).skip_best_epochs == 0
+
+    def test_ema_update_interval_default_is_one(self, tmp_path):
+        """ema_update_interval defaults to 1 (update every step)."""
+        assert self._tc(tmp_path).ema_update_interval == 1
+
+    def test_compute_val_loss_default_is_true(self, tmp_path):
+        """compute_val_loss defaults to True."""
+        assert self._tc(tmp_path).compute_val_loss is True
+
+    def test_compute_test_loss_default_is_true(self, tmp_path):
+        """compute_test_loss defaults to True."""
+        assert self._tc(tmp_path).compute_test_loss is True
+
+    # --- promoted fields: accept explicit values ---
+
+    @pytest.mark.parametrize(
+        "field, value",
+        [
+            pytest.param("clip_max_norm", 0.5, id="clip_max_norm"),
+            pytest.param("seed", 42, id="seed"),
+            pytest.param("sync_bn", True, id="sync_bn"),
+            pytest.param("fp16_eval", True, id="fp16_eval"),
+            pytest.param("lr_scheduler", "cosine", id="lr_scheduler_cosine"),
+            pytest.param("lr_min_factor", 0.01, id="lr_min_factor"),
+            pytest.param("dont_save_weights", True, id="dont_save_weights"),
+            pytest.param("run_test", True, id="run_test"),
+            pytest.param("eval_interval", 3, id="eval_interval"),
+            pytest.param("skip_best_epochs", 3, id="skip_best_epochs"),
+            pytest.param("ema_update_interval", 4, id="ema_update_interval"),
+            pytest.param("compute_val_loss", False, id="compute_val_loss"),
+            pytest.param("compute_test_loss", False, id="compute_test_loss"),
+            pytest.param("train_log_sync_dist", True, id="train_log_sync_dist"),
+            pytest.param("train_log_on_step", True, id="train_log_on_step"),
+            pytest.param("log_per_class_metrics", False, id="log_per_class_metrics"),
+            pytest.param("prefetch_factor", 4, id="prefetch_factor"),
+            pytest.param("pin_memory", False, id="pin_memory"),
+            pytest.param("persistent_workers", False, id="persistent_workers"),
+        ],
+    )
+    def test_promoted_field_accepts_explicit_value(self, tmp_path, field, value):
+        """Each promoted field accepts an explicit value."""
+        tc = self._tc(tmp_path, **{field: value})
+        assert getattr(tc, field) == value
+
+    def test_lr_scheduler_rejects_invalid_value(self, tmp_path):
+        """lr_scheduler must reject values other than 'step' and 'cosine'."""
+        with pytest.raises((ValueError, ValidationError)):
+            self._tc(tmp_path, lr_scheduler="cyclic")
+
+    @pytest.mark.parametrize(
+        ("field", "value"),
+        [
+            pytest.param("eval_interval", 0, id="eval_interval_zero"),
+            pytest.param("skip_best_epochs", -1, id="skip_best_epochs_negative"),
+            pytest.param("ema_update_interval", 0, id="ema_update_interval_zero"),
+            pytest.param("prefetch_factor", 0, id="prefetch_factor_zero"),
+        ],
+    )
+    def test_interval_and_prefetch_reject_non_positive_values(self, tmp_path, field, value):
+        """Eval/EMA intervals and prefetch_factor must be >= 1 when provided."""
+        with pytest.raises((ValueError, ValidationError)):
+            self._tc(tmp_path, **{field: value})
+
+    def test_batch_size_auto_is_accepted(self, tmp_path):
+        """batch_size accepts the special 'auto' value."""
+        tc = self._tc(tmp_path, batch_size="auto")
+        assert tc.batch_size == "auto"
+
+    @pytest.mark.parametrize(
+        "field,value",
+        [
+            ("batch_size", 0),
+            ("grad_accum_steps", 0),
+            ("auto_batch_target_effective", 0),
+            ("auto_batch_max_targets_per_image", 0),
+        ],
+    )
+    def test_auto_batch_related_fields_reject_non_positive_values(self, tmp_path, field, value):
+        """batch/accum/target-effective/max_targets fields must be >= 1 (except batch_size='auto')."""
+        with pytest.raises((ValueError, ValidationError)):
+            self._tc(tmp_path, **{field: value})
+
+    @pytest.mark.parametrize("ema_headroom", [0.0, 1.5])
+    def test_auto_batch_ema_headroom_must_be_in_open_one(self, tmp_path, ema_headroom):
+        """auto_batch_ema_headroom must be in (0, 1]."""
+        with pytest.raises((ValueError, ValidationError)):
+            self._tc(tmp_path, auto_batch_ema_headroom=ema_headroom)
+
+
+class TestBuildTrainerUsesRealFields:
+    """build_trainer() must read clip_max_norm, seed, sync_bn from real TrainConfig fields."""
+
+    def _tc(self, tmp_path, **kwargs):
+        defaults = dict(
+            dataset_dir=str(tmp_path),
+            output_dir=str(tmp_path),
+            tensorboard=False,
+            wandb=False,
+            mlflow=False,
+            clearml=False,
+            use_ema=False,
+        )
+        defaults.update(kwargs)
+        return TrainConfig(**defaults)
+
+    def _kp_tc(self, tmp_path, **kwargs):
+        defaults = dict(
+            dataset_dir=str(tmp_path),
+            output_dir=str(tmp_path),
+            tensorboard=False,
+            wandb=False,
+            mlflow=False,
+            clearml=False,
+            use_ema=False,
+        )
+        defaults.update(kwargs)
+        return KeypointTrainConfig(**defaults)
+
+    def _mc(self, **kwargs):
+        from rfdetr.config import RFDETRBaseConfig
+
+        defaults = dict(pretrain_weights=None, device="cpu", num_classes=3)
+        defaults.update(kwargs)
+        return RFDETRBaseConfig(**defaults)
+
+    def test_clip_max_norm_forwarded_to_trainer_for_detection(self, tmp_path):
+        """Detection models use Lightning's automatic optimization, so ``gradient_clip_val`` flows through to the
+        Trainer from ``TrainConfig.clip_max_norm`` unchanged."""
+        from rfdetr.training import build_trainer
+
+        trainer = build_trainer(self._tc(tmp_path, clip_max_norm=0.25), self._mc())
+        assert trainer.gradient_clip_val == pytest.approx(0.25)
+
+    def test_clip_max_norm_owned_by_model_module_for_keypoints(self, tmp_path):
+        """Keypoint models use manual optimization; trainer-owned clipping is disabled and ``clip_max_norm`` is applied
+        inside ``RFDETRModelModule._step_optimizer`` instead."""
+        from rfdetr.training import build_trainer
+
+        trainer = build_trainer(
+            self._kp_tc(tmp_path, clip_max_norm=0.25),
+            self._mc(use_grouppose_keypoints=True),
+        )
+        assert trainer.gradient_clip_val is None
+
+    def test_seed_not_applied_in_build_trainer_factory(self, tmp_path):
+        """Seeding is deferred to RFDETRModule.on_fit_start, not build_trainer()."""
+        import unittest.mock as mock
+
+        from rfdetr.training import build_trainer
+
+        with mock.patch("pytorch_lightning.seed_everything") as mock_seed:
+            build_trainer(self._tc(tmp_path, seed=99), self._mc())
+        mock_seed.assert_not_called()
+
+    def test_sync_bn_forwarded_to_trainer(self, tmp_path):
+        """sync_batchnorm=True is passed to Trainer when TrainConfig.sync_bn is True."""
+        import unittest.mock as mock
+
+        from rfdetr.training import build_trainer
+
+        captured_kwargs = {}
+
+        real_trainer_init = __import__("pytorch_lightning").Trainer.__init__
+
+        def _capture_init(self_t, **kwargs):
+            captured_kwargs.update(kwargs)
+            real_trainer_init(self_t, **kwargs)
+
+        with mock.patch("rfdetr.training.trainer.Trainer.__init__", _capture_init):
+            build_trainer(self._tc(tmp_path, sync_bn=True), self._mc())
+
+        assert captured_kwargs.get("sync_batchnorm") is True
+
+
+class TestDeprecatedTrainConfigFields:
+    """Item #3 Phase A: TrainConfig fields deprecated in favour of ModelConfig ownership."""
+
+    def _tc(self, **kwargs):
+        defaults = dict(dataset_dir="/tmp")
+        defaults.update(kwargs)
+        return TrainConfig(**defaults)
+
+    @pytest.mark.parametrize(
+        "field,value",
+        [
+            pytest.param("group_detr", 5, id="group_detr"),
+            pytest.param("ia_bce_loss", False, id="ia_bce_loss"),
+            pytest.param("segmentation_head", True, id="segmentation_head"),
+            pytest.param("num_select", 100, id="num_select"),
+        ],
+    )
+    def test_explicitly_set_deprecated_field_emits_warning(self, field, value) -> None:
+        """Setting a deprecated TrainConfig field explicitly must emit DeprecationWarning."""
+        with pytest.warns(DeprecationWarning, match=f"TrainConfig\\.{field} is deprecated"):
+            self._tc(**{field: value})
+
+    def test_default_group_detr_no_warning(self, recwarn) -> None:
+        """TrainConfig() without explicit group_detr must NOT warn."""
+        self._tc()
+        depr_warnings = [w for w in recwarn.list if issubclass(w.category, DeprecationWarning)]
+        assert not depr_warnings, f"Unexpected DeprecationWarning: {depr_warnings}"
+
+    def test_segmentation_train_config_no_warning_on_default_fields(self, recwarn) -> None:
+        """SegmentationTrainConfig() must NOT warn for its class-level defaults.
+
+        segmentation_head=True and num_select=None are SegmentationTrainConfig defaults, not explicitly set by the user
+        — they must not trigger DeprecationWarning.
+        """
+        SegmentationTrainConfig(dataset_dir="/tmp")
+        depr_warnings = [w for w in recwarn.list if issubclass(w.category, DeprecationWarning)]
+        assert not depr_warnings, f"Unexpected DeprecationWarning: {depr_warnings}"
+
+
+class TestDeprecatedModelConfigClsLossCoef:
+    """Item #3 Phase A: ModelConfig.cls_loss_coef deprecated in favour of TrainConfig ownership."""
+
+    def test_explicit_cls_loss_coef_emits_warning(self) -> None:
+        """Setting cls_loss_coef on ModelConfig explicitly must emit DeprecationWarning."""
+        sample = dict(
+            encoder="dinov2_windowed_small",
+            out_feature_indexes=[1, 2, 3],
+            dec_layers=3,
+            projector_scale=["P3"],
+            hidden_dim=256,
+            patch_size=14,
+            num_windows=2,
+            sa_nheads=8,
+            ca_nheads=8,
+            dec_n_points=4,
+            resolution=384,
+            positional_encoding_size=256,
+        )
+        with pytest.warns(DeprecationWarning, match="ModelConfig\\.cls_loss_coef is deprecated"):
+            ModelConfig(**sample, cls_loss_coef=2.0)
+
+    def test_default_cls_loss_coef_no_warning(self, recwarn) -> None:
+        """RFDETRBaseConfig() without explicit cls_loss_coef must NOT warn."""
+        RFDETRBaseConfig(pretrain_weights=None, device="cpu")
+        depr_warnings = [w for w in recwarn.list if issubclass(w.category, DeprecationWarning)]
+        assert not depr_warnings, f"Unexpected DeprecationWarning: {depr_warnings}"
+
+
+class TestSyncPEWithResolutionAtConstruction:
+    """Tests for the _sync_pe_with_resolution model_validator.
+
+    When a user provides a custom resolution at construction time (e.g., ``RFDETRLarge(resolution=640)``),
+    positional_encoding_size must be updated proportionally for configs where the default PE is formula-derived
+    (``default_pe == default_resolution // patch_size``).
+    """
+
+    @pytest.mark.parametrize(
+        "config_cls, new_resolution, expected_pe",
+        [
+            pytest.param(RFDETRLargeConfig, 640, 640 // 16, id="large_640"),
+            pytest.param(RFDETRLargeConfig, 576, 576 // 16, id="large_576"),
+            pytest.param(RFDETRSmallConfig, 640, 640 // 16, id="small_640"),
+            pytest.param(RFDETRMediumConfig, 640, 640 // 16, id="medium_640"),
+            pytest.param(RFDETRNanoConfig, 416, 416 // 16, id="nano_416"),
+            pytest.param(RFDETRSegNanoConfig, 360, 360 // 12, id="seg_nano_360"),
+            pytest.param(RFDETRSegSmallConfig, 480, 480 // 12, id="seg_small_480"),
+            pytest.param(RFDETRSegMediumConfig, 480, 480 // 12, id="seg_medium_480"),
+            pytest.param(RFDETRSegLargeConfig, 576, 576 // 12, id="seg_large_576"),
+            pytest.param(RFDETRSegXLargeConfig, 576, 576 // 12, id="seg_xlarge_576"),
+            pytest.param(RFDETRSeg2XLargeConfig, 720, 720 // 12, id="seg_2xlarge_720"),
+        ],
+    )
+    def test_positional_encoding_size_updated_for_formula_derived_configs(
+        self,
+        config_cls: type,
+        new_resolution: int,
+        expected_pe: int,
+    ) -> None:
+        """PE is auto-derived from the custom resolution for formula-derived model configs."""
+        cfg = config_cls(resolution=new_resolution, pretrain_weights=None)
+        assert cfg.positional_encoding_size == expected_pe
+
+    def test_explicit_positional_encoding_size_is_not_overridden(self) -> None:
+        """When positional_encoding_size is explicitly provided, the validator must not override it."""
+        cfg = RFDETRLargeConfig(resolution=640, positional_encoding_size=50, pretrain_weights=None)
+        assert cfg.positional_encoding_size == 50
+
+    def test_default_resolution_preserves_default_pe(self) -> None:
+        """Constructing with default resolution (no explicit resolution) must not change PE."""
+        cfg = RFDETRLargeConfig(pretrain_weights=None)
+        assert cfg.resolution == 704
+        assert cfg.positional_encoding_size == 44  # 704 // 16
+
+
+class TestDetectDevice:
+    """Tests for _detect_device() covering PyTorch accelerator detection paths."""
+
+    @patch("rfdetr.config.torch")
+    def test_falls_back_to_cuda_when_accelerator_module_absent(self, mock_torch: MagicMock) -> None:
+        """Returns 'cuda' via legacy fallback when torch.accelerator lacks current_accelerator (PyTorch < 2.4)."""
+        mock_torch.accelerator = MagicMock(spec=[])  # no current_accelerator → hasattr returns False → fallback
+        mock_torch.cuda.is_available.return_value = True
+        mock_torch.backends.mps.is_available.return_value = False
+        assert _detect_device() == "cuda"
+
+    @patch("rfdetr.config.torch")
+    def test_returns_cpu_when_current_accelerator_raises(self, mock_torch: MagicMock) -> None:
+        """Returns 'cpu' directly from the except handler when current_accelerator() raises RuntimeError."""
+        mock_torch.accelerator.current_accelerator.side_effect = RuntimeError("no device")
+        assert _detect_device() == "cpu"
+
+    @patch("rfdetr.config.torch")
+    def test_returns_cpu_when_no_gpu_available(self, mock_torch: MagicMock) -> None:
+        """Returns 'cpu' when accelerator is absent and neither CUDA nor MPS is available."""
+        mock_torch.accelerator = MagicMock(spec=[])  # no current_accelerator → fallback branch
+        mock_torch.cuda.is_available.return_value = False
+        mock_torch.backends.mps.is_available.return_value = False
+        assert _detect_device() == "cpu"
+
+    @patch("rfdetr.config.torch")
+    def test_returns_cpu_when_accelerator_compiled_in_but_unavailable(self, mock_torch: MagicMock) -> None:
+        """Returns 'cpu' when torch was compiled with CUDA but no driver is present at runtime.
+
+        Without ``check_available=True``, ``current_accelerator()`` reports the compile-time accelerator, so the default
+        CUDA wheel on a driverless machine yields ``device("cuda")`` and every model build crashes with "Found no NVIDIA
+        driver". The runtime availability check must win.
+        """
+
+        def fake_current_accelerator(check_available: bool = False) -> "torch.device | None":
+            return None if check_available else torch.device("cuda")
+
+        mock_torch.accelerator.current_accelerator = fake_current_accelerator
+        assert _detect_device() == "cpu"
+
+    @patch("rfdetr.config.torch")
+    def test_returns_accelerator_when_runtime_available(self, mock_torch: MagicMock) -> None:
+        """Returns the accelerator when it passes the runtime availability check."""
+
+        def fake_current_accelerator(check_available: bool = False) -> "torch.device | None":
+            return torch.device("cuda") if check_available else None
+
+        mock_torch.accelerator.current_accelerator = fake_current_accelerator
+        assert _detect_device() == "cuda"
+
+    @patch("rfdetr.config.torch")
+    def test_legacy_signature_unavailable_accelerator_returns_cpu(self, mock_torch: MagicMock) -> None:
+        """Falls back to ``is_available()`` when ``current_accelerator`` lacks ``check_available`` (PyTorch < 2.7)."""
+
+        def legacy_current_accelerator() -> "torch.device":
+            return torch.device("cuda")
+
+        mock_torch.accelerator.current_accelerator = legacy_current_accelerator
+        mock_torch.accelerator.is_available.return_value = False
+        assert _detect_device() == "cpu"
+
+    @patch("rfdetr.config.torch")
+    def test_legacy_signature_available_accelerator_is_kept(self, mock_torch: MagicMock) -> None:
+        """Keeps the accelerator on pre-``check_available`` builds when ``is_available()`` confirms it."""
+
+        def legacy_current_accelerator() -> "torch.device":
+            return torch.device("cuda")
+
+        mock_torch.accelerator.current_accelerator = legacy_current_accelerator
+        mock_torch.accelerator.is_available.return_value = True
+        assert _detect_device() == "cuda"
+
+    @patch("rfdetr.config.torch")
+    def test_legacy_signature_runtime_error_on_fallback_returns_cpu(self, mock_torch: MagicMock) -> None:
+        """Outer RuntimeError handler catches error from legacy fallback call.
+
+        Control-flow: ``current_accelerator(check_available=True)`` raises ``TypeError`` (inner except),
+        then ``current_accelerator()`` raises ``RuntimeError`` (outer except catches) → ``"cpu"``.
+        """
+        call_count = 0
+
+        def raises_on_fallback(**kwargs: object) -> "torch.device":
+            nonlocal call_count
+            call_count += 1
+            if "check_available" in kwargs:
+                raise TypeError("unexpected keyword argument 'check_available'")
+            raise RuntimeError("NVML error on legacy fallback")
+
+        mock_torch.accelerator.current_accelerator = raises_on_fallback
+        assert _detect_device() == "cpu"
+
+
+class TestPretrainWeightsCompatibilityWarning:
+    """Config-time warning for overrides that prevent pretrained weights from loading.
+
+    These tests instantiate the variant *config* directly (not the wrapper class) so they do not touch the network, the
+    cache, or any model construction.
+    """
+
+    def _capture(self, config_cls: type, **kwargs: object) -> list[warnings.WarningMessage]:
+        """Instantiate ``config_cls(**kwargs)`` and return only the pretrain-compat warnings."""
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            config_cls(**kwargs)
+            return [w for w in caught if issubclass(w.category, PretrainWeightsCompatibilityWarning)]
+
+    def test_default_construction_emits_no_warning(self) -> None:
+        """Default variant construction must not warn — defaults match the published checkpoint."""
+        assert self._capture(RFDETRNanoConfig) == []
+
+    def test_encoder_registers_override_warns(self) -> None:
+        """The dinov2-with-registers footgun: switching encoder away from the variant default."""
+        captured = self._capture(RFDETRNanoConfig, encoder="dinov2_registers_windowed_small")
+        assert len(captured) == 1
+        message = str(captured[0].message)
+        assert "encoder" in message
+        assert "dinov2_registers_windowed_small" in message
+        assert "dinov2_windowed_small" in message
+
+    @pytest.mark.parametrize(
+        "field, value",
+        [
+            pytest.param("hidden_dim", 384, id="hidden_dim"),
+            pytest.param("dec_layers", 6, id="dec_layers"),
+            pytest.param("num_windows", 4, id="num_windows"),
+            pytest.param("sa_nheads", 4, id="sa_nheads"),
+            pytest.param("ca_nheads", 8, id="ca_nheads"),
+            pytest.param("dec_n_points", 4, id="dec_n_points"),
+            pytest.param("out_feature_indexes", [2, 5, 8, 11], id="out_feature_indexes"),
+            pytest.param("projector_scale", ["P3", "P4"], id="projector_scale"),
+            pytest.param("bbox_reparam", False, id="bbox_reparam"),
+            pytest.param("lite_refpoint_refine", False, id="lite_refpoint_refine"),
+            pytest.param("layer_norm", False, id="layer_norm"),
+            pytest.param("two_stage", False, id="two_stage"),
+            pytest.param("num_channels", 1, id="num_channels"),
+        ],
+    )
+    def test_load_breaking_override_warns(self, field: str, value: object) -> None:
+        """Each load-breaking architecture override fires the warning."""
+        captured = self._capture(RFDETRNanoConfig, **{field: value})
+        assert len(captured) == 1
+        assert field in str(captured[0].message)
+
+    def test_mask_downsample_ratio_warns_on_seg_variant(self) -> None:
+        """``mask_downsample_ratio`` change is silently miscalibrating; must warn at config time."""
+        captured = self._capture(RFDETRSegNanoConfig, mask_downsample_ratio=2)
+        assert len(captured) == 1
+        assert "mask_downsample_ratio" in str(captured[0].message)
+
+    def test_patch_size_override_warns_defense_in_depth(self) -> None:
+        """patch_size already raises in load_pretrain_weights; the new warning is defense-in-depth.
+
+        We change patch_size to a value that differs from RFDETRNanoConfig's default (16).
+        """
+        captured = self._capture(RFDETRNanoConfig, patch_size=14)
+        assert len(captured) == 1
+        assert "patch_size" in str(captured[0].message)
+
+    def test_segmentation_head_override_warns(self) -> None:
+        """segmentation_head also raises at load time but warning fires first."""
+        # RFDETRNanoConfig has segmentation_head=False; flipping it to True is the override.
+        captured = self._capture(RFDETRNanoConfig, segmentation_head=True)
+        assert len(captured) == 1
+        assert "segmentation_head" in str(captured[0].message)
+
+    @pytest.mark.parametrize(
+        "field, value",
+        [
+            pytest.param("num_queries", 200, id="num_queries_decrease"),
+            pytest.param("num_queries", 300, id="num_queries_equal"),
+            pytest.param("group_detr", 8, id="group_detr_decrease"),
+            pytest.param("num_classes", 5, id="num_classes"),
+            pytest.param("resolution", 448, id="resolution"),
+            pytest.param("positional_encoding_size", 20, id="positional_encoding_size"),
+        ],
+    )
+    def test_silent_field_overrides(self, field: str, value: object) -> None:
+        """Fields that are auto-handled at load time must not emit a warning at config construction."""
+        assert self._capture(RFDETRNanoConfig, **{field: value}) == []
+
+    @pytest.mark.parametrize(
+        "field, value",
+        [
+            pytest.param("num_queries", 400, id="num_queries"),
+            pytest.param("group_detr", 20, id="group_detr"),
+        ],
+    )
+    def test_increase_field_warns(self, field: str, value: object) -> None:
+        """Increasing an integer field above the variant default warns — extra slots are randomly initialised."""
+        captured = self._capture(RFDETRNanoConfig, **{field: value})
+        assert len(captured) == 1
+        assert field in str(captured[0].message)
+
+    def test_pretrain_weights_none_warns(self) -> None:
+        """Explicitly opting out of pretrained weights warns about training from scratch."""
+        captured = self._capture(RFDETRNanoConfig, pretrain_weights=None)
+        assert len(captured) == 1
+        message = str(captured[0].message)
+        assert "from scratch" in message
+        assert "rf-detr-nano.pth" in message
+
+    def test_pretrain_weights_none_only_one_warning(self) -> None:
+        """When pretrain_weights=None, the architecture-overrides warning is suppressed.
+
+        The from-scratch warning is the dominant message; we don't pile on with arch warnings.
+        """
+        captured = self._capture(
+            RFDETRNanoConfig,
+            pretrain_weights=None,
+            encoder="dinov2_registers_windowed_small",
+            hidden_dim=384,
+        )
+        assert len(captured) == 1
+        assert "from scratch" in str(captured[0].message)
+
+    def test_custom_pretrain_weights_path_suppresses_arch_warning(self) -> None:
+        """Custom pretrain_weights path → defer to load-time detector — no config-time arch warning."""
+        captured = self._capture(
+            RFDETRNanoConfig,
+            pretrain_weights="/tmp/my_custom.pth",
+            encoder="dinov2_registers_windowed_small",
+        )
+        assert captured == []
+
+    def test_multiple_overrides_consolidated_into_one_warning(self) -> None:
+        """All overrides are listed in a single warning, not one warning per field."""
+        captured = self._capture(
+            RFDETRNanoConfig,
+            encoder="dinov2_registers_windowed_small",
+            hidden_dim=384,
+            num_queries=400,
+        )
+        assert len(captured) == 1
+        message = str(captured[0].message)
+        for needle in ("encoder", "hidden_dim", "num_queries"):
+            assert needle in message, f"expected {needle!r} in consolidated warning message"
+
+    def test_warning_is_user_warning_subclass(self) -> None:
+        """Confirms downstream filtering via UserWarning works."""
+        assert issubclass(PretrainWeightsCompatibilityWarning, UserWarning)
+
+    def test_modelconfig_with_required_fields_does_not_warn(self, sample_model_config: dict[str, object]) -> None:
+        """Constructing the abstract ModelConfig with required fields cannot compare to defaults — no warning."""
+        assert self._capture(ModelConfig, **sample_model_config) == []
+
+    def test_breaking_field_with_default_factory_skips_comparison(self) -> None:
+        """A subclass whose breaking field uses ``default_factory`` (so ``.default`` is ``PydanticUndefined``) must be
+        silently skipped — we have nothing to compare against."""
+        from pydantic import Field
+
+        class _DefaultFactoryConfig(RFDETRNanoConfig):
+            # Field uses default_factory → FieldInfo.default is PydanticUndefined,
+            # but is_required() is False.  Hits the `continue` on the
+            # PydanticUndefined check.
+            encoder: str = Field(default_factory=lambda: "dinov2_windowed_small")
+
+        assert self._capture(_DefaultFactoryConfig, encoder="dinov2_registers_windowed_small") == []
+
+    def test_increase_field_when_required_skips_comparison(self) -> None:
+        """A subclass where ``num_queries`` becomes required (no default) must be skipped."""
+
+        class _RequiredNumQueriesConfig(RFDETRNanoConfig):
+            num_queries: int  # type: ignore[misc]  # no default → required
+
+        assert self._capture(_RequiredNumQueriesConfig, num_queries=400) == []
+
+    def test_increase_field_with_non_int_default_skips_comparison(self) -> None:
+        """A subclass where ``num_queries`` has a non-int default must be skipped (can't ``>`` compare)."""
+        from typing import Any
+
+        class _NonIntDefaultConfig(RFDETRNanoConfig):
+            num_queries: Any = "300"  # type: ignore[assignment]  # non-int default
+
+        assert self._capture(_NonIntDefaultConfig, num_queries="400") == []
+
+    def test_explicit_variant_default_path_runs_arch_override_check(self) -> None:
+        """Passing the variant's own published-default path string must still check arch overrides.
+
+        Before the case-2 fix, any non-None explicit pretrain_weights bypassed the architecture-override check entirely
+        — including when the user passed the exact variant default string such as "rf-detr-nano.pth".
+        """
+        captured = self._capture(
+            RFDETRNanoConfig,
+            pretrain_weights="rf-detr-nano.pth",
+            encoder="dinov2_registers_windowed_small",
+        )
+        assert len(captured) == 1
+        assert "encoder" in str(captured[0].message)
+
+    def test_product_preserving_group_detr_increase_still_warns(self) -> None:
+        """Increasing group_detr while halving num_queries still warns — check is per-field, not product-aware.
+
+        This documents known current behaviour: the validator compares each field to its variant default independently,
+        not the combined query-slot product.  A product- preserving change (group_detr=26, num_queries=150 vs defaults
+        13, 300) warns for group_detr because 26 > 13, regardless of whether total slots are the same.
+        """
+        captured = self._capture(RFDETRNanoConfig, num_queries=150, group_detr=26)
+        assert len(captured) == 1
+        assert "group_detr" in str(captured[0].message)
+
+
+class TestBreakingListIntegrity:
+    """Guards against stale entries in the pretrain-compatibility breaking-field lists."""
+
+    def test_all_breaking_fields_exist_in_model_config(self) -> None:
+        """Every field guarded by the pretrain-compatibility check must exist in ModelConfig.model_fields.
+
+        Catches typos and fields renamed/removed without updating the breaking lists.
+        """
+        all_breaking = {
+            "encoder",
+            "hidden_dim",
+            "dec_layers",
+            "num_windows",
+            "sa_nheads",
+            "ca_nheads",
+            "dec_n_points",
+            "out_feature_indexes",
+            "projector_scale",
+            "bbox_reparam",
+            "lite_refpoint_refine",
+            "layer_norm",
+            "two_stage",
+            "patch_size",
+            "segmentation_head",
+            "num_channels",
+            "num_queries",
+            "group_detr",
+        }
+        stale = all_breaking - set(ModelConfig.model_fields.keys())
+        assert not stale, f"Fields in breaking lists not in ModelConfig.model_fields: {stale}"

tests/models/test_config_keypoints.py

@@ -0,0 +1,112 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for keypoint config defaults and namespace forwarding."""
+
+import pytest
+
+from rfdetr._namespace import _namespace_from_configs
+from rfdetr.config import (
+    KeypointTrainConfig,
+    RFDETRBaseConfig,
+    RFDETRKeypointPreviewConfig,
+    SegmentationTrainConfig,
+)
+
+
+def test_keypoint_config_defaults() -> None:
+    """Default model/train keypoint configuration values should match the preview contract."""
+    model = RFDETRKeypointPreviewConfig()
+    train = KeypointTrainConfig(dataset_dir="/tmp")
+
+    assert model.use_grouppose_keypoints is True
+    assert model.dual_projector is True
+    assert model.dual_projector_kp_only is True
+    assert model.num_keypoints_per_class == [17]
+    assert model.positional_encoding_size == 576 // 12
+
+    assert train.keypoint_l1_loss_coef == pytest.approx(1.0)
+    assert train.keypoint_findable_loss_coef == pytest.approx(1.0)
+    assert train.keypoint_visible_loss_coef == pytest.approx(1.0)
+    assert train.keypoint_nll_loss_coef == pytest.approx(1.0)
+    assert train.cls_loss_coef == pytest.approx(2.0)
+
+
+def test_keypoint_preview_config_person_schema() -> None:
+    """Person-keypoint preview config must expose a person-only schema."""
+    model = RFDETRKeypointPreviewConfig()
+
+    assert model.num_keypoints_per_class == [17]
+    assert sum(model.num_keypoints_per_class) == 17
+    assert model.out_feature_indexes == [3, 6, 9, 12]
+    assert model.num_windows == 2
+    assert model.dec_layers == 4
+    assert model.patch_size == 12
+    assert model.resolution == 576
+    assert model.pretrain_weights == "rf-detr-keypoint-preview-xlarge.pth"
+
+
+def test_keypoint_fields_propagate_to_namespace(tmp_path) -> None:
+    """All keypoint config fields are forwarded through _namespace_from_configs."""
+    model = RFDETRKeypointPreviewConfig()
+    train = KeypointTrainConfig(
+        dataset_dir=str(tmp_path),
+        keypoint_flip_pairs=[0, 1, 2, 3],
+        keypoint_l1_loss_coef=1.5,
+        keypoint_findable_loss_coef=2.5,
+        keypoint_visible_loss_coef=3.5,
+        keypoint_nll_loss_coef=4.5,
+    )
+
+    namespace = _namespace_from_configs(model, train)
+
+    assert namespace.use_grouppose_keypoints is True
+    assert namespace.keypoint_cross_attn is True
+    assert namespace.inter_instance_kp_attn is False
+    assert namespace.grouppose_keypoint_dim_downscale == 1
+    assert namespace.dual_projector is True
+    assert namespace.dual_projector_kp_only is True
+    assert namespace.num_keypoints_per_class == [17]
+    assert namespace.keypoint_flip_pairs == [0, 1, 2, 3]
+    assert namespace.keypoint_l1_loss_coef == pytest.approx(1.5)
+    assert namespace.keypoint_findable_loss_coef == pytest.approx(2.5)
+    assert namespace.keypoint_visible_loss_coef == pytest.approx(3.5)
+    assert namespace.keypoint_nll_loss_coef == pytest.approx(4.5)
+
+
+def test_keypoint_nll_loss_coef_default_restored_to_1_0() -> None:
+    """keypoint_nll_loss_coef must default to 1.0 after the 0.5 revert.
+
+    The 0.5 default was introduced to dampen OKS@75 oscillation.  It was later reverted to 1.0 to align with all other
+    keypoint loss terms (l1, findable, visible).  This test guards against silent regressions.
+    """
+    train = KeypointTrainConfig(dataset_dir="/tmp")
+    assert train.keypoint_nll_loss_coef == pytest.approx(1.0)
+
+
+def test_segmentation_train_config_cls_loss_coef_default() -> None:
+    """SegmentationTrainConfig.cls_loss_coef must default to 1.0, not the erroneous 5.0.
+
+    The 5.0 value was always present in SegmentationTrainConfig but was dead code pre-v1.7 (namespace builder read from
+    ModelConfig=1.0).  The v1.7 TrainConfig ownership migration silently activated it.  This test guards against re-
+    introducing that regression.
+    """
+    tc = SegmentationTrainConfig(dataset_dir="/tmp")
+    assert tc.cls_loss_coef == pytest.approx(1.0)
+
+
+def test_unknown_keypoint_fields_are_not_public_config_fields() -> None:
+    """Private keypoint implementation fields are not accepted as public model config."""
+    with pytest.raises(ValueError, match="Unknown parameter"):
+        RFDETRBaseConfig(num_classes=1, keypoint_private_hidden_dim=256)
+
+    # KeypointTrainConfig (a TrainConfig subclass) uses extra="forbid", so unknown
+    # kwargs raise with a helpful message rather than being silently dropped.
+    with pytest.raises(ValueError, match="Unknown parameter"):
+        KeypointTrainConfig(
+            dataset_dir="/tmp",
+            keypoint_private_hidden_dim=256,
+            keypoint_private_loss_coef=1.0,
+        )

tests/models/test_criterion.py

@@ -0,0 +1,134 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for SetCriterion edge paths: _output_device and num_boxes_for_targets."""
+
+import pytest
+import torch
+
+from rfdetr.models.criterion import SetCriterion
+
+
+class _MatcherStub:
+    """Minimal matcher that returns identity indices for every target in the batch."""
+
+    def __call__(self, outputs, targets, group_detr=1):
+        return [(torch.arange(len(t["labels"])), torch.arange(len(t["labels"]))) for t in targets]
+
+
+def _bare_criterion() -> SetCriterion:
+    """Return a SetCriterion with no losses so forward() is a no-op."""
+    criterion = SetCriterion.__new__(SetCriterion)
+    criterion.training = True
+    criterion.group_detr = 1
+    criterion.sum_group_losses = False
+    criterion.losses = []
+    criterion.weight_dict = {}
+    criterion.matcher = _MatcherStub()
+    criterion.num_keypoints_per_class = []
+    return criterion
+
+
+class TestOutputDevice:
+    """Tests for SetCriterion._output_device — probes top-level tensor values only."""
+
+    def test_returns_device_of_first_tensor(self):
+        """Device inferred from the first tensor value in outputs."""
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+
+        device = SetCriterion._output_device(outputs)
+
+        assert device == torch.device("cpu")
+
+    def test_raises_when_no_tensor_present(self):
+        """ValueError raised when no top-level value is a tensor."""
+        outputs = {"meta": "string_value", "count": 42}
+
+        with pytest.raises(ValueError, match="at least one tensor"):
+            SetCriterion._output_device(outputs)
+
+    def test_skips_non_tensor_values(self):
+        """Non-tensor entries at the top level are skipped; first tensor wins."""
+        outputs = {"meta": "ignored", "pred_logits": torch.zeros(1, 1, 1)}
+
+        device = SetCriterion._output_device(outputs)
+
+        assert device == torch.device("cpu")
+
+
+class TestNumBoxesForTargets:
+    """Tests for SetCriterion.num_boxes_for_targets — clamp and empty-target edge cases."""
+
+    def test_returns_tensor_gte_one(self):
+        """Result must be clamped to >= 1.0 to prevent division by zero."""
+        criterion = _bare_criterion()
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+        targets = [{"labels": torch.tensor([0, 1])}]
+
+        result = criterion.num_boxes_for_targets(outputs, targets)
+
+        assert result.item() >= 1.0
+
+    def test_clamps_zero_box_count_to_one(self):
+        """Empty targets (no labels) must clamp to 1.0 to avoid zero denominator."""
+        criterion = _bare_criterion()
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+        targets = [{"labels": torch.zeros(0, dtype=torch.int64)}]
+
+        result = criterion.num_boxes_for_targets(outputs, targets)
+
+        assert result.item() == pytest.approx(1.0)
+
+    def test_clamps_empty_target_list(self):
+        """Empty target list (batch_size=0 edge case) must also clamp to 1.0."""
+        criterion = _bare_criterion()
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+        targets = []
+
+        result = criterion.num_boxes_for_targets(outputs, targets)
+
+        assert result.item() == pytest.approx(1.0)
+
+    def test_counts_labels_correctly(self):
+        """Box count equals total number of labels across all targets in the batch."""
+        criterion = _bare_criterion()
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+        targets = [
+            {"labels": torch.tensor([0, 1])},
+            {"labels": torch.tensor([0])},
+        ]
+
+        result = criterion.num_boxes_for_targets(outputs, targets)
+
+        # 2 + 1 = 3 boxes; single-process so no all-reduce
+        assert result.item() == pytest.approx(3.0)
+
+
+class TestLossMasksEmptyMatch:
+    """Tests for the dict-path zero-GT branch of SetCriterion.loss_masks."""
+
+    def test_dict_path_zero_gt_stays_connected_to_graph(self):
+        """Zero-match dict path returns a loss that back-propagates to every segmentation-head output."""
+        criterion = _bare_criterion()
+        spatial_features = torch.randn(1, 4, 8, 8, requires_grad=True)
+        query_features = torch.randn(1, 5, 4, requires_grad=True)
+        bias = torch.randn(1, requires_grad=True)
+        outputs = {
+            "pred_masks": {
+                "spatial_features": spatial_features,
+                "query_features": query_features,
+                "bias": bias,
+            }
+        }
+        empty = torch.empty(0, dtype=torch.long)
+        indices = [(empty, empty)]
+
+        losses = criterion.loss_masks(outputs, targets=[{}], indices=indices, num_boxes=1)
+
+        assert losses["loss_mask_ce"].requires_grad
+        (losses["loss_mask_ce"] + losses["loss_mask_dice"]).backward()
+        assert spatial_features.grad is not None
+        assert query_features.grad is not None
+        assert bias.grad is not None

tests/models/test_criterion_keypoints.py

@@ -0,0 +1,161 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for keypoint losses in SetCriterion."""
+
+import torch
+
+from rfdetr.models.criterion import SetCriterion
+
+
+class _MatcherStub:
+    """Matcher stub used to avoid depending on Hungarian matching internals."""
+
+    def __call__(self, outputs, targets, group_detr=1):
+        indices = []
+        for target in targets:
+            num_targets = int(target["labels"].shape[0])
+            idx = torch.arange(num_targets, dtype=torch.int64)
+            indices.append((idx, idx))
+        return indices
+
+
+def _make_outputs(
+    batch_size: int,
+    num_queries: int,
+    num_keypoints: int,
+) -> dict[str, torch.Tensor]:
+    return {
+        "pred_logits": torch.zeros(batch_size, num_queries, 2),
+        "pred_boxes": torch.rand(batch_size, num_queries, 4).clamp(0.05, 0.95),
+        "pred_keypoints": torch.randn(batch_size, num_queries, num_keypoints, 8),
+    }
+
+
+def test_loss_keypoints_list_of_dicts_targets() -> None:
+    """Keypoint loss should consume list-of-dicts targets used by public training."""
+    criterion = SetCriterion(
+        num_classes=2,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[17],
+    )
+    outputs = _make_outputs(batch_size=1, num_queries=1, num_keypoints=17)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.4, 0.4]], dtype=torch.float32),
+            "keypoints": torch.cat(
+                [
+                    torch.rand(1, 17, 2),
+                    torch.full((1, 17, 1), 2.0),
+                ],
+                dim=-1,
+            ),
+        }
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert "loss_keypoints_l1" in losses
+    assert "loss_keypoints_findable" in losses
+    assert "loss_keypoints_visible" in losses
+    assert "loss_keypoints_nll" in losses
+    assert all(torch.isfinite(value) for value in losses.values())
+
+
+def test_loss_keypoints_empty_targets() -> None:
+    """Empty target batches should produce finite zero-valued keypoint losses."""
+    criterion = SetCriterion(
+        num_classes=2,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[17],
+    )
+    outputs = _make_outputs(batch_size=1, num_queries=1, num_keypoints=17)
+    targets = [
+        {
+            "labels": torch.zeros((0,), dtype=torch.int64),
+            "boxes": torch.zeros((0, 4), dtype=torch.float32),
+            "keypoints": torch.zeros((0, 17, 3), dtype=torch.float32),
+        }
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert losses["loss_keypoints_l1"].item() == 0.0
+    assert losses["loss_keypoints_findable"].item() == 0.0
+    assert losses["loss_keypoints_visible"].item() == 0.0
+    assert losses["loss_keypoints_nll"].item() == 0.0
+
+
+def test_loss_keypoints_person_schema_shape() -> None:
+    """Person-only schema `[17]` should be consumed without shape mismatches."""
+    criterion = SetCriterion(
+        num_classes=2,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[17],
+    )
+    outputs = _make_outputs(batch_size=2, num_queries=2, num_keypoints=17)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.4, 0.4]], dtype=torch.float32),
+            "keypoints": torch.rand(1, 17, 3),
+        },
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.4, 0.6, 0.3, 0.5]], dtype=torch.float32),
+            "keypoints": torch.rand(1, 17, 3),
+        },
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert losses["loss_keypoints_l1"].ndim == 0
+    assert losses["loss_keypoints_findable"].ndim == 0
+    assert losses["loss_keypoints_visible"].ndim == 0
+    assert losses["loss_keypoints_nll"].ndim == 0
+
+
+def test_loss_keypoints_multiclass_schema_kmax_targets() -> None:
+    """Heterogeneous keypoint classes should consume Kmax-padded targets."""
+    criterion = SetCriterion(
+        num_classes=3,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[2, 1],
+    )
+    outputs = _make_outputs(batch_size=1, num_queries=2, num_keypoints=4)
+    targets = [
+        {
+            "labels": torch.tensor([0, 1], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.4, 0.4], [0.4, 0.6, 0.3, 0.5]], dtype=torch.float32),
+            "keypoints": torch.tensor(
+                [
+                    [[0.2, 0.3, 2.0], [0.4, 0.5, 2.0]],
+                    [[0.6, 0.7, 2.0], [0.0, 0.0, 0.0]],
+                ],
+                dtype=torch.float32,
+            ),
+        }
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert losses["loss_keypoints_l1"].ndim == 0
+    assert losses["loss_keypoints_findable"].ndim == 0
+    assert losses["loss_keypoints_visible"].ndim == 0
+    assert losses["loss_keypoints_nll"].ndim == 0
+    assert all(torch.isfinite(value) for value in losses.values())

tests/models/test_lwdetr_backbone_contract.py

@@ -0,0 +1,60 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests proving LWDETR forward contract survives joiner return-shape changes."""
+
+from unittest.mock import MagicMock
+
+import torch
+
+from rfdetr.models.lwdetr import LWDETR
+from rfdetr.utilities.tensors import NestedTensor
+
+
+def test_lwdetr_default_detection_forward_after_backbone_change() -> None:
+    """LWDETR should accept the updated 3-tuple backbone output in non-keypoint mode."""
+    batch_size = 2
+    num_queries = 3
+    hidden_dim = 4
+    num_classes = 7
+
+    features = [
+        NestedTensor(
+            torch.zeros(batch_size, hidden_dim, 4, 4),
+            torch.zeros(batch_size, 4, 4, dtype=torch.bool),
+        )
+    ]
+    poss = [torch.zeros(batch_size, 4, 4, dtype=torch.bool)]
+
+    backbone = MagicMock()
+    backbone.return_value = (features, poss, None)
+
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer_out = (
+        torch.zeros(1, batch_size, num_queries, hidden_dim),
+        torch.zeros(1, batch_size, num_queries, hidden_dim),
+        torch.zeros(batch_size, num_queries, hidden_dim),
+        torch.zeros(batch_size, num_queries, hidden_dim),
+    )
+    transformer.return_value = transformer_out
+
+    model = LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=num_queries,
+        aux_loss=False,
+        group_detr=1,
+        two_stage=False,
+        lite_refpoint_refine=False,
+        bbox_reparam=False,
+    )
+
+    outputs = model(torch.ones(batch_size, 3, 8, 8))
+
+    assert outputs["pred_logits"].shape == (batch_size, num_queries, num_classes)
+    assert outputs["pred_boxes"].shape == (batch_size, num_queries, 4)

tests/models/test_lwdetr_keypoints.py

@@ -0,0 +1,226 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for GroupPose keypoint output wiring in LWDETR."""
+
+from unittest.mock import MagicMock
+
+import torch
+from torch import nn
+
+from rfdetr.models.heads import ConditionalQueryInitializer
+from rfdetr.models.lwdetr import LWDETR
+from rfdetr.utilities.tensors import NestedTensor
+
+
+def _build_feature_batch(batch_size: int, hidden_dim: int) -> list[NestedTensor]:
+    return [
+        NestedTensor(
+            torch.zeros(batch_size, hidden_dim, 4, 4),
+            torch.zeros(batch_size, 4, 4, dtype=torch.bool),
+        )
+    ]
+
+
+class _DummyKeypointDecoder(nn.Module):
+    """Minimal decoder surface needed for keypoint schema resizing."""
+
+    def __init__(self, hidden_dim: int, num_keypoints_per_class: list[int]) -> None:
+        super().__init__()
+        self.num_keypoints_per_class = num_keypoints_per_class
+        self.keypoint_pos_embed = nn.Parameter(torch.randn(sum(num_keypoints_per_class), hidden_dim))
+        self.register_buffer(
+            "keypoint_class_mask",
+            torch.zeros(1 + sum(num_keypoints_per_class), 1 + sum(num_keypoints_per_class), dtype=torch.bool),
+        )
+
+
+class _DummyKeypointTransformer(nn.Module):
+    """Minimal transformer surface needed for LWDETR construction and keypoint schema resizing."""
+
+    def __init__(self, hidden_dim: int, num_keypoints_per_class: list[int]) -> None:
+        super().__init__()
+        self.d_model = hidden_dim
+        self.num_keypoints_per_class = num_keypoints_per_class
+        self.decoder = _DummyKeypointDecoder(hidden_dim, num_keypoints_per_class)
+        self.keypoint_query_initializer = ConditionalQueryInitializer(hidden_dim, sum(num_keypoints_per_class))
+        self.keypoint_query_initializer_enc = ConditionalQueryInitializer(hidden_dim, sum(num_keypoints_per_class))
+
+
+def test_lwdetr_keypoint_forward_outputs() -> None:
+    """GroupPose mode should expose keypoint tensors in model outputs."""
+    batch_size = 2
+    num_queries = 3
+    hidden_dim = 8
+    num_classes = 6
+
+    features = _build_feature_batch(batch_size=batch_size, hidden_dim=hidden_dim)
+    poss = [torch.zeros(batch_size, hidden_dim, 4, 4)]
+
+    backbone = MagicMock()
+    backbone.return_value = (features, poss, None)
+
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer.return_value = (
+        torch.zeros(2, batch_size, num_queries, hidden_dim),  # hs
+        torch.zeros(2, batch_size, num_queries, 4),  # ref_unsigmoid
+        torch.zeros(batch_size, num_queries, hidden_dim),  # hs_enc
+        torch.zeros(batch_size, num_queries, 4),  # ref_enc
+        torch.zeros(2, batch_size, num_queries, 17, hidden_dim),  # keypoint_hs
+        torch.zeros(batch_size, num_queries, 17, 8),  # enc_kp_predictions
+        torch.zeros(batch_size, num_queries, 17, hidden_dim),  # unused keypoint encoder hidden state
+    )
+
+    model = LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=num_queries,
+        aux_loss=True,
+        group_detr=1,
+        two_stage=False,
+        lite_refpoint_refine=False,
+        bbox_reparam=False,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+        grouppose_keypoint_dim_downscale=1,
+    )
+
+    outputs = model(torch.ones(batch_size, 3, 8, 8))
+
+    assert outputs["pred_logits"].shape == (batch_size, num_queries, num_classes)
+    assert outputs["pred_boxes"].shape == (batch_size, num_queries, 4)
+    assert outputs["pred_keypoints"].shape == (batch_size, num_queries, 17, 8)
+    assert "keypoint_hidden_states" not in outputs
+    assert "pred_keypoints" in outputs["aux_outputs"][0]
+    assert "keypoint_hidden_states" not in outputs["aux_outputs"][0]
+
+
+def test_lwdetr_reinitialize_keypoint_head_updates_schema_dependent_state() -> None:
+    """Keypoint schema reinit should resize masks and learned keypoint query embeddings."""
+    hidden_dim = 8
+    transformer = _DummyKeypointTransformer(hidden_dim=hidden_dim, num_keypoints_per_class=[17])
+    model = LWDETR(
+        backbone=MagicMock(),
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=3,
+        num_queries=2,
+        aux_loss=False,
+        group_detr=1,
+        two_stage=True,
+        lite_refpoint_refine=True,
+        bbox_reparam=False,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+        grouppose_keypoint_dim_downscale=1,
+    )
+
+    model.reinitialize_keypoint_head([2, 1])
+
+    assert model.num_keypoints_per_class == [2, 1]
+    assert model.get_num_keypoints_per_class() == [2, 1]
+    assert model._kp_active_mask.shape == (2, 2)
+    assert model._kp_active_mask.tolist() == [[True, True], [True, False]]
+    assert transformer.num_keypoints_per_class == [2, 1]
+    assert transformer.decoder.num_keypoints_per_class == [2, 1]
+    assert transformer.decoder.keypoint_pos_embed.shape == (3, hidden_dim)
+    assert transformer.decoder.keypoint_class_mask.shape == (4, 4)
+    assert transformer.keypoint_query_initializer.queries.shape == (3, hidden_dim)
+    assert transformer.keypoint_query_initializer_enc.queries.shape == (3, hidden_dim)
+
+
+def test_lwdetr_reset_keypoint_gaussian_parameters_preserves_non_gaussian_rows() -> None:
+    """Gaussian reset should only zero precision-Cholesky output rows on decoder and encoder keypoint heads."""
+    hidden_dim = 8
+    transformer = _DummyKeypointTransformer(hidden_dim=hidden_dim, num_keypoints_per_class=[17])
+    model = LWDETR(
+        backbone=MagicMock(),
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=3,
+        num_queries=2,
+        aux_loss=False,
+        group_detr=1,
+        two_stage=True,
+        lite_refpoint_refine=True,
+        bbox_reparam=False,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+        grouppose_keypoint_dim_downscale=1,
+    )
+    with torch.no_grad():
+        model.keypoint_embed.layers[-1].weight.fill_(3.0)
+        model.keypoint_embed.layers[-1].bias.fill_(4.0)
+        model.transformer.enc_out_keypoint_embed[0].layers[-1].weight.fill_(5.0)
+        model.transformer.enc_out_keypoint_embed[0].layers[-1].bias.fill_(6.0)
+
+    model.reset_keypoint_gaussian_parameters()
+
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].weight[:4], torch.full((4, hidden_dim), 3.0))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].weight[4:7], torch.zeros(3, hidden_dim))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].weight[7:], torch.full((1, hidden_dim), 3.0))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].bias[:4], torch.full((4,), 4.0))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].bias[4:7], torch.zeros(3))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].bias[7:], torch.full((1,), 4.0))
+    torch.testing.assert_close(
+        model.transformer.enc_out_keypoint_embed[0].layers[-1].weight[4:7], torch.zeros(3, hidden_dim)
+    )
+    torch.testing.assert_close(model.transformer.enc_out_keypoint_embed[0].layers[-1].bias[4:7], torch.zeros(3))
+
+
+def test_lwdetr_get_num_keypoints_per_class_from_checkpoint() -> None:
+    """Checkpoint keypoint schema should be recoverable from `_kp_active_mask`."""
+    state_dict = {"_kp_active_mask": torch.tensor([[True, True], [True, False]])}
+
+    assert LWDETR.get_num_keypoints_per_class_from_checkpoint(state_dict) == [2, 1]
+
+
+def test_lwdetr_default_detection_contract_unchanged() -> None:
+    """Default detection mode should not expose keypoint outputs."""
+    batch_size = 2
+    num_queries = 3
+    hidden_dim = 8
+    num_classes = 6
+
+    features = _build_feature_batch(batch_size=batch_size, hidden_dim=hidden_dim)
+    poss = [torch.zeros(batch_size, hidden_dim, 4, 4)]
+
+    backbone = MagicMock()
+    backbone.return_value = (features, poss, None)
+
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer.return_value = (
+        torch.zeros(1, batch_size, num_queries, hidden_dim),
+        torch.zeros(1, batch_size, num_queries, 4),
+        torch.zeros(batch_size, num_queries, hidden_dim),
+        torch.zeros(batch_size, num_queries, 4),
+    )
+
+    model = LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=num_queries,
+        aux_loss=False,
+        group_detr=1,
+        two_stage=False,
+        lite_refpoint_refine=False,
+        bbox_reparam=False,
+        use_grouppose_keypoints=False,
+        num_keypoints_per_class=[],
+        grouppose_keypoint_dim_downscale=1,
+    )
+
+    outputs = model(torch.ones(batch_size, 3, 8, 8))
+
+    assert outputs["pred_logits"].shape == (batch_size, num_queries, num_classes)
+    assert outputs["pred_boxes"].shape == (batch_size, num_queries, 4)
+    assert "pred_keypoints" not in outputs
+    assert "keypoint_hidden_states" not in outputs

tests/models/test_matcher.py

@@ -0,0 +1,401 @@
+# ------------------------------------------------------------------------
+# 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 import matcher as matcher_module
+from rfdetr.models.matcher import HungarianMatcher
+
+
+@pytest.fixture()
+def matcher() -> HungarianMatcher:
+    """Shared HungarianMatcher instance."""
+    return HungarianMatcher()
+
+
+@pytest.fixture()
+def standard_target() -> dict[str, torch.Tensor]:
+    """Single-class target with one box at (0.5, 0.5, 0.2, 0.2)."""
+    return {
+        "labels": torch.tensor([0], dtype=torch.int64),
+        "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+    }
+
+
+class TestHungarianMatcherNonFiniteCosts:
+    """Tests for non-finite cost matrix sanitization in the Hungarian matcher."""
+
+    @pytest.mark.parametrize(
+        "invalid_value",
+        [
+            pytest.param(float("nan"), id="nan"),
+            pytest.param(float("inf"), id="inf"),
+            pytest.param(float("-inf"), id="-inf"),
+        ],
+    )
+    def test_replaces_non_finite_costs_before_assignment(
+        self,
+        matcher: HungarianMatcher,
+        standard_target: dict[str, torch.Tensor],
+        invalid_value: float,
+    ) -> None:
+        """Matcher should sanitize non-finite costs so assignment still succeeds."""
+        outputs = {
+            "pred_logits": torch.tensor([[[0.0], [10.0]]], dtype=torch.float32),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [invalid_value, 0.5, 0.2, 0.2],
+                        [0.5, 0.5, 0.2, 0.2],
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        matched_queries, matched_targets = matcher(outputs, [standard_target])[0]
+
+        assert matched_queries.tolist() == [1]
+        assert matched_targets.tolist() == [0]
+
+    def test_all_nonfinite_produces_valid_assignment(
+        self,
+        matcher: HungarianMatcher,
+        standard_target: dict[str, torch.Tensor],
+    ) -> None:
+        """When ALL costs are non-finite, the fallback sentinel (``dtype_info.max``)
+        should allow ``linear_sum_assignment`` to complete with a valid 1-to-1
+        assignment: exactly one match, query index in [0, num_queries), target index 0.
+
+        This exercises the ``else: replacement_cost = C.new_tensor(dtype_info.max)`` branch.
+        """
+        nan = float("nan")
+        outputs = {
+            "pred_logits": torch.tensor([[[nan], [nan]]], dtype=torch.float32),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [nan, nan, nan, nan],
+                        [nan, nan, nan, nan],
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        matched_queries, matched_targets = matcher(outputs, [standard_target])[0]
+
+        assert len(matched_queries) == len(matched_targets) == 1
+        assert 0 <= matched_queries.item() < 2
+        assert matched_targets.item() == 0
+
+    def test_negative_costs_with_nan_selects_valid_query(
+        self,
+        matcher: HungarianMatcher,
+        standard_target: dict[str, torch.Tensor],
+    ) -> None:
+        """Regression test: when all finite costs are negative and one query produces NaN, the matcher must select the
+        valid query, not the NaN one.
+
+        This guards against the bug where ``max_cost * 2`` (the old replacement formula) could be smaller than
+        ``max_cost`` when all costs are negative, causing the NaN query to appear cheaper than valid queries.
+        """
+        nan = float("nan")
+        # Query 0: NaN box coordinates -> produces non-finite costs
+        # Query 1: valid box, low logit -> all-negative but finite costs
+        outputs = {
+            "pred_logits": torch.tensor([[[0.0], [-10.0]]], dtype=torch.float32),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [nan, nan, nan, nan],
+                        [0.5, 0.5, 0.2, 0.2],
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        matched_queries, matched_targets = matcher(outputs, [standard_target])[0]
+
+        # The valid query (index 1) must be matched, not the NaN query.
+        assert matched_queries.tolist() == [1]
+        assert matched_targets.tolist() == [0]
+
+    @pytest.mark.parametrize(
+        "image_idx, expected_query_idx",
+        [
+            pytest.param(0, 1, id="image0"),
+            pytest.param(1, 0, id="image1"),
+        ],
+    )
+    def test_batch_size_greater_than_one(
+        self,
+        matcher: HungarianMatcher,
+        image_idx: int,
+        expected_query_idx: int,
+    ) -> None:
+        """Exercises the ``C.split(sizes, -1)`` loop with batch_size > 1.
+
+        Each image has 2 queries and 1 target. One query per image has NaN coordinates; the matcher must select the
+        valid query in each case.
+        """
+        nan = float("nan")
+        outputs = {
+            "pred_logits": torch.tensor(
+                [
+                    [[0.0], [10.0]],  # image 0: query 1 is valid
+                    [[10.0], [0.0]],  # image 1: query 0 is valid
+                ],
+                dtype=torch.float32,
+            ),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [nan, 0.5, 0.2, 0.2],  # image 0, query 0: NaN
+                        [0.5, 0.5, 0.2, 0.2],  # image 0, query 1: valid
+                    ],
+                    [
+                        [0.5, 0.5, 0.2, 0.2],  # image 1, query 0: valid
+                        [nan, 0.5, 0.2, 0.2],  # image 1, query 1: NaN
+                    ],
+                ],
+                dtype=torch.float32,
+            ),
+        }
+        targets = [
+            {
+                "labels": torch.tensor([0], dtype=torch.int64),
+                "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+            },
+            {
+                "labels": torch.tensor([0], dtype=torch.int64),
+                "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+            },
+        ]
+
+        results = matcher(outputs, targets)
+
+        assert len(results) == 2
+
+        matched_queries, matched_targets = results[image_idx]
+        assert matched_queries.tolist() == [expected_query_idx]
+        assert matched_targets.tolist() == [0]
+
+    def test_group_detr_with_nonfinite_costs(
+        self,
+        matcher: HungarianMatcher,
+        standard_target: dict[str, torch.Tensor],
+    ) -> None:
+        """Sanitization runs on the full cost matrix before splitting by group, so non-finite entries must be handled
+        correctly when ``group_detr > 1``.
+
+        4 queries, 2 groups of 2. Query 0 has a NaN box; query 2 (the best valid match in group 1) must be selected
+        across groups.
+        """
+        nan = float("nan")
+        outputs = {
+            "pred_logits": torch.tensor(
+                [[[0.0], [10.0], [0.0], [10.0]]],
+                dtype=torch.float32,
+            ),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [nan, nan, nan, nan],  # group 0, query 0: NaN
+                        [0.5, 0.5, 0.2, 0.2],  # group 0, query 1: valid
+                        [nan, nan, nan, nan],  # group 1, query 0: NaN
+                        [0.5, 0.5, 0.2, 0.2],  # group 1, query 1: valid
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        results = matcher(outputs, [standard_target], group_detr=2)
+
+        assert len(results) == 1
+        matched_queries, matched_targets = results[0]
+        # Each group contributes one match; both must map to target 0
+        assert matched_targets.tolist() == [0, 0]
+        # The valid query in each group (indices 1 and 3) must be selected
+        assert set(matched_queries.tolist()) == {1, 3}
+
+    def test_warns_once_per_matcher_instance(
+        self, standard_target: dict[str, torch.Tensor], monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Non-finite-cost warning should be emitted once per matcher instance."""
+        expected_warning = (
+            "Non-finite values detected in matcher cost matrix; "
+            "replacing with finite sentinel. "
+            "Check for numerical instability."
+        )
+        warning_messages: list[str] = []
+
+        def record_warning(msg: str, *args: object, **kwargs: object) -> None:
+            warning_messages.append(msg)
+
+        monkeypatch.setattr(matcher_module.logger, "warning", record_warning)
+
+        outputs = {
+            "pred_logits": torch.tensor([[[0.0], [10.0]]], dtype=torch.float32),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [float("nan"), 0.5, 0.2, 0.2],
+                        [0.5, 0.5, 0.2, 0.2],
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        first_matcher = HungarianMatcher()
+        second_matcher = HungarianMatcher()
+
+        first_matcher(outputs, [standard_target])
+        first_matcher(outputs, [standard_target])
+        second_matcher(outputs, [standard_target])
+
+        assert warning_messages == [expected_warning, expected_warning]
+
+
+class TestHungarianMatcherSanitization:
+    """Unit tests for the private matcher cost sanitization helper."""
+
+    def test_sanitize_cost_matrix_replaces_non_finite_entries(self) -> None:
+        """Non-finite entries should be replaced with a larger finite sentinel."""
+        cost_matrix = torch.tensor(
+            [
+                [1.0, float("nan")],
+                [float("inf"), -2.0],
+            ],
+            dtype=torch.float32,
+        )
+
+        sanitized = HungarianMatcher._sanitize_cost_matrix(cost_matrix)
+
+        assert torch.isfinite(sanitized).all()
+        assert sanitized[0, 1] == 4.0
+        assert sanitized[1, 0] == 4.0
+        assert sanitized[0, 0] == 1.0
+        assert sanitized[1, 1] == -2.0
+
+    def test_sanitize_cost_matrix_all_non_finite_fallback(self) -> None:
+        """All-non-finite matrices should fall back to the dtype maximum."""
+        cost_matrix = torch.tensor(
+            [
+                [float("nan"), float("inf")],
+                [float("-inf"), float("nan")],
+            ],
+            dtype=torch.float32,
+        )
+
+        sanitized = HungarianMatcher._sanitize_cost_matrix(cost_matrix)
+
+        assert torch.isfinite(sanitized).all()
+        assert torch.all(sanitized == torch.finfo(cost_matrix.dtype).max)
+
+    def test_sanitize_cost_matrix_clamps_overflowing_replacement_cost(self) -> None:
+        """Overflow in the computed replacement cost should clamp to dtype max."""
+        dtype_max = torch.finfo(torch.float32).max
+        cost_matrix = torch.tensor(
+            [
+                [dtype_max, float("nan")],
+                [0.0, 1.0],
+            ],
+            dtype=torch.float32,
+        )
+
+        sanitized = HungarianMatcher._sanitize_cost_matrix(cost_matrix)
+
+        assert torch.isfinite(sanitized).all()
+        assert sanitized[0, 1] == dtype_max
+
+
+class TestHungarianMatcherFocalAlpha:
+    """The configured ``focal_alpha`` must drive the classification matching cost."""
+
+    def test_focal_alpha_changes_assignment(self) -> None:
+        """Two matchers differing only in ``focal_alpha`` must be able to produce different assignments.
+
+        ``focal_alpha`` is accepted, documented as "used in the classification cost", and stored on the matcher, so it
+        must actually influence matching. This input is chosen so the optimal query->target pairing flips between
+        ``focal_alpha=0.25`` and ``focal_alpha=0.90``; if the cost ignores the configured alpha, both assignments
+        collapse to the same result.
+        """
+        outputs = {
+            "pred_logits": torch.tensor(
+                [[[2.3936, -1.4217], [2.3731, -2.1974]]],
+                dtype=torch.float32,
+            ),
+            "pred_boxes": torch.tensor(
+                [[[0.3898, 0.4340, 0.5331, 0.1901], [0.4256, 0.1002, 0.6955, 0.7815]]],
+                dtype=torch.float32,
+            ),
+        }
+        targets = [
+            {
+                "labels": torch.tensor([0, 1], dtype=torch.int64),
+                "boxes": torch.tensor(
+                    [[0.2111, 0.6630, 0.7569, 0.8855], [0.7750, 0.4393, 0.8838, 0.8792]],
+                    dtype=torch.float32,
+                ),
+            }
+        ]
+
+        def assignment(focal_alpha: float) -> list[int]:
+            matcher = HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0, focal_alpha=focal_alpha)
+            matched_queries, matched_targets = matcher(outputs, targets)[0]
+            # Queries ordered by the target index they are matched to.
+            return matched_queries[matched_targets.argsort()].tolist()
+
+        assert assignment(0.25) != assignment(0.90)
+        # Pin the exact expected mappings so a misapplied-alpha refactor is caught even when
+        # the two values remain different for unrelated reasons.
+        assert assignment(0.25) == [0, 1]
+        assert assignment(0.90) == [1, 0]
+
+    @pytest.mark.parametrize(
+        "focal_alpha, expected",
+        [
+            pytest.param(0.0, [0, 1], id="alpha_zero_pos_cost_zeroed"),
+            pytest.param(1.0, [1, 0], id="alpha_one_neg_cost_zeroed"),
+        ],
+    )
+    def test_focal_alpha_boundary_values_no_nan(self, focal_alpha: float, expected: list[int]) -> None:
+        """Degenerate focal_alpha values (0.0 and 1.0) must not produce NaN and must yield a valid assignment.
+
+        focal_alpha=0.0 zeroes ``pos_cost_class``; focal_alpha=1.0 zeroes ``neg_cost_class``. Neither path touches
+        ``log(prob)`` directly (formula uses logsigmoid of logits), so no division-by-zero or NaN can occur.
+        """
+        outputs = {
+            "pred_logits": torch.tensor(
+                [[[2.3936, -1.4217], [2.3731, -2.1974]]],
+                dtype=torch.float32,
+            ),
+            "pred_boxes": torch.tensor(
+                [[[0.3898, 0.4340, 0.5331, 0.1901], [0.4256, 0.1002, 0.6955, 0.7815]]],
+                dtype=torch.float32,
+            ),
+        }
+        targets = [
+            {
+                "labels": torch.tensor([0, 1], dtype=torch.int64),
+                "boxes": torch.tensor(
+                    [[0.2111, 0.6630, 0.7569, 0.8855], [0.7750, 0.4393, 0.8838, 0.8792]],
+                    dtype=torch.float32,
+                ),
+            }
+        ]
+
+        matcher = HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0, focal_alpha=focal_alpha)
+        matched_queries, matched_targets = matcher(outputs, targets)[0]
+
+        assert not matcher._warned_non_finite_costs, "boundary focal_alpha produced non-finite costs"
+        result = matched_queries[matched_targets.argsort()].tolist()
+        assert result == expected

tests/models/test_matcher_keypoints.py

@@ -0,0 +1,112 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for keypoint matching costs in HungarianMatcher."""
+
+import torch
+
+from rfdetr.models.matcher import HungarianMatcher
+
+
+def _base_outputs(num_queries: int = 2) -> dict[str, torch.Tensor]:
+    """Build minimal detection outputs used across matcher keypoint tests."""
+    pred_logits = torch.full((1, num_queries, 1), 5.0, dtype=torch.float32)
+    pred_boxes = torch.tensor([0.5, 0.5, 0.2, 0.2], dtype=torch.float32).view(1, 1, 4).repeat(1, num_queries, 1)
+    return {
+        "pred_logits": pred_logits,
+        "pred_boxes": pred_boxes,
+    }
+
+
+def test_matcher_keypoint_cost_list_of_dicts_targets() -> None:
+    """Keypoint matching costs should work with public list-of-dicts targets."""
+    matcher = HungarianMatcher(
+        cost_class=0.0,
+        cost_bbox=1.0,
+        cost_giou=0.0,
+        num_keypoints_per_class=[1],
+        keypoint_l1_loss_coef=10.0,
+        keypoint_findable_loss_coef=0.0,
+        keypoint_visible_loss_coef=0.0,
+        keypoint_nll_loss_coef=0.0,
+    )
+    outputs = _base_outputs()
+    outputs["pred_keypoints"] = torch.zeros((1, 2, 1, 8), dtype=torch.float32)
+    outputs["pred_keypoints"][0, 0, 0, :2] = torch.tensor([0.5, 0.5], dtype=torch.float32)
+    outputs["pred_keypoints"][0, 1, 0, :2] = torch.tensor([0.0, 0.0], dtype=torch.float32)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+            "keypoints": torch.tensor([[[0.5, 0.5, 2.0]]], dtype=torch.float32),
+        }
+    ]
+
+    matched_queries, matched_targets = matcher(outputs, targets)[0]
+
+    assert matched_queries.tolist() == [0]
+    assert matched_targets.tolist() == [0]
+
+
+def test_matcher_keypoint_cost_coefficients_off() -> None:
+    """Zero keypoint coefficients should preserve non-keypoint matching behavior."""
+    base_matcher = HungarianMatcher(cost_class=1.0, cost_bbox=1.0, cost_giou=1.0)
+    keypoint_matcher = HungarianMatcher(
+        cost_class=1.0,
+        cost_bbox=1.0,
+        cost_giou=1.0,
+        num_keypoints_per_class=[1],
+        keypoint_l1_loss_coef=0.0,
+        keypoint_findable_loss_coef=0.0,
+        keypoint_visible_loss_coef=0.0,
+        keypoint_nll_loss_coef=0.0,
+    )
+    outputs = _base_outputs()
+    outputs["pred_logits"][0, 0, 0] = 10.0
+    outputs["pred_logits"][0, 1, 0] = -10.0
+    outputs["pred_boxes"][0, 1, :] = torch.tensor([0.1, 0.1, 0.1, 0.1], dtype=torch.float32)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+            "keypoints": torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32),
+        }
+    ]
+    outputs_with_keypoints = dict(outputs)
+    outputs_with_keypoints["pred_keypoints"] = torch.zeros((1, 2, 1, 8), dtype=torch.float32)
+
+    base_indices = base_matcher(outputs, targets)[0]
+    keypoint_indices = keypoint_matcher(outputs_with_keypoints, targets)[0]
+
+    assert base_indices[0].tolist() == keypoint_indices[0].tolist()
+    assert base_indices[1].tolist() == keypoint_indices[1].tolist()
+
+
+def test_matcher_keypoint_empty_targets() -> None:
+    """Empty keypoint targets should return valid empty match results."""
+    matcher = HungarianMatcher(
+        cost_class=1.0,
+        cost_bbox=1.0,
+        cost_giou=1.0,
+        num_keypoints_per_class=[1],
+        keypoint_l1_loss_coef=1.0,
+        keypoint_findable_loss_coef=1.0,
+        keypoint_visible_loss_coef=1.0,
+        keypoint_nll_loss_coef=1.0,
+    )
+    outputs = _base_outputs(num_queries=3)
+    outputs["pred_keypoints"] = torch.zeros((1, 3, 1, 8), dtype=torch.float32)
+    targets = [
+        {
+            "labels": torch.zeros((0,), dtype=torch.int64),
+            "boxes": torch.zeros((0, 4), dtype=torch.float32),
+            "keypoints": torch.zeros((0, 1, 3), dtype=torch.float32),
+        }
+    ]
+
+    matched_queries, matched_targets = matcher(outputs, targets)[0]
+
+    assert matched_queries.numel() == 0
+    assert matched_targets.numel() == 0

tests/models/test_math.py

@@ -0,0 +1,102 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for rfdetr.models.math utility functions."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from rfdetr.models.math import accuracy, interpolate, inverse_sigmoid
+
+
+class TestInterpolate:
+    """Verify interpolate() delegates to F.interpolate across torchvision versions."""
+
+    def test_resizes_to_target_size(self) -> None:
+        """Interpolate() upsamples a 4-D tensor to the requested spatial size."""
+        x = torch.randn(2, 3, 4, 4)
+
+        out = interpolate(x, size=[8, 8], mode="bilinear", align_corners=False)
+
+        assert out.shape == (2, 3, 8, 8)
+
+    def test_handles_empty_batch(self) -> None:
+        """Interpolate() supports an empty batch dimension without error."""
+        x = torch.randn(0, 3, 4, 4)
+
+        out = interpolate(x, size=[8, 8], mode="nearest")
+
+        assert out.shape == (0, 3, 8, 8)
+
+
+class TestAccuracy:
+    """Verify accuracy() computes precision@k correctly."""
+
+    def test_top1_perfect_batch(self) -> None:
+        """All predictions correct returns top-1 accuracy of 100.0."""
+        output = torch.tensor([[0.0, 10.0], [10.0, 0.0], [0.0, 10.0]])
+        target = torch.tensor([1, 0, 1])
+        result = accuracy(output, target, topk=(1,))
+        assert len(result) == 1
+        assert result[0].item() == pytest.approx(100.0)
+
+    def test_top1_zero_accuracy(self) -> None:
+        """All predictions wrong returns top-1 accuracy of 0.0."""
+        output = torch.tensor([[10.0, 0.0], [0.0, 10.0]])
+        target = torch.tensor([1, 0])
+        result = accuracy(output, target, topk=(1,))
+        assert result[0].item() == pytest.approx(0.0)
+
+    def test_topk_returns_list_of_correct_length(self) -> None:
+        """Topk=(1, 5) returns a list of length 2."""
+        output = torch.randn(10, 10)
+        target = torch.zeros(10, dtype=torch.long)
+        result = accuracy(output, target, topk=(1, 5))
+        assert len(result) == 2
+
+    def test_empty_target_returns_single_zero_regardless_of_topk(self) -> None:
+        """Empty target returns list of length 1 with value 0 regardless of topk length."""
+        output = torch.zeros(0, 5)
+        target = torch.zeros(0, dtype=torch.long)
+        result = accuracy(output, target, topk=(1, 5))
+        assert len(result) == 1
+        assert result[0].item() == pytest.approx(0.0)
+
+
+class TestInverseSigmoid:
+    """Verify inverse_sigmoid() computes the logit function correctly."""
+
+    def test_identity_at_half(self) -> None:
+        """inverse_sigmoid(0.5) equals 0.0 since sigmoid(0.0) = 0.5."""
+        x = torch.tensor([0.5])
+        result = inverse_sigmoid(x)
+        assert result.item() == pytest.approx(0.0, abs=1e-5)
+
+    def test_clamping_at_zero_is_finite(self) -> None:
+        """inverse_sigmoid(0.0) is finite due to eps clamping."""
+        x = torch.tensor([0.0])
+        result = inverse_sigmoid(x)
+        assert torch.isfinite(result).all()
+
+    def test_clamping_at_one_is_finite(self) -> None:
+        """inverse_sigmoid(1.0) is finite due to eps clamping."""
+        x = torch.tensor([1.0])
+        result = inverse_sigmoid(x)
+        assert torch.isfinite(result).all()
+
+    def test_output_shape_matches_input(self) -> None:
+        """Output shape matches input shape for a multi-dimensional tensor."""
+        x = torch.rand(3, 4)
+        result = inverse_sigmoid(x)
+        assert result.shape == x.shape
+
+    def test_gradient_flows_for_non_saturated_input(self) -> None:
+        """Gradients are non-zero for a non-saturated input value."""
+        x = torch.tensor([0.3], requires_grad=True)
+        inverse_sigmoid(x).sum().backward()
+        assert x.grad is not None
+        assert x.grad.abs().item() > 0.0

tests/models/test_model.py

@@ -0,0 +1,54 @@
+# ------------------------------------------------------------------------
+# 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 import RFDETRBase, RFDETRLarge
+
+
+def _get_patch_embed_projection(model) -> torch.nn.Conv2d:
+    """Return the patch-embedding projection layer for an RF-DETR model.
+
+    RFDETR wrappers are not nn.Module; the underlying PyTorch module lives at ``model.model.model``.  Walk
+    named_modules() on that object.
+
+    Args:
+        model: Instantiated RF-DETR wrapper (RFDETRBase / RFDETRLarge).
+
+    Returns:
+        The convolution used to project image channels into patch embeddings.
+
+    Raises:
+        AssertionError: If the patch-embedding projection cannot be located.
+    """
+    # model.model → model context; model.model.model → nn.Module
+    nn_model = model.model.model
+    proj = nn_model.backbone[0].encoder.encoder.embeddings.patch_embeddings.projection
+    if isinstance(proj, torch.nn.Conv2d):
+        return proj
+
+    # Fallback: scan named_modules on the underlying nn.Module
+    for name, module in nn_model.named_modules():
+        if "patch_embeddings" in name and "projection" in name and isinstance(module, torch.nn.Conv2d):
+            return module
+
+    msg = "Could not find patch embedding projection on model"
+    raise AssertionError(msg)
+
+
+@pytest.mark.parametrize("model_class", [RFDETRBase, RFDETRLarge])
+@pytest.mark.parametrize("channels", [1, 4])
+def test_multispectral_support(model_class, channels: int) -> None:
+    model = model_class(
+        num_channels=channels,
+        device="cpu",
+        pretrain_weights=None,
+    )
+
+    patch_embed_projection = _get_patch_embed_projection(model)
+
+    assert patch_embed_projection.in_channels == channels

tests/models/test_namespace.py

@@ -0,0 +1,189 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for _namespace_from_configs() config forwarding."""
+
+import sys
+from typing import Any
+
+import pytest
+
+from rfdetr._namespace import _namespace_from_configs
+from rfdetr.config import RFDETRBaseConfig, RFDETRSegNanoConfig, SegmentationTrainConfig, TrainConfig
+from rfdetr.models._types import BuilderArgs
+
+
+class TestNamespaceForwarding:
+    """Verify that _namespace_from_configs() forwards TrainConfig fields that were previously hardcoded to wrong
+    defaults."""
+
+    def _make_ns(self: "TestNamespaceForwarding", **tc_kwargs: Any) -> Any:
+        """Build a namespace for tests with minimal default TrainConfig values."""
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc_kwargs.setdefault("dataset_dir", "/tmp")
+        tc = TrainConfig(**tc_kwargs)
+        return _namespace_from_configs(mc, tc)
+
+    def test_aug_config_forwarded_when_set(self: "TestNamespaceForwarding") -> None:
+        aug = {"hsv_h": 0.015, "hsv_s": 0.7}
+        ns = self._make_ns(aug_config=aug)
+        assert ns.aug_config == aug
+
+    def test_aug_config_none_by_default(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns()
+        assert ns.aug_config is None
+
+    def test_use_ema_forwarded_true(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns(use_ema=True)
+        assert ns.use_ema is True
+
+    def test_use_ema_forwarded_false(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns(use_ema=False)
+        assert ns.use_ema is False
+
+    def test_early_stopping_use_ema_forwarded_true(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns(early_stopping_use_ema=True)
+        assert ns.early_stopping_use_ema is True
+
+    def test_early_stopping_use_ema_forwarded_false(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns(early_stopping_use_ema=False)
+        assert ns.early_stopping_use_ema is False
+
+
+class TestNamespaceProtocol:
+    """_namespace_from_configs() output must satisfy the BuilderArgs Protocol."""
+
+    def _make_ns(self, mc=None, tc=None):
+        mc = mc or RFDETRBaseConfig(num_classes=80)
+        tc = tc or TrainConfig(dataset_dir="/tmp")
+        return _namespace_from_configs(mc, tc)
+
+    @pytest.mark.skipif(
+        sys.version_info < (3, 12),
+        reason="Runtime Protocol attribute checks require Python 3.12+",
+    )
+    def test_namespace_satisfies_builderargs_protocol_py312(self) -> None:
+        """On Python 3.12+, isinstance() verifies data-attribute presence."""
+        ns = self._make_ns()
+        assert isinstance(ns, BuilderArgs)
+
+    def test_namespace_is_builderargs_instance(self) -> None:
+        """Isinstance() check passes on all supported Python versions.
+
+        On Python 3.10/3.11 this is a structural no-op (no method members to check).  On 3.12+ it verifies attribute
+        presence.  The test documents the intent regardless of Python version.
+        """
+        ns = self._make_ns()
+        assert isinstance(ns, BuilderArgs)
+
+
+class TestNamespaceFieldOwnership:
+    """Verify that the namespace reads each field from the authoritative owner."""
+
+    def _make_ns(self, mc=None, tc=None):
+        mc = mc or RFDETRBaseConfig(num_classes=80)
+        tc = tc or TrainConfig(dataset_dir="/tmp")
+        return _namespace_from_configs(mc, tc)
+
+    # --- cls_loss_coef must come from TrainConfig ---
+
+    def test_cls_loss_coef_from_train_config(self) -> None:
+        """ns.cls_loss_coef must reflect TrainConfig.cls_loss_coef, not ModelConfig."""
+        mc = RFDETRBaseConfig(num_classes=80)  # cls_loss_coef=1.0 (ModelConfig default)
+        tc = TrainConfig(dataset_dir="/tmp", cls_loss_coef=2.5)
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(2.5)
+
+    def test_cls_loss_coef_segmentation_default_matches_pre_1_7_effective_value(self) -> None:
+        """SegmentationTrainConfig default must preserve the pre-1.7 effective loss_ce weight."""
+        mc = RFDETRSegNanoConfig()
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(1.0)
+
+    def test_cls_loss_coef_segmentation_explicit_train_config_value_wins(self) -> None:
+        """Explicit SegmentationTrainConfig.cls_loss_coef values must propagate to namespace."""
+        mc = RFDETRSegNanoConfig()
+        tc = SegmentationTrainConfig(dataset_dir="/tmp", cls_loss_coef=5.0)
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(5.0)
+
+    def test_cls_loss_coef_train_config_wins_over_explicit_model_config(self) -> None:
+        """When both are explicitly set, TrainConfig.cls_loss_coef takes precedence."""
+        with pytest.warns(DeprecationWarning, match="ModelConfig\\.cls_loss_coef is deprecated"):
+            mc = RFDETRBaseConfig(num_classes=80, cls_loss_coef=0.5)
+        tc = TrainConfig(dataset_dir="/tmp", cls_loss_coef=3.0)
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(3.0)
+
+    def test_cls_loss_coef_model_config_explicit_is_preserved_during_deprecation(self) -> None:
+        """Explicit ModelConfig.cls_loss_coef remains effective until removal."""
+        with pytest.warns(DeprecationWarning, match="ModelConfig\\.cls_loss_coef is deprecated"):
+            mc = RFDETRBaseConfig(num_classes=80, cls_loss_coef=2.5)
+        tc = TrainConfig(dataset_dir="/tmp")
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(2.5)
+
+    # --- num_select must come from ModelConfig unconditionally ---
+
+    def test_num_select_from_model_config(self) -> None:
+        """ns.num_select must equal mc.num_select regardless of tc.num_select."""
+        mc = RFDETRSegNanoConfig()  # num_select=100
+        tc = TrainConfig(dataset_dir="/tmp")  # num_select=300 (default — was the bug)
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.num_select == 100
+
+    @pytest.mark.parametrize(
+        "config_class, expected_num_select",
+        [
+            pytest.param(RFDETRSegNanoConfig, 100, id="seg_nano"),
+            pytest.param(RFDETRBaseConfig, 300, id="base"),
+        ],
+    )
+    def test_num_select_matches_model_config_variant(self, config_class, expected_num_select) -> None:
+        """ns.num_select must equal the model config's num_select for each variant."""
+        mc = config_class()
+        tc = TrainConfig(dataset_dir="/tmp")
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.num_select == expected_num_select
+
+
+class TestBuildNamespaceDeprecated:
+    """build_namespace() is a deprecated shim — verify the warning fires."""
+
+    def test_emits_deprecation_warning(self, reset_build_namespace_warning_state) -> None:
+        """Every call to build_namespace() must emit a DeprecationWarning."""
+        from rfdetr._namespace import build_namespace
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+
+        with pytest.warns(FutureWarning, match="build_namespace"):
+            build_namespace(mc, tc)
+
+    def test_result_identical_to_namespace_from_configs(self, reset_build_namespace_warning_state) -> None:
+        """build_namespace output must equal _namespace_from_configs output."""
+        from rfdetr._namespace import build_namespace
+        from rfdetr.models._defaults import MODEL_DEFAULTS
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+
+        with pytest.warns(FutureWarning):
+            ns_legacy = build_namespace(mc, tc)
+        ns_new = _namespace_from_configs(mc, tc, MODEL_DEFAULTS)
+
+        legacy_attrs = vars(ns_legacy)
+        new_attrs = vars(ns_new)
+
+        assert set(legacy_attrs.keys()) == set(new_attrs.keys()), (
+            f"Key mismatch: "
+            f"legacy_only={set(legacy_attrs) - set(new_attrs)}, "
+            f"new_only={set(new_attrs) - set(legacy_attrs)}"
+        )
+        for key in sorted(legacy_attrs):
+            assert legacy_attrs[key] == new_attrs[key], (
+                f"Value mismatch for '{key}': legacy={legacy_attrs[key]!r}, new={new_attrs[key]!r}"
+            )

tests/models/test_position_encoding.py

@@ -0,0 +1,88 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for rfdetr.models.position_encoding.build_position_encoding.
+
+Covers:
+- Supported aliases ``"sine"`` / ``"v2"`` return a ``PositionEmbeddingSine`` instance.
+- Unsupported but previously accepted aliases ``"learned"`` / ``"v3"`` now raise
+  ``ValueError`` with a message that names supported alternatives.
+- Fully unsupported values raise ``ValueError`` with the same pattern.
+"""
+
+import pytest
+
+from rfdetr.models.position_encoding import (
+    PositionEmbeddingSine,
+    build_position_encoding,
+)
+
+
+class TestBuildPositionEncodingSupportedValues:
+    """build_position_encoding returns valid modules for supported aliases."""
+
+    @pytest.mark.parametrize(
+        "alias",
+        [
+            pytest.param("sine", id="sine"),
+            pytest.param("v2", id="v2"),
+        ],
+    )
+    def test_returns_sine_embedding(self, alias: str) -> None:
+        """Supported aliases produce a PositionEmbeddingSine with normalized=True."""
+        enc = build_position_encoding(hidden_dim=256, position_embedding=alias)
+        assert isinstance(enc, PositionEmbeddingSine)
+        assert enc.normalize is True
+
+    @pytest.mark.parametrize(
+        "hidden_dim, expected_num_pos_feats",
+        [
+            pytest.param(256, 128, id="dim256"),
+            pytest.param(512, 256, id="dim512"),
+        ],
+    )
+    def test_num_pos_feats_is_half_hidden_dim(self, hidden_dim: int, expected_num_pos_feats: int) -> None:
+        """The sine encoding uses hidden_dim // 2 positional feature dimensions."""
+        enc = build_position_encoding(hidden_dim=hidden_dim, position_embedding="sine")
+        assert enc.num_pos_feats == expected_num_pos_feats
+
+
+class TestBuildPositionEncodingUnsupportedValues:
+    """build_position_encoding raises ValueError for broken or unknown aliases."""
+
+    @pytest.mark.parametrize(
+        "alias",
+        [
+            pytest.param("learned", id="learned"),
+            pytest.param("v3", id="v3"),
+        ],
+    )
+    def test_learned_raises_value_error(self, alias: str) -> None:
+        """'learned' and 'v3' are doubly broken and must raise ValueError immediately.
+
+        The PositionEmbeddingLearned class has two bugs:
+        1. forward() signature is incompatible with Joiner.forward() (no align_dim_orders param).
+        2. h, w = x.shape[:2] unpacks batch and channels instead of height and width.
+        Rejecting them at build time is preferable to a silent or confusing runtime failure.
+        """
+        with pytest.raises(ValueError, match="not supported"):
+            build_position_encoding(hidden_dim=256, position_embedding=alias)
+
+    def test_unknown_value_raises_value_error(self) -> None:
+        """A fully unknown alias raises ValueError naming the supported alternatives."""
+        with pytest.raises(ValueError, match="not supported"):
+            build_position_encoding(hidden_dim=256, position_embedding="unknown_variant")
+
+    @pytest.mark.parametrize(
+        "alias",
+        [
+            pytest.param("learned", id="learned"),
+            pytest.param("v3", id="v3"),
+        ],
+    )
+    def test_error_message_mentions_supported_alternatives(self, alias: str) -> None:
+        """Error message for 'learned'/'v3' mentions at least one supported alternative."""
+        with pytest.raises(ValueError, match="sine"):
+            build_position_encoding(hidden_dim=256, position_embedding=alias)

tests/models/test_postprocess.py

@@ -0,0 +1,170 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for PostProcess box clamping behaviour."""
+
+import pytest
+import torch
+
+from rfdetr.models.postprocess import PostProcess
+
+
+class TestGatherAndScaleBoxes:
+    """Tests for :meth:`PostProcess._gather_and_scale_boxes`."""
+
+    def test_clamps_boxes_to_image_bounds(self):
+        """Boxes that extrapolate beyond [0, 1] in normalized space are clamped to pixel-space image dimensions after
+        scaling."""
+        # Three synthetic boxes in cxcywh normalized coords:
+        #   [0] cx=0.01, w=0.10 → x1 = (0.01 - 0.05) * 640 = -25.6  ← negative
+        #   [1] cx=0.99, w=0.10 → x2 = (0.99 + 0.05) * 640 = 665.6  ← overflow
+        #   [2] cx=0.50, w=0.20 → fully in-bounds
+        out_bbox = torch.tensor(
+            [
+                [
+                    [0.01, 0.01, 0.10, 0.10],  # negative x1, y1 after scale
+                    [0.99, 0.99, 0.10, 0.10],  # x2 > img_w, y2 > img_h after scale
+                    [0.50, 0.50, 0.20, 0.20],  # in-bounds control
+                ]
+            ]
+        )  # shape (B=1, Q=3, 4)
+
+        topk_boxes = torch.tensor([[0, 1, 2]])  # select all three
+        target_sizes = torch.tensor([[480, 640]])  # (h, w)
+
+        boxes = PostProcess._gather_and_scale_boxes(out_bbox, topk_boxes, target_sizes)
+
+        img_h, img_w = 480, 640
+
+        # All coords must be >= 0
+        assert (boxes >= 0).all(), f"Negative coords present: {boxes[boxes < 0]}"
+        # x1, x2 must be <= image width
+        assert (boxes[..., 0] <= img_w).all()
+        assert (boxes[..., 2] <= img_w).all()
+        # y1, y2 must be <= image height
+        assert (boxes[..., 1] <= img_h).all()
+        assert (boxes[..., 3] <= img_h).all()
+
+        # Exact clamped values — bounds-only check cannot catch a clamp returning e.g. 1.0 instead of 0.0
+        # box [0]: x1_raw=-25.6, y1_raw=-19.2 → clamped to 0.0
+        assert boxes[0, 0, 0].item() == pytest.approx(0.0), "x1 of underflowing box must clamp to 0"
+        assert boxes[0, 0, 1].item() == pytest.approx(0.0), "y1 of underflowing box must clamp to 0"
+        # box [1]: x2_raw=665.6 → clamped to img_w=640.0; y2_raw=499.2 → clamped to img_h=480.0
+        assert boxes[0, 1, 2].item() == pytest.approx(640.0), "x2 of overflowing box must clamp to img_w"
+        assert boxes[0, 1, 3].item() == pytest.approx(480.0), "y2 of overflowing box must clamp to img_h"
+
+    def test_in_bounds_boxes_unchanged(self):
+        """Boxes already within image bounds are not altered by clamping."""
+        out_bbox = torch.tensor(
+            [
+                [
+                    [0.30, 0.30, 0.20, 0.20],
+                    [0.70, 0.60, 0.30, 0.40],
+                ]
+            ]
+        )
+
+        topk_boxes = torch.tensor([[0, 1]])
+        target_sizes = torch.tensor([[480, 640]])
+
+        boxes = PostProcess._gather_and_scale_boxes(out_bbox, topk_boxes, target_sizes)
+
+        # Manually computed expected values (no clamping needed)
+        expected = torch.tensor(
+            [
+                [
+                    [128.0, 96.0, 256.0, 192.0],  # cx=0.30,cy=0.30,w=0.20,h=0.20
+                    [352.0, 192.0, 544.0, 384.0],  # cx=0.70,cy=0.60,w=0.30,h=0.40
+                ]
+            ]
+        )
+
+        assert torch.allclose(boxes, expected, atol=1e-4), (
+            f"In-bounds boxes were altered.\nExpected:\n{expected}\nGot:\n{boxes}"
+        )
+
+    def test_multiple_images_in_batch(self):
+        """Clamping works correctly across a batch of mixed image sizes."""
+        out_bbox = torch.tensor(
+            [
+                [
+                    [0.01, 0.50, 0.10, 0.20],  # image 0: negative x1
+                ],
+                [
+                    [0.99, 0.50, 0.10, 0.20],  # image 1: x2 overflow
+                ],
+            ]
+        )
+
+        topk_boxes = torch.tensor([[0], [0]])
+        target_sizes = torch.tensor(
+            [
+                [300, 400],  # image 0: 400×300
+                [600, 800],  # image 1: 800×600
+            ]
+        )
+
+        boxes = PostProcess._gather_and_scale_boxes(out_bbox, topk_boxes, target_sizes)
+
+        # Image 0: all coords must be in [0, 400]×[0, 300]
+        assert (boxes[0, :, 0] >= 0).all(), "img0 x1: expected >= 0"
+        assert (boxes[0, :, 0] <= 400).all(), "img0 x1: expected <= img_w (400)"
+        assert (boxes[0, :, 1] >= 0).all(), "img0 y1: expected >= 0"
+        assert (boxes[0, :, 1] <= 300).all(), "img0 y1: expected <= img_h (300)"
+        assert (boxes[0, :, 2] >= 0).all(), "img0 x2: expected >= 0"
+        assert (boxes[0, :, 2] <= 400).all(), "img0 x2: expected <= img_w (400)"
+        assert (boxes[0, :, 3] >= 0).all(), "img0 y2: expected >= 0"
+        assert (boxes[0, :, 3] <= 300).all(), "img0 y2: expected <= img_h (300)"
+
+        # Image 1: all coords must be in [0, 800]×[0, 600]
+        assert (boxes[1, :, 0] >= 0).all(), "img1 x1: expected >= 0"
+        assert (boxes[1, :, 0] <= 800).all(), "img1 x1: expected <= img_w (800)"
+        assert (boxes[1, :, 1] >= 0).all(), "img1 y1: expected >= 0"
+        assert (boxes[1, :, 1] <= 600).all(), "img1 y1: expected <= img_h (600)"
+        assert (boxes[1, :, 2] >= 0).all(), "img1 x2: expected >= 0"
+        assert (boxes[1, :, 2] <= 800).all(), "img1 x2: expected <= img_w (800)"
+        assert (boxes[1, :, 3] >= 0).all(), "img1 y2: expected >= 0"
+        assert (boxes[1, :, 3] <= 600).all(), "img1 y2: expected <= img_h (600)"
+
+
+class TestPostProcessForward:
+    """Integration tests for :meth:`PostProcess.forward`."""
+
+    def test_forward_clamps_edge_boxes_to_bounds(self):
+        """PostProcess.forward returns non-negative in-bounds boxes for edge-hugging predictions."""
+        postprocess = PostProcess(num_select=2)
+        outputs = {
+            "pred_logits": torch.tensor([[[10.0, -10.0], [9.0, -10.0]]]),
+            "pred_boxes": torch.tensor([[[0.01, 0.01, 0.10, 0.10], [0.99, 0.99, 0.10, 0.10]]]),
+        }
+        target_sizes = torch.tensor([[480, 640]])
+        results = postprocess(outputs, target_sizes)
+        boxes = results[0]["boxes"]
+        assert (boxes >= 0).all(), f"Negative coords present: {boxes[boxes < 0]}"
+        assert (boxes[..., 0] <= 640).all(), "x1 exceeds img_w (640)"
+        assert (boxes[..., 2] <= 640).all(), "x2 exceeds img_w (640)"
+        assert (boxes[..., 1] <= 480).all(), "y1 exceeds img_h (480)"
+        assert (boxes[..., 3] <= 480).all(), "y2 exceeds img_h (480)"
+
+
+class TestPostProcessMasks:
+    """Tests for :meth:`PostProcess._postprocess_masks` mask resizing."""
+
+    def test_chunked_upsample_preserves_shape_for_large_k(self):
+        """Chunked upsampling of K=64 masks returns full-resolution boolean masks of shape [K, 1, H, W]."""
+        batch, num_queries, mask_h, mask_w = 1, 64, 16, 16
+        num_select, img_h, img_w = 64, 512, 512
+        out_masks = torch.randn(batch, num_queries, mask_h, mask_w)
+        scores = torch.rand(batch, num_select)
+        labels = torch.zeros(batch, num_select, dtype=torch.long)
+        boxes = torch.zeros(batch, num_select, 4)
+        topk_boxes = torch.arange(num_select).unsqueeze(0)
+        target_sizes = torch.tensor([[img_h, img_w]])
+
+        results = PostProcess._postprocess_masks(out_masks, scores, labels, boxes, topk_boxes, target_sizes)
+
+        masks = results[0]["masks"]
+        assert masks.shape == (num_select, 1, img_h, img_w)
+        assert masks.dtype == torch.bool

tests/models/test_postprocess_keypoints.py

@@ -0,0 +1,174 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for keypoint decoding in PostProcess."""
+
+import pytest
+import torch
+
+from rfdetr.models.postprocess import PostProcess
+
+
+def test_postprocess_keypoints_shape_and_scores() -> None:
+    """PostProcess should emit keypoints and raw precision parameters for top detections."""
+    postprocess = PostProcess(num_select=2, num_keypoints_per_class=[17])
+    outputs = {
+        "pred_logits": torch.tensor([[[10.0, -10.0], [9.0, -10.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5], [0.4, 0.6, 0.2, 0.3]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 2, 17, 8), dtype=torch.float32),
+    }
+    outputs["pred_keypoints"][0, :, :, 0] = 0.5
+    outputs["pred_keypoints"][0, :, :, 1] = 0.25
+    outputs["pred_keypoints"][0, :, :, 2] = 3.0
+    outputs["pred_keypoints"][0, :, :, 4] = 0.25
+    outputs["pred_keypoints"][0, :, :, 5] = 0.5
+    outputs["pred_keypoints"][0, :, :, 6] = -0.25
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+    keypoints = results[0]["keypoints"]
+    keypoint_precision = results[0]["keypoint_precision_cholesky"]
+
+    assert keypoints.shape == (2, 17, 3)
+    assert torch.allclose(keypoints[:, :, 0], torch.full((2, 17), 100.0))
+    assert torch.allclose(keypoints[:, :, 1], torch.full((2, 17), 25.0))
+    assert torch.all((keypoints[:, :, 2] > 0) & (keypoints[:, :, 2] < 1))
+    assert keypoint_precision.shape == (2, 17, 3)
+    torch.testing.assert_close(keypoint_precision[:, :, 0], torch.full((2, 17), 0.25))
+    torch.testing.assert_close(keypoint_precision[:, :, 1], torch.full((2, 17), 0.5))
+    torch.testing.assert_close(keypoint_precision[:, :, 2], torch.full((2, 17), -0.25))
+
+
+def test_postprocess_keypoints_class_filtering() -> None:
+    """Class-specific keypoint slots should be selected from padded per-class keypoint tensors."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[2, 1])
+    outputs = {
+        "pred_logits": torch.tensor([[[0.0, 10.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 4, 8), dtype=torch.float32),
+    }
+    # class 0 slots: [0, 1], class 1 slots: [2, 3]
+    outputs["pred_keypoints"][0, 0, 2, 0] = 0.25
+    outputs["pred_keypoints"][0, 0, 2, 1] = 0.4
+    outputs["pred_keypoints"][0, 0, 2, 2] = 2.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+    keypoints = results[0]["keypoints"]
+    keypoint_precision = results[0]["keypoint_precision_cholesky"]
+
+    assert keypoints.shape == (1, 2, 3)
+    assert torch.allclose(keypoints[0, 0, 0], torch.tensor(50.0))
+    assert torch.allclose(keypoints[0, 0, 1], torch.tensor(40.0))
+    assert 0.0 < keypoints[0, 0, 2].item() < 1.0
+    torch.testing.assert_close(keypoints[0, 1], torch.zeros(3))
+    torch.testing.assert_close(keypoint_precision[0, 1], torch.full((3,), float("nan")), equal_nan=True)
+
+
+def test_postprocess_keypoints_trace_alpha_rescores_active_keypoints_only() -> None:
+    """Trace fusion should use active keypoints for the predicted class and ignore padded slots."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[2, 1], trace_alpha=1.0)
+    outputs = {
+        "pred_logits": torch.tensor([[[-10.0, 0.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 4, 8), dtype=torch.float32),
+    }
+    # class 1 has one active slot at flat index 2 and one padded inactive slot at flat index 3.
+    outputs["pred_keypoints"][0, 0, 2, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 2, 4] = 0.0
+    outputs["pred_keypoints"][0, 0, 2, 5] = 0.0
+    outputs["pred_keypoints"][0, 0, 2, 6] = 0.0
+    outputs["pred_keypoints"][0, 0, 3, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 3, 4] = -2.0
+    outputs["pred_keypoints"][0, 0, 3, 6] = -2.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+
+    expected_score = torch.tensor([0.2], dtype=torch.float32)
+    torch.testing.assert_close(results[0]["scores"], expected_score, rtol=1e-4, atol=1e-6)
+
+
+def test_postprocess_keypoints_trace_alpha_normalizes_large_fused_scores() -> None:
+    """Trace-fused keypoint scores should be bounded after empirical normalization."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[1], trace_alpha=1.0)
+    outputs = {
+        "pred_logits": torch.tensor([[[10.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 1, 8), dtype=torch.float32),
+    }
+    outputs["pred_keypoints"][0, 0, 0, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 0, 4] = 2.0
+    outputs["pred_keypoints"][0, 0, 0, 5] = 0.0
+    outputs["pred_keypoints"][0, 0, 0, 6] = 2.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+
+    original_score = torch.sigmoid(torch.tensor([10.0], dtype=torch.float32))
+    mean_trace = torch.tensor([2.0], dtype=torch.float32) * torch.exp(torch.tensor([-4.0], dtype=torch.float32))
+    fused_score = original_score * mean_trace.pow(-1.0)
+    expected_score = fused_score / (1.0 + fused_score)
+    assert fused_score.item() > 1.0
+    assert 0.0 < results[0]["scores"].item() < 1.0
+    torch.testing.assert_close(results[0]["scores"], expected_score, rtol=1e-4, atol=1e-6)
+
+
+def test_postprocess_keypoints_trace_alpha_clamps_overflowing_fused_scores() -> None:
+    """Trace fusion should stay finite and strictly below 1.0 when the raw fused score overflows."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[1], trace_alpha=1.0)
+    outputs = {
+        "pred_logits": torch.tensor([[[0.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 1, 8), dtype=torch.float32),
+    }
+    outputs["pred_keypoints"][0, 0, 0, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 0, 4] = 50.0
+    outputs["pred_keypoints"][0, 0, 0, 5] = 0.0
+    outputs["pred_keypoints"][0, 0, 0, 6] = 50.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+
+    score = results[0]["scores"]
+    expected_score = torch.nextafter(torch.ones_like(score), torch.zeros_like(score))
+    assert torch.isfinite(score).all()
+    assert 0.0 < score.item() < 1.0
+    torch.testing.assert_close(score, expected_score, rtol=0.0, atol=0.0)
+
+
+def test_postprocess_keypoints_trace_alpha_uses_log_space_for_extreme_trace() -> None:
+    """Trace fusion should stay finite for extreme covariance terms."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[1])
+    outputs = {
+        "pred_logits": torch.tensor([[[0.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 1, 8), dtype=torch.float32),
+    }
+    outputs["pred_keypoints"][0, 0, 0, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 0, 4] = -50.0
+    outputs["pred_keypoints"][0, 0, 0, 5] = 0.0
+    outputs["pred_keypoints"][0, 0, 0, 6] = 0.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+
+    expected_score = torch.tensor([0.5], dtype=torch.float32) * torch.exp(torch.tensor([-20.0], dtype=torch.float32))
+    torch.testing.assert_close(results[0]["scores"], expected_score, rtol=1e-4, atol=1e-12)
+
+
+def test_postprocess_validate_outputs_raises_when_masks_and_keypoints_both_present() -> None:
+    """PostProcess should raise ValueError when both pred_masks and pred_keypoints are present."""
+    postprocess = PostProcess(num_select=10)
+    outputs = {
+        "pred_logits": torch.zeros((1, 2, 2)),
+        "pred_boxes": torch.zeros((1, 2, 4)),
+        "pred_masks": torch.zeros((1, 2, 4, 4)),
+        "pred_keypoints": torch.zeros((1, 2, 17, 8)),
+    }
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+
+    with pytest.raises(ValueError, match="cannot be used together"):
+        postprocess(outputs, target_sizes)

tests/models/test_safe_torch_load.py

@@ -0,0 +1,175 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for the ``_safe_torch_load`` helper in ``rfdetr.util.io``.
+
+Covers the three-stage safe-load strategy: strict weights_only, safe-globals fallback, and opt-in pickle fallback.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from rfdetr.util.io import _safe_torch_load
+
+# ---------------------------------------------------------------------------
+# Fixtures / helpers
+# ---------------------------------------------------------------------------
+
+
+def _write_tensor_only_checkpoint(path: Path) -> None:
+    """Save a checkpoint containing only tensors and plain dicts to *path*."""
+    ckpt = {"model": {"weight": torch.tensor([1.0, 2.0]), "bias": torch.tensor([0.0])}}
+    torch.save(ckpt, path)
+
+
+def _write_namespace_checkpoint(path: Path) -> None:
+    """Save a checkpoint with an ``argparse.Namespace`` args value to *path*.
+
+    Legacy RF-DETR engine.py checkpoints embed a Namespace; strict ``weights_only=True`` (without safe globals) would
+    reject these.
+    """
+    ckpt = {
+        "model": {"weight": torch.tensor([1.0])},
+        "args": argparse.Namespace(pretrain_weights="rf-detr-small.pth", num_classes=80),
+    }
+    torch.save(ckpt, path)
+
+
+def _write_simple_namespace_checkpoint(path: Path) -> None:
+    """Save a checkpoint with a ``types.SimpleNamespace`` to *path*."""
+    ckpt = {
+        "model": {"weight": torch.tensor([1.0])},
+        "args": SimpleNamespace(pretrain_weights="rf-detr-small.pth"),
+    }
+    torch.save(ckpt, path)
+
+
+class _ArbitraryObject:
+    """Module-level object that torch.save can pickle but weights_only=True rejects.
+
+    Must be defined at module scope so pickle can resolve its fully-qualified name during serialisation (local/nested
+    classes cannot be pickled by torch.save).
+    """
+
+    value = 42
+
+
+def _write_arbitrary_pickle_checkpoint(path: Path) -> None:
+    """Save a checkpoint that embeds an arbitrary class (requires pickle)."""
+    ckpt = {"model": {"weight": torch.tensor([1.0])}, "extra": _ArbitraryObject()}
+    torch.save(ckpt, path)
+
+
+# ---------------------------------------------------------------------------
+# Safe path (weights_only=True)
+# ---------------------------------------------------------------------------
+
+
+class TestSafeTorchLoadSafePath:
+    """Tensor-only checkpoints load without trust=True."""
+
+    def test_tensor_only_checkpoint_loads(self, tmp_path: Path) -> None:
+        """Pure-tensor checkpoint succeeds on the first safe-load attempt."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_tensor_only_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path)
+
+        assert "model" in result
+        assert torch.allclose(result["model"]["weight"], torch.tensor([1.0, 2.0]))
+
+    def test_accepts_pathlib_path(self, tmp_path: Path) -> None:
+        """Helper accepts a :class:`pathlib.Path` argument without error."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_tensor_only_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path)
+
+        assert "model" in result
+
+    def test_accepts_string_path(self, tmp_path: Path) -> None:
+        """Helper accepts a :class:`str` path argument without error."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_tensor_only_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(str(ckpt_path))
+
+        assert "model" in result
+
+
+# ---------------------------------------------------------------------------
+# Safe-globals fallback (legacy Namespace checkpoints)
+# ---------------------------------------------------------------------------
+
+
+class TestSafeTorchLoadSafeGlobals:
+    """Checkpoints with argparse.Namespace / SimpleNamespace load without trust=True."""
+
+    def test_argparse_namespace_loads_without_trust(self, tmp_path: Path) -> None:
+        """argparse.Namespace checkpoint succeeds via the safe-globals retry."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_namespace_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path)
+
+        assert isinstance(result["args"], argparse.Namespace)
+        assert result["args"].num_classes == 80
+
+    def test_simple_namespace_loads_without_trust(self, tmp_path: Path) -> None:
+        """SimpleNamespace checkpoint succeeds via the safe-globals retry."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_simple_namespace_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path)
+
+        assert isinstance(result["args"], SimpleNamespace)
+
+
+# ---------------------------------------------------------------------------
+# Arbitrary pickle — trust=False must raise, trust=True must succeed
+# ---------------------------------------------------------------------------
+
+
+class TestSafeTorchLoadTrustGate:
+    """Arbitrary-pickle checkpoints require explicit trust=True."""
+
+    def test_arbitrary_pickle_raises_without_trust(self, tmp_path: Path) -> None:
+        """Checkpoint with unknown Python object raises RuntimeError when trust=False."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_arbitrary_pickle_checkpoint(ckpt_path)
+
+        with pytest.raises(RuntimeError, match="trust_checkpoint=True"):
+            _safe_torch_load(ckpt_path, trust=False)
+
+    def test_arbitrary_pickle_raises_by_default(self, tmp_path: Path) -> None:
+        """Checkpoint with unknown Python object raises RuntimeError when trust omitted (default=False)."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_arbitrary_pickle_checkpoint(ckpt_path)
+
+        with pytest.raises(RuntimeError, match="trust_checkpoint=True"):
+            _safe_torch_load(ckpt_path)
+
+    def test_arbitrary_pickle_succeeds_with_trust(self, tmp_path: Path) -> None:
+        """Checkpoint with unknown Python object loads when trust=True."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_arbitrary_pickle_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path, trust=True)
+
+        assert "model" in result
+
+    def test_trust_true_emits_warning(self, tmp_path: Path) -> None:
+        """Trust=True triggers a UserWarning about unsafe loading."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_arbitrary_pickle_checkpoint(ckpt_path)
+
+        with pytest.warns(UserWarning, match="weights_only=False"):
+            _safe_torch_load(ckpt_path, trust=True)

tests/models/test_transformer.py

@@ -0,0 +1,501 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for transformer utilities, MS deformable attention core, and MSDeformAttn module."""
+
+import io
+
+import numpy as np
+import pytest
+import torch
+
+from rfdetr.models.ops.functions import ms_deform_attn_core_pytorch
+from rfdetr.models.ops.modules.ms_deform_attn import MSDeformAttn
+from rfdetr.models.transformer import gen_encoder_output_proposals, gen_sineembed_for_position
+
+
+@pytest.fixture(autouse=True)
+def _reset_random_seeds() -> None:
+    """Ensure reproducible random state for every test."""
+    torch.manual_seed(42)
+    torch.cuda.manual_seed_all(42)
+
+
+_MSDeformInputs = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[tuple[int, int]]]
+
+
+def _build_ms_deform_inputs(
+    bsz: int = 1,
+    n_heads: int = 2,
+    head_dim: int = 4,
+    len_q: int = 3,
+    npts: int = 1,
+    levels: list[tuple[int, int]] | None = None,
+) -> _MSDeformInputs:
+    """Build minimal valid inputs for ms_deform_attn_core_pytorch.
+
+    Args:
+        bsz: Batch size.
+        n_heads: Number of attention heads.
+        head_dim: Dimension per head.
+        len_q: Number of query elements.
+        npts: Number of sampling points per level.
+        levels: List of (H, W) int pairs; defaults to [(4, 4), (2, 2)].
+
+    Returns:
+        Tuple of (value, spatial_shapes_tensor, sampling_locations,
+                  attention_weights, spatial_shapes_hw).
+    """
+    if levels is None:
+        levels = [(4, 4), (2, 2)]
+    nlvl = len(levels)
+
+    total_hw = sum(ht * wd for ht, wd in levels)
+    spatial_shapes_tensor = torch.tensor(levels, dtype=torch.long)
+    value = torch.randn(bsz, n_heads, head_dim, total_hw)
+    # sampling_locations: (bsz, len_q, n_heads, nlvl, npts, 2) in [0, 1]
+    sampling_locations = torch.rand(bsz, len_q, n_heads, nlvl, npts, 2)
+    # attention_weights: (bsz, len_q, n_heads, nlvl * npts)
+    attention_weights = torch.softmax(torch.randn(bsz, len_q, n_heads, nlvl * npts), dim=-1)
+
+    return value, spatial_shapes_tensor, sampling_locations, attention_weights, levels
+
+
+def test_gen_encoder_output_proposals_passes_ij_indexing_to_meshgrid(monkeypatch) -> None:
+    """`gen_encoder_output_proposals` should call `torch.meshgrid` with explicit ij indexing."""
+    original_meshgrid = torch.meshgrid
+    call_count = 0
+
+    def _meshgrid_with_indexing_assertion(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if kwargs.get("indexing") != "ij":
+            raise AssertionError("torch.meshgrid must be called with indexing='ij'")
+        return original_meshgrid(*args, **kwargs)
+
+    monkeypatch.setattr(torch, "meshgrid", _meshgrid_with_indexing_assertion)
+
+    memory = torch.randn(1, 4, 8)
+    spatial_shapes = torch.tensor([[2, 2]], dtype=torch.long)
+
+    output_memory, output_proposals = gen_encoder_output_proposals(
+        memory,
+        spatial_shapes=spatial_shapes,
+    )
+
+    assert call_count == 1
+
+
+def test_gen_sineembed_for_position_keeps_box_dimensions_in_sin_cos_order() -> None:
+    """4D box positional embeddings must use the pretrained sin/cos order for all dimensions."""
+    pos_tensor = torch.tensor([[[0.125, 0.25, 0.5, 0.75]]], dtype=torch.float32)
+    dim = 4
+    scale = 2 * torch.pi
+    dim_t = torch.arange(dim, dtype=pos_tensor.dtype)
+    dim_t = 10000 ** (2 * (dim_t // 2) / dim)
+
+    expected_parts = []
+    for coord_idx in (1, 0, 2, 3):
+        coord = pos_tensor[:, :, coord_idx] * scale
+        encoded = coord[:, :, None] / dim_t
+        expected_parts.append(torch.stack((encoded[:, :, 0::2].sin(), encoded[:, :, 1::2].cos()), dim=3).flatten(2))
+    expected = torch.cat(expected_parts, dim=2)
+
+    actual = gen_sineembed_for_position(pos_tensor, dim=dim)
+
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-6)
+
+
+def test_gen_encoder_output_proposals_rejects_non_square_ij_indexing(monkeypatch) -> None:
+    """Wrong meshgrid indexing (xy vs ij) produces different proposals for non-square spatial shapes."""
+    original_meshgrid = torch.meshgrid
+
+    def _meshgrid_wrong_indexing(*args, **kwargs):
+        kwargs["indexing"] = "xy"
+        return original_meshgrid(*args, **kwargs)
+
+    # Use non-square spatial shapes so that ij vs xy indexing produces observably different outputs.
+    memory = torch.randn(1, 8, 8)
+    spatial_shapes = torch.tensor([[2, 4]], dtype=torch.long)
+
+    correct_memory, correct_proposals = gen_encoder_output_proposals(memory, spatial_shapes=spatial_shapes)
+
+    monkeypatch.setattr(torch, "meshgrid", _meshgrid_wrong_indexing)
+
+    wrong_memory, wrong_proposals = gen_encoder_output_proposals(memory, spatial_shapes=spatial_shapes)
+
+    assert not torch.allclose(correct_proposals, wrong_proposals), (
+        "xy indexing must produce different proposals than ij indexing for non-square spatial shapes"
+    )
+
+
+def test_gen_encoder_output_proposals_accepts_int_tuple_spatial_shapes() -> None:
+    """`gen_encoder_output_proposals` must accept `spatial_shapes` as a tensor of int pairs."""
+    batch = 2
+    ht, wd = 4, 4
+    memory = torch.randn(batch, ht * wd, 8)
+    spatial_shapes = torch.tensor([[ht, wd]], dtype=torch.long)
+
+    output_memory, output_proposals = gen_encoder_output_proposals(memory, spatial_shapes=spatial_shapes)
+
+    assert output_memory.shape == memory.shape
+    assert output_proposals.shape == (batch, ht * wd, 4)
+
+
+def test_gen_encoder_output_proposals_accepts_python_int_pair_spatial_shapes() -> None:
+    """`gen_encoder_output_proposals` must accept `spatial_shapes` as `list[tuple[int, int]]` with no padding mask.
+
+    Regression: `Transformer.forward` passes Python int pairs derived from `src.shape`, so the
+    export-driven call path uses `list[tuple[int, int]]` rather than a tensor.
+    """
+    batch, ht, wd, dim = 2, 4, 4, 8
+    memory = torch.randn(batch, ht * wd, dim)
+    spatial_shapes = [(ht, wd)]  # Python int pairs, as produced by Transformer.forward()
+
+    output_memory, output_proposals = gen_encoder_output_proposals(
+        memory,
+        memory_padding_mask=None,
+        spatial_shapes=spatial_shapes,
+    )
+
+    assert output_memory.shape == memory.shape
+    assert output_proposals.shape == (batch, ht * wd, 4)
+
+
+class TestMSDeformAttnCorePytorch:
+    """Tests for ms_deform_attn_core_pytorch with Python int pair spatial shapes.
+
+    Regression suite for torch.export.export compatibility: iterating over a spatial_shapes tensor yields FakeTensor
+    scalars during FakeTensor tracing, which cannot be used as Python int split/view sizes.  The function now accepts an
+    optional ``value_spatial_shapes_hw`` list of Python int pairs that bypasses tensor iteration.
+    """
+
+    @pytest.fixture
+    def make_inputs(self) -> _MSDeformInputs:
+        """Default two-level inputs: levels=[(4, 4), (2, 2)]."""
+        return _build_ms_deform_inputs()
+
+    @pytest.fixture
+    def single_level_inputs(self) -> _MSDeformInputs:
+        """Single-level inputs: levels=[(8, 8)]."""
+        return _build_ms_deform_inputs(levels=[(8, 8)])
+
+    def test_with_tensor_spatial_shapes(self, make_inputs: _MSDeformInputs) -> None:
+        """Baseline: passing only the tensor spatial_shapes still works."""
+        value, spatial_shapes_tensor, sampling_locations, attention_weights, _ = make_inputs
+
+        output = ms_deform_attn_core_pytorch(value, spatial_shapes_tensor, sampling_locations, attention_weights)
+
+        bsz, n_heads, head_dim, _ = value.shape
+        len_q = sampling_locations.shape[1]
+        assert output.shape == (bsz, len_q, n_heads * head_dim)
+
+    def test_with_python_int_pair_spatial_shapes(self, make_inputs: _MSDeformInputs) -> None:
+        """Regression: value_spatial_shapes_hw list of Python int pairs must be accepted.
+
+        This is the torch.export.export-compatible code path: tensor scalar values (from iterating over a FakeTensor)
+        cannot be used as split/view sizes, so the caller passes explicit Python int pairs via value_spatial_shapes_hw
+        instead.
+        """
+        value, spatial_shapes_tensor, sampling_locations, attention_weights, levels = make_inputs
+
+        output = ms_deform_attn_core_pytorch(
+            value,
+            spatial_shapes_tensor,
+            sampling_locations,
+            attention_weights,
+            value_spatial_shapes_hw=levels,
+        )
+
+        bsz, n_heads, head_dim, _ = value.shape
+        len_q = sampling_locations.shape[1]
+        assert output.shape == (bsz, len_q, n_heads * head_dim)
+
+    def test_tensor_and_hw_paths_produce_identical_outputs(self, make_inputs: _MSDeformInputs) -> None:
+        """Python int pair path and tensor iteration path must produce the same result."""
+        value, spatial_shapes_tensor, sampling_locations, attention_weights, levels = make_inputs
+
+        out_tensor_path = ms_deform_attn_core_pytorch(
+            value, spatial_shapes_tensor, sampling_locations, attention_weights
+        )
+        out_hw_path = ms_deform_attn_core_pytorch(
+            value,
+            spatial_shapes_tensor,
+            sampling_locations,
+            attention_weights,
+            value_spatial_shapes_hw=levels,
+        )
+
+        torch.testing.assert_close(out_tensor_path, out_hw_path)
+
+    def test_single_level(self, single_level_inputs: _MSDeformInputs) -> None:
+        """Single-level case with Python int pair path must not crash."""
+        value, spatial_shapes_tensor, sampling_locations, attention_weights, levels = single_level_inputs
+
+        output = ms_deform_attn_core_pytorch(
+            value,
+            spatial_shapes_tensor,
+            sampling_locations,
+            attention_weights,
+            value_spatial_shapes_hw=levels,
+        )
+
+        assert output.shape[0] == 1
+
+
+class TestMSDeformAttnModule:
+    """Tests for MSDeformAttn.forward covering the export-compatibility changes.
+
+    Validates the module-level parameter threading and export-mode assert guard introduced in the torch.export.export
+    compatibility fix.
+    """
+
+    _d_model = 32
+    _n_heads = 4
+    _n_levels = 2
+    _n_points = 1
+    _hw_pairs: list[tuple[int, int]] = [(4, 4), (2, 2)]
+
+    def _make_module_inputs(
+        self,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        list[tuple[int, int]],
+    ]:
+        """Build minimal valid inputs for MSDeformAttn.forward.
+
+        Returns:
+            Tuple of (query, reference_points, input_flatten,
+                      input_spatial_shapes, input_level_start_index, hw_pairs).
+        """
+        hw_pairs = self._hw_pairs
+        total_len = sum(ht * wd for ht, wd in hw_pairs)
+        bsz, len_q = 1, 3
+
+        query = torch.randn(bsz, len_q, self._d_model)
+        reference_points = torch.rand(bsz, len_q, self._n_levels, 2)
+        input_flatten = torch.randn(bsz, total_len, self._d_model)
+        input_spatial_shapes = torch.tensor(hw_pairs, dtype=torch.long)
+        # Cumulative start index per level: [0, H0*W0]
+        starts = [sum(ht * wd for ht, wd in hw_pairs[:idx]) for idx in range(self._n_levels)]
+        input_level_start_index = torch.tensor(starts, dtype=torch.long)
+
+        return query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, hw_pairs
+
+    def test_forward_without_hw_param_backward_compat(self) -> None:
+        """MSDeformAttn.forward without hw param produces correct output shape."""
+        module = MSDeformAttn(
+            d_model=self._d_model, n_levels=self._n_levels, n_heads=self._n_heads, n_points=self._n_points
+        )
+        query, ref_pts, input_flatten, spatial_shapes, level_start_index, _ = self._make_module_inputs()
+
+        output = module(query, ref_pts, input_flatten, spatial_shapes, level_start_index)
+
+        bsz, len_q, _ = query.shape
+        assert output.shape == (bsz, len_q, self._d_model)
+
+    def test_forward_with_hw_param_produces_correct_shape(self) -> None:
+        """MSDeformAttn.forward with input_spatial_shapes_hw produces correct output shape."""
+        module = MSDeformAttn(
+            d_model=self._d_model, n_levels=self._n_levels, n_heads=self._n_heads, n_points=self._n_points
+        )
+        query, ref_pts, input_flatten, spatial_shapes, level_start_index, hw_pairs = self._make_module_inputs()
+
+        output = module(
+            query, ref_pts, input_flatten, spatial_shapes, level_start_index, input_spatial_shapes_hw=hw_pairs
+        )
+
+        bsz, len_q, _ = query.shape
+        assert output.shape == (bsz, len_q, self._d_model)
+
+    def test_export_mode_forward_with_hw_param(self) -> None:
+        """MSDeformAttn.forward in export mode with hw param must not raise."""
+        module = MSDeformAttn(
+            d_model=self._d_model, n_levels=self._n_levels, n_heads=self._n_heads, n_points=self._n_points
+        )
+        module.export()
+        query, ref_pts, input_flatten, spatial_shapes, level_start_index, hw_pairs = self._make_module_inputs()
+
+        output = module(
+            query, ref_pts, input_flatten, spatial_shapes, level_start_index, input_spatial_shapes_hw=hw_pairs
+        )
+
+        bsz, len_q, _ = query.shape
+        assert output.shape == (bsz, len_q, self._d_model)
+
+    def test_export_flag_set_after_export_call(self) -> None:
+        """Calling .export() must set _export=True, enabling the torch._assert guard path."""
+        module = MSDeformAttn(
+            d_model=self._d_model, n_levels=self._n_levels, n_heads=self._n_heads, n_points=self._n_points
+        )
+        assert not module._export
+
+        module.export()
+
+        assert module._export
+
+
+class TestGenEncoderOutputProposalsDynamicBatch:
+    """Regression tests for dynamic batch support in gen_encoder_output_proposals.
+
+    Ensures that the ONNX-symbolic refactoring (PR #950 / issue #949) does not bake a fixed batch dimension into
+    proposals and that output shapes are correct for varying batch sizes.
+    """
+
+    @pytest.mark.parametrize("batch_size", [1, 2, 4, 8])
+    def test_output_shape_invariant_across_batch_sizes(self, batch_size: int) -> None:
+        """Output shapes must scale correctly with batch size, with no baked constants.
+
+        Args:
+            batch_size: Number of images in the batch.
+        """
+        ht, wd, dim = 4, 4, 8
+        memory = torch.randn(batch_size, ht * wd, dim)
+        spatial_shapes = [(ht, wd)]
+
+        output_memory, output_proposals = gen_encoder_output_proposals(
+            memory, memory_padding_mask=None, spatial_shapes=spatial_shapes
+        )
+
+        assert output_memory.shape == (batch_size, ht * wd, dim)
+        assert output_proposals.shape == (batch_size, ht * wd, 4)
+
+    def test_proposals_semantically_equivalent_across_batch_sizes(self) -> None:
+        """Proposals for batch=1 and batch=4 must be identical per image.
+
+        Regression: if batch_size were baked as a constant, repeating the same image
+        N times would produce different proposals for each copy.
+        """
+        ht, wd, dim = 4, 4, 8
+        memory_single = torch.randn(1, ht * wd, dim)
+        memory_multi = memory_single.expand(4, -1, -1).contiguous()
+        spatial_shapes = [(ht, wd)]
+
+        _, proposals_single = gen_encoder_output_proposals(
+            memory_single, memory_padding_mask=None, spatial_shapes=spatial_shapes
+        )
+        _, proposals_multi = gen_encoder_output_proposals(
+            memory_multi, memory_padding_mask=None, spatial_shapes=spatial_shapes
+        )
+
+        torch.testing.assert_close(proposals_single.expand(4, -1, -1), proposals_multi)
+
+    @pytest.mark.parametrize("batch_size", [1, 4])
+    def test_output_shape_invariant_with_padding_mask(self, batch_size: int) -> None:
+        """Output shapes must be correct when memory_padding_mask is provided with varying batch sizes.
+
+        Regression for PR #950 / issue #949: the masked branch used .reshape(-1, h, w, 1) to infer the batch dimension
+        dynamically; this test verifies the branch handles varying batch sizes without error.
+
+        Args:
+            batch_size: Number of images in the batch.
+        """
+        ht, wd, dim = 4, 4, 8
+        total_hw = ht * wd
+        memory = torch.randn(batch_size, total_hw, dim)
+        # Mask shape: (batch, sum_hw) — True means padding (invalid position)
+        memory_padding_mask = torch.zeros(batch_size, total_hw, dtype=torch.bool)
+        spatial_shapes = [(ht, wd)]
+
+        output_memory, output_proposals = gen_encoder_output_proposals(
+            memory, memory_padding_mask=memory_padding_mask, spatial_shapes=spatial_shapes
+        )
+
+        assert output_memory.shape == (batch_size, total_hw, dim)
+        assert output_proposals.shape == (batch_size, total_hw, 4)
+
+    @pytest.mark.parametrize("batch_size", [1, 4, 8])
+    def test_onnx_export_with_dynamic_batch_axis(self, batch_size: int) -> None:
+        """ONNX export with dynamic batch axis must run inference for batch sizes other than the trace batch.
+
+        Regression for issue #949: exporting with a fixed trace batch baked `Reshape([8,...])` as a constant ONNX node,
+        causing TRT engines to fail at inference for any batch != 8. Skipped when onnx or onnxruntime is not installed.
+        """
+        pytest.importorskip("onnx")
+        onnxruntime = pytest.importorskip("onnxruntime")
+
+        ht, wd, dim = 4, 4, 8
+        spatial_shapes_list = [(ht, wd)]
+
+        class _ProposalModule(torch.nn.Module):
+            """Thin wrapper to export gen_encoder_output_proposals via torch.onnx."""
+
+            def forward(self, memory: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+                """Forward pass delegating to gen_encoder_output_proposals."""
+                return gen_encoder_output_proposals(
+                    memory, memory_padding_mask=None, spatial_shapes=spatial_shapes_list
+                )
+
+        module = _ProposalModule()
+        trace_memory = torch.randn(2, ht * wd, dim)
+
+        buf = io.BytesIO()
+        torch.onnx.export(
+            module,
+            (trace_memory,),
+            buf,
+            input_names=["memory"],
+            output_names=["output_memory", "output_proposals"],
+            dynamic_axes={"memory": {0: "batch"}},
+            opset_version=17,
+        )
+        buf.seek(0)
+        onnx_bytes = buf.read()
+
+        session = onnxruntime.InferenceSession(onnx_bytes, providers=["CPUExecutionProvider"])
+        memory_np = np.random.randn(batch_size, ht * wd, dim).astype(np.float32)
+        out_memory, out_proposals = session.run(None, {"memory": memory_np})
+        assert out_memory.shape == (batch_size, ht * wd, dim), f"wrong memory shape for batch={batch_size}"
+        assert out_proposals.shape == (batch_size, ht * wd, 4), f"wrong proposals shape for batch={batch_size}"
+
+
+def test_ms_deform_attn_core_pytorch_export_compatible() -> None:
+    """torch.export.export must succeed on a module using ms_deform_attn_core_pytorch with hw param.
+
+    Regression test for the FakeTensor tracing failure: iterating over spatial_shapes and using the scalar elements as
+    split/view sizes fails during torch.export.export because FakeTensor data is not allocated. Passing
+    value_spatial_shapes_hw (concrete Python ints from a module attribute) bypasses the tensor iteration entirely.
+    """
+    levels: list[tuple[int, int]] = [(4, 4), (2, 2)]
+    bsz, n_heads, head_dim = 1, 2, 4
+    total_hw = sum(ht * wd for ht, wd in levels)
+    len_q, nlvl, npts = 3, len(levels), 1
+
+    class _MinimalDeformAttn(torch.nn.Module):
+        """Minimal wrapper to test torch.export.export on the hw-param code path."""
+
+        def __init__(self, hw: list[tuple[int, int]]) -> None:
+            super().__init__()
+            self.hw = hw
+
+        def forward(
+            self,
+            value: torch.Tensor,
+            spatial_shapes: torch.Tensor,
+            sampling_locations: torch.Tensor,
+            attention_weights: torch.Tensor,
+        ) -> torch.Tensor:
+            """Forward using concrete Python int pairs for export compatibility."""
+            return ms_deform_attn_core_pytorch(
+                value,
+                spatial_shapes,
+                sampling_locations,
+                attention_weights,
+                value_spatial_shapes_hw=self.hw,
+            )
+
+    value = torch.randn(bsz, n_heads, head_dim, total_hw)
+    spatial_shapes = torch.tensor(levels, dtype=torch.long)
+    sampling_locations = torch.rand(bsz, len_q, n_heads, nlvl, npts, 2)
+    attention_weights = torch.softmax(torch.randn(bsz, len_q, n_heads, nlvl * npts), dim=-1)
+
+    module = _MinimalDeformAttn(hw=levels)
+
+    exported = torch.export.export(module, args=(value, spatial_shapes, sampling_locations, attention_weights))
+    assert exported is not None

tests/models/test_transformer_keypoints.py

@@ -0,0 +1,250 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for GroupPose-oriented transformer streams."""
+
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from rfdetr.models.transformer import Transformer, TransformerDecoder, TransformerDecoderLayer, build_transformer
+
+
+def _build_transformer_inputs(
+    batch_size: int = 2,
+    hidden_dim: int = 16,
+    num_levels: int = 2,
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], torch.Tensor, torch.Tensor]:
+    """Build minimal synthetic multi-scale inputs for `Transformer.forward`.
+
+    Args:
+        batch_size: Mini-batch size.
+        hidden_dim: Transformer and input channel size.
+        num_levels: Number of feature pyramid levels.
+
+    Returns:
+        srcs, masks, pos_embeds, refpoint_embed, query_feat.
+    """
+    spatial_shapes = [(4, 4), (2, 2)]
+    srcs = [
+        torch.randn(batch_size, hidden_dim, spatial_shapes[idx][0], spatial_shapes[idx][1]) for idx in range(num_levels)
+    ]
+    masks = [torch.zeros(batch_size, h, w, dtype=torch.bool) for h, w in spatial_shapes[:num_levels]]
+    pos_embeds = [
+        torch.randn(batch_size, hidden_dim, spatial_shapes[idx][0], spatial_shapes[idx][1]) for idx in range(num_levels)
+    ]
+    refpoint_embed = torch.rand(6, 4)
+    query_feat = torch.randn(6, hidden_dim)
+    return srcs, masks, pos_embeds, refpoint_embed, query_feat
+
+
+def test_transformer_keypoint_disabled_matches_default_contract() -> None:
+    """Transformer without GroupPose should keep the 4-item return contract."""
+    srcs, masks, pos_embeds, refpoint_embed, query_feat = _build_transformer_inputs()
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        use_grouppose_keypoints=False,
+    )
+
+    outputs = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=None)
+    assert len(outputs) == 4, f"Expected 4 outputs, got {len(outputs)}"
+    hs, references, memory_ts, boxes_ts = outputs
+    assert hs is not None and references is not None
+    assert memory_ts is None and boxes_ts is None
+
+
+def test_transformer_keypoint_enabled_shapes() -> None:
+    """GroupPose path should emit keypoint hidden states and encoder keypoint slots."""
+    srcs, masks, pos_embeds, refpoint_embed, query_feat = _build_transformer_inputs()
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        two_stage=True,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+    )
+    transformer.enc_out_class_embed = nn.ModuleList([nn.Linear(16, 2)])
+    transformer.enc_out_bbox_embed = nn.ModuleList([nn.Linear(16, 4)])
+
+    outputs = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=None)
+    assert len(outputs) == 7, f"Expected 7 outputs, got {len(outputs)}"
+    hs, references, memory_ts, boxes_ts, keypoint_hs, enc_kp_predictions, keypoint_memory_ts = outputs
+
+    assert isinstance(hs, torch.Tensor)
+    assert hs.shape[-1] == 16
+    assert references.shape[-1] == 4
+    assert memory_ts is not None and boxes_ts is not None
+    assert keypoint_hs is not None
+    assert keypoint_hs.shape[3] == 17
+    assert enc_kp_predictions is not None
+    assert enc_kp_predictions.shape[-1] == 16
+    assert keypoint_memory_ts is not None
+
+
+def test_build_transformer_defaults_inter_instance_keypoint_attention_to_config_default() -> None:
+    """Older args objects without `inter_instance_kp_attn` should keep the preview topology default."""
+    args = SimpleNamespace(
+        hidden_dim=16,
+        sa_nheads=4,
+        ca_nheads=4,
+        num_queries=6,
+        dropout=0.0,
+        dim_feedforward=32,
+        dec_layers=1,
+        group_detr=1,
+        num_feature_levels=2,
+        dec_n_points=1,
+        lite_refpoint_refine=True,
+        decoder_norm="LN",
+        bbox_reparam=False,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+    )
+
+    transformer = build_transformer(args)
+
+    decoder_layer = transformer.decoder.layers[0]
+    assert isinstance(decoder_layer, TransformerDecoderLayer)
+    assert decoder_layer.enable_keypoint_processing
+    assert not decoder_layer.inter_instance_kp_attn
+
+
+def test_keypoint_class_mask_person_only() -> None:
+    """Person-only schema `[17]` should build a keypoint class mask with only self-class tokens."""
+    layer = TransformerDecoderLayer(
+        d_model=16,
+        sa_nhead=4,
+        ca_nhead=4,
+        dim_feedforward=32,
+        num_feature_levels=2,
+        enable_keypoint_processing=True,
+        grouppose_keypoint_dim_downscale=1,
+        keypoint_cross_attn=False,
+        inter_instance_kp_attn=False,
+    )
+    decoder = TransformerDecoder(
+        decoder_layer=layer,
+        num_layers=1,
+        return_intermediate=True,
+        d_model=16,
+        lite_refpoint_refine=True,
+        enable_keypoint_processing=True,
+        num_keypoints_per_class=[17],
+        grouppose_keypoint_dim_downscale=1,
+    )
+
+    assert decoder.keypoint_pos_embed is not None
+    assert decoder.keypoint_class_mask.shape == (18, 18)
+    assert decoder.keypoint_class_mask.dtype == torch.bool
+    assert not decoder.keypoint_class_mask.any()
+
+
+def test_enc_keypoint_embed_eval_uses_only_head_zero() -> None:
+    """Encoder keypoint path must use a single head (head 0) in eval mode.
+
+    Regression: ``group_detr = len(self.enc_out_keypoint_embed)`` without a
+    ``self.training`` guard caused eval mode to split ``num_queries`` across all
+    group heads instead of routing every query through head 0. Fix:
+    ``group_detr = len(...) if self.training else 1``.
+
+    Strategy: zero all head weights/biases; set head-0 last-layer bias to
+    ``sentinel``. In eval mode every query routes through head 0, so
+    ``kp_pred[..., 2:]`` (the pure-delta dims unaffected by ref_xy/wh) must all
+    equal ``sentinel``. In training mode only the first 1/group_detr queries go
+    through head 0 (the rest equal 0.0).
+    """
+    group_detr = 3
+    num_queries = 6  # divisible by group_detr
+    hidden_dim = 16
+    batch_size = 1
+    sentinel = 50.0
+
+    srcs, masks, pos_embeds, _, _ = _build_transformer_inputs(batch_size=batch_size, hidden_dim=hidden_dim)
+    refpoint_embed = torch.rand(num_queries, 4)
+    query_feat = torch.randn(num_queries, hidden_dim)
+
+    transformer = Transformer(
+        d_model=hidden_dim,
+        num_queries=num_queries,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        two_stage=True,
+        group_detr=group_detr,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[2],
+    )
+    transformer.enc_out_class_embed = nn.ModuleList([nn.Linear(hidden_dim, 2) for _ in range(group_detr)])
+    transformer.enc_out_bbox_embed = nn.ModuleList([nn.Linear(hidden_dim, 4) for _ in range(group_detr)])
+
+    # Zero all keypoint head weights and biases; give head 0 a distinctive output.
+    with torch.no_grad():
+        for _, head in enumerate(transformer.enc_out_keypoint_embed):
+            for layer in head.layers:
+                layer.weight.zero_()
+                layer.bias.zero_()
+        transformer.enc_out_keypoint_embed[0].layers[-1].bias.fill_(sentinel)
+
+    transformer.eval()
+    with torch.no_grad():
+        outputs = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=None)
+
+    _, _, _, _, _, enc_kp_predictions, _ = outputs
+    assert enc_kp_predictions is not None, "enc_kp_predictions should not be None in keypoint mode"
+
+    # kp_pred = [kp_xy(2 dims), kp_delta[2:]]; dims 2: are pure MLP output unaffected by ref_xy/wh.
+    kp_beyond_xy = enc_kp_predictions[..., 2:]
+    assert (kp_beyond_xy == sentinel).all(), (
+        f"Eval mode must route all {num_queries} queries through head 0 (bias={sentinel}). "
+        f"Got min={kp_beyond_xy.min().item():.2f}, max={kp_beyond_xy.max().item():.2f}. "
+        "Bug: group_detr not guarded by self.training in enc_out_keypoint_embed loop."
+    )
+
+
+def test_cross_attn_srcs_none_backward_compat() -> None:
+    """`cross_attn_srcs=None` must remain equivalent to passing the primary feature stream."""
+    srcs, masks, pos_embeds, refpoint_embed, query_feat = _build_transformer_inputs()
+
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        use_grouppose_keypoints=False,
+    )
+    outputs_default = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=None)
+    outputs_explicit = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=srcs)
+
+    assert len(outputs_default) == len(outputs_explicit) == 4
+    for default_part, explicit_part in zip(outputs_default, outputs_explicit):
+        if default_part is None:
+            assert explicit_part is None
+        else:
+            torch.testing.assert_close(default_part, explicit_part)

tests/models/test_transformer_onnx_spatial_shapes.py

@@ -0,0 +1,409 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""ONNX-export regression tests for the shared Transformer's spatial_shapes path.
+
+These guard the fix that builds ``spatial_shapes`` from symbolic feature-map shapes (``Shape`` -> ``Concat``) instead of
+``torch.empty`` + in-place index assignment. The latter (added in #871 to keep the trace symbolic for dynamic-batch
+export) emitted a ``ScatterND`` that fed a shape tensor, which TensorRT rejects ("IScatterLayer cannot be used to
+compute a shape tensor"). The constant-baking ``torch.as_tensor`` alternative avoids the ScatterND but regresses the
+symbolic trace back to a baked constant.
+
+The Transformer is shared by detection, segmentation and keypoint models, so a single low-level export here covers the
+spatial_shapes path for all of them.
+"""
+
+import inspect
+import io
+
+import numpy as np
+import pytest
+import torch
+from torch import nn
+
+onnx = pytest.importorskip("onnx", reason="onnx not installed; skip ONNX export tests")
+
+from rfdetr.models.transformer import Transformer  # noqa: E402
+
+# CI guard: torch._shape_as_tensor is a private ATen API used on the live forward path in
+# Transformer.forward(). If a future PyTorch upgrade removes it, this assertion fails
+# loudly in CI before users hit an AttributeError at inference time.
+assert hasattr(torch, "_shape_as_tensor"), (
+    "torch._shape_as_tensor not found — update spatial_shapes construction in "
+    "Transformer.forward() for the new PyTorch API."
+)
+
+# dynamo kwarg was added in PyTorch 2.1; our minimum is >=2.2.0, so it is always present.
+# Check via inspect so that a future signature removal surfaces here rather than as a
+# confusing TypeError inside torch.onnx.export.
+_DYNAMO_KWARG: dict[str, bool] = (
+    {"dynamo": False} if "dynamo" in inspect.signature(torch.onnx.export).parameters else {}
+)
+
+
+class _TransformerExportWrapper(nn.Module):
+    """Wrap Transformer.forward with flat tensor args for 2-level ONNX export.
+
+    ``torch.onnx.export`` (TorchScript/non-dynamo path) cannot trace Python list arguments; this wrapper flattens the
+    list args of ``Transformer.forward`` into positional tensor arguments.
+    """
+
+    def __init__(self, transformer: Transformer) -> None:
+        super().__init__()
+        self.transformer = transformer
+
+    def forward(
+        self,
+        s0: torch.Tensor,
+        s1: torch.Tensor,
+        p0: torch.Tensor,
+        p1: torch.Tensor,
+        m0: torch.Tensor,
+        m1: torch.Tensor,
+        refpoint_embed: torch.Tensor,
+        query_feat: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run Transformer with two feature levels and return the first decoder output.
+
+        Args:
+            s0: First-level feature map of shape ``(B, C, H0, W0)``.
+            s1: Second-level feature map of shape ``(B, C, H1, W1)``.
+            p0: Positional embeddings for level 0, same shape as ``s0``.
+            p1: Positional embeddings for level 1, same shape as ``s1``.
+            m0: Boolean padding mask for level 0 of shape ``(B, H0, W0)``.
+            m1: Boolean padding mask for level 1 of shape ``(B, H1, W1)``.
+            refpoint_embed: Reference point embeddings of shape ``(num_queries, 4)``.
+            query_feat: Query feature embeddings of shape ``(num_queries, C)``.
+
+        Returns:
+            First intermediate decoder output tensor.
+        """
+        outputs = self.transformer([s0, s1], [m0, m1], [p0, p1], refpoint_embed, query_feat, cross_attn_srcs=None)
+        return outputs[0]
+
+
+class _TransformerExportWrapper1Level(nn.Module):
+    """Wrap Transformer.forward with flat tensor args for 1-level ONNX export.
+
+    Production RF-DETR models (Base, Small, Nano, Medium, Large) use a single feature level
+    (``projector_scale=["P4"]``). This wrapper exercises that path.
+    """
+
+    def __init__(self, transformer: Transformer) -> None:
+        super().__init__()
+        self.transformer = transformer
+
+    def forward(
+        self,
+        s0: torch.Tensor,
+        p0: torch.Tensor,
+        m0: torch.Tensor,
+        refpoint_embed: torch.Tensor,
+        query_feat: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run Transformer with one feature level and return the first decoder output.
+
+        Args:
+            s0: Feature map of shape ``(B, C, H0, W0)``.
+            p0: Positional embeddings for level 0, same shape as ``s0``.
+            m0: Boolean padding mask for level 0 of shape ``(B, H0, W0)``.
+            refpoint_embed: Reference point embeddings of shape ``(num_queries, 4)``.
+            query_feat: Query feature embeddings of shape ``(num_queries, C)``.
+
+        Returns:
+            First intermediate decoder output tensor.
+        """
+        outputs = self.transformer([s0], [m0], [p0], refpoint_embed, query_feat, cross_attn_srcs=None)
+        return outputs[0]
+
+
+# ---------------------------------------------------------------------------
+# Session-scoped fixtures — build and export once per test session
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="session")
+def transformer_wrapper_2lvl() -> _TransformerExportWrapper:
+    """Build a 2-level Transformer and wrap it for ONNX export.
+
+    Returns:
+        Eval-mode ``_TransformerExportWrapper`` with ``d_model=16``.
+    """
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        use_grouppose_keypoints=False,
+    )
+    return _TransformerExportWrapper(transformer.eval()).eval()
+
+
+@pytest.fixture(scope="session")
+def transformer_wrapper_1lvl() -> _TransformerExportWrapper1Level:
+    """Build a 1-level Transformer and wrap it for ONNX export.
+
+    Returns:
+        Eval-mode ``_TransformerExportWrapper1Level`` with ``d_model=16``.
+    """
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=1,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        use_grouppose_keypoints=False,
+    )
+    return _TransformerExportWrapper1Level(transformer.eval()).eval()
+
+
+@pytest.fixture(scope="session")
+def example_inputs_2lvl() -> tuple[torch.Tensor, ...]:
+    """Return fixed 2-level example input tensors for ``_TransformerExportWrapper``.
+
+    Level 0: spatial size 4×4. Level 1: spatial size 2×2. Batch size 1, ``d_model=16``.
+
+    Returns:
+        8-tuple ``(s0, s1, p0, p1, m0, m1, refpoint_embed, query_feat)``.
+    """
+    return (
+        torch.randn(1, 16, 4, 4),
+        torch.randn(1, 16, 2, 2),
+        torch.randn(1, 16, 4, 4),
+        torch.randn(1, 16, 2, 2),
+        torch.zeros(1, 4, 4, dtype=torch.bool),
+        torch.zeros(1, 2, 2, dtype=torch.bool),
+        torch.rand(6, 4),
+        torch.randn(6, 16),
+    )
+
+
+@pytest.fixture(scope="session")
+def example_inputs_1lvl() -> tuple[torch.Tensor, ...]:
+    """Return fixed 1-level example input tensors for ``_TransformerExportWrapper1Level``.
+
+    Level 0: spatial size 4×4. Batch size 1, ``d_model=16``.
+
+    Returns:
+        5-tuple ``(s0, p0, m0, refpoint_embed, query_feat)``.
+    """
+    return (
+        torch.randn(1, 16, 4, 4),
+        torch.randn(1, 16, 4, 4),
+        torch.zeros(1, 4, 4, dtype=torch.bool),
+        torch.rand(6, 4),
+        torch.randn(6, 16),
+    )
+
+
+@pytest.fixture(scope="session")
+def exported_static_onnx(
+    tmp_path_factory: pytest.TempPathFactory,
+    transformer_wrapper_2lvl: _TransformerExportWrapper,
+    example_inputs_2lvl: tuple[torch.Tensor, ...],
+) -> "onnx.ModelProto":
+    """Export the 2-level Transformer to ONNX (static batch) and return the loaded proto.
+
+    Returns:
+        Loaded ``onnx.ModelProto`` for structural graph assertions.
+    """
+    out = tmp_path_factory.mktemp("onnx_static") / "transformer.onnx"
+    torch.onnx.export(
+        transformer_wrapper_2lvl,
+        example_inputs_2lvl,
+        str(out),
+        input_names=["s0", "s1", "p0", "p1", "m0", "m1", "refpoint_embed", "query_feat"],
+        output_names=["hs"],
+        opset_version=17,
+        **_DYNAMO_KWARG,
+    )
+    return onnx.load(str(out))
+
+
+@pytest.fixture(scope="session")
+def exported_dynamic_onnx_bytes(
+    transformer_wrapper_2lvl: _TransformerExportWrapper,
+    example_inputs_2lvl: tuple[torch.Tensor, ...],
+) -> bytes:
+    """Export the 2-level Transformer with dynamic batch axis and return ONNX bytes.
+
+    Returns:
+        Serialized ONNX model bytes for onnxruntime inference with variable batch sizes.
+    """
+    buf = io.BytesIO()
+    torch.onnx.export(
+        transformer_wrapper_2lvl,
+        example_inputs_2lvl,
+        buf,
+        input_names=["s0", "s1", "p0", "p1", "m0", "m1", "refpoint_embed", "query_feat"],
+        output_names=["hs"],
+        dynamic_axes={
+            "s0": {0: "batch"},
+            "s1": {0: "batch"},
+            "p0": {0: "batch"},
+            "p1": {0: "batch"},
+            "m0": {0: "batch"},
+            "m1": {0: "batch"},
+        },
+        opset_version=17,
+        **_DYNAMO_KWARG,
+    )
+    return buf.getvalue()
+
+
+@pytest.fixture(scope="session")
+def exported_1lvl_onnx(
+    tmp_path_factory: pytest.TempPathFactory,
+    transformer_wrapper_1lvl: _TransformerExportWrapper1Level,
+    example_inputs_1lvl: tuple[torch.Tensor, ...],
+) -> "onnx.ModelProto":
+    """Export the 1-level Transformer to ONNX and return the loaded proto.
+
+    Returns:
+        Loaded ``onnx.ModelProto`` for structural graph assertions.
+    """
+    out = tmp_path_factory.mktemp("onnx_1lvl") / "transformer_1lvl.onnx"
+    torch.onnx.export(
+        transformer_wrapper_1lvl,
+        example_inputs_1lvl,
+        str(out),
+        input_names=["s0", "p0", "m0", "refpoint_embed", "query_feat"],
+        output_names=["hs"],
+        opset_version=17,
+        **_DYNAMO_KWARG,
+    )
+    return onnx.load(str(out))
+
+
+# ---------------------------------------------------------------------------
+# Tests: structural ONNX graph assertions (2-level static export)
+# ---------------------------------------------------------------------------
+
+
+def test_spatial_shapes_export_has_no_scatternd(exported_static_onnx: "onnx.ModelProto") -> None:
+    """The exported Transformer must not contain a ScatterND (TRT shape-tensor killer)."""
+    op_types = [n.op_type for n in exported_static_onnx.graph.node]
+    assert "ScatterND" not in op_types, (
+        "ScatterND reintroduced in Transformer export — spatial_shapes is no longer "
+        "built from symbolic Shape ops; this breaks TensorRT engine building."
+    )
+
+
+def test_spatial_shapes_export_is_shape_derived(exported_static_onnx: "onnx.ModelProto") -> None:
+    """Sanity-check that the exported 2-level Transformer graph contains Shape ops.
+
+    Note: ``spatial_shapes`` itself traces as a ``Constant`` node in TorchScript ONNX export —
+    the tracer records concrete H,W values at trace time, not ``Shape`` ops. The ``Shape`` ops
+    present here originate from other model computations (e.g. batch-size extraction). The true
+    regression guards are ``test_spatial_shapes_export_has_no_scatternd`` (no ScatterND) and
+    ``test_spatial_shapes_dynamic_batch_inference`` (runtime correctness at variable batch size).
+    """
+    op_types = [n.op_type for n in exported_static_onnx.graph.node]
+    assert "Shape" in op_types, "No Shape op in 2-level Transformer ONNX graph — unexpected graph structure change."
+
+
+# ---------------------------------------------------------------------------
+# Tests: dynamic-batch export (2-level)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("batch_size", [pytest.param(1, id="batch1"), pytest.param(2, id="batch2")])
+def test_spatial_shapes_dynamic_batch_inference(
+    exported_dynamic_onnx_bytes: bytes,
+    batch_size: int,
+) -> None:
+    """Dynamic-batch Transformer ONNX must run onnxruntime inference at batch != trace batch.
+
+    Regression guard: a baked batch constant in ``spatial_shapes`` or any upstream tensor
+    would cause shape mismatches at runtime for any batch size other than the trace batch (1).
+    """
+    onnxruntime = pytest.importorskip("onnxruntime", reason="onnxruntime not installed")
+    session = onnxruntime.InferenceSession(exported_dynamic_onnx_bytes, providers=["CPUExecutionProvider"])
+    # The TorchScript tracer may constant-fold positional embeddings (p0, p1) into the
+    # graph; query the actual session inputs rather than assuming all 8 are present.
+    actual_names = {inp.name for inp in session.get_inputs()}
+    candidate_feeds = {
+        "s0": np.random.randn(batch_size, 16, 4, 4).astype(np.float32),
+        "s1": np.random.randn(batch_size, 16, 2, 2).astype(np.float32),
+        "p0": np.random.randn(batch_size, 16, 4, 4).astype(np.float32),
+        "p1": np.random.randn(batch_size, 16, 2, 2).astype(np.float32),
+        "m0": np.zeros((batch_size, 4, 4), dtype=bool),
+        "m1": np.zeros((batch_size, 2, 2), dtype=bool),
+        "refpoint_embed": np.random.rand(6, 4).astype(np.float32),
+        "query_feat": np.random.randn(6, 16).astype(np.float32),
+    }
+    feeds = {k: v for k, v in candidate_feeds.items() if k in actual_names}
+    (hs,) = session.run(None, feeds)
+    assert hs.shape[1] == batch_size, f"Expected batch dim=={batch_size} at index 1, got shape {hs.shape}"
+
+
+# ---------------------------------------------------------------------------
+# Tests: single-level export (production models use 1 feature level)
+# ---------------------------------------------------------------------------
+
+
+def test_spatial_shapes_single_level_export_has_no_scatternd(
+    exported_1lvl_onnx: "onnx.ModelProto",
+) -> None:
+    """Single-level Transformer ONNX export must not contain ScatterND."""
+    op_types = [n.op_type for n in exported_1lvl_onnx.graph.node]
+    assert "ScatterND" not in op_types, (
+        "ScatterND present in single-level Transformer export — torch.stack on a "
+        "one-element list must still produce a Shape-derived result."
+    )
+
+
+def test_spatial_shapes_single_level_export_is_shape_derived(
+    exported_1lvl_onnx: "onnx.ModelProto",
+) -> None:
+    """Sanity-check that the exported 1-level Transformer graph contains Shape ops.
+
+    Note: ``spatial_shapes`` itself traces as a ``Constant`` node in TorchScript ONNX export —
+    the tracer records the concrete H,W value at trace time. The ``Shape`` ops present here
+    originate from other model computations. The true regression guards are
+    ``test_spatial_shapes_single_level_export_has_no_scatternd`` and the dynamic-batch
+    inference test.
+    """
+    op_types = [n.op_type for n in exported_1lvl_onnx.graph.node]
+    assert "Shape" in op_types, "No Shape op in 1-level Transformer ONNX graph — unexpected graph structure change."
+
+
+# ---------------------------------------------------------------------------
+# Tests: level_start_index numerical correctness
+# ---------------------------------------------------------------------------
+
+
+def test_level_start_index_correctness_two_levels() -> None:
+    """spatial_shapes construction must extract correct (H, W) for non-square feature maps.
+
+    Uses H0=8, W0=6 and H1=4, W1=3 — non-square and H≠W so a transposed [2:4] slice would produce wrong values. Verifies
+    that the ``torch.stack(_shape_as_tensor)`` formula in ``Transformer.forward()`` yields
+    ``spatial_shapes=[[8,6],[4,3]]`` and ``level_start_index=[0, 48]`` (cumulative H*W: 8*6=48, then 48+4*3=60 but index
+    stops at level boundaries → [0, 48]).
+    """
+    s0 = torch.randn(1, 16, 8, 6)
+    s1 = torch.randn(1, 16, 4, 3)
+    srcs = [s0, s1]
+
+    spatial_shapes = torch.stack([torch._shape_as_tensor(src)[2:4] for src in srcs]).to(
+        device=s0.device, dtype=torch.long
+    )
+    level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
+
+    assert torch.equal(spatial_shapes, torch.tensor([[8, 6], [4, 3]], dtype=torch.long)), (
+        f"spatial_shapes mismatch: got {spatial_shapes.tolist()}"
+    )
+    assert torch.equal(level_start_index, torch.tensor([0, 48], dtype=torch.long)), (
+        f"level_start_index expected [0, 48], got {level_start_index.tolist()}"
+    )

tests/models/test_validate_shape_dims.py

@@ -0,0 +1,152 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for :func:`rfdetr.detr._validate_shape_dims` and :func:`rfdetr.detr._resolve_patch_size`.
+
+Tests call each helper directly so each validation path has a single focused test without the export/predict scaffolding
+overhead.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from rfdetr.detr import _resolve_patch_size, _validate_shape_dims
+
+
+class TestValidateShapeDimsHappyPath:
+    """_validate_shape_dims returns normalised (height, width) for valid inputs."""
+
+    def test_exact_plain_ints(self) -> None:
+        """Plain int dims divisible by block_size are returned unchanged."""
+        assert _validate_shape_dims((56, 112), 14, 14, 1) == (56, 112)
+
+    def test_returns_plain_int_tuple(self) -> None:
+        """Return type is always a tuple of plain Python int."""
+        h, w = _validate_shape_dims((56, 56), 14, 14, 1)
+        assert type(h) is int and type(w) is int
+
+    def test_numpy_int_accepted(self) -> None:
+        """numpy.int64 dims are accepted via operator.index and normalised."""
+        import numpy as np
+
+        h, w = _validate_shape_dims((np.int64(56), np.int64(112)), 14, 14, 1)
+        assert (h, w) == (56, 112)
+        assert type(h) is int and type(w) is int
+
+    def test_non_square_shape(self) -> None:
+        """Non-square (H != W) shapes are returned correctly."""
+        assert _validate_shape_dims((112, 224), 14, 14, 1) == (112, 224)
+
+    def test_block_size_from_num_windows(self) -> None:
+        """block_size = patch_size * num_windows; both dims divisible by it."""
+        # patch_size=16, num_windows=2 → block_size=32
+        assert _validate_shape_dims((64, 128), 32, 16, 2) == (64, 128)
+
+
+class TestValidateShapeDimsArityErrors:
+    """_validate_shape_dims raises ValueError for non-two-element shapes."""
+
+    def test_one_element_raises(self) -> None:
+        """Single-element tuple must raise ValueError."""
+        with pytest.raises(ValueError, match="shape must be a sequence"):
+            _validate_shape_dims((56,), 14, 14, 1)
+
+    def test_three_element_raises(self) -> None:
+        """Three-element tuple must raise ValueError."""
+        with pytest.raises(ValueError, match="shape must be a sequence"):
+            _validate_shape_dims((56, 56, 3), 14, 14, 1)
+
+    def test_scalar_raises(self) -> None:
+        """Bare scalar (not a sequence) must raise ValueError."""
+        with pytest.raises(ValueError, match="shape must be a sequence"):
+            _validate_shape_dims(56, 14, 14, 1)  # type: ignore[arg-type]
+
+
+class TestValidateShapeDimsInvalidDim:
+    """_validate_shape_dims rejects bool, float, and non-positive dimension values."""
+
+    @pytest.mark.parametrize("shape,match", [((True, 56), "height"), ((56, False), "width")])
+    def test_bool_dim_raises(self, shape: tuple, match: str) -> None:
+        """Bool dims must raise ValueError even though bool is an int subtype."""
+        with pytest.raises(ValueError, match=f"{match} must be an integer"):
+            _validate_shape_dims(shape, 14, 14, 1)  # type: ignore[arg-type]
+
+    @pytest.mark.parametrize("shape", [(56.0, 56.0), (56.0, 56), (56, 56.0)])
+    def test_float_dim_raises(self, shape: tuple) -> None:
+        """Float dims must raise ValueError (operator.index rejects them)."""
+        with pytest.raises(ValueError, match="must be an integer"):
+            _validate_shape_dims(shape, 14, 14, 1)
+
+    @pytest.mark.parametrize("shape", [(0, 56), (56, 0), (-14, 56), (56, -14)])
+    def test_non_positive_dim_raises(self, shape: tuple[int, int]) -> None:
+        """Zero or negative dims must raise ValueError."""
+        with pytest.raises(ValueError, match="positive integers"):
+            _validate_shape_dims(shape, 14, 14, 1)
+
+
+class TestValidateShapeDimsDivisibilityCheck:
+    """_validate_shape_dims enforces divisibility by block_size."""
+
+    @pytest.mark.parametrize(
+        "shape, block_size",
+        [
+            pytest.param((55, 56), 14, id="height_not_divisible"),
+            pytest.param((56, 55), 14, id="width_not_divisible"),
+            pytest.param((48, 64), 32, id="height_not_divisible_large_block"),
+        ],
+    )
+    def test_indivisible_shape_raises(self, shape: tuple[int, int], block_size: int) -> None:
+        """Shapes not divisible by block_size must raise ValueError."""
+        with pytest.raises(ValueError, match=f"divisible by {block_size}"):
+            _validate_shape_dims(shape, block_size, 14, 1)
+
+    def test_error_message_includes_patch_size_and_num_windows(self) -> None:
+        """Error message must name patch_size and num_windows for debuggability."""
+        with pytest.raises(ValueError, match="patch_size=16") as exc_info:
+            _validate_shape_dims((48, 64), 32, 16, 2)
+        assert "num_windows=2" in str(exc_info.value)
+
+
+class TestResolvePatchSize:
+    """_resolve_patch_size resolves and validates patch_size for export()/predict()."""
+
+    def _cfg(self, patch_size: int) -> SimpleNamespace:
+        """Return a minimal model_config stub with the given patch_size."""
+        return SimpleNamespace(patch_size=patch_size)
+
+    def test_none_reads_from_model_config(self) -> None:
+        """patch_size=None resolves to model_config.patch_size."""
+        assert _resolve_patch_size(None, self._cfg(16), "export") == 16
+
+    def test_none_falls_back_to_14_when_config_missing(self) -> None:
+        """patch_size=None falls back to 14 when model_config has no patch_size."""
+        assert _resolve_patch_size(None, SimpleNamespace(), "export") == 14
+
+    def test_explicit_matching_config_accepted(self) -> None:
+        """Providing patch_size equal to model_config.patch_size succeeds."""
+        assert _resolve_patch_size(14, self._cfg(14), "export") == 14
+
+    def test_explicit_mismatch_raises(self) -> None:
+        """Providing patch_size != model_config.patch_size must raise ValueError."""
+        with pytest.raises(ValueError, match="does not match"):
+            _resolve_patch_size(16, self._cfg(14), "export")
+
+    def test_mismatch_error_includes_caller_name(self) -> None:
+        """Mismatch error message includes the caller name for context."""
+        with pytest.raises(ValueError, match="predict"):
+            _resolve_patch_size(16, self._cfg(14), "predict")
+
+    @pytest.mark.parametrize("bad", [0, -1, True, False])
+    def test_invalid_explicit_patch_size_raises(self, bad: int) -> None:
+        """Non-positive-int patch_size must raise ValueError before the mismatch check."""
+        cfg = SimpleNamespace(patch_size=bad)
+        with pytest.raises(ValueError, match="patch_size must be a positive integer"):
+            _resolve_patch_size(bad, cfg, "export")
+
+    def test_invalid_config_patch_size_raises(self) -> None:
+        """Bad patch_size in model_config (when caller passes None) must raise ValueError."""
+        with pytest.raises(ValueError, match="patch_size must be a positive integer"):
+            _resolve_patch_size(None, SimpleNamespace(patch_size=0), "export")