Add torch._check hints for torch.export

Summary:
Pull Request resolved: #5468

Fixes #5347

This allows for torch.export.export ImageList.from_tensors.

Reviewed By: wat3rBro

Differential Revision: D74835454

fbshipit-source-id: 0da4c30cb7fdebcab60313e43a41f73f21533d72

Author: yushangdi

detectron2/structures/image_list.py

@@ -6,6 +6,7 @@
 from torch.nn import functional as F
 
 from detectron2.layers.wrappers import move_device_like, shapes_to_tensor
+from detectron2.utils.torch_version_utils import min_torch_version
 
 
 class ImageList:
@@ -111,7 +112,13 @@ def from_tensors(
             # This seems slightly (2%) faster.
             # TODO: check whether it's faster for multiple images as well
             image_size = image_sizes[0]
-            padding_size = [0, max_size[-1] - image_size[1], 0, max_size[-2] - image_size[0]]
+            u0 = max_size[-1] - image_size[1]
+            u1 = max_size[-2] - image_size[0]
+            padding_size = [0, u0, 0, u1]
+            if not torch.jit.is_scripting():
+                if min_torch_version("2.6.0") and torch.compiler.is_compiling():
+                    torch._check(u0.item() >= 0)
+                    torch._check(u1.item() >= 0)
             batched_imgs = F.pad(tensors[0], padding_size, value=pad_value).unsqueeze_(0)
         else:
             # max_size can be a tensor in tracing mode, therefore convert to list

detectron2/utils/testing.py

@@ -8,7 +8,6 @@
 from typing import Callable
 import torch
 import torch.onnx.symbolic_helper as sym_help
-from packaging import version
 from torch._C import ListType
 from torch.onnx import register_custom_op_symbolic
 
@@ -19,6 +18,7 @@
 from detectron2.modeling import build_model
 from detectron2.structures import Boxes, Instances, ROIMasks
 from detectron2.utils.file_io import PathManager
+from detectron2.utils.torch_version_utils import min_torch_version
 
 
 """
@@ -162,20 +162,6 @@ def reload_lazy_config(cfg):
         return LazyConfig.load(fname)
 
 
-def min_torch_version(min_version: str) -> bool:
-    """
-    Returns True when torch's  version is at least `min_version`.
-    """
-    try:
-        import torch
-    except ImportError:
-        return False
-
-    installed_version = version.parse(torch.__version__.split("+")[0])
-    min_version = version.parse(min_version)
-    return installed_version >= min_version
-
-
 def has_dynamic_axes(onnx_model):
     """
     Return True when all ONNX input/output have only dynamic axes for all ranks

detectron2/utils/torch_version_utils.py

@@ -0,0 +1,15 @@
+from packaging import version
+
+
+def min_torch_version(min_version: str) -> bool:
+    """
+    Returns True when torch's  version is at least `min_version`.
+    """
+    try:
+        import torch
+    except ImportError:
+        return False
+
+    installed_version = version.parse(torch.__version__.split("+")[0])
+    min_version = version.parse(min_version)
+    return installed_version >= min_version