mindspore 2.0.0a0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/rewrite/symbol_tree.py

@@ -21,15 +21,14 @@ import ast
 import importlib
 import types
 import time
-
 import astunparse
 
 from mindspore.nn import Cell
 from mindspore import log as logger
 from mindspore.rewrite.ast_creator_register import ast_creator_registry
-from .node import Node, TreeNode, PASS_THROUGH_METHOD
+from .node import Node, TreeNode
 from .api.node_type import NodeType
-from .ast_helpers import AstModifier, AstReplacer, StrChecker, AstFinder
+from .ast_helpers import AstModifier, AstReplacer, StrChecker, AstFinder, CheckPropertyIsUsed
 from .api.scoped_value import ScopedValue, ValueType
 from .symbol_tree_dumper import SymbolTreeDumper
 from .topological_manager import TopoManager
@@ -160,7 +159,6 @@ class SymbolTree(Observer, Observable):
         self._topo_mgr = TopoManager()
         self._topo_mgr.reg_observer(self)
 
-        self._global_vars: {str, object} = {origin_network_key: origin_network}
         self._nodes: {str, Node} = {}
         # parameters of forward method
         self._inputs: [Node] = []
@@ -171,6 +169,10 @@ class SymbolTree(Observer, Observable):
         self._class_ast: Optional[ast.ClassDef] = None
         self._root_ast: Optional[ast.FunctionDef] = None
         self._init_func_ast: Optional[ast.FunctionDef] = None
+        self._deleted_field = {}
+        self._deleted_node = []
+        self._external_func_ast = []
+        self._father_class_ast = []
 
         # head node is always point to the first node(in source code order) of SymbolTree
         self._head = None
@@ -263,6 +265,8 @@ class SymbolTree(Observer, Observable):
         for node in stree.nodes():
             if not isinstance(node, TreeNode):
                 continue
+            if node.symbol_tree._class_ast is None:
+                continue
             sub_stree: SymbolTree = node.symbol_tree
             SymbolTree._find_all_class_in_symboltree(sub_stree, seen_class, allow_class_name, replacers)
             # all modified ast.ClassDef should export to code
@@ -281,31 +285,7 @@ class SymbolTree(Observer, Observable):
         """Add Event.TopologicalChangeEvent event when build is finished."""
         self.add_event(Event.TopologicalChangeEvent)
 
-    def create_assign_node(self, targets, func_name, args, kwargs):
-        """
-        Create a ast.Assign type node.
-
-        Args:
-            targets (list): _description_
-            func_name (_type_): _description_
-            args (_type_): _description_
-            kwargs (_type_): _description_
-
-        Returns:
-            _type_: _description_
-        """
-        # create targets
-        ast_targets = [ast_creator_registry.get("Name")(targets)]
-        # create call
-        ast_func = ast_creator_registry.get("Attribute")(func_name)
-        ast_args = ast_creator_registry.get("Args")(args)
-        ast_kwargs = ast_creator_registry.get("KwArgs")(kwargs) if kwargs else []
-        ast_value = ast_creator_registry.get("Call")(func=ast_func, args=ast_args, keywords=ast_kwargs)
-        # create assign
-        ast_node = ast_creator_registry.get("Assign")(targets=ast_targets, value=ast_value)
-        return ast_node
-
-    def create_call_function(self, func, targets, args, kwargs):
+    def _create_call_function(self, func, targets, args, kwargs):
         """
         Create a Node object and generate the execution code to insert into the source code.
         The source code calls the 'func' function with 'args' and' kwargs' as parameters.
@@ -345,6 +325,30 @@ class SymbolTree(Observer, Observable):
                                                call_kwargs)
         return node
 
+    def create_assign_node(self, targets, func_name, args, kwargs):
+        """
+        Create a ast.Assign type node.
+
+        Args:
+            targets (list): _description_
+            func_name (_type_): _description_
+            args (_type_): _description_
+            kwargs (_type_): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        # create targets
+        ast_targets = [ast_creator_registry.get("Name")(targets)]
+        # create call
+        ast_func = ast_creator_registry.get("Attribute")(func_name)
+        ast_args = ast_creator_registry.get("Args")(args)
+        ast_kwargs = ast_creator_registry.get("KwArgs")(kwargs) if kwargs else []
+        ast_value = ast_creator_registry.get("Call")(func=ast_func, args=ast_args, keywords=ast_kwargs)
+        # create assign
+        ast_node = ast_creator_registry.get("Assign")(targets=ast_targets, value=ast_value)
+        return ast_node
+
     def inner_create_call_function(self, node_name, ast_node, func_name, func, targets, args, kwargs):
         '''
         Instantiate an instance of node whose type is `CallFunction`.
@@ -458,12 +462,6 @@ class SymbolTree(Observer, Observable):
         self._init_func_ast = ast_node
 
     def get_inputs(self):
-        """
-        Getter of `_inputs` which represents parameters of current forward method.
-
-        Returns:
-            A list of instance of Node whose node_type is NodeType.Input as input nodes.
-        """
         return self._inputs
 
     def get_head_node(self):
@@ -484,17 +482,6 @@ class SymbolTree(Observer, Observable):
         """
         return self._origin_network
 
-    def get_global_vars(self):
-        """Get global variables."""
-        return self._global_vars
-
-    def add_global_vars(self, key: str, value):
-        """Add global variables."""
-        if self._global_vars.get(key) is not None:
-            logger.info(f"The key '{key}' is duplicated")
-            return
-        self._global_vars[key] = value
-
     def get_nodes_dict(self):
         """Get dict of nodes"""
         return self._nodes
@@ -614,7 +601,6 @@ class SymbolTree(Observer, Observable):
             RuntimeError: If 'node_or_name' is not belong to this SymbolTree or any sub-SymbolTree of current
                 SymbolTree.
         """
-
         node = self._get_real_node(node_or_name)
         if node is None:
             raise RuntimeError("Node is not belong to current SymbolTree: ", node_or_name)
@@ -653,7 +639,12 @@ class SymbolTree(Observer, Observable):
             RuntimeError: If 'position' is not in current SymbolTree.
             RuntimeError: If corresponding ast node is not an ast.Assign when 'insert_to_ast' is True.
         """
-
+        if position is not None and hasattr(position.node, "container"):
+            cellcontainer = getattr(position.node, "container")
+            index = cellcontainer.node_list.index(position.node)
+            index = index if position.before_node else index + 1
+            cellcontainer.insert(index, node)
+            return node
         # if position in current SymbolTree
         if position is not None and position.symbol_tree is not self:
             raise RuntimeError("Position is not in current SymbolTree:", position)
@@ -678,37 +669,7 @@ class SymbolTree(Observer, Observable):
             self._node_visitor.append_node(node)
         # update init-function-ast and construct-function-ast
         if insert_to_ast:
-            node.set_func(ScopedValue.create_naming_value(node_name, "self"))
-            node_ast = node.get_ast()
-            if not isinstance(node_ast, ast.Assign):
-                raise RuntimeError("Only support insert cell op now")
-            if isinstance(node, TreeNode):
-                global_vars_key = node.get_name() + "_args"
-                self.add_global_vars(global_vars_key, node.symbol_tree.get_global_vars())
-                args_call = AstModifier.create_call(ScopedValue.create_naming_value("get", "global_vars"),
-                                                    [ScopedValue.create_variable_value(global_vars_key)])
-                value = ast.Call(func=ast.Name(node.symbol_tree.get_opt_cls_name(), ast.Store(), lineno=0,
-                                               col_offset=0), args=[args_call], keywords=[], lineno=0, col_offset=0)
-
-                ast_target = ast.Name("self." + node.get_name(), ast.Store(), lineno=0, col_offset=0)
-                assign = ast.Assign(targets=[ast_target], value=value, lineno=0, col_offset=0)
-                AstModifier.insert_assign_ast_to_function(self._init_func_ast, assign)
-
-                AstModifier.insert_assign_ast_to_function(self._root_ast, node_ast,
-                                                          None if position is None else position.node.get_ast(),
-                                                          position.before_node)
-                sub_stree: SymbolTree = node.symbol_tree
-                from .symbol_tree_builder import SymbolTreeBuilder
-                SymbolTreeBuilder.merge_module_of_subtree(self, sub_stree)
-            else:
-                AstModifier.insert_assign_to_function(self._init_func_ast,
-                                                      targets=[ScopedValue(ValueType.NamingValue, "self", node_name)],
-                                                      expr=ScopedValue(ValueType.NamingValue, "global_vars", "get"),
-                                                      args=[ScopedValue(ValueType.StringValue, "", node_name)])
-                AstModifier.insert_assign_ast_to_function(self._root_ast, node_ast,
-                                                          None if position is None else position.node.get_ast(),
-                                                          position.before_node)
-            self._global_vars[node_name] = node.get_instance()
+            self._insert_to_ast_while_insert_node(node, position)
         return node
 
     def append_node(self, node: Node, append_to_ast: bool = True) -> Node:
@@ -807,8 +768,9 @@ class SymbolTree(Observer, Observable):
         Returns:
             An instance of python node which has been appended to SymbolTree.
         """
-        logger.warning("Ignoring unsupported node(%s) in %s.", type(ast_node).__name__, type(ast_scope).__name__)
+        logger.info("Ignoring unsupported node (%s) (%s).", type(ast_node).__name__, type(ast_scope).__name__)
         node_name = self._node_name_namer.get_name(type(ast_node).__name__)
+        self._update_names_for_unique(ast_node)
         node = Node.create_python_node(ast_node, node_name)
         self._insert_node(Position.create(self, self._tail, False), node)
         return node
@@ -851,6 +813,10 @@ class SymbolTree(Observer, Observable):
         node = self._get_real_node(node_or_name)
         if node is None:
             raise RuntimeError("Node is not belong to current SymbolTree: ", node_or_name)
+        if hasattr(node, "container"):
+            cellcontainer = getattr(node, "container")
+            cellcontainer.erase(node)
+            return node
         ret = AstModifier.erase_ast_from_function(self._root_ast, node.get_ast())
         if not ret:
             raise RuntimeError("node not in function ast tree.")
@@ -860,6 +826,7 @@ class SymbolTree(Observer, Observable):
                 value.isolate()
                 break
         self._topo_mgr.on_erase_node(node)
+        self._deleted_node.append(node.get_name())
         return node
 
     def replace(self, old_node: Node, new_nodes: [Node]) -> Node:
@@ -884,6 +851,9 @@ class SymbolTree(Observer, Observable):
             RuntimeError: If 'old_node' is not belong to current SymbolTree.
         """
 
+        if hasattr(old_node, "container"):
+            self._replace_container_node(old_node, new_nodes)
+            return new_nodes[0]
         real_old_node = self._get_real_node(old_node)
         if real_old_node is None:
             raise RuntimeError("Old node is not belong to current SymbolTree:", old_node)
@@ -981,6 +951,13 @@ class SymbolTree(Observer, Observable):
         dump_st = SymbolTreeDumper(self)
         dump_st.dump()
 
+    def update_module_ast(self):
+        for node in self._external_func_ast:
+            self._module_ast.body.append(node)
+        for node in self._father_class_ast:
+            index = self._module_ast.body.index(self._class_ast)
+            self._module_ast.body.insert(index, node)
+
     def get_code(self) -> str:
         """
         Get source code of modified network.
@@ -992,6 +969,7 @@ class SymbolTree(Observer, Observable):
         if self._init_func_ast:
             self._remove_unused_field()
         self._remove_duplicated_import()
+        self.update_module_ast()
         ast.fix_missing_locations(self._module_ast)
         # Find all ast.ClassDef which can be export to code
         # Replace duplicated ast.ClassDef reference in main-ClassDef
@@ -1026,21 +1004,20 @@ class SymbolTree(Observer, Observable):
             A network object.
         """
         cls = self._get_cls_through_file()
-        return cls(self._global_vars)
+        new_net = cls(self._origin_network)
+        self._merge_origin_property(new_net)
+        return new_net
 
     def set_saved_file_name(self, file_name: str):
-        """Sets the filename used to save the network."""
         if file_name.endswith(".py"):
             self._saved_file_name = file_name
         else:
             self._saved_file_name = file_name + ".py"
 
     def get_saved_file_name(self):
-        """Gets the filename used to save the network."""
         return self._saved_file_name
 
     def save_network_to_file(self):
-        """Save the modified network to a file."""
         abs_path = os.path.abspath(self._saved_file_name)
         if os.path.isfile(abs_path):
             os.remove(abs_path)
@@ -1049,6 +1026,58 @@ class SymbolTree(Observer, Observable):
             f.write(source.encode('utf-8'))
             f.flush()
 
+    def update_scope_for_unique(self, node: Union[ast.Attribute, ast.Call, ast.Subscript]):
+        """ Update scope of ast node because of unique-ing of targets of other nodes. """
+        if isinstance(node, ast.Call):
+            self.update_scope_for_unique(node.func)
+            return
+        if not isinstance(node, (ast.Attribute, ast.Subscript)):
+            logger.warning(f"Cannot update node {astunparse.unparse(node)} for unique, type of node should "
+                           f"be one of (ast.Attribute, ast.Subscript).")
+            return
+        scope = node.value
+        if not isinstance(scope, ast.Name):
+            self.update_scope_for_unique(scope)
+            return
+        scope_name = scope.id
+        scope_name_unique = self._target_namer.get_real_arg(scope_name)
+        scope.id = scope_name_unique
+
+    def _insert_to_ast_while_insert_node(self, node: Node, position: Optional[Position]):
+        """ insert_to_ast_while_insert_node. """
+        node.set_func(ScopedValue.create_naming_value(node.get_name(), "self"))
+        node_ast = node.get_ast()
+        if not isinstance(node_ast, ast.Assign):
+            raise RuntimeError("Only support insert cell op now")
+        if isinstance(node, TreeNode):
+            setattr(self._origin_network, node.get_name(), node.get_instance())
+            args_call = AstModifier.create_call(ScopedValue(ValueType.NamingValue, "", "getattr"),
+                                                [ScopedValue(ValueType.NamingValue, "", "obj"),
+                                                 ScopedValue(ValueType.StringValue, "", node.get_name())])
+            value = ast.Call(func=ast.Name(node.symbol_tree.get_opt_cls_name(), ast.Store(), lineno=0,
+                                           col_offset=0), args=[args_call], keywords=[], lineno=0, col_offset=0)
+
+            ast_target = ast.Name("self." + node.get_name(), ast.Store(), lineno=0, col_offset=0)
+            assign = ast.Assign(targets=[ast_target], value=value, lineno=0, col_offset=0)
+            AstModifier.insert_assign_ast_to_function(self._init_func_ast, assign)
+
+            AstModifier.insert_assign_ast_to_function(self._root_ast, node_ast,
+                                                      None if position is None else position.node.get_ast(),
+                                                      position.before_node)
+            sub_stree: SymbolTree = node.symbol_tree
+            from .symbol_tree_builder import SymbolTreeBuilder
+            SymbolTreeBuilder.merge_module_of_subtree(self, sub_stree)
+        else:
+            AstModifier.insert_assign_to_function(self._init_func_ast,
+                                                  targets=[ScopedValue(ValueType.NamingValue, "self", node.get_name())],
+                                                  expr=ScopedValue(ValueType.NamingValue, "", "getattr"),
+                                                  args=[ScopedValue(ValueType.NamingValue, "", "obj"),
+                                                        ScopedValue(ValueType.StringValue, "", node.get_name())])
+            AstModifier.insert_assign_ast_to_function(self._root_ast, node_ast,
+                                                      None if position is None else position.node.get_ast(),
+                                                      position.before_node)
+        setattr(self._origin_network, node.get_name(), node.get_instance())
+
     def _remove_unused_import(self):
         """remove unused import in self._module_ast"""
         str_checker = StrChecker(self._module_ast)
@@ -1070,49 +1099,43 @@ class SymbolTree(Observer, Observable):
                     else:
                         body.names.remove(alias)
 
+    def _replace_container_node(self, old_node, new_nodes):
+        cellcontainer = getattr(old_node, "container")
+        index = cellcontainer.node_list.index(old_node)
+        for n in reversed(new_nodes):
+            cellcontainer.insert(index, n)
+        index = cellcontainer.node_list.index(old_node)
+        cellcontainer.erase(old_node)
+
     def _filter_out_to_delete_field(self, to_delete_field):
         """filter out used field from `to_delete_field`"""
-        # filter _handler field
-        if to_delete_field.get("_handler"):
-            to_delete_field.pop("_handler")
-        # filter field used in node of construct
-        for node in self._nodes.values():
-            if node.get_node_type() in (NodeType.CallCell, NodeType.CallPrimitive, NodeType.Tree):
-                func: ScopedValue = node.get_func()
-                if func.scope == "self" and to_delete_field.get(func.value):
-                    to_delete_field.pop(func.value)
-            if node.get_node_type() == NodeType.CallMethod and node.get_func() == PASS_THROUGH_METHOD:
-                var_name = node.get_args()[0].value
-                if to_delete_field.get(var_name):
-                    to_delete_field.pop(var_name)
-        # filter field used in test-of-if of construct function
-        for body in self._root_ast.body:
-            if not isinstance(body, ast.If):
+        for func_def in self._class_ast.body:
+            if not isinstance(func_def, ast.FunctionDef):
                 continue
-            test = body.test
-            field_finder = FieldFinder(test)
-            to_delete_to_delete_keys = []
-            for key, _ in to_delete_field.items():
-                if field_finder.check(key):
-                    to_delete_to_delete_keys.append(key)
-            for key in to_delete_to_delete_keys:
-                to_delete_field.pop(key)
-        # filter field used in test-of-if of init function
-        for body in self._init_func_ast.body:
-            if not isinstance(body, ast.If):
-                continue
-            test = body.test
-            field_finder = FieldFinder(test)
-            to_delete_to_delete_keys = []
-            for key, _ in to_delete_field.items():
-                if field_finder.check(key):
-                    to_delete_to_delete_keys.append(key)
-            for key in to_delete_to_delete_keys:
-                to_delete_field.pop(key)
+            if func_def.name != "__init__":
+                to_delete_to_delete_keys = []
+                property_checker = CheckPropertyIsUsed(func_def)
+                for key, _ in self._deleted_field.items():
+                    if property_checker.check("self", key):
+                        to_delete_to_delete_keys.append(key)
+                        property_checker = CheckPropertyIsUsed(func_def)
+                for key in to_delete_to_delete_keys:
+                    self._deleted_field.pop(key)
+            else:
+                for body in func_def.body:
+                    if not isinstance(body, ast.If):
+                        continue
+                    test = body.test
+                    field_finder = FieldFinder(test)
+                    to_delete_to_delete_keys = []
+                    for key, _ in self._deleted_field.items():
+                        if field_finder.check(key):
+                            to_delete_to_delete_keys.append(key)
+                    for key in to_delete_to_delete_keys:
+                        self._deleted_field.pop(key)
 
     def _remove_unused_field(self):
         """remove unused field in __init__ function"""
-        to_delete_field = {}
         multi_targets = []
         for index, body in enumerate(self._init_func_ast.body):
             if not isinstance(body, ast.Assign):
@@ -1121,12 +1144,12 @@ class SymbolTree(Observer, Observable):
             for target in targets:
                 if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) \
                         and target.value.id == "self":
-                    to_delete_field[target.attr] = index
+                    self._deleted_field[target.attr] = index
                     if len(targets) > 1:
                         multi_targets.append(index)
-        self._filter_out_to_delete_field(to_delete_field)
+        self._filter_out_to_delete_field(self._deleted_field)
         for i in range(len(self._init_func_ast.body) - 1, -1, -1):
-            if i in to_delete_field.values():
+            if i in self._deleted_field.values():
                 if i in multi_targets:
                     raise RuntimeError("Can not erase field ast node in __init__ function because of multi-targets")
                 AstModifier.erase_ast_from_function(self._init_func_ast, self._init_func_ast.body[i])
@@ -1144,12 +1167,9 @@ class SymbolTree(Observer, Observable):
                     self._module_ast.body.remove(body)
 
     def _get_real_node(self, node_or_name: Union[Node, str]) -> Optional[Node]:
-        if isinstance(node_or_name, Node):
-            result = self.get_node(node_or_name.get_name())
-            return result if result is node_or_name else None
         if isinstance(node_or_name, str):
             return self.get_node(node_or_name)
-        return None
+        return node_or_name
 
     def _insert_tree(self, position: Position, root: Node, insert_to_ast: bool = True) -> Node:
         """
@@ -1298,7 +1318,7 @@ class SymbolTree(Observer, Observable):
                 raise TypeError("value should be ScopedValue, got: ", type(value))
             if value.type == ValueType.CustomObjValue:
                 field = self._node_name_namer.get_name(f"var_{type(value.value).__name__}")
-                self._global_vars[field] = value.value
+                setattr(self._origin_network, field, value.value)
                 init_targets = [ScopedValue.create_naming_value(field, "self")]
                 AstModifier.append_global_vars_expr_to_init(self._init_func_ast, init_targets, field)
                 result[arg] = init_targets[0]
@@ -1316,15 +1336,34 @@ class SymbolTree(Observer, Observable):
         Returns:
             A class handle.
         """
-        file_name = "new_network_{0}.py".format(int(time.time() * 10000))
-        with os.fdopen(os.open(file_name, os.O_WRONLY | os.O_CREAT, stat.S_IRWXU), 'wb') as f:
+        self._update_container()
+        file_path = os.getcwd()
+        file_path = os.path.join(file_path, "rewritten_network")
+        if not os.path.exists(file_path):
+            os.mkdir(file_path)
+        file_name = "{0}_{1}.py".format(self._opt_cls_name, id(self))
+        network_file = os.path.join(file_path, file_name)
+        with os.fdopen(os.open(network_file, os.O_WRONLY | os.O_CREAT, stat.S_IRWXU), 'wb') as f:
             source = self.get_code()
             f.write(source.encode('utf-8'))
             f.flush()
-        tmp_module_path, tmp_module_file = os.path.split(file_name)
+            os.fsync(f)
+        tmp_module_path, tmp_module_file = os.path.split(network_file)
         tmp_module_name = tmp_module_file[:-3]
         sys.path.append(tmp_module_path)
-        tmp_module = importlib.import_module(tmp_module_name)
+        tmp_module = None
+
+        i = 0
+        while not tmp_module:
+            try:
+                tmp_module = importlib.import_module(tmp_module_name)
+            except ModuleNotFoundError:
+                if i > 10:
+                    break
+                time.sleep(0.1)
+                i += 1
+        if not tmp_module:
+            logger.error(f"load module {tmp_module_name} failed.")
         network_cls = getattr(tmp_module, self._opt_cls_name)
         if network_cls is None:
             raise RuntimeError("Can not find network class:", self._opt_cls_name)
@@ -1333,3 +1372,87 @@ class SymbolTree(Observer, Observable):
     def _on_change(self, event: Event):
         self._modified = True
         self.changed(event)
+
+    def _update_container(self):
+        """Update instance of node in container."""
+        for node in self.nodes():
+            index = 0
+            if node.get_node_type() == NodeType.CellContainer:
+                for n in node.node_list:
+                    if not n.valid:
+                        continue
+                    if n.get_node_type() == NodeType.Tree:
+                        obj = n.symbol_tree.get_network()
+                        node.get_instance()[index] = obj
+                    else:
+                        node.get_instance()[index] = n.get_instance()
+                    index += 1
+
+    def _cal_difference_set(self, input, other):
+        """Calculate different set of two sets."""
+        set1 = set(input)
+        set2 = set(other)
+        return set1 - set2
+
+    def _merge_origin_property(self, new_net):
+        """Merge property of two network."""
+        tmp = self._cal_difference_set(dir(self._origin_network), dir(new_net))
+        new_attr_names = self._cal_difference_set(tmp, self._deleted_field.keys())
+        for name in new_attr_names:
+            setattr(new_net, name, getattr(self._origin_network, name))
+        # merger cells
+        cells = self._cal_difference_set(self._origin_network.name_cells().keys(), new_net.name_cells().keys())
+        cells = self._cal_difference_set(cells, self._deleted_node)
+        for c in cells:
+            new_net.insert_child_to_cell(c, self._origin_network.name_cells()[c])
+        # merge primitives
+        primitives = self._cal_difference_set(self._origin_network._primitives.keys(), new_net._primitives.keys())
+        for p in primitives:
+            new_net._primitives[p] = self._origin_network._primitives[p]
+
+    def _update_names_for_unique(self, node: ast.AST):
+        """ Update names of ast nodes for unique. """
+        if isinstance(node, (ast.For, ast.If, ast.While)):
+            self._update_names_for_unique_branchs(node)
+        elif isinstance(node, ast.Assign):
+            self._update_names_for_unique(node.value)
+            for target in node.targets:
+                self._update_names_for_unique(target)
+        elif isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Attribute):
+                self._update_names_for_unique(node.func.value)
+            for arg in node.args:
+                self._update_names_for_unique(arg)
+            for keyword in node.keywords:
+                self._update_names_for_unique(keyword)
+        elif isinstance(node, ast.UnaryOp):
+            self._update_names_for_unique(node.operand)
+        elif isinstance(node, ast.BinOp):
+            self._update_names_for_unique(node.left)
+            self._update_names_for_unique(node.right)
+        elif isinstance(node, (ast.Attribute, ast.Subscript, ast.Return)):
+            self._update_names_for_unique(node.value)
+        elif isinstance(node, (ast.List, ast.Tuple)):
+            for elt in node.elts:
+                self._update_names_for_unique(elt)
+        elif isinstance(node, ast.Compare):
+            for comparator in node.comparators:
+                self._update_names_for_unique(comparator)
+        elif isinstance(node, ast.Name):
+            node.id = self._target_namer.get_real_arg(node.id)
+
+    def _update_names_for_unique_branchs(self, node: Union[ast.For, ast.If, ast.While]):
+        """ Update names of ast nodes for unique with ast.For, ast.If or ast.While """
+        if isinstance(node, ast.For):
+            self._update_names_for_unique(node.target)
+            self._update_names_for_unique(node.iter)
+            for body in node.body:
+                self._update_names_for_unique(body)
+            for body in node.orelse:
+                self._update_names_for_unique(body)
+        elif isinstance(node, (ast.If, ast.While)):
+            self._update_names_for_unique(node.test)
+            for body in node.body:
+                self._update_names_for_unique(body)
+            for body in node.orelse:
+                self._update_names_for_unique(body)

mindspore/rewrite/symbol_tree_builder.py

@@ -28,6 +28,41 @@ from .ast_helpers import AstModifier
 from .ast_helpers import AstFinder
 
 
+class FunctionSymbolTreeBuilder:
+    """Create function SymbolTree"""
+    def __init__(self, network: Cell, ast_root):
+        self._origin_net = network
+        self._ast_root: ast.Module = ast_root
+        self._root_tree: Optional[SymbolTree] = None
+
+    @staticmethod
+    def _ast_transform(ast_root: ast.AST) -> ast.AST:
+        """
+        Optimize ast before parse.
+
+        Args:
+             ast_root (ast.AST): An instance of ast to be optimized.
+
+        Returns:
+             An instance of ast been optimized.
+        """
+        transform_list = [FlattenRecursiveStmt()]
+        for transformer in transform_list:
+            ast_root = transformer.transform(ast_root)
+        return ast_root
+
+    def build(self) -> SymbolTree:
+        """
+        Build SymbolTree.
+
+        Returns:
+             An instance of SymbolTree.
+        """
+        self._root_tree: SymbolTree = SymbolTree(self._origin_net, self._ast_root)
+        self._root_tree.finish_build()
+        return self._root_tree
+
+
 class SymbolTreeBuilder:
     """
     `SymbolTreeBuilder` for building a SymbolTree from network.
@@ -43,6 +78,8 @@ class SymbolTreeBuilder:
         network_str = inspect.getsource(type(network))
         self._ast_root: ast.Module = ast.parse(network_str)
         self._root_tree: Optional[SymbolTree] = None
+        if isinstance(network, Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     @staticmethod
     def merge_module_of_subtree(main_tree: SymbolTree, sub_stree: SymbolTree):
@@ -140,7 +177,7 @@ class SymbolTreeBuilder:
         """
 
         for node in self._root_tree.nodes():
-            if isinstance(node, TreeNode):
+            if isinstance(node, TreeNode) and node.get_instance():
                 SymbolTreeBuilder.merge_module_of_subtree(self._root_tree, node.symbol_tree)
 
     def _reduce_redundant_import(self):