[dynamo] update call map to allow multiple input parameters (pytorch#130748)

Fixes pytorch#128072.

Commandeering pytorch#128282 since the issue is now hi pri.

Pull Request resolved: pytorch#130748
Approved by: https://github.com/Skylion007, https://github.com/anijain2305

Author: williamwen42

test/dynamo/test_repros.py

@@ -5327,6 +5327,27 @@ def random_op(tensor, params):
         tensor = torch.randn([2, 3])
         res = random_op(tensor, params)
 
+    # https://github.com/pytorch/pytorch/issues/128072
+    def test_map_with_multiple_args(self):
+        def f(a, b):
+            return a[0] * b[0] + a[1] * b[1]
+
+        def gen_inps(len_x, len_y):
+            x = [torch.randn(5) for _ in range(len_x)]
+            y = [torch.randn(5) for _ in range(len_y)]
+            return x, y
+
+        def g(x, y):
+            return tuple(map(f, x, y))
+
+        opt_g = torch.compile(g, fullgraph=True, backend="eager")
+
+        inps = gen_inps(3, 3)
+        self.assertEqual(g(*inps), opt_g(*inps))
+
+        inps = gen_inps(3, 5)
+        self.assertEqual(g(*inps), opt_g(*inps))
+
 
 instantiate_parametrized_tests(ReproTests)
 

torch/_dynamo/variables/builtin.py

@@ -1470,9 +1470,10 @@ def call_hasattr(self, tx, obj, attr):
                 return variables.ConstantVariable(hasattr(obj.fn, name))
             return obj.call_hasattr(tx, name)
 
-    def call_map(self, tx, fn, seq):
-        if seq.has_unpack_var_sequence(tx):
-            items = [fn.call_function(tx, [x], {}) for x in seq.unpack_var_sequence(tx)]
+    def call_map(self, tx, fn, *seqs):
+        if all(seq.has_unpack_var_sequence(tx) for seq in seqs):
+            unpacked = [seq.unpack_var_sequence(tx) for seq in seqs]
+            items = [fn.call_function(tx, list(args), {}) for args in zip(*unpacked)]
             return variables.TupleVariable(items)
 
     def call_sum(self, tx, seq, start=_SENTINEL):