chore: import upstream snapshot with attribution

tests/utilities/__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/utilities/test_box_ops.py

@@ -0,0 +1,73 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import torch
+
+from rfdetr.utilities.box_ops import box_iou, generalized_box_iou, masks_to_boxes
+
+
+def test_box_iou_zero_area_boxes_are_finite() -> None:
+    """Zero-area boxes yield finite IoU/union instead of a 0/0 NaN."""
+    zero_box = torch.tensor([[10.0, 10.0, 10.0, 10.0]])  # w = h = 0
+
+    iou, union = box_iou(zero_box, zero_box)
+
+    assert torch.isfinite(iou).all()
+    assert torch.isfinite(union).all()
+
+
+def test_generalized_box_iou_zero_area_boxes_are_finite() -> None:
+    """Degenerate zero-area boxes give finite GIoU instead of NaN/inf."""
+    zero_box = torch.tensor([[10.0, 10.0, 10.0, 10.0]])  # w = h = 0
+
+    giou = generalized_box_iou(zero_box, zero_box)
+
+    assert torch.isfinite(giou).all()
+
+
+def test_masks_to_boxes_passes_ij_indexing_to_meshgrid(monkeypatch) -> None:
+    """`masks_to_boxes` 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)
+
+    masks = torch.zeros((1, 2, 3), dtype=torch.bool)
+    masks[0, 0, 1] = True
+    masks[0, 1, 2] = True
+
+    boxes = masks_to_boxes(masks)
+
+    assert call_count == 1
+    assert boxes.shape == (1, 4)
+
+
+def test_masks_to_boxes_builds_grid_on_masks_device(monkeypatch) -> None:
+    """`masks_to_boxes` should construct arange tensors on the same device as masks."""
+    original_arange = torch.arange
+    observed_devices = []
+
+    def _arange_with_device_capture(*args, **kwargs):
+        observed_devices.append(kwargs.get("device"))
+        return original_arange(*args, **kwargs)
+
+    monkeypatch.setattr(torch, "arange", _arange_with_device_capture)
+
+    masks = torch.zeros((1, 2, 3), dtype=torch.bool)
+    masks[0, 1, 2] = True
+
+    boxes = masks_to_boxes(masks)
+
+    assert boxes.shape == (1, 4)
+    assert observed_devices
+    assert all(device == masks.device for device in observed_devices)

tests/utilities/test_console.py

@@ -0,0 +1,245 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for rfdetr.utilities.console — Rich console helpers for callbacks."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from rfdetr.utilities.console import (
+    _IS_RICH_AVAILABLE,
+    _build_summary_renderable,
+    _get_rich_console,
+    _has_progress_bar,
+    _render_overall_merged,
+    _render_summary_tables,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_trainer(callbacks: list[object] | None = None) -> MagicMock:
+    """Return a minimal mock Trainer."""
+    trainer = MagicMock(name="trainer")
+    trainer.callbacks = callbacks or []
+    return trainer
+
+
+def _minimal_overall(max_dets: int = 500) -> dict:
+    """Return an overall dict with the minimal keys _render_overall_merged expects."""
+    return {
+        "mAP 50:95": 0.4,
+        "mAP 50": 0.6,
+        "mAP 75": 0.3,
+        f"mAR @{max_dets}": 0.5,
+        "F1": 0.55,
+        "Precision": 0.6,
+        "Recall": 0.5,
+    }
+
+
+# ---------------------------------------------------------------------------
+# _render_overall_merged
+# ---------------------------------------------------------------------------
+
+
+class TestRenderOverallMerged:
+    """_render_overall_merged renders a multi-line ASCII table string."""
+
+    def test_returns_string(self) -> None:
+        """Result is a non-empty multi-line string."""
+        result = _render_overall_merged("Val", _minimal_overall(), 500)
+        assert isinstance(result, str)
+        assert "\n" in result
+
+    def test_title_prefix_in_output(self) -> None:
+        """Title prefix appears in the rendered string."""
+        result = _render_overall_merged("Test", _minimal_overall(), 500)
+        assert "Test" in result
+
+    def test_metric_values_in_output(self) -> None:
+        """Formatted metric values appear in the output."""
+        result = _render_overall_merged("Val", _minimal_overall(500), 500)
+        assert "0.4000" in result
+
+    def test_nan_renders_as_em_dash(self) -> None:
+        """NaN values render as '—' (em-dash)."""
+        overall = _minimal_overall()
+        overall["mAP 50:95"] = float("nan")
+        result = _render_overall_merged("Val", overall, 500)
+        assert "—" in result
+
+    def test_negative_sentinel_renders_as_em_dash(self) -> None:
+        """Pycocotools sentinel -1 renders as '—'."""
+        overall = _minimal_overall()
+        overall["mAP 50"] = -1.0
+        result = _render_overall_merged("Val", overall, 500)
+        assert "—" in result
+
+    def test_segm_group_present_when_key_exists(self) -> None:
+        """Segm mAP group rendered when segm keys present."""
+        overall = _minimal_overall()
+        overall["segm mAP 50:95"] = 0.3
+        overall["segm mAP 50"] = 0.5
+        result = _render_overall_merged("Val", overall, 500)
+        assert "segm mAP" in result
+
+    def test_segm_group_absent_when_key_missing(self) -> None:
+        """Segm mAP group not rendered when keys absent."""
+        result = _render_overall_merged("Val", _minimal_overall(), 500)
+        assert "segm mAP" not in result
+
+    def test_mar_label_uses_max_dets(self) -> None:
+        """MAR column label contains the max_dets value."""
+        result = _render_overall_merged("Val", _minimal_overall(100), 100)
+        assert "@100" in result
+
+
+# ---------------------------------------------------------------------------
+# _has_progress_bar
+# ---------------------------------------------------------------------------
+
+
+class TestHasProgressBar:
+    """_has_progress_bar detects any callback whose class name ends with ProgressBar."""
+
+    def test_returns_false_with_no_callbacks(self) -> None:
+        """Returns False when trainer has no callbacks."""
+        assert not _has_progress_bar(_make_trainer())
+
+    def test_returns_true_for_tqdm_progress_bar(self) -> None:
+        """Returns True when a TQDMProgressBar callback present."""
+        tqdm_bar_cls = type("TQDMProgressBar", (), {})
+        trainer = _make_trainer(callbacks=[tqdm_bar_cls()])
+        assert _has_progress_bar(trainer)
+
+    def test_returns_true_for_rich_progress_bar(self) -> None:
+        """Returns True when a RichProgressBar callback present."""
+        rich_bar_cls = type("RichProgressBar", (), {})
+        trainer = _make_trainer(callbacks=[rich_bar_cls()])
+        assert _has_progress_bar(trainer)
+
+    def test_returns_false_for_non_progress_bar_callback(self) -> None:
+        """Returns False when callbacks don't end with ProgressBar."""
+        trainer = _make_trainer(callbacks=[MagicMock(name="SomeOtherCallback")])
+        assert not _has_progress_bar(trainer)
+
+
+# ---------------------------------------------------------------------------
+# _get_rich_console
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _IS_RICH_AVAILABLE, reason="Rich not installed")
+class TestGetRichConsole:
+    """_get_rich_console returns the PTL RichProgressBar console or a fresh one."""
+
+    def test_returns_fresh_console_with_no_callbacks(self) -> None:
+        """Returns a Console instance when no RichProgressBar present."""
+        from rich.console import Console
+
+        result = _get_rich_console(_make_trainer())
+        assert isinstance(result, Console)
+
+    def test_returns_rich_progress_bar_console_when_active(self) -> None:
+        """Returns _console from RichProgressBar when callback present and _console set."""
+        rich_bar_cls = type("RichProgressBar", (), {})
+        expected_console = MagicMock(name="expected_console")
+        cb = rich_bar_cls()
+        cb._console = expected_console  # type: ignore[attr-defined]
+        trainer = _make_trainer(callbacks=[cb])
+
+        result = _get_rich_console(trainer)
+
+        assert result is expected_console
+
+    def test_falls_back_when_console_attribute_is_none(self) -> None:
+        """Falls back to fresh Console when _console is None (outside active stage)."""
+        from rich.console import Console
+
+        rich_bar_cls = type("RichProgressBar", (), {})
+        cb = rich_bar_cls()
+        cb._console = None  # type: ignore[attr-defined]
+        trainer = _make_trainer(callbacks=[cb])
+
+        result = _get_rich_console(trainer)
+
+        assert isinstance(result, Console)
+
+    def test_mro_subclass_detected(self) -> None:
+        """Subclass of RichProgressBar is detected via MRO name check."""
+        rich_bar_cls = type("RichProgressBar", (), {})
+        themed_bar_cls = type("ThemedProgressBar", (rich_bar_cls,), {})
+        expected_console = MagicMock(name="expected_console")
+        cb = themed_bar_cls()
+        cb._console = expected_console  # type: ignore[attr-defined]
+        trainer = _make_trainer(callbacks=[cb])
+
+        result = _get_rich_console(trainer)
+
+        assert result is expected_console
+
+
+# ---------------------------------------------------------------------------
+# _build_summary_renderable
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _IS_RICH_AVAILABLE, reason="Rich not installed")
+class TestBuildSummaryRenderable:
+    """_build_summary_renderable returns a Rich Group for console.print()."""
+
+    def test_returns_group_without_per_class(self) -> None:
+        """Returns a Group renderable when per_class is empty."""
+        from rich.console import Group
+
+        result = _build_summary_renderable("Val", "overall-text", [])
+        assert isinstance(result, Group)
+
+    def test_returns_group_with_per_class(self) -> None:
+        """Returns a Group renderable when per_class rows present."""
+        from rich.console import Group
+
+        per_class = [{"name": "cat", "ap": 0.5, "ar": 0.6, "f1": 0.55, "precision": 0.6, "recall": 0.5}]
+        result = _build_summary_renderable("Val", "overall-text", per_class)
+        assert isinstance(result, Group)
+
+    def test_nan_per_class_renders_as_em_dash(self) -> None:
+        """NaN in per-class metric renders without error."""
+        from rich.console import Console, Group
+
+        per_class = [{"name": "cat", "ap": float("nan"), "ar": -1.0, "f1": 0.0, "precision": 0.0, "recall": 0.0}]
+        result = _build_summary_renderable("Val", "overall-text", per_class)
+        assert isinstance(result, Group)
+        console = Console(force_terminal=True)
+        with console.capture() as capture:
+            console.print(result)
+        assert "—" in capture.get()
+
+
+# ---------------------------------------------------------------------------
+# _render_summary_tables
+# ---------------------------------------------------------------------------
+
+
+class TestRenderSummaryTables:
+    """_render_summary_tables delegates to console.print when Rich available."""
+
+    @pytest.mark.skipif(not _IS_RICH_AVAILABLE, reason="Rich not installed")
+    def test_calls_console_print_once(self) -> None:
+        """console.print called exactly once with the Group renderable."""
+        console = MagicMock(name="console")
+        _render_summary_tables(console, "Val", "overall-text", [])
+        console.print.assert_called_once()
+
+    def test_no_op_when_rich_unavailable(self) -> None:
+        """No call made when Rich not installed."""
+        console = MagicMock(name="console")
+        with patch("rfdetr.utilities.console._IS_RICH_AVAILABLE", False):
+            _render_summary_tables(console, "Val", "overall-text", [])
+        console.print.assert_not_called()

tests/utilities/test_distributed.py

@@ -0,0 +1,27 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for distributed utility helpers."""
+
+from unittest.mock import patch
+
+from rfdetr.utilities.distributed import all_gather
+
+
+def test_all_gather_supports_cpu_without_tensor_truthiness_error() -> None:
+    """all_gather should work on CPU-only setups and return gathered objects."""
+
+    def _fake_all_gather(output_tensors, input_tensor) -> None:
+        for out in output_tensors:
+            out.copy_(input_tensor)
+
+    with (
+        patch("rfdetr.utilities.distributed.get_world_size", return_value=2),
+        patch("rfdetr.utilities.distributed.dist.all_gather", side_effect=_fake_all_gather),
+        patch("rfdetr.utilities.distributed.torch.cuda.is_available", return_value=False),
+    ):
+        result = all_gather({"value": 7})
+
+    assert result == [{"value": 7}, {"value": 7}]

tests/utilities/test_file_validation.py

@@ -0,0 +1,284 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import os
+import tempfile
+from pathlib import Path
+from typing import Iterable, Iterator, Literal, Optional
+from unittest.mock import Mock, patch
+
+import pytest
+import requests
+
+from rfdetr.utilities.files import _compute_file_md5, _download_file, _validate_file_md5
+
+
+class _DummyTqdm:
+    """Minimal tqdm stand-in for download tests.
+
+    This avoids real progress bars while preserving the context manager and `update` calls used by the downloader.
+    """
+
+    def __init__(self, **kwargs: object) -> None:
+        """Store initialization kwargs for optional inspection."""
+        self.kwargs = kwargs
+
+    def __enter__(self) -> "_DummyTqdm":
+        """Return self to satisfy context manager protocol."""
+        return self
+
+    def __exit__(self, exc_type: object, exc: object, tb: object) -> bool:
+        """Propagate exceptions raised inside the context."""
+        return False
+
+    def update(self, size: int) -> None:
+        """No-op progress update for compatibility with tqdm."""
+        return None
+
+
+class _FakeResponse:
+    """Test double for requests responses used by the downloader.
+
+    Provides headers, iterable content chunks, and optional HTTP error behavior via `raise_for_status`.
+    """
+
+    def __init__(
+        self,
+        content_chunks: Iterable[bytes],
+        headers: Optional[dict[str, str]] = None,
+        raise_error: Optional[Exception] = None,
+    ) -> None:
+        """Initialize the fake response with content and metadata."""
+        self._content_chunks = list(content_chunks)
+        self.headers = headers or {}
+        self._raise_error = raise_error
+
+    def raise_for_status(self) -> None:
+        """Raise the configured HTTP error, if any."""
+        if self._raise_error is not None:
+            raise self._raise_error
+
+    def iter_content(self, chunk_size: int = 1024) -> Iterator[bytes]:
+        """Yield the configured content chunks."""
+        for chunk in self._content_chunks:
+            yield chunk
+
+
+def _assert_no_download_temp_files(tmp_path: Path, filename: str = "weights.bin") -> None:
+    """Assert that randomized temporary download files were cleaned up."""
+    assert not list(tmp_path.glob(f"{filename}.*.tmp"))
+
+
+class TestFileMD5Validation:
+    """Test MD5 hash computation and validation."""
+
+    def test_compute_file_md5(self):
+        """Test MD5 hash computation for a simple file."""
+        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+            f.write("Hello, World!")
+            temp_file = f.name
+
+        try:
+            # Known MD5 hash for "Hello, World!"
+            expected_hash = "65a8e27d8879283831b664bd8b7f0ad4"
+            actual_hash = _compute_file_md5(temp_file)
+            assert actual_hash == expected_hash
+        finally:
+            os.unlink(temp_file)
+
+    def test_validate_file_md5_success(self):
+        """Test successful MD5 validation."""
+        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+            f.write("Test content")
+            temp_file = f.name
+
+        try:
+            # Compute the actual hash first
+            expected_hash = _compute_file_md5(temp_file)
+
+            # Validation should succeed
+            assert _validate_file_md5(temp_file, expected_hash) is True
+        finally:
+            os.unlink(temp_file)
+
+    def test_validate_file_md5_failure(self):
+        """Test MD5 validation failure with wrong hash."""
+        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+            f.write("Test content")
+            temp_file = f.name
+
+        try:
+            # Use a wrong hash
+            wrong_hash = "0" * 32
+
+            # Validation should fail
+            assert _validate_file_md5(temp_file, wrong_hash) is False
+        finally:
+            os.unlink(temp_file)
+
+    @pytest.mark.parametrize(
+        "hash_case", [pytest.param("lower", id="lowercase"), pytest.param("upper", id="uppercase")]
+    )
+    def test_validate_file_md5_case_insensitive(self, hash_case: Literal["lower", "upper"]) -> None:
+        """Test that MD5 validation is case-insensitive."""
+        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+            f.write("Test content")
+            temp_file = f.name
+
+        try:
+            # Get hash in lowercase
+            hash_lower = _compute_file_md5(temp_file)
+            test_hash = hash_lower.upper() if hash_case == "upper" else hash_lower
+
+            # Each case variant (lower/upper) should validate successfully
+            assert _validate_file_md5(temp_file, test_hash) is True
+        finally:
+            os.unlink(temp_file)
+
+    def test_validate_nonexistent_file(self):
+        """Test validation of non-existent file."""
+        nonexistent_file = "/tmp/nonexistent_file_xyz.txt"
+        assert _validate_file_md5(nonexistent_file, "abc123") is False
+
+    def test_compute_file_md5_empty_file(self):
+        """Test MD5 hash computation for empty file."""
+        with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+            # Create empty file
+            temp_file = f.name
+
+        try:
+            # Known MD5 hash for empty file
+            expected_hash = "d41d8cd98f00b204e9800998ecf8427e"
+            actual_hash = _compute_file_md5(temp_file)
+            assert actual_hash == expected_hash
+        finally:
+            os.unlink(temp_file)
+
+    def test_compute_file_md5_large_file(self):
+        """Test MD5 computation for larger file (tests chunking)."""
+        with tempfile.NamedTemporaryFile(mode="wb", delete=False) as f:
+            # Write 1MB of data
+            data = b"A" * (1024 * 1024)
+            f.write(data)
+            temp_file = f.name
+
+        try:
+            # Compute hash (should handle chunking correctly)
+            hash_value = _compute_file_md5(temp_file)
+
+            # Verify it's a valid MD5 hash format
+            assert len(hash_value) == 32
+            assert all(c in "0123456789abcdef" for c in hash_value)
+        finally:
+            os.unlink(temp_file)
+
+
+class TestDownloadFile:
+    """Test download helper behavior and failure cleanup."""
+
+    @patch("rfdetr.utilities.files.tqdm", _DummyTqdm)
+    @patch("rfdetr.utilities.files.requests.get")
+    def test_download_file_missing_content_length(self, mock_get: Mock, tmp_path: Path):
+        """Download succeeds when content-length is missing."""
+        target_path = tmp_path / "weights.bin"
+        response = _FakeResponse([b"hello", b"world"], headers={})
+        mock_get.return_value = response
+
+        _download_file("https://example.com/file.bin", str(target_path))
+
+        assert target_path.exists()
+        assert target_path.read_bytes() == b"helloworld"
+        _assert_no_download_temp_files(tmp_path)
+        mock_get.assert_called_once_with("https://example.com/file.bin", stream=True, timeout=30.0)
+
+    @patch("rfdetr.utilities.files.requests.get")
+    def test_download_file_http_error(self, mock_get: Mock, tmp_path: Path):
+        """HTTP errors raise and do not create files."""
+        target_path = tmp_path / "weights.bin"
+        response = _FakeResponse([], raise_error=requests.HTTPError("bad request"))
+        mock_get.return_value = response
+
+        with pytest.raises(requests.HTTPError):
+            _download_file("https://example.com/file.bin", str(target_path))
+
+        assert not target_path.exists()
+        _assert_no_download_temp_files(tmp_path)
+
+    @patch("rfdetr.utilities.files.tqdm", _DummyTqdm)
+    @patch("rfdetr.utilities.files.requests.get")
+    def test_download_file_stream_error_cleans_temp(self, mock_get: Mock, tmp_path: Path):
+        """Streaming errors clean up temp files."""
+        target_path = tmp_path / "weights.bin"
+
+        class _StreamErrorResponse(_FakeResponse):
+            def iter_content(self, chunk_size: int = 1024) -> Iterator[bytes]:
+                yield b"partial"
+                raise RuntimeError("stream failure")
+
+        response = _StreamErrorResponse([b"partial"], headers={"content-length": "7"})
+        mock_get.return_value = response
+
+        with pytest.raises(RuntimeError):
+            _download_file("https://example.com/file.bin", str(target_path))
+
+        assert not target_path.exists()
+        _assert_no_download_temp_files(tmp_path)
+
+    @patch("rfdetr.utilities.files.tqdm", _DummyTqdm)
+    @patch("rfdetr.utilities.files.requests.get")
+    def test_download_file_md5_failure_cleans_temp(self, mock_get: Mock, tmp_path: Path):
+        """MD5 failure removes temp file and target is not created."""
+        target_path = tmp_path / "weights.bin"
+        response = _FakeResponse([b"data"], headers={"content-length": "4"})
+        mock_get.return_value = response
+
+        with pytest.raises(ValueError):
+            _download_file(
+                "https://example.com/file.bin",
+                str(target_path),
+                expected_md5="0" * 32,
+            )
+
+        assert not target_path.exists()
+        _assert_no_download_temp_files(tmp_path)
+
+    @patch("rfdetr.utilities.files.tqdm", _DummyTqdm)
+    @patch("rfdetr.utilities.files.requests.get")
+    def test_download_file_replace_failure_cleans_temp(self, mock_get: Mock, tmp_path: Path):
+        """Replace failure removes temp file and target is not created."""
+        target_path = tmp_path / "weights.bin"
+        response = _FakeResponse([b"data"], headers={"content-length": "4"})
+        mock_get.return_value = response
+
+        with (
+            patch("rfdetr.utilities.files.os.replace", side_effect=PermissionError("replace denied")),
+            pytest.raises(PermissionError, match="replace denied"),
+        ):
+            _download_file("https://example.com/file.bin", str(target_path))
+
+        assert not target_path.exists()
+        _assert_no_download_temp_files(tmp_path)
+
+    @patch("rfdetr.utilities.files.tqdm", _DummyTqdm)
+    @patch("rfdetr.utilities.files.open", side_effect=AssertionError("must not reopen temp filename"))
+    @patch("rfdetr.utilities.files.requests.get")
+    def test_download_file_uses_fdopen_without_reopen(
+        self,
+        mock_get: Mock,
+        mock_open: Mock,
+        tmp_path: Path,
+    ) -> None:
+        """Download writes through mkstemp file descriptor to avoid TOCTOU reopen."""
+        target_path = tmp_path / "weights.bin"
+        response = _FakeResponse([b"data"], headers={"content-length": "4"})
+        mock_get.return_value = response
+
+        _download_file("https://example.com/file.bin", str(target_path))
+
+        assert target_path.exists()
+        assert target_path.read_bytes() == b"data"
+        _assert_no_download_temp_files(tmp_path)
+        mock_open.assert_not_called()

tests/utilities/test_keypoints.py

@@ -0,0 +1,163 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for keypoint utility functions in rfdetr.utilities.keypoints."""
+
+import numpy as np
+import pytest
+
+from rfdetr.utilities.keypoints import (
+    _is_bg_first_schema,
+    _to_active_first,
+    _to_bg_first,
+    precision_cholesky_to_pixel_covariance,
+    schemas_semantically_equal,
+)
+
+
+class TestIsBgFirstSchema:
+    """Group: is_bg_first_schema — schema classification."""
+
+    @pytest.mark.parametrize(
+        ("schema", "expected"),
+        [
+            pytest.param([0, 17], True, id="bg-first-single-class"),
+            pytest.param([0, 17, 4], True, id="bg-first-multi-class"),
+            pytest.param([0], True, id="bg-only-slot"),
+            pytest.param([17], False, id="active-first-single"),
+            pytest.param([17, 4], False, id="active-first-multi"),
+            pytest.param([], False, id="empty-schema"),
+        ],
+    )
+    def test_classification(self, schema: list[int], expected: bool) -> None:
+        """is_bg_first_schema returns expected bool for each schema form."""
+        assert _is_bg_first_schema(schema) == expected
+
+
+class TestToActiveFirst:
+    """Group: to_active_first — strip leading background slot."""
+
+    @pytest.mark.parametrize(
+        ("schema", "expected"),
+        [
+            pytest.param([0, 17], [17], id="bg-first-to-active"),
+            pytest.param([0, 17, 4], [17, 4], id="bg-first-multi-class"),
+            pytest.param([0], [], id="bg-only-to-empty"),
+            pytest.param([17], [17], id="already-active-first"),
+            pytest.param([17, 4], [17, 4], id="already-active-multi"),
+            pytest.param([], [], id="empty-schema"),
+            pytest.param([0, 0, 17], [0, 17], id="multi-leading-zero-strips-one"),
+        ],
+    )
+    def test_conversion(self, schema: list[int], expected: list[int]) -> None:
+        """to_active_first strips only the first leading zero slot."""
+        assert _to_active_first(schema) == expected
+
+    def test_returns_new_list(self) -> None:
+        """to_active_first always returns a new list, never the input object."""
+        schema = [17]
+        result = _to_active_first(schema)
+        assert result is not schema
+
+
+class TestToBgFirst:
+    """Group: to_bg_first — prepend background slot."""
+
+    @pytest.mark.parametrize(
+        ("schema", "expected"),
+        [
+            pytest.param([17], [0, 17], id="active-first-to-bg"),
+            pytest.param([17, 4], [0, 17, 4], id="active-first-multi"),
+            pytest.param([0, 17], [0, 17], id="already-bg-first-no-op"),
+            pytest.param([], [], id="empty-schema-no-op"),
+        ],
+    )
+    def test_conversion(self, schema: list[int], expected: list[int]) -> None:
+        """to_bg_first prepends 0 only when schema is active-first and non-empty."""
+        assert _to_bg_first(schema) == expected
+
+    def test_returns_new_list(self) -> None:
+        """to_bg_first always returns a new list, never the input object."""
+        schema = [0, 17]
+        result = _to_bg_first(schema)
+        assert result is not schema
+
+
+class TestSchemasSemanticallyEqual:
+    """Group: schemas_semantically_equal — cross-form equality."""
+
+    @pytest.mark.parametrize(
+        ("a", "b", "expected"),
+        [
+            pytest.param([0, 17], [17], True, id="bg-first-eq-active-first"),
+            pytest.param([17], [17], True, id="identical-active-first"),
+            pytest.param([0, 17], [0, 17], True, id="identical-bg-first"),
+            pytest.param([0, 17], [0, 33], False, id="different-keypoint-counts"),
+            pytest.param([17], [33], False, id="active-first-mismatch"),
+            pytest.param([0], [], True, id="bg-only-eq-empty"),
+            pytest.param([], [], True, id="empty-eq-empty"),
+            pytest.param([0, 17, 4], [17, 4], True, id="multi-class-cross-form"),
+        ],
+    )
+    def test_equality(self, a: list[int], b: list[int], expected: bool) -> None:
+        """schemas_semantically_equal returns expected result for each pair."""
+        assert schemas_semantically_equal(a, b) == expected
+
+    def test_symmetric(self) -> None:
+        """schemas_semantically_equal(a, b) == schemas_semantically_equal(b, a)."""
+        assert schemas_semantically_equal([0, 17], [17]) == schemas_semantically_equal([17], [0, 17])
+
+
+class TestPrecisionCholeskyToPixelCovariance:
+    """Group: precision_cholesky_to_pixel_covariance — non-finite input handling."""
+
+    def test_nan_in_single_slot_produces_nan_only_in_that_slot(self) -> None:
+        """NaN params in one detection slot should propagate NaN only to that slot's output."""
+        # N=2, K=1: first slot valid, second slot has NaN in all three params.
+        params = np.array(
+            [[[0.0, 0.0, 0.0]], [[np.nan, 0.0, 0.0]]],
+            dtype=np.float32,
+        )
+        source_shape = np.array([[10.0, 20.0], [10.0, 20.0]], dtype=np.float32)
+
+        covariance = precision_cholesky_to_pixel_covariance(
+            precision_cholesky=params,
+            source_shape=source_shape,
+        )
+
+        # First slot (valid) should be all-finite.
+        assert np.isfinite(covariance[0, 0]).all(), f"First slot expected all-finite, got {covariance[0, 0]}"
+        # Second slot (NaN input) should be all-NaN.
+        assert np.isnan(covariance[1, 0]).all(), f"Second slot expected all-NaN, got {covariance[1, 0]}"
+
+    def test_all_inf_params_produce_all_nan_covariance(self) -> None:
+        """Infinite precision params should produce all-NaN pixel covariances."""
+        params = np.full((1, 1, 3), np.inf, dtype=np.float32)
+        source_shape = np.array([[10.0, 20.0]], dtype=np.float32)
+
+        covariance = precision_cholesky_to_pixel_covariance(
+            precision_cholesky=params,
+            source_shape=source_shape,
+        )
+
+        assert np.isnan(covariance).all(), f"Expected all-NaN output for all-inf inputs, got {covariance}"
+
+    def test_mixed_valid_and_nan_rows_isolates_nan_to_bad_row(self) -> None:
+        """First detection valid, second detection NaN — only second row should be NaN."""
+        params = np.array(
+            [[[0.0, 0.0, 0.0]], [[np.nan, np.nan, np.nan]]],
+            dtype=np.float32,
+        )
+        source_shape = np.array([[10.0, 20.0], [5.0, 8.0]], dtype=np.float32)
+
+        covariance = precision_cholesky_to_pixel_covariance(
+            precision_cholesky=params,
+            source_shape=source_shape,
+        )
+
+        # Row 0 — valid identity input, covariance should be finite.
+        assert np.isfinite(covariance[0]).all(), f"Row 0 expected all-finite, got {covariance[0]}"
+        # Row 1 — NaN input, covariance should be all-NaN.
+        assert np.isnan(covariance[1]).all(), f"Row 1 expected all-NaN, got {covariance[1]}"

tests/utilities/test_misc.py

@@ -0,0 +1,90 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from rfdetr.utilities.state_dict import strip_checkpoint
+
+
+class TestStripCheckpoint:
+    def test_strip_checkpoint_keeps_only_model_and_args(self, tmp_path):
+        checkpoint_path = tmp_path / "checkpoint_best_total.pth"
+        torch.save(
+            {
+                "model": {"weight": torch.tensor([1.0])},
+                "args": SimpleNamespace(class_names=["a"]),
+                "optimizer": {"lr": 1e-4},
+            },
+            checkpoint_path,
+        )
+
+        strip_checkpoint(str(checkpoint_path))
+
+        stripped = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+        assert set(stripped.keys()) == {"model", "args"}
+
+    def test_strip_checkpoint_preserves_model_name_when_present(self, tmp_path):
+        checkpoint_path = tmp_path / "checkpoint_best_total.pth"
+        torch.save(
+            {
+                "model": {"weight": torch.tensor([1.0])},
+                "args": SimpleNamespace(class_names=["a"]),
+                "model_name": "RFDETRSmall",
+                "optimizer": {"lr": 1e-4},
+            },
+            checkpoint_path,
+        )
+
+        strip_checkpoint(str(checkpoint_path))
+
+        stripped = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+        assert set(stripped.keys()) == {"model", "args", "model_name"}
+        assert stripped["model_name"] == "RFDETRSmall"
+
+    def test_strip_checkpoint_omits_model_name_when_absent(self, tmp_path):
+        checkpoint_path = tmp_path / "checkpoint_best_total.pth"
+        torch.save(
+            {
+                "model": {"weight": torch.tensor([1.0])},
+                "args": SimpleNamespace(class_names=["a"]),
+                "optimizer": {"lr": 1e-4},
+            },
+            checkpoint_path,
+        )
+
+        strip_checkpoint(str(checkpoint_path))
+
+        stripped = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+        assert "model_name" not in stripped
+
+    def test_strip_checkpoint_is_atomic_when_save_fails(self, tmp_path, monkeypatch):
+        checkpoint_path = tmp_path / "checkpoint_best_total.pth"
+        original_checkpoint = {
+            "model": {"weight": torch.tensor([1.0])},
+            "args": SimpleNamespace(class_names=["a"]),
+            "optimizer": {"lr": 1e-4},
+        }
+        torch.save(original_checkpoint, checkpoint_path)
+
+        original_torch_save = torch.save
+
+        def failing_torch_save(obj, destination, *args, **kwargs):
+            if str(destination) != str(checkpoint_path):
+                raise RuntimeError("simulated save failure")
+            return original_torch_save(obj, destination, *args, **kwargs)
+
+        monkeypatch.setattr(torch, "save", failing_torch_save)
+
+        with pytest.raises(RuntimeError, match="simulated save failure"):
+            strip_checkpoint(str(checkpoint_path))
+
+        recovered = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+        assert set(recovered.keys()) == set(original_checkpoint.keys())
+        assert recovered["model"]["weight"].equal(original_checkpoint["model"]["weight"])
+        assert recovered["optimizer"] == original_checkpoint["optimizer"]

tests/utilities/test_package.py

@@ -0,0 +1,269 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for package metadata helpers and structural import paths."""
+
+import subprocess
+import sys
+from unittest.mock import patch
+
+import pytest
+
+from rfdetr.utilities.package import get_sha
+
+
+def test_get_sha_marks_dirty_worktree_when_diff_command_returns_exit_code_1() -> None:
+    """A diff exit code of 1 should report uncommitted changes, not unknown."""
+
+    def _fake_check_output(command: list[str], cwd: str | None = None) -> bytes:
+        if command[:3] == ["git", "rev-parse", "HEAD"]:
+            return b"abc123\n"
+        if command[:4] == ["git", "rev-parse", "--abbrev-ref", "HEAD"]:
+            return b"feature/test\n"
+        raise AssertionError(f"Unexpected command: {command!r}")
+
+    class _RunResult:
+        def __init__(self, returncode: int) -> None:
+            self.returncode = returncode
+
+    with (
+        patch("rfdetr.utilities.package.subprocess.check_output", side_effect=_fake_check_output),
+        patch("rfdetr.utilities.package.subprocess.run", return_value=_RunResult(returncode=1)),
+    ):
+        sha = get_sha()
+
+    assert sha == "sha: abc123, status: has uncommitted changes, branch: feature/test"
+
+
+def test_peft_not_imported_eagerly_on_backbone_import_characterization() -> None:
+    """Importing backbone.backbone must NOT pull peft into sys.modules (peft is optional).
+
+    This characterization test captures the invariant introduced in PR 1 (chore/packaging-peft-lora): after the lazy-
+    import refactor, importing backbone at module-load time must not trigger a top-level ``from peft import PeftModel``.
+    """
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            (
+                "import sys; "
+                "import rfdetr.models.backbone.backbone; "
+                "assert 'peft' not in sys.modules, "
+                "'peft was eagerly imported by backbone.backbone'"
+            ),
+        ],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert result.returncode == 0, (
+        "Subprocess for backbone import failed:\n"
+        f"return code: {result.returncode}\n"
+        f"stdout:\n{result.stdout}\n"
+        f"stderr:\n{result.stderr}"
+    )
+
+
+class TestImportPaths:
+    """Structural tests for the detr.py → inference.py + variants.py split (PR 6).
+
+    After the split:
+    - ``rfdetr.inference`` exports ``ModelContext`` and ``_build_model_context``
+    - ``rfdetr.variants`` exports all 15 concrete model classes
+    - ``rfdetr.detr`` re-exports both for backward compatibility
+    - ``rfdetr`` (top-level) continues to export public names unchanged
+    """
+
+    def test_model_context_importable_from_inference(self) -> None:
+        """ModelContext must be importable from the new rfdetr.inference module."""
+        from rfdetr.inference import ModelContext
+
+        assert ModelContext is not None
+
+    def test_build_model_context_importable_from_inference(self) -> None:
+        """_build_model_context must be importable from rfdetr.inference."""
+        from rfdetr.inference import _build_model_context
+
+        assert callable(_build_model_context)
+
+    def test_rfdetr_large_importable_from_variants(self) -> None:
+        """RFDETRLarge must be importable from the new rfdetr.variants module."""
+        from rfdetr.variants import RFDETRLarge
+
+        assert RFDETRLarge is not None
+
+    def test_variants_import_first_does_not_trigger_circular_import(self) -> None:
+        """Importing variants before detr must still preserve shared class identity."""
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                (
+                    "from rfdetr.variants import RFDETRLarge; "
+                    "from rfdetr.detr import RFDETRLarge as FromDetr; "
+                    "assert RFDETRLarge is FromDetr"
+                ),
+            ],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        assert result.returncode == 0, (
+            "Subprocess for variants-first import failed:\n"
+            f"return code: {result.returncode}\n"
+            f"stdout:\n{result.stdout}\n"
+            f"stderr:\n{result.stderr}"
+        )
+
+    @pytest.mark.parametrize(
+        "class_name",
+        [
+            pytest.param("RFDETRBase", id="base"),
+            pytest.param("RFDETRKeypointPreview", id="keypoint-preview"),
+            pytest.param("RFDETRNano", id="nano"),
+            pytest.param("RFDETRSmall", id="small"),
+            pytest.param("RFDETRMedium", id="medium"),
+            pytest.param("RFDETRLargeDeprecated", id="large-deprecated"),
+            pytest.param("RFDETRLarge", id="large"),
+            pytest.param("RFDETRSeg", id="seg-base"),
+            pytest.param("RFDETRSegPreview", id="seg-preview"),
+            pytest.param("RFDETRSegNano", id="seg-nano"),
+            pytest.param("RFDETRSegSmall", id="seg-small"),
+            pytest.param("RFDETRSegMedium", id="seg-medium"),
+            pytest.param("RFDETRSegLarge", id="seg-large"),
+            pytest.param("RFDETRSegXLarge", id="seg-xlarge"),
+            pytest.param("RFDETRSeg2XLarge", id="seg-2xlarge"),
+        ],
+    )
+    def test_all_variant_classes_importable_from_variants(self, class_name: str) -> None:
+        """Every concrete variant class must be importable from rfdetr.variants."""
+        import rfdetr.variants as variants_mod
+
+        cls = getattr(variants_mod, class_name, None)
+        assert cls is not None, f"{class_name} not found in rfdetr.variants"
+
+    def test_model_context_backward_compat_from_detr(self) -> None:
+        """ModelContext must remain importable from rfdetr.detr (backward compat)."""
+        from rfdetr.detr import ModelContext
+
+        assert ModelContext is not None
+
+    def test_rfdetr_large_backward_compat_from_detr(self) -> None:
+        """RFDETRLarge must remain importable from rfdetr.detr (backward compat)."""
+        from rfdetr.detr import RFDETRLarge
+
+        assert RFDETRLarge is not None
+
+    def test_rfdetr_large_importable_from_top_level(self) -> None:
+        """RFDETRLarge must remain importable from rfdetr (top-level package)."""
+        from rfdetr import RFDETRLarge
+
+        assert RFDETRLarge is not None
+
+    def test_model_context_importable_from_top_level(self) -> None:
+        """ModelContext must remain importable from rfdetr (top-level package)."""
+        from rfdetr import ModelContext
+
+        assert ModelContext is not None
+
+    def test_top_level_import_sets_numpy_complex_alias(self) -> None:
+        """Verify rfdetr shim sets np.complex_ in a fresh interpreter with complex_ absent.
+
+        Runs in a subprocess so the shim code path is actually executed — importing rfdetr inside pytest is a no-op
+        because rfdetr is already cached in sys.modules.
+        """
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                (
+                    "import numpy\n"
+                    "if hasattr(numpy, 'complex_'):\n"
+                    "    del numpy.complex_\n"
+                    "import rfdetr\n"
+                    "assert hasattr(numpy, 'complex_'), 'shim did not set numpy.complex_'\n"
+                    "assert numpy.complex_ is numpy.complex128, 'complex_ must alias complex128'\n"
+                ),
+            ],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        assert result.returncode == 0, (
+            "Subprocess for NumPy complex_ shim failed:\n"
+            f"return code: {result.returncode}\n"
+            f"stdout:\n{result.stdout}\n"
+            f"stderr:\n{result.stderr}"
+        )
+
+    def test_identity_across_import_paths(self) -> None:
+        """The same class object must be returned regardless of import path.
+
+        This ensures re-exports are true re-exports (not copies) so that isinstance() checks work across all import
+        paths.
+        """
+        import rfdetr
+        from rfdetr.detr import ModelContext as FromDetr
+        from rfdetr.detr import RFDETRLarge as LargeFromDetr
+        from rfdetr.inference import ModelContext as FromInference
+        from rfdetr.variants import RFDETRLarge as LargeFromVariants
+
+        assert FromDetr is FromInference, (
+            "rfdetr.detr.ModelContext and rfdetr.inference.ModelContext must be the same object"
+        )
+        assert LargeFromDetr is LargeFromVariants, (
+            "rfdetr.detr.RFDETRLarge and rfdetr.variants.RFDETRLarge must be the same object"
+        )
+        assert rfdetr.RFDETRLarge is LargeFromVariants, (
+            "top-level rfdetr.RFDETRLarge and rfdetr.variants.RFDETRLarge must be the same object"
+        )
+        assert rfdetr.ModelContext is FromInference, (
+            "top-level rfdetr.ModelContext and rfdetr.inference.ModelContext must be the same object"
+        )
+
+    def test_build_model_context_backward_compat_from_detr(self) -> None:
+        """_build_model_context must remain importable from rfdetr.detr (backward compat)."""
+        from rfdetr.detr import _build_model_context
+
+        assert callable(_build_model_context)
+
+    def test_getattr_raises_for_unknown_names(self) -> None:
+        """Accessing an unknown name via rfdetr.detr must raise AttributeError."""
+        import rfdetr.detr as detr_mod
+
+        with pytest.raises(AttributeError, match="has no attribute"):
+            _ = detr_mod._this_name_does_not_exist_12345
+
+    def test_dir_includes_lazy_exports(self) -> None:
+        """dir(rfdetr.detr) must include all lazy re-export names."""
+        import rfdetr.detr as detr_mod
+
+        names = dir(detr_mod)
+        assert "ModelContext" in names, "ModelContext missing from dir(rfdetr.detr)"
+        assert "RFDETRLarge" in names, "RFDETRLarge missing from dir(rfdetr.detr)"
+        assert "RFDETRBase" in names, "RFDETRBase missing from dir(rfdetr.detr)"
+
+    def test_detr_first_import_order_preserves_identity(self) -> None:
+        """Importing detr before variants must preserve object identity for variant classes."""
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                (
+                    "from rfdetr.detr import RFDETRLarge; "
+                    "from rfdetr.variants import RFDETRLarge as FromVariants; "
+                    "assert RFDETRLarge is FromVariants"
+                ),
+            ],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        assert result.returncode == 0, (
+            "Subprocess for detr-first import failed:\n"
+            f"return code: {result.returncode}\n"
+            f"stdout:\n{result.stdout}\n"
+            f"stderr:\n{result.stderr}"
+        )

tests/utilities/test_state_dict.py

@@ -0,0 +1,493 @@
+# ------------------------------------------------------------------------
+# 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.utilities.state_dict."""
+
+import logging
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from pytorch_lightning import LightningModule, Trainer
+
+from rfdetr.utilities.state_dict import (
+    _make_fit_loop_state,
+    remap_projector_to_cross_attn,
+    strip_checkpoint,
+    validate_checkpoint_compatibility,
+)
+
+# ---------------------------------------------------------------------------
+# _make_fit_loop_state
+# ---------------------------------------------------------------------------
+
+
+class TestMakeFitLoopState:
+    """Tests for _make_fit_loop_state epoch counter encoding."""
+
+    @pytest.mark.parametrize(
+        "epoch,expected_n",
+        [
+            pytest.param(0, 1, id="epoch_0"),
+            pytest.param(4, 5, id="epoch_4"),
+            pytest.param(9, 10, id="epoch_9"),
+        ],
+    )
+    def test_epoch_progress_completed_is_epoch_plus_one(self, epoch: int, expected_n: int) -> None:
+        """epoch_progress.current.completed == epoch + 1 so PTL sets current_epoch correctly."""
+        state = _make_fit_loop_state(epoch)
+        assert state["epoch_progress"]["current"]["completed"] == expected_n
+        assert state["epoch_progress"]["total"]["completed"] == expected_n
+
+    def test_epoch_progress_all_counters_equal(self) -> None:
+        """All four counters in epoch_progress should be equal (epoch fully completed)."""
+        state = _make_fit_loop_state(7)
+        for scope in ("total", "current"):
+            ep = state["epoch_progress"][scope]
+            vals = [ep["ready"], ep["started"], ep["processed"], ep["completed"]]
+            assert len(set(vals)) == 1, f"epoch_progress.{scope} counters differ: {ep}"
+
+    def test_batches_that_stepped_is_zero(self) -> None:
+        """Optimizer/scheduler state should start fresh; _batches_that_stepped must be 0."""
+        state = _make_fit_loop_state(3)
+        assert state["epoch_loop.state_dict"]["_batches_that_stepped"] == 0
+
+    def test_batch_progress_is_zero(self) -> None:
+        """Batch progress counters should be zeroed out (not mid-batch resume)."""
+        state = _make_fit_loop_state(5)
+        for key in ("epoch_loop.batch_progress", "epoch_loop.val_loop.batch_progress"):
+            bp = state[key]
+            assert bp["is_last_batch"] is False
+            for scope in ("total", "current"):
+                assert all(v == 0 for v in bp[scope].values()), f"{key}.{scope} not zero: {bp[scope]}"
+
+    def test_ptl_accepts_fit_loop_state(self) -> None:
+        """PTL's _FitLoop.load_state_dict must not raise with our synthesised state dict."""
+
+        class _DummyModule(LightningModule):
+            def training_step(self, batch, idx):
+                return torch.tensor(0.0, requires_grad=True)
+
+            def configure_optimizers(self):
+                return torch.optim.SGD(self.parameters(), lr=1e-3)
+
+        trainer = Trainer(max_epochs=10, accelerator="cpu", enable_progress_bar=False, logger=False)
+        trainer.strategy.connect(_DummyModule())
+
+        epoch = 4
+        state = _make_fit_loop_state(epoch)
+        trainer.fit_loop.load_state_dict(state)
+        assert trainer.current_epoch == epoch + 1
+
+    def test_required_top_level_keys_present(self) -> None:
+        """State dict must contain all keys the FitLoop accesses during load."""
+        required = {
+            "state_dict",
+            "epoch_loop.state_dict",
+            "epoch_loop.batch_progress",
+            "epoch_loop.scheduler_progress",
+            "epoch_loop.automatic_optimization.state_dict",
+            "epoch_loop.automatic_optimization.optim_progress",
+            "epoch_loop.manual_optimization.state_dict",
+            "epoch_loop.manual_optimization.optim_step_progress",
+            "epoch_loop.val_loop.state_dict",
+            "epoch_loop.val_loop.batch_progress",
+            "epoch_progress",
+        }
+        state = _make_fit_loop_state(0)
+        missing = required - set(state.keys())
+        assert not missing, f"Missing keys: {missing}"
+
+
+# ---------------------------------------------------------------------------
+# validate_checkpoint_compatibility
+# ---------------------------------------------------------------------------
+
+
+class TestValidateCheckpointCompatibility:
+    """Direct unit tests for validate_checkpoint_compatibility."""
+
+    # ------------------------------------------------------------------
+    # Early-return / silent-skip cases
+    # ------------------------------------------------------------------
+
+    def test_no_args_key_returns_without_raising(self):
+        """Checkpoint without 'args' key must return silently."""
+        checkpoint = {"model": {}}
+        model_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        validate_checkpoint_compatibility(checkpoint, model_args)  # must not raise
+
+    def test_ckpt_has_segmentation_head_model_does_not_skips(self):
+        """One-sided: ckpt has segmentation_head, model_args lacks it — skip, no error."""
+        ckpt_args = SimpleNamespace(segmentation_head=True, patch_size=14)
+        checkpoint = {"args": ckpt_args}
+        model_args = SimpleNamespace(patch_size=14)  # no segmentation_head attribute
+        validate_checkpoint_compatibility(checkpoint, model_args)  # must not raise
+
+    def test_ckpt_lacks_patch_size_model_has_it_skips(self):
+        """One-sided: ckpt has no patch_size, model has it — skip that check, no error."""
+        ckpt_args = SimpleNamespace(segmentation_head=False)  # no patch_size
+        checkpoint = {"args": ckpt_args}
+        model_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        validate_checkpoint_compatibility(checkpoint, model_args)  # must not raise
+
+    def test_compatible_checkpoint_no_exception(self):
+        """Checkpoint with matching segmentation_head and patch_size must not raise."""
+        ckpt_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        checkpoint = {"args": ckpt_args}
+        model_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        validate_checkpoint_compatibility(checkpoint, model_args)  # must not raise
+
+    def test_compatible_segmentation_checkpoint_no_exception(self):
+        """Matching segmentation model (seg_head=True both sides) must not raise."""
+        ckpt_args = SimpleNamespace(segmentation_head=True, patch_size=16)
+        checkpoint = {"args": ckpt_args}
+        model_args = SimpleNamespace(segmentation_head=True, patch_size=16)
+        validate_checkpoint_compatibility(checkpoint, model_args)  # must not raise
+
+    # ------------------------------------------------------------------
+    # segmentation_head mismatch
+    # ------------------------------------------------------------------
+
+    def test_seg_ckpt_into_detection_model_raises(self):
+        """Segmentation checkpoint loaded into a detection model must raise ValueError."""
+        ckpt_args = SimpleNamespace(segmentation_head=True, patch_size=14)
+        checkpoint = {"args": ckpt_args}
+        model_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        with pytest.raises(ValueError, match="segmentation head"):
+            validate_checkpoint_compatibility(checkpoint, model_args)
+
+    def test_detection_ckpt_into_seg_model_raises(self):
+        """Detection checkpoint loaded into a segmentation model must raise ValueError."""
+        ckpt_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        checkpoint = {"args": ckpt_args}
+        model_args = SimpleNamespace(segmentation_head=True, patch_size=14)
+        with pytest.raises(ValueError, match="segmentation head"):
+            validate_checkpoint_compatibility(checkpoint, model_args)
+
+    # ------------------------------------------------------------------
+    # patch_size mismatch
+    # ------------------------------------------------------------------
+
+    def test_patch_size_mismatch_raises_with_both_sizes(self):
+        """patch_size mismatch must raise ValueError and mention both sizes."""
+        ckpt_args = SimpleNamespace(segmentation_head=False, patch_size=12)
+        checkpoint = {"args": ckpt_args}
+        model_args = SimpleNamespace(segmentation_head=False, patch_size=16)
+        with pytest.raises(ValueError, match=r"patch_size=12.*patch_size=16|patch_size=16.*patch_size=12"):
+            validate_checkpoint_compatibility(checkpoint, model_args)
+
+    # ------------------------------------------------------------------
+    # patch_size inferred from projection weight (no "args" key)
+    # ------------------------------------------------------------------
+
+    @pytest.mark.parametrize(
+        "ckpt_patch_size,model_patch_size,should_raise",
+        [
+            pytest.param(16, 12, True, id="ckpt_16_model_12_raises"),
+            pytest.param(14, 16, True, id="ckpt_14_model_16_raises"),
+            pytest.param(16, 16, False, id="matching_16_no_raise"),
+        ],
+    )
+    def test_patch_size_inferred_from_projection_weight(
+        self, ckpt_patch_size: int, model_patch_size: int, should_raise: bool
+    ) -> None:
+        """Projection weight shape used to infer ckpt patch_size when 'args' key absent.
+
+        Regression test for #965 — pretrained COCO weights lack 'args', so the shape-based fallback must fire before
+        load_state_dict raises a cryptic RuntimeError.
+        """
+        proj_key = "backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight"
+        proj_weight = torch.zeros(384, 3, ckpt_patch_size, ckpt_patch_size)
+        checkpoint = {"model": {proj_key: proj_weight}}  # no "args" key
+        model_args = SimpleNamespace(patch_size=model_patch_size)
+
+        if should_raise:
+            with pytest.raises(
+                ValueError,
+                match=rf"patch_size={ckpt_patch_size}.*patch_size={model_patch_size}"
+                rf"|patch_size={model_patch_size}.*patch_size={ckpt_patch_size}",
+            ):
+                validate_checkpoint_compatibility(checkpoint, model_args)
+        else:
+            validate_checkpoint_compatibility(checkpoint, model_args)  # must not raise
+
+    @pytest.mark.parametrize(
+        "checkpoint,model_args_kwargs",
+        [
+            pytest.param(
+                {},
+                {"patch_size": 16},
+                id="no_model_key_skips",
+            ),
+            pytest.param(
+                {"model": {}},
+                {"patch_size": 16},
+                id="no_projection_key_skips",
+            ),
+            pytest.param(
+                {
+                    "model": {
+                        "backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight": torch.zeros(
+                            384, 3, 16, 16
+                        )
+                    }
+                },
+                {},
+                id="model_no_patch_size_attr_skips",
+            ),
+            pytest.param(
+                {
+                    "model": {
+                        "backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight": torch.zeros(
+                            384, 3
+                        )  # 2D — not a Conv2d weight; rank guard must skip cleanly
+                    }
+                },
+                {"patch_size": 16},
+                id="proj_weight_2d_skips",
+            ),
+            pytest.param(
+                {
+                    "model": {
+                        "backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight": torch.zeros(
+                            384, 3, 16
+                        )  # 3D — not a Conv2d weight; rank guard must skip cleanly
+                    }
+                },
+                {"patch_size": 16},
+                id="proj_weight_3d_skips",
+            ),
+            pytest.param(
+                {
+                    "model": {
+                        "backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight": torch.zeros(
+                            384, 3, 16, 16, 16
+                        )  # 5D — Conv3d-like; rank guard (== 4) must skip cleanly
+                    }
+                },
+                {"patch_size": 8},
+                id="proj_weight_5d_skips",
+            ),
+            pytest.param(
+                {
+                    "args": SimpleNamespace(patch_size=14),
+                    "model": {
+                        "backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight": torch.zeros(
+                            384, 3, 16, 16
+                        )  # projection suggests 16, but args.patch_size=14 takes precedence
+                    },
+                },
+                {"patch_size": 14},
+                id="args_patch_size_suppresses_projection_inference",
+            ),
+            pytest.param(
+                {
+                    "args": {"patch_size": 14},  # PTL-style dict args (not SimpleNamespace)
+                    "model": {
+                        "backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight": torch.zeros(
+                            384, 3, 16, 16
+                        )  # projection suggests 16, but dict args["patch_size"]=14 takes precedence
+                    },
+                },
+                {"patch_size": 14},
+                id="dict_args_patch_size_suppresses_projection_inference",
+            ),
+        ],
+    )
+    def test_projection_inference_silently_skips_when_incomplete(
+        self, checkpoint: dict, model_args_kwargs: dict
+    ) -> None:
+        """Shape-based patch_size check is skipped when key or attribute is absent.
+
+        Verifies backward compatibility: missing projection key, missing model key, model_args without patch_size
+        attribute, non-4D projection weights, or an explicit args.patch_size (SimpleNamespace or dict) must all be
+        handled without error.
+        """
+        model_args = SimpleNamespace(**model_args_kwargs)
+        validate_checkpoint_compatibility(checkpoint, model_args)  # must not raise
+
+    def test_non_square_projection_kernel_skips_check(self) -> None:
+        """Non-square patch projection kernel is skipped — patch_size cannot be inferred reliably.
+
+        Guards against hypothetical future backbones with non-square Conv2d kernels where shape[-1] would not equal
+        patch_size.
+        """
+        proj_key = "backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight"
+        proj_weight = torch.zeros(384, 3, 16, 14)  # non-square: h=16, w=14
+        checkpoint = {"model": {proj_key: proj_weight}}
+        model_args = SimpleNamespace(patch_size=16)
+        validate_checkpoint_compatibility(checkpoint, model_args)  # must not raise
+
+    # ------------------------------------------------------------------
+    # class-count mismatch warnings
+    # ------------------------------------------------------------------
+
+    def test_class_count_mismatch_backbone_pretrain_warns(self, caplog):
+        """Backbone pretrain scenario: checkpoint 91 classes, model 2 — warns about re-init."""
+        ckpt_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        checkpoint = {
+            "args": ckpt_args,
+            "model": {"class_embed.bias": torch.randn(91)},
+        }
+        model_args = SimpleNamespace(segmentation_head=False, patch_size=14, num_classes=2)
+
+        rf_detr_logger = logging.getLogger("rf-detr")
+        prev_propagate = rf_detr_logger.propagate
+        rf_detr_logger.propagate = True
+        try:
+            with caplog.at_level(logging.WARNING, logger="rf-detr"):
+                validate_checkpoint_compatibility(checkpoint, model_args)
+        finally:
+            rf_detr_logger.propagate = prev_propagate
+
+        warning_msgs = [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING]
+        assert any("re-initialized to 2 classes" in msg for msg in warning_msgs), (
+            f"Expected 're-initialized to 2 classes' warning, got: {warning_msgs}"
+        )
+
+    def test_class_count_mismatch_finetune_checkpoint_warns(self, caplog):
+        """Fine-tuned checkpoint scenario: checkpoint 3 classes, model 90 — warns with num_classes hint."""
+        ckpt_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        checkpoint = {
+            "args": ckpt_args,
+            "model": {"class_embed.bias": torch.randn(3)},
+        }
+        model_args = SimpleNamespace(segmentation_head=False, patch_size=14, num_classes=90)
+
+        rf_detr_logger = logging.getLogger("rf-detr")
+        prev_propagate = rf_detr_logger.propagate
+        rf_detr_logger.propagate = True
+        try:
+            with caplog.at_level(logging.WARNING, logger="rf-detr"):
+                validate_checkpoint_compatibility(checkpoint, model_args)
+        finally:
+            rf_detr_logger.propagate = prev_propagate
+
+        warning_msgs = [r.getMessage() for r in caplog.records if r.name == "rf-detr" and r.levelno >= logging.WARNING]
+        assert any("Pass num_classes=2" in msg for msg in warning_msgs), (
+            f"Expected 'Pass num_classes=2' warning, got: {warning_msgs}"
+        )
+
+    def test_class_count_match_no_warning(self, caplog):
+        """Matching class count — no warning emitted."""
+        ckpt_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        checkpoint = {
+            "args": ckpt_args,
+            "model": {"class_embed.bias": torch.randn(91)},
+        }
+        model_args = SimpleNamespace(segmentation_head=False, patch_size=14, num_classes=90)
+
+        rf_detr_logger = logging.getLogger("rf-detr")
+        prev_propagate = rf_detr_logger.propagate
+        rf_detr_logger.propagate = True
+        try:
+            with caplog.at_level(logging.WARNING, logger="rf-detr"):
+                validate_checkpoint_compatibility(checkpoint, model_args)
+        finally:
+            rf_detr_logger.propagate = prev_propagate
+
+        warning_msgs = [r.getMessage() for r in caplog.records if r.name == "rf-detr" and r.levelno >= logging.WARNING]
+        assert not warning_msgs, f"Expected no warnings, got: {warning_msgs}"
+
+
+class TestRemapProjectorToCrossAttn:
+    """Tests for dual-projector checkpoint key remapping."""
+
+    def test_clones_projector_weights_when_dual_projector_enabled(self) -> None:
+        """Dual-projector models clone projector keys into cross_attn_projector when missing."""
+        state_dict = {
+            "backbone.0.projector.0.weight": torch.randn(4, 4, 1, 1),
+            "backbone.0.projector.0.bias": torch.randn(4),
+        }
+        model = SimpleNamespace(backbone=[SimpleNamespace(dual_projector=True)])
+
+        remapped = remap_projector_to_cross_attn(state_dict, model)
+
+        assert remapped is state_dict
+        assert "backbone.0.cross_attn_projector.0.weight" in remapped
+        assert "backbone.0.cross_attn_projector.0.bias" in remapped
+        assert torch.equal(
+            remapped["backbone.0.cross_attn_projector.0.weight"],
+            state_dict["backbone.0.projector.0.weight"],
+        )
+        assert torch.equal(
+            remapped["backbone.0.cross_attn_projector.0.bias"],
+            state_dict["backbone.0.projector.0.bias"],
+        )
+
+    def test_skips_when_cross_attn_keys_already_present(self) -> None:
+        """No remap is applied when cross_attn_projector keys already exist."""
+        state_dict = {
+            "backbone.0.projector.0.weight": torch.randn(4, 4, 1, 1),
+            "backbone.0.cross_attn_projector.0.weight": torch.randn(4, 4, 1, 1),
+        }
+        model = SimpleNamespace(backbone=[SimpleNamespace(dual_projector=True)])
+
+        remapped = remap_projector_to_cross_attn(state_dict, model)
+
+        assert remapped is state_dict
+        assert len([key for key in remapped if key.startswith("backbone.0.cross_attn_projector.")]) == 1
+
+    def test_class_count_missing_model_key_no_warning(self, caplog):
+        """Checkpoint without 'model' key — no warning (backward compat)."""
+        ckpt_args = SimpleNamespace(segmentation_head=False, patch_size=14)
+        checkpoint = {"args": ckpt_args}
+        model_args = SimpleNamespace(segmentation_head=False, patch_size=14, num_classes=90)
+
+        rf_detr_logger = logging.getLogger("rf-detr")
+        prev_propagate = rf_detr_logger.propagate
+        rf_detr_logger.propagate = True
+        try:
+            with caplog.at_level(logging.WARNING, logger="rf-detr"):
+                validate_checkpoint_compatibility(checkpoint, model_args)
+        finally:
+            rf_detr_logger.propagate = prev_propagate
+
+        warning_msgs = [r.getMessage() for r in caplog.records if r.name == "rf-detr" and r.levelno >= logging.WARNING]
+        assert not warning_msgs, f"Expected no warnings, got: {warning_msgs}"
+
+
+class TestStripCheckpoint:
+    """Tests for strip_checkpoint loop-stub backfill."""
+
+    def _make_minimal_ckpt(self, tmp_path, extra: dict | None = None) -> Path:
+        """Write a minimal checkpoint to a temp file."""
+        payload = {"model": {"w": torch.tensor(1.0)}, "args": {"lr": 1e-4}}
+        if extra:
+            payload.update(extra)
+        ckpt_path = Path(tmp_path) / "ckpt.pth"
+        torch.save(payload, ckpt_path)
+        return ckpt_path
+
+    def test_strip_adds_validate_loop_stub_when_loops_present_but_missing_key(self, tmp_path) -> None:
+        """Old checkpoints with loops but no validate_loop/test_loop get stubs backfilled."""
+        ckpt_path = self._make_minimal_ckpt(
+            tmp_path,
+            extra={"loops": {"fit_loop": {"state_dict": {}}}},
+        )
+        strip_checkpoint(ckpt_path)
+        result = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        assert result["loops"]["validate_loop"] == {"state_dict": {}}
+        assert result["loops"]["test_loop"] == {"state_dict": {}}
+
+    def test_strip_preserves_existing_validate_loop_stub(self, tmp_path) -> None:
+        """Checkpoints with validate_loop already present are not overwritten."""
+        original_stub = {"state_dict": {"some_key": 1}}
+        ckpt_path = self._make_minimal_ckpt(
+            tmp_path,
+            extra={"loops": {"fit_loop": {"state_dict": {}}, "validate_loop": original_stub}},
+        )
+        strip_checkpoint(ckpt_path)
+        result = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        assert result["loops"]["validate_loop"] == original_stub
+
+    def test_strip_no_loops_key_leaves_loops_absent(self, tmp_path) -> None:
+        """Checkpoints without a loops key must not gain one after stripping."""
+        ckpt_path = self._make_minimal_ckpt(tmp_path)
+        strip_checkpoint(ckpt_path)
+        result = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        assert "loops" not in result

tests/utilities/test_tensors.py

@@ -0,0 +1,560 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for rfdetr.utilities.tensors.
+
+Covers:
+- ``_bilinear_grid_sample`` parity (manual gather path vs ``F.grid_sample``).
+- ``nested_tensor_from_tensor_list`` with ``block_size`` (backbone-aware batch rounding).
+- ``make_collate_fn`` factory.
+"""
+
+import pickle
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn.functional as F  # noqa: N812
+import torch.testing
+
+from rfdetr.utilities.tensors import (
+    _bilinear_grid_sample,
+    make_collate_fn,
+    nested_tensor_from_tensor_list,
+)
+
+
+def _grid_sample_reference(
+    input: torch.Tensor,
+    grid: torch.Tensor,
+    padding_mode: str = "zeros",
+    align_corners: bool = False,
+) -> torch.Tensor:
+    """Ground-truth output from F.grid_sample for comparison."""
+    return F.grid_sample(
+        input,
+        grid,
+        mode="bilinear",
+        padding_mode=padding_mode,
+        align_corners=align_corners,
+    )
+
+
+def _call_manual_path(
+    input: torch.Tensor,
+    grid: torch.Tensor,
+    padding_mode: str = "zeros",
+    align_corners: bool = False,
+) -> torch.Tensor:
+    """Force the manual gather-based code path by mocking input.device.type.
+
+    The function checks ``input.device.type != "mps"`` to decide which branch to take.  We patch ``torch.Tensor.device``
+    to return an object whose ``.type`` is ``"mps"`` so the manual path runs on a normal CPU tensor.
+    """
+
+    class _FakeMPSDevice:
+        type = "mps"
+
+        def __eq__(self, other):
+            return False
+
+        def __repr__(self):
+            return "device(type='mps')"
+
+    with patch.object(torch.Tensor, "device", new_callable=lambda: property(lambda self: _FakeMPSDevice())):
+        return _bilinear_grid_sample(input, grid, padding_mode=padding_mode, align_corners=align_corners)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def seed():
+    """Fix random seed for reproducible grid/input generation."""
+    torch.manual_seed(42)
+
+
+# ---------------------------------------------------------------------------
+# Test scenarios as parametrize parameters
+# ---------------------------------------------------------------------------
+
+_PADDING_ALIGN_COMBOS = [
+    pytest.param("zeros", False, id="zeros-no_align"),
+    pytest.param("border", False, id="border-no_align"),
+    pytest.param("zeros", True, id="zeros-align_corners"),
+]
+
+_LOW_PRECISION_DTYPES = [
+    pytest.param(torch.float16, id="float16"),
+    pytest.param(torch.bfloat16, id="bfloat16"),
+]
+
+_LOW_PRECISION_GRAD_TOLERANCES = {
+    torch.float16: (1e-2, 2e-2),
+    torch.bfloat16: (3e-2, 1e-1),
+}
+
+
+def _require_grid_sample_dtype_support(dtype: torch.dtype) -> None:
+    """Skip test when current backend does not support grid_sample for dtype."""
+    input = torch.randn(1, 1, 2, 2, dtype=dtype, requires_grad=True)
+    grid = (torch.rand(1, 1, 1, 2, dtype=dtype) * 1.6 - 0.8).requires_grad_(True)
+    try:
+        out = F.grid_sample(input, grid, mode="bilinear", padding_mode="zeros", align_corners=False)
+        out.backward(torch.ones_like(out))
+    except RuntimeError as error:
+        pytest.skip(f"grid_sample dtype support missing for {dtype}: {error}")
+
+
+class TestBilinearGridSampleParity:
+    """Manual gather path must match F.grid_sample for all grid/padding combos."""
+
+    @pytest.mark.parametrize(
+        "padding_mode, align_corners",
+        _PADDING_ALIGN_COMBOS,
+    )
+    def test_interior_grid_coordinates(self, seed, padding_mode, align_corners):
+        """Grid values well inside [-1, 1] -- pure interpolation, no boundary effects."""
+        input = torch.randn(1, 3, 8, 8)
+        # Grid in [-0.8, 0.8] -- safely inside
+        grid = torch.rand(1, 4, 4, 2) * 1.6 - 0.8
+
+        expected = _grid_sample_reference(input, grid, padding_mode, align_corners)
+        actual = _call_manual_path(input, grid, padding_mode, align_corners)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+    @pytest.mark.parametrize(
+        "padding_mode, align_corners",
+        _PADDING_ALIGN_COMBOS,
+    )
+    def test_partially_outside_grid_coordinates(self, seed, padding_mode, align_corners):
+        """Grid values spanning [-1.5, 1.5] -- some samples fall outside the image."""
+        input = torch.randn(1, 3, 8, 8)
+        grid = torch.rand(1, 6, 6, 2) * 3.0 - 1.5
+
+        expected = _grid_sample_reference(input, grid, padding_mode, align_corners)
+        actual = _call_manual_path(input, grid, padding_mode, align_corners)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+    @pytest.mark.parametrize(
+        "padding_mode, align_corners",
+        _PADDING_ALIGN_COMBOS,
+    )
+    def test_exact_boundary_grid_values(self, seed, padding_mode, align_corners):
+        """Grid values at exact boundaries: -1.0, 0.0, 1.0."""
+        input = torch.randn(1, 2, 4, 4)
+        # Manually craft grid with boundary values
+        coords = torch.tensor([-1.0, 0.0, 1.0])
+        grid_y, grid_x = torch.meshgrid(coords, coords, indexing="ij")
+        grid = torch.stack([grid_x, grid_y], dim=-1).unsqueeze(0)  # (1, 3, 3, 2)
+
+        expected = _grid_sample_reference(input, grid, padding_mode, align_corners)
+        actual = _call_manual_path(input, grid, padding_mode, align_corners)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+    @pytest.mark.parametrize(
+        "padding_mode, align_corners",
+        _PADDING_ALIGN_COMBOS,
+    )
+    def test_single_pixel_input(self, padding_mode, align_corners):
+        """1x1 spatial input -- extreme edge case for index arithmetic."""
+        input = torch.tensor([[[[3.14]]]])  # (1, 1, 1, 1)
+        grid = torch.tensor([[[[0.0, 0.0]]]])  # (1, 1, 1, 2) -- center
+
+        expected = _grid_sample_reference(input, grid, padding_mode, align_corners)
+        actual = _call_manual_path(input, grid, padding_mode, align_corners)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+    @pytest.mark.parametrize(
+        "padding_mode, align_corners",
+        _PADDING_ALIGN_COMBOS,
+    )
+    def test_batch_and_multichannel(self, seed, padding_mode, align_corners):
+        """Batch size > 1 and multiple channels."""
+        input = torch.randn(3, 5, 10, 12)
+        grid = torch.rand(3, 7, 9, 2) * 2.0 - 1.0  # [-1, 1]
+
+        expected = _grid_sample_reference(input, grid, padding_mode, align_corners)
+        actual = _call_manual_path(input, grid, padding_mode, align_corners)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+    @pytest.mark.parametrize(
+        "padding_mode, align_corners",
+        _PADDING_ALIGN_COMBOS,
+    )
+    def test_all_out_of_bounds(self, padding_mode, align_corners):
+        """All grid coordinates far outside [-1, 1] -- tests OOB handling."""
+        input = torch.randn(1, 2, 4, 4)
+        # All coordinates at +5.0 -- far outside
+        grid = torch.full((1, 3, 3, 2), 5.0)
+
+        expected = _grid_sample_reference(input, grid, padding_mode, align_corners)
+        actual = _call_manual_path(input, grid, padding_mode, align_corners)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+    @pytest.mark.parametrize(
+        "padding_mode, align_corners",
+        _PADDING_ALIGN_COMBOS,
+    )
+    def test_negative_out_of_bounds(self, padding_mode, align_corners):
+        """All grid coordinates at -5.0 -- far outside on the negative side."""
+        input = torch.randn(1, 2, 4, 4)
+        grid = torch.full((1, 3, 3, 2), -5.0)
+
+        expected = _grid_sample_reference(input, grid, padding_mode, align_corners)
+        actual = _call_manual_path(input, grid, padding_mode, align_corners)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+    @pytest.mark.parametrize(
+        "padding_mode, align_corners",
+        [
+            pytest.param("zeros", False, id="zeros-no_align"),
+            pytest.param("border", False, id="border-no_align"),
+        ],
+    )
+    def test_non_square_spatial_dimensions(self, seed, padding_mode, align_corners):
+        """Non-square H != W input -- tests that x/y coordinate handling is correct."""
+        input = torch.randn(1, 2, 5, 13)  # tall vs wide
+        grid = torch.rand(1, 4, 6, 2) * 2.0 - 1.0
+
+        expected = _grid_sample_reference(input, grid, padding_mode, align_corners)
+        actual = _call_manual_path(input, grid, padding_mode, align_corners)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+
+class TestBilinearGridSampleDelegation:
+    """On non-MPS devices, the function delegates directly to F.grid_sample."""
+
+    def test_cpu_delegates_to_grid_sample(self, seed):
+        """On CPU, output should match F.grid_sample exactly (same code path)."""
+        input = torch.randn(1, 3, 8, 8)
+        grid = torch.rand(1, 4, 4, 2) * 2.0 - 1.0
+
+        expected = _grid_sample_reference(input, grid, "zeros", False)
+        actual = _bilinear_grid_sample(input, grid, padding_mode="zeros", align_corners=False)
+
+        torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+    def test_cpu_border_delegates_to_grid_sample(self, seed):
+        """On CPU with border padding, output matches F.grid_sample exactly."""
+        input = torch.randn(2, 4, 6, 6)
+        grid = torch.rand(2, 3, 3, 2) * 3.0 - 1.5
+
+        expected = _grid_sample_reference(input, grid, "border", False)
+        actual = _bilinear_grid_sample(input, grid, padding_mode="border", align_corners=False)
+
+        torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
+class TestBilinearGridSampleOutputShape:
+    """Output shape must be (N, C, Hg, Wg) for all inputs."""
+
+    @pytest.mark.parametrize(
+        "n, c, h, w, hg, wg",
+        [
+            pytest.param(1, 1, 1, 1, 1, 1, id="minimal"),
+            pytest.param(1, 3, 8, 8, 4, 4, id="standard"),
+            pytest.param(2, 5, 10, 12, 7, 9, id="batch_multichannel"),
+            pytest.param(1, 1, 3, 7, 5, 5, id="non_square"),
+        ],
+    )
+    def test_output_shape(self, n, c, h, w, hg, wg):
+        """Manual path output shape is (N, C, Hg, Wg)."""
+        input = torch.randn(n, c, h, w)
+        grid = torch.rand(n, hg, wg, 2) * 2.0 - 1.0
+
+        actual = _call_manual_path(input, grid)
+        assert actual.shape == (n, c, hg, wg), f"Expected shape ({n}, {c}, {hg}, {wg}), got {actual.shape}"
+
+
+class TestBilinearGridSampleGradient:
+    """Gradient correctness for the manual gather path."""
+
+    @pytest.mark.parametrize(
+        "padding_mode, align_corners",
+        [
+            pytest.param("zeros", False, id="zeros-no_align"),
+            pytest.param("border", False, id="border-no_align"),
+            pytest.param("zeros", True, id="zeros-align_corners"),
+        ],
+    )
+    def test_gradient_matches_grid_sample(self, seed, padding_mode, align_corners):
+        """Gradients from manual path match those from F.grid_sample."""
+        input_ref = torch.randn(1, 2, 6, 6, requires_grad=True)
+        grid_ref = (torch.rand(1, 4, 4, 2) * 1.6 - 0.8).requires_grad_(True)
+
+        # Clone for manual path
+        input_man = input_ref.detach().clone().requires_grad_(True)
+        grid_man = grid_ref.detach().clone().requires_grad_(True)
+
+        # Forward
+        out_ref = _grid_sample_reference(input_ref, grid_ref, padding_mode, align_corners)
+        out_man = _call_manual_path(input_man, grid_man, padding_mode, align_corners)
+
+        # Backward with same upstream gradient
+        upstream = torch.randn_like(out_ref)
+        out_ref.backward(upstream)
+        out_man.backward(upstream)
+
+        torch.testing.assert_close(
+            input_man.grad,
+            input_ref.grad,
+            atol=1e-5,
+            rtol=1e-5,
+            msg="Input gradient mismatch between manual path and F.grid_sample",
+        )
+        torch.testing.assert_close(
+            grid_man.grad,
+            grid_ref.grad,
+            atol=1e-5,
+            rtol=1e-5,
+            msg="Grid gradient mismatch between manual path and F.grid_sample",
+        )
+
+    def test_gradcheck_manual_path(self, seed):
+        """torch.autograd.gradcheck passes on the manual path (double precision)."""
+        input = torch.randn(1, 1, 4, 4, dtype=torch.float64, requires_grad=True)
+        grid = (torch.rand(1, 3, 3, 2, dtype=torch.float64) * 1.6 - 0.8).requires_grad_(True)
+
+        assert torch.autograd.gradcheck(
+            lambda inp, grd: _call_manual_path(inp, grd, padding_mode="zeros", align_corners=False),
+            (input, grid),
+            eps=1e-6,
+            atol=1e-4,
+            rtol=1e-3,
+        ), "gradcheck failed for manual bilinear grid sample path"
+
+
+class TestBilinearGridSampleLowPrecision:
+    """Low-precision parity and gradients stay aligned with F.grid_sample."""
+
+    @pytest.mark.parametrize("dtype", _LOW_PRECISION_DTYPES)
+    def test_low_precision_parity(self, seed, dtype):
+        """Manual path output matches F.grid_sample for low-precision inputs."""
+        _require_grid_sample_dtype_support(dtype)
+
+        input = torch.randn(2, 3, 6, 6, dtype=dtype)
+        grid = torch.rand(2, 4, 4, 2, dtype=dtype) * 3.0 - 1.5
+
+        expected = _grid_sample_reference(input, grid, padding_mode="zeros", align_corners=False)
+        actual = _call_manual_path(input, grid, padding_mode="zeros", align_corners=False)
+
+        torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3)
+        assert actual.dtype == dtype
+
+    @pytest.mark.parametrize("dtype", _LOW_PRECISION_DTYPES)
+    def test_low_precision_gradient_parity(self, seed, dtype):
+        """Manual path gradients match F.grid_sample gradients for low precision."""
+        _require_grid_sample_dtype_support(dtype)
+        atol, rtol = _LOW_PRECISION_GRAD_TOLERANCES[dtype]
+
+        input_ref = torch.randn(1, 2, 6, 6, dtype=dtype, requires_grad=True)
+        grid_ref = (torch.rand(1, 4, 4, 2, dtype=dtype) * 1.6 - 0.8).requires_grad_(True)
+
+        input_man = input_ref.detach().clone().requires_grad_(True)
+        grid_man = grid_ref.detach().clone().requires_grad_(True)
+
+        out_ref = _grid_sample_reference(input_ref, grid_ref, padding_mode="zeros", align_corners=False)
+        out_man = _call_manual_path(input_man, grid_man, padding_mode="zeros", align_corners=False)
+
+        upstream = torch.randn_like(out_ref)
+        out_ref.backward(upstream)
+        out_man.backward(upstream)
+
+        torch.testing.assert_close(input_man.grad, input_ref.grad, atol=atol, rtol=rtol)
+        torch.testing.assert_close(grid_man.grad, grid_ref.grad, atol=atol, rtol=rtol)
+        assert input_man.grad is not None
+        assert grid_man.grad is not None
+        assert input_man.grad.dtype == dtype
+        assert grid_man.grad.dtype == dtype
+
+
+class TestBilinearGridSampleRealUseCases:
+    """Parity tests matching the actual call sites in the codebase."""
+
+    def test_ms_deform_attn_pattern(self, seed):
+        """Matches ms_deform_attn_func: padding_mode='zeros', align_corners=False.
+
+        The attention function passes (B*n_heads, head_dim, H, W) input and (B*n_heads, Len_q, P, 2) grid.
+        """
+        # Simulate B=2, n_heads=8, head_dim=32
+        input = torch.randn(16, 32, 14, 14)
+        grid = torch.rand(16, 100, 4, 2) * 2.0 - 1.0
+
+        expected = _grid_sample_reference(input, grid, "zeros", False)
+        actual = _call_manual_path(input, grid, "zeros", False)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+    def test_point_sample_pattern(self, seed):
+        """Matches point_sample in segmentation: padding_mode='border', align_corners=False.
+
+        point_sample transforms point_coords via ``2.0 * point_coords - 1.0`` to map [0, 1] -> [-1, 1].
+        """
+        input = torch.randn(4, 256, 28, 28)
+        # Simulate point_coords in [0, 1], transformed to [-1, 1]
+        point_coords_01 = torch.rand(4, 12544, 1, 2)
+        grid = 2.0 * point_coords_01 - 1.0
+
+        expected = _grid_sample_reference(input, grid, "border", False)
+        actual = _call_manual_path(input, grid, "border", False)
+
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+
+class TestNestedTensorBlockSize:
+    """``nested_tensor_from_tensor_list`` with block_size rounds batch max H/W up.
+
+    This is the collator-level pad for backbone divisibility.  The rounded-up strip must be marked as padding in the
+    mask so downstream attention skips it.  See
+    https://github.com/roboflow/rf-detr/issues/983
+    for context.
+    """
+
+    @staticmethod
+    def _image(c: int, h: int, w: int, fill: float = 1.0) -> torch.Tensor:
+        """Return a ``(C, H, W)`` float32 tensor filled with the given value."""
+        return torch.full((c, h, w), fill, dtype=torch.float32)
+
+    def test_block_size_none_preserves_old_behavior(self) -> None:
+        """Without block_size, the batch tensor is exactly batch-max H/W."""
+        images = [self._image(3, 100, 200), self._image(3, 150, 180)]
+        nested = nested_tensor_from_tensor_list(images)
+        _, _, h, w = nested.tensors.shape
+        assert (h, w) == (150, 200)
+        # Mask reflects per-image sizes (no block rounding).
+        assert nested.mask[0, :100, :200].any().item() is False
+        assert nested.mask[0, 100:, :].all().item() is True
+        assert nested.mask[1, :150, :180].any().item() is False
+        assert nested.mask[1, :, 180:].all().item() is True
+
+    def test_block_size_rounds_up(self) -> None:
+        """Batch-max is rounded up to the next multiple of block_size."""
+        images = [self._image(3, 100, 200), self._image(3, 150, 180)]
+        nested = nested_tensor_from_tensor_list(images, block_size=32)
+        _, _, h, w = nested.tensors.shape
+        # max_h=150 -> 160, max_w=200 -> 224
+        assert (h, w) == (160, 224)
+
+    def test_block_size_equal_to_max_is_noop(self) -> None:
+        """When batch max already matches a multiple of block_size, no extra rounding."""
+        images = [self._image(3, 128, 256)]
+        nested = nested_tensor_from_tensor_list(images, block_size=32)
+        _, _, h, w = nested.tensors.shape
+        assert (h, w) == (128, 256)
+
+    def test_divisor_pad_marked_in_mask(self) -> None:
+        """All padded cells (both batch-level and divisor round-up) are marked True in the mask."""
+        images = [self._image(3, 100, 200)]
+        nested = nested_tensor_from_tensor_list(images, block_size=32)
+        tensor = nested.tensors[0]
+        mask = nested.mask[0]
+
+        # Content region is the original 100x200; mask[:100, :200] must be False.
+        assert mask[:100, :200].any().item() is False
+        # The rounded-up strip (100:128 rows, 200:224 cols) must be True.
+        assert mask[100:, :].all().item() is True
+        assert mask[:, 200:].all().item() is True
+
+        # Content region is the original fill; pad region is zero.
+        assert torch.all(tensor[:, :100, :200] == 1.0)
+        assert torch.all(tensor[:, 100:, :] == 0.0)
+        assert torch.all(tensor[:, :, 200:] == 0.0)
+
+    @pytest.mark.parametrize(
+        "block_size,shape,expected",
+        [
+            pytest.param(32, (100, 100), (128, 128), id="both-rounded"),
+            pytest.param(32, (128, 200), (128, 224), id="h-aligned-w-rounded"),
+            pytest.param(32, (100, 256), (128, 256), id="h-rounded-w-aligned"),
+            pytest.param(56, (100, 100), (112, 112), id="patch14-num-windows4"),
+            pytest.param(64, (100, 100), (128, 128), id="block-size-64"),
+        ],
+    )
+    def test_single_image_rounding_parametrized(self, block_size: int, shape: tuple, expected: tuple) -> None:
+        """Single-image batch; round-up applied correctly for various block sizes."""
+        images = [self._image(3, shape[0], shape[1])]
+        nested = nested_tensor_from_tensor_list(images, block_size=block_size)
+        _, _, h, w = nested.tensors.shape
+        assert (h, w) == expected
+
+
+class TestMakeCollateFn:
+    """``make_collate_fn`` returns a picklable collate callable with block_size rounding baked in."""
+
+    @staticmethod
+    def _batch(*shapes: tuple[int, ...]) -> list[tuple[torch.Tensor, dict]]:
+        """Build a list of ``(tensor, target_dict)`` pairs with given shapes.
+
+        Args:
+            *shapes: Variadic sequence of ``(C, H, W)`` shapes, one per image.
+
+        Returns:
+            List of ``(image_tensor, target_dict)`` pairs ready to pass to a collate callable.
+        """
+        batch = []
+        for shape in shapes:
+            img = torch.full(shape, 1.0, dtype=torch.float32)
+            target = {"boxes": torch.zeros((0, 4)), "labels": torch.zeros((0,), dtype=torch.long)}
+            batch.append((img, target))
+        return batch
+
+    def test_default_block_size_none_behaves_like_collate_fn(self) -> None:
+        """With block_size=None, the factory returns a collate equivalent to the default."""
+        collate = make_collate_fn()  # block_size=None
+        samples, targets = collate(self._batch((3, 100, 200), (3, 150, 180)))
+        _, _, h, w = samples.tensors.shape
+        assert (h, w) == (150, 200)  # exact batch max
+        assert len(targets) == 2
+
+    def test_block_size_rounds_up_batch_max(self) -> None:
+        """Factory with block_size=32 rounds batch-max up to 32-multiples."""
+        collate = make_collate_fn(block_size=32)
+        samples, _ = collate(self._batch((3, 100, 200), (3, 150, 180)))
+        _, _, h, w = samples.tensors.shape
+        assert (h, w) == (160, 224)
+
+    def test_targets_passed_through(self) -> None:
+        """Factory collator preserves the list-of-targets second element."""
+        collate = make_collate_fn(block_size=32)
+        samples, targets = collate(self._batch((3, 100, 200), (3, 150, 180)))
+        assert isinstance(targets, tuple)
+        assert len(targets) == 2
+        for t in targets:
+            assert set(t.keys()) == {"boxes", "labels"}
+
+    def test_mixed_landscape_portrait_batch_masked_correctly(self) -> None:
+        """Mixed-orientation batch: all pad (batch + divisor) correctly marked True in mask."""
+        # landscape (H=100, W=200) and portrait (H=200, W=100).  block_size=32 rounds
+        # batch max (200, 200) to (224, 224).
+        collate = make_collate_fn(block_size=32)
+        samples, _ = collate(self._batch((3, 100, 200), (3, 200, 100)))
+        _, _, h, w = samples.tensors.shape
+        assert (h, w) == (224, 224)
+
+        # Each image's content region equals its original shape; everything else is pad.
+        mask_a = samples.mask[0]
+        mask_b = samples.mask[1]
+        assert mask_a[:100, :200].any().item() is False
+        assert mask_a[100:, :].all().item() is True
+        assert mask_a[:, 200:].all().item() is True
+        assert mask_b[:200, :100].any().item() is False
+        assert mask_b[200:, :].all().item() is True
+        assert mask_b[:, 100:].all().item() is True
+
+    def test_make_collate_fn_is_picklable(self) -> None:
+        """make_collate_fn returns a functools.partial picklable for num_workers > 0."""
+        collate = make_collate_fn(block_size=32)
+        assert pickle.dumps(collate) is not None