mindspore 2.4.10__cp39-cp39-win_amd64.whl → 2.6.0__cp39-cp39-win_amd64.whl

mindspore/{ops_generate → ops/_utils}/arg_dtype_cast.py

@@ -20,14 +20,133 @@ import mindspore as ms
 from mindspore import ops
 from mindspore.common.tensor import Tensor
 from mindspore.ops.operations._sequence_ops import TensorToScalar, TensorToTuple
-from mindspore.ops_generate.gen_ops_inner_prim import TupleToList, ListToTuple
 from mindspore._c_expression import OpDtype
+from mindspore._c_expression import typing
+from mindspore._c_expression import op_enum
+from mindspore.ops.primitive import Primitive, prim_attr_register, prim_arg_register
 
 tensor_to_tuple_ = TensorToTuple()
+tensor_to_scalar_ = TensorToScalar()
+
+
+class TupleToList(Primitive):
+    r"""
+    Convert tuple to list.
+
+    Inputs:
+        - **x** (tuple) - The input
+
+    Outputs:
+        List, has the same elements as the `input`.
+
+    Supported Platforms:
+        ``CPU``
+
+    Examples:
+        >>> from mindspore.ops._utils.arg_dtype_cast import TupleToList
+        >>> x = (1, 2, 3)
+        >>> result = TupleToList()(x)
+        >>> print(result)
+        [1, 2, 3]
+    """
+    @prim_arg_register
+    def __init__(self):
+        """Initialize TupleToList"""
+
+    def __call__(self, input):
+        return list(input)
+
+
+class ListToTuple(Primitive):
+    r"""
+    Convert list to tuple.
+
+    Inputs:
+        - **x** (list) - The input
+
+    Outputs:
+        Tuple, has the same elements as the `input`.
+
+    Supported Platforms:
+        ``CPU``
+
+    Examples:
+        >>> from mindspore.ops._utils.arg_dtype_cast import ListToTuple
+        >>> x = [1, 2, 3]
+        >>> result = ListToTuple()(x)
+        >>> print(result)
+        (1, 2, 3)
+    """
+    @prim_arg_register
+    def __init__(self):
+        """Initialize TupleToList"""
+
+    def __call__(self, input):
+        return tuple(input)
+
+
 tuple_to_list = TupleToList()
 list_to_tuple = ListToTuple()
 
 
+class DtypeToEnum(Primitive):
+    r"""
+    Convert mindspore dtype to enum.
+
+    Inputs:
+        - **op_name** (str) - The op name
+        - **arg_name** (str) - The arg name
+        - **dtype** (mindspore.dtype) - The data type.
+
+    Outputs:
+        An integer.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        """Initialize"""
+
+    def __call__(self, op_name, arg_name, dtype):
+        """Run in PyNative mode"""
+        if not isinstance(dtype, typing.Type):
+            raise TypeError(
+                f"For '{op_name}', the input '{arg_name}' should be mindspore dtype, but got {dtype}.")
+        return typing.type_to_type_id(dtype)
+
+
+class StringToEnum(Primitive):
+    r"""
+    Convert string to enum.
+
+    Inputs:
+        - **op_name** (str) - The op name
+        - **arg_name** (str) - The arg name
+        - **enum_str** (str) - The str data.
+
+    Outputs:
+        An integer.
+
+    Supported Platforms:
+        ``CPU``
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        """Initialize"""
+
+    def __call__(self, op_name, arg_name, enum_str):
+        """Run in PyNative mode"""
+        if enum_str is None:
+            return None
+        if not isinstance(enum_str, str):
+            raise TypeError(
+                f"For '{op_name}', the input '{arg_name}' should be a str, but got {type(enum_str)}.")
+        return op_enum.str_to_enum(op_name, arg_name, enum_str)
+
+
 def int_to_float(data):
     return float(data)
 
@@ -184,7 +303,7 @@ def get_support_dtype_list(src_type, dst_type):
     return support_list
 
 
-def to_py_number(data, dst_type):
+def tensor_to_number(data, dst_type):
     """Convert tensor to python number"""
     if dst_type == DT_INT_VAL:
         data = ops.cast(data, ms.int64)
@@ -197,7 +316,7 @@ def to_py_number(data, dst_type):
             data = ops.cast(data, ms.int64)
         elif src_type in (ms.bfloat16, ms.float16, ms.float32, ms.float64):
             data = ops.cast(data, ms.float32)
-    return TensorToScalar()(data)
+    return tensor_to_scalar_(data)
 
 
 def do_type_cast(data, dst_type):
@@ -230,7 +349,7 @@ def do_type_cast(data, dst_type):
             return list_to_tensor(data)
     elif is_number(dst_type):
         if isinstance(data, Tensor):
-            return to_py_number(data, dst_type)
+            return tensor_to_number(data, dst_type)
     raise TypeError("Type conversion failed.")
 
 

mindspore/{ops_generate → ops/_utils}/arg_handler.py

@@ -14,13 +14,11 @@
 # ============================================================================
 """Operator argument handle function."""
 
-from mindspore.ops_generate.gen_ops_inner_prim import DtypeToEnum, StringToEnum
-# Enum Class:
-from mindspore._c_expression import FormatEnum as Format
-from mindspore._c_expression import ReductionEnum as Reduction
 from mindspore.common import Tensor
 from mindspore.common import dtype as mstype
 
+from .arg_dtype_cast import DtypeToEnum, StringToEnum
+
 
 def arg_invalid_info(op_name, arg_name, arg_val):
     """
@@ -116,67 +114,6 @@ def to_2d_paddings(op_name, arg_name, pad):
     raise ValueError(arg_invalid_info(op_name, arg_name, pad))
 
 
-def to_paddings(op_name, arg_name, pad):
-    """
-    convert paddings: int -> tuple[int*4].
-    """
-    if isinstance(pad, int):
-        return (pad,) * 4
-    if isinstance(pad, (tuple, list)):
-        return pad
-    raise ValueError(arg_invalid_info(op_name, arg_name, pad))
-
-
-def to_3d_kernel_size(op_name, arg_name, kernel_size):
-    """
-    convert 3d kernel_size: int/tuple[int*6] -> tuple[int*3].
-    """
-    if isinstance(kernel_size, int):
-        return (kernel_size, kernel_size, kernel_size)
-    if isinstance(kernel_size, (tuple, list)):
-        if len(kernel_size) == 5:
-            return (kernel_size[2], kernel_size[3], kernel_size[4])
-        return kernel_size
-    raise ValueError(arg_invalid_info(op_name, arg_name, kernel_size))
-
-
-def to_3d_strides(op_name, arg_name, stride):
-    """
-    convert 3d stride: int/tuple[int*6] -> tuple[int*3].
-    """
-    if isinstance(stride, int):
-        return (stride, stride, stride)
-    if isinstance(stride, (tuple, list)):
-        if len(stride) == 5:
-            return (stride[2], stride[3], stride[4])
-        return stride
-    raise ValueError(arg_invalid_info(op_name, arg_name, stride))
-
-
-def to_3d_dilations(op_name, arg_name, dilation):
-    """
-    convert 3d dilation: int/tuple[int*6] -> tuple[int*3].
-    """
-    if isinstance(dilation, int):
-        return (dilation, dilation, dilation)
-    if isinstance(dilation, (tuple, list)):
-        if len(dilation) == 5:
-            return (dilation[2], dilation[3], dilation[4])
-        return dilation
-    raise ValueError(arg_invalid_info(op_name, arg_name, dilation))
-
-
-def to_3d_paddings(op_name, arg_name, pad):
-    """
-    convert 3d paddings: int -> tuple[int*6].
-    """
-    if isinstance(pad, int):
-        return (pad,) * 6
-    if isinstance(pad, (tuple, list)):
-        return pad
-    raise ValueError(arg_invalid_info(op_name, arg_name, pad))
-
-
 def generator_handler(op_name, arg_name, inputs):
     """
     convert constant value in tuple to tensor
@@ -189,6 +126,7 @@ def generator_handler(op_name, arg_name, inputs):
             new_inputs.append(input_)
     return tuple(new_inputs)
 
+
 dtype_to_type_id = DtypeToEnum()
 
 # string to enum

mindspore/ops/_vmap/vmap_array_ops.py

@@ -15,12 +15,13 @@
 
 """array_ops vmap impl."""
 from __future__ import absolute_import
+from enum import Enum
 
 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._c_expression import TensorPy as Tensor_
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
 from mindspore.ops.primitive import constexpr, _primexpr
@@ -140,6 +141,8 @@ def _get_prefix(indices_shape, axis_size, indices_dtype):
     the generated prefix is a Tensor([[[0], [0]],
                                       [[1], [1]]])
     """
+    cast_op = P.Cast()
+
     def _check(indices_shape):
         if not indices_shape:
             raise ValueError("indices_shape is empty in _get_prefix.")
@@ -147,8 +150,8 @@ def _get_prefix(indices_shape, axis_size, indices_dtype):
     _check(indices_shape)
     indices_len = len(indices_shape)
     if indices_len == 1:
-        prefix = P.Range()(Tensor(0, indices_dtype), Tensor(axis_size, indices_dtype), Tensor(1, indices_dtype))
-        return prefix
+        prefix = P.Range()(0, axis_size, 1)
+        return cast_op(prefix, indices_dtype)
 
     indices_end = indices_len - 1
     prefix_shape = ()
@@ -163,9 +166,8 @@ def _get_prefix(indices_shape, axis_size, indices_dtype):
         else:
             expand_shape = expand_shape + (1,)
 
-    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
+    prefix = P.BroadcastTo(prefix_shape)(P.Reshape()(P.Range()(0, axis_size, 1), expand_shape))
+    return cast_op(prefix, indices_dtype)
 
 
 @vmap_rules_getters.register(P.Transpose)
@@ -212,6 +214,31 @@ def get_transpose_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register("TransposeExtView")
+def get_transpose_ext_vmap_rule(prim, axis_size):
+    """VmapRule for `TransposeExtView` operation."""
+    if isinstance(prim, str):
+        prim = Primitive(prim)
+
+    def vmap_rule(x_bdim, dim1_bdim, dim2_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x_bdim, dim1_bdim, dim2_bdim)
+        if is_all_none:
+            return result
+
+        x, dim = x_bdim
+        dim1, dim1_dim = dim1_bdim
+        dim2, dim2_dim = dim2_bdim
+        if dim1_dim is not None or dim2_dim is not None:
+            _raise_value_error("The source axis of dim1_dim and dim2_dim in `TransposeExtView` must be None, "
+                               "but got {} and {}.".format(dim1_dim, dim2_dim))
+        batch_dim1 = dim1 if dim1 < dim else dim1 + 1
+        batch_dim2 = dim2 if dim2 < dim else dim2 + 1
+        out = prim(x, batch_dim1, batch_dim2)
+        return out, dim
+
+    return vmap_rule
+
+
 @vmap_rules_getters.register("Tile")
 def get_tile_vmap_rule(prim, axis_size):
     """VmapRule for `P.Tile` operation."""
@@ -1488,16 +1515,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 +1548,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 +1563,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 +1921,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 +1939,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

@@ -25,7 +25,9 @@ from mindspore.ops.primitive import _primexpr
 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
+from mindspore._c_expression import FormatEnum as Format
+from mindspore._c_expression import ReductionEnum as Reduction
 
 
 @vmap_rules_getters.register(G.NLLLossGrad)
@@ -225,33 +227,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 +677,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 +691,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 +705,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)

mindspore/ops/_vmap/vmap_math_ops.py

@@ -19,6 +19,7 @@ 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 auto_generate as gen
 from mindspore.ops.auto_generate import MatMulExt
 from mindspore.ops.primitive import _primexpr
 from mindspore.common import Tensor
@@ -29,7 +30,7 @@ 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, 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
+    _bdim_at_any, _get_reduce_batch_axis, _get_reduce_out_dim
 from mindspore.ops.operations.math_ops import Bernoulli, BesselI0, BesselI1, BesselJ0, BesselJ1, \
     BesselK0, BesselK0e, BesselY0, BesselY1, BesselK1, BesselK1e, Median
 
@@ -128,28 +129,29 @@ def get_addcxxx_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
-@vmap_rules_getters.register(P.Cdist)
+@vmap_rules_getters.register(gen.Cdist)
 def get_cdist_vmap_rule(prim, axis_size):
     """VmapRule for `cdist` 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(x_bdim, y_bdim):
+    def vmap_rule(x_bdim, y_bdim, p_bdim):
         x, x_dim = x_bdim
         y, y_dim = y_bdim
+        p, _ = p_bdim
 
-        if x_dim is None and y_dim is None:
+        if x_dim is None and y_dim is None and p is None:
             out = prim(x, y)
             return (out, None)
         x = _bdim_at_front(x, x_dim, axis_size)
         y = _bdim_at_front(y, y_dim, axis_size)
 
-        out = batch_prim(x, y)
+        out = prim(x, y, p)
         return out, 0
 
     return vmap_rule
@@ -559,20 +561,17 @@ def get_index_add_vmap_rule(prim, axis_size):
 @vmap_rules_getters.register(linalg_ops.Svd)
 def get_svd_vmap_rule(prim, axis_size):
     """VmapRule for 'Svd' operation."""
-    if isinstance(prim, str):
-        prim = Primitive(prim)
-        compute_uv = True
-    else:
-        compute_uv = prim.compute_uv
 
-    def vmap_rule(x_bdim):
+    def vmap_rule(x_bdim, full_matrices_bdim, compute_uv_bdim):
         is_all_none, result = vmap_general_preprocess(prim, x_bdim)
         if is_all_none:
             return result
 
         x, x_dim = x_bdim
+        full_matrices, _ = full_matrices_bdim
+        compute_uv, _ = compute_uv_bdim
         x = _bdim_at_front(x, x_dim, axis_size)
-        s, u, v = prim(x)
+        s, u, v = prim(x, full_matrices, compute_uv)
         if compute_uv:
             return (s, 0), (u, 0), (v, 0)
         return (s, 0), (u, None), (v, None)