Repository: pytorch/pytorch

Issue:
[ONNX] Support device().type() string comparison with constant
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #86216
* #86169
* __->__ #86168

Fixes `batched_nms` export for torch vision.



Discussion:

Resolving patch:
diff --git a/torch/onnx/_internal/jit_utils.py b/torch/onnx/_internal/jit_utils.py
--- a/torch/onnx/_internal/jit_utils.py
+++ b/torch/onnx/_internal/jit_utils.py
@@ -5,7 +5,8 @@
 
 import dataclasses
 import re
-from typing import Any, Dict, Iterable, Sequence, Tuple, Union
+import typing
+from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union
 
 import torch
 from torch import _C
@@ -306,3 +307,17 @@ def _add_attribute(node: _C.Node, key: str, value: Any, aten: bool):
             else:
                 kind = "i"
     return getattr(node, f"{kind}_")(name, value)
+
+
+# TODO: Expose this to user when migrating symbolic helper functions to here.
+@_beartype.beartype
+def _is_tensor(x: _C.Value) -> bool:
+    return x.type().isSubtypeOf(_C.TensorType.get())
+
+
+@_beartype.beartype
+def get_device_from_value(value: _C.Value) -> Optional[torch.device]:
+    if not _is_tensor(value):
+        return None
+    tensor_type = typing.cast(_C.TensorType, value.type())
+    return tensor_type.device()
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
--- a/torch/onnx/symbolic_opset9.py
+++ b/torch/onnx/symbolic_opset9.py
@@ -205,6 +205,7 @@
     "prim_shape",
     "prim_tolist",
     "prim_tuple_construct",
+    "prim_type",
     "prim_unchecked_cast",
     "prim_uninitialized",
     "rand_like",
@@ -2067,6 +2068,21 @@ def eq(g: jit_utils.GraphContext, self, other):
         # ONNX doesn't have devices, so consider them all to be equal.
         # The no-op check for equality will get constant-folded.
         return g.op("Constant", value_t=torch.tensor(True, dtype=torch.bool))
+    self_node = self.node()
+    other_node = other.node()
+    if self_node.kind() == other_node.kind() == "onnx::Constant":
+        if self_node.kindOf("value") == other_node.kindOf("value") == "s":
+            # Exporting strings to ONNX is not supported.
+            # If both strings are constant, we can compare them directly.
+            # The no-op check for equality will get constant-folded.
+            return g.op(
+                "Constant",
+                value_t=torch.tensor(
+                    self_node.s("value") == other_node.s("value"),
+                    dtype=torch.bool,
+                ),
+            )
+
     return g.op("Equal", self, other)
 
 
@@ -6621,6 +6637,21 @@ def prim_constant(g: jit_utils.GraphContext, *inputs, **attrs):
     )
 
 
+@_onnx_symbolic("prim::type")
+@_beartype.beartype
+def prim_type(g: jit_utils.GraphContext, device_value: _C.Value, *args, **kwargs):
+    if device_value.node().kind() == "prim::device":
+        device = jit_utils.get_device_from_value(device_value.node().input())
+        if device is not None:
+            return g.op("Constant", value_s=str(device))
+
+    return symbolic_helper._unimplemented(
+        "prim::type",
+        "Device type cannot be statically determined.",
+        device_value,
+    )
+
+
 @_onnx_symbolic("onnx::Placeholder")
 @_beartype.beartype
 def onnx_placeholder(g: jit_utils.GraphContext, *inputs, **attrs):


Test patch:
diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
--- a/test/onnx/test_pytorch_onnx_onnxruntime.py
+++ b/test/onnx/test_pytorch_onnx_onnxruntime.py
@@ -9974,11 +9974,7 @@ def forward(self, boxes, scores):
 
         self.run_test(Module(), (boxes, scores))
 
-    @unittest.skip(
-        "Broken in recent TorchVision, see https://github.com/pytorch/pytorch/issues/81121"
-    )
     @skipIfUnsupportedMinOpsetVersion(11)
-    # TODO: Fails with vision 0.13. See #77671
     def test_batched_nms(self):
         num_boxes = 100
         boxes = torch.rand(num_boxes, 4)