[SOT] Standardize the from_iterator method of MapVariable and ZipVariable (#71407)

GoldenStain · web-flow · commit d7fa8e4189f8 · 2025-03-05T12:17:24.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py b/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
@@ -646,10 +646,10 @@ def create_zip(*var: VariableBase):
 
 # map
 @Dispatcher.register_decorator(map)
-def create_map(func: CallableVariable, *var: VariableBase):
-    tracked_vars = [func, *var]
+def create_map(fn: CallableVariable, *vars: VariableBase):
+    tracked_vars = [fn, *vars]
     return MapVariable.from_iterator(
-        func, var, graph=Dispatcher.graph, tracker=DummyTracker(tracked_vars)
+        fn, vars, graph=Dispatcher.graph, tracker=DummyTracker(tracked_vars)
     )
 
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py b/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py
@@ -232,41 +232,36 @@ class MapVariable(SequenceIterVariable):
     MapVariable holds a SequenceIterVariable and return a Iterable Variable after map function
     """
 
-    def __init__(self, func, iters, graph, tracker):
-        # iters may contain only one iter.
+    def __init__(self, fn, iters: TupleVariable, graph, tracker):
         super().__init__(iters, graph, tracker)
-        self.func = func
+        self.fn = fn
 
     def next(self):
-        values = []
-        for iter_var in self.hold:
-            next_var = iter_var.next()
-            values.append(next_var)
-        return self.func(*values)
+        return self.fn(*[iter_var.next() for iter_var in self.hold])
 
     def to_list(self) -> list:
-        lists = [iter_vars.to_list() for iter_vars in self.hold]
+        lists = [iter_var.to_list() for iter_var in self.hold]
         min_len = min(len(l) for l in lists)
         result = []
         for i in range(min_len):
-            result.append(self.func(*(l[i] for l in lists)))
+            result.append(self.fn(*(l[i] for l in lists)))
         return result
 
     def has_side_effect(self) -> bool:
-        return self.hold.has_side_effect()
+        return any(iter_var.has_side_effect() for iter_var in self.hold)
 
     def _reconstruct(self, codegen: PyCodeGen):
         if self.has_side_effect():
             super()._reconstruct(codegen)
         else:
             codegen.gen_load_global("map", push_null=True)
-            self.func.reconstruct(codegen)
+            self.fn.reconstruct(codegen)
             self.hold.reconstruct(codegen)
             codegen.gen_call_function(2)
 
     @staticmethod
     def from_iterator(
-        func,
+        fn,
         value: Sequence[VariableBase],
         graph: FunctionGraph | None,
         tracker: Tracker,
@@ -279,7 +274,7 @@ def from_iterator(
                 return UserDefinedIterVariable(value, graph, tracker)
             map_targets.append(iter_variable)
 
-        return MapVariable(func, map_targets, graph, tracker)
+        return MapVariable(fn, map_targets, graph, tracker)
 
 
 # what UserDefinedIterVariable holds doesn't matter, because use user defined iterator will trigger break graph
diff --git a/test/sot/test_builtin_map.py b/test/sot/test_builtin_map.py
@@ -19,7 +19,7 @@
 
 from test_case_base import TestCaseBase, test_with_faster_guard
 
-from paddle import to_tensor
+from paddle import Tensor, to_tensor
 from paddle.jit import sot
 from paddle.jit.sot.psdb import check_no_breakgraph
 from paddle.jit.sot.utils import strict_mode_guard
@@ -100,9 +100,9 @@ def test_map_for_loop(x: list):
 
 
 @check_no_breakgraph
-def test_map_multi_input(func, tensor_, tuple_):
-    x, y, z = map(func, tensor_, tuple_)
-    return x
+def test_map_multi_input(func, a: Tensor, b: tuple[int, ...]):
+    x, y, z = map(func, a, b)
+    return x, y, z
 
 
 @check_no_breakgraph
