mindspore 2.3.0__cp39-cp39-win_amd64.whl → 2.4.1__cp39-cp39-win_amd64.whl

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.

mindspore/nn/layer/transformer.py

@@ -16,6 +16,7 @@
 Transformer Cells module, include TransformerEncoderLayer, TransformerDecoderLayer,
 TransformerEncoder, TransformerDecoder, Transformer.
 """
+import copy
 import math
 from typing import Union, Optional
 import mindspore
@@ -31,7 +32,6 @@ from .basic import Dense, Dropout
 from .activation import ReLU, GELU
 from .normalization import LayerNorm
 from .container import CellList
-
 __all__ = ['MultiheadAttention', 'TransformerEncoderLayer', 'TransformerDecoderLayer',
            'TransformerEncoder', 'TransformerDecoder', 'Transformer']
 
@@ -588,7 +588,7 @@ class TransformerEncoder(Cell):
                                          encoder_layer.dropout_num, encoder_layer.activation1,
                                          encoder_layer.layernorm_eps, encoder_layer.batch_first,
                                          encoder_layer.norm_first, dtype=encoder_layer.dtype)
-        self.layers = CellList([layers for _ in range(num_layers)])
+        self.layers = CellList([copy.deepcopy(layers) for _ in range(num_layers)])
         self.num_layers = num_layers
         self.norm = norm
 
@@ -663,7 +663,7 @@ class TransformerDecoder(Cell):
                                          decoder_layer.dropout_num, decoder_layer.activation1,
                                          decoder_layer.layernorm_eps, decoder_layer.batch_first,
                                          decoder_layer.norm_first, dtype=decoder_layer.dtype)
-        self.layers = CellList([layers for _ in range(num_layers)])
+        self.layers = CellList([copy.deepcopy(layers) for _ in range(num_layers)])
         self.num_layers = num_layers
         self.norm = norm
 

mindspore/nn/loss/__init__.py

@@ -25,7 +25,7 @@ from mindspore.nn.loss.loss import LossBase, L1Loss, CTCLoss, MSELoss, SmoothL1L
     SampledSoftmaxLoss, TripletMarginWithDistanceLoss,\
     PoissonNLLLoss, MultiLabelSoftMarginLoss, DiceLoss, BCEWithLogitsLoss, MultiClassDiceLoss, \
     RMSELoss, MAELoss, HuberLoss, CrossEntropyLoss, NLLLoss, KLDivLoss, MarginRankingLoss, GaussianNLLLoss, \
-    HingeEmbeddingLoss, MultilabelMarginLoss, TripletMarginLoss
+    HingeEmbeddingLoss, MultilabelMarginLoss, TripletMarginLoss, L1LossExt
 
 
 __all__ = ['LossBase', 'L1Loss', 'CTCLoss', 'MSELoss', 'SmoothL1Loss', 'SoftMarginLoss', 'FocalLoss',
@@ -33,4 +33,4 @@ __all__ = ['LossBase', 'L1Loss', 'CTCLoss', 'MSELoss', 'SmoothL1Loss', 'SoftMarg
            'CosineEmbeddingLoss', 'SampledSoftmaxLoss', 'TripletMarginWithDistanceLoss', 'PoissonNLLLoss',
            'MultiLabelSoftMarginLoss', 'DiceLoss', 'MultiClassDiceLoss', 'MultilabelMarginLoss',
            'RMSELoss', 'MAELoss', 'HuberLoss', 'CrossEntropyLoss', 'NLLLoss', 'KLDivLoss', 'MarginRankingLoss',
-           'GaussianNLLLoss', 'HingeEmbeddingLoss', 'TripletMarginLoss']
+           'GaussianNLLLoss', 'HingeEmbeddingLoss', 'TripletMarginLoss', 'L1LossExt']

mindspore/nn/loss/loss.py

@@ -33,6 +33,7 @@ from mindspore.nn.cell import Cell
 from mindspore.nn.layer.activation import get_activation
 from mindspore import _checkparam as validator
 from mindspore import context
+from mindspore.ops.auto_generate import l1_loss_ext_op
 
 
 class LossBase(Cell):
@@ -247,6 +248,80 @@ class L1Loss(LossBase):
         return F.l1_loss(logits, labels, self.reduction)
 
 
+class L1LossExt(LossBase):
+    r"""
+    L1Loss is used to calculate the mean absolute error between the predicted value and the target value.
+
+    Assuming that the :math:`x` and :math:`y` are 1-D Tensor, length :math:`N`, then calculate the loss of :math:`x` and
+    :math:`y` without dimensionality reduction (the reduction parameter is set to ``'none'`` ). The formula is as
+    follows:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad \text{with } l_n = \left| x_n - y_n \right|,
+
+    where :math:`N` is the batch size. If `reduction` is not ``'none'`` , then:
+
+    .. math::
+        \ell(x, y) =
+        \begin{cases}
+            \operatorname{mean}(L), & \text{if reduction} = \text{'mean';}\\
+            \operatorname{sum}(L),  & \text{if reduction} = \text{'sum'.}
+        \end{cases}
+
+    Args:
+        reduction (str, optional): Apply specific reduction method to the output: ``'none'`` , ``'mean'`` ,
+            ``'sum'`` . Default: ``'mean'`` .
+
+            - ``'none'``: no reduction will be applied.
+            - ``'mean'``: compute and return the mean of elements in the output.
+            - ``'sum'``: the output elements will be summed.
+
+    Inputs:
+        - **logits** (Tensor) - Predicted value, Tensor of any dimension.
+        - **labels** (Tensor) - Target value, same shape as the `logits` in common cases.
+          However, it supports the shape of `logits` is different from the shape of `labels`
+          and they should be broadcasted to each other.
+
+    Outputs:
+        Tensor, data type is float.
+
+    Raises:
+        ValueError: If `reduction` is not one of ``'none'`` , ``'mean'`` or ``'sum'`` .
+        ValueError: If `logits` and `labels` have different shapes and cannot be broadcasted to each other.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, nn
+        >>> import numpy as np
+        >>> # Case 1: logits.shape = labels.shape = (3,)
+        >>> loss = nn.L1LossExt()
+        >>> logits = Tensor(np.array([1, 2, 3]), mindspore.float32)
+        >>> labels = Tensor(np.array([1, 2, 2]), mindspore.float32)
+        >>> output = loss(logits, labels)
+        >>> print(output)
+        0.33333334
+        >>> # Case 2: logits.shape = (3,), labels.shape = (2, 3)
+        >>> loss = nn.L1LossExt(reduction='none')
+        >>> logits = Tensor(np.array([1, 2, 3]), mindspore.float32)
+        >>> labels = Tensor(np.array([[1, 1, 1], [1, 2, 2]]), mindspore.float32)
+        >>> output = loss(logits, labels)
+        >>> print(output)
+        [[0. 1. 2.]
+         [0. 0. 1.]]
+    """
+
+    def __init__(self, reduction='mean'):
+        """Initialize L1LossExt."""
+        super(L1LossExt, self).__init__(reduction)
+        self.reduction = reduction
+
+    def construct(self, logits, labels):
+        return l1_loss_ext_op(logits, labels, self.reduction)
+
+
 class MSELoss(LossBase):
     r"""
     Calculates the mean squared error between the predicted value and the label value.
@@ -287,6 +362,7 @@ class MSELoss(LossBase):
     Raises:
         ValueError: If `reduction` is not one of ``'none'``, ``'mean'`` or ``'sum'``.
         ValueError: If `logits` and `labels` have different shapes and cannot be broadcasted.
+        TypeError: if `logits` and `labels` have different data types.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -1580,7 +1656,7 @@ class BCELoss(LossBase):
     The formula is as follow:
 
     .. math::
-        L = \{l_1,\dots,l_N\}^\top, \quad
+        L = \{l_1,\dots,l_n,\dots,l_N\}^\top, \quad
         l_n = - w_n \left[ y_n \cdot \log x_n + (1 - y_n) \cdot \log (1 - x_n) \right]
 
     where N is the batch size. Then,
@@ -1850,14 +1926,16 @@ class BCEWithLogitsLoss(LossBase):
 
         weight (Tensor, optional): A rescaling weight applied to the loss of each batch element.
             If not None, it can be broadcast to a tensor with shape of `input`,
-            data type must be float16 or float32. Default: ``None`` .
+            data type must be float16, float32 or bfloat16(only Atlas A2 series products are supported).
+            Default: ``None`` .
         pos_weight (Tensor, optional): A weight of positive examples. Must be a vector with length equal to the
             number of classes. If not None, it must be broadcast to a tensor with shape of `input`, data type
-            must be float16 or float32. Default: ``None`` .
+            must be float16, float32 or bfloat16(only Atlas A2 series products are supported). Default: ``None`` .
 
     Inputs:
         - **input** (Tensor) - Input `input` with shape :math:`(N, *)` where :math:`*` means, any number
-          of additional dimensions. The data type must be float16 or float32.
+          of additional dimensions. The data type must be float16, float32 or bfloat16(only Atlas A2 series products
+          are supported).
         - **target** (Tensor) - Ground truth label with shape :math:`(N, *)` where :math:`*` means, any number
           of additional dimensions. The same shape and data type as `input`.
 
@@ -1867,9 +1945,9 @@ class BCEWithLogitsLoss(LossBase):
 
     Raises:
         TypeError: If input `input` or `target` is not Tensor.
-        TypeError: If data type of `input` or `target` is neither float16 nor float32.
+        TypeError: If data type of `input` or `target` is not float16, float32 or bfloat16.
         TypeError: If `weight` or `pos_weight` is a parameter.
-        TypeError: If data type of `weight` or `pos_weight` is neither float16 nor float32.
+        TypeError: If data type of `weight` or `pos_weight` is not float16 , float32 or bfloat16.
         TypeError: If data type of `reduction` is not string.
         ValueError: If `weight` or `pos_weight` can not be broadcast to a tensor with shape of `input`.
         ValueError: If `reduction` is not one of ``'none'``, ``'mean'``, ``'sum'``.

mindspore/nn/optim/__init__.py

@@ -38,7 +38,8 @@ from mindspore.nn.optim.adafactor import AdaFactor
 from mindspore.nn.optim.adasum import AdaSumByDeltaWeightWrapCell, AdaSumByGradWrapCell
 from mindspore.nn.optim.adamax import AdaMax
 from mindspore.nn.optim.adadelta import Adadelta
+from mindspore.nn.optim.tft_wrapper import OptTFTWrapper
 
 __all__ = ['Optimizer', 'Momentum', 'LARS', 'Adam', 'AdamWeightDecay', 'LazyAdam', 'AdamOffload',
            'Lamb', 'SGD', 'ASGD', 'Rprop', 'FTRL', 'RMSProp', 'ProximalAdagrad', 'Adagrad', 'thor', 'AdaFactor',
-           'AdaSumByDeltaWeightWrapCell', 'AdaSumByGradWrapCell', 'AdaMax', 'Adadelta']
+           'AdaSumByDeltaWeightWrapCell', 'AdaSumByGradWrapCell', 'AdaMax', 'Adadelta', 'OptTFTWrapper']

mindspore/nn/optim/adadelta.py

@@ -55,7 +55,7 @@ class Adadelta(Optimizer):
             w_{t} = w_{t-1} - \gamma * update_{t}
         \end{array}
 
-    where :math:`g` represents `grads`, :math:`\gamma` represents `learning_rate`, :math:`p` represents `rho`,
+    where :math:`g` represents `grads`, :math:`\gamma` represents `learning_rate`, :math:`\rho` represents `rho`,
     :math:`\epsilon` represents `epsilon`, :math:`w` represents `params`,
     :math:`accum` represents accumulation, :math:`accum\_update` represents accumulation update,
     :math:`t` represents current step.

mindspore/nn/optim/adam.py

@@ -906,7 +906,7 @@ class AdamWeightDecay(Optimizer):
         There is usually no connection between a optimizer and mixed precision. But when `FixedLossScaleManager` is used
         and `drop_overflow_update` in `FixedLossScaleManager` is set to False, optimizer needs to set the 'loss_scale'.
         As this optimizer has no argument of `loss_scale`, so `loss_scale` needs to be processed by other means, refer
-        document `LossScale <https://www.mindspore.cn/tutorials/en/master/advanced/mixed_precision.html>`_ to
+        document `LossScale <https://www.mindspore.cn/tutorials/en/master/beginner/mixed_precision.html>`_ to
         process `loss_scale` correctly.
 
         If parameters are not grouped, the `weight_decay` in optimizer will be applied on the network parameters without

mindspore/nn/optim/lamb.py

@@ -132,7 +132,7 @@ class Lamb(Optimizer):
         There is usually no connection between a optimizer and mixed precision. But when `FixedLossScaleManager` is used
         and `drop_overflow_update` in `FixedLossScaleManager` is set to False, optimizer needs to set the 'loss_scale'.
         As this optimizer has no argument of `loss_scale`, so `loss_scale` needs to be processed by other means. Refer
-        document `LossScale <https://www.mindspore.cn/tutorials/en/master/advanced/mixed_precision.html>`_ to
+        document `LossScale <https://www.mindspore.cn/tutorials/en/master/beginner/mixed_precision.html>`_ to
         process `loss_scale` correctly.
 
         If parameters are not grouped, the `weight_decay` in optimizer will be applied on the network parameters without

mindspore/nn/optim/tft_wrapper.py

@@ -0,0 +1,124 @@
+# 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.
+# ============================================================================
+"""OptTFTWrapper"""
+from __future__ import absolute_import
+
+import os
+from mindspore.common.tensor import Tensor
+from mindspore.nn.optim.optimizer import Optimizer
+from mindspore.ops.operations.manually_defined._inner import TensorReport
+from mindspore import ops, context
+from mindspore.common.parameter import Parameter
+import mindspore.common.dtype as mstype
+
+class OptTFTWrapper(Optimizer):
+    r"""
+    Implements TFT optimizer wrapper, this wrapper is used to report status to MindIO TFT before optimizer updating.
+
+    Note:
+        This optimizer is depend on MindIO TFT feature. Currently only support ascend graph mode and
+        sink_size must be less than 1.
+
+    Args:
+        opt (Optimizer): Must be sub-class of Optimizer.
+
+    Inputs:
+        - **gradients** (tuple[Tensor]) - The gradients of opt's `params`, the shape is the same as opt's `params`.
+
+    Outputs:
+        Tensor, result of executing optimizer 'opt'.
+
+    Raises:
+        TypeError: If the parameter opt is not an subclass of Optimizer.
+        ValueError: If the platform is not Ascend graph mode, or customer doesn't switch on TFT feature.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore import nn
+        >>>
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> #1) All parameters use the same learning rate and weight decay
+        >>> optim = nn.SGD(params=net.trainable_params())
+        >>> optim_wrapper = nn.OptTFTWrapper(optim)
+        >>>
+        >>> loss = nn.SoftmaxCrossEntropyWithLogits()
+        >>> model = ms.train.Model(net, loss_fn=loss, optimizer=optim)
+    """
+
+    def __init__(self, opt, **kwargs):
+        if not isinstance(opt, Optimizer):
+            raise TypeError(f"For 'OptTFTWrapper', the argument 'opt' must be Optimizer type, " f"but got {type(opt)}.")
+        super(OptTFTWrapper, self).__init__(opt.learning_rate, opt._parameters) # pylint: disable=W0212
+        tft_env = os.getenv("MS_ENABLE_TFT", "")
+        if ("TTP:1" not in tft_env) and ("UCE:1" not in tft_env):
+            raise ValueError("MindIO TFT regitster need custom switch on[MS_ENABLE_TFT='{TTP:1,UCE:1}']!")
+        mode = context.get_context("mode")
+        device_target = context.get_context("device_target")
+        if device_target != "Ascend" or mode != context.GRAPH_MODE:
+            raise ValueError("MindIO adataper only support on Ascend device with GRAPH Mode!")
+        self.opt = opt
+        self.report = TensorReport()
+        self.depend = ops.Depend()
+        self.allreduce_sum = ops.AllReduce()
+        self.allreduce_sum.add_prim_attr("tft_report_before", True)
+        self.tft_g_one_flag = Parameter(Tensor([1], dtype=mstype.int32))
+
+        self.param_rank = opt.param_rank
+        self.optim_filter = opt.optim_filter
+        self.loss_scale = opt.loss_scale
+        self.dynamic_weight_decay = opt.dynamic_weight_decay
+        self.grad_centralization = opt.grad_centralization
+
+        self.dynamic_lr = opt.dynamic_lr
+        self.global_step = opt.global_step
+        self.is_group = opt.is_group
+        self.is_group_lr = opt.is_group_lr
+        self.is_group_params_ordered = opt.is_group_params_ordered
+        self.use_parallel = opt.use_parallel
+        if self.is_group:
+            self.group_params = opt.group_params
+            self.group_lr = opt.group_lr
+            self.group_weight_decay = opt.group_weight_decay
+            self.group_grad_centralization = opt.group_grad_centralization
+            self.grad_centralization_flags = opt.grad_centralization_flags
+
+        self.skip_auto_parallel_compile = opt.skip_auto_parallel_compile
+
+        self.learning_rate = opt.learning_rate
+        self.parameters = opt.parameters
+        self.decay_flags = opt.decay_flags
+        self.dynamic_decay_flags = opt.dynamic_decay_flags
+        self.weight_decay = opt.weight_decay
+        self.exec_weight_decay = opt.exec_weight_decay
+        self.ps_parameters = opt.ps_parameters
+        self.cache_enable = opt.cache_enable
+        self.reciprocal_scale = opt.reciprocal_scale
+        self.need_scale = opt.need_scale
+        self.global_step_increase_tensor = opt.global_step_increase_tensor
+        self.param_length = opt.param_length
+        self.enable_tuple_broaden = opt.enable_tuple_broaden
+
+    def construct(self, gradients):
+        tft_g_one_flag = self.depend(self.tft_g_one_flag, gradients)
+        self.tft_g_one_flag = self.allreduce_sum(tft_g_one_flag)
+
+        grads = self.depend(gradients, self.report("tft_report", self.tft_g_one_flag))
+        opt_ret = self.opt(grads)
+        return opt_ret