[PIR]Fix co_flags (#66358)

0x45f · web-flow · commit 2f2fb28ee3bb · 2024-07-23T16:45:53.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -121,6 +121,15 @@ def func(x):
     return map_variables(func, inputs)
 
 
+def convert_to_py_value(inputs):
+    def func(x):
+        if isinstance(x, VariableBase):
+            return x.get_py_value()
+        return x
+
+    return map_variables(func, inputs)
+
+
 def record_symbols(SIR, *args, **kwargs):
     symbol_meta_map = {}
     params = set()
diff --git a/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py b/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py
@@ -18,6 +18,7 @@
 
 from __future__ import annotations
 
+import inspect
 import opcode
 import random
 import sys
@@ -465,6 +466,9 @@ def set_function_outputs(self, outputs: list[str]):
 
     def create_function(self) -> types.FunctionType:
         self.update_code_name(self.fn_name, is_resumed_fn=True)
+        self._code_options['co_flags'] &= ~(
+            inspect.CO_VARARGS | inspect.CO_VARKEYWORDS
+        )
         new_code = self.gen_pycode()
         if len(new_code.co_freevars) + len(new_code.co_cellvars) > 0:
             raise FallbackError("Break graph in closure is not support.")
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -731,7 +731,13 @@ def get_py_value(self, allow_tensor=False):
         return self.value
 
     def call_function(self, /, *args, **kwargs):
-        new_object = self.value.__new__(self.value)
+        from ..function_graph import convert_to_py_value
+
+        new_object = self.value.__new__(
+            self.value,
+            *convert_to_py_value(args),
+            **convert_to_py_value(kwargs),
+        )
 
         # do not have init function
         if self.value.__init__ is object.__init__:
