Custom Op handle 1-element tuples (pytorch#155447)

Fixes pytorch#150472

Modification of [PR 151408](pytorch#151408). This PR modifies the return parsing in `infer_schema` to handle the case of a Tuple with a single element.
Pull Request resolved: pytorch#155447
Approved by: https://github.com/bdhirsh, https://github.com/zou3519

Author: azahed98

test/test_custom_ops.py

@@ -225,6 +225,31 @@ def _(x: torch.Tensor) -> torch.Tensor:
         example = torch.zeros([10, 20], device=device)
         torch.library.opcheck(f, args=[example])
 
+    # https://github.com/pytorch/pytorch/issues/150472
+    def test_single_element_tuple_output(self, device):
+        # Helper function to register id_tuple custom and the fake tensor implementation
+        # so that Dynamo has the fake tensor implementation
+        def get_id_tuple():
+            @torch.library.custom_op("test::id_tuple", mutates_args=[])
+            def id_tuple(x: torch.Tensor) -> Tuple[torch.Tensor]:
+                return (x.clone(),)
+
+            @id_tuple.register_fake
+            def _(
+                x: torch.Tensor,
+            ) -> Tuple[torch.Tensor]:
+                return (x.clone(),)
+
+            return id_tuple
+
+        id_tuple = get_id_tuple()
+        x = torch.randn(3, device=device)
+        ret = id_tuple(x)
+        # Check if ret is a tuple and has exactly one and the same element
+        self.assertIsInstance(ret, tuple)
+        self.assertEqual(len(ret), 1)
+        self.assertEqual(x, ret[0])
+
     def test_missing_abstract_impl(self, device):
         lib = self.lib()
         lib.define("foo(Tensor x) -> Tensor")

torch/_library/infer_schema.py

@@ -77,7 +77,7 @@ def convert_type_string(annotation_type: str):
             )
 
     def unstringify_types(
-        tys: tuple[Union[type[object], str], ...]
+        tys: tuple[Union[type[object], str], ...],
     ) -> tuple[tuple[typing.Any, ...], bool]:
         res = []
         changed = False
@@ -282,8 +282,12 @@ def parse_return(annotation, error_fn):
                 f"Return has unsupported type {annotation}. "
                 f"The valid types are: {SUPPORTED_RETURN_TYPES}."
             )
+    output_ty = ", ".join([SUPPORTED_RETURN_TYPES[arg] for arg in args])
 
-    return "(" + ", ".join([SUPPORTED_RETURN_TYPES[arg] for arg in args]) + ")"
+    # use (()) to represent tuple with single element
+    if len(args) == 1:
+        output_ty = "(" + output_ty + ")"
+    return "(" + output_ty + ")"
 
 
 SUPPORTED_PARAM_TYPES = get_supported_param_types()