mindspore 2.1.0__cp39-none-any.whl → 2.2.11__cp39-none-any.whl

mindspore/rewrite/api/symbol_tree.py

@@ -13,8 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """Rewrite module api: SymbolTree."""
-from typing import Optional, Union
-from types import FunctionType
+from typing import Optional, Union, List
 import mindspore as ms
 
 from mindspore.nn import Cell
@@ -53,7 +52,61 @@ class SymbolTree:
 
         This interface parses the `network` instance, expands each source
         code statement of the forward computation process, and parses it into nodes,
-        which is stored in the SymbolTree.
+        which is stored in the SymbolTree. The specific process is as follows:
+
+        1. Obtain the source code of the network instance.
+        2. Perform AST parsing on the network and obtain the AST nodes (abstract syntax trees) of each
+           statement in the network.
+        3. Expand complex statements in the network forward evaluation process into multiple simple statements.
+        4. Create a SymbolTree object. Each SymbolTree corresponds to one network instance.
+        5. Use the rewrite node to store each statement of the network forward computation process. The node records
+           the input, output, and other information of the statement.
+        6. Save the rewrite node to the SymbolTree, and update and maintain the topological connection between
+           the nodes.
+        7. Return the SymbolTree object corresponding to the network instance.
+
+        If a user-defined network of type :class:`mindspore.nn.Cell` is called in the forward computation process
+        of the network, rewrite will generate a node of type `NodeType.Tree` for the corresponding statement. This
+        type of node stores a new SymbolTree, which parses and maintains the node information of the user-defined
+        network.
+
+        If the following types of statements are called in the forward computation process of the network, rewrite
+        will parse the internal statements in the statement and generate corresponding nodes:
+
+        - :class:`mindspore.nn.SequentialCell`
+        - Functions within classes
+        - Control flow statements, such as `if` statements
+
+        Note:
+            Because the specific execution branch of control flows are still unknown during the rewrite operation
+            of the network, no topology information will be established between the nodes inside the control flow
+            and the nodes outside.
+            Users cannot obtain nodes inside the control flow when they acquire nodes outside the control flow using
+            interfaces like :func:`mindspore.rewrite.Node.get_inputs` and :func:`mindspore.rewrite.Node.get_users` .
+            Users also cannot obtain nodes outside the control flow, if they use these interfaces inside the control
+            flow.
+            Therefore, when users modify the network, they need to manually handle the node information inside and
+            outside the control flow.
+
+        The current rewrite module has the following syntax limitations:
+
+        - Only networks of type :class:`mindspore.nn.Cell` are supported as input to the rewrite module.
+        - Parsing assignment statements with multiple output values is not currently supported.
+        - Parsing loop statements is not currently supported.
+        - Parsing decorator syntax is not currently supported.
+        - Parsing class variable syntax is not currently supported. If class variable uses external data,
+            the network after rewrite may be missing data.
+        - Parsing local classes and embedded classes is not currently supported, that is, the definition
+            of classes need to be placed on the outermost layer.
+        - Parsing closure syntax is not currently supported, that is, the definition of out-of-class
+            functions need to be placed at the outermost layer.
+        - Parsing lambda expression syntax is not currently supported.
+
+        For statements that do not support parsing, rewrite will generate nodes of type `NodeType.Python`
+        for corresponding statements to ensure that the network after rewrite can run normally.
+        The `Python` node does not support modifying the input and output of statements, and there may be
+        a problem between variable names and those generated by the rewrite. In this case, users need to
+        adjust the variable names manually.
 
         Args:
             network (Cell): `network` used to create SymbolTree.
@@ -67,7 +120,7 @@ class SymbolTree:
         Examples:
             >>> from mindspore.rewrite import SymbolTree
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> print(type(stree))
@@ -90,42 +143,37 @@ class SymbolTree:
             if v not in MsDtypes and not isinstance(v, ParamTypes):
                 raise TypeError(f"For call-function Node, got unsupported kwarg value: {v}, type: {type(v)}")
 
-    def create_call_function(self, func, targets, *args, **kwargs): # pylint: disable=C0111
-        Validator.check_value_type("func", func, [FunctionType], "SymbolTree node")
-        Validator.check_element_type_of_iterable("targets", targets, [str], "SymbolTree node")
-        args_ = list(args)
-        SymbolTree._check_args_type(args_)
-        for i, arg in enumerate(args_):
-            if isinstance(arg, Node):
-                args_[i] = arg.get_handler()
-        SymbolTree._check_kwargs_type(kwargs)
-        for key, value in kwargs.items():
-            if isinstance(value, Node):
-                kwargs[key] = value.get_handler()
-        return Node(self._symbol_tree._create_call_function(func, targets, args_, kwargs)) # pylint: disable=W0212
-
     def get_handler(self) -> SymbolTreeImpl:
         return self._symbol_tree
 
-    def nodes(self):
+    def nodes(self, all_nodes: bool = False):
         """
         Get the generator of the node in the current SymbolTree, which is used to iterate
         through the nodes in SymbolTree.
 
+        Args:
+            all_nodes (bool): Get all nodes including nodes in CallFunction node, CellContainer node
+                and sub symbol tree. Default: ``False`` .
+
         Returns:
-            A generator for node of current SymbolTree.
+            A generator for nodes in SymbolTree.
+
+        Raises:
+            TypeError: If `all_nodes` is not bool.
 
         Examples:
             >>> from mindspore.rewrite import SymbolTree
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> print([node.get_name() for node in stree.nodes()])
             ['input_x', 'Expr', 'conv1', 'relu', 'max_pool2d', 'conv2', 'relu_1', 'max_pool2d_1',
              'flatten', 'fc1', 'relu_2', 'fc2', 'relu_3', 'fc3', 'return']
         """
-        for node in self._symbol_tree.nodes():
+        Validator.check_value_type("all_nodes", all_nodes, [bool], "nodes")
+        nodes = self._symbol_tree.all_nodes() if all_nodes else self._symbol_tree.nodes()
+        for node in nodes:
             yield Node(node)
 
     def get_node(self, node_name: str) -> Optional[Node]:
@@ -141,7 +189,7 @@ class SymbolTree:
         Examples:
             >>> from mindspore.rewrite import SymbolTree
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> node = stree.get_node('conv1')
@@ -149,12 +197,12 @@ class SymbolTree:
             conv1
         """
         Validator.check_value_type("node_name", node_name, [str], "SymbolTree")
-        node_impl = self._symbol_tree.get_node(node_name)
+        node_impl = self._symbol_tree.get_node_from_name(node_name)
         if node_impl is None:
             return None
         return Node(node_impl)
 
-    def get_inputs(self) -> [Node]:
+    def get_inputs(self) -> List[Node]:
         return [Node(node_impl) for node_impl in self._symbol_tree.get_inputs()]
 
     def before(self, node: Union[Node, str]):
@@ -174,15 +222,17 @@ class SymbolTree:
         Examples:
             >>> from mindspore.rewrite import SymbolTree
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> for node in stree.nodes():
             ...     if node.get_name() == "conv1":
             ...         position = stree.before(node)
         """
-        Validator.check_value_type("node", node, [Node], "SymbolTree")
-        return self._symbol_tree.before(node.get_handler())
+        Validator.check_value_type("node", node, [Node, str], "SymbolTree")
+        if isinstance(node, Node):
+            node = node.get_handler()
+        return self._symbol_tree.before(node)
 
     def after(self, node: Union[Node, str]):
         """
@@ -201,15 +251,17 @@ class SymbolTree:
         Examples:
             >>> from mindspore.rewrite import SymbolTree
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> for node in stree.nodes():
             ...     if node.get_name() == "conv1":
             ...         position = stree.after(node)
         """
-        Validator.check_value_type("node", node, [Node], "SymbolTree")
-        return self._symbol_tree.after(node.get_handler())
+        Validator.check_value_type("node", node, [Node, str], "SymbolTree")
+        if isinstance(node, Node):
+            node = node.get_handler()
+        return self._symbol_tree.after(node)
 
     def insert(self, position, node: Node) -> Node:
         """
@@ -233,7 +285,7 @@ class SymbolTree:
             >>> from mindspore.rewrite import SymbolTree, ScopedValue
             >>> import mindspore.nn as nn
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> node = stree.get_node("conv1")
@@ -244,7 +296,7 @@ class SymbolTree:
         """
         Validator.check_value_type("position", position, [Position], "SymbolTree")
         Validator.check_value_type("node", node, [Node], "SymbolTree")
-        return Node(self._symbol_tree.insert_node(position, node.get_handler()))
+        return Node(self._symbol_tree.insert_node(node.get_handler(), position.node, position.before_node))
 
     def erase(self, node: Union[Node, str]) -> Optional[Node]:
         """
@@ -262,16 +314,18 @@ class SymbolTree:
         Examples:
             >>> from mindspore.rewrite import SymbolTree
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> node = stree.get_node("conv1")
             >>> stree.erase(node)
         """
-        Validator.check_value_type("node", node, [Node], "SymbolTree")
-        return Node(self._symbol_tree.erase_node(node.get_handler()))
+        Validator.check_value_type("node", node, [Node, str], "SymbolTree")
+        if isinstance(node, Node):
+            node = node.get_handler()
+        return Node(self._symbol_tree.erase_node(node))
 
-    def replace(self, old_node: Node, new_nodes: [Node]) -> Node:
+    def replace(self, old_node: Node, new_nodes: List[Node]) -> Node:
         """
         Replace the `old_node` with nodes in the `new_nodes` list.
 
@@ -285,7 +339,7 @@ class SymbolTree:
 
         Args:
             old_node (Node): Node to be replaced.
-            new_nodes (list[Node]): Nodes of the node_tree to replace in.
+            new_nodes (List[Node]): Nodes of the node_tree to replace in.
 
         Returns:
             An instance of Node represents root of node_tree been replaced in.
@@ -299,7 +353,7 @@ class SymbolTree:
             >>> from mindspore.rewrite import SymbolTree, ScopedValue
             >>> import mindspore.nn as nn
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> node = stree.get_node("conv1")
@@ -320,16 +374,38 @@ class SymbolTree:
     def dump(self):
         self._symbol_tree.dump()
 
-    def print_node_tabulate(self):
-        """
+    def print_node_tabulate(self, all_nodes: bool = False):
+        r"""
         Print the topology information of nodes in SymbolTree, including node type, node name, node code,
         and node input-output relationship.
-        The information is output to the screen using the print interface.
 
-        .. warning::
-            This is an experimental API that is subject to change or deletion.
+        The information is output to the screen using the print interface, including the following information:
+
+        - **node type** (str): The type of node, refer to class:`mindspore.rewrite.NodeType` .
+        - **name** (str): The name of node.
+        - **codes** (str): The source code statement corresponding to the node.
+        - **arg providers** (Dict[int, Tuple[str, int]]): The format is `{[idx, (n, k)]}` , which means the
+          `idx` th parameter of the node is provided by the `k` th output of node `n` .
+        - **target users** (Dict[int, List[Tuple[str, int]]]): The format is '{[idx, [(n, k)]]}' , which means
+          the `idx` th output of the node is used as the `k` th parameter of node `n` .
+
+        Args:
+            all_nodes (bool): Print information of all nodes, including nodes in CallFunction
+                node, CellContainer node and sub symbol tree. Default: ``False`` .
+
+        Raises:
+            TypeError: If `all_nodes` is not bool.
+
+        Examples:
+            >>> from mindspore.rewrite import SymbolTree
+            >>> # Define the network structure of LeNet5. Refer to
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+            >>> net = LeNet5()
+            >>> stree = SymbolTree.create(net)
+            >>> stree.print_node_tabulate()
         """
-        self._symbol_tree.print_node_tabulate()
+        Validator.check_value_type("all_nodes", all_nodes, [bool], "print_node_tabulate")
+        self._symbol_tree.print_node_tabulate(all_nodes)
 
     def get_code(self) -> str:
         """
@@ -342,7 +418,7 @@ class SymbolTree:
         Examples:
             >>> from mindspore.rewrite import SymbolTree
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> codes = stree.get_code()
@@ -355,13 +431,21 @@ class SymbolTree:
         Get the network object generated based on SymbolTree.
         The source code is saved to a file in the 'rewritten_network' folder of the current directory.
 
+        Note:
+            - The modification of network by rewrite module is based on the modification of AST tree of
+              original network instance, and the new network instance will obtain attribute information
+              from original network instance, so the new network instance and the original network instance
+              have data association, and the original network should no longer be used.
+            - Due to the data association between the new network and the original network instance, manually creating
+              a network instance using the source code file generated by rewrite is not currently supported.
+
         Returns:
             A network object generated from SymbolTree.
 
         Examples:
             >>> from mindspore.rewrite import SymbolTree
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> stree = SymbolTree.create(net)
             >>> new_net = stree.get_network()

mindspore/rewrite/ast_helpers/__init__.py

@@ -17,7 +17,8 @@
 Define some ast helpers for manipulating python ast.
 """
 
-from .ast_finder import AstFinder, StrChecker, CheckPropertyIsUsed, GetPropertyOfObj
+from .ast_finder import AstFinder, StrChecker, CheckPropertyIsUsed, GetPropertyOfObj, \
+    AstAssignFinder, AstClassFinder, AstFunctionFinder
 from .ast_replacer import AstReplacer
 from .ast_modifier import AstModifier
 from .ast_creator import ast_args_creator, ast_assign_creator, ast_attributer_creator, ast_call_creator, \

mindspore/rewrite/ast_helpers/ast_finder.py

@@ -225,3 +225,132 @@ class GetPropertyOfObj(ast.NodeVisitor):
         self._property = set()
         self.generic_visit(self._context)
         return self._property
+
+
+class AstAssignFinder(ast.NodeVisitor):
+    """
+    Get assign definition ast of specifical parameter in specific scope.
+
+    Args:
+        node (ast.AST): An instance of ast node as check scope.
+    """
+    def __init__(self, node: ast.AST):
+        self._context = node
+        self._scope = ""
+        self._value = ""
+        self._target = None
+
+    def visit_Assign(self, node: ast.Assign):
+        if self._scope and isinstance(node.targets[0], ast.Attribute):
+            if node.targets[0].attr == self._value and isinstance(node.targets[0].value, ast.Name) \
+                and node.targets[0].value.id == self._scope:
+                self._target = node
+        elif not self._scope and isinstance(node.targets[0], ast.Name):
+            if node.targets[0].id == self._value:
+                self._target = node
+
+    def get_ast(self, value: str, scope: str = "") -> bool:
+        """
+        Get assign ast of specifical parameter in specific ast.
+
+        Args:
+            value (str): A string indicates assign target value.
+            scope (str): A string indicates assign target scope.
+
+        Returns:
+            An assign ast with the same target name as `scope.value` .
+        """
+        self._scope = scope
+        self._value = value
+        self.generic_visit(self._context)
+        return self._target
+
+
+class AstClassFinder(ast.NodeVisitor):
+    """
+    Find all specific name of ast class node in specific scope.
+
+    Args:
+        node (ast.AST): An instance of ast node as search scope.
+    """
+
+    def __init__(self, node: ast.AST):
+        self._scope: ast.AST = node
+        self._target: str = ""
+        self._results: [ast.ClassDef] = []
+
+    def visit_ClassDef(self, node):
+        """
+        An override method, iterating over all ClassDef nodes and save target ast nodes.
+
+        Args:
+            node (ast.AST): An instance of ast node which is visited currently.
+        """
+
+        if node.name == self._target:
+            self._results.append(node)
+
+    def find_all(self, class_name: str) -> [ast.AST]:
+        """
+        Find all matched ast node.
+
+        Args:
+            class_name (str): Name of class to be found.
+
+        Returns:
+            A list of instance of ast.ClassDef as matched result.
+
+        Raises:
+            TypeError: If input `class_name` is not str.
+        """
+        if not isinstance(class_name, str):
+            raise TypeError("Input class_name should be a str")
+        self._target = class_name
+        self._results.clear()
+        self.visit(self._scope)
+        return self._results
+
+
+class AstFunctionFinder(ast.NodeVisitor):
+    """
+    Find all specific name of ast function node in specific scope.
+
+    Args:
+        node (ast.AST): An instance of ast node as search scope.
+    """
+
+    def __init__(self, node: ast.AST):
+        self._scope: ast.AST = node
+        self._target: str = ""
+        self._results: [ast.ClassDef] = []
+
+    def visit_FunctionDef(self, node):
+        """
+        An override method, iterating over all FunctionDef nodes and save target ast nodes.
+
+        Args:
+            node (ast.AST): An instance of ast node which is visited currently.
+        """
+
+        if node.name == self._target:
+            self._results.append(node)
+
+    def find_all(self, func_name: str) -> [ast.AST]:
+        """
+        Find all matched ast node.
+
+        Args:
+            func_name (str): Name of function to be found.
+
+        Returns:
+            A list of instance of ast.FunctionDef as matched result.
+
+        Raises:
+            TypeError: If input `func_name` is not str.
+        """
+        if not isinstance(func_name, str):
+            raise TypeError("Input func_name should be a str")
+        self._target = func_name
+        self._results.clear()
+        self.visit(self._scope)
+        return self._results