mindspore 2.3.0__cp310-cp310-win_amd64.whl → 2.4.0__cp310-cp310-win_amd64.whl

mindspore/nn/layer/embedding.py

@@ -16,13 +16,14 @@
 from __future__ import absolute_import
 
 import mindspore.common.dtype as mstype
+import mindspore.ops as ops
 from mindspore import log as logger
 from mindspore.common.tensor import Tensor
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
 from mindspore.common.parameter import Parameter
 from mindspore.common.parameter import _get_unique_parameter_key
-from mindspore.common.initializer import initializer
+from mindspore.common.initializer import initializer, Normal
 from mindspore.communication.management import get_group_size, get_rank
 from mindspore.context import ParallelMode
 from mindspore.parallel._utils import _get_parallel_mode, _get_full_batch
@@ -34,7 +35,7 @@ from mindspore.ops.primitive import constexpr, _primexpr
 from mindspore.nn.layer.basic import ClipByNorm
 from mindspore.nn.cell import Cell
 
-__all__ = ['Embedding', 'EmbeddingLookup', 'MultiFieldEmbeddingLookup']
+__all__ = ['Embedding', 'EmbeddingExt', 'EmbeddingLookup', 'MultiFieldEmbeddingLookup']
 
 
 @_primexpr
@@ -161,6 +162,123 @@ class Embedding(Cell):
             f'embedding_table={self.embedding_table}, dtype={self.dtype}, padding_idx={self.padding_idx}'
 
 
+class EmbeddingExt(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.mint.nn.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}'
+
+
 @_primexpr
 def _make_axis_range(start, end):
     axis = tuple(range(start, end))

mindspore/nn/layer/normalization.py

@@ -19,7 +19,8 @@ from __future__ import division
 import itertools
 import numbers
 import hashlib
-
+import numpy as np
+import mindspore.ops as ops
 from mindspore.ops import operations as P
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.common.parameter import Parameter
@@ -37,7 +38,7 @@ from mindspore.nn.cell import Cell
 from mindspore import log as logger
 from mindspore.ops import group_norm
 
-__all__ = ['BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d', 'LayerNorm', 'GroupNorm',
+__all__ = ['BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d', 'LayerNorm', 'LayerNormExt', 'GroupNorm',
            'SyncBatchNorm', 'InstanceNorm1d', 'InstanceNorm2d', 'InstanceNorm3d']
 
 
@@ -558,14 +559,14 @@ class SyncBatchNorm(_BatchNorm):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/master/parallel/rank_table.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/rank_table.html>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/master/parallel/mpirun.html>`_ .
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/mpirun.html>`_ .
 
             For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/master/parallel/dynamic_cluster.html>`_ .
+            Startup <https://www.mindspore.cn/docs/en/master/model_train/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with multiple devices.
 
@@ -774,6 +775,101 @@ class LayerNorm(Cell):
             self.normalized_shape, self.begin_norm_axis, self.begin_params_axis, self.gamma, self.beta)
 
 
+class LayerNormExt(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` is the scale value learned through training and :math:`\beta` is the shift value.
+    It can be described using the following formula:
+
+    .. math::
+        y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        normalized_shape (Union(tuple[int], list[int], int)): The normalized shape of `x` for LayerNorm
+        eps (float): A value added to the denominator for numerical stability(:math:`\epsilon`). Default: ``1e-5`` .
+        elementwise_affine (bool): Whether affine transformation is required. When this parameter is set to ``True``,
+            the weight parameter is initialized to 1 and the offset is initialized to 0. Default: ``True``.
+        bias (bool): If set to ``False``, the layer will not learn an additive bias (only relevant if
+             `elementwise_affine` is ``True``). Default: ``True``.
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``None`` .
+
+    Inputs:
+        - **x** (Tensor) - The shape is :math:`(N, *)`, where :math:`*` is equal to normalized_shape.
+
+    Outputs:
+        Tensor, the normalized and scaled offset tensor, has the same shape and data type as the `x`.
+
+    Raises:
+        TypeError: If `eps` 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.LayerNormExt(shape1)
+        >>> output = m(x).shape
+        >>> print(output)
+        (20, 5, 10, 10)
+    """
+
+    def __init__(self,
+                 normalized_shape,
+                 eps=1e-5,
+                 elementwise_affine=True,
+                 bias=True,
+                 dtype=None
+                 ):
+        """Initialize LayerNormExt."""
+        super(LayerNormExt, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            # mypy error: incompatible types in assignment
+            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
+        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 = tuple(normalized_shape)
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        ms_dtype = mstype.float32 if dtype is None else dtype
+        if self.elementwise_affine:
+            self.weight = Parameter(Tensor(np.ones(normalized_shape), ms_dtype), name="weight")
+            if bias:
+                self.bias = Parameter(Tensor(np.zeros(normalized_shape), ms_dtype), name="bias")
+            else:
+                self.bias = None
+        else:
+            self.weight = None
+            self.bias = None
+
+    def construct(self, input):
+        y = ops.layer_norm(input, self.normalized_shape, self.weight,
+                           self.bias, self.eps)
+        return y
+
+    def extend_repr(self):
+        return 'normalized_shape={}, eps={}, elementwise_affine={}'.format(
+            self.normalized_shape, self.eps, self.elementwise_affine)
+
+
 class _InstanceNorm(Cell):
     """Instance Normalization base class."""
     @cell_attr_register

mindspore/nn/layer/padding.py

@@ -279,20 +279,6 @@ class ConstantPad1d(_ConstantPadNd):
            [0.5 1.  1.  1.  1.  0.5]]]]
         >>> print(out.shape)
         (1, 2, 3, 6)
-        >>> # padding is negative
-        >>> padding = (-1, 0)
-        >>> value = 0.5
-        >>> pad1d = ms.nn.ConstantPad1d(padding, value)
-        >>> out = pad1d(x)
-        >>> print(out)
-        [[[[1. 1. 1.]
-           [1. 1. 1.]
-           [1. 1. 1.]]
-          [[1. 1. 1.]
-           [1. 1. 1.]
-           [1. 1. 1.]]]]
-        >>> print(out.shape)
-        (1, 2, 3, 3)
     """
 
     def __init__(self, padding, value):
@@ -336,21 +322,21 @@ class ConstantPad2d(_ConstantPadNd):
         >>> import mindspore as ms
         >>> x = np.ones(shape=(1, 2, 3, 4)).astype(np.float32)
         >>> x = ms.Tensor(x)
-        >>> padding = (-1, 1, 0, 1)
+        >>> padding = (1, 1, 0, 1)
         >>> value = 0.5
         >>> pad2d = ms.nn.ConstantPad2d(padding, value)
         >>> out = pad2d(x)
         >>> print(out)
-        [[[[1.  1.  1.  0.5]
-           [1.  1.  1.  0.5]
-           [1.  1.  1.  0.5]
-           [0.5 0.5 0.5 0.5]]
-          [[1.  1.  1.  0.5]
-           [1.  1.  1.  0.5]
-           [1.  1.  1.  0.5]
-           [0.5 0.5 0.5 0.5]]]]
+        [[[[0.5  1.  1.  1.  1.  0.5]
+           [0.5  1.  1.  1.  1.  0.5]
+           [0.5  1.  1.  1.  1.  0.5]
+           [0.5  0.5 0.5 0.5 0.5 0.5]]
+          [[0.5  1.  1.  1.  1.  0.5]
+           [0.5  1.  1.  1.  1.  0.5]
+           [0.5  1.  1.  1.  1.  0.5]
+           [0.5  0.5 0.5 0.5 0.5 0.5]]]]
         >>> print(out.shape)
-        (1, 2, 4, 4)
+        (1, 2, 4, 6)
     """
 
     def __init__(self, padding, value):
@@ -396,25 +382,25 @@ class ConstantPad3d(_ConstantPadNd):
         >>> import mindspore as ms
         >>> x = np.ones(shape=(1, 2, 3, 4)).astype(np.float32)
         >>> x = ms.Tensor(x)
-        >>> padding = (-1, 1, 0, 1, 1, 0)
+        >>> padding = (1, 1, 0, 1, 1, 0)
         >>> value = 0.5
         >>> pad3d = ms.nn.ConstantPad3d(padding, value)
         >>> out = pad3d(x)
         >>> print(out)
-        [[[[0.5 0.5 0.5 0.5]
-           [0.5 0.5 0.5 0.5]
-           [0.5 0.5 0.5 0.5]
-           [0.5 0.5 0.5 0.5]]
-          [[1.  1.  1.  0.5]
-           [1.  1.  1.  0.5]
-           [1.  1.  1.  0.5]
-           [0.5 0.5 0.5 0.5]]
-          [[1.  1.  1.  0.5]
-           [1.  1.  1.  0.5]
-           [1.  1.  1.  0.5]
-           [0.5 0.5 0.5 0.5]]]]
+        [[[[0.5 0.5 0.5 0.5 0.5 0.5]
+           [0.5 0.5 0.5 0.5 0.5 0.5]
+           [0.5 0.5 0.5 0.5 0.5 0.5]
+           [0.5 0.5 0.5 0.5 0.5 0.5]]
+          [[0.5 1.  1.  1.  1.  0.5]
+           [0.5 1.  1.  1.  1.  0.5]
+           [0.5 1.  1.  1.  1.  0.5]
+           [0.5 0.5 0.5 0.5 0.5 0.5]]
+          [[0.5 1.  1.  1.  1.  0.5]
+           [0.5 1.  1.  1.  1.  0.5]
+           [0.5 1.  1.  1.  1.  0.5]
+           [0.5 0.5 0.5 0.5 0.5 0.5]]]]
         >>> print(out.shape)
-        (1, 3, 4, 4)
+        (1, 3, 4, 6)
     """
 
     def __init__(self, padding, value):
@@ -652,20 +638,20 @@ class ZeroPad2d(_ConstantPadNd):
         >>> import mindspore as ms
         >>> x = np.ones(shape=(1, 2, 3, 4)).astype(np.float32)
         >>> x = ms.Tensor(x)
-        >>> padding = (-1, 1, 0, 1)
+        >>> padding = (1, 1, 0, 1)
         >>> pad = ms.nn.ZeroPad2d(padding)
         >>> out = pad(x)
         >>> print(out)
-        [[[[1. 1. 1. 0.]
-           [1. 1. 1. 0.]
-           [1. 1. 1. 0.]
-           [0. 0. 0. 0.]]
-          [[1. 1. 1. 0.]
-           [1. 1. 1. 0.]
-           [1. 1. 1. 0.]
-           [0. 0. 0. 0.]]]]
+        [[[[0. 1. 1. 1. 1. 0.]
+           [0. 1. 1. 1. 1. 0.]
+           [0. 1. 1. 1. 1. 0.]
+           [0. 0. 0. 0. 0. 0.]]
+          [[0. 1. 1. 1. 1. 0.]
+           [0. 1. 1. 1. 1. 0.]
+           [0. 1. 1. 1. 1. 0.]
+           [0. 0. 0. 0. 0. 0.]]]]
         >>> print(out.shape)
-        (1, 2, 4, 4)
+        (1, 2, 4, 6)
     """
 
     def __init__(self, padding):

mindspore/nn/layer/pooling.py

@@ -18,6 +18,7 @@ from __future__ import absolute_import
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
 import mindspore.ops as ops
+from mindspore.ops.function.nn_func import avg_pool2d_ext
 from mindspore._checkparam import _check_3d_int_or_tuple
 from mindspore import _checkparam as validator
 from mindspore.ops.primitive import constexpr, _primexpr
@@ -26,13 +27,14 @@ import mindspore.context as context
 from mindspore.common import dtype as mstype
 from mindspore.ops.operations.nn_ops import AdaptiveMaxPool2D
 from mindspore.ops.operations.nn_ops import AdaptiveMaxPool3D, AdaptiveAvgPool3D
+from mindspore.ops.auto_generate.gen_ops_prim import MaxPoolWithIndices, MaxPoolWithMask
 from mindspore.nn.cell import Cell
 from mindspore._c_expression import MSContext
 
 __all__ = ['AvgPool3d', 'MaxPool3d', 'AvgPool2d', 'MaxPool2d', 'AvgPool1d', 'MaxPool1d', 'FractionalMaxPool2d',
            'FractionalMaxPool3d', 'AdaptiveAvgPool1d', 'AdaptiveMaxPool1d', 'AdaptiveMaxPool2d', 'AdaptiveMaxPool3d',
            'AdaptiveAvgPool2d', 'AdaptiveAvgPool3d', 'MaxUnpool1d', 'MaxUnpool2d', 'MaxUnpool3d', 'LPPool1d',
-           'LPPool2d']
+           'LPPool2d', 'AvgPool2dExt', 'MaxPool2dExt']
 
 
 class _PoolNd(Cell):
@@ -96,6 +98,9 @@ class LPPool1d(Cell):
     .. math::
         f(X) = \sqrt[p]{\sum_{x \in X} x^{p}}
 
+    Note:
+        This interface currently does not support Atlas A2 training series products.
+
     Args:
         norm_type (Union[int, float]): Type of normalization, represents :math:`p` in the formula, can not be 0.
 
@@ -169,6 +174,9 @@ class LPPool2d(Cell):
     .. math::
         f(X) = \sqrt[p]{\sum_{x \in X} x^{p}}
 
+    Note:
+        This interface currently does not support Atlas A2 training series products.
+
     Args:
         norm_type(Union[int, float]): Type of normalization, represents :math:`p` in the formula, can not be 0.
 
@@ -374,6 +382,7 @@ class MaxPool3d(_PoolNd):
     Examples:
         >>> import mindspore as ms
         >>> import mindspore.nn as nn
+        >>> from mindspore import Tensor
         >>> import numpy as np
         >>> np_x = np.random.randint(0, 10, [5, 3, 4, 6, 7])
         >>> x = Tensor(np_x, ms.float32)
@@ -592,6 +601,102 @@ class MaxPool2d(_PoolNd):
         return out
 
 
+class MaxPool2dExt(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.mint.nn.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(MaxPool2dExt, 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
+
+
 class MaxPool1d(_PoolNd):
     r"""
     Applies a 1D max pooling over an input Tensor which can be regarded as a composition of 1D planes.
@@ -793,6 +898,9 @@ class AvgPool3d(_PoolNd):
         \frac{1}{d_{ker} * h_{ker} * w_{ker}} \sum_{l=0}^{d_{ker}-1} \sum_{m=0}^{h_{ker}-1} \sum_{n=0}^{w_{ker}-1}
         \text{input}(N_i, C_j, s_0 \times d + l, s_1 \times h + m, s_2 \times w + n)
 
+    Note:
+        This interface currently does not support Atlas A2 training series products.
+
     Args:
         kernel_size (Union[int, tuple[int]], optional): The size of kernel used to take the average value,
             can be an int number or a single element tuple that represents depth, height and width, or a tuple of three
@@ -910,6 +1018,46 @@ class AvgPool3d(_PoolNd):
         return out
 
 
+class AvgPool2dExt(Cell):
+    r"""
+    Applies a 2D average pooling over an input Tensor which can be regarded as
+    a composition of 2D input planes.
+
+    For details, please refer to :func:`mindspore.mint.nn.functional.avg_pool2d`.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import numpy as np
+        >>> from mindspore import Tensor, nn
+        >>> from mindspore import dtype as mstype
+        >>> x = Tensor(np.arange(1 * 3 * 3 * 4).reshape(1, 3, 3, 4), mstype.float32)
+        >>> m =  nn.AvgPool2dExt(x, kernel_size=2, stride=1)
+        >>> output = m(x)
+        >>> print(output)
+        [[[[ 2.5   3.5   4.5]
+           [ 6.5   7.5   8.5]]
+          [[14.5  15.5  16.5]
+           [18.5  19.5  20.5]]
+          [[26.5  27.5  28.5]
+           [30.5  31.5  32.5]]]]
+    """
+    def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False,
+                 count_include_pad=True, divisor_override=None):
+        super(AvgPool2dExt, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.ceil_mode = ceil_mode
+        self.count_include_pad = count_include_pad
+        self.divisor_override = divisor_override
+
+    def construct(self, input):
+        return avg_pool2d_ext(input, self.kernel_size, self.stride, self.padding,
+                              self.ceil_mode, self.count_include_pad, self.divisor_override)
+
+
 class AvgPool2d(_PoolNd):
     r"""
     Applies a 2D average pooling over an input Tensor which can be regarded as a composition of 2D input planes.
@@ -922,6 +1070,9 @@ class AvgPool2d(_PoolNd):
         \text{output}(N_i, C_j, h, w) = \frac{1}{h_{ker} * w_{ker}} \sum_{m=0}^{h_{ker}-1} \sum_{n=0}^{w_{ker}-1}
         \text{input}(N_i, C_j, s_0 \times h + m, s_1 \times w + n)
 
+    Note:
+        This interface currently does not support Atlas A2 training series products.
+
     Args:
         kernel_size (Union[int, tuple[int]]): The size of kernel used to take the average value.
             The data type of kernel_size must be int or a single element tuple and the value represents the height
@@ -1015,12 +1166,12 @@ class AvgPool2d(_PoolNd):
                  data_format="NCHW"):
         """Initialize AvgPool2d."""
         super(AvgPool2d, self).__init__(kernel_size, stride, pad_mode, data_format)
-        self.ascend_910bc_target = (MSContext.get_instance().get_ascend_soc_version() in ['ascend910b', 'ascend910c'])
+        self.ascend_910b_target = (MSContext.get_instance().get_ascend_soc_version() in ['ascend910b', 'ascend910_93'])
         if pad_mode.upper() == 'PAD' or padding != 0 or ceil_mode or not count_include_pad \
                 or divisor_override is not None:
-            if self.ascend_910bc_target:
-                raise ValueError(f"For '{self.cls_name}, the pad_mod 'PAD' is not support in 910B now, "
-                                 f"it will be supported in the future.")
+            if self.ascend_910b_target:
+                raise ValueError(f"For '{self.cls_name}, the pad_mod 'PAD' is not support in Ascend910B or Ascend910_93"
+                                 f" now, it will be supported in the future.")
             if self.format == "NHWC":
                 raise ValueError(f"For '{self.cls_name}, the 'NHWC' format are not support when 'pad_mode' is 'pad' or "
                                  f"'padding' is not 0 or 'ceil_mode' is not False or 'count_include_pad' is not True"
@@ -1083,6 +1234,9 @@ class AvgPool1d(_PoolNd):
         \text{output}(N_i, C_j, l) = \frac{1}{l_{ker}} \sum_{n=0}^{l_{ker}-1}
         \text{input}(N_i, C_j, s_0 \times l + n)
 
+    Note:
+        This interface currently does not support Atlas A2 training series products.
+
     Args:
         kernel_size (int): The size of kernel window used to take the average value, Default: ``1`` .
         stride (int): The distance of kernel moving, an int number that represents
@@ -1682,7 +1836,7 @@ class AdaptiveMaxPool3d(Cell):
 
 class FractionalMaxPool2d(Cell):
     r"""
-    Applies the 2D FractionalMaxPool operatin over input. The output Tensor shape can be determined by either
+    Applies the 2D FractionalMaxPool operation over input. The output Tensor shape can be determined by either
     `output_size` or `output_ratio`, and the step size is determined by `_random_samples`. `output_size` will take
     effect when `output_size` and `output_ratio` are set at the same time.
     And `output_size` and `output_ratio` can not be ``None`` at the same time.
@@ -1783,7 +1937,7 @@ class FractionalMaxPool2d(Cell):
 
 class FractionalMaxPool3d(Cell):
     r"""
-    Applies the 3D FractionalMaxPool operatin over `input`. The output Tensor shape can be determined by either
+    Applies the 3D FractionalMaxPool operation over `input`. The output Tensor shape can be determined by either
     `output_size` or `output_ratio`, and the step size is determined by `_random_samples`. `output_size` will take
     effect when `output_size` and `output_ratio` are set at the same time.
     And `output_size` and `output_ratio` can not be ``None`` at the same time.