[Dy2St] Remove utils_helper (#61236)

SigureMo · web-flow · commit 11f5d3d66d52 · 2024-01-28T00:38:45.000+08:00
diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py
@@ -1889,6 +1889,21 @@ def get_feed_fetch(all_vars, partial_vars):
             )
 
 
+def set_dynamic_shape(variable, shape_list):
+    if paddle.base.dygraph.base.in_to_static_mode():
+        if isinstance(variable, paddle.base.framework.Variable):
+            variable.desc.set_shape(shape_list)
+        elif isinstance(variable, paddle.pir.Value):
+            variable.set_shape(shape_list)
+        else:
+            raise TypeError(
+                "In to_static mode, variable must be a Variable or Value"
+            )
+    else:
+        # in dygraph mode, dynamic shape is not needed, just do nothing.
+        return
+
+
 def get_ast_static_function(function):
     if isinstance(function, SymbolicStaticFunction):
         if function._class_instance:
diff --git a/python/paddle/jit/dy2static/error.py b/python/paddle/jit/dy2static/error.py
@@ -18,12 +18,12 @@
 import sys
 import traceback
 
-import numpy as np  # noqa: F401
+import numpy as np
 
 from .origin_info import Location, OriginInfo, global_origin_info_map
 from .utils import (
     RE_PYMODULE,
-    _is_api_in_module_helper,
+    is_api_in_module_helper,
 )
 
 __all__ = []
@@ -214,8 +214,9 @@ def numpy_api_check(self, format_exception, error_line):
                 func_str = searched_name.group(0)
                 break
         try:
-            fn = eval(func_str)
-            module_result = _is_api_in_module_helper(fn, "numpy")
+            globals = {'np': np}
+            fn = eval(func_str, globals)
+            module_result = is_api_in_module_helper(fn, "numpy")
             is_numpy_api_err = module_result or (
                 func_str.startswith("numpy.") or func_str.startswith("np.")
             )
diff --git a/python/paddle/jit/dy2static/transformers/call_transformer.py b/python/paddle/jit/dy2static/transformers/call_transformer.py
@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from paddle.jit.dy2static.utils import ast_to_source_code, is_paddle_api
+from paddle.jit.dy2static.utils import ast_to_source_code
 from paddle.utils import gast
 
 from ..utils import is_builtin
 from .base import BaseTransformer
+from .utils import is_paddle_api
 
 PDB_SET = "pdb.set_trace"
 
diff --git a/python/paddle/jit/dy2static/transformers/utils.py b/python/paddle/jit/dy2static/transformers/utils.py
@@ -16,9 +16,13 @@
 
 import copy
 
+import numpy as np
+
 from paddle.jit.dy2static.ast_utils import ast_to_source_code
 from paddle.utils import gast
 
+from ..utils import PADDLE_MODULE_PREFIX, is_api_in_module_helper
+
 
 def index_in_list(array_list, item):
     try:
@@ -151,3 +155,37 @@ def get_attribute_full_name(node):
         node, gast.Attribute
     ), "Input non-Attribute node to get attribute full name"
     return ast_to_source_code(node).strip()
+
+
+def is_api_in_module(node, module_prefix):
+    assert isinstance(
+        node, gast.Call
+    ), "Input non-Call node for is_api_in_module"
+
+    # Python can have gast.Call as function, for example: convert_call(func)(x)
+    # We only check the most outside function
+    func_node = node.func
+    while isinstance(func_node, gast.Call):
+        func_node = func_node.func
+
+    func_str = ast_to_source_code(func_node).strip()
+    try:
+        import paddle
+        import paddle.jit.dy2static as _jst
+        from paddle import to_tensor
+
+        globals = {
+            'np': np,
+            'paddle': paddle,
+            '_jst': _jst,
+            'to_tensor': to_tensor,
+        }
+
+        fn = eval(func_str, globals)
+        return is_api_in_module_helper(fn, module_prefix)
+    except Exception:
+        return False
+
+
+def is_paddle_api(node):
+    return is_api_in_module(node, PADDLE_MODULE_PREFIX)
diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py
@@ -39,14 +39,13 @@
 from paddle.utils import flatten, gast
 
 from .ast_utils import ast_to_source_code
-from .utils_helper import (  # noqa: F401
-    PADDLE_MODULE_PREFIX,
-    _is_api_in_module_helper,
-    is_paddle_api,
-)
 
 __all__ = []
 
+# Note(Aurelius): Do not forget the dot `.` to distinguish other
+# module such as paddlenlp.
+PADDLE_MODULE_PREFIX = 'paddle.'
+
 # Note(Aurelius): Do not forget the dot `.` to distinguish other
 # module such as paddlenlp.
 GET_ARGS_FUNC_PREFIX = 'get_args'
@@ -81,9 +80,6 @@
 FOR_CONDITION_PREFIX = 'for_loop_condition'
 FOR_BODY_PREFIX = 'for_loop_body'
 
-GRAD_PREFIX = 'grad/'
-GRAD_SUFFIX = '@GRAD'
-
 NO_SHAPE_VAR_TYPE = [
     core.VarDesc.VarType.READER,
     core.VarDesc.VarType.STEP_SCOPES,
@@ -970,6 +966,11 @@ def prim_is_enabled():
     return core._is_bwd_prim_enabled() or core._is_fwd_prim_enabled()
 
 
+def is_api_in_module_helper(obj, module_prefix):
+    m = inspect.getmodule(obj)
+    return m is not None and m.__name__.startswith(module_prefix)
+
+
 def is_builtin(func, name=None):
     """predict whether a function is a builtin function with name={name}.
     if name == None, then any builtin function will return True
diff --git a/python/paddle/jit/dy2static/utils_helper.py b/python/paddle/jit/dy2static/utils_helper.py
diff --git a/test/dygraph_to_static/test_set_dynamic_shape.py b/test/dygraph_to_static/test_set_dynamic_shape.py
@@ -29,7 +29,7 @@ class TestSetDynamicShape(Dy2StTestBase):
     def test_start(self):
         def dygraph_func(loop_number):
             mask = paddle.randn([2, 2])
-            paddle.jit.dy2static.utils_helper.set_dynamic_shape(mask, [-1, 2])
+            paddle.jit.api.set_dynamic_shape(mask, [-1, 2])
             n = paddle.randn([1, 2])
             for i in range(loop_number):
                 mask = paddle.concat([mask, n], axis=0)
@@ -47,7 +47,7 @@ def dygraph_func(loop_number):
     def test_pir_ast(self):
         def dygraph_func():
             mask = paddle.randn([2, 2])
-            paddle.jit.dy2static.utils_helper.set_dynamic_shape(mask, [-1, 2])
+            paddle.jit.api.set_dynamic_shape(mask, [-1, 2])
             return mask.shape[0]
 
         out = paddle.jit.to_static(dygraph_func)()
