mindspore 2.1.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.11__cp38-cp38-manylinux1_x86_64.whl

mindspore/nn/probability/distribution/transformed_distribution.py

@@ -16,6 +16,7 @@
 import numpy as np
 from mindspore import _checkparam as validator
 from mindspore.ops import operations as P
+from mindspore.ops import functional as F
 from mindspore.common import dtype as mstype
 import mindspore.nn as nn
 from .distribution import Distribution
@@ -125,7 +126,6 @@ class TransformedDistribution(Distribution):
         self.cast_base = P.Cast()
         self.equal_base = P.Equal()
         self.select_base = P.Select()
-        self.fill_base = P.Fill()
 
         # broadcast bijector batch_shape and distribution batch_shape
         self._broadcast_shape = self._broadcast_bijector_dist()
@@ -176,9 +176,9 @@ class TransformedDistribution(Distribution):
         """
         if self.batch_shape is None or self.bijector.batch_shape is None:
             return None
-        bijector_shape_tensor = self.fill_base(
+        bijector_shape_tensor = F.fill(
             self.dtype, self.bijector.batch_shape, 0.0)
-        dist_shape_tensor = self.fill_base(self.dtype, self.batch_shape, 0.0)
+        dist_shape_tensor = F.fill(self.dtype, self.batch_shape, 0.0)
         return (bijector_shape_tensor + dist_shape_tensor).shape
 
     def _cdf(self, value, *args, **kwargs):

mindspore/nn/probability/distribution/uniform.py

@@ -14,6 +14,7 @@
 # ============================================================================
 """Uniform Distribution"""
 import numpy as np
+from mindspore.ops import functional as F
 from mindspore.ops import operations as P
 from mindspore.ops import composite as C
 from mindspore import _checkparam as Validator
@@ -170,7 +171,6 @@ class Uniform(Distribution):
         self.cast = P.Cast()
         self.const = P.ScalarToTensor()
         self.dtypeop = P.DType()
-        self.fill = P.Fill()
         self.less = P.Less()
         self.lessequal = P.LessEqual()
         self.logicaland = P.LogicalAnd()
@@ -287,10 +287,10 @@ class Uniform(Distribution):
         value = self._check_value(value, 'value')
         value = self.cast(value, self.dtype)
         low, high = self._check_param_type(low, high)
-        neg_ones = self.fill(self.dtype, self.shape(value), -1.0)
+        neg_ones = F.fill(self.dtype, self.shape(value), -1.0)
         prob = self.exp(neg_ones * self.log(high - low))
         broadcast_shape = self.shape(prob)
-        zeros = self.fill(self.dtypeop(prob), broadcast_shape, 0.0)
+        zeros = F.fill(self.dtypeop(prob), broadcast_shape, 0.0)
         comp_lo = self.less(value, low)
         comp_hi = self.lessequal(value, high)
         less_than_low = self.select(comp_lo, zeros, prob)
@@ -316,7 +316,7 @@ class Uniform(Distribution):
         kl = self.log(high_b - low_b) - self.log(high_a - low_a)
         comp = self.logicaland(self.lessequal(
             low_b, low_a), self.lessequal(high_a, high_b))
-        inf = self.fill(self.dtypeop(kl), self.shape(kl), np.inf)
+        inf = F.fill(self.dtypeop(kl), self.shape(kl), np.inf)
         return self.select(comp, kl, inf)
 
     def _cdf(self, value, low=None, high=None):
@@ -338,8 +338,8 @@ class Uniform(Distribution):
         low, high = self._check_param_type(low, high)
         prob = (value - low) / (high - low)
         broadcast_shape = self.shape(prob)
-        zeros = self.fill(self.dtypeop(prob), broadcast_shape, 0.0)
-        ones = self.fill(self.dtypeop(prob), broadcast_shape, 1.0)
+        zeros = F.fill(self.dtypeop(prob), broadcast_shape, 0.0)
+        ones = F.fill(self.dtypeop(prob), broadcast_shape, 1.0)
         comp_lo = self.less(value, low)
         comp_hi = self.less(value, high)
         less_than_low = self.select(comp_lo, zeros, prob)

mindspore/nn/wrap/__init__.py

@@ -20,7 +20,8 @@ Use the Wrapper to combine the loss or build the training steps.
 from __future__ import absolute_import
 
 from mindspore.nn.wrap.cell_wrapper import ForwardValueAndGrad, TrainOneStepCell, WithLossCell, WithGradCell, \
-    WithEvalCell, ParameterUpdate, GetNextSingleOp, VirtualDatasetCellTriple, MicroBatchInterleaved, PipelineCell
+    WithEvalCell, ParameterUpdate, GetNextSingleOp, VirtualDatasetCellTriple, MicroBatchInterleaved, PipelineCell, \
+    GradAccumulationCell
 from mindspore.nn.wrap.loss_scale import TrainOneStepWithLossScaleCell,\
     DynamicLossScaleUpdateCell, FixedLossScaleUpdateCell
 from mindspore.nn.wrap.grad_reducer import DistributedGradReducer
@@ -42,5 +43,6 @@ __all__ = [
     "ParameterUpdate",
     "DynamicLossScaleUpdateCell",
     "FixedLossScaleUpdateCell",
-    "VirtualDatasetCellTriple"
+    "VirtualDatasetCellTriple",
+    "GradAccumulationCell"
     ]

mindspore/nn/wrap/cell_wrapper.py

@@ -99,7 +99,7 @@ class WithLossCell(Cell):
         >>> from mindspore import Tensor, nn
         >>> import numpy as np
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=False)
         >>> net_with_criterion = nn.WithLossCell(net, loss_fn)
@@ -115,8 +115,7 @@ class WithLossCell(Cell):
         super(WithLossCell, self).__init__(auto_prefix=False)
         self._backbone = backbone
         self._loss_fn = loss_fn
-        if isinstance(backbone, Cell) and backbone.jit_config_dict:
-            self._jit_config_dict = backbone.jit_config_dict
+        self._get_attr_from_cell(backbone)
 
     def construct(self, data, label):
         out = self._backbone(data)
@@ -133,7 +132,7 @@ class WithLossCell(Cell):
         Examples:
             >>> from mindspore import nn
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=False)
             >>> net_with_criterion = nn.WithLossCell(net, loss_fn)
@@ -176,7 +175,7 @@ class WithGradCell(Cell):
         >>> import mindspore as ms
         >>> from mindspore import nn
         >>> # Defined a network without loss function, taking LeNet5 as an example.
-        >>> # Refer to https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # Refer to https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits()
         >>> grad_net = nn.WithGradCell(net, loss_fn)
@@ -199,8 +198,7 @@ class WithGradCell(Cell):
         else:
             self.network_with_loss = WithLossCell(self.network, self.loss_fn)
         self.network_with_loss.set_train()
-        if isinstance(network, Cell) and network.jit_config_dict:
-            self._jit_config_dict = network.jit_config_dict
+        self._get_attr_from_cell(network)
 
     def construct(self, *inputs):
         weights = self.weights
@@ -219,7 +217,7 @@ class ForwardValueAndGrad(Cell):
     The backward graph will be created in the gradient function to calculating gradient.
 
     Args:
-        network (Cell): The training network.
+        network (Union[Cell, Function, MethodType]): The training network.
         weights (ParameterTuple): The parameters of the training network that need to calculate the gradient.
             Default: ``None`` .
         get_all (bool): If ``True`` , get all the gradients with respect to inputs. Default: ``False`` .
@@ -302,8 +300,7 @@ class ForwardValueAndGrad(Cell):
         self.get_by_list = get_by_list
         self.sens_param = sens_param
         self.grad = C.GradOperation(get_all=self.get_all, get_by_list=self.get_by_list, sens_param=self.sens_param)
-        if isinstance(network, Cell) and network.jit_config_dict:
-            self._jit_config_dict = network.jit_config_dict
+        self._get_attr_from_cell(network)
 
     def construct(self, *inputs):
         grad_inputs = inputs
@@ -349,7 +346,7 @@ class TrainOneStepCell(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits()
         >>> optim = nn.Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
@@ -414,8 +411,7 @@ class TrainOneStepCell(Cell):
                 create_group(server_group_name, group_list[current_index])
                 group = server_group_name
             self.grad_reducer = DistributedGradReducer(self.weights, self.mean, self.degree, group=group)
-        if isinstance(network, Cell) and network.jit_config_dict:
-            self._jit_config_dict = network.jit_config_dict
+        self._get_attr_from_cell(network)
 
     def construct(self, *inputs):
         if not self.sense_flag:
@@ -514,8 +510,7 @@ class _VirtualDatasetCell(Cell):
         super(_VirtualDatasetCell, self).__init__(auto_prefix=False)
         self._backbone = backbone
         self._virtual_dataset = _VirtualDataset()
-        if isinstance(backbone, Cell) and backbone.jit_config_dict:
-            self._jit_config_dict = backbone.jit_config_dict
+        self._get_attr_from_cell(backbone)
 
     def construct(self, *inputs):
         output = self._virtual_dataset(*inputs)
@@ -524,6 +519,8 @@ class _VirtualDatasetCell(Cell):
 
 @_primexpr
 def _check_shape_value_on_axis_divided_by_target_value(input_shape, micro_size):
+    if F.isconstant(input_shape[0]) is False:
+        return
     if input_shape[0] % micro_size != 0:
         raise ValueError(f"For micro batch initialization, the 0th dimension shape of input({input_shape[0]}) must be "
                          f"divided by micro size({micro_size})")
@@ -548,8 +545,8 @@ class _MicroBatch(Cell):
         for each_input in inputs:
             input_shape = self.shape(each_input)
             _check_shape_value_on_axis_divided_by_target_value(input_shape, self.micro_size)
-            micro_batch_begin = i * input_shape[0] // self.micro_size
-            micro_batch_end = (i + 1) * input_shape[0] // self.micro_size
+            micro_batch_begin = (input_shape[0] // self.micro_size) * i
+            micro_batch_end = (input_shape[0] // self.micro_size) * (i + 1)
             strided_slice_begin = (micro_batch_begin,)
             strided_slice_strides = (1,)
             for _ in range(len(input_shape) - 1):
@@ -589,7 +586,7 @@ class MicroBatchInterleaved(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> net = nn.MicroBatchInterleaved(net, 2)
     """
@@ -610,8 +607,7 @@ class MicroBatchInterleaved(Cell):
             interleave_data.strided_slice.add_prim_attr("strided_slice_flag", True)
             interleave_data.strided_slice.add_prim_attr("interleave_num", interleave_num)
             self.interleave_inputs.append(interleave_data)
-        if isinstance(network, Cell) and network.jit_config_dict:
-            self._jit_config_dict = network.jit_config_dict
+        self._get_attr_from_cell(network)
 
     def construct(self, *inputs):
         output = 0.0
@@ -638,7 +634,7 @@ class PipelineCell(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> net = nn.PipelineCell(net, 4)
     """
@@ -648,13 +644,73 @@ class PipelineCell(Cell):
         self.micro_inputs = nn.CellList()
         self.micro_size = micro_size
         self.add_list = []
+        if not isinstance(network, Cell):
+            raise TypeError("For 'PipelineCell', the argument 'network' must cell type, "
+                            "but got the type : {}.".format(type(network)))
+        if not isinstance(micro_size, int):
+            raise TypeError("For 'PipelineCell', the argument 'micro_size' must be integer, "
+                            "but got the type : {}.".format(type(micro_size)))
+        if micro_size <= 0:
+            raise ValueError("For 'PipelineCell', the argument 'micro_size' must be large than 0, "
+                             "but got {}.".format(micro_size))
         for i in range(micro_size):
             micro_input = _MicroBatch(micro_size)
             self.micro_inputs.append(micro_input)
             self.add = P.Add().add_prim_attr("pipeline_end", i)
             self.add_list.append(self.add)
-        if isinstance(network, Cell) and network.jit_config_dict:
-            self._jit_config_dict = network.jit_config_dict
+        self._get_attr_from_cell(network)
+
+    def construct(self, *inputs):
+        ret = None
+        for i in range(self.micro_size):
+            micro_input = self.micro_inputs[i](i, *inputs)
+            output = self.network(*micro_input)
+            if ret is not None:
+                ret = self.add_list[i](ret, output)
+            else:
+                ret = output
+        return ret
+
+class GradAccumulationCell(Cell):
+    """
+    Wrap the network with Micro Batch to enable the grad accumulation in semi_auto_parallel/auto_parallel mode.
+
+    Args:
+        network (Cell): The target network to wrap.
+        micro_size (int): MicroBatch size.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        >>> import mindspore.nn as nn
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> net = nn.GradAccumulationCell(net, 4)
+    """
+    def __init__(self, network, micro_size):
+        super(GradAccumulationCell, self).__init__(auto_prefix=False)
+        self.network = network
+        self.micro_inputs = nn.CellList()
+        self.micro_size = micro_size
+        self.add_list = []
+        if not isinstance(network, Cell):
+            raise TypeError("For 'GradAccumulationCell', the argument 'network' must cell type, "
+                            "but got the type : {}.".format(type(network)))
+        if not isinstance(micro_size, int):
+            raise TypeError("For 'GradAccumulationCell', the argument 'micro_size' must be integer, "
+                            "but got the type : {}.".format(type(micro_size)))
+        if micro_size <= 0:
+            raise ValueError("For 'GradAccumulationCell', the argument 'micro_size' must be large than 0, "
+                             "but got {}.".format(micro_size))
+        for i in range(micro_size):
+            micro_input = _MicroBatch(micro_size)
+            micro_input.strided_slice.add_prim_attr("grad_accu_num", micro_size)
+            self.micro_inputs.append(micro_input)
+            self.add = P.Add().add_prim_attr("forward_end", i)
+            self.add_list.append(self.add)
+        self._get_attr_from_cell(network)
 
     def construct(self, *inputs):
         ret = None
@@ -674,23 +730,22 @@ def _pipeline_clear_grad(accu_grad, grad):
     return F.assign(accu_grad, zeros)
 
 
-class _TrainPipelineAccuStepCell(TrainOneStepCell):
+class _TrainGradAccuStepCell(TrainOneStepCell):
     """
     Wraps the network with an optimizer in pipeline mode.
     """
     def __init__(self, network, optimizer, sens=None):
-        super(_TrainPipelineAccuStepCell, self).__init__(network, optimizer, sens)
+        super(_TrainGradAccuStepCell, self).__init__(network, optimizer, sens)
         self.accu_grads = self.weights.clone(prefix="accu_grads", init="zeros")
         self.hyper_map = ops.HyperMap()
         self.opt_shard = _get_enable_parallel_optimizer()
-        if isinstance(network, Cell) and network.jit_config_dict:
-            self._jit_config_dict = network.jit_config_dict
+        self._get_attr_from_cell(network)
 
     def construct(self, *inputs):
         if not self.sense_flag:
             return self._no_sens_impl(*inputs)
         loss = self.network(*inputs)
-        sens = ops.Fill()(ops.DType()(loss), ops.Shape()(loss), self.sens)
+        sens = ops.fill(ops.DType()(loss), ops.Shape()(loss), self.sens)
         grads = self.grad(self.network, self.weights)(*inputs, sens)
         accu_grads = ops.depend(self.accu_grads, grads)
         if self.opt_shard:
@@ -735,7 +790,7 @@ class VirtualDatasetCellTriple(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> net = nn.VirtualDatasetCellTriple(net)
     """
@@ -744,8 +799,7 @@ class VirtualDatasetCellTriple(Cell):
         super(VirtualDatasetCellTriple, self).__init__(auto_prefix=False)
         logger.warning("WARN_DEPRECATED: The usage of VirtualDatasetCellTriple is deprecated.")
         self._backbone = backbone
-        if isinstance(backbone, Cell) and backbone.jit_config_dict:
-            self._jit_config_dict = backbone.jit_config_dict
+        self._get_attr_from_cell(backbone)
 
     def construct(self, a, b, c):
         return self._backbone(a, b, c)
@@ -779,7 +833,7 @@ class WithEvalCell(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define a forward network without loss function, taking LeNet5 as an example.
-        >>> # Refer to https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # Refer to https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits()
         >>> eval_net = nn.WithEvalCell(net, loss_fn)
@@ -790,8 +844,7 @@ class WithEvalCell(Cell):
         self._network = network
         self._loss_fn = loss_fn
         self.add_cast_fp32 = validator.check_value_type("add_cast_fp32", add_cast_fp32, [bool], self.cls_name)
-        if isinstance(network, Cell) and network.jit_config_dict:
-            self._jit_config_dict = network.jit_config_dict
+        self._get_attr_from_cell(network)
 
     def construct(self, data, label):
         outputs = self._network(data)

mindspore/nn/wrap/grad_reducer.py

@@ -314,12 +314,15 @@ class DistributedGradReducer(Cell):
             Before running the following examples, you need to configure the communication environment variables.
 
             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/r2.1/parallel/train_ascend.html#preparations>`_
+            Please see the `rank table Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.1/parallel/train_gpu.html#preparation>`_ .
+            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/r2.2/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/r2.2/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with multiple devices.
 
@@ -356,7 +359,7 @@ class DistributedGradReducer(Cell):
         ...     def construct(self, *args):
         ...         weights = self.weights
         ...         loss = self.network(*args)
-        ...         sens = ops.Fill()(ops.DType()(loss), ops.Shape()(loss), self.sens)
+        ...         sens = F.fill(ops.DType()(loss), ops.Shape()(loss), self.sens)
         ...         grads = self.grad(self.network, weights)(*args, sens)
         ...         if self.reducer_flag:
         ...             # apply grad reducer on grads

mindspore/nn/wrap/loss_scale.py

@@ -15,6 +15,7 @@
 """Loss scale cell for loss scale training."""
 from __future__ import absolute_import
 
+import os
 import mindspore.context as context
 from mindspore.context import ParallelMode
 from mindspore.parallel._utils import _get_enable_parallel_optimizer
@@ -30,6 +31,7 @@ from mindspore.ops import composite as C
 from mindspore.ops import operations as P
 from mindspore.common import dtype as mstype
 from mindspore.common.api import jit
+from mindspore._c_expression import MSContext
 
 _grad_scale = C.MultitypeFuncGraph("grad_scale")
 reciprocal = P.Reciprocal()
@@ -60,6 +62,28 @@ def _tensor_grad_overflow_row_tensor(grad):
     return grad_overflow(grad.values)
 
 
+_ascend_grad_overflow = C.MultitypeFuncGraph("_ascend_grad_overflow")
+ascend_grad_overflow = P.IsFinite()
+
+
+@_ascend_grad_overflow.register("Tensor")
+def _tensor_ascend_grad_overflow(grad):
+    status = ascend_grad_overflow(grad)
+    base = Tensor(1.0, dtype=mstype.float32)
+    output = base - status.all()
+    output = P.Reshape()(output, ((1,)))
+    return output
+
+
+@_ascend_grad_overflow.register("RowTensor")
+def _tensor_ascend_grad_overflow_row_tensor(grad):
+    status = ascend_grad_overflow(grad.values)
+    base = Tensor(1.0, dtype=mstype.float32)
+    output = base - status.all()
+    output = P.Reshape()(output, ((1,)))
+    return output
+
+
 class DynamicLossScaleUpdateCell(Cell):
     r"""
     Dynamic Loss scale update cell.
@@ -296,16 +320,18 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         >>> size, in_features, out_features = 16, 16, 10
         >>> #1) when the type of scale_sense is Cell:
         >>> net = Net(in_features, out_features)
-        >>> loss = nn.MSELoss()
+        >>> loss_fn = nn.MSELoss()
         >>> optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
-        >>> net_with_loss = nn.WithLossCell(net, loss)
-        >>> manager = nn.DynamicLossScaleUpdateCell(loss_scale_value=2**12, scale_factor=2, scale_window=1000)
-        >>> train_network = nn.TrainOneStepWithLossScaleCell(net_with_loss, optimizer, scale_sense=manager)
+        >>> net_with_loss = nn.WithLossCell(net, loss_fn)
         >>> input = Tensor(np.ones([out_features, in_features]), mindspore.float32)
         >>> labels = Tensor(np.ones([out_features,]), mindspore.float32)
-        >>> output = train_network(input, labels)
-        >>> status, scaling_sens = train_network.start_overflow_check(loss, train_network.scaling_sens)
-        >>> grads = train_network.grad(train_network.network, weights)(*inputs, scaling_sens_filled)
+        >>> loss = net_with_loss(input, labels)
+        >>> manager = nn.DynamicLossScaleUpdateCell(loss_scale_value=2**12, scale_factor=2, scale_window=1000)
+        >>> train_network = nn.TrainOneStepWithLossScaleCell(net_with_loss, optimizer, scale_sense=manager)
+        >>> status = Tensor([0] * 8, mindspore.int32)
+        >>> scaling_sens = train_network.scale_sense
+        >>> scaling_sens_filled = ops.ones_like(loss) * ops.cast(scaling_sens, ops.dtype(loss))
+        >>> grads = train_network.grad(train_network.network, train_network.weights)(input, labels, scaling_sens_filled)
         >>> grads = train_network.grad_reducer(grads)
         >>> cond = train_network.get_overflow_status(status, grads)
         >>> overflow = train_network.process_loss_scale(cond)
@@ -341,7 +367,12 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         self.allreduce = P.AllReduce()
         self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
         self.gpu_target = (context.get_context("device_target") == "GPU")
+        self.ascend_910a_target = (MSContext.get_instance().get_ascend_soc_version() == 'ascend910')
+        self.ascend_910b_target = (MSContext.get_instance().get_ascend_soc_version() == 'ascend910b')
         self.loss_scaling_manager = None
+        self._ascend910b_check_overflow_status_mode = os.environ.get('MS_ASCEND_CHECK_OVERFLOW_MODE')
+
+
         if isinstance(scale_sense, Cell):
             self.loss_scaling_manager = scale_sense
             self.scale_sense = Parameter(Tensor(scale_sense.get_loss_scale(), dtype=mstype.float32),
@@ -358,6 +389,7 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
                             "the 'scale_sense' must be Cell or Tensor, but got 'scale_sense' type: {}."
                             .format(type(scale_sense)))
         self.enable_tuple_broaden = True
+        self._get_attr_from_cell(network)
 
     def construct(self, *inputs):
         weights = self.weights
@@ -418,13 +450,68 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
             is cleaned up when the function returns.
         """
         status = Tensor([0] * 8, mstype.int32)
-        if not self.gpu_target:
+        if self.ascend_910a_target or (self.ascend_910b_target and \
+                                       self._ascend910b_check_overflow_status_mode != "INFNAN_MODE"):
             status = F.depend(status, pre_cond)
             # clear overflow buffer
             clear_status = NPUClearFloatStatusV2()(status)
             compute_input = F.depend(compute_input, clear_status)
         return status, compute_input
 
+    def _check_overflow_status_on_infnan_mode(self, grad_overflow_check_func, compute_output):
+        """check overflow status on infnan mode."""
+        flag_sum = self.hyper_map(F.partial(grad_overflow_check_func), compute_output)
+        flag_sum = P.AddN()(flag_sum)
+        # convert flag_sum to scalar
+        flag_sum = P.Reshape()(flag_sum, (()))
+        return flag_sum
+
+    def _get_distributed_overflow_status_on_infnan_mode(self, grad_overflow_check_func, compute_output):
+        """converge the distributed overflow status on infnan mode."""
+        flag_sum = self._check_overflow_status_on_infnan_mode(grad_overflow_check_func, compute_output)
+
+        if self.is_distributed:
+            # sum overflow flag over devices
+            flag_reduce = self.allreduce(flag_sum)
+            overflow = self.less_equal(self.base, flag_reduce)
+        else:
+            overflow = self.less_equal(self.base, flag_sum)
+        return overflow
+
+    def _get_gpu_overflow_status(self, compute_output):
+        """get overflow status of gpu."""
+        overflow = self._get_distributed_overflow_status_on_infnan_mode(_grad_overflow, compute_output)
+        return overflow
+
+    def _get_ascend_overflow_status_on_infnan_mode(self, compute_output):
+        """get overflow status of ascend on infnan mode."""
+        overflow = self._get_distributed_overflow_status_on_infnan_mode(_ascend_grad_overflow, compute_output)
+        return overflow
+
+    def _get_ascend_overflow_status_on_saturation_mode(self, status, compute_output):
+        """get overflow status of ascend on saturation mode"""
+        status = F.depend(status, compute_output)
+        get_status = NPUGetFloatStatusV2()(status)
+
+        if self.is_distributed:
+            # sum overflow flag over devices
+            flag_reduce = self.allreduce(get_status)
+            # get_status not equal to [0]*8 means overflow
+            flag = self.equal(self.base0, flag_reduce)
+            status = F.depend(status, flag)
+            # distributed needs to skip allreduce to avoid its overflow affecting the next step
+            clear_status = NPUClearFloatStatusV2()(status)
+            flag = F.depend(flag, clear_status)
+            overall_finite = self.reduce_all(flag)
+        else:
+            status = F.depend(status, get_status)
+            clear_status = NPUClearFloatStatusV2()(status)
+            get_status = F.depend(get_status, clear_status)
+            flag = self.equal(self.base0, get_status)
+            overall_finite = self.reduce_all(flag)
+        overflow = self.logic_not(overall_finite)
+        return overflow
+
     @jit
     def get_overflow_status(self, status, compute_output):
         """
@@ -442,39 +529,15 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         Returns:
             bool, whether the overflow occurs or not.
         """
-        if not self.gpu_target:
-            status = F.depend(status, compute_output)
-            get_status = NPUGetFloatStatusV2()(status)
-
-            if self.is_distributed:
-                # sum overflow flag over devices
-                flag_reduce = self.allreduce(get_status)
-                # get_status not equal to [0]*8 means overflow
-                flag = self.equal(self.base0, flag_reduce)
-                status = F.depend(status, flag)
-                # distributed needs to skip allreduce to avoid its overflow affecting the next step
-                clear_status = NPUClearFloatStatusV2()(status)
-                flag = F.depend(flag, clear_status)
-                overall_finite = self.reduce_all(flag)
-            else:
-                status = F.depend(status, get_status)
-                clear_status = NPUClearFloatStatusV2()(status)
-                get_status = F.depend(get_status, clear_status)
-                flag = self.equal(self.base0, get_status)
-                overall_finite = self.reduce_all(flag)
-            overflow = self.logic_not(overall_finite)
-        else:
-            flag_sum = self.hyper_map(F.partial(_grad_overflow), compute_output)
-            flag_sum = P.AddN()(flag_sum)
-            # convert flag_sum to scalar
-            flag_sum = P.Reshape()(flag_sum, (()))
-
-            if self.is_distributed:
-                # sum overflow flag over devices
-                flag_reduce = self.allreduce(flag_sum)
-                overflow = self.less_equal(self.base, flag_reduce)
+        if self.gpu_target:
+            overflow = self._get_gpu_overflow_status(compute_output)
+        elif self.ascend_910b_target:
+            if self._ascend910b_check_overflow_status_mode != "INFNAN_MODE":
+                overflow = self._get_ascend_overflow_status_on_saturation_mode(status, compute_output)
             else:
-                overflow = self.less_equal(self.base, flag_sum)
+                overflow = self._get_ascend_overflow_status_on_infnan_mode(compute_output)
+        else:  # ascend_910a_target
+            overflow = self._get_ascend_overflow_status_on_saturation_mode(status, compute_output)
         return overflow
 
     def process_loss_scale(self, overflow):
@@ -517,7 +580,7 @@ def tensor_shard_grad_scale_pipeline(scale, grad, accu_grad):
     return new_grad
 
 
-class _TrainPipelineWithLossScaleCell(TrainOneStepCell):
+class _TrainGradAccuWithLossScaleCell(TrainOneStepCell):
     """
     Append an optimizer to the training network after that the construct
     function can be called to create the backward graph.
@@ -528,7 +591,7 @@ class _TrainPipelineWithLossScaleCell(TrainOneStepCell):
         scale_sense (Cell): Cell to do the loss scale.
     """
     def __init__(self, network, optimizer, scale_sense):
-        super(_TrainPipelineWithLossScaleCell, self).__init__(network, optimizer, sens=None)
+        super(_TrainGradAccuWithLossScaleCell, self).__init__(network, optimizer, sens=None)
         self.network = network
         self.network.add_flags(defer_inline=True)
         self.weights = optimizer.parameters

mindspore/numpy/array_creations.py

@@ -1304,7 +1304,7 @@ def triu(m, k=0):
     if rank < 1:
         _raise_value_error("input m's rank should be larger than 0")
     elif rank == 1:
-        mask = tri(m.shape[0], k=k-1, dtype=mstype.bool_)
+        mask = tri(m.shape[0], k=k - 1, dtype=mstype.bool_)
         return where(mask, zeros(1, m.dtype), m)
     # Only Ascend hardware will reduce accuracy
     if device_target == "Ascend":
@@ -2587,7 +2587,6 @@ def _limit_stat_length(stat_length, shape):
     return tuple((min(stat_pair[0], shape[i]), min(stat_pair[1], shape[i])) for i, stat_pair in enumerate(stat_length))
 
 
-@constexpr
 def _convert_pad_to_nd(pad_values, ndim):
     """broadcasts the pad_values to (ndim * 2)"""
     if not isinstance(pad_values, (int, list, tuple, Tensor)):

mindspore/numpy/array_ops.py

@@ -2585,8 +2585,9 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False):
     """
     def unique_w_ind(arr):
         array, sort_indices = arr.ravel().sort()
-        cmp_array1 = F.cat((array, Tensor([0], dtype=array.dtype)))
-        cmp_array2 = F.cat((Tensor([0], dtype=array.dtype), array))
+        array_type = array.dtype
+        cmp_array1 = F.cat((array, Tensor([0], dtype=array_type)))
+        cmp_array2 = F.cat((Tensor([0], dtype=array_type), array))
         mask = cmp_array1 != cmp_array2
         mask[0] = True
         array = F.masked_select(array, mask[:-1])