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

mindspore/ops/_vmap/vmap_grad_math_ops.py

@@ -16,12 +16,15 @@
 """The vmap implement of grad operator corresponding to math_ops."""
 from __future__ import absolute_import
 
+import mindspore.numpy as mnp
 from mindspore.ops import functional as F
 from mindspore.ops.primitive import _primexpr
 from mindspore.ops.operations import _grad_ops as G
+from mindspore.ops import auto_generate as gen
 from mindspore.ops.function import _VmapGeneralRule
 from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_preprocess, _bdim_at_front, \
-    _handle_broadcasting, get_unary_grad_vmap_rule, _get_broadcasting_with_front_axis_additional_axis
+    _handle_broadcasting, get_unary_grad_vmap_rule, _get_broadcasting_with_front_axis_additional_axis, \
+    _broadcast_by_axis
 
 
 @vmap_rules_getters.register('MaximumGrad')
@@ -29,8 +32,8 @@ from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_prepr
 def get_broadcast_binary_op_grad_vmap_rule(prim, axis_size):
     """VmapRule for grad of binary operations with broadcasting"""
     broadcast_binary_op_grad_map = {
-        "MinimumGrad": G.MinimumGrad,
-        "MaximumGrad": G.MaximumGrad
+        "MinimumGrad": gen.MinimumGrad,
+        "MaximumGrad": gen.MaximumGrad
     }
 
     if isinstance(prim, str):
@@ -49,14 +52,17 @@ def get_broadcast_binary_op_grad_vmap_rule(prim, axis_size):
                 return y_shape
         return g_shape
 
-    def vmap_rule(x_bdim, y_bdim, grad_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, x_bdim, y_bdim, grad_bdim)
+    def vmap_rule(x_bdim, y_bdim, grad_bdim, grad_x_bdim, grad_y_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, y_bdim, grad_bdim, grad_x_bdim,
+                                                      grad_y_bdim)
         if is_all_none:
             return result
 
         x, x_dim = x_bdim
         y, y_dim = y_bdim
         g, g_dim = grad_bdim
+        g_x, _ = grad_x_bdim
+        g_y, _ = grad_y_bdim
 
         x_shape = F.shape(x)
         y_shape = F.shape(y)
@@ -65,7 +71,7 @@ def get_broadcast_binary_op_grad_vmap_rule(prim, axis_size):
         is_dim_ok = x_dim == y_dim and x_dim == g_dim
         is_shape_ok = x_shape == y_shape and x_shape == g_shape
         if is_dim_ok and is_shape_ok:
-            dx, dy = prim(x, y, g)
+            dx, dy = prim(x, y, g, g_x, g_y)
             return (dx, x_dim), (dy, y_dim)
 
         x = _bdim_at_front(x, x_dim, axis_size)
@@ -84,7 +90,7 @@ def get_broadcast_binary_op_grad_vmap_rule(prim, axis_size):
         x_axis_for_reduce = _get_broadcasting_with_front_axis_additional_axis(x_shape, longest_shape)
         y_axis_for_reduce = _get_broadcasting_with_front_axis_additional_axis(y_shape, longest_shape)
 
-        dx, dy = prim(x, y, g)
+        dx, dy = prim(x, y, g, g_x, g_y)
         if x_axis_for_reduce:
             dx = F.reduce_sum(dx, x_axis_for_reduce)
 
@@ -95,7 +101,6 @@ def get_broadcast_binary_op_grad_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
-@vmap_rules_getters.register(G.MaximumGradGrad)
 @vmap_rules_getters.register(G.MinimumGradGrad)
 def get_broadcast_grad_grad_vmap_rule(prim, axis_size):
     """VmapRule for GradGrad operations with broadcasting."""
@@ -142,6 +147,55 @@ def get_broadcast_grad_grad_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(G.MaximumGradGrad)
+def get_maximum_grad_grad_vmap_rule(prim, axis_size):
+    """VmapRule for GradGrad operations with broadcasting."""
+
+    def vmap_rule(x1_bdim, x2_bdim, dx1_bdim, dx2_bdim, grad_x_bdim, grad_y_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x1_bdim, x2_bdim, dx1_bdim, dx2_bdim, grad_x_bdim,
+                                                      grad_y_bdim)
+        if is_all_none:
+            return result
+
+        x1, x1_dim = x1_bdim
+        x2, x2_dim = x2_bdim
+        dx1, dx1_dim = dx1_bdim
+        dx2, dx2_dim = dx2_bdim
+        grad_x, _ = grad_x_bdim
+        grad_y, _ = grad_y_bdim
+        x1_shape = F.shape(x1)
+        x2_shape = F.shape(x2)
+        dx1_shape = F.shape(dx1)
+        dx2_shape = F.shape(dx2)
+
+        is_dim_ok = x1_dim == x2_dim and dx1_dim == dx2_dim and x1_dim == dx1_dim
+        is_shape_ok = x1_shape == x2_shape and dx1_shape == dx2_shape
+        if is_dim_ok and is_shape_ok:
+            sopd_x1, sopd_x2, sopd_grad = prim(x1, x2, dx1, dx2, grad_x, grad_y)
+            return (sopd_x1, x1_dim), (sopd_x2, x1_dim), (sopd_grad, x1_dim)
+
+        if F.rank(x1):
+            x1 = _bdim_at_front(x1, x1_dim, 1)
+        if F.rank(x2):
+            x2 = _bdim_at_front(x2, x2_dim, 1)
+        if F.rank(dx1):
+            dx1 = _bdim_at_front(dx1, dx2_dim, 1)
+        if F.rank(dx2):
+            dx2 = _bdim_at_front(dx2, dx2_dim, 1)
+        x1_shape = F.shape(x1)
+        x2_shape = F.shape(x2)
+        dx1_shape = F.shape(dx1)
+        dx2_shape = F.shape(dx2)
+        x1 = _handle_broadcasting(x1, x1_shape, x2_shape)
+        x2 = _handle_broadcasting(x2, x2_shape, x1_shape)
+        dx1 = _handle_broadcasting(dx1, dx1_shape, dx2_shape)
+        dx2 = _handle_broadcasting(dx2, dx2_shape, dx1_shape)
+        sopd_x1, sopd_x2, sopd_grad = prim(x1, x2, dx1, dx2, grad_x, grad_y)
+        return (sopd_x1, 0), (sopd_x2, 0), (sopd_grad, 0)
+
+    return vmap_rule
+
+
 @vmap_rules_getters.register(G.MedianGrad)
 def get_median_grad_vmap_rule(prim, axis_size):
     """VmapRule for MedianGrad."""
@@ -180,9 +234,41 @@ def get_median_grad_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(G.LogitGrad)
+def get_logit_grad_vmap_rule(prim_func, axis_size):
+    """VmapRule for `LogitGrad`."""
+    if isinstance(prim_func, str):
+        raise TypeError("prim_func can't be str.")
+
+    def vmap_rule(grad_bdim, x_bdim, eps_bdim):
+        grad, grad_dim = grad_bdim
+        x, x_dim = x_bdim
+        eps, _ = eps_bdim
+        x_shape = F.shape(x)
+        grad_shape = F.shape(grad)
+        if x_dim == grad_dim and x_shape == grad_shape:
+            out = prim_func(grad, x, eps)
+            return (out, x_dim)
+
+        # This branch means (x_dim is None) and (grad_dim is not None).
+        if x_dim is None:
+            x = _broadcast_by_axis(x, grad_dim, axis_size)
+            out_dim = grad_dim
+        # This branch means (x_dim is not None) and (grad_dim is None).
+        elif grad_dim is None:
+            grad = _broadcast_by_axis(grad, x_dim, axis_size)
+            out_dim = x_dim
+        # This branch means (x_dim is not None) and (grad_dim is not None).
+        else:
+            grad = mnp.moveaxis(grad, grad_dim, x_dim)
+            out_dim = x_dim
+        out = prim_func(grad, x, eps)
+        return out, out_dim
+
+    return vmap_rule
+
 # UnaryGrad vmap
 get_unary_grad_vmap_rule = vmap_rules_getters.register(G.InvGrad)(get_unary_grad_vmap_rule)
-get_unary_grad_vmap_rule = vmap_rules_getters.register(G.LogitGrad)(get_unary_grad_vmap_rule)
 get_unary_grad_vmap_rule = vmap_rules_getters.register('AbsGrad')(get_unary_grad_vmap_rule)
 get_unary_grad_vmap_rule = vmap_rules_getters.register('ReciprocalGrad')(get_unary_grad_vmap_rule)
 get_unary_grad_vmap_rule = vmap_rules_getters.register('SqrtGrad')(get_unary_grad_vmap_rule)

mindspore/ops/_vmap/vmap_grad_nn_ops.py

@@ -21,12 +21,11 @@ from functools import reduce
 import mindspore.numpy as mnp
 from mindspore.ops.operations import _grad_ops as G
 from mindspore.ops import functional as F
-from mindspore.ops import constexpr
 from mindspore.ops.primitive import _primexpr
-from mindspore.ops.primitive import Primitive
 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
 
 
 @vmap_rules_getters.register(G.NLLLossGrad)
@@ -45,15 +44,10 @@ def get_nll_loss_grad_vmap_rule(prim, axis_size):
             lambda x, y: x * y, shape if keep_dim == 0 else shape[:-keep_dim])
         return (new_batch_size,) if keep_dim == 0 else (new_batch_size, *shape[-keep_dim:])
 
-    if isinstance(prim, str):
-        prim = Primitive(prim)
-        reduction_type = "none"
-    else:
-        reduction_type = prim.reduction
-
-    def vmap_rule(x_bdim, loss_grad_bdim, target_bdim, weight_bdim, total_weight_bdim):
-        is_all_none, result = vmap_general_preprocess(
-            prim, x_bdim, loss_grad_bdim, target_bdim, weight_bdim, total_weight_bdim)
+    def vmap_rule(x_bdim, loss_grad_bdim, target_bdim, weight_bdim, total_weight_bdim, reduction_bdim,
+                  ignore_index_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, loss_grad_bdim, target_bdim, weight_bdim,
+                                                      total_weight_bdim, reduction_bdim, ignore_index_bdim)
         if is_all_none:
             return result
 
@@ -62,6 +56,8 @@ def get_nll_loss_grad_vmap_rule(prim, axis_size):
         target, target_dim = target_bdim
         weight, w_dim = weight_bdim
         total_weight, tw_dim = total_weight_bdim
+        reduction, _ = reduction_bdim
+        ignore_index, _ = ignore_index_bdim
 
         x_shape = F.shape(x)
         loss_grad_shape = F.shape(loss_grad)
@@ -71,10 +67,10 @@ def get_nll_loss_grad_vmap_rule(prim, axis_size):
         if w_dim is not None or tw_dim is not None:
             _raise_value_error("The source axis of weight and total_weight in `NLLLossGrad` must be None for now, "
                                "but got {} and {}.".format(w_dim, tw_dim))
-        if lg_dim is not None and reduction_type != "none":
+        if lg_dim is not None and reduction != Reduction.NONE:
             _raise_value_error("The source axis of loss_grad in `NLLLossGrad` can be not None "
                                "just when reduction type is none for vmap, "
-                               "but reduction type is {}.".format(reduction_type))
+                               "but reduction type is {}.".format((Reduction(reduction)).name))
 
         # If stacked, move vmap_dim to first dim and reshape to right shape.
         if x_dim is not None:
@@ -97,7 +93,7 @@ def get_nll_loss_grad_vmap_rule(prim, axis_size):
                 target_shape = F.shape(target)
             target = F.reshape(target, _get_reshape_shape(target_shape))
 
-        out = prim(x, loss_grad, target, weight, total_weight)
+        out = prim(x, loss_grad, target, weight, total_weight, reduction, ignore_index)
         output = F.reshape(out, x_shape)
         out_dim = 0
         return output, out_dim
@@ -106,9 +102,8 @@ def get_nll_loss_grad_vmap_rule(prim, axis_size):
 
 
 @vmap_rules_getters.register(G.MaxPoolGrad)
-@vmap_rules_getters.register(G.AvgPoolGrad)
-def get_avg_pool_grad_vmap_rule(prim, axis_size):
-    """VmapRule for `AvgPoolGrad`."""
+def get_max_pool_grad_vmap_rule(prim, axis_size):
+    """VmapRule for `MaxPoolGrad`."""
     chw_reverse_index = -3
 
     def vmap_rule(x_bdim, y_bdim, dy_bdim):
@@ -135,6 +130,41 @@ def get_avg_pool_grad_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(G.AvgPoolGrad)
+def get_avg_pool_grad_vmap_rule(prim, axis_size):
+    """VmapRule for `AvgPoolGrad`."""
+    chw_reverse_index = -3
+
+    def vmap_rule(x_bdim, y_bdim, dy_bdim, kernel_size_bdim,
+                  strides_bdim, pad_mode_bdim, data_format_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, y_bdim, dy_bdim, kernel_size_bdim,
+                                                      strides_bdim, pad_mode_bdim, data_format_bdim)
+        if is_all_none:
+            return result
+
+        x, x_dim = x_bdim
+        y, y_dim = y_bdim
+        dy, dy_dim = dy_bdim
+        kernel_size, _ = kernel_size_bdim
+        strides, _ = strides_bdim
+        pad_mode, _ = pad_mode_bdim
+        data_format, _ = data_format_bdim
+        x = _bdim_at_front(x, x_dim, axis_size)
+        y = _bdim_at_front(y, y_dim, axis_size)
+        dy = _bdim_at_front(dy, dy_dim, axis_size)
+        x_shape = F.shape(x)
+        y_shape = F.shape(y)
+        dy_shape = F.shape(dy)
+        x = F.reshape(x, (-1,) + x_shape[chw_reverse_index:])
+        y = F.reshape(y, (-1,) + y_shape[chw_reverse_index:])
+        dy = F.reshape(dy, (-1,) + dy_shape[chw_reverse_index:])
+        out = prim(x, y, dy, kernel_size, strides, pad_mode, data_format)
+        out = F.reshape(out, x_shape)
+        return out, 0
+
+    return vmap_rule
+
+
 @vmap_rules_getters.register(G.AvgPool3DGrad)
 def get_avg_pool3d_grad_vmap_rule(prim, axis_size):
     """VmapRule for `AvgPool3DGrad`."""
@@ -276,16 +306,22 @@ def get_adaptive_avgpool2d_vmap_rule(prim, axis_size):
 @vmap_rules_getters.register(G.BatchNormGradGrad)
 def get_batchnorm_grad_grad_vmap_rule(prim, axis_size):
     """VmapRule for `BatchNormGradGrad` operation."""
-    data_format = prim.format
+    NCHW = Format.NCHW
 
     def vmap_rule(x_bdim, dy_bdim, scale_bdim, mean_bdim, variance_bdim, dout_dx_bdim,
-                  dout_dscale_bdim, dout_dbias_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, x_bdim, dy_bdim, scale_bdim, mean_bdim, variance_bdim,
-                                                      dout_dx_bdim, dout_dscale_bdim, dout_dbias_bdim)
+                  dout_dscale_bdim, dout_dbias_bdim, is_training_bdim, epsilon_bdim, data_format_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, dy_bdim, scale_bdim, mean_bdim,
+                                                      variance_bdim, dout_dx_bdim, dout_dscale_bdim,
+                                                      dout_dbias_bdim, is_training_bdim, epsilon_bdim,
+                                                      data_format_bdim)
         if is_all_none:
             return result
 
-        dst_dim = 1 if data_format == "NCHW" else 3
+        is_training, _ = is_training_bdim
+        epsilon, _ = epsilon_bdim
+        data_format, _ = data_format_bdim
+
+        dst_dim = 1 if data_format == NCHW else 3
         x = _bdim_at_any(*x_bdim, dst_dim, axis_size)
         dy = _bdim_at_any(*dy_bdim, dst_dim, axis_size)
         dout_dx = _bdim_at_any(*dout_dx_bdim, dst_dim, axis_size)
@@ -298,11 +334,11 @@ def get_batchnorm_grad_grad_vmap_rule(prim, axis_size):
 
         x_shape = x.shape
         scale_shape = scale.shape
-        shape = (x_shape[0], -1,) + x_shape[3:] if data_format == "NCHW" else x_shape[:-2] + (-1,)
+        shape = (x_shape[0], -1,) + x_shape[3:] if data_format == NCHW else x_shape[:-2] + (-1,)
         dx, ddy, dscale = prim(x.reshape(shape), dy.reshape(shape), scale.flatten(), mean.flatten(),
                                variance.flatten(), dout_dx.reshape(shape), dout_dscale.flatten(),
-                               dout_dbias.flatten())
-        pos = 1 if data_format == "NCHW" else 3
+                               dout_dbias.flatten(), is_training, epsilon, data_format)
+        pos = 1 if data_format == NCHW else 3
         return (dx.reshape(x_shape), pos), (ddy.reshape(x_shape), pos), (dscale.reshape(scale_shape), 0)
     return vmap_rule
 
@@ -347,25 +383,29 @@ def get_batchnorm_grad_vmap_rule(prim, axis_size):
     """VmapRule for `BatchNormGrad` operation."""
     bn_min_dim = 3
     bn_max_dim = 5
-    data_format = prim.data_format
     prim_name = prim.name
-    if data_format == "NHWC":
-        batchnorm_grad_nhwc_vmap = _VmapGeneralRule(prim, axis_size)
+    NHWC = Format.NHWC
 
-    def vmap_rule(grad_bdim, x_bdim, scale_bdim, rsv_1_bdim, rsv_2_bdim, rsv_3_bdim):
-        is_all_none, result = \
-            vmap_general_preprocess(prim, grad_bdim, x_bdim, scale_bdim, rsv_1_bdim, rsv_2_bdim, rsv_3_bdim)
+    def vmap_rule(grad_bdim, x_bdim, scale_bdim, rsv_1_bdim, rsv_2_bdim,
+                  rsv_3_bdim, training_bdim, epsilon_bdim, format_bdim):
+        is_all_none, result = vmap_general_preprocess(
+            prim, grad_bdim, x_bdim, scale_bdim, rsv_1_bdim, rsv_2_bdim,
+            rsv_3_bdim, training_bdim, epsilon_bdim, format_bdim)
         if is_all_none:
             return result
-        if data_format == "NHWC":
-            # BatchNormGrad with NHWC format is a GPU backend operation and not supported for now.
-            return batchnorm_grad_nhwc_vmap(grad_bdim, x_bdim, scale_bdim, rsv_1_bdim, rsv_2_bdim, rsv_3_bdim)
         grad, grad_dim = grad_bdim
         input_x, input_x_dim = x_bdim
         scale, scale_dim = scale_bdim
         rsv_1, rsv_1_dim = rsv_1_bdim
         rsv_2, rsv_2_dim = rsv_2_bdim
         rsv_3, rsv_3_dim = rsv_3_bdim
+        training, _ = training_bdim
+        epsilon, _ = epsilon_bdim
+        data_format, _ = format_bdim
+        if data_format == NHWC:
+            batchnorm_grad_nhwc_vmap = _VmapGeneralRule(prim, axis_size)
+            # BatchNormGrad with NHWC format is a GPU backend operation and not supported for now.
+            return batchnorm_grad_nhwc_vmap(grad_bdim, x_bdim, scale_bdim, rsv_1_bdim, rsv_2_bdim, rsv_3_bdim)
         x_ndim = F.rank(input_x)
         if x_ndim not in (bn_min_dim, bn_max_dim):
             raise ValueError("The dim of `input_x` in `{}` must be equal to {} or {}, "
@@ -387,7 +427,8 @@ def get_batchnorm_grad_vmap_rule(prim, axis_size):
         rsv_1 = rsv_1.flatten()
         rsv_2 = rsv_2.flatten()
         rsv_3 = rsv_3.flatten()
-        grad_x, grad_scale, grad_offset = prim(grad, input_x, scale, rsv_1, rsv_2, rsv_3)
+        grad_x, grad_scale, grad_offset =\
+            prim(grad, input_x, scale, rsv_1, rsv_2, rsv_3, training, epsilon, data_format)
         grad_x = F.reshape(grad_x, x_shape)
         grad_scale = F.reshape(grad_scale, other_shape)
         grad_offset = F.reshape(grad_offset, other_shape)
@@ -553,7 +594,6 @@ def get_mirror_pad_grad_grad_vmap_rule(prim, axis_size):
 @vmap_rules_getters.register('LayerNormGrad')
 def get_layernormgrad_vmap_rule(prim, axis_size):
     """VmapRule for `LayerNormGrad` operation."""
-    @constexpr
     def process_attr_axis(prim_attr_axis):
         if prim_attr_axis < 0:
             return prim_attr_axis
@@ -571,13 +611,11 @@ def get_layernormgrad_vmap_rule(prim, axis_size):
     def get_logical_shape(var_shape):
         return var_shape[1:]
 
-    norm_axis = process_attr_axis(prim.begin_norm_axis)
-    params_axis = process_attr_axis(prim.begin_params_axis)
-    batch_prim = G.LayerNormGrad(norm_axis, params_axis)
     eps = 1e-12
 
-    def vmap_rule(x_bdim, dy_bdim, var_bdim, mean_bdim, gamma_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, x_bdim, dy_bdim, var_bdim, mean_bdim, gamma_bdim)
+    def vmap_rule(x_bdim, dy_bdim, var_bdim, mean_bdim, gamma_bdim, begin_norm_axis_bdim, begin_params_axis_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, dy_bdim, var_bdim, mean_bdim, gamma_bdim,
+                                                      begin_norm_axis_bdim, begin_params_axis_bdim)
         if is_all_none:
             return result
 
@@ -586,6 +624,11 @@ def get_layernormgrad_vmap_rule(prim, axis_size):
         var, var_dim = var_bdim
         mean, mean_dim = mean_bdim
         gamma, gamma_dim = gamma_bdim
+        begin_norm_axis, _ = begin_norm_axis_bdim
+        begin_params_axis, _ = begin_params_axis_bdim
+
+        begin_norm_axis = process_attr_axis(begin_norm_axis)
+        begin_params_axis = process_attr_axis(begin_params_axis)
 
         x = _bdim_at_front(x, x_dim, axis_size)
         dy = _bdim_at_front(dy, dy_dim, axis_size)
@@ -594,7 +637,7 @@ def get_layernormgrad_vmap_rule(prim, axis_size):
         gamma = _bdim_at_front(gamma, gamma_dim, axis_size)
 
         dy_shape = F.shape(dy)
-        batch_params_reduce_axes = get_batch_params_reduce_axes(params_axis, dy_shape)
+        batch_params_reduce_axes = get_batch_params_reduce_axes(begin_norm_axis, dy_shape)
 
         if not batch_params_reduce_axes:
             d_beta = dy
@@ -612,7 +655,7 @@ def get_layernormgrad_vmap_rule(prim, axis_size):
         dy = dy * gamma
         gamma_logical_shape = get_logical_shape(gamma_shape)
         ones_like_gamma = F.ones(gamma_logical_shape, F.dtype(gamma))
-        dx, _, _ = batch_prim(x, dy, var, mean, ones_like_gamma)
+        dx, _, _ = prim(x, dy, var, mean, ones_like_gamma, begin_norm_axis, begin_params_axis)
 
         return (dx, 0), (d_gamma, 0), (d_beta, 0)
     return vmap_rule
@@ -625,17 +668,24 @@ def get_grid_sampler_grad_vmap_rule(prim, axis_size):
     prim_name = prim.name
     if prim_name == "GridSampler2DGrad":
         non_batch_dim_index = -3
-    else:
+    elif prim_name == "GridSampler3DGrad":
         non_batch_dim_index = -4
+    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):
-        is_all_none, result = vmap_general_preprocess(prim, grad_bdim, input_x_bdim, grid_bdim)
+    def vmap_rule(grad_bdim, input_x_bdim, grid_bdim, interpolation_mode_bdim, padding_mode_bdim, align_corners_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)
         if is_all_none:
             return result
 
         grad, grad_dim = grad_bdim
         input_x, input_x_dim = input_x_bdim
         grid, grid_dim = grid_bdim
+        interpolation_mode, _ = interpolation_mode_bdim
+        padding_mode, _ = padding_mode_bdim
+        align_corners, _ = align_corners_bdim
 
         grad = _bdim_at_front(grad, grad_dim, axis_size)
         grad_shape = F.shape(grad)
@@ -649,7 +699,7 @@ 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)
+        dx, dgrid = prim(grad, input_x, grid, interpolation_mode, padding_mode, align_corners)
         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)
@@ -661,8 +711,7 @@ def get_grid_sampler_grad_vmap_rule(prim, axis_size):
 
 
 @vmap_rules_getters.register(G.UpsampleNearest3DGrad)
-@vmap_rules_getters.register(G.UpsampleTrilinear3DGrad)
-def get_upsample_grad_vmap_rule(prim, axis_size):
+def get_upsample_nearesst3d_grad_vmap_rule(prim, axis_size):
     """VmapRule for `UpsampleNearest3DGrad` and `UpsampleTrilinear3DGrad`."""
     cdhw_reverse_index = -4
 
@@ -696,26 +745,62 @@ def get_upsample_grad_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(G.UpsampleTrilinear3DGrad)
+def get_upsample_trilinear3d_grad_vmap_rule(prim, axis_size):
+    """VmapRule for `UpsampleNearest3DGrad` and `UpsampleTrilinear3DGrad`."""
+    cdhw_reverse_index = -4
+
+    def vmap_rule(grad_bdim, isize_bdim, osize_bdim, scales_bdim, align_corners_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, grad_bdim)
+        if is_all_none:
+            return result
+
+        grad, grad_dim = grad_bdim
+        grad = _bdim_at_front(grad, grad_dim, axis_size)
+        grad_shape = F.shape(grad)
+        input_shape = (-1,) + grad_shape[cdhw_reverse_index:]
+        grad = F.reshape(grad, input_shape)
+        real_in_shape = F.shape(grad)
+
+        isize, isize_dim = isize_bdim
+        osize, osize_dim = osize_bdim
+        scales, scales_dim = scales_bdim
+        align_corners, align_corners_dim = align_corners_bdim
+        check_args = [isize_dim, osize_dim, scales_dim, align_corners_dim]
+        is_all_none = True
+        for arg in check_args:
+            is_all_none = is_all_none and arg is None
+        if not is_all_none:
+            _raise_value_error(
+                "The source axis of `input_size`, `output_size` and `scales` must be None, but got {0}, {1} and {2}."
+                .format(isize_dim, osize_dim, scales_dim))
+        # update batch dimension of input_size
+        new_isize = (real_in_shape[0],) + isize[1:]
+
+        out = prim(grad, new_isize, osize, scales, align_corners)
+        out_shape = F.shape(out)
+        real_out_shape = grad_shape[:cdhw_reverse_index] + out_shape[cdhw_reverse_index:]
+        out = F.reshape(out, real_out_shape)
+        return out, 0
+    return vmap_rule
+
+
 @vmap_rules_getters.register(G.LogSoftmaxGrad)
-def get_log_softmax_vmap_rule(prim, axis_size):
+def get_log_softmax_grad_vmap_rule(prim_func, axis_size):
     """VmapRule for 'LogSoftmaxGrad' operation."""
-    if isinstance(prim, str):
-        axis = -1
-    else:
-        axis = prim.axis
-
-    def vmap_rule(x_bdim, grad_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, x_bdim)
+    def vmap_rule(x_bdim, grad_bdim, axis_bdim):
+        is_all_none, result = vmap_general_preprocess(prim_func, x_bdim, axis_bdim)
         if is_all_none:
             return result
         x, x_dim = x_bdim
         grad, _ = grad_bdim
+        axis, _ = axis_bdim
         x_ndim = F.rank(x) - 1
 
         batch_axis = axis + x_ndim if axis < 0 else axis
         batch_axis = batch_axis if batch_axis < x_dim else batch_axis + 1
 
-        dx = G.LogSoftmaxGrad(axis=batch_axis)(x, grad)
+        dx = prim_func(x, grad, batch_axis)
         return dx, x_dim
 
     return vmap_rule

mindspore/ops/_vmap/vmap_image_ops.py

@@ -27,18 +27,20 @@ from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_prepr
 
 
 @vmap_rules_getters.register(IMG.ResizeBilinearV2)
-@vmap_rules_getters.register(IMG.ResizeLinear1D)
-def get_resize_dynamic_input_rule(prim, axis_size):
-    """VmapRule for `Resize` operation."""
+def get_resize_bilinear_v2_vmap_rule(prim, axis_size):
+    """VmapRule for `ResizeBilinearV2` operation."""
     prim_name = prim.name
 
-    def vmap_rule(x_bdim, size_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, x_bdim, size_bdim)
+    def vmap_rule(x_bdim, size_bdim, align_corners_bdim, half_pixel_centers_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, size_bdim,
+                                                      align_corners_bdim, half_pixel_centers_bdim)
         if is_all_none:
             return result
 
         x, x_dim = x_bdim
         size, size_dim = size_bdim
+        align_corners, _ = align_corners_bdim
+        half_pixel_centers, _ = half_pixel_centers_bdim
         if size_dim is not None:
             _raise_value_error("The source axis of `size` in `{}` must be None, "
                                "but got {}.".format(prim_name, size_dim))
@@ -46,11 +48,39 @@ def get_resize_dynamic_input_rule(prim, axis_size):
         x = _bdim_at_front(x, x_dim, axis_size)
         x_shape = F.shape(x)
         # (b, n, c, i_h, i_w) -> (b*n, c, i_h, i_w) for 4-D input
-        # (b, n, c, i_w) -> (b*n, c, i_w) for 3-D input
         x = F.reshape(x, (-1,) + x_shape[2:])
-        out = prim(x, size)
+        out = prim(x, size, align_corners, half_pixel_centers)
         out_shape = F.shape(out)
         # (b*n, c, o_h, o_w) -> (b, n, c, o_h, o_w) for 4-D input
+        out = F.reshape(out, x_shape[:2] + out_shape[1:])
+        return out, 0
+
+    return vmap_rule
+
+
+@vmap_rules_getters.register(IMG.ResizeLinear1D)
+def get_resize_linear_1d_vmap_rule(prim, axis_size):
+    """VmapRule for `ResizeLinear1D` operation."""
+    prim_name = prim.name
+
+    def vmap_rule(x_bdim, size_bdim, coordinate_transformation_mode_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, size_bdim, coordinate_transformation_mode_bdim)
+        if is_all_none:
+            return result
+
+        x, x_dim = x_bdim
+        size, size_dim = size_bdim
+        coordinate_transformation_mode, _ = coordinate_transformation_mode_bdim
+        if size_dim is not None:
+            _raise_value_error("The source axis of `size` in `{}` must be None, "
+                               "but got {}.".format(prim_name, size_dim))
+
+        x = _bdim_at_front(x, x_dim, axis_size)
+        x_shape = F.shape(x)
+        # (b, n, c, i_w) -> (b*n, c, i_w) for 3-D input
+        x = F.reshape(x, (-1,) + x_shape[2:])
+        out = prim(x, size, coordinate_transformation_mode)
+        out_shape = F.shape(out)
         # (b*n, c, o_w) -> (b, n, c, o_w) for 3-D input
         out = F.reshape(out, x_shape[:2] + out_shape[1:])
         return out, 0
@@ -59,28 +89,55 @@ def get_resize_dynamic_input_rule(prim, axis_size):
 
 
 @vmap_rules_getters.register(G.ResizeBilinearGrad)
+def get_resize_bilinear_grad_vmap_rule(prim, axis_size):
+    """VmapRule for `ResizeBilinearGrad` operation."""
+
+    def vmap_rule(grad_bdim, img_bdim, align_corners_bdim, half_pixel_centers_bdim):
+        is_all_none, result = vmap_general_preprocess(grad_bdim, img_bdim, align_corners_bdim, half_pixel_centers_bdim)
+        if is_all_none:
+            return result
+
+        grad, grad_dim = grad_bdim
+        img, img_dim = img_bdim
+        align_corners, _ = align_corners_bdim
+        half_pixel_centers, _ = half_pixel_centers_bdim
+
+        grad = _bdim_at_front(grad, grad_dim, axis_size)
+        grad_shape = F.shape(grad)
+        img = _bdim_at_front(img, img_dim, axis_size)
+        img_shape = F.shape(img)
+        # (b, n, c, i_h, i_w) -> (b*n, c, i_h, i_w) for 4-D input
+        grad = F.reshape(grad, (-1,) + grad_shape[2:])
+        img = F.reshape(img, (-1,) + img_shape[2:])
+        out = prim(grad, img, align_corners, half_pixel_centers)
+        # (b*n, c, o_h, o_w) -> (b, n, c, o_h, o_w) for 4-D input
+        out = F.reshape(out, img_shape)
+        return out, 0
+
+    return vmap_rule
+
+
 @vmap_rules_getters.register(G.ResizeLinear1DGrad)
-def get_resize_grad_dynamic_rule(prim, axis_size):
-    """VmapRule for `ResizeGrad` operation."""
+def get_resize_linear_1d_grad_vmap_rule(prim, axis_size, coordinate_transformation_mode_bdim):
+    """VmapRule for `ResizeLinear1DGrad` operation."""
 
     def vmap_rule(grad_bdim, img_bdim):
-        is_all_none, result = vmap_general_preprocess(grad_bdim, img_bdim)
+        is_all_none, result = vmap_general_preprocess(grad_bdim, img_bdim, coordinate_transformation_mode_bdim)
         if is_all_none:
             return result
 
         grad, grad_dim = grad_bdim
         img, img_dim = img_bdim
+        coordinate_transformation_mode, _ = coordinate_transformation_mode_bdim
 
         grad = _bdim_at_front(grad, grad_dim, axis_size)
         grad_shape = F.shape(grad)
         img = _bdim_at_front(img, img_dim, axis_size)
         img_shape = F.shape(img)
-        # (b, n, c, i_h, i_w) -> (b*n, c, i_h, i_w) for 4-D input
         # (b, n, c, i_w) -> (b*n, c, i_w) for 3-D input
         grad = F.reshape(grad, (-1,) + grad_shape[2:])
         img = F.reshape(img, (-1,) + img_shape[2:])
-        out = prim(grad, img)
-        # (b*n, c, o_h, o_w) -> (b, n, c, o_h, o_w) for 4-D input
+        out = prim(grad, img, coordinate_transformation_mode)
         # (b*n, c, o_w) -> (b, n, c, o_w) for 3-D input
         out = F.reshape(out, img_shape)
         return out, 0