mindspore 2.0.0a0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/ops/composite/base.py

@@ -27,15 +27,15 @@ from mindspore.parallel._utils import _grads_divided_by_device_num_if_recomputat
 from mindspore._c_expression import GradOperation_, HyperMap_, Map_, MultitypeFuncGraph_, Tail_, \
     TupleAdd_, UnpackCall_, ZipOperation_, ListAppend_, TupleGetItemTensor_, ListInsert_, \
     SequenceSliceGetItem_, ListSliceSetItem_, VmapOperation_, TaylorOperation_, ListPop_, \
-    ListClear_, ListReverse_, ListExtend_, ListCount_, DictClear_, DictHasKey_, DictUpdate_, \
-    DictFromKeys_
+    ListClear_, ListReverse_, ListExtend_, DictClear_, DictHasKey_, DictUpdate_, DictFromKeys_, \
+    ZerosLike_
 from mindspore.common import dtype as mstype
 from mindspore.common.api import jit, _pynative_executor, _wrap_func
 from mindspore.common.api import _add_flags, _core
 from mindspore.ops.primitive import Primitive
 from mindspore.ops import signature as sig
 
-__all__ = [TupleAdd_, UnpackCall_, TupleGetItemTensor_, SequenceSliceGetItem_, ListSliceSetItem_]
+__all__ = [TupleAdd_, UnpackCall_, TupleGetItemTensor_, SequenceSliceGetItem_, ListSliceSetItem_, ZerosLike_]
 
 
 def add_flags(fn=None, **flags):
@@ -329,7 +329,7 @@ class GradOperation(GradOperation_):
         self.get_all = get_all
         self.get_by_list = get_by_list
         self.sens_param = sens_param
-        GradOperation_.__init__(self, 'grad', get_all, get_by_list, sens_param, False, False, False)
+        GradOperation_.__init__(self, 'grad', get_all, get_by_list, sens_param, False, False, False, False)
         self.grad_fn = None
         self.fn = None
         self.weights_id = None
@@ -353,16 +353,16 @@ class GradOperation(GradOperation_):
                 fn.grad_ops_label = True
             if self.get_by_list:
                 @jit(input_signature=dynamic_shape_inputs)
-                def after_grad(*args):
-                    return grad_(fn, weights)(*args)
+                def after_grad(*args, **kwargs):
+                    return grad_(fn, weights)(*args, **kwargs)
             else:
                 @jit(input_signature=dynamic_shape_inputs)
-                def after_grad(*args):
-                    return grad_(fn)(*args)
+                def after_grad(*args, **kwargs):
+                    return grad_(fn)(*args, **kwargs)
         elif self.pynative_:
             @_wrap_func
             def after_grad(*args, **kwargs):
-                self._pynative_forward_run(fn, grad_, args, kwargs)
+                self._pynative_forward_run(fn, grad_, weights, args, kwargs)
                 _pynative_executor.grad(fn, grad_, weights, self.grad_position, *args, **kwargs)
                 out = _pynative_executor()
                 out = _grads_divided_by_device_num_if_recomputation(out)
@@ -382,7 +382,7 @@ class GradOperation(GradOperation_):
         self.weights_id = weights_id
         return self.grad_fn
 
-    def _pynative_forward_run(self, fn, grad, args, kwargs):
+    def _pynative_forward_run(self, fn, grad, weights, args, kwargs):
         """ Pynative forward run to build grad graph. """
         new_kwargs = kwargs
         if self.sens_param:
@@ -392,14 +392,14 @@ class GradOperation(GradOperation_):
                 new_kwargs = kwargs.copy()
                 new_kwargs.pop('sens')
         if isinstance(fn, (FunctionType, MethodType)):
-            if not _pynative_executor.check_run(grad, fn, self.weights_id, *args, **new_kwargs):
+            if not _pynative_executor.check_run(grad, fn, weights, None, *args, **new_kwargs):
                 _pynative_executor.set_grad_flag(True)
                 _pynative_executor.new_graph(fn, *args, **new_kwargs)
                 output = fn(*args, **new_kwargs)
                 _pynative_executor.end_graph(fn, output, *args, **new_kwargs)
         else:
             # Check if fn have run already
-            if not _pynative_executor.check_run(grad, fn, self.weights_id, *args, **new_kwargs):
+            if not _pynative_executor.check_run(grad, fn, weights, None, *args, **new_kwargs):
                 fn.set_grad()
                 fn(*args, **new_kwargs)
                 fn.set_grad(False)
@@ -432,12 +432,77 @@ class _TaylorOperation(TaylorOperation_):
         return self.grad_fn
 
 
+def _combine_weight(grad_position, weights, out, out_with_ids):
+    """ Making resulting tuple for weight, when return_ids is set to True. """
+    weight_tuple = []
+    position = 0
+    if isinstance(weights, (list, ParameterTuple, tuple)) and grad_position:
+        for weight in weights:
+            weight_tuple.append((weight.name, out[1][position]))
+            position += 1
+    elif isinstance(weights, (list, ParameterTuple, tuple)):
+        for weight in weights:
+            weight_tuple.append((weight.name, out[position]))
+            position += 1
+    elif grad_position:
+        weight_tuple.append(weights.name)
+        weight_tuple.append(out[1])
+    else:
+        weight_tuple.append(weights.name)
+        weight_tuple.append(out)
+    if grad_position:
+        out_with_ids.append(tuple(weight_tuple))
+    else:
+        out_with_ids = weight_tuple
+    return out_with_ids
+
+
+def _combine_position(grad_position, weights, out, out_with_ids):
+    """ Making resulting tuple for position, when return_ids is set to True. """
+    position_tuple = []
+    position = 0
+    if grad_position == (0,) and weights is not None:
+        position_tuple.append(0)
+        position_tuple.append(out[0])
+    elif grad_position == (0,):
+        position_tuple.append(0)
+        position_tuple.append(out)
+    elif weights is not None:
+        for index in grad_position:
+            position_tuple.append((index, out[0][position]))
+            position += 1
+    else:
+        for index in grad_position:
+            position_tuple.append((index, out[position]))
+            position += 1
+    if weights:
+        out_with_ids.append(tuple(position_tuple))
+    else:
+        out_with_ids = position_tuple
+    return out_with_ids
+
+
+def _combine_with_ids(grad_position, weights, out):
+    """ Making resulting tuple, when return_ids is set to True. """
+    out_with_ids = []
+    if grad_position:
+        out_with_ids = _combine_position(
+            grad_position, weights, out, out_with_ids)
+    if weights is not None:
+        out_with_ids = _combine_weight(
+            grad_position, weights, out, out_with_ids)
+    if not out_with_ids:
+        raise ValueError(f"output tuple should not be a empty tuple.")
+    return tuple(out_with_ids)
+
+
 class _Grad(GradOperation_):
     """
     A higher-order function which is used to generate the gradient function by position for the input function.
     """
 
-    def __init__(self, get_by_list=False, sens_param=False, get_by_position=False, has_aux=False, get_value=False):
+    def __init__(self, get_by_list=False, sens_param=False, get_by_position=False, has_aux=False, get_value=False,
+                 return_ids=False):
         """Initialize _Grad."""
         if not isinstance(get_by_position, bool):
             raise TypeError(f"For '_Grad', the 'get_by_position' should be bool, "
@@ -454,18 +519,22 @@ class _Grad(GradOperation_):
         if not isinstance(get_value, bool):
             raise TypeError(f"For '_Grad', the 'get_value' should be bool, "
                             f"but got {type(get_value).__name__}")
+        if not isinstance(return_ids, bool):
+            raise TypeError(f"For '_Grad', the 'return_ids' should be bool, "
+                            f"but got {type(return_ids).__name__}")
         self.get_by_position = get_by_position
         self.get_by_list = get_by_list
         self.sens_param = sens_param
         self.has_aux = has_aux
         self.get_value = get_value
-        GradOperation_.__init__(self, 'grad', False, get_by_list, sens_param, get_by_position, has_aux, get_value)
+        self.return_ids = return_ids
+        GradOperation_.__init__(self, 'grad', False, get_by_list, sens_param, get_by_position, has_aux, get_value,
+                                return_ids)
         self.grad_fn = None
         self.fn = None
         self.pynative_ = False
         self.grad_position = None
         self.weights_id = None
-        self.grad_hash_id = None
 
     def __call__(self, fn, weights=None, grad_position=0):
         weights_id = _get_grad_weights_id(weights)
@@ -483,7 +552,8 @@ class _Grad(GradOperation_):
                 res += (stop_gradient(item),)
             return res
 
-        grad_ = _Grad(self.get_by_list, self.sens_param, self.get_by_position, self.has_aux, self.get_value)
+        grad_ = _Grad(self.get_by_list, self.sens_param, self.get_by_position, self.has_aux, self.get_value,
+                      self.return_ids)
         # If calling Grad in GRAPH_MODE or calling Grad in functions decorated with 'jit', do grad in GRAPH_MODE
         # If calling Grad in pure PYNATIVE_MODE do grad in PYNATIVE_MODE
         #   In pure PYNATIVE_MODE the out layer after_grad just used to set pynative flag for inner GradOperation.
@@ -509,10 +579,12 @@ class _Grad(GradOperation_):
         elif self.pynative_:
             @_wrap_func
             def after_grad(*args, **kwargs):
-                res = self._pynative_forward_run(fn, grad_, args, kwargs)
+                res = self._pynative_forward_run(fn, grad_, weights, args, kwargs)
                 _pynative_executor.grad(fn, grad_, weights, grad_position, *args, **kwargs)
                 out = _pynative_executor()
                 out = _grads_divided_by_device_num_if_recomputation(out)
+                if self.return_ids and out:
+                    out = _combine_with_ids(grad_position, weights, out)
                 if self.get_value:
                     return res, out
                 if self.has_aux:
@@ -539,10 +611,9 @@ class _Grad(GradOperation_):
         self.fn = fn
         self.grad_position = grad_position
         self.weights_id = weights_id
-        self.grad_hash_id = (grad_position, weights_id)
         return self.grad_fn
 
-    def _pynative_forward_run(self, fn, grad, args, kwargs):
+    def _pynative_forward_run(self, fn, grad, weights, args, kwargs):
         """ Pynative forward runs to build grad graph. """
         new_kwargs = kwargs
         outputs = ()
@@ -553,7 +624,7 @@ class _Grad(GradOperation_):
             else:
                 args = args[:-1]
         if isinstance(fn, (FunctionType, MethodType)):
-            if not _pynative_executor.check_run(grad, fn, self.grad_hash_id, *args, **new_kwargs):
+            if not _pynative_executor.check_run(grad, fn, weights, self.grad_position, *args, **new_kwargs):
                 _pynative_executor.set_grad_flag(True)
                 _pynative_executor.new_graph(fn, *args, **new_kwargs)
                 outputs = fn(*args, **new_kwargs)
@@ -561,7 +632,7 @@ class _Grad(GradOperation_):
                 return outputs
         else:
             # Check if fn has run already.
-            if not _pynative_executor.check_run(grad, fn, self.grad_hash_id, *args, **new_kwargs):
+            if not _pynative_executor.check_run(grad, fn, weights, self.grad_position, *args, **new_kwargs):
                 fn.set_grad()
                 outputs = fn(*args, **new_kwargs)
                 fn.set_grad(False)
@@ -591,8 +662,8 @@ class _Vmap(VmapOperation_):
         vmap_ = self
 
         @jit
-        def after_vmap(*args):
-            return vmap_(fn, in_axes, out_axes)(*args)
+        def after_vmap(*args, **kwargs):
+            return vmap_(fn, in_axes, out_axes)(*args, **kwargs)
 
         self.vmap_fn = after_vmap
         self.fn = fn
@@ -612,6 +683,7 @@ class MultitypeFuncGraph(MultitypeFuncGraph_):
         name (str): Operator name.
         read_value (bool, optional): If the registered function do not need to set value on Parameter,
             and all inputs will pass by value, set `read_value` to True. Default: False.
+        doc_url (str, optional): The official document link corresponding to the registered function. Default:"".
 
     Raises:
         ValueError: If failed to find a matching function for the given arguments.
@@ -643,9 +715,9 @@ class MultitypeFuncGraph(MultitypeFuncGraph_):
         [0.2 1.2 2.4]
     """
 
-    def __init__(self, name, read_value=False):
+    def __init__(self, name, read_value=False, doc_url=""):
         """Initialize MultitypeFuncGraph."""
-        MultitypeFuncGraph_.__init__(self, name)
+        MultitypeFuncGraph_.__init__(self, name, doc_url)
         self.entries = list()
         if read_value:
             self.set_signatures((
@@ -710,10 +782,11 @@ class HyperMap(HyperMap_):
           Only supported in graph mode. Default is False.
 
     Inputs:
-        - **args** (Tuple[sequence]) - If `ops` is not `None`, all the inputs should be sequences with the same length.
-          And each row of the sequences will be the inputs of the operation.
+        - **args** (Tuple[sequence]) -
 
-          If `ops` is `None`, the first input is the operation, and the others are inputs.
+          - If `ops` is not `None`, all the inputs should be sequences with the same length.
+            And each row of the sequences will be the inputs of the operation.
+          - If `ops` is `None`, the first input is the operation, and the others are inputs.
 
     Note:
         Except for the operation input, the number of inputs should be equal to the number of inputs to `ops`.
@@ -956,25 +1029,6 @@ class _ListExtend(ListExtend_):
 _extend = _ListExtend("extend")
 
 
-class _ListCount(ListCount_):
-    """
-    A metafuncgraph class that count the number of times an element appears in list.
-
-    Args:
-        name (str): The name of the metafuncgraph object.
-    """
-
-    def __init__(self, name):
-        """Initialize _ListCount."""
-        ListCount_.__init__(self, name)
-
-    def __call__(self, *args):
-        pass
-
-
-_count = _ListCount("count")
-
-
 class _DictClear(DictClear_):
     """
     A metafuncgraph class that clear the dict.