mindspore 1.10.0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops/_vmap/vmap_grad_nn_ops.py

@@ -14,16 +14,17 @@
 # ============================================================================
 
 """The vmap implement of grad operator corresponding to nn_ops."""
-
 from __future__ import absolute_import
+
 from __future__ import division
 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.composite import _VmapGeneralRule
+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, _vmap_update_prim_attr, _bdim_at_any, _handle_broadcasting
 
@@ -38,7 +39,7 @@ def get_nll_loss_grad_vmap_rule(prim, axis_size):
     2. And weight only support shape as (C,), while total_weight should be a scalar.
     """
 
-    @constexpr
+    @_primexpr
     def _get_reshape_shape(shape, keep_dim=0):
         new_batch_size = reduce(
             lambda x, y: x * y, shape if keep_dim == 0 else shape[:-keep_dim])
@@ -104,6 +105,7 @@ def get_nll_loss_grad_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@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`."""
@@ -225,11 +227,15 @@ def get_cdist_grad_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(G.AdaptiveMaxPool3DGrad)
 @vmap_rules_getters.register(G.AdaptiveMaxPool2DGrad)
 def get_adaptive_avgpool2d_vmap_rule(prim, axis_size):
-    """VmapRule for `AdaptiveMaxPool2DGrad` operation."""
+    """VmapRule for `AdaptiveMaxPool2DGrad` and `AdaptiveMaxPool3DGrad` operation."""
     chw_reverse_index = -3
-    hw_reverse_index = -2
+    if prim.name == "AdaptiveMaxPool2DGrad":
+        hw_reverse_index = -2
+    else:
+        hw_reverse_index = -3
 
     def vmap_rule(ygrad_bdim, x_bdim, max_index_bdim):
         is_all_none, result = vmap_general_preprocess(prim, ygrad_bdim, x_bdim, max_index_bdim)
@@ -352,7 +358,7 @@ def get_batchnorm_grad_vmap_rule(prim, axis_size):
         if is_all_none:
             return result
         if data_format == "NHWC":
-            #BatchNormGrad with NHWC format is a GPU backend operation and not supported for now.
+            # 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
@@ -392,8 +398,9 @@ def get_batchnorm_grad_vmap_rule(prim, axis_size):
 
 @vmap_rules_getters.register(G.MaxPoolGradGrad)
 @vmap_rules_getters.register(G.MaxPoolGradGradWithArgmax)
+@vmap_rules_getters.register(G.MaxPoolGradWithArgmaxV2)
 def get_maxpool_grad_grad_vmap_rule(prim, axis_size):
-    """VmapRule for `MaxPoolGradGrad` and `MaxPoolGradGradWithArgmax`."""
+    """VmapRule for `MaxPoolGradGrad`, `MaxPoolGradGradWithArgmax` and `MaxPoolGradWithArgmaxV2`."""
     chw_reverse_index = -3
 
     def vmap_rule(in0_bdim, in1_bdim, in2_bdim):
@@ -552,7 +559,7 @@ def get_layernormgrad_vmap_rule(prim, axis_size):
             return prim_attr_axis
         return prim_attr_axis + 1
 
-    @constexpr
+    @_primexpr
     def get_batch_params_reduce_axes(begin_params_axis, x_shape):
         if begin_params_axis < 0:
             x_rank = len(x_shape)
@@ -560,7 +567,7 @@ def get_layernormgrad_vmap_rule(prim, axis_size):
         batch_params_reduce_axes = tuple(range(1, begin_params_axis))
         return batch_params_reduce_axes
 
-    @constexpr
+    @_primexpr
     def get_logical_shape(var_shape):
         return var_shape[1:]
 
@@ -682,3 +689,28 @@ def get_upsample_grad_vmap_rule(prim, axis_size):
         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):
+    """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)
+        if is_all_none:
+            return result
+        x, x_dim = x_bdim
+        grad, _ = grad_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)
+        return dx, x_dim
+
+    return vmap_rule

mindspore/ops/_vmap/vmap_image_ops.py

@@ -16,9 +16,12 @@
 """image_ops vmap impl."""
 from __future__ import absolute_import
 
+import numpy as np
+from mindspore import Tensor
 from mindspore.ops import functional as F
 from mindspore.ops.operations import _grad_ops as G
 from mindspore.ops.operations import image_ops as IMG
+from mindspore.ops import constexpr
 from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_preprocess, _bdim_at_front, \
     _raise_value_error
 
@@ -83,3 +86,52 @@ def get_resize_grad_dynamic_rule(prim, axis_size):
         return out, 0
 
     return vmap_rule
+
+
+@vmap_rules_getters.register(IMG.CropAndResize)
+def get_crop_and_resize_vmap_rule(prim, axis_size):
+    """VmapRule for `CropAndResize` operation."""
+
+    @constexpr
+    def get_box_indices_offsets(axis_size, batch_size, num_boxes):
+        offsets = np.arange(0, axis_size * batch_size, batch_size).astype(np.int32)
+        offsets = np.reshape(offsets, (axis_size, 1))
+        offsets = np.broadcast_to(offsets, (axis_size, num_boxes))
+        return Tensor(offsets)
+
+    def vmap_rule(x_bdim, boxes_bdim, box_indices_bdim, crop_size_bdim):
+        is_all_none, result = vmap_general_preprocess(x_bdim, boxes_bdim, box_indices_bdim, crop_size_bdim)
+        if is_all_none:
+            return result
+
+        x, x_dim = x_bdim
+        boxes, boxes_dim = boxes_bdim
+        box_indices, box_indices_dim = box_indices_bdim
+        crop_size, crop_size_dim = crop_size_bdim
+        if crop_size_dim is not None:
+            _raise_value_error(
+                "The axis of `crop_size` in `{}` must be None, but got {}.".format(prim.name, crop_size_dim))
+
+        boxes = _bdim_at_front(boxes, boxes_dim, axis_size)
+        box_indices = _bdim_at_front(box_indices, box_indices_dim, axis_size)
+        boxes = F.reshape(boxes, (-1, 4))
+        num_boxes = F.shape(box_indices)[-1]
+
+        if x_dim is None:
+            box_indices = F.reshape(box_indices, (-1,))
+            out = prim(x, boxes, box_indices, crop_size)
+        else:
+            x = _bdim_at_front(x, x_dim, axis_size)
+            x_shape = F.shape(x)
+            x = F.reshape(x, (-1,) + x_shape[2:])
+            offsets = get_box_indices_offsets(axis_size, x_shape[1], num_boxes)
+            box_indices = F.add(box_indices, offsets)
+            box_indices = F.reshape(box_indices, (-1,))
+            out = prim(x, boxes, box_indices, crop_size)
+
+        out_shape = F.shape(out)
+        out = F.reshape(out, (-1, num_boxes) + out_shape[1:])
+        return out, 0
+
+
+    return vmap_rule

mindspore/ops/_vmap/vmap_math_ops.py

@@ -19,13 +19,13 @@ from __future__ import absolute_import
 import mindspore.numpy as mnp
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
-from mindspore.ops import constexpr
+from mindspore.ops.primitive import _primexpr
 from mindspore.common import Tensor
 from mindspore.ops.operations import math_ops
 from mindspore.ops.operations import linalg_ops
 from mindspore.ops.operations import _inner_ops
 from mindspore.ops.primitive import Primitive
-from mindspore.ops.composite import _VmapGeneralRule
+from mindspore.ops.function import _VmapGeneralRule
 from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_preprocess, get_assign_vmap_rule, \
     get_unop_vmap_rule, _raise_value_error, _bdim_at_front, _broadcast_by_axis, _handle_broadcasting, \
     _vmap_clone_prim, _bdim_at_any, _get_reduce_batch_axis, _get_reduce_out_dim
@@ -33,7 +33,7 @@ from mindspore.ops.operations.math_ops import Bernoulli, BesselI0, BesselI1, Bes
     BesselK0, BesselK0e, BesselY0, BesselY1, BesselK1, BesselK1e, Median
 
 
-@constexpr
+@_primexpr
 def _broadcast_shape(nd, x_ndim, x_shape):
     return x_shape + (1,) * (nd - x_ndim)
 
@@ -97,6 +97,35 @@ def get_broadcast_binary_op_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(P.Addcdiv)
+@vmap_rules_getters.register(P.Addcmul)
+def get_addcxxx_vmap_rule(prim, axis_size):
+    """VmapRule for addcxxx, such as `Addcdiv` and `Addcmul`."""
+
+    def vmap_rule(input_data_bdim, x1_bdim, x2_bdim, value_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, input_data_bdim, x1_bdim, x2_bdim, value_bdim)
+        if is_all_none:
+            return result
+
+        input_data, input_data_dim = input_data_bdim
+        x1, x1_dim = x1_bdim
+        x2, x2_dim = x2_bdim
+        value, value_dim = value_bdim
+        if input_data_dim == x1_dim and x1_dim == x2_dim and x2_dim == value_dim:
+            out = prim(input_data, x1, x2, value)
+            return out, input_data_dim
+
+        input_data = _bdim_at_front(input_data, input_data_dim, axis_size)
+        x1 = _bdim_at_front(x1, x1_dim, axis_size)
+        x2 = _bdim_at_front(x2, x2_dim, axis_size)
+        value = _bdim_at_front(value, value_dim, axis_size)
+
+        out = prim(input_data, x1, x2, value)
+        return out, 0
+
+    return vmap_rule
+
+
 @vmap_rules_getters.register(P.Cdist)
 def get_cdist_vmap_rule(prim, axis_size):
     """VmapRule for `cdist` operation."""
@@ -358,6 +387,8 @@ def get_inplace_ops_vmap_rule(prim, axis_size):
 @vmap_rules_getters.register(P.ReduceMin)
 @vmap_rules_getters.register(P.ReduceMean)
 @vmap_rules_getters.register(P.ReduceProd)
+@vmap_rules_getters.register(P.ReduceAll)
+@vmap_rules_getters.register(P.ReduceAny)
 def get_reducer_vmap_rule(prim, axis_size):
     """VmapRule for reduce operations, such as `ReduceSum`."""
     reduce_op_map = {
@@ -365,7 +396,9 @@ def get_reducer_vmap_rule(prim, axis_size):
         "ReduceMax": P.ReduceMax,
         "ReduceMin": P.ReduceMin,
         "ReduceMean": P.ReduceMean,
-        "ReduceProd": P.ReduceProd
+        "ReduceProd": P.ReduceProd,
+        "ReduceAll": P.ReduceAll,
+        "ReduceAny": P.ReduceAny,
     }
 
     if isinstance(prim, str):
@@ -403,7 +436,7 @@ def get_median_vmap_rule(prim, axis_size):
     axis = prim.axis
     keep_dims = prim.keep_dims
 
-    @constexpr
+    @_primexpr
     def trans_axis(axis, rank, dim, keep_dims):
         if axis < 0:
             axis += rank - 1
@@ -431,7 +464,7 @@ def get_index_add_vmap_rule(prim, axis_size):
     """VmapRule for IndexAdd."""
     axis = prim.axis
 
-    @constexpr
+    @_primexpr
     def _get_index_add_batch_axis(axis, x_dim, x_ndim):
         """get batch_axis for IndexAdd."""
         # case1: batch not exists
@@ -770,6 +803,44 @@ def get_square_sum_all_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(math_ops.FFTWithSize)
+def get_fft_with_size_vmap_rule(prim, axis_size):
+    """VmapRule for `FFTWithSize` operation"""
+    if isinstance(prim, str):
+        prim_name = prim
+        prim = Primitive(prim)
+        signal_ndim = 1
+        inverse = False
+        real = False
+        norm = "backward"
+        oneside = True
+        signal_sizes = ()
+    else:
+        prim_name = prim.name
+        signal_ndim = prim.signal_ndim
+        inverse = prim.inverse
+        real = prim.real
+        norm = prim.norm
+        oneside = prim.oneside
+        signal_sizes = prim.signal_sizes
+
+    fft = math_ops.FFTWithSize(signal_ndim, inverse, real, norm, oneside, signal_sizes)
+
+    def vmap_rule(x_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim)
+        if is_all_none:
+            return result
+        x, x_dim = x_bdim
+        x_ndim = F.rank(x)
+        if x_dim < 0 or x_dim >= x_ndim - signal_ndim:
+            _raise_value_error("The source axi of `x` in `{} must be`in range of ({} {}), "
+                               "but got {}.".format(prim_name, 0, x_ndim - signal_ndim, x_dim))
+        out = fft(x)
+        return (out, x_dim)
+
+    return vmap_rule
+
+
 get_assign_vmap_rule = vmap_rules_getters.register(P.AssignAdd)(get_assign_vmap_rule)
 get_assign_vmap_rule = vmap_rules_getters.register(P.AssignSub)(get_assign_vmap_rule)
 

mindspore/ops/_vmap/vmap_nn_ops.py

@@ -23,6 +23,7 @@ from mindspore.ops.operations import _grad_ops as G
 from mindspore.ops.operations import nn_ops as NN
 from mindspore.ops import functional as F
 from mindspore.ops import constexpr
+from mindspore.ops.primitive import _primexpr
 from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_preprocess, get_unop_vmap_rule, \
     _bdim_at_any, _bdim_at_front, _bdim_at_back, _handle_broadcasting, get_unary_grad_vmap_rule, _raise_value_error, \
     _vmap_clone_prim, _get_reduce_batch_axis
@@ -375,7 +376,7 @@ def get_bias_add_vmap_rule(prim, axis_size):
     def get_channal_pos_in_x(d_format):
         return d_format.find('C') + 1
 
-    @constexpr
+    @_primexpr
     def get_bias_dst_shape(x_shape, n_dims, d_format, has_b_dim: bool):
         pos = get_channal_pos_in_x(d_format)
 
@@ -430,7 +431,7 @@ def get_bias_add_grad_vmap_rule(prim, axis_size):
     def get_channal_pos(d_format):
         return d_format.find('C') + 1
 
-    @constexpr
+    @_primexpr
     def get_axis_for_reduce(x_shape_rank, data_format):
         channal_pos = get_channal_pos(data_format)
         axis_list = ()
@@ -1072,24 +1073,24 @@ def get_pad_v3_vmap_rule(prim, axis_size):
         if is_all_none:
             return result
         if len(params_bdim) < 2:
-            _raise_value_error("The input params in `{}` must >= 2, "
-                               "but got {}.".format(prim.name, len(params_bdim)))
+            _raise_value_error("The input params in `PadV3` must >= 2, "
+                               "but got {}.".format(len(params_bdim)))
         input_x, input_x_dim = params_bdim[0]
         paddings, paddings_dim = params_bdim[1]
         values = None
         out = None
         x = _bdim_at_front(input_x, input_x_dim, axis_size)
         if paddings_dim is not None:
-            _raise_value_error("The source axis of `paddings` in `{}` must be None, "
-                               "but got {}.".format(prim.name, paddings_dim))
+            _raise_value_error("The source axis of `paddings` in `PadV3` must be None, "
+                               "but got {}.".format(paddings_dim))
         if mode == "constant":
             if len(params_bdim) != 3:
-                _raise_value_error("The input params in `{}` of constant mode must be 3, "
-                                   "but got {}.".format(prim.name, len(params_bdim)))
+                _raise_value_error("The input params in `PadV3` of constant mode must be 3, "
+                                   "but got {}.".format(len(params_bdim)))
             values, values_dim = params_bdim[2]
             if values_dim is not None:
-                _raise_value_error("The source axis of `values_dim` in `{}` must be None, "
-                                   "but got {}.".format(prim.name, values_dim))
+                _raise_value_error("The source axis of `values_dim` in `PadV3` must be None, "
+                                   "but got {}.".format(values_dim))
         if isinstance(paddings, Tensor):
             pad_dim = F.shape(paddings)[0] / pad_pair
         else:
@@ -1101,7 +1102,7 @@ def get_pad_v3_vmap_rule(prim, axis_size):
                 out = prim(x, paddings, values)
             else:
                 out = prim(x, paddings)
-        elif x_ndim > input_max_dim:
+        elif x_ndim >= input_max_dim:
             # reshape to 4 dims
             x_shape = F.shape(x)
             diff_dim = x_ndim - input_max_dim
@@ -1118,8 +1119,8 @@ def get_pad_v3_vmap_rule(prim, axis_size):
             real_out_shape = x_shape[:diff_dim + 1] + out_shape[1:]
             out = F.reshape(out, real_out_shape)
         else:
-            _raise_value_error("The dim of `input_x` in `{}` must be bigger than {}, "
-                               "but got {}.".format(prim.name, pad_dim, x_ndim))
+            _raise_value_error("The dim of `input_x` in `PadV3` must be bigger than {}, "
+                               "but got {}.".format(pad_dim, x_ndim))
         return out, 0
 
     return vmap_rule
@@ -1308,6 +1309,60 @@ def get_apply_adam_with_amsgrad_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(P.Adam)
+def get_adam_rule(prim, axis_size):
+    """VmapRule for `Adam` operation"""
+    if hasattr(prim, "batch_rank"):
+        batch_rank = prim.batch_rank + 1
+    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, m_bdim, v_bdim, beta1_power_bdim, beta2_power_bdim, lr_bdim, beta1_bdim,
+                  beta2_bdim, epsilon_bdim, grad_bdim, u_monad):
+        var, var_dim = var_bdim
+        m, m_dim = m_bdim
+        v, v_dim = v_bdim
+        beta1_power, beta1_power_dim = beta1_power_bdim
+        beta2_power, beta2_power_dim = beta2_power_bdim
+        lr, lr_dim = lr_bdim
+        beta1, beta1_dim = beta1_bdim
+        beta2, beta2_dim = beta2_bdim
+        epsilon, epsilon_dim = epsilon_bdim
+        grad, grad_dim = grad_bdim
+
+        all_dim = [m_dim, v_dim, beta1_power_dim, beta2_power_dim, lr_dim, beta1_dim, beta2_dim, epsilon_dim, grad_dim]
+        if var_dim is None:
+            if any(dim is not None for dim in all_dim):
+                raise ValueError("The source axis of `var` is None, "
+                                 "but the source axis of `m/v/vhat/beta1_power/beta2_power/lr/beta1/beta2/epsilon grad"
+                                 " is not None. The execution of operator `{}` cannot be guaranteed.".format(prim_name))
+            out_var, out_m, out_v = prim(
+                var, m, v, beta1_power, beta2_power, lr, beta1, beta2, epsilon, grad, u_monad)
+            return ((out_var, None), (out_m, None), (out_v, None))
+
+        if any(dim != 0 for dim in [var_dim, m_dim, v_dim]):
+            raise ValueError("For `{}`, the source axis of `var/m/v` must be 0, "
+                             "but get `var`: {}, `m`: {}, `v`: {}".format(prim_name, var_dim,
+                                                                          m_dim, v_dim))
+
+        beta1_power = _bdim_at_front(beta1_power, beta1_power_dim, axis_size)
+        beta2_power = _bdim_at_front(beta2_power, beta2_power_dim, axis_size)
+        lr = _bdim_at_front(lr, lr_dim, axis_size)
+        beta1 = _bdim_at_front(beta1, beta1_dim, axis_size)
+        beta2 = _bdim_at_front(beta2, beta2_dim, axis_size)
+        epsilon = _bdim_at_front(epsilon, epsilon_dim, axis_size)
+        grad = _bdim_at_front(grad, grad_dim, axis_size)
+
+        out_var, out_m, out_v = batch_prim(
+            var, m, v, beta1_power, beta2_power, lr, beta1, beta2, epsilon, grad, u_monad)
+        return ((out_var, 0), (out_m, 0), (out_v, 0))
+
+    return vmap_rule
+
+
 @vmap_rules_getters.register(P.ApplyPowerSign)
 def get_apply_power_sign_rule(prim, axis_size):
     """VmapRule for `ApplyPowerSign` operation."""
@@ -1461,10 +1516,9 @@ def get_adaptive_max_pool_2d_vmap_rule(prim, axis_size):
     nchw_index = 4
     chw_reverse_index = -3
     hw_size = 2
-    return_indices = prim.return_indices
     output_size = prim.output_size
 
-    @constexpr
+    @_primexpr
     def get_output_shape(x_ori_shape, output_size):
         if isinstance(output_size, tuple):
             h_out, w_out = output_size
@@ -1499,20 +1553,14 @@ def get_adaptive_max_pool_2d_vmap_rule(prim, axis_size):
             x_ori_shape = F.shape(x)
             x = F.reshape(x, (-1,) + x_ori_shape[chw_reverse_index:])
             output_shape = get_output_shape(x_ori_shape, output_size)
-            if return_indices:
-                out, indices = prim(x)
-                out = F.reshape(out, output_shape)
-                indices = F.reshape(indices, output_shape)
-                return (out, 0), (indices, 0)
-            out = prim(x)
-            out = F.reshape(out, output_shape)
-            return out, 0
-        # for the case of CHW
-        if return_indices:
             out, indices = prim(x)
+            out = F.reshape(out, output_shape)
+            indices = F.reshape(indices, output_shape)
             return (out, 0), (indices, 0)
-        out = prim(x)
-        return out, 0
+
+        # for the case of CHW
+        out, indices = prim(x)
+        return (out, 0), (indices, 0)
 
     return vmap_rule
 
@@ -1657,6 +1705,7 @@ def get_apply_centered_rmsprop_vmap_rule(prim, axis_size):
 
 @vmap_rules_getters.register(P.MaxPool)
 @vmap_rules_getters.register(P.MaxPoolWithArgmax)
+@vmap_rules_getters.register(P.MaxPoolWithArgmaxV2)
 def get_max_pool_vmap_rule(prim, axis_size):
     """VmapRule for `MaxPool` operation."""
     if isinstance(prim, str):
@@ -1664,7 +1713,7 @@ def get_max_pool_vmap_rule(prim, axis_size):
 
     prim_name = prim.name
 
-    @constexpr
+    @_primexpr
     def get_original_shape(x_shape, out_shape):
         h_new = out_shape[2]
         w_new = out_shape[3]
@@ -1709,7 +1758,7 @@ def get_layernorm_vmap_rule(prim, axis_size):
     params_axis = process_attr_axis(prim.begin_params_axis)
     batch_prim = P.LayerNorm(norm_axis, params_axis, prim.epsilon)
 
-    @constexpr
+    @_primexpr
     def get_logical_shape(var_shape):
         return var_shape[1:]
 

mindspore/ops/_vmap/vmap_other_ops.py

@@ -83,7 +83,9 @@ def get_partical_vmap_rule(prim, axis_size):
             else:
                 val, dim = val_bdim
                 if dim is not None:
-                    _raise_value_error("The source axis of args in {} must be None, "
+                    _raise_value_error("In the scenario where vmap contains control flow, currently only the "
+                                       "case of each batch branch with the same processing operations is "
+                                       "supported, so that the source axis of args in {} must be None, "
                                        "but got {}.".format(prim_name, dim))
                 vals = vals + (val,)
 

mindspore/ops/_vmap/vmap_random_ops.py

@@ -1,4 +1,3 @@
-
 # Copyright 2022 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -17,8 +16,11 @@
 """random_ops vmap impl."""
 from __future__ import absolute_import
 
-from mindspore.ops.operations.random_ops import UniformCandidateSampler, RandomShuffle
-from mindspore.ops._vmap.vmap_base import vmap_rules_getters, _bdim_at_front, _vmap_clone_prim, vmap_general_preprocess
+from mindspore.ops.operations.random_ops import UniformCandidateSampler, RandomShuffle, Multinomial, \
+    RandomChoiceWithMask
+from mindspore.ops.function import _VmapGeneralRule
+from mindspore.ops._vmap.vmap_base import vmap_rules_getters, _bdim_at_front, _vmap_clone_prim, \
+    vmap_general_preprocess, _raise_value_error
 
 
 @vmap_rules_getters.register(UniformCandidateSampler)
@@ -68,3 +70,53 @@ def get_random_shuffle_vmap_rule(prim, axis_size):
         return out, 0
 
     return vmap_rule
+
+
+@vmap_rules_getters.register(Multinomial)
+def get_multinomial_vmap_rule(prim, axis_size):
+    """VmapRule for `Multinomial` operation."""
+    prim_name = prim.name
+    prim_vmap = _VmapGeneralRule(prim, axis_size)
+
+    def vmap_rule(x_bdim, num_samples_bdim):
+        is_all_none, result = vmap_general_preprocess(
+            prim, x_bdim, num_samples_bdim)
+        if is_all_none:
+            return result
+
+        x, x_dim = x_bdim
+        num_samples, num_samples_dim = num_samples_bdim
+        if len(x.shape) > 2:
+            out = prim_vmap(x_bdim, num_samples_bdim)
+            return out
+        if num_samples_dim is not None:
+            _raise_value_error("The source axis of args in {} must be None, "
+                               "but got {}.".format(prim_name, num_samples_dim))
+        x = _bdim_at_front(x, x_dim, axis_size)
+        out = prim(x, num_samples)
+        return (out, 0)
+
+    return vmap_rule
+
+
+@vmap_rules_getters.register(RandomChoiceWithMask)
+def get_random_choice_with_mask(prim, axis_size):
+    """VmapRule for 'RandomChoiceWithMask' operation."""
+    if hasattr(prim, 'batch_rank'):
+        batch_rank = prim.batch_rank + 1
+    else:
+        batch_rank = 1
+
+    batch_prim = _vmap_clone_prim(prim)
+    batch_prim.add_prim_attr('batch_rank', batch_rank)
+
+    def vmap_rule(x_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim)
+        if is_all_none:
+            return result
+        x_data, x_dim = x_bdim
+        x = _bdim_at_front(x_data, x_dim, axis_size)
+        index, mask = batch_prim(x)
+        return (index, 0), (mask, 0)
+
+    return vmap_rule

mindspore/ops/_vmap/vmap_sparse_ops.py

@@ -14,6 +14,7 @@
 # ============================================================================
 
 """sparse_ops vmap impl."""
+from __future__ import absolute_import
 
 from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_preprocess, _raise_value_error
 from mindspore.ops.primitive import Primitive

mindspore/ops/bprop_mindir/ApproximateEqual_bprop.mindir

@@ -1,20 +1,19 @@
 
-0.1.1	MindSpore*1.9.0:�
-�

-
-bprop.33:xbprop.33:[CNode]34:1bprop.33:[CNode]34:1"(REF::S-Prim-hyper_map[zeros_like_leaf]:2:.Default/S-Prim-hyper_map[zeros_like_leaf]-op23
-�

-
-bprop.33:ybprop.33:[CNode]35:3bprop.33:[CNode]35:3"(REF::S-Prim-hyper_map[zeros_like_leaf]:2:.Default/S-Prim-hyper_map[zeros_like_leaf]-op24
-�

-bprop.33:[CNode]34:1
-bprop.33:[CNode]35:3bprop.33:[CNode]36:4bprop.33:[CNode]36:4"REF::S-Prim-MakeTuple:5:Default/S-Prim-MakeTuple-op25bprop.33*
-
-bprop.33:x*
-
-bprop.33:y*
-bprop.33:out*
-
-bprop.33:[CNode]36:4:@157abcda70ad669686a320be427dedf2a17498d59a38b65d04c272c1a40296e0P
b&
-S-Prim-MakeTuple:5S-Prim-MakeTuplebH
-#S-Prim-hyper_map[zeros_like_leaf]:2!S-Prim-hyper_map[zeros_like_leaf]h

+0.1.1	MindSpore*2.0.0:�
+�
.get_bprop_approximate_equal.1184:[CNode]1185:1.get_bprop_approximate_equal.1184:[CNode]1185:1"REF::bprop.1186:Default/bprop.1186-op927 get_bprop_approximate_equal.1184*'
+%get_bprop_approximate_equal.1184:self*$
+"get_bprop_approximate_equal.1184:x*$
+"get_bprop_approximate_equal.1184:y*&
+$get_bprop_approximate_equal.1184:out*'
+%get_bprop_approximate_equal.1184:dout20
+.get_bprop_approximate_equal.1184:[CNode]1185:1:@7fb54a66e55c2c40cd92783044880b792666a0d7fc794bb717bee3544337d6a0J/grad_math_ops.pyB�
+�

+"get_bprop_approximate_equal.1184:xbprop.1186:[CNode]1187:2bprop.1186:[CNode]1187:2".REF::MetaFuncGraph::hyper_map[zeros_like_leaf]:/Default/S-Prim-hyper_map[zeros_like_leaf]-op928
+�

+"get_bprop_approximate_equal.1184:ybprop.1186:[CNode]1188:3bprop.1186:[CNode]1188:3".REF::MetaFuncGraph::hyper_map[zeros_like_leaf]:/Default/S-Prim-hyper_map[zeros_like_leaf]-op929
+�

+bprop.1186:[CNode]1187:2
+bprop.1186:[CNode]1188:3bprop.1186:[CNode]1189:4bprop.1186:[CNode]1189:4"REF::S-Prim-MakeTuple:5:Default/S-Prim-MakeTuple-op930
+bprop.11862
+bprop.1186:[CNode]1189:4P
b&
+S-Prim-MakeTuple:5S-Prim-MakeTupleh