mindspore 2.6.0__cp311-none-any.whl → 2.7.0rc1__cp311-none-any.whl

mindspore/nn/wrap/loss_scale.py

@@ -26,33 +26,31 @@ from mindspore.common import Tensor
 from mindspore.common.sparse_tensor import RowTensorInner
 from mindspore.common.parameter import Parameter
 from mindspore.ops.operations.math_ops import NPUGetFloatStatusV2, NPUClearFloatStatusV2
-from mindspore.ops import functional as F
-from mindspore.ops import composite as C
-from mindspore.ops import operations as P
+from mindspore import ops
 from mindspore.ops.operations.nn_ops import AllFinite
 from mindspore.common import dtype as mstype
 from mindspore._c_expression import MSContext
 from mindspore.run_check._check_version import AscendEnvChecker
 from mindspore import log as logger
 
-_grad_scale = C.MultitypeFuncGraph("grad_scale")
-reciprocal = P.Reciprocal()
+_grad_scale = ops.MultitypeFuncGraph("grad_scale")
+reciprocal = ops.Reciprocal()
 
 
 @_grad_scale.register("Tensor", "Tensor")
 def tensor_grad_scale(scale, grad):
-    return grad * F.cast(reciprocal(scale), F.dtype(grad))
+    return grad * ops.cast(reciprocal(scale), ops.dtype(grad))
 
 
 @_grad_scale.register("Tensor", "RowTensor")
 def tensor_grad_scale_row_tensor(scale, grad):
     return RowTensorInner(grad.indices,
-                          grad.values * F.cast(reciprocal(scale), F.dtype(grad.values)),
+                          grad.values * ops.cast(reciprocal(scale), ops.dtype(grad.values)),
                           grad.dense_shape)
 
 
-_grad_overflow = C.MultitypeFuncGraph("_grad_overflow")
-grad_overflow = P.FloatStatus()
+_grad_overflow = ops.MultitypeFuncGraph("_grad_overflow")
+grad_overflow = ops.FloatStatus()
 
 
 @_grad_overflow.register("Tensor")
@@ -65,8 +63,8 @@ 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 = ops.MultitypeFuncGraph("_ascend_grad_overflow")
+ascend_grad_overflow = ops.IsFinite()
 
 
 @_ascend_grad_overflow.register("Tensor")
@@ -74,7 +72,7 @@ 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,)))
+    output = ops.Reshape()(output, ((-1,)))
     return output
 
 
@@ -83,7 +81,7 @@ 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,)))
+    output = ops.Reshape()(output, ((1,)))
     return output
 
 
@@ -154,14 +152,14 @@ class DynamicLossScaleUpdateCell(Cell):
 
         self.cur_iter = Parameter(Tensor(1, dtype=mstype.int32), name="current_iterator_step")
         self.last_overflow_iter = Parameter(Tensor(0, dtype=mstype.int32), name="last_overflow_iterator_step")
-        self.select = P.Select()
-        self.max = P.Maximum()
+        self.select = ops.Select()
+        self.max = ops.Maximum()
         self.minimum_loss_scale = Tensor(1.0, dtype=mstype.float32)
-        self.reciprocal = P.Reciprocal()
-        self.less_equal = P.LessEqual()
-        self.logic_and = P.LogicalAnd()
-        self.logic_not = P.LogicalNot()
-        self.logic_or = P.LogicalOr()
+        self.reciprocal = ops.Reciprocal()
+        self.less_equal = ops.LessEqual()
+        self.logic_and = ops.LogicalAnd()
+        self.logic_not = ops.LogicalNot()
+        self.logic_or = ops.LogicalOr()
         self.const_true = Tensor(True, dtype=mstype.bool_)
 
     def get_loss_scale(self):
@@ -187,14 +185,14 @@ class DynamicLossScaleUpdateCell(Cell):
         should_inc = self.less_equal(self.scale_window, self.cur_iter - self.last_overflow_iter)
         last_iter_cond = self.logic_or(overflow_cond, should_inc)
         last_overflow_iter = self.select(last_iter_cond, self.cur_iter, self.last_overflow_iter)
-        last_iter = F.assign(self.last_overflow_iter, last_overflow_iter)
+        last_iter = ops.assign(self.last_overflow_iter, last_overflow_iter)
         update_scale_cond = self.logic_and(should_inc, self.logic_not(overflow_cond))
         scale_mul_res = loss_scale_on_overflow * self.scale_factor
         scaled_loss_scale = self.select(update_scale_cond, scale_mul_res, loss_scale_on_overflow)
-        F.assign(loss_scale, scaled_loss_scale)
+        ops.assign(loss_scale, scaled_loss_scale)
         inc_cur_iter = self.cur_iter + 1
-        inc_cur_iter = F.depend(inc_cur_iter, last_iter)
-        F.assign(self.cur_iter, inc_cur_iter)
+        inc_cur_iter = ops.depend(inc_cur_iter, last_iter)
+        ops.assign(self.cur_iter, inc_cur_iter)
         return overflow
 
 
@@ -360,19 +358,19 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
 
     def __init__(self, network, optimizer, scale_sense):
         super(TrainOneStepWithLossScaleCell, self).__init__(network, optimizer, sens=None)
-        self.hyper_map = C.HyperMap()
+        self.hyper_map = ops.HyperMap()
         self.base = Tensor(1, mstype.float32)
         self.base0 = Tensor(0, mstype.int32)
-        self.reduce_sum = P.ReduceSum(keep_dims=False)
-        self.reduce_all = P.ReduceAll(keep_dims=False)
-        self.less_equal = P.LessEqual()
-        self.equal = P.Equal()
-        self.logic_not = P.LogicalNot()
-        self.allreduce = P.AllReduce()
+        self.reduce_sum = ops.ReduceSum(keep_dims=False)
+        self.reduce_all = ops.ReduceAll(keep_dims=False)
+        self.less_equal = ops.LessEqual()
+        self.equal = ops.Equal()
+        self.logic_not = ops.LogicalNot()
+        self.allreduce = ops.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() in ['ascend910b', 'ascend910_93'])
+        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() in ['ascend910b', 'ascend910_93']
         self.loss_scaling_manager = None
         self._ascend_check_overflow_mode = os.environ.get('MS_ASCEND_CHECK_OVERFLOW_MODE')
 
@@ -420,9 +418,9 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         scaling_sens = self.scale_sense
         status = Tensor([0] * 8, mstype.int32)
 
-        scaling_sens_filled = C.ones_like(loss) * F.cast(scaling_sens, F.dtype(loss))
+        scaling_sens_filled = ops.ones_like(loss) * ops.cast(scaling_sens, ops.dtype(loss))
         grads = self.grad(self.network, weights)(*inputs, scaling_sens_filled)
-        grads = self.hyper_map(F.partial(_grad_scale, scaling_sens), grads)
+        grads = self.hyper_map(ops.partial(_grad_scale, scaling_sens), grads)
         # apply grad reducer on grads
         grads = self.grad_reducer(grads)
 
@@ -431,7 +429,7 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         overflow = self.process_loss_scale(cond)
         # if there is no overflow, do optimize
         if not overflow:
-            loss = F.depend(loss, self.optimizer(grads))
+            loss = ops.depend(loss, self.optimizer(grads))
         return loss, cond, scaling_sens
 
     def set_sense_scale(self, sens):
@@ -475,18 +473,18 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         status = Tensor([0] * 8, mstype.int32)
         if self.ascend_910a_target or (self.ascend_910b_target and \
                                        self._ascend_check_overflow_mode == "SATURATION_MODE"):
-            status = F.depend(status, pre_cond)
+            status = ops.depend(status, pre_cond)
             # clear overflow buffer
             clear_status = NPUClearFloatStatusV2()(status)
-            compute_input = F.depend(compute_input, clear_status)
+            compute_input = ops.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)
+        flag_sum = self.hyper_map(ops.partial(grad_overflow_check_func), compute_output)
+        flag_sum = ops.AddN()(flag_sum)
         # convert flag_sum to scalar
-        flag_sum = P.Reshape()(flag_sum, (()))
+        flag_sum = ops.Reshape()(flag_sum, (()))
         return flag_sum
 
     def _get_distributed_overflow_status_on_infnan_mode(self, grad_overflow_check_func, compute_output):
@@ -506,8 +504,8 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         overflow = AllFinite()(compute_output)
 
         if self.is_distributed:
-            overflow = P.Cast()(overflow, mstype.float32)
-            overflow = P.Cast()(self.allreduce(overflow), mstype.bool_)
+            overflow = ops.Cast()(overflow, mstype.float32)
+            overflow = ops.Cast()(self.allreduce(overflow), mstype.bool_)
         return overflow
 
     def _get_gpu_overflow_status(self, compute_output):
@@ -526,7 +524,7 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
 
     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)
+        status = ops.depend(status, compute_output)
         get_status = NPUGetFloatStatusV2()(status)
 
         if self.is_distributed:
@@ -534,15 +532,15 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
             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)
+            status = ops.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)
+            flag = ops.depend(flag, clear_status)
             overall_finite = self.reduce_all(flag)
         else:
-            status = F.depend(status, get_status)
+            status = ops.depend(status, get_status)
             clear_status = NPUClearFloatStatusV2()(status)
-            get_status = F.depend(get_status, clear_status)
+            get_status = ops.depend(get_status, clear_status)
             flag = self.equal(self.base0, get_status)
             overall_finite = self.reduce_all(flag)
         overflow = self.logic_not(overall_finite)
@@ -592,26 +590,26 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         return overflow
 
 
-grad_scale = C.MultitypeFuncGraph("grad_scale")
-shard_grad_scale = C.MultitypeFuncGraph("shard_grad_scale")
-reciprocal = P.Reciprocal()
+grad_scale = ops.MultitypeFuncGraph("grad_scale")
+shard_grad_scale = ops.MultitypeFuncGraph("shard_grad_scale")
+reciprocal = ops.Reciprocal()
 
 
 @grad_scale.register("Tensor", "Tensor", "Tensor")
 def tensor_grad_scale_pipeline(scale, grad, accu_grad):
-    accu_grad = F.depend(accu_grad, grad)
+    accu_grad = ops.depend(accu_grad, grad)
     new_grad = accu_grad * reciprocal(scale)
-    accu_grad = F.depend(accu_grad, new_grad)
-    zeros = F.tensor_mul(accu_grad, 0.0)
-    new_grad = F.depend(new_grad, F.assign(accu_grad, zeros))
+    accu_grad = ops.depend(accu_grad, new_grad)
+    zeros = ops.tensor_mul(accu_grad, 0.0)
+    new_grad = ops.depend(new_grad, ops.assign(accu_grad, zeros))
     return new_grad
 
 
 @shard_grad_scale.register("Tensor", "Tensor", "Tensor")
 def tensor_shard_grad_scale_pipeline(scale, grad, accu_grad):
     new_grad = grad * reciprocal(scale)
-    accu_grad = F.depend(accu_grad, new_grad)
-    new_grad = F.depend(new_grad, F.assign(accu_grad, F.zeros_like(accu_grad)))
+    accu_grad = ops.depend(accu_grad, new_grad)
+    new_grad = ops.depend(new_grad, ops.assign(accu_grad, ops.zeros_like(accu_grad)))
     return new_grad
 
 
@@ -633,23 +631,23 @@ class _TrainGradAccuWithLossScaleCell(TrainOneStepCell):
         self.weights = optimizer.parameters
         self.accu_grads = self.weights.clone(prefix="accu_grads", init="zeros")
         self.optimizer = optimizer
-        self.grad = C.GradOperation(get_by_list=True, sens_param=True)
+        self.grad = ops.GradOperation(get_by_list=True, sens_param=True)
         self.grad_reducer = nn.Identity()
         self.degree = 1
-        self.cast = P.Cast()
-        self.alloc_status = P.NPUAllocFloatStatus()
-        self.get_status = P.NPUGetFloatStatus()
-        self.clear_before_grad = P.NPUClearFloatStatus()
-        self.reduce_sum = P.ReduceSum(keep_dims=False)
+        self.cast = ops.Cast()
+        self.alloc_status = ops.NPUAllocFloatStatus()
+        self.get_status = ops.NPUGetFloatStatus()
+        self.clear_before_grad = ops.NPUClearFloatStatus()
+        self.reduce_sum = ops.ReduceSum(keep_dims=False)
         if self.parallel_mode not in [ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL]:
             raise ValueError(f"ParallelMode must be one of "
                              f"[ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL], but found "
                              f"{self.parallel_mode}.")
-        self.allreduce = P.AllReduce()
+        self.allreduce = ops.AllReduce()
         self.base = Tensor(1, mstype.float32)
-        self.less_equal = P.LessEqual()
-        self.hyper_map = C.HyperMap()
-        self.reshape = P.Reshape()
+        self.less_equal = ops.LessEqual()
+        self.hyper_map = ops.HyperMap()
+        self.reshape = ops.Reshape()
         self.loss_scaling_manager = None
         if isinstance(scale_sense, Cell):
             self.loss_scaling_manager = scale_sense
@@ -669,19 +667,19 @@ class _TrainGradAccuWithLossScaleCell(TrainOneStepCell):
         loss = self.network(*inputs)
         scaling_sens = self.scale_sense
         init = self.alloc_status()
-        scaling_sens_filled = C.ones_like(loss) * F.cast(scaling_sens, F.dtype(loss))
-        scaling_sens_filled = F.depend(scaling_sens_filled, self.clear_before_grad(init))
+        scaling_sens_filled = ops.ones_like(loss) * ops.cast(scaling_sens, ops.dtype(loss))
+        scaling_sens_filled = ops.depend(scaling_sens_filled, self.clear_before_grad(init))
         grads = self.grad(self.network, self.weights)(*inputs, scaling_sens_filled)
-        init = F.depend(init, grads)
+        init = ops.depend(init, grads)
         get_status = self.get_status(init)
-        init = F.depend(init, get_status)
+        init = ops.depend(init, get_status)
         flag_sum = self.reduce_sum(init, (0,))
         if self.opt_shard:
             grads = self.grad_reducer(grads)
-            grads = self.hyper_map(F.partial(shard_grad_scale, scaling_sens * self.degree), grads, self.accu_grads)
+            grads = self.hyper_map(ops.partial(shard_grad_scale, scaling_sens * self.degree), grads, self.accu_grads)
         else:
             accu_grads = self.grad_reducer(self.accu_grads)
-            grads = self.hyper_map(F.partial(grad_scale, scaling_sens * self.degree), grads, accu_grads)
+            grads = self.hyper_map(ops.partial(grad_scale, scaling_sens * self.degree), grads, accu_grads)
         # sum overflow flag over devices
         flag_reduce = self.allreduce(flag_sum)
         cond = self.less_equal(self.base, flag_reduce)

mindspore/numpy/array_creations.py

@@ -1384,7 +1384,7 @@ def trace(a, offset=0, axis1=0, axis2=1, dtype=None):
 
     Note:
         - `trace` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `trace` is not supported on Windows platform yet.
 
     Args:
@@ -2512,8 +2512,8 @@ def _pad_symmetric(arr, pad_width, reflect_type):
     for i in range(arr.ndim):
         array_length = arr.shape[i]
 
-        has_pad_before = (pad_width[i][0] > 0)
-        has_pad_after = (pad_width[i][1] > 0)
+        has_pad_before = pad_width[i][0] > 0
+        has_pad_after = pad_width[i][1] > 0
 
         times_to_pad_before = pad_width[i][0] // array_length + 1
         additional_pad_before = pad_width[i][0] % array_length
@@ -2541,8 +2541,8 @@ def _pad_reflect(arr, pad_width, reflect_type):
             total_repeats = pad_width[i][0] + pad_width[i][1] + 1
             arr = ops.tile(arr, _tuple_setitem((1,) * arr.ndim, i, total_repeats))
         else:
-            has_pad_before = (pad_width[i][0] > 0)
-            has_pad_after = (pad_width[i][1] > 0)
+            has_pad_before = pad_width[i][0] > 0
+            has_pad_after = pad_width[i][1] > 0
 
             pad_size = array_length - 1
             times_to_pad_before = pad_width[i][0] // pad_size + 1

mindspore/numpy/fft.py

@@ -684,7 +684,7 @@ def ihfft(a, n=None, axis=-1, norm=None):
 
     Note:
         - `ihfft` is currently only used in `mindscience` scientific computing scenarios and
-        dose not support other usage scenarios.
+        does not support other usage scenarios.
         - `ihfft` is not supported on Windows platform yet.
 
     Args:

mindspore/numpy/math_ops.py

@@ -5809,7 +5809,7 @@ def correlate(a, v, mode='valid'):
 
     Note:
         - `correlate` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `correlate` is not supported on Windows platform yet.
 
     Args:

mindspore/ops/_grad_experimental/grad_comm_ops.py

@@ -16,6 +16,7 @@
 """Generate bprop for comm ops"""
 from __future__ import division
 from __future__ import absolute_import
+import os
 from mindspore import Tensor, Parameter
 import mindspore.common.dtype as mstype
 from mindspore.ops import functional as F
@@ -27,7 +28,8 @@ from mindspore.ops.operations._inner_ops import issubclass_
 from mindspore.common.sparse_tensor import RowTensorInner
 from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
 from mindspore.ops.operations.comm_ops import (AllGather, _MiniStepAllGather, _HostAllGather, AllReduce,
-                                               NeighborExchange, AlltoAll, AlltoAllV, NeighborExchangeV2, Broadcast,
+                                               NeighborExchange, AlltoAll, AlltoAllV, NeighborExchangeV2,
+                                               Broadcast, AllGatherV, ReduceScatterV,
                                                _GetTensorSlice, _MirrorOperator, _MirrorMiniStepOperator, ReduceOp,
                                                ReduceScatter, _HostReduceScatter, _VirtualDiv, _VirtualAdd, _AllSwap,
                                                _VirtualAssignAdd, _VirtualAccuGrad, _MirrorMicroStepOperator,
@@ -236,7 +238,6 @@ def get_bprop_mirror_micro_step_operator(self):
     allgather for sparse feature.
     """
     group = self.group
-    global_rank = get_rank()
     dev_num = self.dev_num
     mean_flag = self.mean_flag
     param_name = " "
@@ -270,6 +271,9 @@ def get_bprop_mirror_micro_step_operator(self):
     dump_local_norm = ms.get_auto_parallel_context("dump_local_norm")
     dump_local_norm_path = ms.get_auto_parallel_context("dump_local_norm_path")
     dump_device_local_norm = ms.get_auto_parallel_context("dump_device_local_norm")
+    if dump_local_norm_path:
+        global_rank = get_rank()
+        file = os.path.join(dump_local_norm_path, "rank_" + str(global_rank), "local_norm__" + param_name)
     if dump_device_local_norm:
         # init _squared _squared_device_local_norm
         squared_device_local_norm = get_squared_device_local_norm_param()
@@ -279,8 +283,7 @@ def get_bprop_mirror_micro_step_operator(self):
             squared_norm = reduce_sum(square((z)))
             if dump_local_norm:
                 if dump_local_norm_path:
-                    z = F.depend(z, tensor_dump(dump_local_norm_path + "/rank_" + str(global_rank) +
-                                                "/local_norm__" + param_name, sqrt(squared_norm)))
+                    z = F.depend(z, tensor_dump(file, sqrt(squared_norm)))
                 else:
                     z = F.depend(z, ln_print("dump local norm: ", param_name, sqrt(squared_norm)))
             if dump_device_local_norm:
@@ -336,13 +339,15 @@ def get_bprop_all_gather(self):
     dump_local_norm_path = ms.get_auto_parallel_context("dump_local_norm_path")
     dump_device_local_norm = ms.get_auto_parallel_context("dump_device_local_norm")
     if param_name and (dump_local_norm or dump_device_local_norm):
-        global_rank = get_rank()
         cast = P.Cast()
         ln_print = P.Print()
         tensor_dump = P.TensorDump()
         reduce_sum = P.ReduceSum(keep_dims=False)
         square = P.Square()
         sqrt = P.Sqrt()
+    if dump_local_norm_path:
+        global_rank = get_rank()
+        file = os.path.join(dump_local_norm_path, "rank_" + str(global_rank), "local_norm__" + param_name)
     if dump_device_local_norm:
         # init _squared _squared_device_local_norm
         squared_device_local_norm = get_squared_device_local_norm_param()
@@ -352,8 +357,7 @@ def get_bprop_all_gather(self):
             squared_norm = reduce_sum(square((dout)))
             if dump_local_norm:
                 if dump_local_norm_path:
-                    dout = F.depend(dout, tensor_dump(dump_local_norm_path + "/rank_" + str(global_rank) +
-                                                      "/local_norm__" + param_name, sqrt(squared_norm)))
+                    dout = F.depend(dout, tensor_dump(file, sqrt(squared_norm)))
                 else:
                     dout = F.depend(dout, ln_print("dump local norm: ", param_name, sqrt(squared_norm)))
             if dump_device_local_norm:
@@ -430,7 +434,6 @@ def get_bprop_micro_step_all_gather(self):
         if self.instance_name:
             instance_name = "grad_" + self.instance_name
             reduce_scatter.set_prim_instance_name(instance_name)
-    global_rank = get_rank()
     cast = P.Cast()
     dtype = P.DType()
     out_tensor = Tensor(1.0, mstype.float16)
@@ -443,6 +446,9 @@ def get_bprop_micro_step_all_gather(self):
     dump_local_norm = ms.get_auto_parallel_context("dump_local_norm")
     dump_local_norm_path = ms.get_auto_parallel_context("dump_local_norm_path")
     dump_device_local_norm = ms.get_auto_parallel_context("dump_device_local_norm")
+    if dump_local_norm_path:
+        global_rank = get_rank()
+        file = os.path.join(dump_local_norm_path, "rank_" + str(global_rank), "local_norm__" + param_name)
     if dump_device_local_norm:
         # init _squared _squared_device_local_norm
         squared_device_local_norm = get_squared_device_local_norm_param()
@@ -460,8 +466,7 @@ def get_bprop_micro_step_all_gather(self):
             squared_norm = reduce_sum(square((z)))
             if dump_local_norm:
                 if dump_local_norm_path:
-                    z = F.depend(z, tensor_dump(dump_local_norm_path + "/rank_" + str(global_rank) +
-                                                "/local_norm__" + param_name, sqrt(squared_norm)))
+                    z = F.depend(z, tensor_dump(file, sqrt(squared_norm)))
                 else:
                     z = F.depend(z, ln_print("dump local norm: ", param_name, sqrt(squared_norm)))
             if dump_device_local_norm:
@@ -650,6 +655,38 @@ def get_bprop_all_to_all_v(self):
     return bprop
 
 
+@bprop_getters.register(AllGatherV)
+def get_bprop_all_gather_v(self):
+    """Generate bprop for AllGatherV."""
+    all_gather_v_grad = ReduceScatterV(ReduceOp.SUM, self.group)
+    if hasattr(self, "instance_name") and self.instance_name:
+        instance_name = "grad" + self.instance_name
+        all_gather_v_grad.set_prim_instance_name(instance_name)
+
+    def bprop(x, output_split_sizes, out, dout):
+        dx = all_gather_v_grad(dout, output_split_sizes)
+        return (dx, zeros_like(output_split_sizes))
+
+    return bprop
+
+
+@bprop_getters.register(ReduceScatterV)
+def get_bprop_reduce_scatter_v(self):
+    """Generate bprop for ReduceScatterV."""
+    reduce_scatter_v_grad = AllGatherV(self.group)
+    if hasattr(self, "instance_name") and self.instance_name:
+        instance_name = "grad" + self.instance_name
+        reduce_scatter_v_grad.set_prim_instance_name(instance_name)
+    if self.op != ReduceOp.SUM:
+        raise RuntimeError("The reducescatter bprop only support ReduceOp.SUM until now.")
+
+    def bprop(x, input_split_sizes, out, dout):
+        dx = reduce_scatter_v_grad(dout, input_split_sizes)
+        return (dx, zeros_like(input_split_sizes))
+
+    return bprop
+
+
 @bprop_getters.register(NeighborExchangeV2)
 def get_bprop_neighborexchangev2(self):
     """Generate bprop for NeighborExchangeV2."""
@@ -685,11 +722,13 @@ def get_bprop_mirror_operator(self):
     dump_local_norm = ms.get_auto_parallel_context("dump_local_norm")
     dump_local_norm_path = ms.get_auto_parallel_context("dump_local_norm_path")
     dump_device_local_norm = ms.get_auto_parallel_context("dump_device_local_norm")
+    if dump_local_norm_path:
+        global_rank = get_rank()
+        file = os.path.join(dump_local_norm_path, "rank_" + str(global_rank), "local_norm__" + param_name)
     if dump_device_local_norm:
         # init _squared _squared_device_local_norm
         squared_device_local_norm = get_squared_device_local_norm_param()
     if dev_num > 1:
-        global_rank = get_rank()
         dev_num_r = 1.0 / dev_num
         all_reduce = AllReduce(group=group)
         all_gather = AllGather(group=group)
@@ -717,8 +756,7 @@ def get_bprop_mirror_operator(self):
             squared_norm = reduce_sum(square((dout)))
             if dump_local_norm:
                 if dump_local_norm_path:
-                    dout = F.depend(dout, tensor_dump(dump_local_norm_path + "/rank_" + str(global_rank) +
-                                                      "/local_norm__" + param_name, sqrt(squared_norm)))
+                    dout = F.depend(dout, tensor_dump(file, sqrt(squared_norm)))
                 else:
                     dout = F.depend(dout, ln_print("dump local norm: ", param_name, sqrt(squared_norm)))
             if dump_device_local_norm:

mindspore/ops/_grad_experimental/grad_debug_ops.py

@@ -41,3 +41,17 @@ def get_bprop_insert_gradient_of(self):
 def bprop_tensor_dump(file, input_x, out, dout):
     """Generate bprop for TensorDump"""
     return file, C.zeros_like(input_x)
+
+
+@bprop_getters.register(P.DumpGradient)
+def get_bprop_dump_gradient(self):
+    """Generate bprop for DumpGradient"""
+    td = P.TensorDump()
+    td.add_prim_attr("side_effect_io", False)
+    td.add_prim_attr("td_flag", True)
+
+    def bprop(path, x, input_output, out, dout):
+        tded = td(path, dout)
+        fdout = F.depend(dout, tded)
+        return C.zeros_like(path), fdout, C.zeros_like(input_output)
+    return bprop

mindspore/ops/_vmap/vmap_array_ops.py

@@ -1527,14 +1527,13 @@ def get_meshgrid_vmap_rule(prim, axis_size):
 
         if not isinstance(inputs_bdim, (tuple)):
             _raise_value_error("The inputs of P.Meshgrid is not tuple.")
-        args = inputs_bdim
-        if len(args) <= 1:
+        if len(inputs_bdim) <= 1:
             _raise_value_error(
                 "The input number of P.Meshgrid must be greater than 1.")
 
         output_shape = []
         ones_shape = []
-        for each_arg in args:
+        for each_arg in inputs_bdim:
             x, bdim = each_arg
             if bdim is None:
                 _raise_value_error(
@@ -1548,22 +1547,16 @@ def get_meshgrid_vmap_rule(prim, axis_size):
         output_shape.insert(0, axis_size)
         ones_shape.insert(0, axis_size)
 
-        indexing, _ = indexing_bdim
-
-        if indexing == Indexing.xy.value:
+        if indexing_bdim[0] == Indexing.xy.value:
             output_shape[1], output_shape[2] = output_shape[2], output_shape[1]
-        shape = tuple(output_shape)
-
-        input_0, _ = args[0]
-        dtype = F.dtype(input_0)
-        ones_tensor = F.fill(dtype, shape, 1)
+        ones_tensor = F.fill(F.dtype(inputs_bdim[0][0]), tuple(output_shape), 1)
 
         index = 0
         vals_out_tuple = ()
-        for each_arg in args:
+        for each_arg in inputs_bdim:
             x, bdim = each_arg
             x = _bdim_at_front(x, bdim, axis_size)
-            shape_index = (1 - index) if (index <= 1 and indexing == Indexing.xy.value) else index
+            shape_index = (1 - index) if (index <= 1 and indexing_bdim[0] == Indexing.xy.value) else index
             ones_shape[shape_index + 1] = output_shape[shape_index + 1]
             x = P.Reshape()(x, tuple(ones_shape))
             output = P.Mul()(x, ones_tensor)

mindspore/ops/_vmap/vmap_nn_ops.py

@@ -2068,14 +2068,15 @@ def get_sparse_apply_adagrad_vmap_rule(prim, axis_size):
         indices, indices_dim = indices_bdim
         if var_dim is None:
             if any(dim is not None for dim in [accum_dim, grad_dim, indices_dim]):
-                ValueError("The source axis of `var` is None, but the source "
-                           "axis of `accum/grad/indices` is not None. The execution order of "
-                           "operator `{}` cannot be guaranteed.".format(prim_name))
+                _raise_value_error("The source axis of `var` is None, but the source "
+                                   "axis of `accum/grad/indices` is not None. The execution "
+                                   "order of operator `{}` cannot be guaranteed.".format(prim_name))
             var, accum = prim(var, accum, grad, indices, u_monad)
             return (var, None), (accum, None)
         if var_dim != 0 or accum_dim != var_dim:
-            ValueError("For `{}`, the source axis of `var` must be equal to `accum`, and not equal to 0, "
-                       "but got the source axis of `var`: {}, `accum`: {}.".format(prim_name, var_dim, accum_dim))
+            _raise_value_error("For `{}`, the source axis of `var` must be equal to `accum`, "
+                               "and not equal to 0, but got the source axis of `var`: {}, "
+                               "`accum`: {}.".format(prim_name, var_dim, accum_dim))
 
         grad = _bdim_at_front(grad, grad_dim, axis_size)
         indices = _bdim_at_front(indices, indices_dim, axis_size)
@@ -2094,27 +2095,18 @@ def get_sparse_apply_ftrl_vmap_rule(prim, axis_size):
     else:
         batch_rank = 1
 
-    prim_name = prim.name
     batch_prim = _vmap_clone_prim(prim)
     batch_prim.add_prim_attr('batch_rank', batch_rank)
 
     def vmap_rule(var_bdim, accum_bdim, linear_bdim, grad_bdim, indices_bdim, u_monad):
         var, var_dim = var_bdim
-        accum, accum_dim = accum_bdim
-        linear, linear_dim = linear_bdim
+        accum, _ = accum_bdim
+        linear, _ = linear_bdim
         grad, grad_dim = grad_bdim
         indices, indices_dim = indices_bdim
         if var_dim is None:
-            if any(dim is not None for dim in [accum_dim, linear_dim, grad_dim, indices_dim]):
-                ValueError("The source axis of `var` is None, but the source "
-                           "axis of `accum/linear/grad/indices` is not None. The execution order of "
-                           "operator `{}` cannot be guaranteed.".format(prim_name))
             var, accum, linear = prim(var, accum, linear, grad, indices, u_monad)
             return (var, None), (accum, None), (linear, None)
-        if var_dim != 0 or accum_dim != var_dim or linear_dim != var_dim:
-            ValueError("For `{}`, the source axis of `var`, `accum` and `linear` must be equal, and "
-                       "not equal to 0, but got the source axis of `var`: {}, `accum`: {}, "
-                       "`linear`:{}.".format(prim_name, var_dim, accum_dim, linear_dim))
 
         grad = _bdim_at_front(grad, grad_dim, axis_size)
         indices = _bdim_at_front(indices, indices_dim, axis_size)