mindspore 1.10.0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/rewrite/__init__.py

@@ -16,7 +16,9 @@
 MindSpore Rewrite package.
 This is an experimental python package that is subject to change or deletion.
 """
-from __future__ import absolute_import
+from mindspore.rewrite.parsers.attribute_parser import g_attribute_parser
+from mindspore.rewrite.parsers.constant_parser import g_name_parser, g_num_parser, g_str_parser
+from mindspore.rewrite.parsers.container_parser import g_list_parser, g_tuple_parser
 from .parsers.module_parser import g_module_parser
 from .parsers.class_def_parser import g_classdef_parser
 from .parsers.function_def_parser import g_functiondef_parser
@@ -31,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

@@ -13,12 +13,12 @@
 # limitations under the License.
 # ============================================================================
 """Rewrite module api: Node."""
-from __future__ import absolute_import
+
 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
@@ -32,13 +32,11 @@ class Node:
     For the most part, Node represents an operator invoking in forward which could be an instance of `Cell`, an instance
     of `Primitive` or a callable method.
 
-    `NodeImpl` mentioned below is implementation of `Node` which is not an interface of Rewrite. Rewrite recommend
-    invoking specific create method of `Node` to instantiate an instance of Node such as `create_call_cell` rather than
-    invoking constructor of `Node` directly, so don't care about what is `NodeImpl` and use its instance just as a
-    handler.
-
     Args:
-        node (NodeImpl): A handler of `NodeImpl`.
+        node (NodeImpl): A handler of `NodeImpl`. `NodeImpl` mentioned below is implementation of `Node` which is not
+            an interface of Rewrite. Rewrite recommend invoking specific create method of `Node`
+            to instantiate an instance of Node such as `create_call_cell` rather than invoking constructor of `Node`
+            directly, so don't care about what is `NodeImpl` and use its instance just as a handler.
     """
 
     def __init__(self, node: NodeImpl):
@@ -46,7 +44,8 @@ class Node:
 
 
     def __eq__(self, other: 'Node'):
-        Validator.check_value_type("other", other, [Node], "Node")
+        if not isinstance(other, Node):
+            return False
         return self._node == other._node
 
     @staticmethod
@@ -100,32 +99,8 @@ 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_prev(self) -> 'Node':
-        """
-        Get previous node of current node in source code order.
-
-        Returns:
-            An instance of `Node` as previous node.
-        """
-        return Node(self._node.get_prev())
-
-    def get_next(self) -> 'Node':
-        """
-        Get next node of current node in source code order.
-
-        Returns:
-            An instance of `Node` as next node.
-        """
-        return Node(self._node.get_next())
-
     def get_inputs(self) -> ['Node']:
         """
         Get input nodes of current node in topological order.
@@ -141,6 +116,14 @@ class Node:
 
         Returns:
             A list of nodes represents users.
+
+        Examples:
+            >>> from mindspore.rewrite import SymbolTree
+            >>> from lenet import Lenet
+            >>> net = Lenet()
+            >>> stree = SymbolTree.create(net)
+            >>> node = stree.get_node("conv1")
+            >>> users = node.get_users()
         """
         belong_symbol_tree: SymbolTreeImpl = self._node.get_belong_symbol_tree()
         if belong_symbol_tree is None:
@@ -163,6 +146,14 @@ class Node:
         Raises:
             TypeError: If `index` is not a `int` number.
             TypeError: If the type of `arg` is not in [`ScopedValue`, `str`].
+
+        Examples:
+            >>> from mindspore.rewrite import SymbolTree
+            >>> from lenet import Lenet
+            >>> net = Lenet()
+            >>> stree = SymbolTree.create(net)
+            >>> node = stree.get_node("conv1")
+            >>> node.set_arg(0, "x")
         """
         Validator.check_value_type("index", index, [int], "Node")
         Validator.check_value_type("arg", arg, [ScopedValue, str], "Node")
@@ -174,7 +165,7 @@ class Node:
 
     def set_arg_by_node(self, arg_idx: int, src_node: 'Node', out_idx: Optional[int] = None):
         """
-        Set argument of current node by another `Node`.
+        Set argument of current node by another Node.
 
         Args:
             arg_idx (int): Indicate which input being modified.
@@ -184,13 +175,21 @@ 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.
             TypeError: If `out_idx` is not a `int` number.
             ValueError: If `out_idx` is out of range.
             ValueError: If `src_node` has multi-outputs while `out_idx` is None or `out_idx` is not offered.
+
+        Examples:
+            >>> from mindspore.rewrite import SymbolTree
+            >>> from lenet import Lenet
+            >>> net = Lenet()
+            >>> stree = SymbolTree.create(net)
+            >>> src_node = stree.get_node("conv1")
+            >>> dst_node = stree.get_node("conv2")
+            >>> dst_node.set_arg_by_node(0, src_node)
         """
         Validator.check_value_type("arg_idx", arg_idx, [int], "Node")
         Validator.check_value_type("src_node", src_node, [Node], "Node")
@@ -200,22 +199,9 @@ class Node:
         if belong_symbol_tree is None:
             self._node.set_arg_by_node(arg_idx, src_node._node, out_idx)
         else:
-            belong_symbol_tree.set_node_arg_by_node(self._node, arg_idx, src_node._node, out_idx)
+            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.
-        """
         return self._node.get_targets()
 
     def get_name(self) -> str:
@@ -226,6 +212,14 @@ class Node:
 
         Returns:
             A string as name of node.
+
+        Examples:
+            >>> from mindspore.rewrite import SymbolTree
+            >>> from lenet import Lenet
+            >>> net = Lenet()
+            >>> stree = SymbolTree.create(net)
+            >>> node = stree.get_node("conv1")
+            >>> name = node.get_name()
         """
         return self._node.get_name()
 
@@ -235,6 +229,14 @@ class Node:
 
         Returns:
             A NodeType as node_type of node.
+
+        Examples:
+            >>> from mindspore.rewrite import SymbolTree
+            >>> from lenet import Lenet
+            >>> net = Lenet()
+            >>> stree = SymbolTree.create(net)
+            >>> node = stree.get_node("conv1")
+            >>> node_type = node.get_node_type()
         """
         return self._node.get_node_type()
 
@@ -254,82 +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 return values.
-        - When node_type of current node is `Python`, arguments are don't-care.
-
-        Returns:
-            A list of instances of `ScopedValue`.
-        """
         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`.
-        """
         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`.
-        """
         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.
-
-        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

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ============================================================================
 """Rewrite module api: NodeType."""
-from __future__ import absolute_import
 from enum import Enum
 
 
@@ -30,15 +29,20 @@ 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
     # Compute node type
     CallCell = 1
     CallPrimitive = 2
-    CallMethod = 3
+    CallModule = 3
+    CallFunction = 4
+    CallMethod = 5
     # Other node type
-    Python = 4
-    Input = 5
-    Output = 6
-    Tree = 7
+    Python = 6
+    Input = 7
+    Output = 8
+    Tree = 9
+    CellContainer = 10
+    MathOps = 11

mindspore/rewrite/api/pattern_engine.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ============================================================================
 """PatternEngine for modifying SymbolTree by pattern."""
-from __future__ import absolute_import
 from collections import OrderedDict
 from typing import Tuple, Union, List, Type
 import abc
@@ -21,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
@@ -34,8 +33,8 @@ class PatternNode:
 
     Args:
         pattern_node_name (str): Name of current node.
-        match_type (Type): A type represents what type would be matched of current node.
-        inputs (list[PatternNode]): Input nodes of current node.
+        match_type (Type): A type represents what type would be matched of current node. Default: None.
+        inputs (list[PatternNode]): Input nodes of current node. Default: None.
     """
 
     def __init__(self, pattern_node_name: str, match_type: Type = Type[None], inputs: ['PatternNode'] = None):
@@ -158,19 +157,22 @@ class PatternNode:
     def get_inputs(self):
         """
         Getter of inputs.
+
+        Returns:
+            A list of `PatternNode`, the inputs of current node.
         """
 
         return self._inputs
 
     def name(self) -> str:
         """
-        Getter of name.
+        Getter of PatternNode name.
         """
         return self._name
 
     def type(self):
         """
-        Getter of type.
+        Getter of PatternNode type.
         """
         return self._type
 
@@ -190,8 +192,17 @@ class VarNode(PatternNode):
 class Replacement(abc.ABC):
     """
     Interface of replacement function.
-    """
 
+    Examples:
+        >>> from mindspore.rewrite import Replacement, Node
+        >>> from mindspore.nn import nn
+        >>> class BnReplacement(Replacement):
+        ...     def build(self, pattern, is_chain_pattern: bool, matched):
+        ...         bn_node: Node = matched.get(pattern.name())
+        ...         conv = nn.Conv2d(16, 16, 3)
+        ...         conv_node = Node.create_call_cell(conv, ['x1'], bn_node.get_args(), bn_node.get_kwargs())
+        ...         return [conv_node]
+    """
     @abc.abstractmethod
     def build(self, pattern: PatternNode, is_chain_pattern: bool, matched: OrderedDict) -> [Node]:
         """
@@ -250,6 +261,9 @@ class PatternEngine:
     def pattern(self) -> PatternNode:
         """
         Getter of pattern.
+
+        Returns:
+            A instance of `PatternNode`, used to indicate the type that the current pattern needs to match.
         """
 
         return self._pattern
@@ -294,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`.
@@ -345,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
@@ -446,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

@@ -13,10 +13,9 @@
 # limitations under the License.
 # ============================================================================
 """Rewrite module api: ValueType and ScopedValue."""
-from __future__ import absolute_import
 from enum import Enum
 from typing import Optional, Union
-from ..._checkparam import Validator
+from mindspore import _checkparam as Validator
 
 
 class ValueType(Enum):
@@ -62,57 +61,6 @@ class ScopedValue:
         self.scope = scope
         self.value = value
 
-    def __str__(self):
-        if self.type in (ValueType.IntValue, ValueType.FloatValue, ValueType.StringValue):
-            return str(self.value)
-        if self.type == ValueType.NamingValue:
-            return f"{self.scope}.{self.value}" if self.scope else str(self.value)
-        if self.type == ValueType.CustomObjValue:
-            return f"CustomObj: {str(self.value)}"
-        return f"Illegal ValueType: {str(self.type)}"
-
-    def __eq__(self, other):
-        if id(self) == id(other):
-            return True
-        return self.type == other.type and self.scope == other.scope and self.value == other.value
-
-    def __repr__(self):
-        return str(self)
-
-    def __hash__(self):
-        return hash((self.type, self.scope, self.value))
-
-    @staticmethod
-    def create_name_values(names: Union[list, tuple], scopes: Union[list, tuple] = None) -> ['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.
-
-        Returns:
-            An list of instance of `ScopedValue`.
-
-        Raises:
-            RuntimeError: If the length of names is not equal to the length of scopes when scopes are not None.
-            TypeError: If `names` is not `list` or `tuple` and name in `names` is not `str`.
-            TypeError: If `scopes` is not `list` or `tuple` and scope in `scopes` is not `str`.
-        """
-        Validator.check_element_type_of_iterable("names", names, [str], "ScopedValue")
-        if scopes is not None:
-            Validator.check_element_type_of_iterable("scopes", scopes, [str], "ScopedValue")
-            if len(names) != len(scopes):
-                raise RuntimeError("Length of names should be equal to length of scopes")
-        result = []
-        for index, name in enumerate(names):
-            if scopes is not None:
-                scope = scopes[index]
-            else:
-                scope = ""
-            result.append(ScopedValue.create_naming_value(name, scope))
-        return result
-
     @classmethod
     def create_variable_value(cls, value) -> Optional['ScopedValue']:
         """
@@ -125,6 +73,10 @@ class ScopedValue:
 
         Returns:
             An instance of `ScopedValue`.
+
+        Examples:
+            >>> from mindspore.rewrite import ScopedValue
+            >>> variable = ScopedValue.create_variable_value(2)
         """
         if isinstance(value, int):
             return cls(ValueType.IntValue, "", value)
@@ -160,7 +112,66 @@ class ScopedValue:
         Raises:
             TypeError: If `name` is not `str`.
             TypeError: If `scope` is not `str`.
+
+        Examples:
+            >>> from mindspore.rewrite import ScopedValue
+            >>> variable = ScopedValue.create_naming_value("conv", "self")
         """
         Validator.check_value_type("name", name, [str], "ScopedValue")
         Validator.check_value_type("scope", scope, [str], "ScopedValue")
         return cls(ValueType.NamingValue, scope, name)
+
+    @staticmethod
+    def create_name_values(names: Union[list, tuple], scopes: Union[list, tuple] = None) -> ['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.
+
+        Returns:
+            An list of instance of `ScopedValue`.
+
+        Raises:
+            RuntimeError: If the length of names is not equal to the length of scopes when scopes are not None.
+            TypeError: If `names` is not `list` or `tuple` and name in `names` is not `str`.
+            TypeError: If `scopes` is not `list` or `tuple` and scope in `scopes` is not `str`.
+
+        Examples:
+            >>> from mindspore.rewrite import ScopedValue
+            >>> 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:
+            Validator.check_element_type_of_iterable("scopes", scopes, [str], "ScopedValue")
+            if len(names) != len(scopes):
+                raise RuntimeError("Length of names should be equal to length of scopes")
+        result = []
+        for index, name in enumerate(names):
+            if scopes is not None:
+                scope = scopes[index]
+            else:
+                scope = ""
+            result.append(ScopedValue.create_naming_value(name, scope))
+        return result
+
+    def __str__(self):
+        if self.type in (ValueType.IntValue, ValueType.FloatValue, ValueType.StringValue):
+            return str(self.value)
+        if self.type == ValueType.NamingValue:
+            return f"{self.scope}.{self.value}" if self.scope else str(self.value)
+        if self.type == ValueType.CustomObjValue:
+            return f"CustomObj: {str(self.value)}"
+        return f"Illegal ValueType: {str(self.type)}"
+
+    def __eq__(self, other):
+        if id(self) == id(other):
+            return True
+        return self.type == other.type and self.scope == other.scope and self.value == other.value
+
+    def __repr__(self):
+        return str(self)
+
+    def __hash__(self):
+        return hash((self.type, self.scope, self.value))