chore: import upstream snapshot with attribution

python/paddle/jit/dy2static/transformers/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .transform import DygraphToStaticAst  # noqa: F401

python/paddle/jit/dy2static/transformers/assert_transformer.py

@@ -0,0 +1,46 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.jit.dy2static.utils import ast_to_source_code
+from paddle.utils import gast
+
+from .base import BaseTransformer
+
+__all__ = []
+
+
+class AssertTransformer(BaseTransformer):
+    """
+    A class transforms python assert to convert_assert.
+    """
+
+    def __init__(self, root):
+        self.root = root
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_Assert(self, node):
+        convert_assert_node = (
+            gast.parse(
+                '_jst.Assert({test}, {msg})'.format(
+                    test=ast_to_source_code(node.test),
+                    msg=ast_to_source_code(node.msg) if node.msg else "",
+                )
+            )
+            .body[0]
+            .value
+        )
+
+        return gast.Expr(value=convert_assert_node)

python/paddle/jit/dy2static/transformers/base.py

@@ -0,0 +1,566 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.base import unique_name
+from paddle.jit.dy2static.utils import (
+    ORIGIN_INFO,
+    ast_to_source_code,
+)
+from paddle.utils import gast
+
+from .utils import (
+    FOR_ITER_INDEX_PREFIX,
+    FOR_ITER_ITERATOR_PREFIX,
+    FOR_ITER_TARGET_PREFIX,
+    FOR_ITER_VAR_LEN_PREFIX,
+    FOR_ITER_VAR_NAME_PREFIX,
+    FOR_ITER_ZIP_TO_LIST_PREFIX,
+    create_assign_node,
+)
+
+__all__ = []
+
+
+class BaseTransformer(gast.NodeTransformer):
+    def visit(self, node):
+        if not isinstance(node, gast.AST):
+            msg = f'Expected "gast.AST", but got "{type(node)}".'
+            raise ValueError(msg)
+        origin_info = getattr(node, ORIGIN_INFO, None)
+
+        result = super().visit(node)
+
+        iter_result = result
+        if iter_result is not node and iter_result is not None:
+            if not isinstance(iter_result, (list, tuple)):
+                iter_result = (iter_result,)
+            if origin_info is not None:
+                for n in iter_result:
+                    setattr(n, ORIGIN_INFO, origin_info)
+
+        return result
+
+
+class NameNodeReplaceTransformer(BaseTransformer):
+    """
+    This class replaces specified gast.Name node by replace_node.
+    """
+
+    def __init__(self, root_node, target_name, replace_node):
+        assert isinstance(target_name, str)
+
+        # NOTE(liym27):
+        # Use gast.Name to replace gast.Name, otherwise, errors may occur.
+        #
+        # For examples:
+        # If using a gast.Subscript to replace gast.Name, and the original gast.Name
+        # is in the arguments of FunctionDef, an exception will be raised.
+        #
+        # ```
+        # def func(x[i])) # x[i] can not be a argument
+        #    # ...
+        # ```
+
+        assert isinstance(replace_node, gast.Name)
+        self.target_name = target_name
+        self.replace_node = replace_node
+
+        self.visit(root_node)
+
+    def visit_Name(self, node):
+        if node.id == self.target_name:
+            return self.replace_node
+        return node
+
+    def visit_Nonlocal(self, node):
+        names = node.names
+
+        def replace(s):
+            if s == self.target_name:
+                return self.replace_node.id
+            return s
+
+        node.names = list(map(replace, names))
+        return node
+
+
+class ForLoopTuplePreTransformer(BaseTransformer):
+    """pre-process of for loop.
+    >>> for A in B:
+    >>>    C
+
+    will be changed into :
+
+    >>> # make iterator-only to indexable list.
+    >>> UUID_iterator = _jst.Indexable(B)
+    >>> for UUID_target in UUID_iterator:
+    >>>     A = _jst.Unpack(UUID_target, structure)
+    >>>     C
+
+    make the later loop_transform have unified type:
+    >>> for target in iter:
+    >>>     body
+    """
+
+    def __init__(self, root):
+        self.root = root
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_For(self, node):
+        self.generic_visit(node)
+        tuple_target = unique_name.generate(FOR_ITER_TARGET_PREFIX)
+        tuple_iterator = unique_name.generate(FOR_ITER_ITERATOR_PREFIX)
+        origin_tuple_node = node.target
+        assign_iterator_node = gast.parse(
+            f"{tuple_iterator} = _jst.Indexable({ast_to_source_code(node.iter).strip()})"
+        ).body[0]
+        node.target = gast.Name(
+            id=tuple_target,
+            ctx=gast.Store(),
+            annotation=None,
+            type_comment=None,
+        )
+        node.iter = gast.Name(
+            id=tuple_iterator,
+            ctx=gast.Load(),
+            annotation=None,
+            type_comment=None,
+        )
+        node.body[0:0] = self.tuple_to_stmts(origin_tuple_node, tuple_target)
+        # return a list will insert a list of node replace the original for node.
+        return [assign_iterator_node, node]
+
+    def tuple_node_to_unpack_structure(self, node):
+        """Create a sequence to represents the structure of nest.
+        For example: `a, (b,c), [d,e,f]` is represented by
+        `[1, [1,1], [1,1,1]]`. the `1` is just a notation.
+
+        Specially, `a` is represented by `1`.
+        """
+        ret = []
+        if not isinstance(node, (gast.Tuple, gast.List)):
+            return 1
+        for element in node.elts:
+            ret.append(self.tuple_node_to_unpack_structure(element))
+        return ret
+
+    def tuple_to_stmts(self, node, tuple_name):
+        structure_str = str(self.tuple_node_to_unpack_structure(node))
+        node_str = ast_to_source_code(node).strip()
+        assign_node_str = (
+            f"{node_str} = _jst.Unpack({tuple_name}, {structure_str})"
+        )
+        assign_node = gast.parse(assign_node_str).body[0]
+        return [assign_node]
+
+
+class ForNodeVisitor:
+    """
+    This class parses python for statement, get transformed 3 statement components of for node
+    three key statements:
+        1). init_stmts: list[node], prepare nodes of for loop, may not only one
+        2). cond_stmt: node, condition node to judge whether continue loop
+        3). body_stmts: list[node], updated loop body, sometimes we should change
+            the original statement in body, not just append new statement
+
+    In this process, the semantics of for does not change.
+
+    Now only can parse 3 type statements (Here var is Tensor(Tensor) or python variable):
+        1). for x in range(var[*]|var.numpy()[*])
+        2). for x in var|var.numpy()
+        3). for i, x enumerate(var|var.numpy())
+    """
+
+    def __init__(self, for_node):
+        assert isinstance(for_node, gast.For), (
+            "Input node for the initialization of ForNodeVisitor is not gast.For node."
+        )
+        # 1. original for node
+        self.node = for_node
+
+        # 2. gast.For node main parts
+        self.target = for_node.target
+        # NOTE: type may be Node or list[Node]
+        self.iter_args = (
+            for_node.iter if self.is_for_iter() else for_node.iter.args
+        )
+        self.body = for_node.body
+
+        # 3. key shared node or names
+        # - x:
+        #   - for x in range(***)
+        #   - for x in var|var.numpy()
+        #   - for i, x enumerate(var|var.numpy())
+        self.iter_var_name = self._get_iter_var_name()
+
+        # - created index var to slice Variable: __for_loop_var_index_0
+        #   - for x in var|var.numpy()
+        #   - for i, x enumerate(var|var.numpy())
+        self.iter_idx_name = unique_name.generate(FOR_ITER_INDEX_PREFIX)
+
+        # - created shape var to build loop condition: __for_loop_var_len_0
+        #   - for x in var|var.numpy()
+        #   - for i, x enumerate(var|var.numpy())
+        #   - for x in var
+        self.iter_var_len_name = unique_name.generate(FOR_ITER_VAR_LEN_PREFIX)
+        # - created zip to list var : __for_loop_iter_zip_0
+        self.iter_zip_to_list_name = unique_name.generate(
+            FOR_ITER_ZIP_TO_LIST_PREFIX
+        )
+
+        # - var.numpy()/var
+        #   - for x in var|var.numpy()
+        #   - for i, x enumerate(var|var.numpy())
+        self.iter_node = self._get_iter_node()
+
+        # - enumerate i:
+        #   - for i, x enumerate(var|var.numpy())
+        self.enum_idx_name = self._get_enum_idx_name()
+
+        # - range/enumerate args length
+        self.args_length = None
+
+    def parse(self):
+        self._args_check()
+        if self.is_for_range_iter():
+            return self._parse_for_range_stmts()
+        elif self.is_for_iter():
+            return self._parse_for_stmts()
+        elif self.is_for_enumerate_iter():
+            return self._parse_for_enumerate_stmts()
+        else:
+            return None
+
+    def is_for_range_iter(self):
+        return (
+            isinstance(self.node.iter, gast.Call)
+            and isinstance(self.node.iter.func, gast.Name)
+            and self.node.iter.func.id == "range"
+        )
+
+    def is_for_iter(self):
+        if isinstance(
+            self.node.iter, (gast.Name, gast.Attribute, gast.List, gast.Tuple)
+        ):
+            return True
+        elif (
+            isinstance(self.node.iter, gast.Call)
+            and isinstance(self.node.iter.func, gast.Attribute)
+            and self.node.iter.func.attr == 'numpy'
+        ):
+            return True
+        elif isinstance(self.node.iter, gast.Subscript):
+            return True
+        else:
+            return False
+
+    def is_for_enumerate_iter(self):
+        return (
+            isinstance(self.node.iter, gast.Call)
+            and isinstance(self.node.iter.func, gast.Name)
+            and self.node.iter.func.id == "enumerate"
+        )
+
+    def _args_check(self):
+        if self.is_for_range_iter():
+            self.args_length = len(self.iter_args)
+            assert self.args_length >= 1 and self.args_length <= 3, (
+                "range() function takes 1 to 3 arguments"
+            )
+        elif self.is_for_enumerate_iter():
+            self.args_length = len(self.iter_args)
+            assert self.args_length >= 1 and self.args_length <= 2, (
+                "enumerate() function takes 1 to 2 arguments"
+            )
+        else:
+            self.args_length = None
+
+    def _parse_for_range_stmts(self):
+        init_stmts = []
+        init_stmts.append(self._build_index_init_node())
+
+        compare_node = self._build_compare_node()
+        step_node = self._build_step_node()
+        cond_stmt = self._build_cond_stmt(step_node, compare_node)
+
+        body_stmts = self.body
+        body_stmts.append(self._build_index_increase_node(step_node))
+
+        return init_stmts, cond_stmt, body_stmts
+
+    def _parse_for_stmts(self):
+        init_stmts = []
+        init_stmts.extend(self._build_iter_node())
+        init_stmts.append(self._build_index_init_node())
+        init_stmts.append(self._build_var_len_assign_node())
+
+        compare_node = self._build_compare_node()
+        step_node = self._build_step_node()
+        cond_stmt = self._build_cond_stmt(step_node, compare_node)
+
+        body_stmts = self.body
+
+        # NOTE(liym27): Here add a gast.Assign, and the target of it is gast.Name.
+        # In NameNodeReplaceTransformer, using gast.Name to replace gast.Name is safe.
+        target_node, assign_node = self._build_assign_var_slice_node()
+        body_stmts[0:0] = [assign_node]
+        for body_node in body_stmts:
+            NameNodeReplaceTransformer(
+                body_node, self.iter_var_name, target_node
+            )
+        body_stmts.append(self._build_index_increase_node(step_node))
+
+        return init_stmts, cond_stmt, body_stmts
+
+    def _parse_for_enumerate_stmts(self):
+        init_stmts = []
+        init_stmts.extend(self._build_iter_node())
+        init_stmts.append(self._build_index_init_node())
+        init_stmts.append(self._build_var_len_assign_node())
+        init_stmts.append(self._build_enum_init_node())
+
+        compare_node = self._build_compare_node()
+        step_node = self._build_step_node()
+        cond_stmt = self._build_cond_stmt(step_node, compare_node)
+
+        body_stmts = self.body
+
+        target_node, assign_node = self._build_assign_var_slice_node()
+        body_stmts[0:0] = [assign_node]
+        for body_node in body_stmts:
+            NameNodeReplaceTransformer(
+                body_node, self.iter_var_name, target_node
+            )
+
+        body_stmts.append(self._build_index_increase_node(step_node))
+        body_stmts.append(self._build_enum_increase_node())
+
+        return init_stmts, cond_stmt, body_stmts
+
+    def _build_index_init_node(self):
+        if self.is_for_range_iter():
+            if self.args_length == 1:
+                index_init_value_str = '0'
+            else:
+                index_init_value_str = ast_to_source_code(
+                    self.iter_args[0]
+                ).strip()
+
+            index_init_var_name = self.iter_var_name
+        else:
+            index_init_value_str = '0'
+            index_init_var_name = self.iter_idx_name
+
+        index_init_node_source_str = (
+            f"{index_init_var_name} = {index_init_value_str}"
+        )
+
+        index_init_node = gast.parse(index_init_node_source_str).body[0]
+
+        return index_init_node
+
+    def _build_var_len_assign_node(self):
+        # get the length of iterable variable
+        if (
+            isinstance(self.iter_node, gast.Call)
+            and isinstance(self.iter_node.func, gast.Attribute)
+            and self.iter_node.func.attr == 'numpy'
+        ):
+            iter_var_name = ast_to_source_code(
+                self.iter_node.func.value
+            ).strip()
+        else:
+            iter_var_name = ast_to_source_code(self.iter_node).strip()
+
+        convert_len_node_source_str = (
+            f'{self.iter_var_len_name} = _jst.Len({iter_var_name})'
+        )
+
+        convert_len_node = gast.parse(convert_len_node_source_str).body[0]
+
+        return convert_len_node
+
+    def _build_iter_node(self):
+        """
+        Process special cases for iter_node include:
+          - Case 1 (for zip):
+
+            - for i, val in enumerate(zip(x, y))  # original code:
+
+            - __for_loop_iter_zip_0 = list(zip(x, y))
+            - for i, val in enumerate(__for_loop_iter_zip_0)
+        """
+        new_nodes = []
+        if isinstance(self.iter_node, gast.Call) and isinstance(
+            self.iter_node.func, gast.Name
+        ):
+            if self.iter_node.func.id == 'zip':
+                iter_var_name = ast_to_source_code(self.iter_node).strip()
+                zip_to_list_str = (
+                    f"{self.iter_zip_to_list_name} = list({iter_var_name})"
+                )
+                zip_to_list_node = gast.parse(zip_to_list_str).body[0]
+                new_nodes.append(zip_to_list_node)
+
+                self.iter_node = gast.Name(
+                    id=self.iter_zip_to_list_name,
+                    ctx=gast.Load(),
+                    annotation=None,
+                    type_comment=None,
+                )
+
+        return new_nodes
+
+    def _build_enum_init_node(self):
+        if self.is_for_enumerate_iter() and self.args_length != 1:
+            init_value_str = ast_to_source_code(self.iter_args[1]).strip()
+        else:
+            init_value_str = '0'
+
+        enum_init_node_source_str = f"{self.enum_idx_name} = {init_value_str}"
+        enum_init_node = gast.parse(enum_init_node_source_str).body[0]
+        return enum_init_node
+
+    def _build_compare_node(self):
+        if self.is_for_range_iter():
+            compare_node = (
+                self.iter_args[0]
+                if self.args_length == 1
+                else self.iter_args[1]
+            )
+        else:
+            compare_node = gast.Name(
+                id=self.iter_var_len_name,
+                ctx=gast.Load(),
+                annotation=None,
+                type_comment=None,
+            )
+        return compare_node
+
+    def _build_step_node(self):
+        if self.is_for_range_iter():
+            step_node = (
+                self.iter_args[2]
+                if self.args_length == 3
+                else gast.Constant(value=1, kind=None)
+            )
+        else:
+            step_node = gast.Constant(value=1, kind=None)
+        return step_node
+
+    def _build_cond_stmt(self, step_node, compare_node):
+        if not isinstance(step_node, (gast.Constant, gast.UnaryOp)):
+            raise NotImplementedError(
+                "Dynamic-to-Static only supports the step value is a constant or negative constant in 'for-range' statements, "
+                f"such as '2', '-3'. But received: '{ast_to_source_code(step_node).strip()}'. Please fix code to be compatible with Dynamic-to-Static."
+            )
+
+        if isinstance(step_node, gast.UnaryOp) or step_node.value < 0:
+            # eg:
+            # range(max, min, -2)
+            # ->
+            # i > min
+            return gast.Compare(
+                left=gast.Name(
+                    id=(
+                        self.iter_var_name
+                        if self.is_for_range_iter()
+                        else self.iter_idx_name
+                    ),
+                    ctx=gast.Load(),
+                    annotation=None,
+                    type_comment=None,
+                ),
+                ops=[gast.Gt()],
+                comparators=[compare_node],
+            )
+        else:
+            # eg:
+            # range(min, max, 2)
+            # ->
+            # i < max
+            return gast.Compare(
+                left=gast.Name(
+                    id=(
+                        self.iter_var_name
+                        if self.is_for_range_iter()
+                        else self.iter_idx_name
+                    ),
+                    ctx=gast.Load(),
+                    annotation=None,
+                    type_comment=None,
+                ),
+                ops=[gast.Lt()],
+                comparators=[compare_node],
+            )
+
+    def _build_index_increase_node(self, step_node):
+        return gast.AugAssign(
+            target=gast.Name(
+                id=(
+                    self.iter_var_name
+                    if self.is_for_range_iter()
+                    else self.iter_idx_name
+                ),
+                ctx=gast.Store(),
+                annotation=None,
+                type_comment=None,
+            ),
+            op=gast.Add(),
+            value=step_node,
+        )
+
+    def _build_assign_var_slice_node(self):
+        var_slice_str = f"{ast_to_source_code(self.iter_node).strip()}[{self.iter_idx_name}]"
+        var_slice_node = gast.parse(var_slice_str).body[0].value
+        new_iter_var_name = unique_name.generate(FOR_ITER_VAR_NAME_PREFIX)
+        target_node, assign_node = create_assign_node(
+            new_iter_var_name, var_slice_node
+        )
+        return target_node, assign_node
+
+    def _build_enum_increase_node(self):
+        return gast.AugAssign(
+            target=gast.Name(
+                id=self.enum_idx_name,
+                ctx=gast.Store(),
+                annotation=None,
+                type_comment=None,
+            ),
+            op=gast.Add(),
+            value=gast.Constant(value=1, kind=None),
+        )
+
+    def _get_iter_var_name(self):
+        if self.is_for_range_iter():
+            return self.target.id
+        elif self.is_for_iter():
+            return self.target.id
+        elif self.is_for_enumerate_iter():
+            return self.target.elts[1].id
+        return None
+
+    def _get_iter_node(self):
+        if self.is_for_iter():
+            return self.iter_args
+        elif self.is_for_enumerate_iter():
+            return self.iter_args[0]
+        return None
+
+    def _get_enum_idx_name(self):
+        if self.is_for_enumerate_iter():
+            return self.target.elts[0].id
+        return None

python/paddle/jit/dy2static/transformers/break_continue_transformer.py

@@ -0,0 +1,419 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.base import unique_name
+from paddle.utils import gast
+
+from .base import BaseTransformer, ForNodeVisitor
+from .utils import BaseNodeVisitor, create_bool_node, index_in_list
+
+__all__ = []
+
+BREAK_NAME_PREFIX = '__break'
+CONTINUE_NAME_PREFIX = '__continue'
+
+
+class ForToWhileTransformer(BaseTransformer):
+    """
+    Transform python for loop into while loop and add condition node in the
+    loop test
+    """
+
+    def __init__(self, parent_node, loop_node, condition_node):
+        assert isinstance(loop_node, gast.For), (
+            "loop_node is not gast.For in ForToWhileTransformer"
+        )
+        self.parent_node = parent_node
+        self.loop_node = loop_node
+        self.condition_node = condition_node
+
+    def transform(self):
+        if hasattr(self.parent_node, 'body'):
+            body_list = self.parent_node.body
+            i = index_in_list(body_list, self.loop_node)
+            if i != -1:
+                new_stmts = self.get_for_stmt_nodes(body_list[i])
+                body_list[i : i + 1] = new_stmts
+                i += len(new_stmts)
+                return new_stmts
+        if hasattr(self.parent_node, 'orelse'):
+            body_list = self.parent_node.orelse
+            i = index_in_list(body_list, self.loop_node)
+            if i != -1:
+                new_stmts = self.get_for_stmt_nodes(body_list[i])
+                body_list[i : i + 1] = new_stmts
+                i += len(new_stmts)
+                return new_stmts
+        raise ValueError(
+            "parent_node doesn't contain the loop_node in ForToWhileTransformer"
+        )
+
+    def get_for_stmt_nodes(self, node):
+        assert isinstance(node, gast.For), (
+            "Input node is NOT gast.For in get_for_stmt_nodes"
+        )
+
+        # 1. parse current gast.For node
+        current_for_node_parser = ForNodeVisitor(node)
+        stmts_tuple = current_for_node_parser.parse()
+        if stmts_tuple is None:
+            return [node]
+        init_stmts, cond_stmt, body_stmts = stmts_tuple
+
+        # 2. append break statement
+        new_cond_stmt = gast.BoolOp(
+            op=gast.And(), values=[cond_stmt, self.condition_node]
+        )
+
+        # 3. construct gast.While node
+        while_node = gast.While(
+            test=new_cond_stmt, body=body_stmts, orelse=node.orelse
+        )
+        init_stmts.append(while_node)
+        return init_stmts
+
+
+class BreakContinueTransformer(BaseNodeVisitor):
+    """
+    Rewrite 'break' and 'continue' key words in a if-else python way to make
+    it equivalent to original control flow
+
+    The main idea of this class is:
+
+        1. Map the 'break/continue' stmt with an unique boolean variable V.
+
+        2. Find the first ancestor block containing this 'break/continue', a
+        block can be a node containing stmt list. We should remove all stmts
+        after the 'break/continue' and set the V to True here.
+
+        3. Add 'if V' for stmts in ancestor blocks between the first one
+        (exclusive) and the ancestor loop (inclusive)
+
+        4. For 'break' add break into condition of the loop. For 'continue',
+        set continue to False at the beginning of each loop
+
+        TODO: more details should be summarized as design document
+
+    Note: The class is inherited from BaseNodeVisitor instead of NodeTransformer,
+          because ancestor nodes will be modified inplace for `Break/Continue` here.
+          In general, we recommend to inheriting NodeTransformer to modify node!
+    """
+
+    def __init__(self, root):
+        super().__init__()
+
+        self.root = root
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_Break(self, node):
+        function_def_node_index = _find_ancestor_function_def_index(
+            self.ancestor_nodes
+        )
+        loop_node_index = _find_ancestor_loop_index(node, self.ancestor_nodes)
+        assert loop_node_index != -1, "SyntaxError: 'break' outside loop"
+        loop_node = self.ancestor_nodes[loop_node_index]
+        function_def_node = self.ancestor_nodes[function_def_node_index]
+
+        # 1. Map the 'break/continue' stmt with an unique boolean variable V.
+        variable_name = unique_name.generate(BREAK_NAME_PREFIX)
+
+        # 2. Find the first ancestor block containing this 'break/continue', a
+        # block can be a node containing stmt list. We should remove all stmts
+        # after the 'break/continue' and set the V to True here.
+        first_block_index = self._remove_stmts_after_break_continue(
+            node, variable_name, loop_node_index
+        )
+
+        # 3. Add 'if not V' for stmts in ancestor blocks between the first one
+        # (exclusive) and the ancestor loop (inclusive)
+        self._replace_if_stmt(loop_node_index, first_block_index, variable_name)
+
+        # 4. For 'break' add break into condition of the loop.
+        assign_false_node = create_bool_node(variable_name, False)
+        function_def_node.body.insert(0, assign_false_node)
+
+        cond_var_node = gast.UnaryOp(
+            op=gast.Not(),
+            operand=gast.Name(
+                id=variable_name,
+                ctx=gast.Load(),
+                annotation=None,
+                type_comment=None,
+            ),
+        )
+
+        if isinstance(loop_node, gast.While):
+            loop_node.test = gast.BoolOp(
+                op=gast.And(), values=[loop_node.test, cond_var_node]
+            )
+        elif isinstance(loop_node, gast.For):
+            parent_node = self.ancestor_nodes[loop_node_index - 1]
+            for_to_while = ForToWhileTransformer(
+                parent_node, loop_node, cond_var_node
+            )
+            for_to_while.transform()
+
+    def visit_Continue(self, node):
+        function_def_node_index = _find_ancestor_function_def_index(
+            self.ancestor_nodes
+        )
+        loop_node_index = _find_ancestor_loop_index(node, self.ancestor_nodes)
+        assert loop_node_index != -1, "SyntaxError: 'continue' outside loop"
+        loop_node = self.ancestor_nodes[loop_node_index]
+        function_def_node = self.ancestor_nodes[function_def_node_index]
+
+        # 1. Map the 'break/continue' stmt with an unique boolean variable V.
+        variable_name = unique_name.generate(CONTINUE_NAME_PREFIX)
+
+        # 2. Find the first ancestor block containing this 'break/continue', a
+        # block can be a node containing stmt list. We should remove all stmts
+        # after the 'break/continue' and set the V to True here.
+        first_block_index = self._remove_stmts_after_break_continue(
+            node, variable_name, loop_node_index
+        )
+
+        # 3. Add 'if not V' for stmts in ancestor blocks between the first one
+        # (exclusive) and the ancestor loop (inclusive)
+        self._replace_if_stmt(loop_node_index, first_block_index, variable_name)
+
+        # 4. For 'continue', set continue to False at the beginning of each loop
+        assign_false_node = create_bool_node(variable_name, False)
+        loop_node.body.insert(0, assign_false_node)
+        # Add a same assign statement to the beginning of function body to avoid
+        # generate the UndefinedVar
+        function_def_node.body.insert(0, assign_false_node)
+
+    def _remove_stmts_after_break_continue(
+        self, break_continue_node, break_continue_name, loop_node_index
+    ):
+        for first_block_index in range(
+            len(self.ancestor_nodes) - 1, loop_node_index - 1, -1
+        ):
+            first_block = self.ancestor_nodes[first_block_index]
+            if hasattr(
+                first_block, "body"
+            ) and self._replace_break_continue_in_stmt_list(
+                first_block.body, break_continue_node, break_continue_name
+            ):
+                return first_block_index
+
+            if hasattr(
+                first_block, "orelse"
+            ) and self._replace_break_continue_in_stmt_list(
+                first_block.orelse, break_continue_node, break_continue_name
+            ):
+                return first_block_index
+
+        return first_block_index
+
+    def _replace_if_stmt(
+        self, loop_node_index, first_block_index, break_continue_name
+    ):
+        for i in range(first_block_index - 1, loop_node_index - 1, -1):
+            cur_node = self.ancestor_nodes[i]
+            son_node = self.ancestor_nodes[i + 1]
+            if hasattr(
+                cur_node, 'body'
+            ) and self._replace_after_node_to_if_in_stmt_list(
+                cur_node.body, son_node, break_continue_name
+            ):
+                continue
+            if hasattr(
+                cur_node, 'orelse'
+            ) and self._replace_after_node_to_if_in_stmt_list(
+                cur_node.orelse, son_node, break_continue_name
+            ):
+                continue
+
+    def _replace_break_continue_in_stmt_list(
+        self, stmt_list, break_continue_node, break_continue_name
+    ):
+        i = index_in_list(stmt_list, break_continue_node)
+        if i == -1:
+            return False
+        assign_true_node = create_bool_node(break_continue_name, True)
+        stmt_list[i:] = [assign_true_node]
+        return True
+
+    def _replace_after_node_to_if_in_stmt_list(
+        self, stmt_list, node, break_continue_name
+    ):
+        i = index_in_list(stmt_list, node)
+        if i == -1:
+            return False
+
+        if i == len(stmt_list) - 1:
+            # No need to add, we consider this as added successfully
+            return True
+
+        if_stmt = gast.If(
+            test=gast.UnaryOp(
+                op=gast.Not(),
+                operand=gast.Name(
+                    id=break_continue_name,
+                    ctx=gast.Store(),
+                    annotation=None,
+                    type_comment=None,
+                ),
+            ),
+            body=stmt_list[i + 1 :],
+            orelse=[],
+        )
+        stmt_list[i + 1 :] = []
+        stmt_list.append(if_stmt)
+        return True
+
+    def _add_stmt_before_cur_node(self, cur_node_index, stmt_node):
+        cur_node = self.ancestor_nodes[cur_node_index]
+        parent_node = self.ancestor_nodes[cur_node_index - 1]
+        if hasattr(
+            parent_node, "body"
+        ) and self._add_stmt_into_list_before_node(
+            parent_node.body, cur_node, stmt_node
+        ):
+            return True
+        if hasattr(
+            parent_node, "orelse"
+        ) and self._add_stmt_into_list_before_node(
+            parent_node.orelse, cur_node, stmt_node
+        ):
+            return True
+        return False
+
+    def _add_stmt_into_list_before_node(self, stmt_list, node, stmt_node):
+        i = index_in_list(stmt_list, node)
+        if i == -1:
+            return False
+        stmt_list.insert(i, stmt_node)
+        return True
+
+
+def _find_ancestor_loop_index(node, ancestor_nodes):
+    for i in range(len(ancestor_nodes) - 1, -1, -1):
+        if isinstance(ancestor_nodes[i], (gast.For, gast.While)):
+            return i
+    return -1
+
+
+def _find_ancestor_function_def_index(ancestor_nodes):
+    for i in range(len(ancestor_nodes) - 1, -1, -1):
+        if isinstance(ancestor_nodes[i], gast.FunctionDef):
+            return i
+    return -1
+
+
+class BreakTransformOptimizer(BaseNodeVisitor):
+    """
+    In specific pattern, the transformed code could be optimized by joining the
+    If.test with while.test.
+
+    Currently supported pattern is:
+    ```
+        while cond1:            while cond1 and not cond2:
+            if cond2:    --->       do_something()
+                break
+            do_something()
+    ```
+
+    See following example:
+
+    >>> def foo(x):
+    ...     i = paddle.to_tensor(1, dtype='int32')
+    ...     while i < 10:
+    ...         if x.mean() > 5:
+    ...             break
+    ...         x += i
+    ...         i += 1
+    ...     return x
+
+    The generated code after applying optimization will be:
+    ```
+        def foo(x):
+            i = paddle.to_tensor(1, dtype='int32')
+            while i < 10 and not x.mean() > 5:
+                x += i
+                i += 1
+            return x
+    ```
+    It can avoid wrapping all ops after `break` statement into `cond_op` that
+    usually brings very heavy overhead.
+    """
+
+    def __init__(self, root):
+        super().__init__()
+
+        self.root = root
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_Break(self, node):
+        loop_node_index = _find_ancestor_loop_index(node, self.ancestor_nodes)
+        assert loop_node_index != -1, "SyntaxError: 'break' outside loop"
+        loop_node = self.ancestor_nodes[loop_node_index]
+
+        if self._is_break_cond_pattern(node, loop_node):
+            cond_var_node = self._join_with_while_cond(node, loop_node)
+
+            if isinstance(loop_node, gast.While):
+                loop_node.test = gast.BoolOp(
+                    op=gast.And(), values=[loop_node.test, cond_var_node]
+                )
+            elif isinstance(loop_node, gast.For):
+                parent_node = self.ancestor_nodes[loop_node_index - 1]
+                for_to_while = ForToWhileTransformer(
+                    parent_node, loop_node, cond_var_node
+                )
+                for_to_while.transform()
+
+    def _is_break_cond_pattern(self, break_node, loop_node):
+        """
+        Judge whether if match the pattern to join `If.test` with `while.test`
+        """
+        # while/for -> if -> break
+        if len(self.ancestor_nodes) < 3 or self.ancestor_nodes[-3] != loop_node:
+            return False
+
+        assert self.ancestor_nodes[-1] == break_node
+        parent_if_node = self.ancestor_nodes[-2]
+
+        is_matched = False
+        if isinstance(parent_if_node, gast.If):
+            # gast.If only contains `break`
+            break_first_in_if = (
+                parent_if_node.body[0] == break_node
+                and len(parent_if_node.orelse) == 0
+            )
+            # gast.If is first node of loop_node
+            if_first_in_loop = loop_node.body[0] == parent_if_node
+
+            is_matched = if_first_in_loop and break_first_in_if
+
+        return is_matched
+
+    def _join_with_while_cond(self, break_node, loop_node):
+        """
+        Join the `If.test` with `While.test` together.
+        """
+        parent_if_node = self.ancestor_nodes[-2]
+
+        cond_var_node = gast.UnaryOp(op=gast.Not(), operand=parent_if_node.test)
+
+        # remove the gast.If node that contains the gast.Break.
+        assert loop_node.body[0] == parent_if_node
+        loop_node.body.pop(0)
+
+        return cond_var_node

python/paddle/jit/dy2static/transformers/call_transformer.py

@@ -0,0 +1,81 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.utils import gast
+
+from ..utils import ast_to_source_code, is_builtin
+from .base import BaseTransformer
+from .utils import is_paddle_api
+
+PDB_SET = "pdb.set_trace"
+
+__all__ = []
+
+
+class CallTransformer(BaseTransformer):
+    """
+    This class transforms function calls into Static Graph Ast.
+    """
+
+    def __init__(self, root):
+        self.root = root
+
+    def _no_need_convert_call(self, node):
+        """
+        Determines whether a function needs to be transformed by `convert_call`.
+        It doesn't need to be transformed when a function satisfies the following conditions:
+          1. It's a api of paddle
+          2. It's a python builtin function not include `len`, `zip`, `range` and `enumerate`
+        """
+        assert isinstance(node, gast.Call)
+        if is_paddle_api(node):
+            return True
+
+        func_str = ast_to_source_code(node.func).strip()
+        try:
+            need_convert_builtin_func_list = {
+                'len',
+                'zip',
+                'range',
+                'enumerate',
+                'print',
+            }
+            fn = eval(func_str)
+            is_builtin_fn = is_builtin(fn)
+            need_convert = func_str in need_convert_builtin_func_list
+            return is_builtin_fn and not need_convert
+        except Exception:
+            return False
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_Call(self, node):
+        self.generic_visit(node)
+
+        if self._no_need_convert_call(node):
+            return node
+
+        func_str = ast_to_source_code(node.func).strip()
+
+        # NOTE(liym27): Don't convert `pad.set_trace` even if the conversion doesn't work finally, because
+        # it is clearer to see where it is called from.
+        if PDB_SET in func_str:
+            return node
+
+        new_func_str = f"_jst.Call({func_str})"
+        new_func_ast = gast.parse(new_func_str).body[0].value
+        node.func = new_func_ast
+
+        return node

python/paddle/jit/dy2static/transformers/cast_transformer.py

@@ -0,0 +1,44 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.jit.dy2static.utils import ast_to_source_code
+from paddle.utils import gast
+
+from .base import BaseTransformer
+
+__all__ = []
+
+
+class CastTransformer(BaseTransformer):
+    """
+    This class transforms type casting into Static Graph Ast.
+    """
+
+    def __init__(self, root):
+        self.root = root
+        self._castable_type = {'bool', 'int', 'float', 'complex'}
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_Call(self, node):
+        self.generic_visit(node)
+        func_str = ast_to_source_code(node.func).strip()
+        if func_str in self._castable_type and len(node.args) > 0:
+            args_str = ast_to_source_code(node.args[0]).strip()
+            new_func_str = f"_jst.AsDtype({args_str}, '{func_str}')"
+            new_node = gast.parse(new_func_str).body[0].value
+            return new_node
+
+        return node

python/paddle/jit/dy2static/transformers/create_variable_transformer.py

@@ -0,0 +1,41 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .base import BaseTransformer
+from .utils import FunctionNameLivenessAnalysis, create_undefined_var
+
+__all__ = []
+
+
+class CreateVariableTransformer(BaseTransformer):
+    """ """
+
+    def __init__(self, root):
+        self.root = root
+        FunctionNameLivenessAnalysis(self.root)
+
+    def transform(self):
+        """
+        Main function to transform AST.
+        """
+        self.visit(self.root)
+
+    def visit_FunctionDef(self, node):
+        # attributes = set(filter(lambda x: '.' in x, node.pd_scope.modified_vars()))
+        self.generic_visit(node)
+        bodys = node.body
+        names = sorted(node.pd_scope.created_vars())
+        for name in names:
+            bodys[0:0] = [create_undefined_var(name)]
+        return node

python/paddle/jit/dy2static/transformers/decorator_transformer.py

@@ -0,0 +1,140 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+import warnings
+
+from paddle.utils import gast
+
+from ..utils import RE_PYMODULE, RE_PYNAME, ast_to_source_code
+from .base import BaseTransformer
+
+__all__ = []
+
+IGNORE_NAMES = [
+    'declarative',
+    'to_static',
+    'wraps',
+    'staticmethod',
+    'classmethod',
+    'decorator',
+    'inference',
+]
+
+
+class DecoratorTransformer(BaseTransformer):
+    """
+    Transform decorators.
+    """
+
+    def __init__(self, root):
+        self.root = root
+
+    def transform(self):
+        """
+        Main function to transform AST.
+        """
+        self.visit(self.root)
+
+    def visit_FunctionDef(self, node):
+        assert isinstance(node, gast.FunctionDef)
+        self.generic_visit(node)
+
+        deco_list = node.decorator_list
+        node.decorator_list = []
+
+        # every decorator will append a node
+        decofun_nodes = []
+        # func to be decoded next time
+        deco_target = '_orig_' + node.name
+        # last decoded func
+        decoded_func = ''
+
+        for deco in reversed(deco_list):
+            # skip IGNORE_NAMES
+            deco_full_name = ast_to_source_code(deco).strip()
+            if isinstance(deco, gast.Call):
+                # match case like :
+                # 1: @_jst.Call(a.b.c.d.deco)()
+                # 2: @q.w.e.r.deco()
+                re_tmp = re.match(
+                    rf'({RE_PYMODULE})*({RE_PYNAME}\(){{0,1}}({RE_PYMODULE})*({RE_PYNAME})(\)){{0,1}}\(.*$',
+                    deco_full_name,
+                )
+                deco_name = re_tmp.group(4)
+            else:
+                # match case like:
+                # @a.d.g.deco
+                re_tmp = re.match(
+                    rf'({RE_PYMODULE})*({RE_PYNAME})$',
+                    deco_full_name,
+                )
+                deco_name = re_tmp.group(2)
+            if deco_name in IGNORE_NAMES:
+                continue
+            elif deco_name == 'contextmanager':
+                warnings.warn(
+                    "Dy2Static : A context manager decorator is used, this may not work correctly after transform."
+                )
+
+            decoded_func = '_decoedby_' + deco_name
+
+            # get function after decoration
+            if isinstance(deco, gast.Call):
+                if '_jst.Call' in deco_full_name:
+                    # in this case , the deco_full_name will be like:
+                    # '_jst.Call(deco)(5)'
+                    rematch = re.match(
+                        r'\_jst\.Call\((.+?)\)\((.*)\)', deco_full_name
+                    )
+                    re_name = rematch.group(1)
+                    re_args = rematch.group(2)
+                    re_args_with_func = deco_target + ', ' + re_args
+                    decofun_str = f'try:\n\t{decoded_func} = _jst.Call({re_name})({re_args_with_func})\nexcept:\n\t{decoded_func} = _jst.Call({re_name})({re_args})({deco_target})'
+                else:
+                    # paddle api will not be transformed to '_jst.Call'
+                    rematch = re.match(r'(.+?)\((.*)\)', deco_full_name)
+                    re_name = rematch.group(1)
+                    re_args = rematch.group(2)
+                    re_args_with_func = deco_target + ', ' + re_args
+                    decofun_str = f'try:\n\t{decoded_func} = {re_name}({re_args_with_func})\nexcept:\n\t{decoded_func} = {re_name}({re_args})({deco_target})'
+
+            else:
+                decofun_str = f'{decoded_func} = _jst.Call({deco_full_name})({deco_target})'
+
+            decofun_nodes.extend(gast.parse(decofun_str).body)
+            deco_target = decoded_func
+
+        if not decofun_nodes:
+            return node
+
+        orig_func_node = gast.FunctionDef(
+            name='_orig_' + node.name,
+            args=node.args,
+            body=node.body,
+            decorator_list=[],
+            returns=None,
+            type_comment=None,
+            type_params=[],
+        )
+
+        args = [arg.id for arg in node.args.args]
+        arg_str = ','.join(args)
+        callfun_str = f'return {decoded_func}({arg_str})'
+        callfun_node = gast.parse(callfun_str).body[0]
+
+        node.body = [orig_func_node, *decofun_nodes, callfun_node]
+
+        return node

python/paddle/jit/dy2static/transformers/early_return_transformer.py

@@ -0,0 +1,86 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.utils import gast
+
+from .base import BaseTransformer
+
+__all__ = []
+
+
+class EarlyReturnTransformer(BaseTransformer):
+    """
+    Transform if/else return statement of Dygraph into Static Graph.
+    """
+
+    def __init__(self, root):
+        self.root = root
+
+    def transform(self):
+        """
+        Main function to transform AST.
+        """
+        self.visit(self.root)
+
+    def is_define_return_in_if(self, node):
+        assert isinstance(node, gast.If), (
+            f"Type of input node should be gast.If, but received {type(node)}."
+        )
+        for child in node.body:
+            if isinstance(child, gast.Return):
+                return True
+        return False
+
+    def visit_block_nodes(self, nodes):
+        result_nodes = []
+        destination_nodes = result_nodes
+        for node in nodes:
+            rewritten_node = self.visit(node)
+
+            if isinstance(rewritten_node, (list, tuple)):
+                destination_nodes.extend(rewritten_node)
+            else:
+                destination_nodes.append(rewritten_node)
+
+            # append other nodes to if.orelse even though if.orelse is not empty
+            if isinstance(node, gast.If) and self.is_define_return_in_if(node):
+                destination_nodes = node.orelse
+                # handle stmt like `if/elif/elif`
+                while (
+                    len(destination_nodes) > 0
+                    and isinstance(destination_nodes[0], gast.If)
+                    and self.is_define_return_in_if(destination_nodes[0])
+                ):
+                    destination_nodes = destination_nodes[0].orelse
+
+        return result_nodes
+
+    def visit_If(self, node):
+        node.body = self.visit_block_nodes(node.body)
+        node.orelse = self.visit_block_nodes(node.orelse)
+        return node
+
+    def visit_While(self, node):
+        node.body = self.visit_block_nodes(node.body)
+        node.orelse = self.visit_block_nodes(node.orelse)
+        return node
+
+    def visit_For(self, node):
+        node.body = self.visit_block_nodes(node.body)
+        node.orelse = self.visit_block_nodes(node.orelse)
+        return node
+
+    def visit_FunctionDef(self, node):
+        node.body = self.visit_block_nodes(node.body)
+        return node

python/paddle/jit/dy2static/transformers/ifelse_transformer.py

@@ -0,0 +1,447 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from collections import defaultdict
+
+from paddle.base import unique_name
+from paddle.jit.dy2static.utils import (
+    GetterSetterHelper,
+    ast_to_source_code,
+)
+
+# gast is a generic AST to represent Python2 and Python3's Abstract Syntax Tree(AST).
+# It provides a compatibility layer between the AST of various Python versions,
+# as produced by ast.parse from the standard ast module.
+# See details in https://github.com/serge-sans-paille/gast/
+from paddle.utils import gast
+
+from .base import BaseTransformer
+from .utils import (
+    FALSE_FUNC_PREFIX,
+    FOR_ITER_INDEX_PREFIX,
+    FOR_ITER_ITERATOR_PREFIX,
+    FOR_ITER_TARGET_PREFIX,
+    FOR_ITER_TUPLE_INDEX_PREFIX,
+    FOR_ITER_TUPLE_PREFIX,
+    FOR_ITER_VAR_LEN_PREFIX,
+    FOR_ITER_VAR_NAME_PREFIX,
+    FOR_ITER_ZIP_TO_LIST_PREFIX,
+    TRUE_FUNC_PREFIX,
+    FunctionNameLivenessAnalysis,
+    create_function_def_node,
+    create_get_args_node,
+    create_name_str,
+    create_nonlocal_stmt_nodes,
+    create_set_args_node,
+)
+
+__all__ = []
+
+GET_ARGS_FUNC_PREFIX = 'get_args'
+SET_ARGS_FUNC_PREFIX = 'set_args'
+ARGS_NAME = '__args'
+
+
+class IfElseTransformer(BaseTransformer):
+    """
+    Transform if/else statement of Dygraph into Static Graph.
+    """
+
+    def __init__(self, root):
+        self.root = root
+        FunctionNameLivenessAnalysis(
+            self.root
+        )  # name analysis of current ast tree.
+
+    def transform(self):
+        """
+        Main function to transform AST.
+        """
+        self.visit(self.root)
+
+    def visit_If(self, node):
+        self.generic_visit(node)
+        (
+            true_func_node,
+            false_func_node,
+            get_args_node,
+            set_args_node,
+            return_name_ids,
+            push_pop_ids,
+        ) = transform_if_else(node, self.root)
+
+        new_node = create_convert_ifelse_node(
+            return_name_ids,
+            push_pop_ids,
+            node.test,
+            true_func_node,
+            false_func_node,
+            get_args_node,
+            set_args_node,
+        )
+
+        return [
+            get_args_node,
+            set_args_node,
+            true_func_node,
+            false_func_node,
+            new_node,
+        ]
+
+    def visit_Call(self, node):
+        # Remove `numpy()` statement, like `Tensor.numpy()[i]` -> `Tensor[i]`
+        if isinstance(node.func, gast.Attribute):
+            attribute = node.func
+            if attribute.attr == 'numpy':
+                node = attribute.value
+        self.generic_visit(node)
+        return node
+
+    def visit_IfExp(self, node):
+        """
+        Transformation with `true_fn(x) if Tensor > 0 else false_fn(x)`
+        """
+        self.generic_visit(node)
+
+        new_node = create_convert_ifelse_node(
+            None, None, node.test, node.body, node.orelse, None, None, True
+        )
+        # Note: A blank line will be added separately if transform gast.Expr
+        # into source code. Using gast.Expr.value instead to avoid syntax error
+        # in python.
+        if isinstance(new_node, gast.Expr):
+            new_node = new_node.value
+
+        return new_node
+
+
+class NameVisitor(gast.NodeVisitor):
+    def __init__(self, after_node=None, end_node=None):
+        # The start node (exclusive) of the visitor
+        self.after_node = after_node
+        # The terminate node of the visitor.
+        self.end_node = end_node
+        # Dict to store the names and ctxs of vars.
+        self.name_ids = defaultdict(list)
+        # List of current visited nodes
+        self.ancestor_nodes = []
+        # True when in range (after_node, end_node).
+        self._in_range = after_node is None
+        self._candidate_ctxs = (gast.Store, gast.Load, gast.Param)
+        self._def_func_names = set()
+
+    def visit(self, node):
+        """Visit a node."""
+        if self.after_node is not None and node == self.after_node:
+            self._in_range = True
+            return
+        if node == self.end_node:
+            self._in_range = False
+            return
+
+        self.ancestor_nodes.append(node)
+        method = 'visit_' + node.__class__.__name__
+        visitor = getattr(self, method, self.generic_visit)
+        ret = visitor(node)
+        self.ancestor_nodes.pop()
+
+        return ret
+
+    def visit_If(self, node):
+        """
+        For nested `if/else`, the created vars are not always visible for parent node.
+        In addition, the vars created in `if.body` are not visible for `if.orelse`.
+
+        Case 1:
+            x = 1
+            if m > 1:
+                res = new_tensor
+            res = res + 1   # Error, `res` is not visible here.
+
+        Case 2:
+            if x_tensor > 0:
+                res = new_tensor
+            else:
+                res = res + 1   # Error, `res` is not visible here.
+
+        In above two cases, we should consider to manage the scope of vars to parsing
+        the arguments and returned vars correctly.
+        """
+        if not self._in_range or not self.end_node:
+            self.generic_visit(node)
+            return
+        else:
+            before_if_name_ids = copy.deepcopy(self.name_ids)
+            body_name_ids = self._visit_child(node.body)
+            # If traversal process stops early in `if.body`, return the currently seen name_ids.
+            if not self._in_range:
+                self._update_name_ids(before_if_name_ids)
+            else:
+                else_name_ids = self._visit_child(node.orelse)
+                # If traversal process stops early in `if.orelse`, return the currently seen name_ids.
+                if not self._in_range:
+                    self._update_name_ids(before_if_name_ids)
+                else:
+                    # Blocks the vars in `if.body` and only inserts the vars both created in 'if/else' branch
+                    # into name_ids.
+                    new_name_ids = self._find_new_name_ids(
+                        body_name_ids, else_name_ids
+                    )
+                    for new_name_id in new_name_ids:
+                        before_if_name_ids[new_name_id].append(gast.Store())
+
+                    self.name_ids = before_if_name_ids
+
+    def visit_Attribute(self, node):
+        if not self._in_range or not self._is_call_func_name_node(node):
+            self.generic_visit(node)
+
+    def visit_Name(self, node):
+        if not self._in_range:
+            self.generic_visit(node)
+            return
+        blacklist = {'True', 'False', 'None'}
+        if node.id in blacklist:
+            return
+        if node.id in self._def_func_names:
+            return
+        if not self._is_call_func_name_node(node):
+            if isinstance(node.ctx, self._candidate_ctxs):
+                self.name_ids[node.id].append(node.ctx)
+
+    def visit_Assign(self, node):
+        if not self._in_range:
+            self.generic_visit(node)
+            return
+        # Visit `value` firstly.
+        node._fields = ('value', 'targets')
+        self.generic_visit(node)
+
+    def visit_FunctionDef(self, node):
+        # NOTE: We skip to visit names of get_args and set_args, because they contains
+        # nonlocal statement such as 'nonlocal x, self' where 'self' should not be
+        # parsed as returned value in control flow.
+        if (
+            GET_ARGS_FUNC_PREFIX in node.name
+            or SET_ARGS_FUNC_PREFIX in node.name
+        ):
+            return
+
+        if not self._in_range:
+            self.generic_visit(node)
+            return
+        self._def_func_names.add(node.name)
+        if not self.end_node:
+            self.generic_visit(node)
+        else:
+            before_name_ids = copy.deepcopy(self.name_ids)
+            self.name_ids = defaultdict(list)
+            self.generic_visit(node)
+
+            if not self._in_range:
+                self._update_name_ids(before_name_ids)
+            else:
+                self.name_ids = before_name_ids
+
+    def _visit_child(self, node):
+        self.name_ids = defaultdict(list)
+        if isinstance(node, list):
+            for item in node:
+                if isinstance(item, gast.AST):
+                    self.visit(item)
+        elif isinstance(node, gast.AST):
+            self.visit(node)
+
+        return copy.deepcopy(self.name_ids)
+
+    def _find_new_name_ids(self, body_name_ids, else_name_ids):
+        def is_required_ctx(ctxs, required_ctx):
+            for ctx in ctxs:
+                if isinstance(ctx, required_ctx):
+                    return True
+            return False
+
+        candidate_name_ids = set(body_name_ids.keys()) & set(
+            else_name_ids.keys()
+        )
+        store_ctx = gast.Store
+        new_name_ids = set()
+        for name_id in candidate_name_ids:
+            if is_required_ctx(
+                body_name_ids[name_id], store_ctx
+            ) and is_required_ctx(else_name_ids[name_id], store_ctx):
+                new_name_ids.add(name_id)
+
+        return new_name_ids
+
+    def _is_call_func_name_node(self, node):
+        white_func_names = {'append', 'extend'}
+        if len(self.ancestor_nodes) > 1:
+            assert self.ancestor_nodes[-1] == node
+            parent_node = self.ancestor_nodes[-2]
+            if isinstance(parent_node, gast.Call) and parent_node.func == node:
+                # e.g: var_list.append(elem), var_list is also a name_id.
+                should_skip = (
+                    isinstance(node, gast.Attribute)
+                    and node.attr in white_func_names
+                )
+                if not should_skip:
+                    return True
+        return False
+
+    def _update_name_ids(self, new_name_ids):
+        for name_id, ctxs in new_name_ids.items():
+            self.name_ids[name_id] = ctxs + self.name_ids[name_id]
+
+
+def _valid_nonlocal_names(return_name_ids, nonlocal_names):
+    """
+    All var in return_name_ids should be in nonlocal_names.
+    Moreover, we will always put return_name_ids in front of nonlocal_names.
+
+    For Example:
+
+        return_name_ids: [x, y]
+        nonlocal_names : [a, y, b, x]
+
+    Return:
+        nonlocal_names : [x, y, a, b]
+    """
+    assert isinstance(return_name_ids, list)
+    for name in return_name_ids:
+        if name not in nonlocal_names:
+            raise ValueError(
+                f"Required returned var '{name}' must be in 'nonlocal' statement '', but not found."
+            )
+        nonlocal_names.remove(name)
+
+    return return_name_ids + nonlocal_names
+
+
+def transform_if_else(node, root):
+    """
+    Transform ast.If into control flow statement of Paddle static graph.
+    """
+
+    # TODO(liym27): Consider variable like `self.a` modified in if/else node.
+    return_name_ids = sorted(node.pd_scope.modified_vars())
+    push_pop_ids = sorted(node.pd_scope.variadic_length_vars())
+    nonlocal_names = list(return_name_ids)
+    nonlocal_names.sort()
+    # NOTE: All var in return_name_ids should be in nonlocal_names.
+    nonlocal_names = _valid_nonlocal_names(return_name_ids, nonlocal_names)
+
+    # TODO(dev): Need a better way to deal this.
+    # LoopTransformer will create some special vars, which is not visible by users. so we can sure it's safe to remove them.
+    filter_names = [
+        ARGS_NAME,
+        FOR_ITER_INDEX_PREFIX,
+        FOR_ITER_TUPLE_PREFIX,
+        FOR_ITER_TARGET_PREFIX,
+        FOR_ITER_ITERATOR_PREFIX,
+        FOR_ITER_TUPLE_INDEX_PREFIX,
+        FOR_ITER_VAR_LEN_PREFIX,
+        FOR_ITER_VAR_NAME_PREFIX,
+        FOR_ITER_ZIP_TO_LIST_PREFIX,
+    ]
+
+    def remove_if(x):
+        for name in filter_names:
+            if x.startswith(name):
+                return False
+        return True
+
+    nonlocal_names = list(filter(remove_if, nonlocal_names))
+    return_name_ids = nonlocal_names
+
+    nonlocal_stmt_node = create_nonlocal_stmt_nodes(nonlocal_names)
+
+    empty_arg_node = gast.arguments(
+        args=[],
+        posonlyargs=[],
+        vararg=None,
+        kwonlyargs=[],
+        kw_defaults=None,
+        kwarg=None,
+        defaults=[],
+    )
+
+    true_func_node = create_function_def_node(
+        nonlocal_stmt_node + node.body,
+        name=unique_name.generate(TRUE_FUNC_PREFIX),
+        input_args=empty_arg_node,
+        return_name_ids=[],
+    )
+    false_func_node = create_function_def_node(
+        nonlocal_stmt_node + node.orelse,
+        name=unique_name.generate(FALSE_FUNC_PREFIX),
+        input_args=empty_arg_node,
+        return_name_ids=[],
+    )
+
+    helper = GetterSetterHelper(None, None, nonlocal_names, push_pop_ids)
+    get_args_node = create_get_args_node(helper.union())
+    set_args_node = create_set_args_node(helper.union())
+
+    return (
+        true_func_node,
+        false_func_node,
+        get_args_node,
+        set_args_node,
+        return_name_ids,
+        push_pop_ids,
+    )
+
+
+def create_convert_ifelse_node(
+    return_name_ids,
+    push_pop_ids,
+    pred,
+    true_func,
+    false_func,
+    get_args_func,
+    set_args_func,
+    is_if_expr=False,
+):
+    """
+    Create `paddle.jit.dy2static.convert_ifelse(
+            pred, true_fn, false_fn, get_args, set_args, return_name_ids)`
+    to replace original `python if/else` statement.
+    """
+    if is_if_expr:
+        true_func_source = f"lambda : {ast_to_source_code(true_func)}"
+        false_func_source = f"lambda : {ast_to_source_code(false_func)}"
+    else:
+        true_func_source = true_func.name
+        false_func_source = false_func.name
+
+    convert_ifelse_layer = gast.parse(
+        '_jst.IfElse('
+        '{pred}, {true_fn}, {false_fn}, {get_args}, {set_args}, {return_name_ids}, push_pop_names={push_pop_ids})'.format(
+            pred=ast_to_source_code(pred),
+            true_fn=true_func_source,
+            false_fn=false_func_source,
+            get_args=(
+                get_args_func.name if not is_if_expr else 'lambda: None'
+            ),  # TODO: better way to deal with this
+            set_args=(
+                set_args_func.name if not is_if_expr else 'lambda args: None'
+            ),
+            return_name_ids=create_name_str(return_name_ids),
+            push_pop_ids=create_name_str(push_pop_ids),
+        )
+    ).body[0]
+
+    return convert_ifelse_layer

python/paddle/jit/dy2static/transformers/logical_transformer.py

@@ -0,0 +1,102 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.utils import gast
+
+from ..utils import ast_to_source_code
+from .base import BaseTransformer
+
+__all__ = []
+
+cmpop_type_to_str = {
+    gast.Eq: "==",
+    gast.NotEq: "!=",
+    gast.Lt: "<",
+    gast.LtE: "<=",
+    gast.Gt: ">",
+    gast.GtE: ">=",
+    gast.Is: "is",
+    gast.IsNot: "is not",
+    gast.In: "in",
+    gast.NotIn: "not in",
+}
+
+
+def cmpop_node_to_str(node):
+    return cmpop_type_to_str[type(node)]
+
+
+class LogicalTransformer(BaseTransformer):
+    """
+    Transform python boolean op into Paddle logical op.
+
+    For example:
+        a = x > 1 and y < 1
+
+    Transformed code:
+        a = _jst.And(lambda:x>1, lambda:y<1)
+    """
+
+    def __init__(self, root):
+        self.root = root
+
+    def transform(self):
+        return self.visit(self.root)
+
+    def visit_UnaryOp(self, node):
+        self.generic_visit(node)
+        if isinstance(node.op, gast.Not):
+            arg = ast_to_source_code(node.operand)
+            new_node_str = f"_jst.Not({arg})"
+            # NOTE: gast.parse returns Module(body=[expr(value=...)])
+            new_node = gast.parse(new_node_str).body[0].value
+            return new_node
+        return node
+
+    def visit_BoolOp(self, node):
+        self.generic_visit(node)
+        if isinstance(node.op, gast.And):
+            new_node = self._create_bool_op_node(node.values, 'And')
+        elif isinstance(node.op, gast.Or):
+            new_node = self._create_bool_op_node(node.values, 'Or')
+        else:
+            raise TypeError(
+                "Only supports and/or syntax in control flow if statement."
+            )
+        return new_node
+
+    def _create_bool_op_node(self, nodes, api_type):
+        '''
+        NOTE(liym27):
+           The arguments of function convert_logical_XX should be callable so that they can be run
+          according to the actual order. In `convert_logical_and(lambda:x>1, lambda:y<1)`, `lambda:y<1`
+          must be run after `lambda:x>1`, If `x>1` is False, `y<1` should NOT be run.
+        '''
+        assert len(nodes) > 1, (
+            f"The length of BoolOp should be at least 2, but received {len(nodes)}."
+        )
+        if len(nodes) > 2:
+            # Creates logic_and/logic_or node recursively.
+            pre_logic_node = self._create_bool_op_node(nodes[:2], api_type)
+            if len(nodes[2:]) == 1:
+                post_logic_node = nodes[2]
+            else:
+                post_logic_node = self._create_bool_op_node(nodes[2:], api_type)
+            nodes = [pre_logic_node, post_logic_node]
+
+        args = [ast_to_source_code(child) for child in nodes]
+        new_node_str = f"_jst.{api_type}(lambda:{args[0]}, lambda:{args[1]})"
+        # NOTE: gast.parse return Module(body=[expr(...)])
+        new_node = gast.parse(new_node_str).body[0].value
+        return new_node

python/paddle/jit/dy2static/transformers/loop_transformer.py

@@ -0,0 +1,713 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from collections import defaultdict
+
+from paddle.base import unique_name
+from paddle.utils import gast
+
+from ..utils import (
+    GetterSetterHelper,
+    ast_to_source_code,
+)
+from .base import (
+    BaseTransformer,
+    ForLoopTuplePreTransformer,
+    ForNodeVisitor,
+)
+from .utils import (
+    ARGS_NAME,
+    FOR_BODY_PREFIX,
+    FOR_CONDITION_PREFIX,
+    WHILE_BODY_PREFIX,
+    WHILE_CONDITION_PREFIX,
+    FunctionNameLivenessAnalysis,
+    create_get_args_node,
+    create_name_str,
+    create_nonlocal_stmt_nodes,
+    create_set_args_node,
+    get_attribute_full_name,
+    get_parent_mapping,
+)
+
+__all__ = []
+
+
+def create_while_nodes(
+    condition_name,
+    body_name,
+    loop_var_names,
+    push_pop_names,
+    getter_name,
+    setter_name,
+):
+    """
+    Returns a list of gast.Node which represents the calling of Paddle
+    controlflow while_loop.
+
+    Usually, the list just contain 1 statement such as:
+
+    [a, b, c] = paddle.jit.dy2static.convert_while_loop(
+            condition_name, body_name, [a, b, c])
+
+    where a, b, c are in loop_var_names.
+
+    However, if loop_var_names contains property such as foo.x, we cannot
+    assign the property as output of convert_while_loop because Python
+    property is a kind of read-only attribute. To handle the case, we replace
+    the attributes which are output of convert_while_loop with generated
+    variables, then if we know the attribute is not read-only at runtime, we
+    assign the attribute. The created statements are like:
+
+    [a, b, __attribute_variable_1] = paddle.jit.dy2static.convert_while_loop(
+            condition_name, body_name, [a, b, foo.x])
+    if not isinstance(getattr(type(foo), x, None), property): foo.x = __attribute_variable_1
+
+    The number of above statements is not only 1, that's why the return type is
+    a list of gast.Node.
+    """
+    # NOTE(liym27):
+    # It's better to parse the source code into an AST node than to customize an AST node
+    # including child nodes, because it is easy to mistake the ast node type when customizing the node.
+    #
+    # For example: loop_var_names = [a, b, foo.x], the type of `a` or `b` is gast.Name,
+    # but the type of `foo.x` gast.Attribute.
+    # We have to make loop_var_names and assign_loop_var_names with same order
+    # set doesn't have order so we convert it to list
+    loop_var_names = list(loop_var_names)
+    assign_loop_var_names = []
+    for name in loop_var_names:
+        assign_loop_var_names.append(name)
+
+    while_func_name = "_jst.While"
+    while_node_str = f"{while_func_name}({condition_name}, {body_name}, {getter_name}, {setter_name}, return_name_ids={create_name_str(loop_var_names)}, push_pop_names={create_name_str(push_pop_names)})"
+    while_node = gast.parse(while_node_str).body[0]
+
+    ret = [while_node]
+    return ret
+
+
+class NameVisitor(gast.NodeVisitor):
+    '''
+    Analysis name liveness for loop transformer
+    '''
+
+    def __init__(self, root_node):
+        # Set of gast.Name or gast.Attribute for variables
+        self.current_seen_vars = set()
+
+        # List of gast.While/gast.For nodes
+        self.current_loop = []
+
+        # List of nodes that have scope of variables.
+        self.nodes_with_scope = []
+        self.blacklist_names = {"False", "True", "None"}
+
+        # Mapping from gast.While/gast.For to variable nodes
+        self.before_loop_body_vars = defaultdict(set)
+        # NOTE: Use ordered list as dict value
+        self.in_loop_vars = defaultdict(list)
+
+        # Mapping from gast.While/gast.For to variable nodes which is condition
+        # of loop or being modified during the loop
+        self.write_in_loop = defaultdict(set)
+        self.condition_vars = defaultdict(set)
+        self.in_condition = False
+
+        # Some names are types, we shouldn't record them as loop var names.
+        self.type_vars = set()
+
+        self.to_parent_mapping = get_parent_mapping(root_node)
+
+        self.visit(root_node)
+
+    def get_loop_var_names(self, node):
+        assert isinstance(node, (gast.While, gast.For)), (
+            "Input node is not gast loop node"
+        )
+        loop_var_names = set()
+        create_var_names = set()
+        read_context = {type(gast.Load()), type(gast.AugLoad())}
+
+        in_loop_vars_list = self.in_loop_vars[node]
+
+        # get dict `var_name_to_ctxs`
+        var_name_to_ctxs = defaultdict(list)
+        for var_node in in_loop_vars_list:
+            var_name_to_ctxs[self._var_node_to_name(var_node)].append(
+                var_node.ctx
+            )
+
+        in_loop_vars = set(in_loop_vars_list)
+        in_loop_vars = self._remove_unnecessary_vars(in_loop_vars, node)
+        in_loop_name_strs = self._var_nodes_to_names(in_loop_vars)
+
+        before_loop_body_vars = self.before_loop_body_vars[node]
+        before_loop_body_vars = self._remove_unnecessary_vars(
+            before_loop_body_vars, node
+        )
+        before_loop_name_strs = self._var_nodes_to_names(before_loop_body_vars)
+
+        after_loop_vars = (
+            self.current_seen_vars - before_loop_body_vars - in_loop_vars
+        )
+        after_loop_vars = self._remove_unnecessary_vars(after_loop_vars, node)
+        after_loop_name_strs = self._var_nodes_to_names(
+            after_loop_vars, read_context
+        )
+        condition_vars = self.condition_vars[node]
+        condition_names = self._var_nodes_to_names(condition_vars)
+
+        write_vars = self.write_in_loop[node]
+        write_names = self._var_nodes_to_names(write_vars)
+
+        for name in in_loop_name_strs:
+            if name in before_loop_name_strs:
+                # If a variable is used in loop and created before loop
+
+                # If this var is a basic variable and read-only and not
+                # condition var, it may not be loop_var else it should
+                # be in loop_var as input
+                if (name not in condition_names) and (name not in write_names):
+                    continue
+                loop_var_names.add(name)
+
+            elif name in after_loop_name_strs:
+                # If a variable is created in the while loop and read after
+                # loop, it should be in loop_var and we should create it
+
+                # because name in after_loop_name must be initialized in loop
+                # So it is write-only, we don't have to filter read-only basic
+                # vars out
+                loop_var_names.add(name)
+                create_var_names.add(name)
+            else:
+                # If a variable is used and created in loop, but used before created,
+                # it should be in loop_var and we should create it.
+
+                # For example, `var_a` should be in loop_var and we should create it.
+                #
+                #   res = 0
+                #   for i, x in enumerate(x_array):
+                #       if i > 2:
+                #           x = func1(var_a)
+                #       var_a = func2(x)
+                #
+
+                is_created = False
+                for ctx in var_name_to_ctxs[name]:
+                    if isinstance(ctx, gast.Store):
+                        is_created = True
+
+                if (
+                    isinstance(var_name_to_ctxs[name][0], gast.Load)
+                    and is_created
+                ):
+                    loop_var_names.add(name)
+                    create_var_names.add(name)
+
+        return loop_var_names, create_var_names
+
+    def visit_Name(self, node):
+        if self._is_call_func_name_node(node):
+            self.generic_visit(node)
+            return
+        if node.id in self.blacklist_names:
+            self.generic_visit(node)
+            return
+
+        self.current_seen_vars.add(node)
+        write_context = {
+            type(gast.Store()),
+            type(gast.AugStore()),
+            type(gast.Del()),
+        }
+
+        for loop_node in self.current_loop:
+            self.in_loop_vars[loop_node].append(node)
+            if type(node.ctx) in write_context:
+                self.write_in_loop[loop_node].add(node)
+        if self.in_condition:
+            self.condition_vars[loop_node].add(node)
+        self.generic_visit(node)
+
+    def visit_FunctionDef(self, node):
+        self.nodes_with_scope.append(node)
+        self.blacklist_names.add(node.name)
+
+        # The variables in the function are not visible to the outside scope.
+        before_func_seen_vars = copy.copy(self.current_seen_vars)
+
+        self.generic_visit(node)
+        self.nodes_with_scope.pop()
+        # After exiting the scope of the node, variables in this scope
+        # should be removed from self.current_seen_vars.
+        if self.nodes_with_scope:
+            self.current_seen_vars = before_func_seen_vars
+
+    def visit(self, node):
+        method = 'visit_' + node.__class__.__name__
+        visitor = getattr(self, method, self.generic_visit)
+        ret = visitor(node)
+        return ret
+
+    def visit_Attribute(self, node):
+        if self._is_call_func_name_node(node):
+            return
+        attr_full_name = get_attribute_full_name(node)
+        # Class variables are not allowed to appear in the arguments list
+        # of defined function under class methods in Python.
+        """
+        def class_func(self):
+            def while_loop_body(self.x, y) # `self.x` is illegal.
+        """
+        # TODO: If do change the variable with `self.var`, need a better
+        # way to deal with this case.
+        if attr_full_name.startswith("self."):
+            return
+        self.current_seen_vars.add(node)
+
+        for loop_node in self.current_loop:
+            self.in_loop_vars[loop_node].append(node)
+
+        # sub-nodes are visited during get_attribute_full_name and we shouldn't
+        # visit again
+
+    def visit_For(self, node):
+        self.current_loop.append(node)
+        self.in_condition = True
+        self.visit(node.target)
+        self.visit(node.iter)
+        self.in_condition = False
+        self.before_loop_body_vars[node] = copy.copy(self.current_seen_vars)
+        self.generic_visit(node)
+        self.current_loop.pop()
+
+    def visit_While(self, node):
+        self.current_loop.append(node)
+        self.in_condition = True
+        self.visit(node.test)
+        self.in_condition = False
+        self.before_loop_body_vars[node] = copy.copy(self.current_seen_vars)
+        self.generic_visit(node)
+        self.current_loop.pop()
+
+    def visit_Call(self, node):
+        # Store type var names such as "isinstance(x, some_type_names)" and
+        # Remove them later
+        if isinstance(node.func, gast.Name) and node.func.id == 'isinstance':
+            type_node = node.args[1]
+            if isinstance(type_node, gast.Tuple):
+                for element in type_node.elts:
+                    self.type_vars.add(ast_to_source_code(element).strip())
+            else:
+                self.type_vars.add(ast_to_source_code(type_node).strip())
+        self.generic_visit(node)
+
+    def _var_nodes_to_names(self, node_set, ctx_filter_set=None):
+        ret = set()
+        for node in node_set:
+            if ctx_filter_set is None or type(node.ctx) in ctx_filter_set:
+                ret.add(self._var_node_to_name(node))
+        return ret
+
+    def _var_node_to_name(self, node):
+        if isinstance(node, gast.Name):
+            return node.id
+        elif isinstance(node, gast.Attribute):
+            return get_attribute_full_name(node)
+
+    def _is_call_func_name_node(self, node):
+        parent_node = self._get_parent_node(node)
+        if isinstance(parent_node, gast.Call) and parent_node.func == node:
+            return True
+        return False
+
+    def _is_global_or_nonlocal(self, node):
+        return False
+
+    def _is_ancestor_node(self, ancestor_node, node):
+        parent_node = self._get_parent_node(node)
+
+        while parent_node is not None:
+            if parent_node == ancestor_node:
+                return True
+            parent_node = self._get_parent_node(parent_node)
+        return False
+
+    def _get_parent_node(self, node):
+        return self.to_parent_mapping.get(node)
+
+    def _remove_unnecessary_vars(self, loop_vars, loop_node):
+        """
+        Remove unnecessary vars from before_loop_vars, after_loop_vars or in_loop_vars about loop_node.
+            1. Remove target vars of gast.For from before_loop_vars or after_loop_vars.
+            2. Remove vars only in gast.comprehension.
+            3. Remove vars that are type names, for example: "isinstance(x, var_type_name)"
+        :param loop_vars: before_loop_vars, after_loop_vars or in_loop_vars of loop_node.
+        :param loop_node: Current loop node.
+        """
+
+        def filter_name_nodes_from(root_node, target_var_names):
+            """
+            Filter children with gast.Name type from node.(inclusivly)
+            """
+            name_nodes = set()
+            if isinstance(root_node, gast.Name):
+                if node.id in target_var_names:
+                    name_nodes.add(root_node)
+            for child_node in gast.walk(root_node):
+                if isinstance(child_node, gast.Name):
+                    if child_node.id in target_var_names:
+                        name_nodes.add(child_node)
+
+            return name_nodes
+
+        vars_of_list_generator = set()
+        target_vars_of_for_node = set()
+
+        for name_node in loop_vars:
+            if not isinstance(name_node, gast.Name):
+                continue
+
+            parent_node = self._get_parent_node(name_node)
+
+            # NOTE: gast.For.target or gast.comprehension.target can be gast.Tuple.
+            #  For examples:
+            #   1) `for i, j in enumerate(x)` has two target vars: i and j
+            #   2) `[x for x,y in array]` has two target vars: x and y
+            if isinstance(parent_node, gast.Tuple):
+                parent_node = self._get_parent_node(parent_node)
+
+            # 1. Get vars only in gast.comprehension.
+            # For examples:
+            #  1) [x for x,y in array] -> x, x, y
+            #  2) [f(x) for x in array] -> x
+            #  3) [func(x, y) for x in array] -> x, x
+            if isinstance(parent_node, gast.comprehension):
+                # 1.1 target vars in list/set comprehensions
+                target_node = parent_node.target
+                if isinstance(target_node, gast.Tuple):
+                    target_vars = target_node.elts
+                else:
+                    target_vars = [target_node]
+
+                vars_of_list_generator = vars_of_list_generator | set(
+                    target_vars
+                )
+
+                # 1.2 vars from target vars used in elt_node
+                target_var_names = {var.id for var in target_vars}
+                comp_node = self._get_parent_node(parent_node)
+                elt_nodes = []
+                if isinstance(comp_node, gast.ListComp):
+                    elt_nodes.append(comp_node.elt)
+                elif isinstance(comp_node, gast.DictComp):
+                    elt_nodes.extend([comp_node.key, comp_node.value])
+
+                for node in elt_nodes:
+                    vars_of_list_generator |= filter_name_nodes_from(
+                        node, target_var_names
+                    )
+
+            # 2. Get target vars or vars from target vars used in for-loop but the for-loop is
+            #   1) not the "loop_node" itself
+            #   2) not the ancestor of the "loop_node"
+            #
+            # For examples:
+            #   for k in range(x):   # if it's this "loop_node", i or j both should be target vars.
+            #      # do something
+            #
+            #   for i in range(a):   # if it's this "loop_node", k or j should be in target vars but i should not.
+            #     for j in range(a): # if it's this "loop_node", k should be in target_vars but i or j should not.
+            #       x = i+j
+            elif isinstance(parent_node, gast.For):
+                if parent_node is loop_node:
+                    continue
+                if self._is_ancestor_node(parent_node, loop_node):
+                    continue
+                # 2.1 target vars in gast.For node.
+                target_node = parent_node.target
+                if isinstance(target_node, gast.Tuple):
+                    target_vars = target_node.elts
+                else:
+                    target_vars = [target_node]
+
+                target_vars_of_for_node = target_vars_of_for_node | set(
+                    target_vars
+                )
+
+        # 2.2 vars from target vars used in for-loop
+        target_vars_name_strs = {var.id for var in target_vars_of_for_node}
+        for var in loop_vars:
+            if not isinstance(var, gast.Name):
+                continue
+            if (
+                var.id in target_vars_name_strs
+                and var not in self.condition_vars[loop_node]
+            ):
+                target_vars_of_for_node.add(var)
+
+        removed_vars = target_vars_of_for_node | vars_of_list_generator
+
+        # 3. Remove var type names which are stored in self.type_vars
+        for var in loop_vars:
+            if ast_to_source_code(var).strip() in self.type_vars:
+                removed_vars.add(var)
+
+        return loop_vars - removed_vars
+
+
+class LoopTransformer(BaseTransformer):
+    """
+    This class transforms python while/for statement into Static Graph Ast
+    """
+
+    def __init__(self, root):
+        self.root = root
+        FunctionNameLivenessAnalysis(self.root)
+
+    def transform(self):
+        ForLoopTuplePreTransformer(self.root).transform()
+        self.visit(self.root)
+
+    def visit_While(self, node):
+        self.generic_visit(node)
+        new_stmts = self.get_while_stmt_nodes(node)
+        return new_stmts
+
+    def visit_For(self, node):
+        self.generic_visit(node)
+        new_stmts = self.get_for_stmt_nodes(node)
+        return new_stmts
+
+    def replace_stmt_list(self, body_list):
+        if not isinstance(body_list, list):
+            return
+
+        i = 0
+        while i < len(body_list):
+            if isinstance(body_list[i], gast.While):
+                new_stmts = self.get_while_stmt_nodes(body_list[i])
+                body_list[i : i + 1] = new_stmts
+                i += len(new_stmts)
+            elif isinstance(body_list[i], gast.For):
+                new_stmts = self.get_for_stmt_nodes(body_list[i])
+                body_list[i : i + 1] = new_stmts
+                i += len(new_stmts)
+            else:
+                i += 1
+
+    def get_for_stmt_nodes(self, node):
+        # TODO: consider for - else in python
+
+        # 1. get key statements for different cases
+        # NOTE 1: three key statements:
+        #   1). init_stmts: list[node], prepare nodes of for loop, may not only one
+        #   2). cond_stmt: node, condition node to judge whether continue loop
+        #   3). body_stmts: list[node], updated loop body, sometimes we should change
+        #       the original statement in body, not just append new statement
+        #
+        # NOTE 2: The following `for` statements will be transformed to `while` statements:
+        #   1). for x in range(*)
+        #   2). for x in iter_var
+        #   3). for i, x in enumerate(*)
+
+        current_for_node_parser = ForNodeVisitor(node)
+        stmts_tuple = current_for_node_parser.parse()
+        if stmts_tuple is None:
+            return [node]
+        init_stmts, cond_stmt, body_stmts = stmts_tuple
+        # 2. get original loop vars
+        loop_var_names, create_var_names = (
+            node.pd_scope.modified_vars(),
+            node.pd_scope.created_vars(),
+        )
+        push_pop_names = list(node.pd_scope.variadic_length_vars())
+        # TODO: Remove the bunch of code?  We have the unique format `for A in B:`
+        # NOTE: in 'for x in var' or 'for i, x in enumerate(var)' cases,
+        # we need append new loop var & remove useless loop var
+        #   1. for x in var -> x is no need
+        #   2. for i, x in enumerate(var) -> x is no need
+        if current_for_node_parser.is_for_iter():
+            iter_var_name = current_for_node_parser.iter_var_name
+            iter_idx_name = current_for_node_parser.iter_idx_name
+            loop_var_names.add(iter_idx_name)
+            if current_for_node_parser.enum_idx_name is not None:
+                loop_var_names.add(current_for_node_parser.enum_idx_name)
+
+        # 3. prepare result statement list
+        new_stmts = []
+        # Python can create variable in loop and use it out of loop, E.g.
+        #
+        # for x in range(10):
+        #     y += x
+        # print(x) # x = 10
+        #
+        # We don't need to create static variable for them, because
+        # we do this in CreateUndefinedVarTransformer
+
+        # create non-local statement for body and cond.
+        nonlocal_names = list(loop_var_names | create_var_names)
+        nonlocal_names.sort()
+        # TODO(dev): Need a better way to deal this.
+        if ARGS_NAME in nonlocal_names:
+            nonlocal_names.remove(ARGS_NAME)
+
+        nonlocal_stmt_node = create_nonlocal_stmt_nodes(nonlocal_names)
+
+        # 4. append init statements
+        new_stmts.extend(init_stmts)
+
+        # 5. create & append condition function node
+        condition_func_node = gast.FunctionDef(
+            name=unique_name.generate(FOR_CONDITION_PREFIX),
+            args=gast.arguments(
+                args=[],
+                posonlyargs=[],
+                vararg=None,
+                kwonlyargs=[],
+                kw_defaults=None,
+                kwarg=None,
+                defaults=[],
+            ),
+            body=[*nonlocal_stmt_node, gast.Return(value=cond_stmt)],
+            decorator_list=[],
+            returns=None,
+            type_comment=None,
+            type_params=[],
+        )
+        new_stmts.append(condition_func_node)
+
+        # 6. create & append loop body function node
+        # append return values for loop body
+        body_func_node = gast.FunctionDef(
+            name=unique_name.generate(FOR_BODY_PREFIX),
+            args=gast.arguments(
+                args=[],
+                posonlyargs=[],
+                vararg=None,
+                kwonlyargs=[],
+                kw_defaults=None,
+                kwarg=None,
+                defaults=[],
+            ),
+            body=nonlocal_stmt_node + body_stmts,
+            decorator_list=[],
+            returns=None,
+            type_comment=None,
+            type_params=[],
+        )
+        new_stmts.append(body_func_node)
+
+        helper = GetterSetterHelper(None, None, nonlocal_names, push_pop_names)
+        get_args_node = create_get_args_node(helper.union())
+        set_args_node = create_set_args_node(helper.union())
+        # 7. create & append while loop node
+        while_loop_nodes = create_while_nodes(
+            condition_func_node.name,
+            body_func_node.name,
+            nonlocal_names,
+            push_pop_names,
+            get_args_node.name,
+            set_args_node.name,
+        )
+        new_stmts.extend([get_args_node, set_args_node])
+        new_stmts.extend(while_loop_nodes)
+
+        return new_stmts
+
+    def get_while_stmt_nodes(self, node):
+        loop_var_names, create_var_names = (
+            node.pd_scope.modified_vars(),
+            node.pd_scope.created_vars(),
+        )
+        push_pop_names = list(node.pd_scope.variadic_length_vars())
+        new_stmts = []
+
+        # create non-local statement for body and cond.
+        nonlocal_names = list(loop_var_names | create_var_names)
+        nonlocal_names.sort()
+        # TODO(dev): Need a better way to deal this.
+        if ARGS_NAME in nonlocal_names:
+            nonlocal_names.remove(ARGS_NAME)
+
+        nonlocal_stmt_node = create_nonlocal_stmt_nodes(nonlocal_names)
+
+        # Python can create variable in loop and use it out of loop, E.g.
+        #
+        # while x < 10:
+        #     x += 1
+        #     y = x
+        # z = y
+        #
+        # We don't need to create static variable for those variables, because
+        # we do this in CreateUndefinedVarTransformer
+
+        condition_func_node = gast.FunctionDef(
+            name=unique_name.generate(WHILE_CONDITION_PREFIX),
+            args=gast.arguments(
+                args=[],
+                posonlyargs=[],
+                vararg=None,
+                kwonlyargs=[],
+                kw_defaults=None,
+                kwarg=None,
+                defaults=[],
+            ),
+            body=[*nonlocal_stmt_node, gast.Return(value=node.test)],
+            decorator_list=[],
+            returns=None,
+            type_comment=None,
+            type_params=[],
+        )
+
+        new_stmts.append(condition_func_node)
+
+        new_body = node.body
+        body_func_node = gast.FunctionDef(
+            name=unique_name.generate(WHILE_BODY_PREFIX),
+            args=gast.arguments(
+                args=[],
+                posonlyargs=[],
+                vararg=None,
+                kwonlyargs=[],
+                kw_defaults=None,
+                kwarg=None,
+                defaults=[],
+            ),
+            body=nonlocal_stmt_node + new_body,
+            decorator_list=[],
+            returns=None,
+            type_comment=None,
+            type_params=[],
+        )
+        new_stmts.append(body_func_node)
+
+        helper = GetterSetterHelper(None, None, nonlocal_names, push_pop_names)
+        get_args_node = create_get_args_node(helper.union())
+        set_args_node = create_set_args_node(helper.union())
+
+        while_loop_nodes = create_while_nodes(
+            condition_func_node.name,
+            body_func_node.name,
+            nonlocal_names,
+            push_pop_names,
+            get_args_node.name,
+            set_args_node.name,
+        )
+        new_stmts.extend([get_args_node, set_args_node])
+        new_stmts.extend(while_loop_nodes)
+        return new_stmts

python/paddle/jit/dy2static/transformers/name_load_transformer.py

@@ -0,0 +1,130 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from paddle.utils import gast
+
+from ..utils import ast_to_source_code
+from .base import BaseTransformer
+
+__all__ = []
+
+
+class NameloadJstTransformer(BaseTransformer):
+    """
+    change name and attribute load to __jst.Ld(name) pattern.
+    for example:
+        a.dtype -->  __jst.Ld(__jst.Ld(a).dtype)
+
+    In paddle science and deepxde, we have to support changing tensor into variable
+    in arbitrary occasion such as global tensor.
+
+    NOTE: we only deal with ctx=Load() case.
+    """
+
+    def __init__(self, root):
+        self.root = root
+
+    def transform(self):
+        self.visit(self.root)
+        return self.root
+
+    def _surround_with_ld(self, node):
+        node = (
+            gast.parse(f"_jst.Ld({ast_to_source_code(node).strip()})")
+            .body[0]
+            .value
+        )
+        return node
+
+    def visit_Call(self, node):
+        """
+        Can't convert name of function call, because this will affect CallTransformer.
+        """
+        node.args = [self.visit(arg) for arg in node.args]
+        for keyword in node.keywords:
+            keyword.value = self.visit(keyword.value)
+        node.func = self.visit(node.func)
+        return node
+
+    def create_visit_with_convert_load(self, node_type, skip_fn=None):
+        def visit(node):
+            assert isinstance(node, node_type)
+            if skip_fn and skip_fn(node):
+                return node
+            self.generic_visit(node)
+            if isinstance(node.ctx, gast.Load):
+                node = self._surround_with_ld(node)
+            return node
+
+        return visit
+
+    def visit_Attribute(self, node):
+        def skip_fn(node):
+            if isinstance(node.value, gast.Name) and node.value.id == "_jst":
+                return True
+            return False
+
+        return self.create_visit_with_convert_load(gast.Attribute, skip_fn)(
+            node
+        )
+
+    def visit_Subscript(self, node):
+        return self.create_visit_with_convert_load(gast.Subscript)(node)
+
+    def visit_Name(self, node):
+        return self.create_visit_with_convert_load(gast.Name)(node)
+
+
+class AttributeJstTransformer(BaseTransformer):
+    """
+    change some special attribute into __jst.XXX(obj, "attr_name") format.
+    for example:
+        a.size  -->  __jst.attr(a, "size")
+
+    because `size` have different behavior when in dygraph / static graph mode
+    NOTE: we only deal with ctx=Load() case.
+    """
+
+    def __init__(self, node):
+        assert isinstance(node, gast.AST), (
+            "Input non-gast.AST node for the initialization of ToTensorTransformer."
+        )
+        self.interested_name = {
+            'size',
+        }
+        self.root = node
+
+    def transform(self):
+        self.visit(self.root)
+        return self.root
+
+    def visit_Attribute(self, node):
+        assert isinstance(node, gast.Attribute)
+        assert isinstance(node.attr, str)
+        if (
+            isinstance(node.ctx, gast.Load)
+            and node.attr in self.interested_name
+        ):
+            attr = node.attr
+            value = node.value
+            node = (
+                gast.parse(
+                    f'_jst.Attr({ast_to_source_code(value).strip()}, "{attr}")'
+                )
+                .body[0]
+                .value
+            )
+        self.generic_visit(node)
+        return node

python/paddle/jit/dy2static/transformers/return_transformer.py

@@ -0,0 +1,415 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.base import unique_name
+from paddle.utils import gast
+
+from ..utils import (
+    ORIGIN_INFO,
+    Dygraph2StaticException,
+    ast_to_source_code,
+)
+from .base import BaseTransformer
+from .break_continue_transformer import ForToWhileTransformer
+from .utils import create_bool_node, index_in_list
+
+__all__ = []
+
+# Constant for the name of the variable which stores the boolean state that we
+# should return
+RETURN_PREFIX = '__return'
+
+# Constant for the name of the variable which stores the final return value
+RETURN_VALUE_PREFIX = '__return_value'
+
+# Constant for the name of variables to initialize the __return_value
+RETURN_VALUE_INIT_NAME = '__return_value_init'
+
+# Constant magic number representing returning no value. This constant amis to
+# support returning various lengths of variables. Static graph must have fixed
+# size of fetched output while dygraph can have flexible lengths of output, to
+# solve it in dy2stat, we put float64 value with this magic number at Static
+# graph as a place holder to indicate the returning placeholder means no value
+# should return.
+
+
+def get_return_size(return_node):
+    assert isinstance(return_node, gast.Return), "Input is not gast.Return node"
+    return_length = 0
+    if return_node.value is not None:
+        if isinstance(return_node.value, gast.Tuple):
+            return_length = len(return_node.value.elts)
+        else:
+            return_length = 1
+    return return_length
+
+
+class ReplaceReturnNoneTransformer(BaseTransformer):
+    """
+    Replace 'return None' to  'return' because 'None' cannot be a valid input
+    in control flow. In ReturnTransformer single 'Return' will be appended no
+    value placeholder
+    """
+
+    def __init__(self, root_node):
+        self.root = root_node
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_Return(self, node):
+        if isinstance(node.value, gast.Name) and node.value.id == 'None':
+            node.value = None
+            return node
+        if isinstance(node.value, gast.Constant) and node.value.value is None:
+            node.value = None
+            return node
+        return node
+
+
+class ReturnAnalysisVisitor(gast.NodeVisitor):
+    """
+    Visits gast Tree and analyze the information about 'return'.
+    """
+
+    def __init__(self, root_node):
+        self.root = root_node
+        assert isinstance(self.root, gast.FunctionDef), (
+            "Input is not gast.FunctionDef node"
+        )
+
+        # the number of return statements
+        self.count_return = 0
+
+        # maximum number of variables
+        self.max_return_length = 0
+
+        self.visit(self.root)
+
+    def visit_FunctionDef(self, node):
+        """
+        don't analysis closure, just analyze current func def level.
+        """
+        if node == self.root:
+            self.generic_visit(node)
+
+    def visit_Return(self, node):
+        self.count_return += 1
+
+        return_length = get_return_size(node)
+        self.max_return_length = max(self.max_return_length, return_length)
+
+        self.generic_visit(node)
+
+    def get_func_return_count(self):
+        return self.count_return
+
+    def get_func_max_return_length(self):
+        return self.max_return_length
+
+
+class ReturnTransformer(BaseTransformer):
+    """
+    Transforms return statements into equivalent python statements containing
+    only one return statement at last. The basics idea is using a return value
+    variable to store the early return statements and boolean states with
+    if-else to skip the statements after the return.
+
+    Go through all the function definition and call SingleReturnTransformer for each function.
+    SingleReturnTransformer don't care the nested function def.
+    """
+
+    def __init__(self, root):
+        self.root = root
+        pre_transformer = ReplaceReturnNoneTransformer(self.root)
+        pre_transformer.transform()
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_FunctionDef(self, node):
+        node = self.generic_visit(node)
+        node = SingleReturnTransformer(node).transform()
+        return node
+
+
+class SingleReturnTransformer(BaseTransformer):
+    """
+    This function only apply to single function. don't care the nested function_def
+    """
+
+    def __init__(self, root):
+        self.root = root
+        assert isinstance(self.root, gast.FunctionDef), (
+            "Input is not gast.FunctionDef node"
+        )
+
+        self.ancestor_nodes = []
+
+        # The name of return placeholder
+        self.return_value_name = None
+
+        # Every return stmt corresponds to a bool value variable, and return name is the name of the boolean variable
+        self.return_name = []
+
+        self.pre_analysis = None
+
+    def assert_parent_is_not_while(self, parent_node_of_return):
+        if isinstance(parent_node_of_return, (gast.While, gast.For)):
+            raise Dygraph2StaticException(
+                "Found return statement in While or For body and loop "
+                "is meaningless, please check you code and remove return in while/for."
+            )
+
+    def generic_visit(self, node):
+        # Because we change ancestor nodes during visit_Return, not current
+        # node, original generic_visit of NodeTransformer will visit node
+        # which may be deleted. To prevent that node being added into
+        # transformed AST, We self-write a generic_visit and visit
+        for field, value in gast.iter_fields(node):
+            if isinstance(value, list):
+                for item in value:
+                    if isinstance(item, gast.AST):
+                        self.visit(item)
+            elif isinstance(value, gast.AST):
+                self.visit(value)
+
+    def visit(self, node):
+        """
+        Self-defined visit for appending ancestor
+        """
+        self.ancestor_nodes.append(node)
+        ret = super().visit(node)
+        self.ancestor_nodes.pop()
+        return ret
+
+    def visit_FunctionDef(self, node):
+        """
+        don't analysis closure, just analyze current func def level.
+        """
+        if node == self.root:
+            self.generic_visit(node)
+        return node
+
+    def append_assign_to_return_node(
+        self, value, parent_node_of_return, return_name, assign_nodes
+    ):
+        self.assert_parent_is_not_while(parent_node_of_return)
+        assert value in [True, False], "value must be True or False."
+        if isinstance(parent_node_of_return, gast.If):
+            # Prepend control flow boolean nodes such as '__return@1 = True'
+            node_str = f"{return_name} = _jst.create_bool_as_type({ast_to_source_code(parent_node_of_return.test).strip()}, {value})"
+
+            assign_node = gast.parse(node_str).body[0]
+            assign_nodes.append(assign_node)
+
+    def transform(self):
+        node = self.root
+        self.pre_analysis = ReturnAnalysisVisitor(node)
+        max_return_length = self.pre_analysis.get_func_max_return_length()
+        while self.pre_analysis.get_func_return_count() > 0:
+            # every visit will decrease the number of returns.
+            # so we need a while.
+            self.visit(node)
+            self.pre_analysis = ReturnAnalysisVisitor(node)
+
+        if max_return_length == 0:
+            return node
+
+        # Prepend initialization of final return and append final return statement
+        return_flag_names = self.return_name
+        value_name = self.return_value_name
+        if value_name is not None:
+            node.body.append(
+                gast.Return(
+                    value=gast.Name(
+                        id=value_name,
+                        ctx=gast.Load(),
+                        annotation=None,
+                        type_comment=None,
+                    )
+                )
+            )
+            assign_return_value_node = gast.Assign(
+                targets=[
+                    gast.Name(
+                        id=value_name,
+                        ctx=gast.Store(),
+                        annotation=None,
+                        type_comment=None,
+                    )
+                ],
+                value=gast.Constant(kind=None, value=None),
+                type_comment=None,
+            )
+            node.body.insert(0, assign_return_value_node)
+
+        for return_flag_name in return_flag_names:
+            assign_return_flag_node = create_bool_node(return_flag_name, False)
+            node.body.insert(0, assign_return_flag_node)
+
+        # Prepend no value placeholders
+        return node
+
+    def visit_Return(self, node):
+        return_name = unique_name.generate(RETURN_PREFIX)
+        self.return_name.append(return_name)
+        max_return_length = self.pre_analysis.get_func_max_return_length()
+        parent_node_of_return = self.ancestor_nodes[-2]
+
+        for ancestor_index in reversed(range(len(self.ancestor_nodes) - 1)):
+            ancestor = self.ancestor_nodes[ancestor_index]
+            cur_node = self.ancestor_nodes[ancestor_index + 1]
+
+            def _deal_branches(branch_name):
+                if hasattr(ancestor, branch_name):
+                    branch_node = getattr(ancestor, branch_name)
+                    if index_in_list(branch_node, cur_node) != -1:
+                        if cur_node == node:
+                            self._replace_return_in_stmt_list(
+                                branch_node,
+                                cur_node,
+                                return_name,
+                                max_return_length,
+                                parent_node_of_return,
+                            )
+                        self._replace_after_node_to_if_in_stmt_list(
+                            branch_node,
+                            cur_node,
+                            return_name,
+                            parent_node_of_return,
+                        )
+
+            _deal_branches("body")
+            _deal_branches("orelse")
+            # If return node in while loop, add `not return_name` in gast.While.test
+            if isinstance(ancestor, gast.While):
+                cond_var_node = gast.UnaryOp(
+                    op=gast.Not(),
+                    operand=gast.Name(
+                        id=return_name,
+                        ctx=gast.Load(),
+                        annotation=None,
+                        type_comment=None,
+                    ),
+                )
+                ancestor.test = gast.BoolOp(
+                    op=gast.And(), values=[ancestor.test, cond_var_node]
+                )
+                continue
+
+            # If return node in for loop, add `not return_name` in gast.While.test
+            if isinstance(ancestor, gast.For):
+                cond_var_node = gast.UnaryOp(
+                    op=gast.Not(),
+                    operand=gast.Name(
+                        id=return_name,
+                        ctx=gast.Load(),
+                        annotation=None,
+                        type_comment=None,
+                    ),
+                )
+                parent_node = self.ancestor_nodes[ancestor_index - 1]
+                for_to_while = ForToWhileTransformer(
+                    parent_node, ancestor, cond_var_node
+                )
+                new_stmts = for_to_while.transform()
+                while_node = new_stmts[-1]
+                self.ancestor_nodes[ancestor_index] = while_node
+
+            if ancestor == self.root:
+                break
+        # return_node is replaced so we shouldn't return here
+
+    def _replace_return_in_stmt_list(
+        self,
+        stmt_list,
+        return_node,
+        return_name,
+        max_return_length,
+        parent_node_of_return,
+    ):
+        assert max_return_length >= 0, "Input illegal max_return_length"
+        i = index_in_list(stmt_list, return_node)
+        if i == -1:
+            return False
+
+        assign_nodes = []
+        self.append_assign_to_return_node(
+            True, parent_node_of_return, return_name, assign_nodes
+        )
+
+        return_length = get_return_size(return_node)
+        # In this case we should NOT append RETURN_NO_VALUE placeholder
+        if return_node.value is not None:
+            if self.return_value_name is None:
+                self.return_value_name = unique_name.generate(
+                    RETURN_VALUE_PREFIX
+                )
+
+            assign_nodes.append(
+                gast.Assign(
+                    targets=[
+                        gast.Name(
+                            id=self.return_value_name,
+                            ctx=gast.Store(),
+                            annotation=None,
+                            type_comment=None,
+                        )
+                    ],
+                    value=return_node.value,
+                    type_comment=None,
+                )
+            )
+            return_origin_info = getattr(return_node, ORIGIN_INFO, None)
+            setattr(assign_nodes[-1], ORIGIN_INFO, return_origin_info)
+
+        # If there is a return in the body or else of if, the remaining statements
+        # will not be executed, so they can be properly replaced.
+        stmt_list[i:] = assign_nodes
+        return True
+
+    def _replace_after_node_to_if_in_stmt_list(
+        self, stmt_list, node, return_name, parent_node_of_return
+    ):
+        i = index_in_list(stmt_list, node)
+        if i < 0 or i >= len(stmt_list):
+            return False
+        if i == len(stmt_list) - 1:
+            # No need to add, we consider this as added successfully
+            return True
+
+        if_stmt = gast.If(
+            test=gast.UnaryOp(
+                op=gast.Not(),
+                operand=gast.Name(
+                    id=return_name,
+                    ctx=gast.Store(),
+                    annotation=None,
+                    type_comment=None,
+                ),
+            ),
+            body=stmt_list[i + 1 :],
+            orelse=[],
+        )
+
+        stmt_list[i + 1 :] = [if_stmt]
+
+        # Here assume that the parent node of return is gast.If
+        assign_nodes = []
+        self.append_assign_to_return_node(
+            False, parent_node_of_return, return_name, assign_nodes
+        )
+        stmt_list[i:i] = assign_nodes
+        return True

python/paddle/jit/dy2static/transformers/super_transformer.py

@@ -0,0 +1,60 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.utils import gast
+
+from .base import BaseTransformer
+
+__all__ = []
+
+
+class SuperTransformer(BaseTransformer):
+    """
+    This class transforms super() into super(__class__, <first argument>).
+    """
+
+    def __init__(self, root):
+        self.root = root
+        self.first_arg = None
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_FunctionDef(self, node):
+        if self.first_arg is not None:
+            return self.generic_visit(node)
+
+        positional_args = node.args.posonlyargs + node.args.args
+        if not positional_args:
+            return self.generic_visit(node)
+
+        self.first_arg = positional_args[0].id
+
+        return self.generic_visit(node)
+
+    def visit_Call(self, node):
+        # super() -> _jst.WrapSuper(super)(x.__class__, x)
+        self.generic_visit(node)
+        if self.first_arg is None:
+            return node
+        if not isinstance(node.func, gast.Name):
+            return node
+        if node.func.id != "super":
+            return node
+        if node.args:
+            return node
+
+        new_fn_call_str = f"_jst.WrapSuper(super)({self.first_arg}.__class__, {self.first_arg})"
+        new_fn_call_ast = gast.parse(new_fn_call_str).body[0]
+        return new_fn_call_ast

python/paddle/jit/dy2static/transformers/tensor_shape_transformer.py

@@ -0,0 +1,45 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.utils import gast
+
+from ..utils import ast_to_source_code
+from .base import BaseTransformer
+
+__all__ = []
+
+
+class TensorShapeTransformer(BaseTransformer):
+    """
+    This class transforms variable.shape  into Static Graph Ast.
+    All 'xxx.shape' will be converted int '_jst.Shape(x)'.
+    """
+
+    def __init__(self, root):
+        self.root = root
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_Attribute(self, node):
+        self.generic_visit(node)
+        if node.attr == 'shape':
+            args = ast_to_source_code(node.value).strip()
+            # NOTE(dev): we can deal with paddle.shape in this case, but it's
+            # not pretty to modify into 'convert_shape(paddle)(x)[0]'.
+            if args != 'paddle':
+                convert_shape_func = f"_jst.Shape({args})"
+                shape_node = gast.parse(convert_shape_func).body[0].value
+                return shape_node
+        return node

python/paddle/jit/dy2static/transformers/tensorhook_transformer.py

@@ -0,0 +1,111 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from paddle.utils import gast
+
+from ..utils import ast_to_source_code
+from .base import BaseTransformer
+
+
+def get_loads(node: gast.AST):
+    for child in gast.walk(node):
+        if isinstance(
+            child, (gast.Name, gast.Attribute, gast.Subscript)
+        ) and isinstance(child.ctx, gast.Load):
+            yield child
+
+
+class RegisterHookTransformer(BaseTransformer):
+    def __init__(self, root):
+        self.root = root
+
+    def transform(self):
+        """
+        Main function to transform AST.
+        """
+        self.visit(self.root)
+
+    def reorder_block_statements(self, stmts):
+        register_hook_nodes = [
+            n
+            for n in stmts
+            for stmt in gast.walk(n)
+            if isinstance(stmt, gast.Attribute) and stmt.attr == "register_hook"
+        ]
+        # Analyze the register_hook nodes name dependency
+        dependents = {}
+        for n in register_hook_nodes:
+            if n not in stmts:
+                continue
+            for load_node in get_loads(n):
+                load_name = ast_to_source_code(load_node)
+                if load_name not in dependents:
+                    dependents[load_name] = []
+                dependents[load_name].append(n)
+
+        # Reorder the register_hook nodes, insert it before the dependent nodes
+        idx = 0
+        reordered_stmts = list(stmts)
+        while idx < len(reordered_stmts):
+            stmt = reordered_stmts[idx]
+            loads = get_loads(stmt)
+            for load_node in loads:
+                load_name = ast_to_source_code(load_node)
+                if load_name in dependents:
+                    dep_nodes = dependents[load_name]
+                    for dep_node in dep_nodes:
+                        dep_idx = reordered_stmts.index(dep_node)
+                        if dep_idx <= idx:
+                            continue
+                        reordered_stmts.remove(dep_node)
+                        reordered_stmts.insert(idx, dep_node)
+                        idx += 1
+            idx += 1
+        return reordered_stmts
+
+    def visit_FunctionDef(self, node: gast.FunctionDef):
+        node.body = self.reorder_block_statements(node.body)
+        self.generic_visit(node)
+        return node
+
+    def visit_For(self, node: gast.For):
+        node.body = self.reorder_block_statements(node.body)
+        node.orelse = self.reorder_block_statements(node.orelse)
+        self.generic_visit(node)
+        return node
+
+    def visit_While(self, node: gast.While):
+        node.body = self.reorder_block_statements(node.body)
+        node.orelse = self.reorder_block_statements(node.orelse)
+        self.generic_visit(node)
+        return node
+
+    def visit_If(self, node: gast.If):
+        node.body = self.reorder_block_statements(node.body)
+        node.orelse = self.reorder_block_statements(node.orelse)
+        self.generic_visit(node)
+        return node
+
+    def visit_With(self, node: gast.With):
+        node.body = self.reorder_block_statements(node.body)
+        self.generic_visit(node)
+        return node
+
+    def visit_Try(self, node: gast.Try):
+        node.body = self.reorder_block_statements(node.body)
+        node.orelse = self.reorder_block_statements(node.orelse)
+        node.finalbody = self.reorder_block_statements(node.finalbody)
+        self.generic_visit(node)
+        return node

python/paddle/jit/dy2static/transformers/transform.py

@@ -0,0 +1,144 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# gast is a generic AST to represent Python2 and Python3's Abstract Syntax Tree(AST).
+# It provides a compatibility layer between the AST of various Python versions,
+# as produced by ast.parse from the standard ast module.
+# See details in https://github.com/serge-sans-paille/gast/
+
+import os
+
+from paddle.framework import use_pir_api
+
+from .. import logging_utils
+from ..utils import ast_to_source_code
+from .assert_transformer import AssertTransformer
+from .base import BaseTransformer
+from .break_continue_transformer import (
+    BreakContinueTransformer,
+    BreakTransformOptimizer,
+)
+from .call_transformer import CallTransformer
+from .cast_transformer import CastTransformer
+from .create_variable_transformer import CreateVariableTransformer
+from .decorator_transformer import DecoratorTransformer
+from .early_return_transformer import EarlyReturnTransformer
+from .ifelse_transformer import IfElseTransformer
+from .logical_transformer import LogicalTransformer
+from .loop_transformer import LoopTransformer
+from .name_load_transformer import (
+    AttributeJstTransformer,
+    NameloadJstTransformer,
+)
+from .return_transformer import ReturnTransformer
+from .super_transformer import SuperTransformer
+from .tensor_shape_transformer import TensorShapeTransformer
+from .tensorhook_transformer import RegisterHookTransformer
+from .typehint_transformer import TypeHintTransformer
+
+__all__ = []
+
+
+def apply_optimization(transformers):
+    """
+    Judge whether to apply optimized transformation, such as BreakTransformOptimizer.
+    And not all optimized transformations are applied by default. It's controlled by
+    'export FLAGS_optim_transformation=1'
+    """
+    flag = str(os.environ.get('FLAGS_optim_transformation')) in [
+        '1',
+        'True',
+        'true',
+    ]
+    if flag:
+        transformers.insert(3, BreakTransformOptimizer)
+
+
+class DygraphToStaticAst(BaseTransformer):
+    """
+    Main class to transform Dygraph to Static Graph
+    """
+
+    def __init__(self):
+        self.translator_logger = logging_utils.TranslatorLogger()
+
+    def get_static_ast(self, root):
+        self.root = root
+        self.decorate_func_name = None
+
+        # inplace transfer
+        self.transfer_from_node_type(self.root)
+        return self.root
+
+    def _apply(self, transformer, node, log_level):
+        transformer(node).transform()
+        self.translator_logger.log_transformed_code(
+            log_level, self.root, transformer.__name__
+        )
+
+    def transfer_from_node_type(self, node):
+        self.translator_logger.log(
+            1, f"Source code: \n{ast_to_source_code(self.root)}"
+        )
+        # Generic transformation
+        self.visit(node)
+
+        transformers = [
+            TypeHintTransformer,  # remove all typehint
+            SuperTransformer,  # super() -> super(__class__, <first argument>)
+            RegisterHookTransformer,
+            EarlyReturnTransformer,
+            AttributeJstTransformer,  # Tensor.size -> Tensor.size(), it's unnecessary in PIR mode
+            TensorShapeTransformer,  # Tensor.shape -> paddle.shape(Tensor)
+            BreakContinueTransformer,  # break/continue in loops
+            ReturnTransformer,  # return in functions
+            LogicalTransformer,  # logical and/or/not
+            CreateVariableTransformer,  # create undefined var for if / while / for
+            LoopTransformer,  # for/while -> while_op
+            IfElseTransformer,  # if/else -> if_op
+            AssertTransformer,  # assert statement
+            CallTransformer,  # transform call recursively
+            CastTransformer,  # type casting statement
+            DecoratorTransformer,  # transform decorators to function call
+            NameloadJstTransformer,
+        ]
+
+        if use_pir_api():
+            # It's unnecessary in PIR mode
+            transformers.remove(AttributeJstTransformer)
+
+        apply_optimization(transformers)
+
+        for index, transformer in enumerate(transformers):
+            self._apply(transformer, node, log_level=index + 1)
+
+        self.translator_logger.log_transformed_code(
+            logging_utils.LOG_AllTransformer, self.root, "All Transformers"
+        )
+
+    def visit_FunctionDef(self, node):
+        if self.decorate_func_name is None:
+            self.decorate_func_name = node.name
+
+        self.generic_visit(node)
+        return node
+
+    def get_module_name(self):
+        """
+        Return the main function name which will be used as module name
+        in ast_to_func.
+        """
+        # Should consider BaseAPITransformer which add new module name in Yamei's PR.
+        assert self.decorate_func_name, "decorate_func_name shall not be None."
+        return self.decorate_func_name

python/paddle/jit/dy2static/transformers/typehint_transformer.py

@@ -0,0 +1,53 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from paddle.utils import gast
+
+from .base import BaseTransformer
+
+__all__ = []
+
+
+class TypeHintTransformer(BaseTransformer):
+    """
+    A class remove all the typehint in gast.Name(annotation).
+    Please put it behind other transformers because other transformer may relay on typehints.
+    """
+
+    def __init__(self, root):
+        self.root = root
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_FunctionDef(self, node):
+        node.returns = None
+        self.generic_visit(node)
+        return node
+
+    def visit_Name(self, node):
+        node.annotation = None
+        self.generic_visit(node)
+        return node
+
+    def visit_AnnAssign(self, node):
+        if node.value is None:
+            return None
+        assign_node = gast.Assign(
+            targets=[node.target],
+            value=node.value,
+            type_comment=None,
+        )
+        return assign_node

python/paddle/jit/dy2static/transformers/utils.py

@@ -0,0 +1,622 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import copy
+import textwrap
+import warnings
+
+import numpy as np
+
+from paddle.base import unique_name
+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
+
+GET_ARGS_FUNC_PREFIX = 'get_args'
+SET_ARGS_FUNC_PREFIX = 'set_args'
+ARGS_NAME = '__args'
+
+TRUE_FUNC_PREFIX = 'true_fn'
+FALSE_FUNC_PREFIX = 'false_fn'
+
+FOR_ITER_INDEX_PREFIX = '__for_loop_var_index'
+FOR_ITER_TUPLE_PREFIX = '__for_loop_iter_tuple'
+FOR_ITER_TARGET_PREFIX = '__for_loop_iter_target'
+FOR_ITER_ITERATOR_PREFIX = '__for_loop_iter_iterator'
+FOR_ITER_TUPLE_INDEX_PREFIX = '__for_loop_iter_tuple_index'
+FOR_ITER_VAR_LEN_PREFIX = '__for_loop_var_len'
+FOR_ITER_VAR_NAME_PREFIX = '__for_loop_iter_var'
+FOR_ITER_ZIP_TO_LIST_PREFIX = '__for_loop_iter_zip'
+
+WHILE_CONDITION_PREFIX = 'while_condition'
+WHILE_BODY_PREFIX = 'while_body'
+FOR_CONDITION_PREFIX = 'for_loop_condition'
+FOR_BODY_PREFIX = 'for_loop_body'
+
+
+def index_in_list(array_list, item):
+    try:
+        return array_list.index(item)
+    except ValueError:
+        # Item not in array_list
+        return -1
+
+
+class BaseNodeVisitor(gast.NodeVisitor):
+    """
+    Implement customized NodeVisitor inherited from gast.NodeVisitor.
+    Ancestor nodes are traced to easily support more operations of currently
+    visited node.
+    """
+
+    def __init__(self):
+        self.ancestor_nodes = []
+
+    def visit(self, node):
+        """Visit a node."""
+        self.ancestor_nodes.append(node)
+
+        method = 'visit_' + node.__class__.__name__
+        visitor = getattr(self, method, self.generic_visit)
+        ret = visitor(node)
+        self.ancestor_nodes.pop()
+        return ret
+
+
+def create_undefined_var(name):
+    func_code = f"{name} = _jst.UndefinedVar('{name}')"
+    return gast.parse(func_code).body[0]
+
+
+def create_bool_node(name, value):
+    '''
+    Create a assign stmt for name = value .
+    '''
+    assert isinstance(value, bool)
+    node = f"{name} = {value}"
+    return gast.parse(node).body[0]
+
+
+def get_parent_mapping(root):
+    to_parent: dict[gast.AST, gast.AST] = {}
+    for node in gast.walk(root):
+        for child in gast.iter_child_nodes(node):
+            to_parent[child] = node
+    return to_parent
+
+
+def create_name_str(name_ids):
+    """
+    Return "('x', 'y')" for [x, y]
+    """
+    if not name_ids:
+        return 'None'
+
+    names_str = ["'{}'".format(name.replace("'", "\\'")) for name in name_ids]
+    return "({}, )".format(','.join(names_str))
+
+
+def create_function_def_node(nodes, name, input_args, return_name_ids):
+    """
+    Wrapper all statements of nodes into one ast.FunctionDef, which can be
+    called by ast.Call.
+    """
+    nodes = copy.copy(nodes)
+    # add return statement
+    if return_name_ids:
+        nodes.append(gast.Return(value=generate_name_node(return_name_ids)))
+    else:
+        nodes.append(gast.Return(value=None))
+    func_def_node = gast.FunctionDef(
+        name=name,
+        args=input_args,
+        body=nodes,
+        decorator_list=[],
+        returns=None,
+        type_comment=None,
+        type_params=[],
+    )
+    return func_def_node
+
+
+def create_assign_node(name, node):
+    """
+    Creates a `gast.Assign` node by given name_id as target and node as value.
+    """
+    targets = generate_name_node(name, ctx=gast.Store())
+    assign_node = gast.Assign(
+        targets=[targets],
+        value=node,
+        type_comment=None,
+    )
+    return targets, assign_node
+
+
+def create_get_args_node(names):
+    """
+    Create get_args function as follows:
+
+        def get_args_0():
+            nonlocal x, y
+            return x, y
+    """
+
+    def empty_node():
+        func_def = f"""
+        def {unique_name.generate(GET_ARGS_FUNC_PREFIX)}():
+            return
+        """
+        return gast.parse(textwrap.dedent(func_def)).body[0]
+
+    assert isinstance(names, (list, tuple))
+    node = create_nonlocal_stmt_nodes(names)
+    if not names:
+        return empty_node()
+    if node == []:
+        nonlocal_vars = "\n"
+    else:
+        nonlocal_vars = ast_to_source_code(node[0])
+    template = """
+    def {func_name}():
+        {nonlocal_vars}
+        return {vars},
+    """
+    func_def = template.format(
+        func_name=unique_name.generate(GET_ARGS_FUNC_PREFIX),
+        nonlocal_vars=nonlocal_vars,
+        vars=",".join(names),
+    )
+    return gast.parse(textwrap.dedent(func_def)).body[0]
+
+
+def create_set_args_node(names):
+    """
+    Create set_args function as follows:
+
+        def set_args_0(__args):
+            nonlocal x, y
+            x, y = __args
+    """
+
+    def empty_node():
+        func_def = f"""
+        def {unique_name.generate(SET_ARGS_FUNC_PREFIX)}({ARGS_NAME}):
+            pass
+        """
+        return gast.parse(textwrap.dedent(func_def)).body[0]
+
+    assert isinstance(names, (list, tuple))
+    node = create_nonlocal_stmt_nodes(names)
+    if not names:
+        return empty_node()
+    if node == []:
+        nonlocal_vars = "\n"
+    else:
+        nonlocal_vars = ast_to_source_code(node[0])
+    template = """
+    def {func_name}({args}):
+        {nonlocal_vars}
+        {vars}, = {args}
+    """
+    func_def = template.format(
+        func_name=unique_name.generate(SET_ARGS_FUNC_PREFIX),
+        args=ARGS_NAME,
+        nonlocal_vars=nonlocal_vars,
+        vars=",".join(names),
+    )
+    return gast.parse(textwrap.dedent(func_def)).body[0]
+
+
+def create_nonlocal_stmt_nodes(names):
+    assert isinstance(names, (list, tuple))
+
+    mapped = list(filter(lambda n: '.' not in n, names))
+    mapped = list(filter(lambda n: '[' not in n, mapped))
+    names = sorted(
+        mapped, key=mapped.index
+    )  # to keep the order, we can't use set() to unique
+    if not names:
+        return []
+    func_code = "nonlocal {}".format(','.join(names))
+    return [gast.parse(func_code).body[0]]
+
+
+def generate_name_node(name_ids, ctx=gast.Load(), gen_tuple_if_single=False):
+    """
+    If name_ids is list or tuple or set with multiple strings, this function
+    generates gast.Tuple of gast.Name.
+    If the name_ids is single string or contains only 1 string, this function
+    returns gast.Name if gen_tuple_if_single==False else returns gast.Tuple
+    with only one gast.Name
+
+    This function is used at several gast.Return statements.
+    """
+    if isinstance(name_ids, str):
+        name_ids = [name_ids]
+    if not isinstance(name_ids, (list, tuple, set)):
+        raise TypeError(
+            f'name_ids must be list or tuple or set, but received {type(name_ids)}'
+        )
+
+    def create_node_for_name(name):
+        if '.' not in name:
+            return gast.Name(
+                id=name, ctx=ctx, annotation=None, type_comment=None
+            )
+        return gast.parse(name).body[0].value
+
+    gast_names = [create_node_for_name(name_id) for name_id in name_ids]
+    if len(gast_names) == 1 and not gen_tuple_if_single:
+        name_node = gast_names[0]
+    else:
+        name_node = gast.Tuple(elts=gast_names, ctx=ctx)
+    return name_node
+
+
+def get_attribute_full_name(node):
+    assert isinstance(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)
+
+
+class NameScope:
+    def __init__(self):
+        """
+        A NameScope is a object which manager all the variable names.
+        only FunctionDef and Controlflow node will have a namescope property.
+
+        type can be "function" and "controlflow"
+
+        we don't analyze the read only variable because they don't affect the analysis.
+        """
+        self.globals = set()
+        self.nonlocals = set()
+        self.args = set()
+        self.father = None  # point to the nearest function name scope.
+        self.w_vars = set()  # all qualified + normal names been stored
+        self.created = set()  # useful for control flow compatibility
+        # only valid in control_flow nodes
+        # may be remove later.
+        self.push_pop_vars = set()  # we call push and pop in the vars
+
+    def set_father(self, father):
+        self.father = father
+
+    def existed_vars(self):
+        """vars existing in current scope.
+        they must not contain qualified names.
+        """
+        local_vars = self.w_vars - self.globals - self.nonlocals - self.args
+        return set(filter(lambda x: '.' not in x, local_vars))
+
+    def created_vars(self):
+        return self.created
+
+    def modified_vars(self):
+        # may be globals / non-locals / args / qualified names and created_vars
+        return self.w_vars
+
+    def variadic_length_vars(self):
+        """
+        At present, we do not support global append, such as
+
+        import numpy as np
+        a = []
+        def func():
+            a.append() # global names `a`, we will raise a warning.
+            p.append(a, 1) # global names `np`, we will raise a warning.
+        """
+        non_global_push_pop_names = []
+        for var in self.push_pop_vars:
+            if self._is_simple_name(var) and self.is_global_var(var):
+                warnings.warn(
+                    f"Find variable `{var}` defined in global scope"
+                    f" and call `{var}.append() or {var}.pop()`"
+                    f", which will be ignored and never be transferred into"
+                    f" tensor array."
+                )
+            else:
+                non_global_push_pop_names.append(var)
+        return set(non_global_push_pop_names)
+
+    def control_flow_vars(self):
+        valid_names = self.w_vars
+        tmp = (self.father.global_vars & valid_names,)
+        return {"global": tmp, "nonlocal": self.w_vars - tmp}
+
+    def _is_simple_name(self, name):
+        if '.' in name or '[' in name:
+            return False
+        return True
+
+    def is_global_var(self, name):
+        """
+        Return whether the name is a var created in global scope.
+        Search from bottom to top. If it is not created or modified,
+        it means global vars; otherwise, it means local vars.
+        Only valid after FunctionNameLivenessAnalysis visitor.
+        """
+        assert self._is_simple_name(name), (
+            "is_global_var accept a simple name, but get `{name}`."
+        )
+        ancestor = self
+        while ancestor is not None:
+            if name in ancestor.globals:
+                return True
+            if name in (ancestor.nonlocals | ancestor.w_vars):
+                return False
+            ancestor = ancestor.father
+        return True
+
+    def is_local_var(self, name):
+        return not self.is_global_var(name)
+
+    def merge_from(self, name_scope):
+        self.globals |= name_scope.globals
+        self.nonlocals |= name_scope.nonlocals
+        self.args |= name_scope.args
+        self.w_vars |= name_scope.w_vars
+        self.push_pop_vars |= name_scope.push_pop_vars
+
+
+class FunctionNameLivenessAnalysis(gast.NodeVisitor):
+    """analyze the liveness of a function.
+
+    every variables stored in this scope will be collected,
+    in addition with global/nonlocal information and
+    push_pop information.
+
+    1. global variable is stored in node.var_globals.
+    2. nonlocal variable is stored in node.var_nonlocals.
+    3. arguments is stored in node.var_args.
+    4. if a variable's push and pop attribute is called,
+       it will be collected in push_pop_vars. They are
+       used for transformation to tensor_array.
+       NOTE: push_pop_vars **may not** in w_vars.
+       a.push(0) don't modify the variable a, but the content
+       of a.
+
+    For example:
+
+    def func(*args, **kargs):
+        a = 12
+        global i,j
+        nonlocal x,y
+        print(a)
+        i = k
+        b = []
+        c = [1,2,3]
+        for m in range(10):
+            q = 12
+            b.push(1)
+            c.pop()
+
+    After this visitor we have:
+    # node is the FunctionDef node with name: "func"
+    node.pd_scope = NameScope(
+        globals = ['i', 'j'],
+        nonlocals = ['x', 'y'],
+        args = ['args', 'kargs'],
+        wr_vars = ['a', 'i', 'q', 'm', 'c', 'b']
+        push_pop_vars = ['b', 'c']
+    )
+    """
+
+    def __init__(self, root_node):
+        self.scope_node_stack = []  # controlflow, functiondef node
+        self.visit(root_node)
+
+    def _reset_name_scope(self, node):
+        # always reset the node as empty namescope.
+        node.pd_scope = NameScope()
+
+    def _get_name_scope(self, node):
+        if not hasattr(node, "pd_scope"):
+            node.pd_scope = NameScope()
+        return node.pd_scope
+
+    def _current_name_scope(self):
+        return self._get_name_scope(self.scope_node_stack[-1])
+
+    def _father_name_scope(self):
+        if len(self.scope_node_stack) == 1:
+            return None
+        return self._get_name_scope(self.scope_node_stack[-2])
+
+    def _nearest_function_scope(self):
+        if len(self.scope_node_stack) == 1:
+            return None
+        for node in self.scope_node_stack[-2::-1]:
+            if isinstance(node, gast.FunctionDef):
+                return self._get_name_scope(node)
+
+    def visit_ListComp(self, node):
+        """[ i for i in range(10) ]
+        In this case, `i` will not created in FunctionScope.
+        We don't collect `i` by not calling generic_visit.
+        """
+        pass
+
+    def visit_DictComp(self, node):
+        """the same as ListComp."""
+        pass
+
+    def visit_Name(self, node):
+        self.generic_visit(node)
+        write_context = (gast.Store, gast.AugStore, gast.Del)
+        if isinstance(node.ctx, write_context):
+            self._current_name_scope().w_vars.add(node.id)
+
+    def visit_FunctionDef(self, node):
+        def pre_func():
+            self._current_name_scope().args |= set(
+                self._get_argument_names(node)
+            )
+
+        def post_func():
+            """NOTE: why we need merge w_vars and push_pop_vars here ?
+            because we do ifelse_transformer after loop_transformer. Loops will changed into functions. but we know this function will be called in if. so we add w_vars to father function scope.
+            """
+            control_flow_function_def = [
+                WHILE_BODY_PREFIX,
+                WHILE_BODY_PREFIX,
+                FOR_CONDITION_PREFIX,
+                FOR_BODY_PREFIX,
+                TRUE_FUNC_PREFIX,
+                FALSE_FUNC_PREFIX,
+            ]
+
+            def is_control_flow_def_node():
+                for prefix in control_flow_function_def:
+                    if node.name.startswith(prefix):
+                        return True
+                return False
+
+            if self._father_name_scope() and is_control_flow_def_node():
+                self._father_name_scope().w_vars |= (
+                    self._current_name_scope().w_vars
+                )
+                self._father_name_scope().push_pop_vars |= (
+                    self._current_name_scope().push_pop_vars
+                )
+
+        self._visit_scope_node(node, pre_func, post_func)
+
+    def _visit_scope_node(self, node, pre_func, post_func):
+        """scope node main visit logic.
+        pre_func and post_func is callbacks
+        """
+        self._reset_name_scope(node)
+        self.scope_node_stack.append(node)
+        self._current_name_scope().set_father(self._nearest_function_scope())
+        if pre_func:
+            pre_func()
+        self.generic_visit(node)
+        if post_func:
+            post_func()
+        self.scope_node_stack.pop()
+
+    def _visit_controlflow_node(self, node):
+        def post_func():
+            self._father_name_scope().merge_from(self._current_name_scope())
+            self._nearest_function_scope().merge_from(
+                self._current_name_scope()
+            )
+            self._current_name_scope().created = (
+                self._nearest_function_scope().existed_vars()
+                - node.before_created
+            )
+            # gather created vars into father and used in CreateUndefinedVarTransform
+            self._nearest_function_scope().created |= (
+                self._current_name_scope().created
+            )
+
+        def pre_func():
+            node.before_created = self._nearest_function_scope().existed_vars()
+
+        self._visit_scope_node(node, pre_func, post_func)
+
+    def visit_For(self, node):
+        self._visit_controlflow_node(node)
+
+    def visit_While(self, node):
+        self._visit_controlflow_node(node)
+
+    def visit_If(self, node):
+        self._visit_controlflow_node(node)
+
+    def visit_Global(self, node):
+        self._current_name_scope().globals |= set(node.names)
+
+    def visit_Nonlocal(self, node):
+        self._current_name_scope().nonlocals |= set(node.names)
+
+    def visit_Attribute(self, node):
+        self.generic_visit(node)
+        write_context = (gast.Store, gast.AugStore, gast.Del)
+        if isinstance(node.ctx, write_context):
+            name = ast_to_source_code(node).strip()
+            self._current_name_scope().w_vars.add(name)
+
+    def visit_Subscript(self, node):
+        self.generic_visit(node)
+        write_context = (gast.Store, gast.AugStore, gast.Del)
+        if isinstance(node.ctx, write_context):
+            while isinstance(node.value, gast.Subscript):
+                node = node.value
+            if isinstance(node.value, gast.Name):
+                self._current_name_scope().w_vars.add(node.value.id)
+
+    def visit_Call(self, node):
+        self.generic_visit(node)
+        if not isinstance(node.func, gast.Attribute):
+            return
+        variadic_length_method = ['append', 'pop']
+        if node.func.attr not in variadic_length_method:
+            return
+        # we don't treat push and pop as a write operator. such as a[i]=10 is not modify a.
+        name = ast_to_source_code(node.func.value).strip()
+        self._current_name_scope().push_pop_vars.add(name)
+
+    def _get_argument_names(self, node):
+        """get all arguments name in the functiondef node.
+        this node is local to the function and shouldn't
+        be created.
+        """
+        assert isinstance(node, gast.FunctionDef), (
+            "Input node is not function define node"
+        )
+        names = list(node.args.args)
+        names.append(node.args.vararg)
+        names.append(node.args.kwarg)
+        names = [i.id for i in names if i is not None]
+        return names