[SOT] Add dispatch for min/max(*args) and add paddle.geometric to paddle API list (#69746)

SigureMo · web-flow · commit 85279af249f1 · 2024-11-27T22:37:14.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
@@ -1242,6 +1242,32 @@ def dispatch_sum(
 )
 
 
+@Dispatcher.register_decorator(max)
+def dispatch_max_star_args(*args: VariableBase):
+    if not args:
+        raise TypeError("max expected at least 1 arguments, got 0")
+    res = args[0]
+    graph = res.graph
+    for arg in args:
+        gt = BuiltinVariable(operator.gt, graph, DanglingTracker())(arg, res)
+        if gt.get_py_value() is True:
+            res = arg
+    return res
+
+
+@Dispatcher.register_decorator(min)
+def dispatch_min_star_args(*args: VariableBase):
+    if not args:
+        raise TypeError("min expected at least 1 arguments, got 0")
+    res = args[0]
+    graph = res.graph
+    for arg in args:
+        lt = BuiltinVariable(operator.lt, graph, DanglingTracker())(arg, res)
+        if lt.get_py_value() is True:
+            res = arg
+    return res
+
+
 # math functions, e.g. math.log, math.sqrt, math.sin, etc.
 def get_math_unary_functions():
     unary_fns = []
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/container.py b/python/paddle/jit/sot/opcode_translator/executor/variables/container.py
@@ -410,7 +410,7 @@ def max(self):
 
     def min(self):
         if len(self) == 0:
-            raise ValueError("max() arg is an empty sequence")
+            raise ValueError("min() arg is an empty sequence")
         res = self[0]
         getitem = BuiltinVariable(
             operator.getitem, self.graph, DanglingTracker()
diff --git a/python/paddle/jit/sot/utils/paddle_api_config.py b/python/paddle/jit/sot/utils/paddle_api_config.py
@@ -48,6 +48,7 @@ def get_paddle_api():
         paddle.fft,
         paddle.vision.ops,
         paddle.metric,
+        paddle.geometric,
     ]
     special_paddle_apis = [paddle.tensor.fill_constant]
     non_operator_related_apis = [
diff --git a/test/sot/test_builtin_dispatch.py b/test/sot/test_builtin_dispatch.py
@@ -114,6 +114,16 @@ def test_ord(x: str):
     return ord(x)
 
 
+@check_no_breakgraph
+def test_min():
+    return min(9, 8, 2, 4, 1, 7, 3, 5, 6)
+
+
+@check_no_breakgraph
+def test_max():
+    return max(9, 8, 2, 4, 1, 7, 3, 5, 6)
+
+
 @check_no_breakgraph
 def test_sqrt(x: int):
     return math.sqrt(x)
@@ -259,6 +269,12 @@ def test_dispatch_sqrt(self):
     def test_dispatch_log(self):
         self.assert_results(test_log, math.e)
 
+    def test_dispatch_min(self):
+        self.assert_results(test_min)
+
+    def test_dispatch_max(self):
+        self.assert_results(test_max)
+
 
 def run_getattr(x: paddle.Tensor):
     attr = 'dtype'

Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,7 @@ def get_paddle_api():`
`48`	`48`	`paddle.fft,`
`49`	`49`	`paddle.vision.ops,`
`50`	`50`	`paddle.metric,`
	`51`	`+ paddle.geometric,`
`51`	`52`	`]`
`52`	`53`	`special_paddle_apis = [paddle.tensor.fill_constant]`
`53`	`54`	`non_operator_related_apis = [`