mindspore 1.10.0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops/function/grad/grad_func.py

@@ -14,23 +14,31 @@
 # ============================================================================
 
 """Defines gradient related operators with functional form."""
-
 from __future__ import absolute_import
 from functools import partial
-from mindspore.common import ms_function
+import numpy as np
+from mindspore.common import jit, mutable
 from mindspore.common import Tensor
 from mindspore.common import dtype as mstype
-from mindspore.nn.grad.cell_grad import _JvpInner
-from mindspore.nn.grad.cell_grad import _VjpInner
+from mindspore.nn.cell import Cell
 from mindspore.nn.grad.cell_grad import _LinearizeInner
 from mindspore.ops.primitive import constexpr
-from mindspore.ops.function import ones, expand_dims
-from mindspore.ops.composite import _Grad, _TaylorOperation
+from mindspore.ops.function.array_func import ones, expand_dims, size, reshape, broadcast_to, transpose
+from mindspore.ops.composite import _Vmap, _Grad, _TaylorOperation, GradOperation
 from mindspore.ops import operations as P
+from mindspore.ops.operations import _inner_ops as inner
 
 cast = P.Cast()
 dtype = P.DType()
 zeros = P.Zeros()
+oneslike = P.OnesLike()
+
+
+@constexpr
+def _check_has_aux_type(inputs):
+    if not isinstance(inputs, bool):
+        raise TypeError("The 'has_aux' must be bool type.")
+    return True
 
 
 @constexpr
@@ -38,15 +46,27 @@ def _raise_type_error():
     raise TypeError("The inputs type must be a Tensor, tuple or list of Tensors.")
 
 
+@constexpr
+def _check_duplicate_grad_position(grad_position):
+    """Check if `grad_position` has duplicate positions when `grad_position` has more than one numbers."""
+    if len(set(grad_position)) != len(grad_position):
+        raise ValueError("There are duplicate positions in `grad_position`, please check it")
+
+
 @constexpr
 def _convert_grad_position_type(grad_position):
     """Check and convert the type and size of grad position index."""
     if isinstance(grad_position, tuple):
-        for gp in grad_position:
+        _check_duplicate_grad_position(grad_position)
+        _grad_position = list(grad_position)
+        for i, gp in enumerate(_grad_position):
+            if isinstance(gp, bool):
+                _grad_position[i] = int(gp)
             if not isinstance(gp, int):
                 raise TypeError(f"For 'F.grad', the element in 'grad_position' must be int.")
             if gp < 0:
                 raise ValueError("The element in grad_position must be >= 0.")
+        grad_position = tuple(_grad_position)
     elif isinstance(grad_position, int):
         if grad_position < 0:
             raise ValueError("grad_position must be >= 0.")
@@ -57,11 +77,22 @@ def _convert_grad_position_type(grad_position):
 
 
 @constexpr
-def _get_grad_op(get_by_list, get_by_position, has_aux, get_value=False):
-    return _Grad(get_by_list=get_by_list, get_by_position=get_by_position, has_aux=has_aux, get_value=get_value)
+def _check_grad_position(grad_position, args_num):
+    """Check and convert grad position index."""
+    grad_position = _convert_grad_position_type(grad_position)
+    for gp in grad_position:
+        if gp < 0 or gp >= args_num:
+            raise ValueError("The element in grad_position must belong to [0, args_num).")
+    return grad_position
+
+
+@constexpr
+def _get_grad_op(get_by_list, get_by_position, has_aux, get_value=False, return_ids=False):
+    return _Grad(get_by_list=get_by_list, get_by_position=get_by_position, has_aux=has_aux, get_value=get_value,
+                 return_ids=return_ids)
 
 
-def grad(fn, grad_position=0, weights=None, has_aux=False):
+def grad(fn, grad_position=0, weights=None, has_aux=False, return_ids=False):
     """
     A wrapper function to generate the gradient function for the input function.
 
@@ -84,11 +115,19 @@ def grad(fn, grad_position=0, weights=None, has_aux=False):
         has_aux (bool): If True, only the first output of `fn` contributes the gradient of `fn`, while the other outputs
             will be returned straightly. It means the `fn` must return more than one outputs in this case.
             Default: False.
+        return_ids(bool): Whether return the tuple made by gradients and the index to specify which inputs
+            to be differentiated or the name of parameters of the training network that need to calculate the gradient.
+            If True, the output gradients will be replaced by the tuples made by gradients and the index to specify
+            which inputs to be differentiated or the name of parameters of the training network.
+            Default: False.
 
     Returns:
         Function, the gradient function to calculate gradient for the input function or cell.
         For example, as for `out1, out2 = fn(*args)`, when `has_aux` is set True, gradient function will return outputs
         like `(gradient, out2)` and `out2` does not contribute to the differentiation, otherwise `gradient`.
+        When return_ids is set to True, The format of the output will be the same with the output of grad when
+        return_ids is set to false, but every gradient in the output will be replaced by a tuple of position id or
+        parameter name and its gradient.
 
     Raises:
         ValueError: If both `grad_position` and `weights` are None.
@@ -102,7 +141,7 @@ def grad(fn, grad_position=0, weights=None, has_aux=False):
         >>> import mindspore
         >>> import mindspore.nn as nn
         >>> from mindspore import Tensor, ops
-        >>> from mindspore.ops import grad
+        >>> from mindspore import grad
         >>>
         >>> # Cell object to be differentiated
         >>> class Net(nn.Cell):
@@ -131,7 +170,7 @@ def grad(fn, grad_position=0, weights=None, has_aux=False):
         >>> print(aux)
         (Tensor(shape=[2], dtype=Float32, value= [ 5.00000000e+00,  5.00000000e+00]),)
         >>>
-        >>> # For given network to be differentiated with both inputs and weights, there are 3 cases.
+        >>> # For given network to be differentiated with both inputs and weights, there are 4 cases.
         >>> net = nn.Dense(10, 1)
         >>> loss_fn = nn.MSELoss()
         >>> def forward(inputs, labels):
@@ -163,17 +202,36 @@ def grad(fn, grad_position=0, weights=None, has_aux=False):
         >>> inputs_gradient, params_gradient = grad_fn(inputs, labels)
         >>> print(len(weights), len(params_gradient))
         2 2
+        >>> # Case 4: return the gradient with ids.
+        >>> import numpy as np
+        >>> import mindspore
+        >>> import mindspore.nn as nn
+        >>> from mindspore import Tensor, ops
+        >>> from mindspore import grad
+        >>>
+        >>> # Cell object to be differentiated
+        >>> class Net(nn.Cell):
+        ...     def construct(self, x, y, z):
+        ...         return x * y * z
+        >>> x = Tensor([1, 2], mindspore.float32)
+        >>> y = Tensor([-2, 3], mindspore.float32)
+        >>> z = Tensor([0, 3], mindspore.float32)
+        >>> net = Net()
+        >>> output = grad(net, grad_position=(1, 2), return_ids = True)(x, y, z)
+        >>> print(output)
+        ((1, Tensor(shape=[2], dtype=Float32, value=[ 0.00000000e+00,  6.00000000e+00])),
+         (2, Tensor(shape=[2], dtype=Float32, value=[-2.00000000e+00,  6.00000000e+00])))
     """
     if grad_position is None and weights is None:
         raise ValueError("`grad_position` and `weight` can not be None at the same time.")
 
     if grad_position is None:
-        return _get_grad_op(True, False, has_aux)(fn, weights)
+        return _get_grad_op(True, False, has_aux, False, return_ids)(fn, weights)
 
     grad_position = _convert_grad_position_type(grad_position)
     if weights is None:
-        return _get_grad_op(False, True, has_aux)(fn, None, grad_position)
-    return _get_grad_op(True, True, has_aux)(fn, weights, grad_position)
+        return _get_grad_op(False, True, has_aux, False, return_ids)(fn, None, grad_position)
+    return _get_grad_op(True, True, has_aux, False, return_ids)(fn, weights, grad_position)
 
 
 def value_and_grad(fn, grad_position=0, weights=None, has_aux=False):
@@ -216,7 +274,7 @@ def value_and_grad(fn, grad_position=0, weights=None, has_aux=False):
         >>> import numpy as np
         >>> import mindspore
         >>> from mindspore import Tensor, ops, nn
-        >>> from mindspore.ops import value_and_grad
+        >>> from mindspore import value_and_grad
         >>>
         >>> # Cell object to be differentiated
         >>> class Net(nn.Cell):
@@ -300,6 +358,55 @@ def value_and_grad(fn, grad_position=0, weights=None, has_aux=False):
     return _get_grad_op(True, True, has_aux, True)(fn, weights, grad_position)
 
 
+def get_grad(gradients, identifier):
+    """
+    When `return_ids` of :func:`mindspore.grad` is set to True, use its return value as gradients. Then find
+    the specific gradient from `gradients` according to `identifier` .
+
+    As for gradient, two typical cases are included:
+
+    1. `identifier` is the position of the specific tensor to get gradient.
+    2. `identifier` is a parameter of a network.
+
+    Args:
+        gradients (Union[tuple[int, Tensor], tuple[tuple, tuple]]): The return value of :func:`mindspore.grad`
+            when `return_ids` is set to True.
+        identifier (Union[int, Parameter]): The position number of a tensor, or a parameter that is used in
+            :func:`mindspore.grad`.
+
+    Returns:
+        The gradient of the tensor on the position or in the parameter that specified by the `identifier`.
+
+    Raises:
+        RuntimeError: If gradient is not found.
+        TypeError: If type of Args does not belong to required ones.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore
+        >>> import mindspore.nn as nn
+        >>> from mindspore import Tensor, ops
+        >>> from mindspore import grad, get_grad
+        >>>
+        >>>  # Cell object to be differentiated
+        >>> class Net(nn.Cell):
+        ...     def construct(self, x, y, z):
+        ...         return x * y * z
+        >>> x = Tensor([1, 2], mindspore.float32)
+        >>> y = Tensor([-2, 3], mindspore.float32)
+        >>> z = Tensor([0, 3], mindspore.float32)
+        >>> net = Net()
+        >>> out_grad = grad(net, grad_position=(1, 2), return_ids=True)(x, y, z)
+        >>> output = get_grad(out_grad, 1)
+        >>> print(output)
+        [0. 6.]
+    """
+    return inner.GetGrad()(gradients, identifier)
+
+
 def _trans_jet_inputs(primals_item, series_item):
     """Trans inputs of jet"""
     value_type = [mstype.int32, mstype.int64, mstype.float32, mstype.float64]
@@ -376,15 +483,14 @@ def jet(fn, primals, series):
         >>> import numpy as np
         >>> import mindspore.nn as nn
         >>> import mindspore as ms
-        >>> import mindspore.ops as P
+        >>> import mindspore.ops as ops
         >>> from mindspore import Tensor
-        >>> from mindspore.ops.functional import jet
         >>> ms.set_context(mode=ms.GRAPH_MODE)
         >>> class Net(nn.Cell):
         ...     def __init__(self):
         ...         super().__init__()
-        ...         self.sin = P.Sin()
-        ...         self.exp = P.Exp()
+        ...         self.sin = ops.Sin()
+        ...         self.exp = ops.Exp()
         ...     def construct(self, x):
         ...         out1 = self.sin(x)
         ...         out2 = self.exp(out1)
@@ -392,7 +498,7 @@ def jet(fn, primals, series):
         >>> primals = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
         >>> series = Tensor(np.array([[[1, 1], [1, 1]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]]).astype(np.float32))
         >>> net = Net()
-        >>> out_primals, out_series = jet(net, primals, series)
+        >>> out_primals, out_series = ops.jet(net, primals, series)
         >>> print(out_primals, out_series)
         [[2.319777  2.4825778]
          [1.1515628 0.4691642]] [[[ 1.2533808  -1.0331168 ]
@@ -487,15 +593,14 @@ def derivative(fn, primals, order):
         >>> import numpy as np
         >>> import mindspore as ms
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as P
+        >>> import mindspore.ops as ops
         >>> from mindspore import Tensor
-        >>> from mindspore.ops.functional import derivative
         >>> ms.set_context(mode=ms.GRAPH_MODE)
         >>> class Net(nn.Cell):
         ...     def __init__(self):
         ...         super().__init__()
-        ...         self.sin = P.Sin()
-        ...         self.exp = P.Exp()
+        ...         self.sin = ops.Sin()
+        ...         self.exp = ops.Exp()
         ...     def construct(self, x):
         ...         out1 = self.sin(x)
         ...         out2 = self.exp(out1)
@@ -503,7 +608,7 @@ def derivative(fn, primals, order):
         >>> primals = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
         >>> order = 3
         >>> net = Net()
-        >>> out_primals, out_series = derivative(net, primals, order)
+        >>> out_primals, out_series = ops.derivative(net, primals, order)
         >>> print(out_primals, out_series)
         [[2.319777  2.4825778]
          [1.1515628 0.4691642]] [[-4.0515366   3.6724353 ]
@@ -541,10 +646,20 @@ def derivative(fn, primals, order):
     return out_primals, out_series
 
 
-def jvp(fn, inputs, v):
+_grad_single = GradOperation(sens_param=True)
+_grad_all = GradOperation(sens_param=True, get_all=True)
+
+
+@constexpr
+def _check_jvp_input_v_len(inputs_len, v_len):
+    if inputs_len != v_len:
+        raise ValueError(f'v has invalid length: should be {inputs_len}, but got {v_len}')
+
+
+def jvp(fn, inputs, v, has_aux=False):
     """
     Compute the jacobian-vector-product of the given network. `jvp` matches
-    `forward-mode differentiation <https://www.mindspore.cn/docs/en/r1.10/design/auto_gradient.html#forward-mode-ad>`_.
+    `forward-mode differentiation <https://www.mindspore.cn/docs/en/r2.0/design/auto_gradient.html#forward-mode-ad>`_.
 
     Args:
         fn (Union[Function, Cell]): The function or net that takes Tensor inputs and returns single Tensor or tuple of
@@ -552,10 +667,16 @@ def jvp(fn, inputs, v):
         inputs (Union[Tensor, tuple[Tensor], list[Tensor]]): The inputs to `fn` .
         v (Union[Tensor, tuple[Tensor], list[Tensor]]): The vector in jacobian-vector-product. The shape and type of `v`
             should be the same as `inputs` .
+        has_aux (bool): If True, only the first output of `fn` contributes the gradient of `fn`, while the other outputs
+            will be returned straightly. It means the `fn` must return more than one outputs in this case.
+            Default: False.
 
     Returns:
-        - **net_output** (Union[Tensor, tuple[Tensor]]) - The result of `fn(inputs)` .
+        - **net_output** (Union[Tensor, tuple[Tensor]]) - The output of `fn(inputs)` . Specially, when `has_aux` is set
+          True, `netout` is the first output of `fn(inputs)` .
         - **jvp** (Union[Tensor, tuple[Tensor]]) - The result of jacobian-vector-product.
+        - **aux_value** (Union[Tensor, tuple[Tensor]], optional) - When `has_aux` is True, `aux_value` will be returned.
+          It means the second to last outputs of `fn(inputs)` . Specially, `aux_value` does not contribute to gradient.
 
     Raises:
         TypeError: `inputs` or `v` does not belong to required types.
@@ -564,32 +685,102 @@ def jvp(fn, inputs, v):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> from mindspore import ops
+        >>> import numpy as np
+        >>> from mindspore import jvp
         >>> from mindspore import Tensor
+        >>> import mindspore.nn as nn
         >>> class Net(nn.Cell):
         ...     def construct(self, x, y):
         ...         return x**3 + y
         >>> x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
         >>> y = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
         >>> v = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32))
-        >>> output = ops.jvp(Net(), (x, y), (v, v))
+        >>> output = jvp(Net(), (x, y), (v, v))
         >>> print(output[0])
         [[ 2. 10.]
          [30. 68.]]
         >>> print(output[1])
         [[ 4. 13.]
          [28. 49.]]
+        >>>
+        >>> def fn(x, y):
+        ...     return x ** 3 + y, y
+        >>> output, jvp_out, aux = jvp(fn, (x, y), (v, v), has_aux=True)
+        >>> print(output)
+        [[ 2. 10.]
+         [30. 68.]]
+        >>> print(jvp_out)
+        [[ 4. 13.]
+         [28. 49.]]
+        >>> print(aux)
+        [[ 1. 2.]
+         [3. 4.]]
     """
-    jvp_inner = _JvpInner()
-
-    @ms_function(hash_args=fn)
-    def _wrap_container(*arg):
-        args = arg[1:]
+    _check_has_aux_type(has_aux)
+
+    def aux_fn(*args):
+        outputs = fn(*args)
+        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]
+        return res
+
+    def grad_single(u, first_grad_single_value):
+        if has_aux:
+            return _grad_single(aux_fn)(*first_grad_single_value, u)
+        return _grad_single(fn)(*first_grad_single_value, u)
+
+    def grad_all(u, first_grad):
+        if has_aux:
+            return _grad_all(aux_fn)(*first_grad, u)
+        return _grad_all(fn)(*first_grad, u)
+
+    def _wrap_container_inner(*arg):
+        jvp_inputs = arg[1:]
         vectors = arg[0]
-        return jvp_inner(fn, vectors, *args)
+        if has_aux:
+            outputs = aux_fn(*jvp_inputs)
+        else:
+            outputs = fn(*jvp_inputs)
+        if isinstance(outputs, tuple):
+            u = ()
+            for item in outputs:
+                u = u + (mutable(oneslike(item)),)
+        else:
+            u = mutable(oneslike(outputs))
+        if len(jvp_inputs) == 1:
+            second_grad_net = _grad_single(grad_single)
+            gradient_outputs = second_grad_net(u, jvp_inputs, vectors)
+        else:
+            second_grad_net = _grad_single(grad_all)
+            gradient_outputs = second_grad_net(u, jvp_inputs, vectors)
+        if has_aux:
+            res = fn(*jvp_inputs)
+            if len(res) == 2:
+                return res[0], gradient_outputs, res[1]
+            return res[0], gradient_outputs, res[1:]
+        return outputs, gradient_outputs
+
+    if has_aux:
+        @jit(hash_args=aux_fn)
+        def _wrap_container(*arg):
+            return _wrap_container_inner(*arg)
+    else:
+        @jit(hash_args=fn)
+        def _wrap_container(*arg):
+            return _wrap_container_inner(*arg)
 
     if not isinstance(inputs, (Tensor, tuple, list)) or not isinstance(v, (Tensor, tuple, list)):
         _raise_type_error()
+
+    inputs_len = 1
+    v_len = 1
+    if isinstance(inputs, (tuple, list)):
+        inputs_len = len(inputs)
+    if isinstance(v, (tuple, list)):
+        v_len = len(v)
+    _check_jvp_input_v_len(inputs_len, v_len)
+
     if isinstance(v, list):
         v = tuple(v)
     if isinstance(inputs, (tuple, list)):
@@ -647,7 +838,7 @@ def linearize(fn, inputs):
     """
     linearize_inner = _LinearizeInner()
 
-    @ms_function(hash_args=fn)
+    @jit(hash_args=fn)
     def _wrap_container(*arg):
         args = arg[1:-1]
         vectors = arg[-1]
@@ -664,24 +855,38 @@ def linearize(fn, inputs):
     return output, partial(_wrap_container, output, *inputs)
 
 
-def vjp(fn, inputs, v):
+def _check_tensor(inputs):
+    if not isinstance(inputs, (Tensor, tuple)):
+        raise TypeError("The inputs type must be Tensor.")
+    if isinstance(inputs, tuple):
+        for item in inputs:
+            if not isinstance(item, (Tensor, tuple, list)):
+                raise TypeError("The inputs type must be Tensor.")
+    return True
+
+
+def vjp(fn, *inputs, has_aux=False):
     """
     Compute the vector-jacobian-product of the given network. `vjp` matches
-    `reverse-mode differentiation <https://www.mindspore.cn/docs/en/r1.10/design/auto_gradient.html#reverse-mode-ad>`_.
-
-    Note:
-        This function is subjected to change in the future.
+    `reverse-mode differentiation <https://www.mindspore.cn/docs/en/r2.0/design/auto_gradient.html#reverse-mode-ad>`_.
 
     Args:
         fn (Union[Function, Cell]): The function or net that takes Tensor inputs and returns single Tensor or tuple of
             Tensors.
         inputs (Union[Tensor, tuple[Tensor], list[Tensor]]): The inputs to `fn` .
-        v (Union[Tensor, tuple[Tensor], list[Tensor]]): The vector in vector-jacobian-product. The shape and type of `v`
-            should be the same as `fn(inputs)` .
+        has_aux (bool): If True, only the first output of `fn` contributes the gradient of `fn`, while the other outputs
+            will be returned straightly. It means the `fn` must return more than one outputs in this case.
+            Default: False.
 
     Returns:
-        - **net_output** (Union[Tensor, tuple[Tensor]]) - The result of `fn(inputs)` .
-        - **vjp** (Union[Tensor, tuple[Tensor]]) - The result of vector-jacobian-product.
+        Forward outputs and function to calculate vjp.
+
+        - **net_output** (Union[Tensor, tuple[Tensor]]) - The output of `fn(inputs)`. Specially, when `has_aux` is set
+          True, `netout` is the first output of `fn(inputs)`.
+        - **vjp_fn** (Function) - To calculate vector-jacobian-product. Its inputs are the vectors whose shape and
+          type should be the same as `netout` .
+        - **aux_value** (Union[Tensor, tuple[Tensor]], optional) - When `has_aux` is True, `aux_value` will be returned.
+          It means the second to last outputs of `fn(inputs)`. Specially, `aux_value` does not contribute to gradient.
 
     Raises:
         TypeError: `inputs` or `v` does not belong to required types.
@@ -690,7 +895,9 @@ def vjp(fn, inputs, v):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> from mindspore import ops
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> from mindspore import vjp
         >>> from mindspore import Tensor
         >>> class Net(nn.Cell):
         ...     def construct(self, x, y):
@@ -698,41 +905,505 @@ def vjp(fn, inputs, v):
         >>> x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
         >>> y = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
         >>> v = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32))
-        >>> output = ops.vjp(Net(), (x, y), v)
-        >>> print(output[0])
+        >>> outputs, vjp_fn = vjp(Net(), x, y)
+        >>> print(outputs)
         [[ 2. 10.]
          [30. 68.]]
-        >>> print(output[1])
+        >>> gradient = vjp_fn(v)
+        >>> print(gradient)
         (Tensor(shape=[2, 2], dtype=Float32, value=
         [[ 3.00000000e+00,  1.20000000e+01],
          [ 2.70000000e+01,  4.80000000e+01]]), Tensor(shape=[2, 2], dtype=Float32, value=
         [[ 1.00000000e+00,  1.00000000e+00],
          [ 1.00000000e+00,  1.00000000e+00]]))
+        >>> def fn(x, y):
+        ...     return 2 * x + y, y ** 3
+        >>> outputs, vjp_fn, aux = vjp(fn, x, y, has_aux=True)
+        >>> gradient = vjp_fn(v)
+        >>> print(outputs)
+        [[ 3.  6.]
+         [ 9. 12.]]
+        >>> print(aux)
+        [[ 1.  8.]
+         [27. 64.]]
+        >>> print(gradient)
+        (Tensor(shape=[2, 2], dtype=Float32, value=
+        [[ 2.00000000e+00,  2.00000000e+00],
+         [ 2.00000000e+00,  2.00000000e+00]]), Tensor(shape=[2, 2], dtype=Float32, value=
+        [[ 1.00000000e+00,  1.00000000e+00],
+         [ 1.00000000e+00,  1.00000000e+00]]))
     """
-    vjp_inner = _VjpInner()
+    _check_tensor(inputs)
+    _check_has_aux_type(has_aux)
+
+    def aux_fn(*args):
+        outputs = fn(*args)
+        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]
+        return res
+
+    def wrap_container(*v):
+        _check_tensor(v)
+        if has_aux:
+            fn_ = aux_fn
+        else:
+            fn_ = fn
+        if len(v) == 1:
+            return _grad_all(fn_)(*inputs, v[0])
+        return _grad_all(fn_)(*inputs, v)
 
-    @ms_function(hash_args=fn)
-    def wrap_container(*arg):
-        args = arg[:-1]
-        vectors = arg[-1]
-        return vjp_inner(fn, *args, vectors)
+    res = fn(*inputs)
+    if has_aux:
+        if len(res) == 2:
+            return res[0], wrap_container, res[1]
+        return res[0], wrap_container, res[1:]
+    return res, wrap_container
 
-    if not isinstance(inputs, (Tensor, tuple, list)) or not isinstance(v, (Tensor, tuple, list)):
-        _raise_type_error()
-    if isinstance(v, list):
-        v = tuple(v)
-    if isinstance(inputs, (tuple, list)):
-        return wrap_container(*inputs, v)
-    return wrap_container(inputs, v)
+
+@constexpr
+def _jac_generate_target_dimension(x):
+    """For given length = len(x), this method generates target dimension tuple (1, 2, 3,..., length, 0)."""
+    target_dimension = tuple(index + 1 for index, _ in enumerate(x[1:])) + (0,)
+    return target_dimension
+
+
+def _jacfwd_trans_item(item, inputs_shape, grad_position):
+    """transfer origin item to derivative of each output with respect to each input."""
+    output_wrt_input_all = ()
+    for i in grad_position:
+        origin_output_wrt_input = item[inputs_shape[i][1]:inputs_shape[i + 1][1]]
+        target_dimension = _jac_generate_target_dimension(origin_output_wrt_input.shape)
+        temp = transpose(origin_output_wrt_input, target_dimension)
+        output_wrt_input = reshape(temp, temp.shape[:-1] + inputs_shape[i + 1][0])
+        output_wrt_input_all += (output_wrt_input,)
+    return output_wrt_input_all
+
+
+def _jac_postprocess(x, shape, grad_position, mode):
+    """reformat jacobian."""
+
+    if mode == 'forward':
+        func = _jacfwd_trans_item
+        args = (shape, grad_position)
+    else:
+        func = _jacrev_trans_item
+        args = (shape,)
+
+    if isinstance(x, tuple):
+        jacobian = ()
+        for item in x:
+            jacobian += func(item, *args)
+        res = jacobian
+    else:
+        res = func(x, *args)
+    if len(res) == 1:
+        return res[0]
+    input_num = len(grad_position)
+    if len(res) % input_num != 0:
+        raise ValueError("The numbers of inputs and outputs do not match.")
+    output_num = len(res) // input_num
+    if input_num == 1 or output_num == 1:
+        return res
+    jac = ()
+    for i in range(output_num):
+        input_grad = ()
+        for j in range(input_num):
+            if mode == 'forward':
+                grad_increment = (res[i * input_num + j],)
+            else:
+                grad_increment = (res[j * output_num + i],)
+            input_grad += grad_increment
+        jac += (input_grad,)
+    return jac
+
+
+def _jacfwd_postprocess(x, inputs_shape, grad_position):
+    """reformat forward-computed Jacobian."""
+    return _jac_postprocess(x, inputs_shape, grad_position, 'forward')
+
+
+def _jacfwd_construct_v(inputs, grad_position):
+    """
+    For input (x1, x2), x1.shape = (a, b), x2.shape = (c, d), this method generates corresponding v (v1, v2),
+    v1.shape = (N, a, b), v2.shape = (N, c, d), while N = a*b + c*d.
+    """
+    v = ()
+    primals = ()
+    inputs_shape = (((), 0),)
+    num = 0
+    items_num = ()
+    cum_num = (0,)
+    for item in inputs:
+        num += size(item)
+        inputs_shape += ((item.shape, num),)
+        items_num += (size(item),)
+        cum_num += (num,)
+    for i, element in enumerate(inputs):
+        item_size = items_num[i]
+        if i in grad_position:
+            temp2 = Tensor(np.eye(num, item_size, -cum_num[i], np.float32))
+        else:
+            temp2 = zeros((num, item_size), mstype.float32)
+        input_v = reshape(temp2, (num,) + element.shape)
+        primal = broadcast_to(element, (num,) + element.shape)
+        v += (input_v,)
+        primals += (primal,)
+    if len(inputs) == 1:
+        return primals, v[0], inputs_shape
+    return primals, v, inputs_shape
+
+
+_vmap = _Vmap()
+
+
+def jacfwd(fn, grad_position=0, has_aux=False):
+    """
+    Compute Jacobian via forward mode, corresponding to
+    `forward-mode differentiation <https://www.mindspore.cn/docs/en/r2.0/design/auto_gradient.html#forward-mode-ad>`_.
+    When number of outputs is much greater than that of inputs, it's better to calculate Jacobian via forward mode than
+    reverse mode to get better performance.
+
+    Args:
+        fn (Union[Cell, Function]): Function to do GradOperation.
+        grad_position (Union[int, tuple[int]], optional): If int, get the gradient with respect to single input.
+            If tuple, get the gradients with respect to selected inputs. 'grad_position' begins with 0. Default: 0.
+        has_aux (bool, optional): If True, only the first output of `fn` contributes the gradient of `fn`,
+            while the other outputs will be returned straightly. It means the `fn` must return more than one
+            outputs in this case. Default: False.
+
+    Returns:
+        Function, returns the Jacobian function for the input function or cell.
+        For example, as for `out1, out2 = fn(*args)`, when `has_aux` is set True, gradient function will return outputs
+        like `(Jacobian, out2)` and `out2` does not contribute to the differentiation, otherwise `Jacobian` .
+
+    Raises:
+        TypeError: `grad_position` or `has_aux` does not belong to required types.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> from mindspore import jacfwd
+        >>> from mindspore import Tensor
+        >>> class MultipleInputsMultipleOutputsNet(nn.Cell):
+        ...     def construct(self, x, y, z):
+        ...         return x ** 2 + y ** 2 + z ** 2, x * y * z
+        >>> x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+        >>> y = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+        >>> z = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32))
+        >>> net = MultipleInputsMultipleOutputsNet()
+        >>> jac, aux = jacfwd(net, grad_position=0, has_aux=True)(x, y, z)
+        >>> print(jac)
+        [[[[ 2.,  0.]
+           [ 0.,  0.]]
+          [[ 0.,  4.]
+           [ 0.,  0.]]]
+         [[[ 0.,  0.]
+           [ 6.,  0.]]
+          [[ 0.,  0.]
+           [ 0.,  8.]]]]
+        >>> print(aux)
+        [[ 1.  4.]
+         [ 9. 16.]]
+    """
+    _check_has_aux_type(has_aux)
+
+    def aux_fn(*args):
+        outputs = fn(*args)
+        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]
+        return res
+
+    def grad_single(u, first_grad_single_value):
+        if has_aux:
+            return _grad_single(aux_fn)(*first_grad_single_value, u)
+        return _grad_single(fn)(*first_grad_single_value, u)
+
+    def grad_all(u, first_grad):
+        if has_aux:
+            return _grad_all(aux_fn)(*first_grad, u)
+        return _grad_all(fn)(*first_grad, u)
+
+    @jit
+    def wrapped(*args):
+        checked_grad_position = _check_grad_position(grad_position, len(args))
+        primals, v, inputs_shape = _jacfwd_construct_v(args, checked_grad_position)
+
+        def inner_fn(jvp_inputs, vectors):
+            outputs = fn(*jvp_inputs)
+            if isinstance(outputs, tuple):
+                u = ()
+                for item in outputs:
+                    u = u + (mutable(oneslike(item)),)
+            else:
+                u = mutable(oneslike(outputs))
+            if len(jvp_inputs) == 1:
+                second_grad_net = _grad_single(grad_single)
+            else:
+                second_grad_net = _grad_single(grad_all)
+            gradient_outputs = second_grad_net(u, jvp_inputs, vectors)
+            return gradient_outputs
+
+        def inner_aux_fn(jvp_inputs, vectors):
+            outputs = aux_fn(*jvp_inputs)
+            u = mutable(oneslike(outputs))
+            if len(jvp_inputs) == 1:
+                second_grad_net = _grad_single(grad_single)
+            else:
+                second_grad_net = _grad_single(grad_all)
+            gradient_outputs = second_grad_net(u, jvp_inputs, vectors)
+            return gradient_outputs
+
+        if has_aux:
+            res = _vmap(inner_aux_fn)(primals, v)
+            jac_res = _jacfwd_postprocess(res, inputs_shape, checked_grad_position)
+            forward_outputs = fn(*args)
+            if len(forward_outputs) == 2:
+                return jac_res, forward_outputs[1]
+            return jac_res, forward_outputs[1:]
+        res = _vmap(inner_fn)(primals, v)
+        jac_res = _jacfwd_postprocess(res, inputs_shape, checked_grad_position)
+        return jac_res
+
+    return wrapped
+
+
+def _jacrev_trans_item(item, outputs_shape):
+    """transfer origin item to derivative of each output with respect to each input."""
+    output_wrt_input_all = ()
+    length = len(outputs_shape) - 1
+    for i in range(length):
+        origin_output_wrt_input = item[outputs_shape[i][1]:outputs_shape[i + 1][1]]
+        target_dimension = _jac_generate_target_dimension(origin_output_wrt_input.shape)
+        temp = transpose(origin_output_wrt_input, target_dimension)
+        output_wrt_input = reshape(origin_output_wrt_input, outputs_shape[i + 1][0] + temp.shape[:-1])
+        output_wrt_input_all += (output_wrt_input,)
+    return output_wrt_input_all
+
+
+def _jacrev_postprocess(x, outputs_shape, grad_position):
+    """reformat reverse-computed jacobian."""
+    return _jac_postprocess(x, outputs_shape, grad_position, 'reverse')
+
+
+def _jacrev_construct_v(inputs, outputs, has_aux=False):
+    """
+    For outputs (y1, y2), y1.shape = (a, b), y2.shape = (c, d), this method generates corresponding v (v1, v2),
+    v1.shape = (N, a, b), v2.shape = (N, c, d), while N = a*b + c*d.
+    """
+    if isinstance(outputs, Tensor):
+        outputs = (outputs,)
+    if has_aux:
+        outputs = (outputs[0],)
+    v = ()
+    primals = ()
+    outputs_shape = (((), 0),)
+    num = 0
+    items_num = ()
+    cum_num = (0,)
+    for item in outputs:
+        item_num = size(item)
+        num += item_num
+        outputs_shape += ((item.shape, num),)
+        items_num += (item_num,)
+        cum_num += (num,)
+    for element in inputs:
+        primal = broadcast_to(element, (num,) + element.shape)
+        primals += (primal,)
+    for i, element in enumerate(outputs):
+        item_size = items_num[i]
+        temp2 = Tensor(np.eye(num, item_size, -cum_num[i], np.float32))
+        output_v = reshape(temp2, (num,) + element.shape)
+        v += (output_v,)
+    if len(outputs) == 1 or has_aux:
+        return primals, v[0], outputs_shape
+    return primals, v, outputs_shape
+
+
+_grad = _Grad(get_by_position=True, has_aux=False, sens_param=True)
+
+
+def jacrev(fn, grad_position=0, has_aux=False):
+    """
+    Compute Jacobian via reverse mode, corresponding to
+    `reverse-mode differentiation <https://www.mindspore.cn/docs/en/r2.0/design/auto_gradient.html#reverse-mode-ad>`_.
+    When number of inputs is much greater than that of outputs, it's better to calculate Jacobian via reverse mode than
+    forward mode to get better performance.
+
+    Args:
+        fn (Union[Cell, Function]): Function to do GradOperation.
+        grad_position (Union[int, tuple[int]], optional): If int, get the gradient with respect to single input.
+            If tuple, get the gradients with respect to selected inputs. 'grad_position' begins with 0. Default: 0.
+        has_aux (bool, optional): If True, only the first output of `fn` contributes the gradient of `fn`,
+            while the other outputs will be returned straightly. It means the `fn` must return more than
+            one outputs in this case. Default: False.
+
+    Returns:
+        Function, returns the Jacobian function for the input function or cell.
+        For example, as for `out1, out2 = fn(*args)`, when `has_aux` is set True, gradient function will return outputs
+        like `(Jacobian, out2)` and `out2` does not contribute to the differentiation, otherwise `Jacobian` .
+
+    Raises:
+        TypeError: `grad_position` or `has_aux` does not belong to required types.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> from mindspore import jacrev
+        >>> from mindspore import Tensor
+        >>> class MultipleInputsMultipleOutputsNet(nn.Cell):
+        ...     def construct(self, x, y, z):
+        ...         return x ** 2 + y ** 2 + z ** 2, x * y * z
+        >>> x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+        >>> y = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+        >>> z = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32))
+        >>> net = MultipleInputsMultipleOutputsNet()
+        >>> jac, aux = jacrev(net, grad_position=0, has_aux=True)(x, y, z)
+        >>> print(jac)
+        [[[[ 2.,  0.]
+           [ 0.,  0.]]
+          [[ 0.,  4.]
+           [ 0.,  0.]]]
+         [[[ 0.,  0.]
+           [ 6.,  0.]]
+          [[ 0.,  0.]
+           [ 0.,  8.]]]]
+        >>> print(aux)
+        [[ 1.  4.]
+         [ 9. 16.]]
+    """
+    _check_has_aux_type(has_aux)
+
+    def aux_fn(*args):
+        outputs = fn(*args)
+        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]
+        return res
+
+    @jit
+    def wrapped(*args):
+        checked_grad_position = _check_grad_position(grad_position, len(args))
+        outputs = fn(*args)
+        primals, v, outputs_shape = _jacrev_construct_v(args, outputs, has_aux)
+
+        def inner_fn(vjp_inputs, vectors):
+            gradient_outputs = _grad(fn, None, checked_grad_position)(*vjp_inputs, vectors)
+            return gradient_outputs
+
+        def inner_aux_fn(vjp_inputs, vectors):
+            gradient_outputs = _grad(aux_fn, None, checked_grad_position)(*vjp_inputs, vectors)
+            return gradient_outputs
+
+        if has_aux:
+            res = _vmap(inner_aux_fn)(primals, v)
+            jac_res = _jacrev_postprocess(res, outputs_shape, checked_grad_position)
+            forward_outputs = fn(*args)
+            if len(forward_outputs) == 2:
+                return jac_res, forward_outputs[1]
+            return jac_res, forward_outputs[1:]
+
+        res = _vmap(inner_fn)(primals, v)
+        jac_res = _jacrev_postprocess(res, outputs_shape, checked_grad_position)
+        return jac_res
+
+    return wrapped
+
+
+def custom_vjp(fn=None):
+    """
+    Support vjp to custom bprop for function.
+
+    Args:
+        fn (function): The `fn` that need to define custom bprop. Default: None.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+    """
+
+    def deco(fn):
+        class CustomVjp(Cell):
+            """
+            The CustomVjp decorates function into cell to support custom bprop.
+            """
+
+            def __init__(self, fwd):
+                super(CustomVjp, self).__init__()
+                self.fwd = fwd
+                self.bwd = None
+                self.add_flags(custom_vjp=True)
+
+            def construct(self, *args):
+                return self.fwd(*args)
+
+            def defbwd(self, bwd):
+                self.bwd = bwd
+
+            def bprop(self, *args):
+                return self.bwd(*args)
+
+        return CustomVjp(fn)
+
+    if fn is not None:
+        return deco(fn)
+    return deco
+
+
+def stop_gradient(value):
+    """
+    StopGradient is used for eliminating the effect of a value on the gradient, such as truncating
+    the gradient propagation from an output of a function.
+    For more details, please refer to `Stop Gradient
+    <https://www.mindspore.cn/tutorials/en/r2.0/beginner/autograd.html#stop-gradient>`_.
+
+    Args:
+        value (Any): The value whose effect on the gradient to be eliminated.
+
+    Returns:
+        The same as `value`.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore.ops as ops
+        >>> from mindspore import Tensor
+        >>> from mindspore import dtype as mstype
+        >>> def net(x, y):
+        ...     out1 = ops.MatMul()(x, y)
+        ...     out2 = ops.MatMul()(x, y)
+        ...     out2 = ops.stop_gradient(out2)
+        ...     return out1, out2
+        ...
+        >>> x = Tensor([[0.5, 0.6, 0.4], [1.2, 1.3, 1.1]], dtype=mstype.float32)
+        >>> y = Tensor([[0.01, 0.3, 1.1], [0.1, 0.2, 1.3], [2.1, 1.2, 3.3]], dtype=mstype.float32)
+        >>> grad_fn = ops.grad(net)
+        >>> output = grad_fn(x, y)
+        >>> print(output)
+        [[1.4100001 1.6       6.5999994]
+         [1.4100001 1.6       6.5999994]]
+    """
+    return P.StopGradient()(value)
 
 
 __all__ = [
     'grad',
     'value_and_grad',
+    'jacfwd',
+    'jacrev',
     'jet',
     'derivative',
     'jvp',
     'vjp',
-    'linearize'
+    'linearize',
+    'stop_gradient',
+    'get_grad'
 ]
 __all__.sort()