mindspore 1.10.0__cp38-cp38-win_amd64.whl → 2.0.0rc1__cp38-cp38-win_amd64.whl

mindspore/ops/composite/base.py

@@ -18,24 +18,24 @@
 """Basic composite operations."""
 from __future__ import absolute_import
 from functools import partial
+
 from types import FunctionType, MethodType
 import mindspore as ms
-import mindspore.nn as nn
 from mindspore import context
 from mindspore.common.parameter import Parameter, ParameterTuple
-from mindspore import log as logger
-from mindspore._c_expression import GradOperation_, HyperMap_, Map_, MultitypeFuncGraph_, Tail_, Shard_, \
+from mindspore.parallel._utils import _grads_divided_by_device_num_if_recomputation
+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_
+    ListClear_, ListReverse_, ListExtend_, DictClear_, DictHasKey_, DictUpdate_, DictFromKeys_, \
+    ZerosLike_
 from mindspore.common import dtype as mstype
-from mindspore.common.api import ms_function, _pynative_executor, _wrap_func
+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.operations import _grad_ops
-from mindspore.ops import operations as P
 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):
@@ -59,18 +59,7 @@ def add_flags(fn=None, **flags):
         True
     """
 
-    def deco(fn):
-        # need set the attr and access on c++
-        if not hasattr(fn, "_func_graph_flags"):
-            fn._func_graph_flags = {}
-
-        fn._func_graph_flags.update({**flags})
-        return fn
-
-    ret = deco
-    if fn is not None:
-        ret = deco(fn)
-    return ret
+    return _add_flags(fn, **flags)
 
 
 def core(fn=None, **flags):
@@ -81,8 +70,8 @@ def core(fn=None, **flags):
     set flag to a graph.
 
     Args:
-        fn (Function): Function to add flag. Default: None.
-        flags (dict): The following flags can be set core, which indicates that this is a core function or
+        fn (Function, optional): Function to add flag. Default: None.
+        flags (dict, optional): The following flags can be set core, which indicates that this is a core function or
                       other flag. Default: None.
 
     Supported Platforms:
@@ -95,31 +84,18 @@ def core(fn=None, **flags):
         True
     """
 
-    # need set the attr and access on c++
-
-    def deco(fn):
-        fn._func_graph_flags = {
-            'core': True,
-            **flags,
-        }
-        return fn
-
-    if fn is not None:
-        ret = deco(fn)
-    else:
-        ret = deco
-    return ret
+    return _core(fn, **flags)
 
 
 def _get_grad_weights_id(weights=None):
     """generate id of parameters"""
     res = ""
     if isinstance(weights, Parameter):
-        res = weights.name
+        res = weights.name + str(weights.requires_grad)
     if isinstance(weights, ParameterTuple):
-        res = ''.join(item.name for item in weights)
+        res = ''.join(item.name + str(item.requires_grad) for item in weights)
     if isinstance(weights, list):
-        res = ''.join(item.name for item in weights if isinstance(item, Parameter))
+        res = ''.join(item.name + str(item.requires_grad) for item in weights if isinstance(item, Parameter))
     return res
 
 
@@ -130,83 +106,85 @@ class GradOperation(GradOperation_):
     The gradient function generated by `GradOperation` higher-order function can be customized by
     construction arguments.
 
-    Given an input function `net = Net()` that takes `x` and `y` as inputs, and has a parameter `z`,
+    For example, given an input function `net = Net()` that takes `x` and `y` as inputs, and has a parameter `z`,
     see `Net` in Examples.
 
+    - Used to get the derivative of the input:
 
-    To generate a gradient function that returns gradients with respect to the first input
-    (see `GradNetWrtX` in Examples).
+      1. Returns gradients with respect to the first input (see `GradNetWrtX` in Examples).
 
-    1. Construct a `GradOperation` higher-order function with default arguments:
-       `grad_op = GradOperation()`.
+         1) Construct a `GradOperation` higher-order function with default arguments: `grad_op = GradOperation()`.
 
-    2. Call it with input function as argument to get the gradient function: `gradient_function = grad_op(net)`.
+         2) Call it with input function as argument to get the gradient function: `gradient_function = grad_op(net)`.
 
-    3. Call the gradient function with input function's inputs to get the gradients with respect to the first input:
-       `grad_op(net)(x, y)`.
+         3) Call the gradient function with input function's inputs to get the gradients with respect to the first
+            input: `grad_op(net)(x, y)`.
 
+      2. Returns gradients with respect to all inputs (see `GradNetWrtXY` in Examples).
 
-    To generate a gradient function that returns gradients with respect to all inputs (see `GradNetWrtXY` in Examples).
+         1) Construct a `GradOperation` higher-order function with `get_all=True` which indicates getting gradients
+            with respect to all inputs, they are `x` and `y` in example function `Net()`:
+            `grad_op = GradOperation(get_all=True)`.
 
-    1. Construct a `GradOperation` higher-order function with `get_all=True` which
-       indicates getting gradients with respect to all inputs, they are `x` and `y` in example function `Net()`:
-       `grad_op = GradOperation(get_all=True)`.
+         2) Call it with input function as argument to get the gradient function: `gradient_function = grad_op(net)`.
 
-    2. Call it with input function as argument to get the gradient function: `gradient_function = grad_op(net)`.
+         3) Call the gradient function with input function's inputs to get the gradients with respect to all inputs:
+            `gradient_function(x, y)`.
 
-    3. Call the gradient function with input function's inputs to get the gradients with respect to all inputs:
-       `gradient_function(x, y)`.
+    - Used to get the derivative of the parameters:
 
-    To generate a gradient function that returns gradients with respect to given parameters
-    (see `GradNetWithWrtParams` in Examples).
+      Returns gradients with respect to given parameters (see `GradNetWithWrtParams` in Examples).
 
-    1. Construct a `GradOperation` higher-order function with `get_by_list=True`:
-       `grad_op = GradOperation(get_by_list=True)`.
+      1. Construct a `GradOperation` higher-order function with `get_by_list=True`:
+         `grad_op = GradOperation(get_by_list=True)`.
 
-    2. Construct a `ParameterTuple` that will be passed to the input function when constructing
-       `GradOperation` higher-order function, it will be used as a parameter filter that determine
-       which gradient to return: `params = ParameterTuple(net.trainable_params())`.
+      2. Construct a `ParameterTuple` that will be passed to the input function when constructing
+         `GradOperation` higher-order function, it will be used as a parameter filter that determine
+         which gradient to return: `params = ParameterTuple(net.trainable_params())`.
 
-    3. Call it with input function and `params` as arguments to get the gradient function:
-       `gradient_function = grad_op(net, params)`.
+      3. Call it with input function and `params` as arguments to get the gradient function:
+         `gradient_function = grad_op(net, params)`.
 
-    4. Call the gradient function with input function's inputs to get the gradients with
-       respect to given parameters: `gradient_function(x, y)`.
+      4. Call the gradient function with input function's inputs to get the gradients with
+         respect to given parameters: `gradient_function(x, y)`.
 
-    To generate a gradient function that returns gradients with respect to all inputs and given parameters
-    in the format of ((dx, dy), (dz))(see `GradNetWrtInputsAndParams` in Examples).
+    - Used to get the derivative of the inputs and parameters at the same time:
+      Returns gradients with respect to all inputs and given parameters in the format of ((dx, dy), (dz))
+      (see `GradNetWrtInputsAndParams` in Examples).
 
-    1. Construct a `GradOperation` higher-order function with `get_all=True` and `get_by_list=True`:
-       `grad_op = GradOperation(get_all=True, get_by_list=True)`.
+      1. Construct a `GradOperation` higher-order function with `get_all=True` and `get_by_list=True`:
+         `grad_op = GradOperation(get_all=True, get_by_list=True)`.
 
-    2. Construct a `ParameterTuple` that will be passed along input function when constructing
-       `GradOperation` higher-order function: `params = ParameterTuple(net.trainable_params())`.
+      2. Construct a `ParameterTuple` that will be passed along input function when constructing
+         `GradOperation` higher-order function: `params = ParameterTuple(net.trainable_params())`.
 
-    3. Call it with input function and `params` as arguments to get the gradient function:
-       `gradient_function = grad_op(net, params)`.
+      3. Call it with input function and `params` as arguments to get the gradient function:
+         `gradient_function = grad_op(net, params)`.
 
-    4. Call the gradient function with input function's inputs
-       to get the gradients with respect to all inputs and given parameters: `gradient_function(x, y)`.
+      4. Call the gradient function with input function's inputs to get the gradients with respect to
+         all inputs and given parameters: `gradient_function(x, y)`.
 
-    To be noticed, for above gradient functions, the returned gradient result may vary for grad result element number:
-    A single value if only one result, a tuple for multiple results, or an empty tuple for no result.
+    - We can configure the sensitivity(gradient with respect to output) by setting `sens_param` as True and
+      passing an extra sensitivity input to the gradient function, the sensitivity input should has the
+      same shape and type with input function's output(see `GradNetWrtXYWithSensParam` in Examples).
 
-    We can configure the sensitivity(gradient with respect to output) by setting `sens_param` as True and
-    passing an extra sensitivity input to the gradient function, the sensitivity input should has the
-    same shape and type with input function's output(see `GradNetWrtXYWithSensParam` in Examples).
+      1. Construct a `GradOperation` higher-order function with `get_all=True` and `sens_param=True`:
+         `grad_op = GradOperation(get_all=True, sens_param=True)`.
 
-    1. Construct a `GradOperation` higher-order function with `get_all=True` and `sens_param=True`:
-       `grad_op = GradOperation(get_all=True, sens_param=True)`.
+      2. Define `grad_wrt_output` as `sens_param` which works as the gradient with respect to output:
+         `grad_wrt_output = Tensor(np.ones([2, 2]).astype(np.float32))`.
 
-    2. Define `grad_wrt_output` as `sens_param` which works as the gradient with respect to output:
-       `grad_wrt_output = Tensor(np.ones([2, 2]).astype(np.float32))`.
+      3. Call it with input function as argument to get the gradient function: `gradient_function = grad_op(net)`.
 
-    3. Call it with input function as argument to get the gradient function:
-       `gradient_function = grad_op(net)`.
+      4. Call the gradient function with input function's inputs and `sens_param` to
+         get the gradients with respect to all inputs: `gradient_function(x, y, grad_wrt_output)`.
 
-    4. Call the gradient function with input function's inputs and `sens_param` to
-       get the gradients with respect to all inputs:
-       `gradient_function(x, y, grad_wrt_output)`.
+    Note:
+        For above gradient functions, the returned gradient result may vary for grad result element number:
+
+        - Return a single value if only one result.
+        - Return a tuple for multiple results.
+        - Return an empty tuple for no result.
 
     Args:
         get_all (bool): If True, get all the gradients with respect to inputs. Default: False.
@@ -351,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
@@ -363,34 +341,35 @@ class GradOperation(GradOperation_):
         if self.grad_fn is not None and self.fn == fn and self.weights_id == weights_id:
             return self.grad_fn
         grad_ = GradOperation(self.get_all, self.get_by_list, self.sens_param)
-        # If calling Grad in GRAPH_MODE or calling Grad in ms_function, do grad in GRAPH_MODE
+        # 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.
-        #   In PYNATIVE_MODE calling Grad from ms_function, use the out layer after_grad do grad in GRAPH_MODE.
+        #   In PYNATIVE_MODE calling Grad from functions decorated with 'jit', use the out layer after_grad do
+        #   grad in GRAPH_MODE.
         if context.get_context("mode") == context.GRAPH_MODE:
             dynamic_shape_inputs = None
             if isinstance(fn, ms.nn.Cell):
                 dynamic_shape_inputs = fn.get_inputs()
                 fn.grad_ops_label = True
             if self.get_by_list:
-                @ms_function(input_signature=dynamic_shape_inputs)
-                def after_grad(*args):
-                    return grad_(fn, weights)(*args)
+                @jit(input_signature=dynamic_shape_inputs)
+                def after_grad(*args, **kwargs):
+                    return grad_(fn, weights)(*args, **kwargs)
             else:
-                @ms_function(input_signature=dynamic_shape_inputs)
-                def after_grad(*args):
-                    return grad_(fn)(*args)
+                @jit(input_signature=dynamic_shape_inputs)
+                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(fn, grad_.sens_param, *args, **kwargs)
-                _pynative_executor.clear_grad(fn, *args, **kwargs)
+                out = _pynative_executor()
+                out = _grads_divided_by_device_num_if_recomputation(out)
                 return out
         else:
             grad_.pynative_ = True
-            # after_grad of this branch can't use @ms_function, just directly call grad_
+            # after_grad of this branch can't use @jit, just directly call grad_
             if self.get_by_list:
                 def after_grad(*args, **kwargs):
                     return grad_(fn, weights)(*args, **kwargs)
@@ -403,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:
@@ -413,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)
@@ -442,9 +421,9 @@ class _TaylorOperation(TaylorOperation_):
             return self.grad_fn
         taylor_grad_ = _TaylorOperation()
 
-        # If calling Grad in GRAPH_MODE or calling Grad in ms_function, do grad in GRAPH_MODE
+        # If calling Grad in GRAPH_MODE or calling Grad in functions decorated with 'jit', do grad in GRAPH_MODE
 
-        @ms_function
+        @jit
         def after_taylor_grad(*args):
             return taylor_grad_(fn)(*args)
 
@@ -453,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, "
@@ -475,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)
@@ -499,41 +547,44 @@ class _Grad(GradOperation_):
             if not isinstance(outputs, tuple) or len(outputs) < 2:
                 raise ValueError("When has_aux is True, origin fn requires more than one outputs.")
             res = (outputs[0],)
-            stop_gradient = Primitive("stop_gradient")
+            stop_gradient = Primitive("StopGradient")
             for item in outputs[1:]:
                 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)
-        # If calling Grad in GRAPH_MODE or calling Grad in ms_function, do grad in GRAPH_MODE
+        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.
-        #   In PYNATIVE_MODE calling Grad from ms_function, use the out layer after_grad do grad in GRAPH_MODE.
+        #   In PYNATIVE_MODE calling Grad from functions decorated with 'jit', use the out layer after_grad do
+        #   grad in GRAPH_MODE.
         if context.get_context("mode") == context.GRAPH_MODE:
             dynamic_shape_inputs = None
             if isinstance(fn, ms.nn.Cell):
                 dynamic_shape_inputs = fn.get_inputs()
             if self.get_by_position:
-                @ms_function(input_signature=dynamic_shape_inputs)
+                @jit(input_signature=dynamic_shape_inputs)
                 def after_grad(*args):
                     return grad_(fn, weights, grad_position)(*args)
             else:
                 if self.get_by_list:
-                    @ms_function(input_signature=dynamic_shape_inputs)
+                    @jit(input_signature=dynamic_shape_inputs)
                     def after_grad(*args):
                         return grad_(fn, weights)(*args)
                 else:
-                    @ms_function(input_signature=dynamic_shape_inputs)
+                    @jit(input_signature=dynamic_shape_inputs)
                     def after_grad(*args):
                         return grad_(fn)(*args)
         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(fn, grad_.sens_param, *args, **kwargs)
-                _pynative_executor.clear_grad(fn, *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:
@@ -544,7 +595,7 @@ class _Grad(GradOperation_):
             fn_ = fn
             if self.has_aux:
                 fn_ = aux_fn
-            # after_grad of this branch can't use @ms_function, just directly call grad_
+            # after_grad of this branch can't use @jit, just directly call grad_
             if self.get_by_position:
                 def after_grad(*args, **kwargs):
                     return grad_(fn_, weights, grad_position)(*args, **kwargs)
@@ -560,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 = ()
@@ -574,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)
@@ -582,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)
@@ -602,23 +652,28 @@ class _Vmap(VmapOperation_):
         VmapOperation_.__init__(self, 'vmap')
         self.vmap_fn = None
         self.fn = None
+        self.in_axes = None
+        self.out_axes = None
 
     def __call__(self, fn, in_axes=0, out_axes=0):
+        if self.vmap_fn is not None and self.fn == fn and self.in_axes == in_axes and self.out_axes == out_axes:
+            return self.vmap_fn
+
         vmap_ = self
 
-        @ms_function
-        def after_vmap(*args):
-            return vmap_(fn, in_axes, out_axes)(*args)
+        @jit
+        def after_vmap(*args, **kwargs):
+            return vmap_(fn, in_axes, out_axes)(*args, **kwargs)
 
         self.vmap_fn = after_vmap
         self.fn = fn
+        self.in_axes = in_axes
+        self.out_axes = out_axes
         return self.vmap_fn
 
 
 class MultitypeFuncGraph(MultitypeFuncGraph_):
     """
-    Generates overloaded functions.
-
     MultitypeFuncGraph is a class used to generate overloaded functions, considering different types as inputs.
     Initialize an `MultitypeFuncGraph` object with name, and use `register` with input types as the decorator
     for the function to be registered. And the object can be called with different types of inputs,
@@ -626,8 +681,9 @@ class MultitypeFuncGraph(MultitypeFuncGraph_):
 
     Args:
         name (str): Operator name.
-        read_value (bool): If the registered function not need to set value on Parameter,
+        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.
@@ -641,10 +697,10 @@ class MultitypeFuncGraph(MultitypeFuncGraph_):
         >>> from mindspore import Tensor
         >>> from mindspore import ops
         >>> from mindspore import dtype as mstype
-        >>> from mindspore.ops.composite import MultitypeFuncGraph
+        >>> import mindspore.ops as ops
         >>>
         >>> tensor_add = ops.Add()
-        >>> add = MultitypeFuncGraph('add')
+        >>> add = ops.MultitypeFuncGraph('add')
         >>> @add.register("Number", "Number")
         ... def add_scala(x, y):
         ...     return x + y
@@ -659,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((
@@ -726,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`.
@@ -747,23 +804,22 @@ class HyperMap(HyperMap_):
 
     Examples:
         >>> from mindspore import Tensor, ops
-        >>> from mindspore.ops.composite.base import MultitypeFuncGraph, HyperMap
         >>> from mindspore import dtype as mstype
         >>> nest_tensor_list = ((Tensor(1, mstype.float32), Tensor(2, mstype.float32)),
         ...                     (Tensor(3, mstype.float32), Tensor(4, mstype.float32)))
         >>> # square all the tensor in the nested list
         >>>
-        >>> square = MultitypeFuncGraph('square')
+        >>> square = ops.MultitypeFuncGraph('square')
         >>> @square.register("Tensor")
         ... def square_tensor(x):
         ...     return ops.square(x)
         >>>
-        >>> common_map = HyperMap()
+        >>> common_map = ops.HyperMap()
         >>> output = common_map(square, nest_tensor_list)
         >>> print(output)
         ((Tensor(shape=[], dtype=Float32, value= 1), Tensor(shape=[], dtype=Float32, value= 4)),
         (Tensor(shape=[], dtype=Float32, value= 9), Tensor(shape=[], dtype=Float32, value= 16)))
-        >>> square_map = HyperMap(square, False)
+        >>> square_map = ops.HyperMap(square, False)
         >>> output = square_map(nest_tensor_list)
         >>> print(output)
         ((Tensor(shape=[], dtype=Float32, value= 1), Tensor(shape=[], dtype=Float32, value= 4)),
@@ -859,100 +915,6 @@ class Map(Map_):
         return tuple(map(func, *args_list))
 
 
-class Shard(Shard_):
-    """Shard operation"""
-
-    def __init__(self):
-        """Initialize Shard."""
-        Shard_.__init__(self, 'Shard')
-        self.shard_fn = None
-        self.fn = None
-        self.in_strategy = None
-        self.out_strategy = None
-        self.parameter_plan = None
-        self.device = None
-        self.level = None
-
-    def __call__(self, fn, in_strategy, out_strategy, parameter_plan=None, device="Ascend", level=0):
-        if context.get_context("mode") != context.PYNATIVE_MODE or \
-                context.get_auto_parallel_context("parallel_mode") not in ["auto_parallel"]:
-            raise AssertionError(f"'Shard' only supports auto parallel under PyNative mode")
-        if context.get_context("device_target") not in ["Ascend"]:
-            raise AssertionError(f"'Shard' now only supports 'Ascend'")
-        if context.get_auto_parallel_context("full_batch"):
-            raise AssertionError(f"'Shard' doesn't support 'full_batch'. Please set 'full_batch' as False")
-        if context.get_auto_parallel_context("search_mode") != "sharding_propagation":
-            raise AssertionError(f"'search_mode' must be 'sharding_propagation' for 'Shard'")
-        if not isinstance(in_strategy, tuple):
-            raise TypeError(f"For 'Shard', the 'in_strategy' should be a tuple, but got {type(in_strategy).__name__}")
-        if not isinstance(out_strategy, tuple):
-            raise TypeError(f"For 'Shard', the 'out_strategy' should be a tuple, "
-                            f"but got {type(out_strategy).__name__}")
-        if not isinstance(parameter_plan, (dict, type(None))):
-            raise TypeError(f"For 'Shard', the 'parameter_plan' should be a dict or None, "
-                            f"but got {type(parameter_plan).__name__}")
-        if isinstance(parameter_plan, dict):
-            for k in parameter_plan.keys():
-                v = parameter_plan[k]
-                if not isinstance(k, str) or not isinstance(v, tuple):
-                    raise TypeError(f"For 'Shard', the type of each key and value in 'parameter_plan' must be str and "
-                                    f"tuple, but got {type(k).__name__} and {type(parameter_plan[v]).__name__}")
-        parameter_plan = self._parameter_plan_dict2tuple(parameter_plan)
-
-        if not isinstance(device, str):
-            raise TypeError(f"For 'Shard', the 'device' should be a string, "
-                            f"but got {type(device).__name__}")
-        if not isinstance(level, int):
-            raise TypeError(f"For 'Shard', the 'level' should be an integer, "
-                            f"but got {type(level).__name__}")
-
-        if ms.get_algo_parameters("fully_use_devices") is True:
-            logger.warning("After calling 'shard', the environment variable 'fully_use_devices' "
-                           "will be overwritten as False")
-            ms.set_algo_parameters(fully_use_devices=False)
-
-        if self._is_attrs_has_been_set(fn, in_strategy, out_strategy, parameter_plan, device, level):
-            return self.shard_fn
-        shard_ = Shard()
-
-        if isinstance(fn, nn.Cell):
-            for param in fn.trainable_params():
-                param.is_in_shard = True
-
-        def shard_fn(*args):
-            args = (fn,) + args
-
-            @ms_function(hash_args=fn)
-            def after_shard(*args):
-                return shard_(fn, in_strategy, out_strategy, parameter_plan, device, level)(*args)
-
-            return after_shard(*args)
-
-        self.shard_fn = shard_fn
-        self.fn = fn
-        self.in_strategy = in_strategy
-        self.out_strategy = out_strategy
-        self.parameter_plan = parameter_plan
-        self.device = device
-        self.level = level
-        return self.shard_fn
-
-    @staticmethod
-    def _parameter_plan_dict2tuple(parameter_plan):
-        if not isinstance(parameter_plan, dict):
-            return parameter_plan
-
-        parameter_plan_tuple = ()
-        for k in parameter_plan:
-            parameter_plan_tuple += ((k, parameter_plan[k]),)
-        return parameter_plan_tuple
-
-    def _is_attrs_has_been_set(self, fn, in_strategy, out_strategy, parameter_plan, device, level):
-        return self.shard_fn is not None and self.fn == fn and self.in_strategy == in_strategy and \
-               self.out_strategy == out_strategy and self.parameter_plan == parameter_plan and \
-               self.device == device and self.level == level
-
-
 class _ListAppend(ListAppend_):
     """
     A metafuncgraph class that append one element to list.
@@ -1067,23 +1029,80 @@ class _ListExtend(ListExtend_):
 _extend = _ListExtend("extend")
 
 
-class _ListCount(ListCount_):
+class _DictClear(DictClear_):
     """
-    A metafuncgraph class that count the number of times an element appears in list.
+    A metafuncgraph class that clear the dict.
 
     Args:
         name (str): The name of the metafuncgraph object.
     """
 
     def __init__(self, name):
-        """Initialize _ListCount."""
-        ListCount_.__init__(self, name)
+        """Initialize _DictClear."""
+        DictClear_.__init__(self, name)
 
     def __call__(self, *args):
         pass
 
 
-_count = _ListCount("count")
+_dict_clear = _DictClear("clear")
+
+
+class _DictHasKey(DictHasKey_):
+    """
+    A metafuncgraph class that Check if key is in dict.
+
+    Args:
+        name (str): The name of the metafuncgraph object.
+    """
+
+    def __init__(self, name):
+        """Initialize _DictHasKey."""
+        DictHasKey_.__init__(self, name)
+
+    def __call__(self, *args):
+        pass
+
+
+_haskey = _DictHasKey("has_key")
+
+
+class _DictUpdate(DictUpdate_):
+    """
+    A metafuncgraph class that append another dict to the end of the dict.
+
+    Args:
+        name (str): The name of the metafuncgraph object.
+    """
+
+    def __init__(self, name):
+        """Initialize _DictUpdate."""
+        DictUpdate_.__init__(self, name)
+
+    def __call__(self, *args):
+        pass
+
+
+_update = _DictUpdate("update")
+
+
+class _DictFromKeys(DictFromKeys_):
+    """
+    A metafuncgraph class that creates a new dict from the given sequence and value.
+
+    Args:
+        name (str): The name of the metafuncgraph object.
+    """
+
+    def __init__(self, name):
+        """Initialize _DictFromKeys."""
+        DictFromKeys_.__init__(self, name)
+
+    def __call__(self, *args):
+        pass
+
+
+_fromkeys = _DictFromKeys("fromkeys")
 
 
 class _Tail(Tail_):
@@ -1118,15 +1137,3 @@ class _ZipOperation(ZipOperation_):
 
 zip_operation = _ZipOperation('zip_operation')
 """`zip_operation` will generate a tuple of zip iterations of inputs."""
-
-env_get = MultitypeFuncGraph("env_get")
-
-environ_get = Primitive('EnvironGet')
-ref_to_embed = _grad_ops.RefToEmbed()
-zeros_like = P.ZerosLike()
-
-
-@env_get.register("EnvType", "Tensor")
-def _tensor_env_get(env, parameter):
-    """Used to get env."""
-    return environ_get(env, ref_to_embed(parameter), zeros_like(parameter))