chore: import upstream snapshot with attribution

tests/models/heads/__init__.py

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

tests/models/heads/test_detection_head.py

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

tests/models/heads/test_keypoint_heads.py

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

tests/models/heads/test_segmentation_head.py

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