mindspore 2.3.0__cp310-cp310-win_amd64.whl

mindspore/nn/extend/embedding.py

@@ -0,0 +1,143 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""embedding"""
+from __future__ import absolute_import
+
+import mindspore.common.dtype as mstype
+from mindspore.common.initializer import Normal
+from mindspore import _checkparam as Validator
+from mindspore.nn.cell import Cell
+from mindspore import ops
+from mindspore.common.parameter import Parameter
+from mindspore.common.tensor import Tensor
+
+__all__ = ['Embedding']
+
+
+class Embedding(Cell):
+    r"""
+    Embedding layer.
+    Retrieve the word embeddings in weight stored in the layer using indices specified in `input`.
+
+    .. warning::
+        On Ascend, the behavior is unpredictable when the value of `input` is invalid.
+
+    Args:
+        num_embeddings (int): Size of the dictionary of embeddings.
+        embedding_dim (int): The size of each embedding vector.
+        padding_idx (int, optional): If the value is not None, the corresponding row of embedding vector
+            will not be updated in training. The value of embedding vector at `padding_idx` will default
+            to zeros when the Embedding layer is newly constructed. The value should be in range
+            `[-num_embeddings, num_embeddings)` if it's not ``None``. Default ``None``.
+        max_norm (float, optional): If the value is not None, firstly get the p-norm result of the embedding
+            vector specified by `input` where p is specified by `norm_type`; if the result is larger then `max_norm`,
+            update the embedding vector` with :math:`\frac{max\_norm}{result+1e^{-7}}`. Default ``None``.
+        norm_type (float, optional): Indicated the value of p in p-norm. Default ``2.0``.
+        scale_grad_by_freq (bool, optional): If ``True`` the gradients will be scaled by the inverse of frequency
+            of the index in `input`. Default ``False``.
+        _weight (Tensor, optional): Used to initialize the weight of Embedding. If ``None``, the weight will be
+            initialized from normal distribution :math:`{N}(\text{sigma=1.0}, \text{mean=0.0})`. Default ``None``.
+        dtype (mindspore.dtype, optional) : Dtype of Parameters. It is meaningless when `_weight` is not None.
+            Default: ``mindspore.float32``.
+
+    Inputs:
+        - **input** (Tensor) - The indices used to lookup in the embedding vector. The data type must be
+          mindspore.int32 or mindspore.int64, and the value should be in range `[0, num_embeddings)`.
+
+    Outputs:
+        Tensor, has the same data type as weight, the shape is :math:`(*input.shape, embedding\_dim)`.
+
+    Raises:
+        TypeError: If `num_embeddings` is not an int.
+        TypeError: If `embedding_dim` is not an int.
+        ValueError: If `padding_idx` is out of valid range.
+        TypeError: If `max_norm` is not a float.
+        TypeError: If `norm_type` is not a float.
+        TypeError: If `scale_grad_by_freq` is not a bool.
+        TypeError: If `dtype` is not one of mindspore.dtype.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, nn
+        >>> input = Tensor([[1, 0, 1, 1], [0, 0, 1, 0]])
+        >>> embedding = nn.extend.Embedding(num_embeddings=10, embedding_dim=3)
+        >>> output = embedding(input)
+        >>> print(output)
+        [[[-0.0024154  -0.01203444  0.00811537]
+          [ 0.00233847 -0.00596091  0.00536799]
+          [-0.0024154  -0.01203444  0.00811537]
+          [-0.0024154  -0.01203444  0.00811537]]
+         [[ 0.00233847 -0.00596091  0.00536799]
+          [ 0.00233847 -0.00596091  0.00536799]
+          [-0.0024154  -0.01203444  0.00811537]
+          [ 0.00233847 -0.00596091  0.00536799]]]
+    """
+
+    def __init__(self, num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0,
+                 scale_grad_by_freq=False, _weight=None, dtype=mstype.float32):
+        """Initialize Embedding."""
+        super().__init__()
+        self.num_embeddings = Validator.check_value_type(
+            'num_embeddings', num_embeddings, [int], self.cls_name)
+        self.embedding_dim = Validator.check_value_type(
+            'embedding_dim', embedding_dim, [int], self.cls_name)
+        Validator.check_subclass(
+            "dtype", dtype, mstype.number_type, self.cls_name)
+        self.dtype = dtype
+        self.padding_idx = padding_idx
+        if _weight is None:
+            init_tensor = Tensor(shape=[num_embeddings, embedding_dim], dtype=dtype, init=Normal(1, 0))
+            init_tensor = self._zero_weight_by_index(init_tensor)
+            self.weight = Parameter(init_tensor, name='weight')
+        else:
+            self.weight = Parameter(_weight)
+
+        self.max_norm = max_norm
+        if max_norm is not None:
+            self.max_norm = Validator.check_value_type('max_norm', max_norm, [float], self.cls_name)
+
+        self.norm_type = norm_type
+        if norm_type is not None:
+            self.norm_type = Validator.check_value_type('norm_type', norm_type,
+                                                        [float], self.cls_name)
+
+        self.scale_grad_by_freq = scale_grad_by_freq
+        if scale_grad_by_freq is not None:
+            self.scale_grad_by_freq = Validator.check_value_type('scale_grad_by_freq',
+                                                                 scale_grad_by_freq,
+                                                                 [bool], self.cls_name)
+
+    def _zero_weight_by_index(self, init_tensor):
+        if self.padding_idx is not None:
+            self.padding_idx = Validator.check_int_range(self.padding_idx, -self.num_embeddings, self.num_embeddings,
+                                                         Validator.INC_LEFT, "padding_idx", self.cls_name)
+            if isinstance(init_tensor, Tensor) and init_tensor.init is not None:
+                init_tensor = init_tensor.init_data()
+            init_tensor[self.padding_idx] = 0
+
+        return init_tensor
+
+    def construct(self, input):
+        return ops.embedding(input, self.weight, self.padding_idx, self.max_norm,
+                             self.norm_type, self.scale_grad_by_freq)
+
+    def extend_repr(self):
+        return f'num_embeddings={self.num_embeddings}, embedding_dim={self.embedding_dim}, '  \
+               f'padding_idx={self.padding_idx}, max_norm={self.max_norm}, norm_type={self.norm_type}, ' \
+               f'scale_grad_by_freq={self.scale_grad_by_freq}, dtype={self.dtype}'

mindspore/nn/extend/layer/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+Layer.
+
+The high-level components(Cells) used to construct the neural network.
+"""
+from __future__ import absolute_import
+
+from mindspore.nn.extend.layer import normalization
+from mindspore.nn.extend.layer.normalization import *
+
+__all__ = []
+
+__all__.extend(normalization.__all__)

mindspore/nn/extend/layer/normalization.py

@@ -0,0 +1,109 @@
+# Copyright 2020-2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""normalization"""
+from __future__ import absolute_import
+from __future__ import division
+
+from mindspore.ops import functional as F
+from mindspore.common.parameter import Parameter
+from mindspore.common.initializer import initializer
+from mindspore.common import dtype as mstype
+from mindspore.nn.cell import Cell
+
+__all__ = ['LayerNorm']
+
+
+class LayerNorm(Cell):
+    r"""
+    Applies Layer Normalization over a mini-batch of inputs.
+
+    Layer Normalization is widely used in recurrent neural networks. It applies
+    normalization on a mini-batch of inputs for each single training case as described
+    in the paper `Layer Normalization <https://arxiv.org/pdf/1607.06450.pdf>`_. Unlike Batch
+    Normalization, Layer Normalization performs exactly the same computation at training and
+    testing time. It is applied across all channels and pixel but only one batch size.
+    :math:`\gamma` and :math:`\beta` are trainable scale and shift.
+    It can be described using the following formula:
+
+    .. math::
+        y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    Args:
+        normalized_shape (Union(tuple[int], list[int])): The normalized shape of `x` for LayerNorm
+        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\gamma` weight.
+            The values of str refer to the function `initializer` including ``'zeros'`` , ``'ones'`` ,
+            ``'xavier_uniform'`` , ``'he_uniform'`` , etc. Default: ``'ones'`` .
+        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\beta` weight.
+            The values of str refer to the function `initializer` including ``'zeros'`` , ``'ones'`` ,
+            ``'xavier_uniform'`` , ``'he_uniform'`` , etc. Default: ``'zeros'`` .
+        eps (float): A value added to the denominator for numerical stability(:math:`\epsilon`). Default: ``1e-5`` .
+        elementwise_affine (bool): A bool value, When set to True, gamma and beta can be learned. Default: True.
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
+
+    Inputs:
+        - **x** (Tensor) - The shape is :math:`(N, *)`, where :math:`*` means, any number of additional dimensions.
+
+    Outputs:
+        Tensor, the normalized and scaled offset tensor, has the same shape and data type as the `x`.
+
+    Raises:
+        TypeError: If `epsilon` is not a float.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> x = ms.Tensor(np.ones([20, 5, 10, 10]), ms.float32)
+        >>> shape1 = x.shape[1:]
+        >>> m = ms.nn.extend.LayerNorm(shape1)
+        >>> output = m(x).shape
+        >>> print(output)
+        (20, 5, 10, 10)
+    """
+
+    def __init__(self,
+                 normalized_shape,
+                 gamma_init='ones',
+                 beta_init='zeros',
+                 eps=1e-5,
+                 elementwise_affine=True,
+                 dtype=mstype.float32
+                 ):
+        """Initialize LayerNorm."""
+        super(LayerNorm, self).__init__()
+        if not isinstance(normalized_shape, (tuple, list)):
+            raise TypeError(f"For '{self.cls_name}', the type of 'normalized_shape' must be tuple[int] or list[int], "
+                            f"but got {normalized_shape} and the type is {type(normalized_shape)}.")
+        if not normalized_shape:
+            raise ValueError(
+                f"Expected normalized_shape to be at least 1-dimensional, i.e., containing at "
+                f"least one element, but got normalized_shape = {normalized_shape}"
+            )
+        self.normalized_shape = normalized_shape
+        self.epsilon = eps
+        self.gamma = Parameter(initializer(
+            gamma_init, normalized_shape, dtype=dtype), name="gamma", requires_grad=elementwise_affine)
+        self.beta = Parameter(initializer(
+            beta_init, normalized_shape, dtype=dtype), name="beta", requires_grad=elementwise_affine)
+
+    def construct(self, input_x):
+        y = F.layer_norm(input_x, self.normalized_shape, self.gamma.astype(input_x.dtype),
+                         self.beta.astype(input_x.dtype), self.epsilon)
+        return y
+
+    def extend_repr(self):
+        return 'normalized_shape={}, gamma{}, beta={}'.format(self.normalized_shape, self.gamma, self.beta)

mindspore/nn/extend/pooling.py

@@ -0,0 +1,117 @@
+#Copyright 2020-2022 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""pooling"""
+from __future__ import absolute_import
+
+from mindspore.ops.auto_generate.gen_ops_prim import MaxPoolWithIndices, MaxPoolWithMask
+from mindspore.nn.cell import Cell
+
+__all__ = ['MaxPool2d']
+
+
+class MaxPool2d(Cell):
+    r"""
+    Applies a 2D max pooling over an input Tensor which can be regarded as a composition of 2D planes.
+
+    Typically the input is of shape :math:`(N_{in}, C_{in}, H_{in}, W_{in})`, MaxPool2d outputs
+    regional maximum in the :math:`(H_{in}, W_{in})`-dimension. Given kernel size
+    :math:`(h_{ker}, w_{ker})` and stride :math:`(s_0, s_1)`, the operation is as follows.
+
+    .. math::
+        \text{output}(N_i, C_j, h, w) = \max_{m=0, \ldots, h_{ker}-1} \max_{n=0, \ldots, w_{ker}-1}
+        \text{input}(N_i, C_j, s_0 \times h + m, s_1 \times w + n)
+
+    .. warning::
+        Only support on Atlas training series.
+
+    Args:
+        kernel_size (Union[int, tuple[int]]): The size of kernel used to take the max value,
+            is an int number or a single element tuple that represents height and width are both kernel_size,
+            or a tuple of two int numbers that represent height and width respectively.
+            Default: ``1`` .
+        stride (Union[int, tuple[int], None]): The distance of kernel moving, an int number or a single element tuple
+            that represents the height and width of movement are both stride, or a tuple of two int numbers that
+            represent height and width of movement respectively.
+            Default: ``None`` , which indicates the moving step is `kernel_size` .
+        padding (Union(int, tuple[int], list[int])): Specifies the padding value of the pooling operation.
+            Default: ``0`` . `padding` can only be an integer or a tuple/list containing one or two integers. If
+            `padding` is an integer or a tuple/list containing one integer, it will be padded `padding` times in the
+            four directions of the input. If `padding` is a tuple/list containing two integers, it will be padded
+            `padding[0]` times in the up-down direction of the input and `padding[1]` times in the left-right direction
+            of the input.
+        dilation (Union(int, tuple[int])): The spacing between the elements of the kernel in convolution,
+            used to increase the receptive field of the pooling operation. If it is a tuple, it must contain one or two
+            integers. Default: ``1`` .
+        return_indices (bool): If ``True`` , the function will return both the result of max pooling and the indices of
+            the max elements. Default: ``False`` .
+        ceil_mode (bool): If ``True`` , use ceil to compute the output shape instead of floor. Default: ``False`` .
+
+    Inputs:
+        - **input** (Tensor) - Tensor of shape :math:`(N, C_{in}, H_{in}, W_{in})`.
+
+    Outputs:
+        If `return_indices` is ``False`` , return a Tensor `output`, else return a tuple (`output`, `argmax`).
+
+        - **output** (Tensor) - Maxpooling result, with shape :math:`(N_{out}, C_{out}, H_{out}, W_{out})`. It has the
+          same data type as `input`.
+        - **argmax** (Tensor) - Index corresponding to the maximum value. Data type is int32.
+
+        .. math::
+            H_{out} = \left\lfloor\frac{H_{in} + 2 * \text{padding[0]} - \text{dilation[0]}
+                \times (\text{kernel_size[0]} - 1) - 1}{\text{stride[0]}} + 1\right\rfloor
+
+        .. math::
+            W_{out} = \left\lfloor\frac{W_{in} + 2 * \text{padding[1]} - \text{dilation[1]}
+                \times (\text{kernel_size[1]} - 1) - 1}{\text{stride[1]}} + 1\right\rfloor
+
+    Raises:
+        TypeError: If `input` is not a Tensor.
+        ValueError: If length of shape of `input` is not equal to 4.
+        TypeError: If `kernel_size` , `stride` , `padding` or `dilation` is not int or tuple.
+        ValueError: If `kernel_size`, `stride` or `dilation` is less than 1.
+        ValueError: If `dilation` is not all 1.
+        ValueError: If `padding` is less than 0.
+        ValueError: If `padding` is more than half of `kernel_size`.
+        TypeError: If `ceil_mode` is not bool.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> pool = ms.nn.extend.MaxPool2d(kernel_size=3, stride=1)
+        >>> input = ms.Tensor(np.random.randint(0, 10, [1, 2, 4, 4]), ms.float32)
+        >>> output = pool(input)
+        >>> print(output.shape)
+        (1, 2, 2, 2)
+    """
+
+    def __init__(self, kernel_size=1, stride=None, padding=0, dilation=1, return_indices=False,
+                 ceil_mode=False):
+        """Initialize MaxPool2d."""
+        super(MaxPool2d, self).__init__()
+        self.return_indices = return_indices
+        strides = stride if (stride is not None) else kernel_size
+        if return_indices:
+            self.max_pool_func_ = MaxPoolWithIndices(kernel_size, strides, padding, dilation, ceil_mode)
+        else:
+            self.max_pool_func_ = MaxPoolWithMask(kernel_size, strides, padding, dilation, ceil_mode)
+
+    def construct(self, input):
+        out, indices = self.max_pool_func_(input)
+        if self.return_indices:
+            return out, indices
+        return out

mindspore/nn/grad/__init__.py

@@ -0,0 +1,21 @@
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Cells of grad function. Calculate the gradient of input network or function."""
+from __future__ import absolute_import
+
+from mindspore.nn.grad.cell_grad import Jvp, Vjp
+
+
+__all__ = ['Jvp', 'Vjp']

mindspore/nn/grad/cell_grad.py

@@ -0,0 +1,196 @@
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""cell grad"""
+from __future__ import absolute_import
+
+from mindspore.nn.cell import Cell
+from mindspore.ops import composite as C
+from mindspore.ops import operations as P
+from mindspore.ops.operations import _inner_ops as inner
+from mindspore.ops.primitive import Primitive
+from mindspore.common import dtype as mstype
+from mindspore.common.api import jit
+from mindspore.common._decorator import deprecated
+from mindspore.common import mutable
+
+
+class _FirstGrad(Cell):
+    def __init__(self, fn):
+        super(_FirstGrad, self).__init__()
+        self.first_grad_op = C.GradOperation(sens_param=True, get_all=True)
+        self.fn = fn
+
+    def construct(self, u, first_grad_input):
+        return self.first_grad_op(self.fn)(*first_grad_input, u)
+
+
+class _JvpFirstGrad(Cell):
+    def __init__(self):
+        super(_JvpFirstGrad, self).__init__()
+        self.first_grad_op = C.GradOperation(sens_param=True, get_all=True)
+
+    def construct(self, u, fn, first_grad_input):
+        return self.first_grad_op(fn)(*first_grad_input, u)
+
+
+class _FirstGradSingleValue(Cell):
+    def __init__(self, fn):
+        super(_FirstGradSingleValue, self).__init__()
+        self.first_grad_single_value_op = C.GradOperation(sens_param=True)
+        self.fn = fn
+
+    def construct(self, u, first_grad_single_value_input):
+        return self.first_grad_single_value_op(self.fn)(*first_grad_single_value_input, u)
+
+
+class _JvpFirstGradSingleValue(Cell):
+    def __init__(self):
+        super(_JvpFirstGradSingleValue, self).__init__()
+        self.first_grad_single_value_op = C.GradOperation(sens_param=True)
+
+    def construct(self, u, fn, first_grad_single_value_input):
+        return self.first_grad_single_value_op(fn)(*first_grad_single_value_input, u)
+
+
+class Jvp(Cell):
+    """
+    Jvp will be deprecated in the future, please use :func:`mindspore.ops.jvp` instead.
+
+    Supported Platforms:
+        Deprecated
+    """
+
+    @deprecated("1.9", "ops.jvp", True)
+    def __init__(self, fn):
+        super(Jvp, self).__init__()
+        self.fn = fn
+        self.oneslike = P.OnesLike()
+        self.first_grad = _FirstGrad(fn)
+        self.first_grad.add_flags(enable_tuple_grad_first=True)
+        self.first_grad_single_value = _FirstGradSingleValue(fn)
+        self.first_grad_single_value.add_flags(enable_tuple_grad_first=True)
+        self.second_grad_op = C.GradOperation(sens_param=True)
+        self.issubclass_ = inner.IsSubClass()
+        self.typeof = Primitive('typeof')
+        self.make_tuple = Primitive('MakeTuple')
+
+    @jit
+    def construct(self, *args):
+        """construct for jvp."""
+        jvp_input = args[0:-1]
+        v = args[-1]
+        output = self.fn(*jvp_input)
+
+        if self.issubclass_(self.typeof(output), mstype.tuple_):
+            u = self.make_tuple()
+            for _, element in enumerate(output):
+                u = u + self.make_tuple(mutable(self.oneslike(element)))
+        else:
+            u = mutable(self.oneslike(output))
+
+        if len(jvp_input) == 1:
+            second_gradient_net = self.second_grad_op(self.first_grad_single_value)
+            gradient_output = second_gradient_net(u, jvp_input, v)
+        else:
+            second_gradient_net = self.second_grad_op(self.first_grad)
+            gradient_output = second_gradient_net(u, jvp_input, v)
+        return output, gradient_output
+
+
+class _JvpInner(Cell):
+    """
+    Compute the jacobian-vector-product of the given network. Jvp is equivalent to forward mode autodiff.
+    This class implements the inner process of function jvp.
+    """
+
+    def __init__(self):
+        super(_JvpInner, self).__init__()
+        self.oneslike = P.OnesLike()
+        self.first_grad = _JvpFirstGrad()
+        self.first_grad.add_flags(enable_tuple_grad_first=True)
+        self.first_grad_single_value = _JvpFirstGradSingleValue()
+        self.first_grad_single_value.add_flags(enable_tuple_grad_first=True)
+        self.second_grad_op = C.GradOperation(sens_param=True)
+        self.issubclass_ = inner.IsSubClass()
+        self.typeof = Primitive('typeof')
+        self.make_tuple = Primitive('MakeTuple')
+
+    def compute_jvp(self, fn, v, jvp_input, output):
+        """Compute the jacobian-vector-product of the given fn, vector, inputs and outputs."""
+        if self.issubclass_(self.typeof(output), mstype.tuple_):
+            u = self.make_tuple()
+            for _, element in enumerate(output):
+                u = u + self.make_tuple(mutable(self.oneslike(element)))
+        else:
+            u = mutable(self.oneslike(output))
+
+        if len(jvp_input) == 1:
+            second_gradient_net = self.second_grad_op(self.first_grad_single_value)
+            gradient_output = second_gradient_net(u, fn, jvp_input, v)
+        else:
+            second_gradient_net = self.second_grad_op(self.first_grad)
+            gradient_output = second_gradient_net(u, fn, jvp_input, v)
+        return gradient_output
+
+    def construct(self, *args):
+        fn = args[0]
+        v = args[1]
+        jvp_input = args[2:]
+        output = fn(*jvp_input)
+
+        gradient_output = self.compute_jvp(fn, v, jvp_input, output)
+        return output, gradient_output
+
+
+class _LinearizeInner(_JvpInner):
+    """
+    Compute the jacobian-vector-product of the given network. This Class is mainly useful
+    if you want to apply jvp multiple times under the same input and fn.
+    """
+    def construct(self, *args):
+        fn = args[0]
+        v = args[1]
+        output = args[2]
+        jvp_input = args[3]
+        gradient_output = self.compute_jvp(fn, v, jvp_input, output)
+        return gradient_output
+
+
+class Vjp(Cell):
+    """
+    Vjp will be deprecated in the future, please use :func:`mindspore.ops.vjp` instead.
+
+    Supported Platforms:
+        Deprecated
+    """
+
+    @deprecated("1.9", "ops.vjp", True)
+    def __init__(self, fn):
+        super(Vjp, self).__init__()
+        self.fn = fn
+        self.grad = C.GradOperation(get_all=True, sens_param=True)
+        self.grad_single_value = C.GradOperation(sens_param=True)
+        self.issubclass_ = inner.IsSubClass()
+        self.typeof = Primitive('typeof')
+
+    @jit
+    def construct(self, *args):
+        front_input = args[0:-1]
+        output = self.fn(*front_input)
+        if len(front_input) == 1:
+            gradient_output = self.grad_single_value(self.fn)(*args)
+        else:
+            gradient_output = self.grad(self.fn)(*args)
+        return output, gradient_output

mindspore/nn/layer/__init__.py

@@ -0,0 +1,63 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+Layer.
+
+The high-level components(Cells) used to construct the neural network.
+"""
+from __future__ import absolute_import
+
+from mindspore.nn.layer import activation, normalization, container, conv, basic, embedding, pooling, \
+    image, math, combined, timedistributed, thor_layer, rnns, rnn_cells, padding, dense, transformer, \
+    channel_shuffle
+from mindspore.nn.layer.activation import *
+from mindspore.nn.layer.normalization import *
+from mindspore.nn.layer.container import *
+from mindspore.nn.layer.conv import *
+from mindspore.nn.layer.dense import *
+from mindspore.nn.layer.rnns import *
+from mindspore.nn.layer.rnn_cells import *
+from mindspore.nn.layer.basic import *
+from mindspore.nn.layer.embedding import *
+from mindspore.nn.layer.pooling import *
+from mindspore.nn.layer.image import *
+from mindspore.nn.layer.math import *
+from mindspore.nn.layer.combined import *
+from mindspore.nn.layer.timedistributed import *
+from mindspore.nn.layer.transformer import *
+from mindspore.nn.layer.channel_shuffle import ChannelShuffle
+from mindspore.nn.layer.thor_layer import DenseThor, Conv2dThor, EmbeddingThor, EmbeddingLookupThor
+from mindspore.nn.layer.padding import ConstantPad1d, ConstantPad2d, ConstantPad3d, ReflectionPad1d, \
+    ReflectionPad2d, ReflectionPad3d, ZeroPad2d, ReplicationPad1d, ReplicationPad2d, ReplicationPad3d
+
+__all__ = []
+__all__.extend(activation.__all__)
+__all__.extend(normalization.__all__)
+__all__.extend(container.__all__)
+__all__.extend(conv.__all__)
+__all__.extend(dense.__all__)
+__all__.extend(rnn_cells.__all__)
+__all__.extend(rnns.__all__)
+__all__.extend(basic.__all__)
+__all__.extend(embedding.__all__)
+__all__.extend(pooling.__all__)
+__all__.extend(image.__all__)
+__all__.extend(math.__all__)
+__all__.extend(combined.__all__)
+__all__.extend(timedistributed.__all__)
+__all__.extend(transformer.__all__)
+__all__.extend(thor_layer.__all__)
+__all__.extend(padding.__all__)
+__all__.extend(channel_shuffle.__all__)