[GitHub] [tvm] Hzfengsy commented on a diff in pull request #14639: [Unity][Frontend] Add `no_bind_return_tuple` for PyTorch FX Translator

["Hzfengsy (via GitHub)" <gi...@apache.org>]

Hzfengsy commented on code in PR #14639:
URL: https://github.com/apache/tvm/pull/14639#discussion_r1168153551


##########
tests/python/relax/test_frontend_from_fx.py:
##########
@@ -2707,6 +2707,40 @@ def main(
     tvm.ir.assert_structural_equal(mod, Expected)
 
 
+@tvm.testing.requires_gpu
+def test_no_bind_return_tuple():
+    import torch.fx as fx
+    from torch.nn import Module
+    from tvm.relax.frontend.torch import from_fx
+
+    class Identity(Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, x, y):
+            return (x, y)
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(
+            inp_0: R.Tensor((256, 256), dtype="float32"),
+            inp_1: R.Tensor((256, 256), dtype="float32"),
+        ) -> R.Tuple(R.Tensor((256, 256), dtype="float32"), R.Tensor((256, 256), dtype="float32")):
+            with R.dataflow():
+                gv: R.Tensor((256, 256), dtype="float32") = inp_0
+                gv1: R.Tensor((256, 256), dtype="float32") = inp_1
+                R.output(gv, gv1)
+            return (gv, gv1)
+
+    graph_model = fx.symbolic_trace(Identity())
+    mod = from_fx(
+        graph_model, [([256, 256], "float32"), ([256, 256], "float32")], no_bind_return_tuple=True
+    )
+    mod.show()

Review Comment:
   remove printing :)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org