mindspore 2.3.0__cp310-cp310-win_amd64.whl

mindspore/ops/extend/__init__.py

@@ -0,0 +1,53 @@
+# 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.
+# ============================================================================
+
+"""
+
+Operators with better performance
+
+"""
+
+from __future__ import absolute_import
+
+from mindspore.common import Tensor
+from mindspore.ops.primitive import Primitive, PrimitiveWithInfer, PrimitiveWithCheck, prim_attr_register
+from mindspore.ops.vm_impl_registry import get_vm_impl_fn, vm_impl_registry
+from mindspore.ops.op_info_register import op_info_register, custom_info_register, AkgGpuRegOp, AkgAscendRegOp, \
+    AiCPURegOp, TBERegOp, CpuRegOp, CustomRegOp, DataType
+from mindspore.ops.primitive import constexpr
+from . import (
+    array_func,
+    math_func,
+    nn_func,
+)
+
+from .array_func import gather, max, min, one_hot
+from .math_func import (
+    baddbmm,
+    bmm,
+    add,
+    sub
+)
+
+from .nn_func import (
+    conv2d,
+    max_pool2d,
+    leaky_relu_ext
+)
+
+__all__ = []
+__all__.extend(array_func.__all__)
+__all__.extend(math_func.__all__)
+__all__.extend(nn_func.__all__)

mindspore/ops/extend/array_func.py

@@ -0,0 +1,218 @@
+# 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.
+# ============================================================================
+
+"""
+
+Array Operators
+
+"""
+from mindspore.common import Tensor
+from mindspore.ops.operations.array_ops import ArgMaxWithValue, ArgMinWithValue
+from mindspore.ops._primitive_cache import _get_cache_prim
+from mindspore.ops.auto_generate.gen_ops_prim import gather_d_op
+from mindspore.ops.auto_generate.gen_ops_def import max_, min_
+from mindspore.ops.auto_generate.pyboost_inner_prim import _PyboostOneHotExtPrim
+one_hot_ext_impl = _PyboostOneHotExtPrim()
+
+# define Primitive global variables
+
+def gather(input, dim, index):
+    r"""
+    Gather data from a tensor by indices.
+
+    .. math::
+        output[(i_0, i_1, ..., i_{dim}, i_{dim+1}, ..., i_n)] =
+        input[(i_0, i_1, ..., index[(i_0, i_1, ..., i_{dim}, i_{dim+1}, ..., i_n)], i_{dim+1}, ..., i_n)]
+
+    .. warning::
+        On Ascend, the behavior is unpredictable in the following cases:
+
+        - the value of `index` is not in the range `[-input.shape[dim], input.shape[dim])` in forward;
+        - the value of `index` is not in the range `[0, input.shape[dim])` in backward.
+
+    Args:
+        input (Tensor): The target tensor to gather values.
+        dim (int): the axis to index along, must be in range `[-input.rank, input.rank)`.
+        index (Tensor): The index tensor, with int32 or int64 data type. An valid `index` should be:
+
+            - `index.rank == input.rank`;
+            - for `axis != dim`, `index.shape[axis] <= input.shape[axis]`;
+            - the value of `index` is in range `[-input.shape[dim], input.shape[dim])`.
+
+    Returns:
+        Tensor, has the same type as `input` and the same shape as `index`.
+
+    Raises:
+        ValueError: If the shape of `index` is illegal.
+        ValueError: If `dim` is not in `[-input.rank, input.rank)`.
+        ValueError: If the value of `index` is out of the valid range.
+        TypeError: If the type of `index` is illegal.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, ops
+        >>> input = Tensor(np.array([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]), mindspore.float32)
+        >>> index = Tensor(np.array([[0, 0], [1, 1]]), mindspore.int32)
+        >>> output = ops.extend.gather(input, 1, index)
+        >>> print(output)
+        [[-0.1 -0.1]
+         [0.5   0.5]]
+    """
+    return gather_d_op(input, dim, index)
+
+
+def max(input, dim=None, keepdim=False):
+    """
+    Calculates the maximum value along with the given dimension for the input tensor.
+
+    Args:
+        input (Tensor): The input tensor, can be any dimension. Complex tensor is not supported for now.
+        dim (int, optional): The dimension to reduce. Default: ``None`` .
+        keepdim (bool, optional): Whether to reduce dimension, if true, the output will keep same dimension
+            with the input, the output will reduce dimension if false. Default: ``False`` .
+
+    Returns:
+        Tensor if `dim` is the default value ``None`` , the maximum value of input tensor, with the shape :math:`()` ,
+        and same dtype as `input`.
+
+        tuple (Tensor) if `dim` is not the default value ``None`` , tuple of 2 tensors, containing the maximum
+        value of the input tensor along the given dimension `dim` and the corresponding index.
+
+        - **values (Tensor)** - The maximum value of input tensor along the given dimension `dim`, with same dtype as
+          `input`. If `keepdim` is ``True`` , the shape of output tensors is :math:`(input_1, input_2, ...,
+          input_{axis-1}, 1, input_{axis+1}, ..., input_N)` . Otherwise, the shape is :math:`(input_1, input_2, ...,
+          input_{axis-1}, input_{axis+1}, ..., input_N)` .
+        - **index (Tensor)** - The index for the maximum value of the input tensor along the given dimension `dim`, with
+          the same shape as `values`.
+
+    Raises:
+        ValueError: If `dim` is the default value ``None`` and `keepdim` is not ``False`` .
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, ops
+        >>> y = Tensor(np.array([[0.0, 0.3, 0.4, 0.5, 0.1],
+        ...                      [3.2, 0.4, 0.1, 2.9, 4.0]]), mindspore.float32)
+        >>> output, index = ops.extend.max(y, 0, True)
+        >>> print(output, index)
+        [[3.2 0.4 0.4 2.9 4. ]] [[1 1 0 1 1]]
+    """
+    if dim is None:
+        if keepdim is not False:
+            raise ValueError(f"For 'max', the `keepdim` must be False when the `dim` is None, but got {keepdim}")
+        return max_(input)
+    argmax_with_value_op = _get_cache_prim(ArgMaxWithValue)(dim, keepdim)
+    indices, values = argmax_with_value_op(input)
+    return values, indices
+
+
+def min(input, dim=None, keepdim=False):
+    """
+    Calculates the minimum value along with the given dimension for the input tensor.
+
+    Args:
+        input (Tensor): The input tensor, can be any dimension. Complex tensor is not supported for now.
+        dim (int, optional): The dimension to reduce. Default: ``None`` .
+        keepdim (bool, optional): Whether to reduce dimension, if true, the output will keep same dimension
+            with the input, the output will reduce dimension if false. Default: ``False`` .
+
+    Returns:
+        Tensor if `dim` is the default value ``None`` , the minimum value of input tensor, with the shape :math:`()` ,
+        and same dtype as `input`.
+
+        tuple (Tensor) if `dim` is not the default value ``None`` , tuple of 2 tensors, containing the minimum value
+        of the input tensor along the given dimension `dim` and the corresponding index.
+
+        - **values (Tensor)** - The minimum value of input tensor along the given dimension `dim`, with same dtype as
+          `input`. If `keepdim` is ``True`` , the shape of output tensors is :math:`(input_1, input_2, ...,
+          input_{axis-1}, 1, input_{axis+1}, ..., input_N)` . Otherwise, the shape is :math:`(input_1, input_2, ...,
+          input_{axis-1}, input_{axis+1}, ..., input_N)` .
+        - **index (Tensor)** - The index for the minimum value of the input tensor along the given dimension `dim`,
+          with the same shape as `values`.
+
+    Raises:
+        ValueError: If `dim` is the default value ``None`` and `keepdim` is not ``False`` .
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, ops
+        >>> x = Tensor(np.array([0.0, 0.4, 0.6, 0.7, 0.1]), mindspore.float32)
+        >>> output, index = ops.extend.min(x, 0, keepdim=True)
+        >>> print(output, index)
+        [0.0] [0]
+    """
+    if dim is None:
+        if keepdim is not False:
+            raise ValueError(f"For 'min', the `keepdim` must be False when the `dim` is None, but got {keepdim}")
+        return min_(input)
+    argmin_with_value_op = _get_cache_prim(ArgMinWithValue)(dim, keepdim)
+    indices, values = argmin_with_value_op(input)
+    return values, indices
+
+
+def one_hot(tensor, num_classes):
+    r"""
+    Computes a one-hot tensor.
+
+    The locations represented by tensor in `tensor` take value `1`, while all
+    other locations take value `0`.
+
+    Args:
+        tensor (Tensor): A tensor of indices. Tensor of shape :math:`(X_0, \ldots, X_n)`.
+            Data type must be int32 or int64.
+        num_classes (int): A scalar defining the depth of the one-hot dimension.
+
+    Returns:
+        Tensor, one-hot tensor.
+
+    Raises:
+        TypeError: If `num_classes` is not an int.
+        TypeError: If dtype of `tensor` is not int32 or int64.
+        ValueError: If `num_classes` is less than 0.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import ops
+        >>> from mindspore import Tensor
+        >>> tensor = Tensor(np.array([0, 1, 2]), mindspore.int32)
+        >>> num_classes = 3
+        >>> output = ops.extend.one_hot(tensor, num_classes)
+        >>> print(output)
+        [[1. 0. 0.]
+        [0. 1. 0.]
+        [0. 0. 1.]]
+    """
+    on_value = Tensor(1, dtype=tensor.dtype)
+    off_value = Tensor(0, dtype=tensor.dtype)
+    return one_hot_ext_impl(tensor, num_classes, on_value, off_value, -1)
+
+
+__all__ = ['gather', 'max', 'min', 'one_hot']

mindspore/ops/extend/math_func.py

@@ -0,0 +1,76 @@
+# Copyright 2023 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.
+# ============================================================================
+
+"""
+
+Math Operators with better performance
+
+"""
+
+from mindspore.ops import auto_generate as P
+from mindspore.ops.auto_generate.gen_ops_def import add_ext as add, sub_ext as sub, bmm_ext as bmm
+
+
+# define Primitive global variables
+
+
+def baddbmm(input, batch1, batch2, beta=1, alpha=1):
+    r"""
+    The result is the sum of the input and a batch matrix-matrix product of matrices in batch1 and batch2.
+    The formula is defined as follows:
+
+    .. math::
+        \text{out}_{i} = \beta \text{input}_{i} + \alpha (\text{batch1}_{i} \mathbin{@} \text{batch2}_{i})
+
+    Args:
+        input (Tensor): The input Tensor. When batch1 is a :math:`(C, W, T)` Tensor and batch2 is a
+            :math:`(C, T, H)` Tensor, input must be broadcastable with :math:`(C, W, H)` Tensor.
+        batch1 (Tensor): :math:`batch1` in the above formula. Must be 3-D Tensor, dtype is same as input.
+        batch2 (Tensor): :math:`batch2` in the above formula. Must be 3-D Tensor, dtype is same as input.
+        beta (Union[float, int], optional): multiplier for input. Default: ``1`` .
+        alpha (Union[float, int], optional): multiplier for :math:`batch1 @ batch2`. Default: ``1`` .
+            Arguments beta and alpha must be integers when inputs of type not FloatTensor, otherwise they should
+            be a real number.
+
+    Returns:
+        Tensor, has the same dtype as input, shape will be :math:`(C, W, H)`.
+
+    Raises:
+        TypeError: The type of `input`, `batch1`, `batch2` is not Tensor.
+        TypeError: The types of `input`, `batch1`, `batch2` are different.
+        TypeError: For inputs of type FloatTensor or DoubleTensor, \
+                    arguments beta and alpha not be real numbers, otherwise not be integers.
+        TypeError: For Baddbmm, attributes alpha and beta are not real numbers
+        ValueError: If `batch1` and `batch2` are not 3-D tensors.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> from mindspore import Tensor, ops
+        >>> input = Tensor(np.ones([1, 3, 3]).astype(np.float32))
+        >>> batch1 = Tensor(np.ones([1, 3, 4]).astype(np.float32))
+        >>> batch2 = Tensor(np.ones([1, 4, 3]).astype(np.float32))
+        >>> output = ops.baddbmm(input, batch1, batch2)
+        >>> print(output)
+        [[[5. 5. 5.]
+          [5. 5. 5.]
+          [5. 5. 5.]]]
+    """
+    return P.baddbmm(input, batch1, batch2, beta, alpha)
+
+
+__all__ = ['baddbmm', 'add', 'sub', 'bmm']

mindspore/ops/extend/nn_func.py

@@ -0,0 +1,308 @@
+# Copyright 2023 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.
+# ============================================================================
+
+"""
+
+NN Operators with better performance
+
+"""
+from mindspore.ops._primitive_cache import _get_cache_prim
+from mindspore.ops.auto_generate.gen_ops_prim import Convolution, ConstantPadND, MaxPoolWithIndices, MaxPoolWithMask
+from mindspore.ops.auto_generate import leaky_relu_ext
+from mindspore import _checkparam as validator
+
+
+def _check_stride_when_same_mode(stride):
+    """ stride must be 1 when pad mode is same """
+    if isinstance(stride, int):
+        if stride != 1:
+            raise ValueError(f"For conv2d, 'padding=same' is not supported for stride convolution, " \
+                                f"but got {stride}")
+    elif isinstance(stride, tuple):
+        validator.check_int(len(stride), 2, validator.EQ, "stride", 'conv2d')
+        if not all(s == 1 for s in stride):
+            raise ValueError(f"For conv2d, 'padding=same' is not supported for stride convolution, " \
+                                f"but got {stride}")
+    else:
+        raise TypeError(f"For conv2d, the parameter 'stride' must be a int/tuple, but got {type(stride)}")
+
+
+def _get_pad_info(dilation, weight):
+    """ Get pad list by dilation and weight shape """
+    need_pad_nd = False
+    pad_l = ()
+    pad_r = ()
+    for i in range(2):
+        d = dilation[i]
+        weight_size = weight.shape[i + 2]
+        pad = d * (weight_size - 1)
+        pad_l += (int(pad / 2),)
+        pad_r += (int(pad - pad_l[i]),)
+        if pad_l[i] != pad_r[i]:
+            need_pad_nd = True
+    return need_pad_nd, pad_l, pad_r
+
+
+def _get_pad_nd_info(pad_l, pad_r):
+    """ Get pad_nd list if input need to exec pad_nd """
+    pad_nd = ()
+    new_pad_l = ()
+    for i in range(2):
+        delta_pad = pad_r[i] - pad_l[i]
+        if delta_pad > 0:
+            pad_nd = (0, delta_pad,) + pad_nd
+            new_pad_l += (pad_l[i],)
+        else:
+            pad_nd = (delta_pad, 0,) + pad_nd
+            new_pad_l += (pad_r[i],)
+    return pad_nd, new_pad_l
+
+
+def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
+    r"""
+    Applies a 2D convolution over an input tensor. The input tenor is typically of
+    shape :math:`(N, C_{in}, H_{in}, W_{in})`, where :math:`N` is batch size, :math:`C` is
+    channel number, :math:`H` is feature height, :math:`W` is feature width.
+
+    The output is calculated based on formula:
+
+    .. math::
+
+        \text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) +
+        \sum_{k = 0}^{C_{in} - 1} \text{ccor}({\text{weight}(C_{\text{out}_j}, k), \text{X}(N_i, k)})
+
+    where :math:`bias` is the output channel bias, :math:`ccor` is
+    the `cross-correlation <https://en.wikipedia.org/wiki/Cross-correlation>`_,
+    , :math:`weight` is the convolution kernel value and :math:`X` represents the input feature map.
+
+    Here are the indices' meanings:
+
+    - :math:`i` corresponds to the batch number, the range is :math:`[0, N-1]`,
+      where :math:`N` is the batch size of the input.
+
+    - :math:`j` corresponds to the output channel, the range is :math:`[0, C_{out}-1]`,
+      where :math:`C_{out}` is the number of output channels, which is also equal to the number of kernels.
+
+    - :math:`k` corresponds to the input channel, the range is :math:`[0, C_{in}-1]`,
+      where :math:`C_{in}` is the number of
+      input channels, which is also equal to the number of channels in the convolutional kernels.
+
+    Therefore, in the above formula, :math:`{bias}(C_{out_j})` represents the bias of the :math:`j`-th
+    output channel, :math:`{weight}(C_{out_j}, k)` represents the slice of the :math:`j`-th convolutional
+    kernel in the :math:`k`-th channel, and :math:`{X}(N_i, k)` represents the slice of the :math:`k`-th input
+    channel in the :math:`i`-th batch of the input feature map.
+
+    The shape of the convolutional kernel is given by :math:`(\text{kernel_size[0]}, \text{kernel_size[1]})`,
+    where :math:`\text{kernel_size[0]}` and :math:`\text{kernel_size[1]}` are the height and width of the kernel,
+    respectively.
+    If we consider the input and output channels as well as the `group` parameter, the complete kernel shape
+    will be :math:`(C_{out}, C_{in} / \text{group}, \text{kernel_size[0]}, \text{kernel_size[1]})`,
+    where `group` is the number of groups dividing `x`'s input channel when applying group convolution.
+
+    For more details about convolution layer, please refer to `Gradient Based Learning Applied to Document Recognition
+    <http://vision.stanford.edu/cs598_spring07/papers/Lecun98.pdf>`_ and
+    `ConvNets <http://cs231n.github.io/convolutional-networks/>`_.
+
+    Note:
+        On Ascend platform, only group convolution in depthwise convolution scenarios is supported.
+        That is, when `groups>1`, condition :math:`C_{in}` = :math:`C_{out}` = `groups` must be satisfied.
+
+    Args:
+        input (Tensor): Tensor of shape :math:`(N, C_{in}, H_{in}, W_{in})`.
+        weight (Tensor): Tensor of shape
+            :math:`(N, C_{in} / \text{groups}, \text{kernel_size[0]}, \text{kernel_size[1]})`, then the size of kernel
+            is :math:`(\text{kernel_size[0]}, \text{kernel_size[1]})`.
+        bias (Tensor, optional): Bias Tensor with shape :math:`(C_{out})`.
+            When bias is ``None`` , zeros will be used. Default: ``None`` .
+        stride (Union(int, tuple[int]), optional): The distance of kernel moving, an int number that represents
+            the height and width of movement are both strides, or a tuple of two int numbers that
+            represent height and width of movement respectively. Default: ``1`` .
+        padding (Union(int, tuple[int], list[int], str), optional): Implicit paddings on both sides of the input `x`.
+            Can be a string, one integer or a tuple/list with 2 integers.
+            If `padding` is a string, the optional values are ``"same"`` , ``"valid"``.
+
+            - same: Adopts the way of completion. The height and width of the output will be equal to
+              the input `x` divided by stride. The padding will be evenly calculated in top and bottom,
+              left and right possiblily. Otherwise, the last extra padding will be calculated from the bottom
+              and the right side. If this mode is set, `padding` must be 0.
+
+            - valid: Adopts the way of discarding. The possible largest height and width of output will be returned
+              without padding. Extra pixels will be discarded. If this mode is set, `padding` must be 0.
+
+            If `padding` is one integer, the paddings of top, bottom, left and right are the same, equal to padding.
+            If `padding` is a tuple/list with 2 integers, the padding of top adn bottom is padding[0],
+            and the padding of left and right is padding[1]. Default: ``0`` .
+        dilation (Union(int, tuple[int]), optional): Gaps between kernel elements.The data type is int or a tuple of
+            2 integers. Specifies the dilation rate to use for dilated convolution. If set to be :math:`k > 1`,
+            there will be :math:`k - 1` pixels skipped for each sampling location. Its value must
+            be greater than or equal to 1 and bounded by the height and width of the input `x`. Default: ``1`` .
+        groups (int, optional): Splits `input` into groups. Default: ``1`` .
+
+    Returns:
+        Tensor, the value that applied 2D convolution. The shape is :math:`(N, C_{out}, H_{out}, W_{out})`.
+        To see how different pad modes affect the output shape, please refer to
+        :class:`mindspore.nn.Conv2d` for more details.
+
+
+    Raises:
+        TypeError: If `stride`, `padding` or `dilation` is neither an int nor a tuple.
+        TypeError: `groups` is not an int.
+        TypeError: If `bias` is not a Tensor.
+        ValueError: If  the shape of `bias` is not :math:`(C_{out})` .
+        ValueError: If `stride` or `dilation` is less than 1.
+        ValueError: If `pad_mode` is not one of 'same', 'valid' or 'pad'.
+        ValueError: If `padding` is a tuple/list whose length is not equal to 2.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, ops
+        >>> x = Tensor(np.ones([10, 32, 32, 32]), mindspore.float32)
+        >>> weight = Tensor(np.ones([32, 32, 3, 3]), mindspore.float32)
+        >>> output = ops.extend.conv2d(x, weight)
+        >>> print(output.shape)
+        (10, 32, 30, 30)
+    """
+
+    def _convolution_same(input, weight, bias, dilation, groups):
+        """ convolution when mode is 'same' """
+        if isinstance(dilation, int):
+            dilation = (dilation,) * 2
+        validator.check_int(len(weight.shape), 4, validator.EQ, "weight.shape", 'conv2d')
+        validator.check_int(len(dilation), 2, validator.EQ, "dilation", 'conv2d')
+
+        # Calc padding info
+        need_pad_nd, pad_l, pad_r = _get_pad_info(dilation, weight)
+        if not need_pad_nd:
+            conv = _get_cache_prim(Convolution)(stride, pad_l, dilation, False, (0, 0), groups)
+            return conv(input, weight, bias)
+
+        # Calc pad nd info
+        pad_nd, pad_l = _get_pad_nd_info(pad_l, pad_r)
+        pad_nd_op = _get_cache_prim(ConstantPadND)()
+        padded_input = pad_nd_op(input, pad_nd, 0)
+        conv = _get_cache_prim(Convolution)(stride, pad_l, dilation, False, (0, 0), groups)
+        return conv(padded_input, weight, bias)
+
+    if isinstance(padding, int):
+        padding = (padding,) * 2
+
+    if isinstance(padding, (tuple, list)):
+        conv = _get_cache_prim(Convolution)(stride, padding, dilation, False, (0, 0), groups)
+        return conv(input, weight, bias)
+    if isinstance(padding, str):
+        if padding == 'valid':
+            conv = _get_cache_prim(Convolution)(stride, (0, 0), dilation, False, (0, 0), groups)
+            return conv(input, weight, bias)
+        if padding == 'same':
+            _check_stride_when_same_mode(stride)
+            return _convolution_same(input, weight, bias, dilation, groups)
+        raise ValueError(f"For conv2d, the parameter 'padding' must be 'same' or 'valid' when " \
+                            f"the type of 'padding' is string.")
+    raise TypeError(f"For conv2d, the parameter 'padding' must be a tuple/list " \
+                    f"or a string, but got {type(padding)}")
+
+
+def max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, *, ceil_mode=False, return_indices=False):
+    r"""
+    Performs a 2D max pooling on the input Tensor.
+
+    Typically, the input is a Tensor with shape :math:`(N_{in}, C_{in}, H_{in}, W_{in})`, outputs
+    regional maximum in the :math:`(H_{in}, W_{in})`-dimension. Given `kernel_size`
+    :math:`ks = (h_{ker}, w_{ker})` and `stride` :math:`s = (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 A2 training series.
+
+    Args:
+        input (Tensor): Tensor of shape :math:`(N_{in}, C_{in}, H_{in}, W_{in})` with data type of float32
+            in Ascend.
+        kernel_size (Union[int, tuple[int]]): The size of kernel used to take the maximum value and arg
+            value, is an int number that represents height and width of the kernel, or a tuple of
+            two int numbers that represent height and width respectively.
+        stride (Union[int, tuple[int], None]): The distance of kernel moving, an int number 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]]): An int number that represents the height and width of movement are both
+            strides, or a tuple of two int numbers that represent height and width of movement respectively.
+            Default: ``0`` .
+        dilation (Union[int, tuple[int]]): Control the stride of elements in the kernel. Default: ``1`` .
+        ceil_mode (bool): Whether to use ceil instead of floor to calculate output shape. Default: ``False`` .
+        return_indices (bool): Whether to output the indices of max value. Default: ``False`` .
+
+    Returns:
+        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`.
+
+        .. 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
+
+        - **argmax** (Tensor) - Index corresponding to the maximum value. In Ascend, data type is int32.
+          It will be return only when `return_indices` is True.
+
+    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
+        >>> import numpy as np
+        >>> from mindspore import Tensor, ops
+        >>> input = Tensor(np.arange(20 * 16 * 50 * 32).reshape((20, 16, 50, 32)), mindspore.float32)
+        >>> output_tensor, argmax = ops.extend.max_pool2d(input, kernel_size=(3, 2), stride=(2, 1),
+                                                          ceil_mode=False, return_indices=True)
+        >>> print(output_tensor.shape)
+        (20, 16, 24, 31)
+        >>> print(argmax.shape)
+        (20, 16, 24, 31)
+    """
+    strides = stride if (stride is not None) else kernel_size
+    if return_indices:
+        max_pool_func_ = _get_cache_prim(MaxPoolWithIndices)(kernel_size, strides, padding, dilation, ceil_mode)
+        out, indices = max_pool_func_(input)
+    else:
+        max_pool_func_ = _get_cache_prim(MaxPoolWithMask)(kernel_size, strides, padding, dilation, ceil_mode)
+        out, indices = max_pool_func_(input)
+    if return_indices:
+        return out, indices
+    return out
+
+
+__all__ = ['conv2d', 'max_pool2d', 'leaky_relu_ext']