mindspore 2.4.1__cp311-cp311-win_amd64.whl → 2.5.0__cp311-cp311-win_amd64.whl

mindspore/numpy/utils.py

@@ -19,8 +19,8 @@ import types
 
 from mindspore.common import Tensor
 from mindspore._c_expression import Tensor as Tensor_
-from mindspore.ops import functional as F
 from mindspore.common import dtype as mstype
+from mindspore import ops
 
 from mindspore.numpy.utils_const import _tile_size, _add_unit_axes, _raise_type_error, _type_convert, \
     _tuple_setitem, _callable_const, _check_is_float, _get_device
@@ -65,7 +65,7 @@ def _check_input_for_asarray(array_like):
 
 def _is_scalar(shape):
     """check whether input shape is a scalar"""
-    return F.shape_mul(shape) == 1
+    return ops.shape_mul(shape) == 1
 
 
 def _convert_list_tensor_to_tuple_tensor(list_of_tensor):
@@ -80,27 +80,27 @@ def _convert_list_tensor_to_tuple_tensor(list_of_tensor):
 
 def _expand(x, ndim, axis=0):
     """Expand x to ndim from axis, which can be 0 or -1."""
-    shape = _add_unit_axes(F.shape(x), ndim, axis == -1)
-    return F.reshape(x, shape)
+    shape = _add_unit_axes(ops.shape(x), ndim, axis == -1)
+    return ops.reshape(x, shape)
 
 
 def _broadcast_to(x, shape_cur, shape_to, ndim_to):
     """Broadcasts x from shape_cur to shape_to."""
     size = _tile_size(shape_cur, shape_to, ndim_to)
-    return F.tile(x, size)
+    return ops.tile(x, size)
 
 
 def _broadcast_to_shape(x, shape):
     """Broadcasts x from current shape to shape"""
     ndim_to = len(shape)
     x = _expand(x, ndim_to)
-    return _broadcast_to(x, F.shape(x), shape, ndim_to)
+    return _broadcast_to(x, ops.shape(x), shape, ndim_to)
 
 
 def _get_size(x, axis=None):
     """Get the number of elements along the given axis of tensor x."""
-    if axis is None or F.tuple_len(axis) == 0:
-        axis = F.make_range(x.ndim)
+    if axis is None or ops.tuple_len(axis) == 0:
+        axis = ops.make_range(x.ndim)
     nums = 1
     for ax in axis:
         nums *= x.shape[ax]
@@ -110,7 +110,7 @@ def _get_size(x, axis=None):
 def _check_input_tensor(*tensors):
     for tensor in tensors:
         if not isinstance(tensor, Tensor):
-            _raise_type_error('expect Tensor, but got ', F.typeof(tensor))
+            _raise_type_error('expect Tensor, but got ', ops.typeof(tensor))
     return True
 
 
@@ -141,7 +141,7 @@ def _to_tensor(*args):
 
 def _get_dtype_from_scalar(*input_numbers):
     """
-    Get the final dtype from series of input numbers, compared with F.typeof, we
+    Get the final dtype from series of input numbers, compared with ops.typeof, we
     return int32/float32 for python int/float instead.
     """
     bool_flag = True
@@ -184,7 +184,7 @@ def _slice_along_axis(f, axis, slice_start, slice_end):
     slice_size = slice_end - slice_start
     index_start = _tuple_setitem(index_start, axis, slice_start)
     index_end = _tuple_setitem(index_end, axis, slice_size)
-    return F.tensor_slice(f, index_start, index_end)
+    return ops.tensor_slice(f, index_start, index_end)
 
 
 def _to_tensor_origin_dtype(*args):
@@ -203,12 +203,12 @@ def _to_tensor_origin_dtype(*args):
 
 def _callable(tensor, obj):
     """Returns True if `obj` is a function."""
-    if F.isconstant(tensor):
+    if ops.isconstant(tensor):
         return isinstance(obj, types.FunctionType)
-    return _callable_const(F.typeof(obj))
+    return _callable_const(ops.typeof(obj))
 
 
 def _isnan(x):
-    if _get_device() == 'Ascend' and not _check_is_float(F.dtype(x)):
-        return F.fill(mstype.bool_, F.shape(x), False)
-    return F.isnan(x)
+    if _get_device() == 'Ascend' and not _check_is_float(ops.dtype(x)):
+        return ops.fill(mstype.bool_, ops.shape(x), False)
+    return ops.isnan(x)

mindspore/numpy/utils_const.py

@@ -20,7 +20,6 @@ from itertools import accumulate
 import operator
 
 import mindspore.context as context
-from mindspore.ops import functional as F
 from mindspore.ops.primitive import constexpr
 from mindspore.ops.primitive import _primexpr
 from mindspore.common import dtype as mstype
@@ -28,6 +27,7 @@ from mindspore.common import Tensor
 from mindspore._c_expression import Tensor as Tensor_
 from mindspore._c_expression import typing
 from mindspore import _checkparam as validator
+from mindspore import ops
 
 from mindspore.numpy.dtypes import promotion_rule, dtype_tuple, all_types, dtype_map, rule_for_trigonometric
 
@@ -80,13 +80,13 @@ def _check_dtype(dtype):
 @_primexpr
 def _is_shape_empty(shp):
     """Check whether shape contains zero"""
-    if F.is_sequence_shape_unknown(shp):
+    if ops.is_sequence_shape_unknown(shp):
         return False
     if isinstance(shp, int):
         return shp == 0
     if isinstance(shp, (tuple, list)):
         return 0 in shp
-    return F.shape_mul(shp) == 0
+    return ops.shape_mul(shp) == 0
 
 
 @_primexpr
@@ -189,7 +189,7 @@ def _check_axis_valid(axes, ndim):
             raise ValueError('duplicate value in "axis"')
 
     if axes is None:
-        axes = F.make_range(ndim)
+        axes = ops.make_range(ndim)
         return axes
     if isinstance(axes, (tuple, list)):
         axes = tuple(map(lambda x: _check_axis_in_range(x, ndim), axes))

mindspore/ops/__init__.py

@@ -33,6 +33,7 @@ from mindspore.ops import composite, operations, functional, function
 from mindspore.ops import signature
 from mindspore.ops.auto_generate import cpp_create_prim_instance_helper, gen_arg_dtype_cast, gen_arg_handler, \
     gen_extend_func, gen_ops_def, gen_ops_prim, pyboost_inner_prim
+from mindspore.ops.functional_overload import all_gather_matmul, matmul_reduce_scatter
 from mindspore.ops.composite import *
 from mindspore.ops.operations import *
 from mindspore.ops.function import *
@@ -47,7 +48,7 @@ __all__ = ["get_vm_impl_fn", "vm_impl_registry",
            "CpuRegOp", "CustomRegOp", "DataType",
            "constexpr", "reshard",
            "cpp_create_prim_instance_helper", "gen_arg_dtype_cast", "gen_arg_handler", "gen_extend_func", "gen_ops_def",
-           "gen_ops_prim", "pyboost_inner_prim"]
+           "gen_ops_prim", "pyboost_inner_prim", "all_gather_matmul", "matmul_reduce_scatter"]
 __all__.extend(__primitive__)
 __all__.extend(composite.__all__)
 __all__.extend(operations.__all__)

mindspore/ops/_grad_experimental/grad_comm_ops.py

@@ -16,7 +16,7 @@
 """Generate bprop for comm ops"""
 from __future__ import division
 from __future__ import absolute_import
-from mindspore import Tensor
+from mindspore import Tensor, Parameter
 import mindspore.common.dtype as mstype
 from mindspore.ops import functional as F
 from mindspore.communication import get_rank, get_group_size
@@ -37,6 +37,24 @@ from mindspore.ops._grad_experimental.grad_base import bprop_getters
 from mindspore.ops.operations import _grad_ops as G
 import mindspore as ms
 
+_squared_device_local_norm = None
+
+
+def get_squared_device_local_norm_param():
+    """
+    Get Parameter `_squared_device_local_norm`.
+    `_squared_device_local_norm` will accumulate squared local norm of each grad in bprop under GRAPH_MODE.
+    User need to reset it to zero after network propagation each step.
+    """
+    global _squared_device_local_norm
+    if _squared_device_local_norm is None:
+        if ms.get_auto_parallel_context("dump_device_local_norm"):
+            _squared_device_local_norm = Parameter(Tensor(0.0, mstype.float32), name="_squared_device_local_norm",
+                                                   requires_grad=False)
+        else:
+            raise ValueError("The parallel config 'dump_device_local_norm' is False.")
+    return _squared_device_local_norm
+
 
 @bprop_getters.register(AllReduce)
 def get_bprop_all_reduce(self):
@@ -192,7 +210,7 @@ def get_bprop_virtual_assign_kv_cache(self):
         dout_update = dout + y
         kv_equal = F.equal(seq_chunk, 0)
         update_kv = F.select(kv_equal, F.broadcast_to(cast(out_tensor, dtype(y)), F.shape(y)), dout_update)
-        return F.depend((dout_update, cast(out_tensor, dtype(y)),
+        return F.depend((cast(dout_update, dtype(dout)), cast(out_tensor, dtype(y)),
                          cast(out_tensor, dtype(seq_chunk))), assign(y, update_kv))
 
     return bprop
@@ -218,6 +236,7 @@ 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 = " "
@@ -244,13 +263,29 @@ def get_bprop_mirror_micro_step_operator(self):
     out_tensor = Tensor(1.0, mstype.float16)
     opt_shard = _get_enable_parallel_optimizer()
     ln_print = P.Print()
+    tensor_dump = P.TensorDump()
     reduce_sum = P.ReduceSum(keep_dims=False)
     square = P.Square()
+    sqrt = P.Sqrt()
     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_device_local_norm:
+        # init _squared _squared_device_local_norm
+        squared_device_local_norm = get_squared_device_local_norm_param()
 
     def bprop(x, z, out, dout):
-        if dump_local_norm:
-            z = F.depend(z, ln_print("dump local norm: ", param_name, reduce_sum(square((z)))))
+        if dump_local_norm or dump_device_local_norm:
+            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)))
+                else:
+                    z = F.depend(z, ln_print("dump local norm: ", param_name, sqrt(squared_norm)))
+            if dump_device_local_norm:
+                z = F.depend(z, F.assign_add(squared_device_local_norm,
+                                             cast(squared_norm, squared_device_local_norm.dtype)))
         real_grad = z
         assign_out = dout
         if issubclass_(F.typeof(dout), mstype.tensor_type):
@@ -293,8 +328,38 @@ def get_bprop_all_gather(self):
     if self.rank_size == 0:
         raise ValueError(f"The 'rank_size' can not be zero, but got {self.rank_size}.")
     scale = 1.0 / self.rank_size
+    param_name = ""
+    if 'mirror_user_id' in self.get_attr_dict():
+        param_name = self.get_attr_dict()['mirror_user_id']
+    # monitor local norm
+    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 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_device_local_norm:
+        # init _squared _squared_device_local_norm
+        squared_device_local_norm = get_squared_device_local_norm_param()
 
     def bprop(x, out, dout):
+        if param_name and (dump_local_norm or dump_device_local_norm):
+            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)))
+                else:
+                    dout = F.depend(dout, ln_print("dump local norm: ", param_name, sqrt(squared_norm)))
+            if dump_device_local_norm:
+                dout = F.depend(dout, F.assign_add(squared_device_local_norm,
+                                                   cast(squared_norm, squared_device_local_norm.dtype)))
+
         dx = reduce_scatter(dout)
         if mean_flag:
             dx = F.tensor_mul(dx, scale)
@@ -365,14 +430,22 @@ 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)
     with_mirror_operator = self.get_attr_dict()["with_mirror_operator"]
     ln_print = P.Print()
+    tensor_dump = P.TensorDump()
     reduce_sum = P.ReduceSum(keep_dims=False)
     square = P.Square()
+    sqrt = P.Sqrt()
     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_device_local_norm:
+        # init _squared _squared_device_local_norm
+        squared_device_local_norm = get_squared_device_local_norm_param()
 
     def bprop(x, z, out, dout):
         if with_mirror_operator:
@@ -383,8 +456,17 @@ def get_bprop_micro_step_all_gather(self):
                 real_grad = F.tensor_mul(real_grad, scale)
             return (real_grad, cast(out_tensor, dtype(z)))
         z = F.depend(z, dout)
-        if dump_local_norm:
-            z = F.depend(z, ln_print("dump local norm: ", param_name, reduce_sum(square((z)))))
+        if dump_local_norm or dump_device_local_norm:
+            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)))
+                else:
+                    z = F.depend(z, ln_print("dump local norm: ", param_name, sqrt(squared_norm)))
+            if dump_device_local_norm:
+                z = F.depend(z, F.assign_add(squared_device_local_norm,
+                                             cast(squared_norm, squared_device_local_norm.dtype)))
         if not do_mirror:
             return (z, cast(out_tensor, dtype(z)))
         real_grad = reduce_scatter(z)
@@ -586,15 +668,23 @@ def get_bprop_mirror_operator(self):
 
     dev_num_r = 1.0
     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_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)
         mul = P.Mul()
         cast = P.Cast()
         ln_print = P.Print()
+        tensor_dump = P.TensorDump()
         reduce_sum = P.ReduceSum(keep_dims=False)
         square = P.Square()
+        sqrt = P.Sqrt()
 
         fusion = self.get_attr_dict()["fusion"]
         all_reduce.add_prim_attr("fusion", fusion)
@@ -608,8 +698,17 @@ def get_bprop_mirror_operator(self):
             all_reduce.set_prim_instance_name(instance_name)
 
     def bprop(x, out, dout):
-        if dump_local_norm:
-            dout = F.depend(dout, ln_print("dump local norm: ", param_name, reduce_sum(square((dout)))))
+        if dump_local_norm or dump_device_local_norm:
+            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)))
+                else:
+                    dout = F.depend(dout, ln_print("dump local norm: ", param_name, sqrt(squared_norm)))
+            if dump_device_local_norm:
+                dout = F.depend(dout, F.assign_add(squared_device_local_norm,
+                                                   cast(squared_norm, squared_device_local_norm.dtype)))
 
         if dev_num == 1:
             return (dout,)

mindspore/ops/_grad_experimental/grad_debug_ops.py

@@ -15,6 +15,7 @@
 
 """Generate bprop for debug ops"""
 
+import mindspore.ops.functional as F
 from mindspore.ops import operations as P
 from mindspore.ops._grad_experimental.grad_base import bprop_getters
 
@@ -27,5 +28,9 @@ def get_bprop_insert_gradient_of(self):
     f = self.f
 
     def bprop(x, out, dout):
-        return (f(dout),)
+        fdout = f(dout)
+        if fdout is None:
+            dout = F.depend(dout, fdout)
+            return (dout,)
+        return (fdout,)
     return bprop

mindspore/ops/_grad_experimental/grad_inner_ops.py

@@ -23,6 +23,15 @@ from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
 from mindspore.ops._grad_experimental.grad_base import bprop_getters
 
 
+@bprop_getters.register("raise")
+def get_bprop_raise(self):
+    """Grad definition for `raise` operation."""
+    def bprop(x, y, z, out, dout):
+        return x, y, z
+
+    return bprop
+
+
 @bprop_getters.register(inner.ParallelResizeBilinear)
 def get_bprop_parallel_resize_bilinear(self):
     """Grad definition for `ParallelResizeBilinear` operation."""

mindspore/ops/_grad_experimental/grad_math_ops.py

@@ -657,7 +657,8 @@ def get_bprop_fft_with_size(self):
                 dx = rfft_fn(dout)
                 dx = reverse_branch(dx, onesided, dout_shape, offset_shape,
                                     output_type, dout, norm, inverse, signal_ndim, offset_size)
-        return (dx,)
+        return dx, zeros_like(signal_ndim), zeros_like(inverse), zeros_like(real), zeros_like(norm_enum), \
+               zeros_like(onesided), zeros_like(signal_sizes)
 
     return bprop
 

mindspore/ops/_op_impl/cpu/__init__.py

@@ -71,6 +71,7 @@ from .pyexecute import _pyexecute_cpu
 from .pyfunc import _pyfunc_cpu
 from .buffer_append import _buffer_append_cpu
 from .buffer_get import _buffer_get_cpu
+from .raise_op import _raise_cpu
 from .buffer_sample import _buffer_sample_cpu
 from .priority_replay_buffer import _prb_push_op_cpu
 from .priority_replay_buffer import _prb_sample_op_cpu

mindspore/ops/_op_impl/cpu/raise_op.py

@@ -0,0 +1,28 @@
+# 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.
+# ============================================================================
+"""storeattrgrad op"""
+from mindspore.ops.op_info_register import op_info_register, CpuRegOp, DataType
+
+raise_op_info = CpuRegOp("raise") \
+    .input(0, "x", "dynamic") \
+    .output(0, "y", "dynamic") \
+    .dtype_format(DataType.I32_Default, DataType.I32_Default) \
+    .get_op_info()
+
+
+@op_info_register(raise_op_info)
+def _raise_cpu():
+    """_getattrgrad_cpu cpu register"""
+    return

mindspore/ops/_vmap/vmap_array_ops.py

@@ -15,6 +15,7 @@
 
 """array_ops vmap impl."""
 from __future__ import absolute_import
+from enum import Enum
 
 import mindspore
 import mindspore.numpy as mnp
@@ -1488,16 +1489,19 @@ def get_meshgrid_vmap_rule(prim, axis_size):
     """VmapRule for `P.Meshgrid` operation."""
     if isinstance(prim, str):
         prim = Primitive(prim)
-    indexing = prim.indexing
 
-    def vmap_rule(*inputs_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, *inputs_bdim)
+    class Indexing(Enum):
+        ij = 0
+        xy = 1
+
+    def vmap_rule(inputs_bdim, indexing_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, inputs_bdim, indexing_bdim)
         if is_all_none:
             return result
 
         if not isinstance(inputs_bdim, (tuple)):
             _raise_value_error("The inputs of P.Meshgrid is not tuple.")
-        args = inputs_bdim[0]
+        args = inputs_bdim
         if len(args) <= 1:
             _raise_value_error(
                 "The input number of P.Meshgrid must be greater than 1.")
@@ -1518,7 +1522,9 @@ def get_meshgrid_vmap_rule(prim, axis_size):
         output_shape.insert(0, axis_size)
         ones_shape.insert(0, axis_size)
 
-        if indexing == "xy":
+        indexing, _ = indexing_bdim
+
+        if indexing == Indexing.xy.value:
             output_shape[1], output_shape[2] = output_shape[2], output_shape[1]
         shape = tuple(output_shape)
 
@@ -1531,7 +1537,7 @@ def get_meshgrid_vmap_rule(prim, axis_size):
         for each_arg in args:
             x, bdim = each_arg
             x = _bdim_at_front(x, bdim, axis_size)
-            shape_index = (1 - index) if (index <= 1 and indexing == "xy") else index
+            shape_index = (1 - index) if (index <= 1 and indexing == 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)
@@ -1889,10 +1895,6 @@ def get_slice_vmap_rule(prim, axis_size):
 @vmap_rules_getters.register(P.Squeeze)
 def get_squeeze_vmap_rule(prim, axis_size):
     """VmapRule for `Squeeze`."""
-    if hasattr(prim, 'axis'):
-        prim_axis = prim.axis
-    else:
-        prim_axis = None
 
     @_primexpr
     def move_axis(axes):
@@ -1911,27 +1913,26 @@ def get_squeeze_vmap_rule(prim, axis_size):
             new_axis += (i,)
         return new_axis
 
-    def vmap_rule(x_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, x_bdim)
+    def vmap_rule(x_bdim, axis_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, axis_bdim)
         if is_all_none:
             return result
 
         x, x_dim = x_bdim
+        axis, _ = axis_bdim
         x = _bdim_at_front(x, x_dim, axis_size)
 
-        if prim_axis is None:
+        if axis is None:
             if axis_size == 1:
                 new_axis = generate_all_axis_except_first(F.rank(x))
-                batch_squeeze = P.Squeeze(axis=new_axis)
-                out = batch_squeeze(x)
+                out = prim(x, new_axis)
                 return out, 0
 
-            out = prim(x)
+            out = prim(x, axis)
             return out, 0
 
-        new_axis = move_axis(prim_axis)
-        batch_squeeze = P.Squeeze(axis=new_axis)
-        out = batch_squeeze(x)
+        new_axis = move_axis(axis)
+        out = prim(x, new_axis)
         return out, 0
 
     return vmap_rule

mindspore/ops/_vmap/vmap_base.py

@@ -512,8 +512,6 @@ _ops_vmap_clone_prim_dict = {
     "ApplyAdagradV2": P.ApplyAdagradV2,
     "UniformCandidateSampler": UniformCandidateSampler,
     "UniqueWithPad": P.UniqueWithPad,
-    "CdistGrad": G.CdistGrad,
-    "Cdist": P.Cdist,
     "STFT": math_ops.STFT,
     "Conv2D": P.Conv2D,
     "Conv3D": P.Conv3D,

mindspore/ops/_vmap/vmap_grad_nn_ops.py

@@ -26,6 +26,7 @@ from mindspore.ops.function import _VmapGeneralRule
 from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_preprocess, _raise_value_error, \
     _bdim_at_front, _vmap_clone_prim, _bdim_at_any, _handle_broadcasting
 from mindspore.ops.auto_generate.gen_arg_handler import Format, Reduction
+from mindspore.ops import auto_generate as gen
 
 
 @vmap_rules_getters.register(G.NLLLossGrad)
@@ -225,33 +226,35 @@ def get_max_pool3d_grad_with_argmax_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
-@vmap_rules_getters.register(G.CdistGrad)
+@vmap_rules_getters.register(gen.CdistGrad)
 def get_cdist_grad_vmap_rule(prim, axis_size):
     """VmapRule for `cdist grad` operation."""
-    if hasattr(prim, 'batch_rank'):
-        batch_rank = prim.batch_rank + 1
+    if prim.has_label("batch_rank"):
+        batch_rank = prim.get_label("batch_rank") + 1
     else:
         batch_rank = 1
 
-    batch_prim = _vmap_clone_prim(prim)
-    batch_prim.add_prim_attr("batch_rank", batch_rank)
+    prim = prim.clone()
+    prim.set_label('batch_rank', batch_rank)
 
-    def vmap_rule(grad_bdim, x_bdim, y_bdim, cdist_bdim):
-        is_all_none, result = vmap_general_preprocess(prim,
-                                                      grad_bdim, x_bdim, y_bdim, cdist_bdim)
+    def vmap_rule(grad_bdim, x_bdim, y_bdim, cdist_bdim, p_bdim):
+        is_all_none, result = vmap_general_preprocess(
+            prim, grad_bdim, x_bdim, y_bdim, cdist_bdim, p_bdim
+        )
         if is_all_none:
             return result
         grad, grad_dim = grad_bdim
         x, x_dim = x_bdim
         y, y_dim = y_bdim
         cdist, cdist_dim = cdist_bdim
+        p, _ = p_bdim
 
         grad = _bdim_at_front(grad, grad_dim, axis_size)
         x = _bdim_at_front(x, x_dim, axis_size)
         y = _bdim_at_front(y, y_dim, axis_size)
         cdist = _bdim_at_front(cdist, cdist_dim, axis_size)
 
-        out = batch_prim(grad, x, y, cdist)
+        out = prim(grad, x, y, cdist, p)
         return out, 0
 
     return vmap_rule
@@ -673,10 +676,11 @@ def get_grid_sampler_grad_vmap_rule(prim, axis_size):
     else:
         _raise_value_error("The prim name must be `GridSampler2D` or `GridSampler3D`, but got {}.".format(prim_name))
 
-
-    def vmap_rule(grad_bdim, input_x_bdim, grid_bdim, interpolation_mode_bdim, padding_mode_bdim, align_corners_bdim):
+    def vmap_rule(grad_bdim, input_x_bdim, grid_bdim, interpolation_mode_bdim, padding_mode_bdim, align_corners_bdim,
+                  output_mask_bdim):
         is_all_none, result = vmap_general_preprocess(
-            prim, grad_bdim, input_x_bdim, grid_bdim, interpolation_mode_bdim, padding_mode_bdim, align_corners_bdim)
+            prim, grad_bdim, input_x_bdim, grid_bdim, interpolation_mode_bdim, padding_mode_bdim, align_corners_bdim,
+            output_mask_bdim)
         if is_all_none:
             return result
 
@@ -686,6 +690,7 @@ def get_grid_sampler_grad_vmap_rule(prim, axis_size):
         interpolation_mode, _ = interpolation_mode_bdim
         padding_mode, _ = padding_mode_bdim
         align_corners, _ = align_corners_bdim
+        output_mask, _ = output_mask_bdim
 
         grad = _bdim_at_front(grad, grad_dim, axis_size)
         grad_shape = F.shape(grad)
@@ -699,7 +704,8 @@ def get_grid_sampler_grad_vmap_rule(prim, axis_size):
         grid_shape = F.shape(grid)
         grid = F.reshape(grid, (-1,) + grid_shape[non_batch_dim_index:])
 
-        dx, dgrid = prim(grad, input_x, grid, interpolation_mode, padding_mode, align_corners)
+        dx, dgrid = prim(grad, input_x, grid, interpolation_mode,
+                         padding_mode, align_corners, output_mask)
         dx_shape = F.shape(dx)
         dx_return_shape = input_x_shape[:non_batch_dim_index] + dx_shape[non_batch_dim_index:]
         dx = F.reshape(dx, dx_return_shape)