fix(models): build spatial_shapes from Python ints under torch.compile

CHANGELOG.md

@@ -64,6 +64,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 - Empty COCO targets now keep `iscrowd` as `int64` and `area` as `float32`, matching populated targets and allowing mixed empty/populated batches to use lossless packed-target worker transport.
 
+- `compile=True` no longer aborts the training step. Dynamo polyfills `torch._shape_as_tensor` to return a `torch.Size` rather than a tensor, so `Transformer.forward`'s `torch.stack([torch._shape_as_tensor(src)[2:4] for src in srcs])` raised `TypeError: expected Tensor as element 0 in argument 0, but got torch.Size` as soon as `torch.compile` traced it, and `torch._dynamo.config.suppress_errors = True` did not absorb it because the `TypeError` is raised by user code rather than by Dynamo. `spatial_shapes` is now built from the Python-int `(H, W)` pairs the function already collects whenever `torch.compiler.is_compiling()` is true, the same formulation the `torch.export`/ExecuTorch branch has used since #1142. Eager, `torch.jit.trace` and the TorchScript-ONNX/TensorRT export path (#1155) keep the `torch._shape_as_tensor` form unchanged: neither guard is true there. Reproduced on torch 2.9.1 and 2.13.0, at `dynamic=True`, `dynamic=None` and `dynamic=False`. `compile=True` is still inert under the default `multi_scale=True`, which disables compilation before this path is reached.
+
 ---
 
 ## [1.9.4] — 2026-08-24

src/rfdetr/models/transformer.py

@@ -31,14 +31,16 @@
 from rfdetr.models.ops.modules import MSDeformAttn
 
 
-def _not_exporting() -> bool:
-    """Fallback for ``torch.compiler.is_exporting`` on torch<2.6 (which lacks it and never runs ExecuTorch export).
+def _tracer_absent() -> bool:
+    """Fallback for ``torch.compiler.is_exporting``/``is_compiling`` on torch versions that lack them.
 
-    Hoisted to a module-level function so the hot decoder ``forward`` does not allocate a fresh ``lambda`` on every
-    call just to supply this default.
+    ``is_exporting`` was added in torch 2.7, ``is_compiling`` in torch 2.3; each ``getattr(...,
+    _tracer_absent)`` call falls back to this function below its own threshold. Hoisted to a
+    module-level function so the hot decoder ``forward`` does not allocate a fresh ``lambda`` on
+    every call just to supply this default.
 
     Returns:
-        Always ``False`` — torch<2.6 is never in an export trace.
+        Always ``False``.
     """
     return False
 
@@ -337,10 +339,24 @@ def forward(
         # torch.export (ExecuTorch) cannot trace torch._shape_as_tensor — it raises "the tensor has
         # a non-zero number of elements, but its data is not allocated yet". Under that trace build
         # spatial_shapes directly from the concrete Python-int (H, W) pairs instead; ExecuTorch uses
-        # static shapes, so the baked constant is exact. This branch is taken only under
-        # torch.export, leaving the eager and TorchScript-ONNX/TensorRT (#1155) paths untouched.
-        # getattr guards torch<2.6, which lacks is_exporting() and never runs ExecuTorch export.
-        if getattr(torch.compiler, "is_exporting", _not_exporting)():
+        # static shapes, so the baked constant is exact. torch.compile needs the same branch for a
+        # different reason: Dynamo polyfills torch._shape_as_tensor to return a torch.Size, so
+        # torch.stack raises "expected Tensor as element 0 in argument 0, but got torch.Size" and
+        # the compile aborts (suppress_errors=True does not catch it — the TypeError comes from
+        # user code, not from Dynamo). Neither guard is true in eager or under torch.jit.trace, so
+        # the eager and TorchScript-ONNX/TensorRT (#1155) paths keep the _shape_as_tensor form.
+        # getattr guards torch<2.7 (lacks is_exporting()) and torch<2.3 (lacks is_compiling());
+        # both were added to torch.compiler as public API in those releases (verified against
+        # upstream source). Below torch 2.3 the getattr fallback returns False, so compile=True
+        # keeps hitting the same TypeError this branch exists to avoid — there is no substitute:
+        # torch._dynamo.is_compiling() and torch._utils.is_compiling() are hardcoded
+        # `return False` stubs on torch 2.2.x, not real predicates. This project's floor is
+        # torch>=2.2.0 (pyproject.toml), so torch 2.2.x is unaffected by this fix, exactly as it
+        # was already unaffected by is_exporting() before it.
+        if (
+            getattr(torch.compiler, "is_exporting", _tracer_absent)()
+            or getattr(torch.compiler, "is_compiling", _tracer_absent)()
+        ):
             spatial_shapes = torch.as_tensor(spatial_shapes_hw, device=srcs[0].device, dtype=torch.long)
         else:
             spatial_shapes = torch.stack([torch._shape_as_tensor(src)[2:4] for src in srcs]).to(

tests/models/test_transformer_onnx_spatial_shapes.py

@@ -17,15 +17,26 @@
 
 import inspect
 import io
+from typing import TYPE_CHECKING
+from unittest import mock
 
 import numpy as np
 import pytest
 import torch
 from torch import nn
 
-onnx = pytest.importorskip("onnx", reason="onnx not installed; skip ONNX export tests")
+from rfdetr.models.transformer import Transformer
 
-from rfdetr.models.transformer import Transformer  # noqa: E402
+if TYPE_CHECKING:
+    import onnx
+
+# onnx is imported lazily at runtime inside the fixtures that build an ONNX graph
+# (exported_static_onnx, exported_dynamic_onnx_bytes, exported_1lvl_onnx), not at module scope: a
+# module-level pytest.importorskip("onnx") would skip collection of the whole file, including the
+# torch.compile/TorchScript-trace regression tests below, which do not touch onnx at all and must
+# still run in CI environments that install this project without the [onnx] extra (e.g.
+# ci-tests-cpu.yml's "train,augment,cli,visual" set). The TYPE_CHECKING import above only satisfies
+# the "onnx.ModelProto" string annotations used as return/parameter types.
 
 # CI guard: torch._shape_as_tensor is a private ATen API used on the live forward path in
 # Transformer.forward(). If a future PyTorch upgrade removes it, this assertion fails
@@ -218,6 +229,7 @@ def exported_static_onnx(
     Returns:
         Loaded ``onnx.ModelProto`` for structural graph assertions.
     """
+    onnx = pytest.importorskip("onnx", reason="onnx not installed; skip ONNX export tests")
     out = tmp_path_factory.mktemp("onnx_static") / "transformer.onnx"
     torch.onnx.export(
         transformer_wrapper_2lvl,
@@ -241,6 +253,7 @@ def exported_dynamic_onnx_bytes(
     Returns:
         Serialized ONNX model bytes for onnxruntime inference with variable batch sizes.
     """
+    pytest.importorskip("onnx", reason="onnx not installed; skip ONNX export tests")
     buf = io.BytesIO()
     torch.onnx.export(
         transformer_wrapper_2lvl,
@@ -273,6 +286,7 @@ def exported_1lvl_onnx(
     Returns:
         Loaded ``onnx.ModelProto`` for structural graph assertions.
     """
+    onnx = pytest.importorskip("onnx", reason="onnx not installed; skip ONNX export tests")
     out = tmp_path_factory.mktemp("onnx_1lvl") / "transformer_1lvl.onnx"
     torch.onnx.export(
         transformer_wrapper_1lvl,
@@ -407,3 +421,75 @@ def test_level_start_index_correctness_two_levels() -> None:
     assert torch.equal(level_start_index, torch.tensor([0, 48], dtype=torch.long)), (
         f"level_start_index expected [0, 48], got {level_start_index.tolist()}"
     )
+
+
+# ---------------------------------------------------------------------------
+# Tests: torch.compile takes the Python-int branch (Dynamo polyfills _shape_as_tensor)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("wrapper_fixture", "inputs_fixture"),
+    [
+        pytest.param("transformer_wrapper_1lvl", "example_inputs_1lvl", id="1lvl"),
+        pytest.param("transformer_wrapper_2lvl", "example_inputs_2lvl", id="2lvl"),
+    ],
+)
+def test_spatial_shapes_survives_dynamo_shape_as_tensor_polyfill(
+    wrapper_fixture: str,
+    inputs_fixture: str,
+    request: pytest.FixtureRequest,
+) -> None:
+    """``Transformer.forward`` must not call ``torch._shape_as_tensor`` while compiling.
+
+    Dynamo polyfills ``torch._shape_as_tensor`` to return a ``torch.Size`` rather than a tensor, so
+    ``torch.stack([...])`` raises ``TypeError: expected Tensor as element 0 in argument 0, but got torch.Size`` and the
+    whole compile aborts (``suppress_errors=True`` does not catch it: the ``TypeError`` comes from user code, not from
+    Dynamo). This reproduces that polyfill and the ``is_compiling()`` state on CPU, without paying for a real compile in
+    CI, and checks the output still matches the eager one. Both feature-level counts are covered because the two
+    branches build ``spatial_shapes`` from differently shaped sources.
+
+    Args:
+        wrapper_fixture: Name of the Transformer wrapper fixture to exercise.
+        inputs_fixture: Name of the matching example-input fixture.
+        request: Pytest fixture request used to resolve the two names.
+    """
+    wrapper = request.getfixturevalue(wrapper_fixture)
+    inputs = request.getfixturevalue(inputs_fixture)
+    with torch.no_grad():
+        eager = wrapper(*inputs)
+        with (
+            mock.patch.object(torch.compiler, "is_compiling", lambda: True),
+            mock.patch.object(torch, "_shape_as_tensor", lambda t: torch.Size(t.shape)),
+        ):
+            compiling = wrapper(*inputs)
+    assert torch.equal(eager, compiling), "is_compiling() branch changed Transformer output"
+
+
+def test_spatial_shapes_compile_guard_is_false_under_torchscript_trace() -> None:
+    """``is_compiling()`` must stay ``False`` under ``torch.jit.trace``.
+
+    The ``_shape_as_tensor`` form exists because TensorRT accepts the Constant it produces while a
+    ``ScatterND``-producing form is rejected (#1155). If ``is_compiling()`` were true during the TorchScript trace the
+    export path would silently switch formulation.
+    """
+    seen: list[bool] = []
+
+    class _Probe(nn.Module):
+        """Minimal module that records ``is_compiling()`` on every forward call."""
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Record the ``is_compiling()`` state seen during the trace.
+
+            Args:
+                x: Any tensor; only its identity matters here.
+
+            Returns:
+                ``x + 1``, so the traced graph has a real op.
+            """
+            seen.append(bool(getattr(torch.compiler, "is_compiling", lambda: False)()))
+            return x + 1
+
+    probe = _Probe().eval()
+    torch.jit.trace(probe, (torch.randn(1, 3, 4, 4),), check_trace=False)
+    assert seen and not any(seen), "torch.compiler.is_compiling() was true during torch.jit.trace"