Enable tuple operands for cond (#108026)

ydwu4 · pytorchmergebot · commit 138e2895d08a · 2023-08-28T00:17:54.000Z
Pull Request resolved: #108026 Approved by: https://github.com/zou3519 ghstack dependencies: #108025
diff --git a/test/dynamo/test_export.py b/test/dynamo/test_export.py
@@ -1681,6 +1681,25 @@ def false_fn(x):
         test_x = torch.randn(3, 2)
         self.assertEqual(out_graph(test_x), mod(test_x))
 
+    def test_export_with_cond_dynamic_shape_pred_tuple_operands(self):
+        from functorch.experimental.control_flow import cond
+
+        class Module(torch.nn.Module):
+            def forward(self, x):
+                def true_fn(x):
+                    return x + x
+
+                def false_fn(x):
+                    return x[:2]
+
+                return cond(x.shape[0] <= 2, true_fn, false_fn, (x,))
+
+        mod = Module()
+        x = torch.randn(2, 2)
+        out_graph, _ = torch._dynamo.export(mod)(x)
+        test_x = torch.randn(3, 2)
+        self.assertEqual(out_graph(test_x), mod(test_x))
+
     def test_export_with_map_cond(self):
         from functorch.experimental.control_flow import cond, map
 
@@ -2899,7 +2918,7 @@ def f_non_list_operands(x):
         example_inputs = (torch.rand(5),)
         with self.assertRaisesRegex(
             torch._dynamo.exc.UserError,
-            "Expected a list but got",
+            "Expected a tuple but got",
         ):
             torch._dynamo.export(
                 f_non_list_operands,
@@ -2916,7 +2935,7 @@ def f_non_tensor_operands(x):
         example_inputs = (torch.rand(5),)
         with self.assertRaisesRegex(
             torch._dynamo.exc.UserError,
-            "Expected a list of tensors",
+            "Expected a tuple of tensors",
         ):
             torch._dynamo.export(
                 f_non_tensor_operands,
diff --git a/torch/_dynamo/variables/higher_order_ops.py b/torch/_dynamo/variables/higher_order_ops.py
@@ -328,18 +328,18 @@ def call_function(
         tx.output.guards.update(args[0].guards)
 
         # operands
-        if type(args[3]) is not ListVariable:
+        if not isinstance(args[3], (ListVariable, TupleVariable)):
             raise UserError(
                 UserErrorType.DYNAMIC_CONTROL_FLOW,
-                f"Expected a list but got {args[3].python_type()}",
+                f"Expected a tuple but got {args[3].python_type()}",
             )
         operands = args[3].unpack_var_sequence(tx)
         if not all(
             isinstance(operand, (TensorVariable, torch.Tensor)) for operand in operands
         ):
             raise UserError(
                 UserErrorType.DYNAMIC_CONTROL_FLOW,
-                "Expected a list of tensors but got {actual_args}".format(  # noqa: UP032
+                "Expected a tuple of tensors but got {actual_args}".format(  # noqa: UP032
                     actual_args=[
                         str(operand.python_type())
                         if isinstance(operand, VariableTracker)
