[SOT] Add any and all method's dispatch for iterable variables (#71254)

GoldenStain · web-flow · commit baf0a3c83db4 · 2025-02-28T18:50:30.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
@@ -1447,6 +1447,43 @@ def numpy_binary_dispatcher(var: NumpyVariable, other: NumpyVariable):
     ),
 )
 
+
+# any for list
+@Dispatcher.register_decorator(any)
+def dispatch_list_any(var: ContainerVariable | IterVariable):
+    graph = var.graph
+    to_bool = BuiltinVariable(bool, graph, DanglingTracker())
+    it = var.get_iter()
+    while True:
+        try:
+            item = it.next()
+            bool_item = to_bool(item)
+            assert isinstance(bool_item, ConstantVariable)
+            if bool_item.get_py_value():
+                return ConstantVariable(True, graph, DummyTracker([var]))
+        except StopIteration:
+            break
+    return ConstantVariable(False, graph, DummyTracker([var]))
+
+
+# all for list
+@Dispatcher.register_decorator(all)
+def dispatch_list_all(var: ContainerVariable | IterVariable):
+    graph = var.graph
+    to_bool = BuiltinVariable(bool, graph, DanglingTracker())
+    it = var.get_iter()
+    while True:
+        try:
+            item = it.next()
+            bool_item = to_bool(item)
+            assert isinstance(bool_item, ConstantVariable)
+            if not bool_item.get_py_value():
+                return ConstantVariable(False, graph, DummyTracker([var]))
+        except StopIteration:
+            break
+    return ConstantVariable(True, graph, DummyTracker([var]))
+
+
 Dispatcher.register(
     np.number.item,
     ("NumpyNumberVariable",),
diff --git a/test/sot/test_builtin_dispatch.py b/test/sot/test_builtin_dispatch.py
@@ -135,6 +135,26 @@ def test_log(x: int):
     return math.log(x)
 
 
+@check_no_breakgraph
+def test_any(var):
+    return any(var)
+
+
+@check_no_breakgraph
+def test_any_iter(var):
+    return any(iter(var))
+
+
+@check_no_breakgraph
+def test_all(var):
+    return all(var)
+
+
+@check_no_breakgraph
+def test_all_iter(var):
+    return all(iter(var))
+
+
 @check_no_breakgraph
 def test_builtin_type_check_eq():
     a = 1
@@ -301,6 +321,32 @@ def test_dispatch_max(self):
     def test_dispatch_builtin_type_check_eq(self):
         self.assert_results(test_builtin_type_check_eq)
 
+    def test_dispatch_any(self):
+        l_pure_true = [1, True, 5, 6]
+        l_pure_false = [False, 0, 0]
+        l_true_and_false = [1, False, 0, 3]
+        d_true = {"a": 1}
+        d_false = {}
+        self.assert_results(test_any, l_pure_true)
+        self.assert_results(test_any, l_pure_false)
+        self.assert_results(test_any, l_true_and_false)
+        self.assert_results(test_any, d_true)
+        self.assert_results(test_any, d_false)
+        self.assert_results(test_any_iter, l_true_and_false)
+
+    def test_dispatch_all(self):
+        l_pure_true = [1, True, 5, 6]
+        l_pure_false = [False, 0, 0]
+        l_true_and_false = [1, False, 0, 3]
+        d_true = {"a": 1}
+        d_false = {}
+        self.assert_results(test_all, l_pure_true)
+        self.assert_results(test_all, l_pure_false)
+        self.assert_results(test_all, l_true_and_false)
+        self.assert_results(test_all, d_true)
+        self.assert_results(test_all, d_false)
+        self.assert_results(test_all_iter, l_true_and_false)
+
 
 def run_getattr(x: paddle.Tensor):
     attr = 'dtype'
