mindspore 2.0.0a0__cp38-cp38-win_amd64.whl → 2.0.0rc1__cp38-cp38-win_amd64.whl

mindspore/rewrite/__init__.py

@@ -33,6 +33,8 @@ from .api.node import Node
 from .api.node_type import NodeType
 from .api.pattern_engine import PatternEngine, PatternNode, VarNode, Replacement
 from .api.tree_node_helper import TreeNodeHelper
+from .sparsify.sparsify import sparsify
+from .sparsify.utils import ArgType, SparseFunc
 
 __all__ = ["SymbolTree", "Node", "NodeType", "ScopedValue", "ValueType", "PatternEngine", "PatternNode", "VarNode",
-           "Replacement", "TreeNodeHelper"]
+           "Replacement", "TreeNodeHelper", "sparsify", "ArgType", "SparseFunc"]

mindspore/rewrite/api/node.py

@@ -18,7 +18,7 @@ from typing import Union, Optional
 
 from mindspore.nn import Cell
 from mindspore.ops.primitive import Primitive
-from ..._checkparam import Validator
+from mindspore import _checkparam as Validator
 from ..node import Node as NodeImpl
 from ..symbol_tree import SymbolTree as SymbolTreeImpl
 from .node_type import NodeType
@@ -99,12 +99,6 @@ class Node:
                                             args, kwargs, name, is_sub_net))
 
     def get_handler(self) -> NodeImpl:
-        """
-        Get handler of node implementation.
-
-        Returns:
-            An instance of `NodeImpl`.
-        """
         return self._node
 
     def get_inputs(self) -> ['Node']:
@@ -181,7 +175,6 @@ class Node:
 
         Raises:
             RuntimeError: If `src_node` is not belong to current `SymbolTree`.
-            RuntimeError: If current node and `src_node` is not belong to same `SymbolTree`.
             TypeError: If `arg_idx` is not a `int` number.
             ValueError: If `arg_idx` is out of range.
             TypeError: If `src_node` is not a `Node` instance.
@@ -209,27 +202,6 @@ class Node:
             belong_symbol_tree.set_node_arg_by_node(self._node, arg_idx, src_node.get_handler(), out_idx)
 
     def get_targets(self) -> [ScopedValue]:
-        """
-        Get targets of current node.
-
-        - When node_type of current node is `CallCell`, `CallPrimitive`, `CallMethod` or `Tree`, `targets` are strings
-          represents invoke result of the cell-op or primitive-op or function-call which are corresponding to targets of
-          ast.Assign.
-        - When node_type of current node is Input, `targets` should have only one element which is a string represents
-          parameter of function.
-        - When node_type of current node is `Python` or `Output`, `targets` are don't-care.
-
-        Returns:
-            A list of instances of ScopedValue as targets of node.
-
-        Examples:
-            >>> from mindspore.rewrite import SymbolTree
-            >>> from lenet import Lenet
-            >>> net = Lenet()
-            >>> stree = SymbolTree.create(net)
-            >>> node = stree.get_node("conv1")
-            >>> targets = node.get_targets()
-        """
         return self._node.get_targets()
 
     def get_name(self) -> str:
@@ -284,106 +256,21 @@ class Node:
         return self._node.get_instance_type()
 
     def get_instance(self):
-        """
-        Get the instance of current node.
-
-        - When node_type of current node is `CallCell`, instance is an instance of Cell.
-        - When node_type of current node is `CallPrimitive`, instance is an instance of primitive.
-        - When node_type of current node is `Tree`, instance is an instance of network-cell.
-        - When node_type of current node is `Python`, `Input`, `Output` or `CallMethod`, instance should be None.
-
-        Returns:
-            A object represents corresponding instance of current node.
-        """
         return self._node.get_instance()
 
     def get_args(self) -> [ScopedValue]:
-        """
-        Get the arguments of current node.
-
-        - When `node_type` of current node is `CallCell`, `CallPrimitive` or `Tree`, arguments are corresponding to args
-          of ast.Call which represents arguments to invoke forward method of cell-op or primitive-op.
-        - When `node_type` of current node is `Input`, arguments represents default-value of argument of function.
-        - When `node_type` of current node is `Output`, arguments represents the return values of network.
-        - When `node_type` of current node is `Python`, arguments are don't-care.
-
-        Returns:
-            A list of instances of `ScopedValue`.
-
-        Examples:
-            >>> from mindspore.rewrite import SymbolTree
-            >>> from lenet import Lenet
-            >>> net = Lenet()
-            >>> stree = SymbolTree.create(net)
-            >>> node = stree.get_node("conv1")
-            >>> args = node.get_args()
-        """
         return self._node.get_args()
 
     def get_kwargs(self) -> {str: ScopedValue}:
-        """
-        Get the keyword arguments of current node.
-
-        - When node_type of current node is `CallCell`, `CallPrimitive` or `Tree`, keyword arguments are corresponding
-          to kwargs of ast.Call which represents arguments to invoke forward method of cell-op or primitive-op.
-        - When node_type of current node is `Python`, `Input` or `Output`, keyword arguments are don't-care.
-
-        Returns:
-            A dict of str to instance of `ScopedValue`.
-
-        Examples:
-            >>> from mindspore.rewrite import SymbolTree
-            >>> from lenet import Lenet
-            >>> net = Lenet()
-            >>> stree = SymbolTree.create(net)
-            >>> node = stree.get_node("conv1")
-            >>> kwargs = node.get_kwargs()
-        """
         return self._node.get_kwargs()
 
     def set_attribute(self, key: str, value):
-        """
-        Set attribute of current node.
-
-        Args:
-            key (str): Key of attribute.
-            value (object): Value of attribute.
-
-        Raises:
-            TypeError: If `key` is not a `str`.
-
-        Examples:
-            >>> from mindspore.rewrite import SymbolTree
-            >>> from lenet import Lenet
-            >>> net = Lenet()
-            >>> stree = SymbolTree.create(net)
-            >>> node = stree.get_node("conv1")
-            >>> node.set_attribute("channel", 3)
-        """
         Validator.check_value_type("key", key, [str], "Node attribute")
         self._node.set_attribute(key, value)
 
     def get_attributes(self) -> {str: object}:
-        """
-        Get all attributes of current node.
-
-        Returns:
-            A dict of str to instance of object as attributes.
-        """
         return self._node.get_attributes()
 
     def get_attribute(self, key: str):
-        """
-        Get attribute of current node by key.
-
-        Args:
-            key (str): Key of attribute.
-
-        Returns:
-            A object as attribute, can be any type.
-
-        Raises:
-            TypeError: If `key` is not a `str`.
-        """
         Validator.check_value_type("key", key, [str], "Node attribute")
         return self._node.get_attribute(key)

mindspore/rewrite/api/node_type.py

@@ -29,6 +29,7 @@ class NodeType(Enum):
     - Input: `Input` node represents input of `SymbolTree` corresponding to arguments of forward method.
     - Output: `Output` node represents output of SymbolTree corresponding to return statement of forward method.
     - Tree: `Tree` node represents sub-network invoking in forward method.
+    - MathOps: `MathOps` node represents a mathematical operation, such as adding or comparing in forward method.
 
     """
     Unknown = 0
@@ -43,3 +44,5 @@ class NodeType(Enum):
     Input = 7
     Output = 8
     Tree = 9
+    CellContainer = 10
+    MathOps = 11

mindspore/rewrite/api/pattern_engine.py

@@ -20,7 +20,7 @@ import abc
 from mindspore.nn import Cell
 from mindspore.ops.primitive import Primitive
 from mindspore import log as logger
-from ..._checkparam import Validator
+from mindspore import _checkparam as Validator
 from .node_type import NodeType
 from .node import Node
 from .symbol_tree import SymbolTree
@@ -308,6 +308,16 @@ class PatternEngine:
                 queue.extend(inputs_dict.get(cur_node.get_name()))
         return new_root
 
+    @staticmethod
+    def _multi_replace_cellcontainer(stree, cellcontainer, node, matched_dict, new_nodes):
+        """Replace node in CellContainer."""
+        to_erase_list = list(matched_dict.values())
+        stree.replace(Node(node), new_nodes)
+        for n in reversed(to_erase_list):
+            if n.get_handler() is node:
+                continue
+            stree.erase_node(n)
+
     def apply(self, stree: SymbolTree) -> bool:
         """
         Apply current pattern to a `SymbolTree`.
@@ -359,6 +369,9 @@ class PatternEngine:
                 visited.append(cur_node)
                 queue.extend(cur_node.get_users())
                 continue
+            if cur_node.get_node_type() == NodeType.CellContainer:
+                self._process_cellcontainer(stree, cur_node.get_handler())
+                continue
             visited.append(cur_node)
             matched, matched_dict = self._match(self._pattern, cur_node)
             # not matched
@@ -460,3 +473,20 @@ class PatternEngine:
                     logger.debug("Check match failed, pattern leaked")
                     return False
         return True
+
+    def _process_cellcontainer(self, stree, cellcontainer):
+        """Process CellContainer node."""
+        for node in cellcontainer.nodes():
+            if node.get_node_type() == NodeType.Tree:
+                subtree = node.symbol_tree
+                self.apply(SymbolTree(subtree))
+                continue
+            matched, matched_dict = self._match(self._pattern, Node(node))
+            if not matched:
+                continue
+            new_nodes = []
+            if self._replacement is not None:
+                new_nodes = self._replacement(self._pattern, self._is_chain, matched_dict)
+            if not new_nodes:  # if replacement is empty, do nothing
+                continue
+            PatternEngine._multi_replace_cellcontainer(stree, cellcontainer, node, matched_dict, new_nodes)

mindspore/rewrite/api/scoped_value.py

@@ -15,7 +15,7 @@
 """Rewrite module api: ValueType and ScopedValue."""
 from enum import Enum
 from typing import Optional, Union
-from ..._checkparam import Validator
+from mindspore import _checkparam as Validator
 
 
 class ValueType(Enum):
@@ -127,8 +127,8 @@ class ScopedValue:
         Create a list of naming `ScopedValue`.
 
         Args:
-            names: (list[str] or tuple[str]): List or tuple of `str` represents names of referenced variables.
-            scopes: (list[str] or tuple[str]): List or tuple of `str` represents scopes of referenced variables.
+            names (list[str] or tuple[str]): List or tuple of `str` represents names of referenced variables.
+            scopes (list[str] or tuple[str]): List or tuple of `str` represents scopes of referenced variables.
 
         Returns:
             An list of instance of `ScopedValue`.
@@ -140,7 +140,7 @@ class ScopedValue:
 
         Examples:
             >>> from mindspore.rewrite import ScopedValue
-            >>> variables = ScopedValue.create_name_values(["z", "z_1"]), name="subnet")
+            >>> variables = ScopedValue.create_name_values(["z", "z_1"], name="subnet")
         """
         Validator.check_element_type_of_iterable("names", names, [str], "ScopedValue")
         if scopes is not None:

mindspore/rewrite/api/symbol_tree.py

@@ -18,7 +18,7 @@ from types import FunctionType
 import mindspore as ms
 
 from mindspore.nn import Cell
-from ..._checkparam import Validator
+from mindspore import _checkparam as Validator
 from .node import Node
 from ..symbol_tree_builder import SymbolTreeBuilder
 from ..symbol_tree import Position, SymbolTree as SymbolTreeImpl
@@ -70,40 +70,7 @@ 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):
-        r"""
-        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.
-
-        Args:
-            func (FunctionType): The function to be called.
-            targets (list[str]): indicates the output name. As the output of the node in the source code.
-            args (Union[MsDtypes, ParamTypes]): parameter name of the node. Used as a parameter to a code statement in
-                source code. The default value is None, which means there is no parameter input in the cell.
-            kwargs (dict{str,Union[MsDtypes, ParamTypes]}): The key type must be str,
-                and the value must be value or type must be ParamTypes.
-                The input parameter name used to describe the formal parameter with a keyword.
-                Enter the name in the source code as the 'kwargs' in the statement expression.The default value is
-                None, which means there is no 'kwargs' input.
-
-        Returns:
-            An instance of `Node`.
-
-        Raises:
-            TypeError: If `func` is not FunctionType.
-            TypeError: If `targets` is not `list`.
-            TypeError: If the type of `targets` is not str.
-            TypeError: If arg in `args` is not ParamType.
-            TypeError: If key of `kwarg` is not a str or value of kwarg in `kwargs` is not ParamType.
-
-        Examples:
-            >>> from mindspore.rewrite import SymbolTree
-            >>> from lenet import Lenet
-            >>> net = Lenet()
-            >>> stree = SymbolTree.create(net)
-            >>> node = stree.get_node("conv1")
-            >>> new_node = stree.create_call_function(F.abs, ["x"], node)
-        """
+    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)
@@ -115,22 +82,9 @@ class SymbolTree:
         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))
+        return Node(self._symbol_tree._create_call_function(func, targets, args_, kwargs)) # pylint: disable=W0212
 
     def get_handler(self) -> SymbolTreeImpl:
-        """
-        Get handler of `SymbolTree` implementation.
-
-        Returns:
-            An instance of `SymbolTree`.
-
-        Examples:
-            >>> from mindspore.rewrite import SymbolTree
-            >>> from lenet import Lenet
-            >>> net = Lenet()
-            >>> stree = SymbolTree.create(net)
-            >>> handler = stree.get_handler()
-        """
         return self._symbol_tree
 
     def nodes(self):
@@ -152,25 +106,6 @@ class SymbolTree:
             yield Node(node)
 
     def get_node(self, node_name: str) -> Optional[Node]:
-        """
-        Get node by `node_name`.
-
-        Args:
-            node_name (str): A string represents name of node.
-
-        Returns:
-            An instance of node if find else None.
-
-        Raises:
-            TypeError: If `node_name` is not `str`.
-
-        Examples:
-            >>> from mindspore.rewrite import SymbolTree
-            >>> from lenet import Lenet
-            >>> net = Lenet()
-            >>> stree = SymbolTree.create(net)
-            >>> node = stree.get_node("conv1")
-        """
         Validator.check_value_type("node_name", node_name, [str], "SymbolTree")
         node_impl = self._symbol_tree.get_node(node_name)
         if node_impl is None:
@@ -354,16 +289,6 @@ class SymbolTree:
         self._symbol_tree.dump()
 
     def print_node_tabulate(self):
-        """
-        Print node information of graph.
-
-        Examples:
-            >>> from mindspore.rewrite import SymbolTree
-            >>> from lenet import Lenet
-            >>> net = Lenet()
-            >>> stree = SymbolTree.create(net)
-            >>> stree.print_node_tabulate()
-        """
         self._symbol_tree.print_node_tabulate()
 
     def get_code(self) -> str:

mindspore/rewrite/api/tree_node_helper.py

@@ -16,7 +16,7 @@
 from typing import Optional
 
 from mindspore import log as logger
-from ..._checkparam import Validator
+from mindspore import _checkparam as Validator
 from .symbol_tree import SymbolTree
 from .node import Node
 from .node_type import NodeType

mindspore/rewrite/ast_creator_register.py

@@ -23,6 +23,7 @@ class Registry(UserDict):
     """Registry class for registry functions for creating ast node."""
 
     def register(self, obj_str, obj):
+        """Register object by str."""
         if isinstance(obj_str, str):
             self[obj_str] = obj
 

mindspore/rewrite/ast_helpers/__init__.py

@@ -17,11 +17,11 @@
 Define some ast helpers for manipulating python ast.
 """
 
-from .ast_finder import AstFinder, StrChecker, FindConstValueInInit
+from .ast_finder import AstFinder, StrChecker, CheckPropertyIsUsed, GetPropertyOfObj
 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, \
     ast_create_arg_value, ast_index_creator, ast_keyword_creator, ast_kwargs_creator, ast_name_creator, \
     ast_num_creator, ast_str_creator, ast_subscript_creator
 
-__all__ = ["AstFinder", "AstReplacer", "AstModifier", "StrChecker"]
+__all__ = ["AstFinder", "AstReplacer", "AstModifier", "StrChecker", "CheckPropertyIsUsed", "GetPropertyOfObj"]

mindspore/rewrite/ast_helpers/ast_creator.py

@@ -49,14 +49,13 @@ def ast_call_creator(func: ast.AST, args: list, keywords: list):
 
 def ast_create_arg_value(value):
     """Create arg node by type."""
-    from mindspore.rewrite.node import Node
     if isinstance(value, (int, float)):
         ast_value = ast_num_creator(value)
     elif isinstance(value, str):
         ast_value = ast_str_creator(value)
     elif value in (ms.float16, ms.float32, ms.float64):
         ast_value = ast_attributer_creator(".".join(["mindspore", str(value).lower()]))
-    elif isinstance(value, Node):
+    elif isinstance(value, ms.rewrite.node.Node):
         ast_value = ast_str_creator(value.get_targets()[0])
     else:
         raise TypeError("Unsupported arg type: ", type(value))

mindspore/rewrite/ast_helpers/ast_finder.py

@@ -160,3 +160,68 @@ class FindConstValueInInit(ast.NodeVisitor):
         self._hit = False
         self.generic_visit(self._context)
         return self._hit
+
+
+class CheckPropertyIsUsed(ast.NodeVisitor):
+    """
+    Check whether a property is used.
+
+    Args:
+        node (ast.AST): An instance of ast node.
+    """
+    def __init__(self, node: ast.AST):
+        self._context = node
+        self._value = ""
+        self._attr = ""
+        self._hit = False
+
+    def visit_Attribute(self, node: ast.Attribute) -> Any:  # pylint: disable=invalid-name
+        """Visit a node of type ast.Attribute."""
+        if isinstance(node.value, ast.Name) and node.value.id == self._value and node.attr == self._attr:
+            self._hit = True
+        return super(CheckPropertyIsUsed, self).generic_visit(node)
+
+    def generic_visit(self, node: ast.AST) -> Any:
+        """
+        An override method, iterating over all nodes and save target ast nodes.
+        """
+        if self._hit:
+            return
+        super(CheckPropertyIsUsed, self).generic_visit(node)
+
+    def check(self, value, attr) -> bool:
+        """
+        Check whether `value` and `attr` exists.
+        """
+        self._value = value
+        self._attr = attr
+        self._hit = False
+        self.generic_visit(self._context)
+        return self._hit
+
+
+class GetPropertyOfObj(ast.NodeVisitor):
+    """
+    Check whether a property is used.
+
+    Args:
+        node (ast.AST): An instance of ast node.
+    """
+    def __init__(self, node: ast.AST):
+        self._context = node
+        self._property = set()
+
+    def visit_Assign(self, node: ast.Assign) -> Any:  # pylint: disable=invalid-name
+        """Visit a node of type ast.Attribute."""
+        target = node.targets[0]
+        if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) and target.value.id == "self":
+            self._property.add(target.attr)
+        return super(GetPropertyOfObj, self).generic_visit(node)
+
+    def get(self):
+        """
+        Check whether `value` and `attr` exists.
+        """
+        self._property = set()
+        self.generic_visit(self._context)
+        return self._property

mindspore/rewrite/ast_helpers/ast_modifier.py

@@ -241,8 +241,10 @@ class AstModifier(ast.NodeTransformer):
             An instance of ast.Assign which has been appended to 'init_func'.
         """
         return AstModifier.insert_assign_to_function(init_func, targets=targets,
-                                                     args=[ScopedValue.create_variable_value(field)],
-                                                     expr=ScopedValue(ValueType.NamingValue, "global_vars", "get"))
+                                                     expr=ScopedValue(ValueType.NamingValue, "", "setattr"),
+                                                     args=[ScopedValue(ValueType.NamingValue, "obj"),
+                                                           ScopedValue.create_variable_value(field)])
+
 
     @staticmethod
     def create_call_assign(targets: [ScopedValue], expr: ScopedValue, args: [ScopedValue],
@@ -459,7 +461,7 @@ class AstModifier(ast.NodeTransformer):
 
         Args:
             src_argument (ScopedValue): An instance of ScopedValue represents new argument.
-            dst_ast (ast.AST): Targets of ast.Assign.
+            dst_ast (ast.AST): Ast node to be updated by ScopedValue.
 
         Raises:
             TypeError: Input src_argument is not a ScopedValue
@@ -490,6 +492,12 @@ class AstModifier(ast.NodeTransformer):
                                    str(src_argument.type))
             dst_ast.n = src_argument.value
             return
+        if isinstance(dst_ast, ast.Str):
+            if src_argument.type not in [ValueType.StringValue]:
+                raise RuntimeError("src_argument should be a StringValue, but got:",
+                                   str(src_argument.type))
+            dst_ast.s = src_argument.value
+            return
         if isinstance(dst_ast, ast.Name):
             if src_argument.type not in [ValueType.NamingValue, ValueType.StringValue]:
                 raise RuntimeError("src_argument.type should be ValueType.NamingValue or ValueType.StringValue.")

mindspore/rewrite/ast_transformers/flatten_recursive_stmt.py

@@ -17,6 +17,7 @@
 from typing import Any, Tuple
 import ast
 from ast import FunctionDef
+import astunparse
 
 from mindspore import log as logger
 from ..common import error_str
@@ -37,7 +38,8 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
             ast.Call: ["args"],
             ast.BinOp: ["left", "right"],
             ast.BoolOp: ["values"],
-            ast.unaryop: ["operand"],
+            ast.UnaryOp: ["operand"],
+            ast.Compare: ["left", "comparators"],
         }
 
     @staticmethod
@@ -54,7 +56,7 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                 target_name = "function"
         elif isinstance(node, ast.Return):
             target_name = "return_value"
-        elif isinstance(node, (ast.BinOp, ast.boolop, ast.UnaryOp)):
+        elif isinstance(node, (ast.BinOp, ast.BoolOp, ast.UnaryOp)):
             target_name = type(node.op).__name__.lower() + "_var"
         elif isinstance(node, ast.Tuple):
             target_name = type(node).__name__.lower() + "_var"
@@ -180,6 +182,20 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
             child = node.body[index]
             if isinstance(child, ast.Assign):
                 stmt = child.value
+            elif isinstance(child, ast.If):
+                if isinstance(child.body[0], ast.Return) and not isinstance(child.test, ast.UnaryOp):
+                    if isinstance(child.body[0].value, ast.Call):
+                        if_body = child.body
+                        if_func = if_body[0].value
+                        expr = "x = " + astunparse.unparse(if_func)
+                        if_body = ast.parse(expr)
+                        if_body = if_body.body+ast.parse("return x").body
+                        child.body = if_body
+                        stmt = child
+                    else:
+                        stmt = child
+                else:
+                    stmt = child
             elif isinstance(child, ast.Expr):
                 stmt = child.value
             else:

mindspore/rewrite/namespace.py

@@ -24,8 +24,6 @@ _ms_functional_ns = CellNamespace('mindspore.ops.functional')
 
 def is_subtree(cls_name):
     """Determine whether 'cls_name' is a subtree."""
-    if cls_name == "SequentialCell":
-        return True
     if cls_name == "QuantizeWrapperCell":
         return False
     if cls_name in _ms_common_ns or cls_name in _ms_nn_ns or cls_name in _ms_ops_ns: