mindspore 2.4.10__cp310-cp310-win_amd64.whl → 2.6.0__cp310-cp310-win_amd64.whl

mindspore/ops/_vmap/vmap_nn_ops.py

@@ -29,9 +29,9 @@ from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_prepr
     _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
 from mindspore.ops.primitive import Primitive
-from mindspore.ops.auto_generate.gen_arg_handler import Format
 from mindspore.ops.auto_generate import Embedding
-from mindspore.ops.auto_generate import gen_arg_handler as handler
+from mindspore.ops._utils import arg_handler as handler
+from mindspore._c_expression import FormatEnum as Format
 
 
 @vmap_rules_getters.register(P.ApplyAdaMax)
@@ -987,28 +987,16 @@ def get_kl_div_loss_grad_vmap_rule(prim, axis_size):
 @vmap_rules_getters.register(P.SmoothL1Loss)
 def get_smooth_l1_loss_vmap_rule(prim, axis_size):
     """VmapRule for `SmoothL1Loss` operation."""
-    if isinstance(prim, str):
-        prim = Primitive(prim)
-        prim_beta = 1.0
-        prim_reduction = 'none'
-    else:
-        prim_reduction = prim.reduction
-        prim_beta = prim.beta
-
-    smooth_l1_loss_op = P.SmoothL1Loss(prim_beta, 'none')
-    if prim_reduction == 'mean':
-        reduce_op = P.ReduceMean()
-    elif prim_reduction == "sum":
-        reduce_op = P.ReduceSum()
-
-    def vmap_rule(x_bdim, target_bdim):
+    def vmap_rule(x_bdim, target_bdim, beta_bdim, reduction_bdim):
         is_all_none, result = vmap_general_preprocess(
-            prim, x_bdim, target_bdim)
+            prim, x_bdim, target_bdim, beta_bdim, reduction_bdim)
         if is_all_none:
             return result
 
         x, x_dim = x_bdim
         target, target_dim = target_bdim
+        beta, _ = beta_bdim
+        reduction, _ = reduction_bdim
         x_ndim = F.rank(x)
         target_ndim = F.rank(target)
         max_rank = max(x_ndim, target_ndim)
@@ -1020,15 +1008,20 @@ def get_smooth_l1_loss_vmap_rule(prim, axis_size):
             reduce_indexes = tuple(range(1, max_rank))
 
         # elementwise style when reduction='none', otherwise reduce style
-        if prim_reduction == "none":
-            out = prim(x, target)
-        elif prim_reduction in ("mean", "sum"):
-            out = smooth_l1_loss_op(x, target)
+        # reduction is number
+        none_enum = handler.str_to_enum("SmoothL1Loss", "reduction", "none")
+        mean_enum = handler.str_to_enum("SmoothL1Loss", "reduction", "mean")
+        sum_enum = handler.str_to_enum("SmoothL1Loss", "reduction", "sum")
+        if reduction in (none_enum, mean_enum, sum_enum):
+            out = prim(x, target, beta, none_enum)
             if reduce_indexes is not None:
-                out = reduce_op(out, reduce_indexes)
+                if reduction == mean_enum:
+                    out = P.ReduceMean()(out, reduce_indexes)
+                elif reduction == sum_enum:
+                    out = P.ReduceSum()(out, reduce_indexes)
         else:
             raise RuntimeError("For SmoothL1Loss vmap, reduction should be one of "
-                               "['none', 'mean', 'sum'], but got '{}'".format(prim_reduction))
+                               "['none', 'mean', 'sum'], but got '{}'".format(reduction))
         return out, 0
 
     return vmap_rule
@@ -1037,29 +1030,22 @@ def get_smooth_l1_loss_vmap_rule(prim, axis_size):
 @vmap_rules_getters.register(G.SmoothL1LossGrad)
 def get_smooth_l1_loss_grad_vmap_rule(prim, axis_size):
     """VmapRule for `SmoothL1LossGrad`."""
-    if isinstance(prim, str):
-        prim = Primitive(prim)
-        reduction = "none"
-        beta = 1.0
-    else:
-        reduction = prim.reduction
-        beta = prim.beta
-    smooth_l1_loss_grad = G.SmoothL1LossGrad(beta, reduction)
-
-    def vmap_rule(x_bdim, target_bdim, dy_bdim):
+    def vmap_rule(x_bdim, target_bdim, dy_bdim, beta_bdim, reduction_bdim):
         is_all_none, result = vmap_general_preprocess(
-            prim, dy_bdim, x_bdim, target_bdim)
+            prim, dy_bdim, x_bdim, target_bdim, beta_bdim, reduction_bdim)
         if is_all_none:
             return result
 
         dy, dy_dim = dy_bdim
         x, x_dim = x_bdim
         target, target_dim = target_bdim
+        beta, _ = beta_bdim
+        reduction, _ = reduction_bdim
         dy = _bdim_at_front(dy, dy_dim, axis_size)
         x = _bdim_at_front(x, x_dim, axis_size)
         target = _bdim_at_front(target, target_dim, axis_size)
 
-        out = smooth_l1_loss_grad(x, target, dy)
+        out = prim(x, target, dy, beta, reduction)
         return out, 0
 
     return vmap_rule
@@ -1646,13 +1632,12 @@ def get_apply_adagrad_da_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
-@vmap_rules_getters.register(NN.AdaptiveMaxPool2D)
+@vmap_rules_getters.register(P.AdaptiveMaxPool2D)
 def get_adaptive_max_pool_2d_vmap_rule(prim, axis_size):
     """VmapRule for `AdaptiveMaxPool2D`."""
     nchw_index = 4
     chw_reverse_index = -3
     hw_size = 2
-    output_size = prim.output_size
 
     @_primexpr
     def get_output_shape(x_ori_shape, output_size):
@@ -1675,8 +1660,8 @@ def get_adaptive_max_pool_2d_vmap_rule(prim, axis_size):
             output_shape += (w_out,)
         return output_shape
 
-    def vmap_rule(input_x_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, input_x_bdim)
+    def vmap_rule(input_x_bdim, output_size_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, input_x_bdim, output_size_bdim)
         if is_all_none:
             return result
 
@@ -1684,18 +1669,20 @@ def get_adaptive_max_pool_2d_vmap_rule(prim, axis_size):
         x = _bdim_at_front(input_x, input_x_dim, axis_size)
         x_ndim = F.rank(x)
 
+        output_size, _ = output_size_bdim
+
         if x_ndim > nchw_index:
             # for the case of NCHW
             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)
-            out, indices = prim(x)
+            out, indices = prim(x, output_size)
             out = F.reshape(out, output_shape)
             indices = F.reshape(indices, output_shape)
             return (out, 0), (indices, 0)
 
         # for the case of CHW
-        out, indices = prim(x)
+        out, indices = prim(x, output_size)
         return (out, 0), (indices, 0)
 
     return vmap_rule

mindspore/ops/auto_generate/__init__.py

@@ -19,13 +19,14 @@ Primitive operator classes and operator functional.
 A collection of operators to build neural networks or to compute functions.
 """
 
-from . import gen_ops_def, gen_arg_handler, gen_arg_dtype_cast
+from . import gen_ops_def
+from .._utils import arg_handler, arg_dtype_cast
 
 from .gen_ops_prim import *
 from .gen_ops_def import *
-from .gen_arg_handler import *
-from .gen_arg_dtype_cast import *
 from ..operations.manually_defined.ops_def import *
+from .._utils.arg_handler import *
+from .._utils.arg_dtype_cast import *
 
 
 __all__ = []