mindspore 1.10.0__cp37-cp37m-win_amd64.whl → 2.0.0rc1__cp37-cp37m-win_amd64.whl

mindspore/ops/_vmap/vmap_array_ops.py

@@ -16,25 +16,26 @@
 """array_ops vmap impl."""
 from __future__ import absolute_import
 
-import numpy as np
 import mindspore
 import mindspore.numpy as mnp
 from mindspore import ops
 from mindspore.common import Tensor
+from mindspore._c_expression import Tensor as Tensor_
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
-from mindspore.ops import constexpr
+from mindspore.ops.primitive import constexpr, _primexpr
 from mindspore.ops.operations._grad_ops import MaskedSelectGrad
 from mindspore.ops.operations import _grad_ops as G
 from mindspore.ops.operations.array_ops import Fills, UniqueConsecutive, Col2Im, NonZero, IndexFill, \
     TensorScatterElements
 from mindspore.ops.operations.random_ops import RandomPoisson
+from mindspore.ops.operations._inner_ops import DynamicBroadcastTo
 from mindspore.ops.primitive import Primitive
 from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_preprocess, _bdim_at_front, \
     _raise_value_error, _vmap_clone_prim, _handle_broadcasting, get_unsupported_dynamic_vmap_rule, _broadcast_by_axis, \
-    get_unop_vmap_rule, _get_reduce_out_dim, _get_reduce_batch_axis, _vmap_update_prim_attr, \
+    get_unop_vmap_rule, _get_reduce_out_dim, _get_reduce_batch_axis, \
     _bdim_at_any
-from mindspore.ops.composite import _VmapGeneralRule
+from mindspore.ops.function import _VmapGeneralRule
 
 
 @vmap_rules_getters.register(P.NoRepeatNGram)
@@ -137,7 +138,7 @@ def get_arg_min_max_with_value_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
-@constexpr
+@_primexpr
 def _get_prefix(indices_shape, axis_size, indices_dtype):
     """
     Generate prefix by indices shape, whose -1 axis value is the index value of axis 0.
@@ -147,14 +148,16 @@ def _get_prefix(indices_shape, axis_size, indices_dtype):
     the generated prefix is a Tensor([[[0], [0]],
                                       [[1], [1]]])
     """
-    if not indices_shape:
-        raise ValueError("indices_shape is empty in _get_prefix.")
+    def _check(indices_shape):
+        if not indices_shape:
+            raise ValueError("indices_shape is empty in _get_prefix.")
 
+    _check(indices_shape)
     indices_len = len(indices_shape)
-
     if indices_len == 1:
-        prefix = np.arange(axis_size)
-        return Tensor(prefix, indices_dtype)
+        prefix = P.Range()(Tensor(0, indices_dtype), P.Fill()(
+            indices_dtype, (), axis_size), Tensor(1, indices_dtype))
+        return prefix
 
     indices_end = indices_len - 1
     prefix_shape = ()
@@ -169,8 +172,9 @@ def _get_prefix(indices_shape, axis_size, indices_dtype):
         else:
             expand_shape = expand_shape + (1,)
 
-    prefix = np.broadcast_to(np.arange(axis_size).reshape(expand_shape), prefix_shape)
-    return Tensor(prefix, indices_dtype)
+    prefix = P.BroadcastTo(prefix_shape)(P.Reshape()(P.Range()(Tensor(
+        0, indices_dtype), Tensor(axis_size, indices_dtype), Tensor(1, indices_dtype)), expand_shape))
+    return prefix
 
 
 @vmap_rules_getters.register(P.Transpose)
@@ -179,7 +183,7 @@ def get_transpose_vmap_rule(prim, axis_size):
     if isinstance(prim, str):
         prim = Primitive(prim)
 
-    @constexpr
+    @_primexpr
     def _get_transpose_batch_perm(dim, perm, x_rank):
         """Generate batch_perm based on the original perm of transpose operation and dim of the input."""
         if dim < 0:
@@ -223,24 +227,20 @@ def get_tile_vmap_rule(prim, axis_size):
     if isinstance(prim, str):
         prim = Primitive(prim)
 
-    @constexpr
-    def _get_tile_shape(input_shape, multiples):
+    @_primexpr
+    def _get_batch_multiples(input_shape, dim, multiples):
         input_ndim = len(input_shape)
         multiples_ndim = len(multiples)
-        input_shape = (input_shape[0],) + (1,) * (multiples_ndim - input_ndim + 1) + input_shape[1:]
-        multiples = (1,) * (input_ndim - multiples_ndim) + multiples
-        input_expand_shape = (input_shape[0],) + tuple([
-            j
-            for i in input_shape[1:]
-            for j in [1, i]
-        ])
-        repeat_shape = (input_shape[0],) + tuple([
-            k
-            for pair in zip(multiples[1:], input_shape[1:])
-            for k in pair
-        ])
-        output_shape = tuple([a * b for a, b in zip(input_shape, multiples)])
-        return input_expand_shape, repeat_shape, output_shape
+        if multiples_ndim < input_ndim - 1:
+            multiples = (1,) * (input_ndim - 1 - multiples_ndim) + multiples
+
+        rev_dim = input_ndim - 1 - dim
+        if rev_dim == 0:
+            return multiples + (1,), multiples_ndim
+
+        batch_multiples = list(multiples)
+        batch_multiples.insert(-rev_dim, 1)
+        return tuple(batch_multiples), multiples_ndim - rev_dim
 
     def vmap_rule(input_bdim, multiples_bdim):
         is_all_none, result = vmap_general_preprocess(prim, input_bdim, multiples_bdim)
@@ -252,13 +252,10 @@ def get_tile_vmap_rule(prim, axis_size):
         if multiples_dim is not None:
             _raise_value_error("The source axis of shape in `Tile` must be None, but got {}.".format(multiples_dim))
 
-        input_x = _bdim_at_front(input_x, dim, axis_size)
         input_shape = F.shape(input_x)
-        input_expand_shape, repeat_shape, output_shape = _get_tile_shape(input_shape, multiples)
-        expand_input = F.reshape(input_x, input_expand_shape)
-        repeat_tensor = P.BroadcastTo(repeat_shape)(expand_input)
-        output = F.reshape(repeat_tensor, output_shape)
-        return output, 0
+        batch_multiples, out_dim = _get_batch_multiples(input_shape, dim, multiples)
+        repeat_tensor = P.Tile()(input_x, batch_multiples)
+        return repeat_tensor, out_dim
 
     return vmap_rule
 
@@ -359,8 +356,13 @@ def get_unstack_vmap_rule(prim, axis_size):
 def get_reshape_vmap_rule(prim, axis_size):
     """VmapRule for `Reshape` operation."""
 
-    @constexpr
+
+    @_primexpr
     def get_batch_shape(x_shape, x_dim, target_shape, axis_size):
+        def _check(neg_index, target_shape):
+            if neg_index != -1:
+                raise ValueError(f'The shape can only has one -1 at most, but {target_shape}.')
+
         if x_dim == 0:
             return (axis_size,) + target_shape, 0, False
 
@@ -371,19 +373,21 @@ def get_reshape_vmap_rule(prim, axis_size):
         dim_prod = 1
         for i, shp_i in enumerate(target_shape):
             if shp_i == -1:
-                if neg_index != -1:
-                    raise ValueError(f'The shape can only has one -1 at most, but {target_shape}.')
+                _check(neg_index, target_shape)
                 neg_index = i
             else:
                 dim_prod *= shp_i
-        arr_prod = np.prod(x_shape)
+        arr_prod = 1
+        for i in x_shape:
+            arr_prod *= i
         target_shape_list = list(target_shape)
         if neg_index != -1:
             neg_index_size = int(arr_prod // (dim_prod * axis_size))
             target_shape_list[neg_index] = neg_index_size
 
-        arr_prod_before_dim = np.prod(x_shape[:x_dim])
-
+        arr_prod_before_dim = 1
+        for i in x_shape[:x_dim]:
+            arr_prod_before_dim *= i
         dim_prod = 1
         for i, shp_i in enumerate(target_shape_list, start=1):
             dim_prod *= shp_i
@@ -428,7 +432,7 @@ def get_reverse_sequence_vmap_rule(prim, axis_size):
     batch_dim = prim.batch_dim_
     seq_dim = prim.seq_dim_
 
-    @constexpr
+    @_primexpr
     def get_batch_seq_dim(dim, batch_dim_, seq_dim_):
         if dim is None:
             batch_dim_ += 1
@@ -444,7 +448,7 @@ def get_reverse_sequence_vmap_rule(prim, axis_size):
                     seq_dim_ += 1
         return batch_dim_, seq_dim_
 
-    @constexpr
+    @_primexpr
     def get_seq_dim(dim, batch_dim_, seq_dim_):
         if dim is None:
             return seq_dim_
@@ -564,20 +568,19 @@ def get_scatter_nd_vmap_rule(prim, axis_size):
     Reshape the output tensor to `[10, 6, 4, 5]`
     """
 
-    @constexpr
+    @_primexpr
     def _refine_shape(shape, bdim_size):
         offset = shape[0]
         return (bdim_size * shape[0],) + tuple(shape[1:]), offset, (bdim_size,) + tuple(shape)
 
-    @constexpr
+    @_primexpr
     def _gen_indices_offset(shape, offset):
         # original rank(indices.shape) is required >= 2, so indices with batch dim's rank >= 3.
-        shape = [shape[0]] + [1] * (len(shape) - 2) + [shape[-1]]
-        val = np.zeros(shape, np.int32)  # the dtype will be changed when creating Tensor
-        val = np.reshape(val, (shape[0], shape[-1]))
+        shape = (shape[0],) + (1,) * (len(shape) - 2) + (shape[-1],)
+        val = P.Zeros()((shape[0], shape[-1]), mindspore.int32)
         for i in range(shape[0]):
             val[i, 0] = i * offset
-        return np.reshape(val, shape)
+        return P.Reshape()(val, shape)
 
     if isinstance(prim, str):
         prim = Primitive(prim)
@@ -598,7 +601,7 @@ def get_scatter_nd_vmap_rule(prim, axis_size):
         indices_shape = F.shape(indices)
         indices_dtype = F.dtype(indices)
         offset_val = _gen_indices_offset(indices_shape, offset)
-        indices_offset = Tensor(offset_val, indices_dtype)
+        indices_offset = P.Cast()(offset_val, indices_dtype)
         new_indices = P.Add()(indices, indices_offset)
         out = prim(new_indices, updates, new_shape)
         real_out = P.Reshape()(out, out_shape)
@@ -846,6 +849,62 @@ def get_fill_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@constexpr
+def to_tensor_with_type(x, type):
+    """x to Tensor with type"""
+    return Tensor(x, type)
+
+
+@vmap_rules_getters.register(P.FillV2)
+def get_fill_v2_vmap_rule(prim, axis_size):
+    """VmapRule for `FillV2` operation."""
+    if isinstance(prim, str):
+        prim = Primitive(prim)
+
+    def vmap_rule(shape_bdim, value_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, shape_bdim, value_bdim)
+        if is_all_none:
+            return result
+
+        value_shape, shape_dim = shape_bdim
+        if shape_dim is not None:
+            _raise_value_error(
+                "The source axis of `shape` in `P.FillV2` must be None, but got {}."
+                .format(shape_dim))
+
+        value, vdim = value_bdim
+        value_rank = F.rank(value)
+        if value_rank != 1 or vdim != 0:
+            _raise_value_error(
+                "The `value` in `P.FillV2` must be constant value, thus the value only "
+                "can be rank: 1 with source axis: 0 in vmap scope, but got value rank: "
+                "{} with source axis: {}.".format(value_rank, vdim))
+        value = F.reshape(value, (axis_size,) + (1,) * len(value_shape))
+
+        out = None
+        if isinstance(value_shape, (Tensor_, Tensor)):
+            value_shape_rank = F.rank(value_shape)
+            if value_shape_rank != 1:
+                _raise_value_error(
+                    "The `shape` in `P.FillV2` must be 1-D tensor, thus the shape only "
+                    "can be rank: 1, but got shape rank: "
+                    "{}.".format(value_shape_rank))
+            axis_size_tensor = to_tensor_with_type((axis_size,),
+                                                   F.dtype(value_shape))
+            broad_cast_shape = F.concat((axis_size_tensor, value_shape))
+            out = DynamicBroadcastTo()(value, broad_cast_shape)
+        elif isinstance(value_shape, tuple):
+            out = P.BroadcastTo((axis_size,) + value_shape)(value)
+        else:
+            _raise_value_error(
+                f"For `P.FillV2`, the input `shape` should be Tuple or Tensor, but got `shape`: {value_shape}."
+            )
+
+        return out, 0
+
+    return vmap_rule
+
+
 @vmap_rules_getters.register(Fills)
 def get_fills_vmap_rule(prim, axis_size):
     """VmapRule for `Fills` operation."""
@@ -1299,12 +1358,7 @@ def get_gatherd_grad_v2_vmap_rule(prim, axis_size):
     if isinstance(prim, str):
         prim = Primitive(prim)
 
-    dim = 0
-    if hasattr(prim, 'dim'):
-        dim = prim.dim
-
-    @constexpr
-    def _update_attr(x_rank, batch_dim):
+    def _update_dim(dim, x_rank, batch_dim):
         pdim = dim
         if pdim < 0:
             pdim += x_rank
@@ -1312,19 +1366,22 @@ def get_gatherd_grad_v2_vmap_rule(prim, axis_size):
             _raise_value_error(
                 "The `dim` in `GatherDGradV2` must be in range [{}, {}], but got {}.".format(-x_rank, x_rank - 1, dim))
         if pdim >= batch_dim:
-            _vmap_update_prim_attr(prim, 'dim', pdim + 1)
-        elif dim < 0:
-            _vmap_update_prim_attr(prim, 'dim', pdim)
+            return pdim + 1
+        if dim < 0:
+            return pdim
+        return dim
 
-    def vmap_rule(x_bdim, index_bdim, grad_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, x_bdim, index_bdim, grad_bdim)
+    def vmap_rule(x_bdim, dim_bdim, index_bdim, grad_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, dim_bdim, index_bdim, grad_bdim)
         if is_all_none:
             return result
 
         x, x_dim = x_bdim
+        dim, dim_dim = dim_bdim
+        if dim_dim is not None:
+            _raise_value_error("The dim of 'dim' in `GatherDGradV2` must be None, but got {}.".format(dim_dim))
         index, index_dim = index_bdim
         grad, grad_dim = grad_bdim
-
         batch_dim = 0
         if x_dim is not None:
             batch_dim = x_dim
@@ -1336,12 +1393,10 @@ def get_gatherd_grad_v2_vmap_rule(prim, axis_size):
         x = _bdim_at_any(x, x_dim, batch_dim, axis_size)
         index = _bdim_at_any(index, index_dim, batch_dim, axis_size)
         grad = _bdim_at_any(grad, grad_dim, batch_dim, axis_size)
-
-        # Adjust dim-attr if needed
         x_rank = F.rank(x) - 1
-        _update_attr(x_rank, batch_dim)
-
-        out = prim(x, index, grad)
+        # Adjust dim if needed
+        dim = _update_dim(dim, x_rank, batch_dim)
+        out = prim(x, dim, index, grad)
         return (out, batch_dim)
 
     return vmap_rule
@@ -1425,6 +1480,7 @@ def get_meshgrid_vmap_rule(prim, axis_size):
                 "The input number of P.Meshgrid must be greater than 1.")
 
         output_shape = []
+        ones_shape = []
         for each_arg in args:
             x, bdim = each_arg
             if bdim is None:
@@ -1435,19 +1491,30 @@ def get_meshgrid_vmap_rule(prim, axis_size):
                 _raise_value_error(
                     "Each input of Meshgrid must be 1D, but got {}.".format(F.rank(x) - 1))
             output_shape.append(F.shape(x)[-1])
+            ones_shape.append(1)
         output_shape.insert(0, axis_size)
+        ones_shape.insert(0, axis_size)
 
         if indexing == "xy":
             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)
+
+        index = 0
         vals_out_tuple = ()
         for each_arg in args:
             x, bdim = each_arg
             x = _bdim_at_front(x, bdim, axis_size)
-            x = _handle_broadcasting(x, F.shape(x), output_shape)
-            output = P.BroadcastTo(shape)(x)
+            shape_index = (1 - index) if (index <= 1 and indexing == "xy") 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)
             vals_out_tuple = vals_out_tuple + ((output, 0),)
+            ones_shape[shape_index + 1] = 1
+            index = index + 1
 
         return vals_out_tuple
 
@@ -1491,7 +1558,7 @@ def get_gather_vmap_rule(prim, axis_size):
     else:
         prim_name = prim.name
 
-    @constexpr
+    @_primexpr
     def process_axis(axis, x_shape_size, has_xdim: bool, has_idim: bool):
         if has_xdim and has_idim:
             if axis < 0:
@@ -1505,7 +1572,7 @@ def get_gather_vmap_rule(prim, axis_size):
 
         return axis
 
-    @constexpr
+    @_primexpr
     def get_x_dst_shape(x_shape, axis):
         target_axis_size = x_shape[axis + 1]
         x_dst_shape = x_shape[0:axis] + (axis_size * target_axis_size,) + x_shape[axis + 2:]
@@ -1705,7 +1772,7 @@ def get_data_format_dim_map_vmap_rule(prim, axis_size):
 def get_expand_dims_vmap_rule(prim, axis_size):
     """VmapRule for `ExpandDims`."""
 
-    @constexpr
+    @_primexpr
     def process_axis(axis, rank, x_dim):
         if axis < 0:
             axis += rank
@@ -1799,7 +1866,7 @@ def get_squeeze_vmap_rule(prim, axis_size):
     else:
         prim_axis = None
 
-    @constexpr
+    @_primexpr
     def move_axis(axes):
         new_axis = ()
         for axis in axes:
@@ -1809,7 +1876,7 @@ def get_squeeze_vmap_rule(prim, axis_size):
                 new_axis = new_axis + (axis + 1,)
         return new_axis
 
-    @constexpr
+    @_primexpr
     def generate_all_axis_except_first(x_rank):
         new_axis = ()
         for i in range(1, x_rank, 1):
@@ -1838,6 +1905,7 @@ def get_squeeze_vmap_rule(prim, axis_size):
         batch_squeeze = P.Squeeze(axis=new_axis)
         out = batch_squeeze(x)
         return out, 0
+
     return vmap_rule
 
 
@@ -1852,7 +1920,7 @@ def get_stridedslice_vmap_rule(prim, axis_size):
     batch_stridedslice = P.StridedSlice(new_begin_mask, new_end_mask, new_ellipsis_mask, new_new_axis_mask, \
                                         new_shrink_axis_mask)
 
-    @constexpr
+    @_primexpr
     def get_new_begin_end_strided(begin, end, strided):
         new_begin = (0,) + begin
         new_end = (0,) + end
@@ -1891,9 +1959,9 @@ def get_stridedslice_grad_vmap_rule(prim, axis_size):
     new_new_axis_mask = prim.new_axis_mask * 2
     new_shrink_axis_mask = prim.shrink_axis_mask * 2
     batch_stridedslice_grad = G.StridedSliceGrad(new_begin_mask, new_end_mask, new_ellipsis_mask, new_new_axis_mask, \
-                                        new_shrink_axis_mask)
+                                                 new_shrink_axis_mask)
 
-    @constexpr
+    @_primexpr
     def get_new_xshape_begin_end_strided(xshape, begin, end, strided):
         new_xshape = (axis_size,) + xshape
         new_begin = (0,) + begin
@@ -1984,6 +2052,30 @@ def get_im2col_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(P.Split)
+def get_split_vmap_rule(prim, axis_size):
+    """VmapRule for `Split`."""
+
+    axis = prim.axis
+    if axis >= 0:
+        axis += 1
+    batch_prim = P.Split(axis, prim.output_num)
+
+    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 = _bdim_at_front(x, x_dim, axis_size)
+        outputs = batch_prim(x)
+        output = ()
+        for out in outputs:
+            output = output + ((out, 0),)
+        return output
+
+    return vmap_rule
+
+
 get_unsupported_dynamic_vmap_rule = vmap_rules_getters.register(NonZero)(get_unsupported_dynamic_vmap_rule)
 get_unsupported_dynamic_vmap_rule = vmap_rules_getters.register(P.Unique)(get_unsupported_dynamic_vmap_rule)
 get_unsupported_dynamic_vmap_rule = \

mindspore/ops/_vmap/vmap_base.py

@@ -21,12 +21,13 @@ from mindspore.common import Tensor
 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.ops.operations import math_ops
 from mindspore.ops.operations import _grad_ops as G
 from mindspore.ops.operations import nn_ops as nps
-from mindspore.ops.composite import _VmapGeneralPreprocess
-from mindspore.ops.primitive import Primitive
-from mindspore.ops.operations.random_ops import UniformCandidateSampler
+from mindspore.ops.function import _VmapGeneralPreprocess
+from mindspore.ops.primitive import Primitive, _PrimitiveC
+from mindspore.ops.operations.random_ops import UniformCandidateSampler, RandomShuffle
 from mindspore.ops._grad.grad_base import BpropRegistry as VmapRuleRegistry
 
 
@@ -41,7 +42,7 @@ def get_vmap_rule(prim, axis_size):
     return None
 
 
-@constexpr
+@_primexpr
 def _get_broadcast_shape_with_front_axis(x_shape, y_shape):
     """ Explicitly matched with the broadcast shape, that is, 1 is added to the broadcast position. """
     x_len = len(x_shape)
@@ -86,7 +87,7 @@ def _handle_broadcasting(x, x_shape, y_shape):
     return F.reshape(x, broadcast_shape)
 
 
-@constexpr
+@_primexpr
 def _get_broadcasting_with_front_axis_additional_axis(x_shape, y_shape):
     """ Get the axes that are inserted after broadcasting.
     Args:
@@ -129,15 +130,19 @@ def _raise_value_error(info, param=None):
     raise ValueError(info + f"{param}")
 
 
-@constexpr
+@_primexpr
 def _get_broadcast_shape(x_shape, dst, axis_size):
     """Get the target shape for broadcast array."""
+    def _check(dst, broadcast_ndim):
+        if dst < -broadcast_ndim or dst >= broadcast_ndim:
+            _raise_value_error("Destination axis {} is out of bounds for array of dimension"
+                               " [{}, {}).".format(dst, -broadcast_ndim, broadcast_ndim))
+
     x_ndim = len(x_shape)
     broadcast_ndim = x_ndim + 1
 
-    if dst < -broadcast_ndim or dst >= broadcast_ndim:
-        _raise_value_error("Destination axis {} is out of bounds for array of dimension"
-                           " [{}, {}).".format(dst, -broadcast_ndim, broadcast_ndim))
+    _check(dst, broadcast_ndim)
+
     if dst < 0:
         dst = broadcast_ndim + dst
 
@@ -190,6 +195,10 @@ def vmap_unstack(dim, val):
     return P.Unstack(dim)(val)
 
 
+def vmap_stack(val):
+    return P.Stack()(val)
+
+
 def vmap_general_output_process(output):
     """ Match output to axis 0"""
     vals_out_tuple = ()
@@ -416,6 +425,8 @@ def _vmap_clone_prim(prim):
     """
     Cloning a new primitive object same as `prim`.
     """
+    if isinstance(prim, _PrimitiveC):
+        return _PrimitiveC(prim.name, prim.attrs)
     new_ops = _ops_vmap_clone_prim_dict.get(prim.name, None)
     if new_ops is None:
         raise ValueError("Failed to get the primitive object of {} from `_ops_vmap_clone_prim_dict`. Please register "
@@ -433,7 +444,7 @@ def _vmap_clone_prim(prim):
     return cloned
 
 
-@constexpr
+@_primexpr
 def _get_reduce_batch_axis(axis, x_dim, x_ndim):
     """get batch_axis for reduce* operation."""
     # For axis, it's value in Union[int, list, tuple]
@@ -481,6 +492,7 @@ _ops_vmap_clone_prim_dict = {
     "ApplyAdaMax": P.ApplyAdaMax,
     "ApplyAdadelta": P.ApplyAdadelta,
     "ApplyRMSProp": P.ApplyRMSProp,
+    'Adam': P.Adam,
     "ApplyCenteredRMSProp": P.ApplyCenteredRMSProp,
     "ApplyFtrl": P.ApplyFtrl,
     "ApplyGradientDescent": P.ApplyGradientDescent,
@@ -508,4 +520,6 @@ _ops_vmap_clone_prim_dict = {
     "SparseApplyAdagrad": P.SparseApplyAdagrad,
     "SparseApplyAdagradV2": P.SparseApplyAdagradV2,
     "SparseApplyFtrl": P.SparseApplyFtrl,
+    "RandomShuffle": RandomShuffle,
+    "RandomChoiceWithMask": P.RandomChoiceWithMask
 }

mindspore/ops/_vmap/vmap_convolution_ops.py

@@ -16,9 +16,9 @@
 """convolution vmap impl"""
 from __future__ import absolute_import
 
-import numpy as np
 import mindspore.numpy as mnp
 from mindspore.ops import constexpr
+from mindspore.ops.primitive import _primexpr
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
 from mindspore.ops.operations import nn_ops as nps
@@ -142,7 +142,7 @@ def get_conv3d_backprop_filter_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
-@constexpr
+@_primexpr
 def _get_reshape_src_dim(data_dim, cmp_dim):
     """Get source dim for reshape"""
     if data_dim > cmp_dim:
@@ -154,7 +154,7 @@ def _get_reshape_src_dim(data_dim, cmp_dim):
     return expand_dim, merge_dim
 
 
-@constexpr
+@_primexpr
 def _get_merge_shape(src_dim, dst_dim, shape):
     """Get new shape for merging the src_dim and dst_dim. The dst_dim is the value after removing src_dim."""
     new_shape = [shape[i] for i in range(len(shape)) if i != src_dim]
@@ -171,13 +171,10 @@ def _reshape_merge_dims(src_dim, dst_dim, target):
     return output, new_shape
 
 
-@constexpr
+@_primexpr
 def _get_expand_shape(src_dim, dst_size, shape, prim_name):
     """Get new shape for splitting src_dim into dst_size parts."""
-    dst_size2, remainder = np.divmod(shape[src_dim], dst_size)
-    if remainder != 0:
-        _raise_value_error("The remainder of {} / {} should be 0, "
-                           "but got {} in {}.".format(shape[src_dim], dst_size, remainder, prim_name))
+    dst_size2 = shape[src_dim] // dst_size
     new_shape = list(shape)
     new_shape[src_dim:(src_dim + 1)] = [dst_size, dst_size2]
     return tuple(new_shape)
@@ -190,7 +187,7 @@ def _reshape_expand_dims(src_dim, dst_size, target, prim_name):
     return F.reshape(target, new_shape)
 
 
-@constexpr
+@_primexpr
 def _get_new_size_by_index(input_size, batch_size, index):
     """Get the new size of input_size by multiplying input_size[index] by batch_size."""
     new_size = ()
@@ -201,7 +198,7 @@ def _get_new_size_by_index(input_size, batch_size, index):
     return tuple(new_size)
 
 
-@constexpr
+@_primexpr
 def _update_group_attr(prim, groups, batch_size):
     """Set new value for 'group' attribute of the convolution primitive."""
     group = groups * batch_size

mindspore/ops/_vmap/vmap_grad_math_ops.py

@@ -17,9 +17,9 @@
 from __future__ import absolute_import
 
 from mindspore.ops import functional as F
-from mindspore.ops import constexpr
+from mindspore.ops.primitive import _primexpr
 from mindspore.ops.operations import _grad_ops as G
-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, _bdim_at_front, \
     _handle_broadcasting, get_unary_grad_vmap_rule, _get_broadcasting_with_front_axis_additional_axis
 
@@ -36,7 +36,7 @@ def get_broadcast_binary_op_grad_vmap_rule(prim, axis_size):
     if isinstance(prim, str):
         prim = broadcast_binary_op_grad_map.get(prim)()
 
-    @constexpr
+    @_primexpr
     def get_longest_shape(x_shape, y_shape, g_shape):
         x_rank = len(x_shape)
         y_rank = len(y_shape)
@@ -148,7 +148,7 @@ def get_median_grad_vmap_rule(prim, axis_size):
     axis = prim.axis
     keep_dims = prim.keep_dims
 
-    @constexpr
+    @_primexpr
     def trans_grad_axis(axis, rank, dim, keep_dims):
         if axis < 0:
             axis += rank - 1