mindspore 2.3.0__cp39-none-any.whl → 2.3.0rc2__cp39-none-any.whl

mindspore/nn/loss/loss.py

@@ -1820,10 +1820,10 @@ class MultilabelMarginLoss(LossBase):
 
 class BCEWithLogitsLoss(LossBase):
     r"""
-    Adds sigmoid activation function to input `input` as logits, and uses the given logits to compute binary cross
-    entropy between the `input` and the `target`.
+    Adds sigmoid activation function to input logits, and uses the given logits to compute binary cross entropy
+    between the logits and the labels.
 
-    Sets input `input` as :math:`X`, input `target` as :math:`Y`, output as :math:`L`. Then,
+    Sets input `logits` as :math:`X`, input `labels` as :math:`Y`, output as :math:`L`. Then,
 
     .. math::
         p_{ij} = sigmoid(X_{ij}) = \frac{1}{1 + e^{-X_{ij}}}
@@ -1849,29 +1849,29 @@ class BCEWithLogitsLoss(LossBase):
             - ``'sum'``: the output elements will be summed.
 
         weight (Tensor, optional): A rescaling weight applied to the loss of each batch element.
-            If not None, it can be broadcast to a tensor with shape of `input`,
+            If not None, it can be broadcast to a tensor with shape of `logits`,
             data type must be float16 or float32. Default: ``None`` .
         pos_weight (Tensor, optional): A weight of positive examples. Must be a vector with length equal to the
-            number of classes. If not None, it must be broadcast to a tensor with shape of `input`, data type
+            number of classes. If not None, it must be broadcast to a tensor with shape of `logits`, data type
             must be float16 or float32. Default: ``None`` .
 
     Inputs:
-        - **input** (Tensor) - Input `input` with shape :math:`(N, *)` where :math:`*` means, any number
+        - **logits** (Tensor) - Input logits with shape :math:`(N, *)` where :math:`*` means, any number
           of additional dimensions. The data type must be float16 or float32.
-        - **target** (Tensor) - Ground truth label with shape :math:`(N, *)` where :math:`*` means, any number
-          of additional dimensions. The same shape and data type as `input`.
+        - **labels** (Tensor) - Ground truth label with shape :math:`(N, *)` where :math:`*` means, any number
+          of additional dimensions. The same shape and data type as `logits`.
 
     Outputs:
-        Tensor or Scalar, if `reduction` is ``'none'``, its shape is the same as `input`.
+        Tensor or Scalar, if `reduction` is ``'none'``, its shape is the same as `logits`.
         Otherwise, a scalar value will be returned.
 
     Raises:
-        TypeError: If input `input` or `target` is not Tensor.
-        TypeError: If data type of `input` or `target` is neither float16 nor float32.
+        TypeError: If input `logits` or `labels` is not Tensor.
+        TypeError: If data type of `logits` or `labels` is neither float16 nor float32.
         TypeError: If `weight` or `pos_weight` is a parameter.
         TypeError: If data type of `weight` or `pos_weight` is neither float16 nor float32.
         TypeError: If data type of `reduction` is not string.
-        ValueError: If `weight` or `pos_weight` can not be broadcast to a tensor with shape of `input`.
+        ValueError: If `weight` or `pos_weight` can not be broadcast to a tensor with shape of `logits`.
         ValueError: If `reduction` is not one of ``'none'``, ``'mean'``, ``'sum'``.
 
     Supported Platforms:
@@ -1881,10 +1881,10 @@ class BCEWithLogitsLoss(LossBase):
         >>> import mindspore as ms
         >>> import mindspore.nn as nn
         >>> import numpy as np
-        >>> input = ms.Tensor(np.array([[-0.8, 1.2, 0.7], [-0.1, -0.4, 0.7]]).astype(np.float32))
-        >>> target = ms.Tensor(np.array([[0.3, 0.8, 1.2], [-0.6, 0.1, 2.2]]).astype(np.float32))
+        >>> logits = ms.Tensor(np.array([[-0.8, 1.2, 0.7], [-0.1, -0.4, 0.7]]).astype(np.float32))
+        >>> labels = ms.Tensor(np.array([[0.3, 0.8, 1.2], [-0.6, 0.1, 2.2]]).astype(np.float32))
         >>> loss = nn.BCEWithLogitsLoss()
-        >>> output = loss(input, target)
+        >>> output = loss(logits, labels)
         >>> print(output)
         0.3463612
     """
@@ -1900,10 +1900,10 @@ class BCEWithLogitsLoss(LossBase):
         self.weight = weight
         self.pos_weight = pos_weight
 
-    def construct(self, input, target):
-        _check_is_tensor('input', input, self.cls_name)
-        _check_is_tensor('target', target, self.cls_name)
-        loss = ops.binary_cross_entropy_with_logits(input, target, self.weight, self.pos_weight, self.reduction)
+    def construct(self, logits, labels):
+        _check_is_tensor('logits', logits, self.cls_name)
+        _check_is_tensor('labels', labels, self.cls_name)
+        loss = ops.binary_cross_entropy_with_logits(logits, labels, self.weight, self.pos_weight, self.reduction)
         return loss
 
 

mindspore/nn/optim/adasum.py

@@ -29,7 +29,7 @@ from mindspore.parallel._utils import _get_global_rank, _get_stage_device_num
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
 from mindspore.ops import operations as P
-from mindspore.ops import Send, Receive
+from mindspore.ops.operations._inner_ops import Send, Receive
 from mindspore.common.tensor import Tensor
 from mindspore.common import dtype as mstype
 from mindspore.communication.management import create_group

mindspore/nn/optim/sgd.py

@@ -195,9 +195,9 @@ class SGD(Optimizer):
                             "or 'weight_decay' set in grouped 'params' must be float or int type.")
 
         if hasattr(self, "group_weight_decay") and self.group_weight_decay:
-            self.opt = tuple(P.SGD(dampening, 0.0, nesterov) for wd in self.group_weight_decay)
+            self.opt = tuple(P.SGD(dampening, wd, nesterov) for wd in self.group_weight_decay)
         else:
-            self.opt = tuple([P.SGD(dampening, 0.0, nesterov)] * len(self._parameters))
+            self.opt = tuple([P.SGD(dampening, float(weight_decay), nesterov)] * len(self._parameters))
 
         self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")
 
@@ -222,7 +222,6 @@ class SGD(Optimizer):
         params = self._parameters
         accum = self.accum
         stat = self.stat
-        gradients = self.decay_weight(gradients)
         gradients = self.flatten_gradients(gradients)
         gradients = self.gradients_centralization(gradients)
         gradients = self.scale_grad(gradients)

mindspore/nn/probability/distribution/exponential.py

@@ -152,7 +152,7 @@ class Exponential(Distribution):
         if self.rate is not None:
             check_greater_zero(self.rate, 'rate')
 
-        self.minval = np.finfo(np.float_).tiny
+        self.minval = np.finfo(np.float).tiny
 
         # ops needed for the class
         self.exp = exp_generic

mindspore/nn/probability/distribution/geometric.py

@@ -152,7 +152,7 @@ class Geometric(Distribution):
         if self._probs is not None:
             check_prob(self.probs)
 
-        self.minval = np.finfo(np.float_).tiny
+        self.minval = np.finfo(np.float).tiny
 
         # ops needed for the class
         self.exp = exp_generic

mindspore/nn/probability/distribution/logistic.py

@@ -170,7 +170,7 @@ class Logistic(Distribution):
         self.neg = P.Neg()
 
         self.threshold = np.log(np.finfo(np.float32).eps) + 1.
-        self.tiny = np.finfo(np.float_).tiny
+        self.tiny = np.finfo(np.float).tiny
         self.sd_const = np.pi / np.sqrt(3)
 
     def _softplus(self, x):

mindspore/nn/wrap/cell_wrapper.py

@@ -17,18 +17,16 @@
 from __future__ import absolute_import
 from __future__ import division
 
-import os
 from types import FunctionType, MethodType
 
 from mindspore import log as logger
 from mindspore.parallel._utils import _get_device_num, _get_gradients_mean,\
     _get_parallel_mode, _get_enable_parallel_optimizer, _is_pynative_parallel
-from mindspore.context import ParallelMode, GRAPH_MODE, get_context
+from mindspore.context import ParallelMode
 from mindspore import _checkparam as validator
 from mindspore import ops, nn
 from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter, ParameterTuple
-from mindspore.common.tensor import Tensor
 from mindspore.ops.primitive import _primexpr
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
@@ -742,18 +740,6 @@ class _TrainGradAccuStepCell(TrainOneStepCell):
         self.hyper_map = ops.HyperMap()
         self.opt_shard = _get_enable_parallel_optimizer()
         self._get_attr_from_cell(network)
-        self.enable_mindio = False
-        mode = get_context("mode")
-        device_type = get_context("device_target")
-        if device_type != "Ascend" or mode != GRAPH_MODE:
-            return
-        graceful_exit = os.getenv("MS_ENABLE_MINDIO_GRACEFUL_EXIT")
-        ttp_lib_path = os.getenv("MS_MINDIO_TTP_LIB_PATH")
-        ttp_path_check = ttp_lib_path is not None and os.path.isfile(ttp_lib_path)
-        if graceful_exit == "true" and ttp_path_check:
-            self.g_one = Tensor([0.1])
-            self.allreduce_sum = ops.AllReduce()
-            self.enable_mindio = True
 
     def construct(self, *inputs):
         if not self.sense_flag:
@@ -762,11 +748,6 @@ class _TrainGradAccuStepCell(TrainOneStepCell):
         sens = ops.fill(ops.DType()(loss), ops.Shape()(loss), self.sens)
         grads = self.grad(self.network, self.weights)(*inputs, sens)
         accu_grads = ops.depend(self.accu_grads, grads)
-        if self.enable_mindio:
-            g_one = ops.depend(self.g_one, accu_grads)
-            g_one_res = self.allreduce_sum(g_one)
-            accu_grads = ops.depend(accu_grads, g_one_res)
-            grads = ops.depend(grads, g_one_res)
         if self.opt_shard:
             succ = self.optimizer(grads)
         else:
@@ -781,11 +762,6 @@ class _TrainGradAccuStepCell(TrainOneStepCell):
         loss = self.network(*inputs)
         grads = self.grad_no_sens(self.network, self.weights)(*inputs)
         accu_grads = ops.depend(self.accu_grads, grads)
-        if self.enable_mindio:
-            g_one = ops.depend(self.g_one, accu_grads)
-            g_one_res = self.allreduce_sum(g_one)
-            accu_grads = ops.depend(accu_grads, g_one_res)
-            grads = ops.depend(grads, g_one_res)
         if self.opt_shard:
             succ = self.optimizer(grads)
         else:

mindspore/nn/wrap/loss_scale.py

@@ -29,7 +29,6 @@ from mindspore.ops.operations.math_ops import NPUGetFloatStatusV2, NPUClearFloat
 from mindspore.ops import functional as F
 from mindspore.ops import composite as C
 from mindspore.ops import operations as P
-from mindspore.ops.operations.nn_ops import AllFinite
 from mindspore.common import dtype as mstype
 from mindspore.common.api import jit
 from mindspore._c_expression import MSContext
@@ -373,15 +372,6 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         self.loss_scaling_manager = None
         self._ascend_check_overflow_mode = os.environ.get('MS_ASCEND_CHECK_OVERFLOW_MODE')
 
-        self.enable_allfinite = False
-        runtime_conf = os.environ.get('MS_DEV_RUNTIME_CONF')
-        global_jit_config = context.get_jit_config()
-        if runtime_conf is not None and ("all_finite:True" in runtime_conf or "all_finite:true" in runtime_conf):
-            self.enable_allfinite = True
-        elif runtime_conf is not None and ("all_finite:False" in runtime_conf or "all_finite:false" in runtime_conf):
-            self.enable_allfinite = False
-        elif global_jit_config:
-            self.enable_allfinite = global_jit_config["jit_level"] == "O0" or global_jit_config["jit_level"] == "O1"
 
         if isinstance(scale_sense, Cell):
             self.loss_scaling_manager = scale_sense
@@ -488,15 +478,6 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
             overflow = self.less_equal(self.base, flag_sum)
         return overflow
 
-    def _get_distributed_overflow_status_on_infnan_enable_allfinite(self, compute_output):
-        """check overflow status on infnan kernel mode."""
-        overflow = AllFinite()(compute_output)
-
-        if self.is_distributed:
-            overflow = P.Cast()(overflow, mstype.int8)
-            overflow = P.Cast()(self.allreduce(overflow), mstype.bool_)
-        return overflow
-
     def _get_gpu_overflow_status(self, compute_output):
         """get overflow status of gpu."""
         overflow = self._get_distributed_overflow_status_on_infnan_mode(_grad_overflow, compute_output)
@@ -504,11 +485,7 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
 
     def _get_ascend_overflow_status_on_infnan_mode(self, compute_output):
         """get overflow status of ascend on infnan mode."""
-        overflow = False
-        if self.enable_allfinite:
-            overflow = self._get_distributed_overflow_status_on_infnan_enable_allfinite(compute_output)
-        else:
-            overflow = self._get_distributed_overflow_status_on_infnan_mode(_ascend_grad_overflow, compute_output)
+        overflow = self._get_distributed_overflow_status_on_infnan_mode(_ascend_grad_overflow, compute_output)
         return overflow
 
     def _get_ascend_overflow_status_on_saturation_mode(self, status, compute_output):

mindspore/numpy/array_ops.py

@@ -2606,11 +2606,7 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False):
         array1 = ar1.ravel()
         array2 = ar2.ravel()
     concat_array = concatenate((array1, array2))
-    if return_indices:
-        concat_sort_indices = F.argsort(concat_array)
-        concat_array = concat_array[concat_sort_indices]
-    else:
-        concat_array, concat_sort_indices = concat_array.sort()
+    concat_array, concat_sort_indices = concat_array.sort()
 
     mask_res = concat_array[1:] == concat_array[:-1]
     res = F.masked_select(concat_array[1:], mask_res)

mindspore/numpy/dtypes.py

@@ -86,7 +86,7 @@ dtype_map = {
 }
 
 all_types = [
-    'np.int_',
+    'np.int',
     'np.int8',
     'np.int16',
     'np.int32',
@@ -96,11 +96,11 @@ all_types = [
     'np.uint16',
     'np.uint32',
     'np.uint64',
-    'np.float_',
+    'np.float',
     'np.float16',
     'np.float32',
     'np.float64',
-    'np.bool_']
+    'np.bool']
 
 promotion_rule = {
     (uint8, uint16): uint16,

mindspore/numpy/math_ops.py

@@ -4166,18 +4166,18 @@ def multi_dot(arrays):
     Examples:
         >>> import mindspore.numpy as np
         >>> A = np.ones((10000, 100))
-        >>> B = np.ones((100, 100))
-        >>> C = np.ones((100, 5))
+        >>> B = np.ones((100, 1000))
+        >>> C = np.ones((1000, 5))
         >>> D = np.ones((5, 333))
         >>> output = np.multi_dot([A, B, C, D])
         >>> print(output)
-        [[50000. 50000. 50000. ... 50000. 50000. 50000.]
-        [50000. 50000. 50000. ... 50000. 50000. 50000.]
-        [50000. 50000. 50000. ... 50000. 50000. 50000.]
+        [[500000. 500000. 500000. ... 500000. 500000. 500000.]
+        [500000. 500000. 500000. ... 500000. 500000. 500000.]
+        [500000. 500000. 500000. ... 500000. 500000. 500000.]
         ...
-        [50000. 50000. 50000. ... 50000. 50000. 50000.]
-        [50000. 50000. 50000. ... 50000. 50000. 50000.]
-        [50000. 50000. 50000. ... 50000. 50000. 50000.]]
+        [500000. 500000. 500000. ... 500000. 500000. 500000.]
+        [500000. 500000. 500000. ... 500000. 500000. 500000.]
+        [500000. 500000. 500000. ... 500000. 500000. 500000.]]
     """
     if len(arrays) < 2:
         _raise_value_error('Expecting at least 2 arrays')

mindspore/ops/__init__.py

@@ -44,7 +44,7 @@ __primitive__ = [
 __all__ = ["get_vm_impl_fn", "vm_impl_registry",
            "op_info_register", "custom_info_register", "AkgGpuRegOp", "AkgAscendRegOp", "AiCPURegOp", "TBERegOp",
            "CpuRegOp", "CustomRegOp", "DataType",
-           "constexpr", "reshard"]
+           "constexpr"]
 __all__.extend(__primitive__)
 __all__.extend(composite.__all__)
 __all__.extend(operations.__all__)

mindspore/ops/_grad_experimental/grad_comm_ops.py

@@ -22,8 +22,7 @@ from mindspore.ops import functional as F
 from mindspore.communication import get_rank, get_group_size
 from mindspore.parallel._utils import _get_enable_parallel_optimizer, _get_grad_accumulation_shard
 from mindspore.ops import operations as P
-from mindspore.ops import Send, Receive
-from mindspore.ops.operations._inner_ops import issubclass_
+from mindspore.ops.operations._inner_ops import Send, Receive, issubclass_
 from mindspore.common.sparse_tensor import RowTensorInner
 from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
 from mindspore.ops.operations.comm_ops import (AllGather, _MiniStepAllGather, _HostAllGather, AllReduce,
@@ -31,7 +30,7 @@ from mindspore.ops.operations.comm_ops import (AllGather, _MiniStepAllGather, _H
                                                _GetTensorSlice, _MirrorOperator, _MirrorMiniStepOperator, ReduceOp,
                                                ReduceScatter, _HostReduceScatter, _VirtualDiv, _VirtualAdd, _AllSwap,
                                                _VirtualAssignAdd, _VirtualAccuGrad, _MirrorMicroStepOperator,
-                                               _MicroStepAllGather, Reduce, CollectiveGather, CollectiveScatter)
+                                               _MicroStepAllGather)
 from mindspore.ops._grad_experimental.grad_base import bprop_getters
 from mindspore.ops.operations import _grad_ops as G
 
@@ -211,17 +210,21 @@ def get_bprop_mirror_micro_step_operator(self):
     def bprop(x, z, out, dout):
         real_grad = z
         assign_out = dout
-        if issubclass_(F.typeof(dout), mstype.tensor_type):
-            z = F.depend(z, dout)
-            if dev_num > 1:
+        if mean_flag:
+            if issubclass_(F.typeof(dout), mstype.tensor_type):
+                z = F.depend(z, dout)
                 real_grad = all_reduce(z)
-                if mean_flag:
-                    real_grad = F.tensor_mul(real_grad, scale)
-            else:
-                real_grad = z
-            if opt_shard:
-                return (real_grad, cast(out_tensor, dtype(z)))
-            return F.depend((cast(out_tensor, dtype(x)), cast(out_tensor, dtype(z))), assign(z, real_grad))
+                real_grad = F.tensor_mul(real_grad, scale)
+                if opt_shard:
+                    return (real_grad, cast(out_tensor, dtype(z)))
+                return F.depend((cast(out_tensor, dtype(x)), cast(out_tensor, dtype(z))), assign(z, real_grad))
+        else:
+            if issubclass_(F.typeof(dout), mstype.tensor_type):
+                z = F.depend(z, dout)
+                real_grad = all_reduce(z)
+                if opt_shard:
+                    return (real_grad, cast(out_tensor, dtype(z)))
+                return F.depend((cast(out_tensor, dtype(x)), cast(out_tensor, dtype(z))), assign(z, real_grad))
         return F.depend((cast(out_tensor, dtype(x)), cast(out_tensor, dtype(z))), assign_out)
 
     return bprop
@@ -241,13 +244,11 @@ def get_bprop_broad_cast(self):
 def get_bprop_all_gather(self):
     """Generate bprop for AllGather"""
     fusion = self.get_attr_dict()["fusion"]
-    self.group = self.get_attr_dict()["group"]
     reduce_scatter = ReduceScatter(ReduceOp.SUM, self.group).add_prim_attr("fusion", fusion)
     if hasattr(self, "instance_name") and self.instance_name:
         instance_name = "grad_" + self.instance_name
         reduce_scatter.set_prim_instance_name(instance_name)
     mean_flag = self.get_attr_dict()["mean_flag"]
-    self.rank_size = self.get_attr_dict()["rank_size"]
     if self.rank_size == 0:
         raise ValueError(f"The 'rank_size' can not be zero, but got {self.rank_size}.")
     scale = 1.0 / self.rank_size
@@ -377,66 +378,6 @@ def get_bprop_reduce_scatter(self):
     return bprop
 
 
-@bprop_getters.register(Reduce)
-def get_bprop_reduce(self):
-    """Generate bprop for Reduce"""
-    dest_rank = self.get_attr_dict()["dest_rank"]
-    group = self.get_attr_dict()["group"]
-    reduce_grad = Broadcast(dest_rank, group)
-    if hasattr(self, "instance_name") and self.instance_name:
-        instance_name = "grad" + self.instance_name
-        reduce_grad.set_prim_instance_name(instance_name)
-
-    def bprop(x, out, dout):
-        dx = reduce_grad((dout,))
-        return (dx[0],)
-
-    return bprop
-
-
-@bprop_getters.register(CollectiveGather)
-def get_bprop_collective_gather(self):
-    """Generate bprop for CollectiveGather"""
-    group = self.get_attr_dict()["group"]
-    dest_rank = self.get_attr_dict()["dest_rank"]
-    collective_gather_grad = Broadcast(dest_rank, group)
-    rank = get_rank(group)
-    dev_num = self.rank_size
-    split = P.Split(output_num=dev_num)
-    if hasattr(self, "instance_name") and self.instance_name:
-        instance_name = "grad" + self.instance_name
-        collective_gather_grad.set_prim_instance_name(instance_name)
-
-    def bprop(x, out, dout):
-        grad = collective_gather_grad((dout,))
-        dx = split(grad[0])[rank]
-        return (dx,)
-
-    return bprop
-
-
-@bprop_getters.register(CollectiveScatter)
-def get_bprop_collective_scatter(self):
-    """Generate bprop for CollectiveScatter"""
-    group = self.get_attr_dict()["group"]
-    dest_rank = self.get_attr_dict()["src_rank"]
-    rank = get_rank(group)
-    collective_scatter_grad = CollectiveGather(dest_rank, group)
-    if hasattr(self, "instance_name") and self.instance_name:
-        instance_name = "grad" + self.instance_name
-        collective_scatter_grad.set_prim_instance_name(instance_name)
-
-    def bprop(x, out, dout):
-        dx_out = collective_scatter_grad(dout)
-        if rank == dest_rank:
-            dx = dx_out
-        else:
-            dx = F.depend(F.zeros_like(x), dx_out)
-        return (dx,)
-
-    return bprop
-
-
 @bprop_getters.register(_AllSwap)
 def get_bprop_allswap(self):
     """Generate bprop for _AllSwap."""

mindspore/ops/_vmap/vmap_array_ops.py

@@ -2113,33 +2113,6 @@ def get_split_vmap_rule(prim, axis_size):
 
     return vmap_rule
 
-@vmap_rules_getters.register(P.SearchSorted)
-def get_searchsorted_vmap_rule(prim, axis_size):
-    """VmapRule for `SearchSorted`."""
-    def vmap_rule(sequence_bdim, values_bdim, sorter_bdim, dtype_bdim, right_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, sequence_bdim, values_bdim,
-                                                      sorter_bdim, dtype_bdim, right_bdim)
-        if is_all_none:
-            return result
-
-        sequence, sequence_dim = sequence_bdim
-        values, values_dim = values_bdim
-        sorter, sorter_dim = sorter_bdim
-
-        sequence = _bdim_at_front(sequence, sequence_dim, axis_size)
-        values = _bdim_at_front(values, values_dim, axis_size)
-        if sorter is not None and sorter_dim is not None:
-            sorter = _bdim_at_front(sorter, sorter_dim, axis_size)
-
-        dtype, _ = dtype_bdim
-        right, _ = right_bdim
-
-        outputs = prim(sequence, values, sorter, dtype, right)
-
-        return outputs, 0
-
-    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)

mindspore/ops/_vmap/vmap_math_ops.py

@@ -63,6 +63,7 @@ def _broadcast_shape(nd, x_ndim, x_shape):
 @vmap_rules_getters.register(P.BitwiseAnd)
 @vmap_rules_getters.register(P.BitwiseOr)
 @vmap_rules_getters.register(P.BitwiseXor)
+@vmap_rules_getters.register(P.IsClose)
 @vmap_rules_getters.register(P.Xlogy)
 @vmap_rules_getters.register(P.ApproximateEqual)
 @vmap_rules_getters.register(P.TruncateDiv)
@@ -887,35 +888,6 @@ def get_logit_vmap_rule(prim_func, axis_size):
 
     return vmap_rule
 
-
-@vmap_rules_getters.register(P.IsClose)
-def get_isclose_vmap_rule(prim, axis_size):
-    """VmapRule for `IsClose` operation"""
-
-    def vmap_rule(x_bdim, y_bdim, rtol_bdim, atol_bdim, equal_nan_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, x_bdim, x_bdim, rtol_bdim, atol_bdim, equal_nan_bdim)
-        if is_all_none:
-            return result
-
-        x, x_dim = x_bdim
-        y, y_dim = y_bdim
-        rtol, _ = rtol_bdim
-        atol, _ = atol_bdim
-        equal_nan, _ = equal_nan_bdim
-
-        if x_dim == y_dim:
-            out = prim(x, y, rtol, atol, equal_nan)
-            return out, x_dim
-        if y_dim is None:
-            y = _broadcast_by_axis(y, x_dim, axis_size)
-        else:
-            y = mnp.moveaxis(y, y_dim, x_dim)
-
-        out = prim(x, y, rtol, atol, equal_nan)
-        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

@@ -31,7 +31,6 @@ from mindspore.ops._vmap.vmap_base import vmap_rules_getters, vmap_general_prepr
 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
 
 
 @vmap_rules_getters.register(P.ApplyAdaMax)
@@ -299,19 +298,25 @@ def get_bce_with_logits_loss_vamp_rule(prim, axis_size):
 
     if isinstance(prim, str):
         prim = Primitive(prim)
+        prim_reduction = 'none'
+    else:
+        prim_reduction = prim.reduction
     prim_name = prim.name
     bce_logits_with_loss_op = NN.BCEWithLogitsLoss('none')
+    if prim_reduction == 'mean':
+        reduce_op = P.ReduceMean()
+    elif prim_reduction == "sum":
+        reduce_op = P.ReduceSum()
 
-    def vmap_rule(logits_bdim, label_bdim, weight_bdim, pos_weight_bdim, reduction_bdim):
-        is_all_none, result = vmap_general_preprocess(prim, logits_bdim, label_bdim, weight_bdim, pos_weight_bdim,
-                                                      reduction_bdim)
+    def vmap_rule(logits_bdim, label_bdim, weight_bdim, pos_weight_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, logits_bdim, label_bdim,
+                                                      weight_bdim, pos_weight_bdim)
         if is_all_none:
             return result
         logits, logits_dim = logits_bdim
         label, label_dim = label_bdim
         weight, weight_dim = weight_bdim
         pos_weight, pos_weight_dim = pos_weight_bdim
-        prim_reduction, _ = reduction_bdim
         logits_rank = F.rank(logits)
         label_rank = F.rank(label)
         weight_rank = F.rank(weight)
@@ -327,14 +332,11 @@ def get_bce_with_logits_loss_vamp_rule(prim, axis_size):
         shape = F.shape(logits)
         shape_ok = shape == F.shape(label) and shape == F.shape(weight) and shape == F.shape(pos_weight)
         if logits_dim_ok and shape_ok:
-            if prim_reduction == handler.str_to_enum("BCEWithLogitsLoss", "reduction", 'none'):
-                output = prim(logits, label, weight, pos_weight, prim_reduction)
-            elif prim_reduction == handler.str_to_enum("BCEWithLogitsLoss", "reduction", 'mean'):
+            if prim_reduction == 'none':
+                output = prim(logits, label, weight, pos_weight)
+            elif prim_reduction in ('mean', 'sum'):
                 out = bce_logits_with_loss_op(logits, label, weight, pos_weight)
-                output = P.ReduceMean()(out, reduce_indexes)
-            elif prim_reduction == handler.str_to_enum("BCEWithLogitsLoss", "reduction", 'sum'):
-                out = bce_logits_with_loss_op(logits, label, weight, pos_weight)
-                output = P.ReduceSum()(out, reduce_indexes)
+                output = reduce_op(out, reduce_indexes)
             else:
                 raise RuntimeError("For {} vmap, the attribute of reduction must in "
                                    "('none', 'mean', 'sum'), but got {}."
@@ -350,14 +352,11 @@ def get_bce_with_logits_loss_vamp_rule(prim, axis_size):
         pos_weight_shape = F.shape(pos_weight)
         weight = _handle_broadcasting(weight, weight_shape, logits_shape)
         pos_weight = _handle_broadcasting(pos_weight, pos_weight_shape, logits_shape)
-        if prim_reduction == handler.str_to_enum("BCEWithLogitsLoss", "reduction", 'none'):
-            output = prim(logits, label, weight, pos_weight, prim_reduction)
-        elif prim_reduction == handler.str_to_enum("BCEWithLogitsLoss", "reduction", 'mean'):
-            out = bce_logits_with_loss_op(logits, label, weight, pos_weight)
-            output = P.ReduceMean()(out, reduce_indexes)
-        elif prim_reduction == handler.str_to_enum("BCEWithLogitsLoss", "reduction", 'sum'):
+        if prim_reduction == 'none':
+            output = prim(logits, label, weight, pos_weight)
+        elif prim_reduction in ('mean', 'sum'):
             out = bce_logits_with_loss_op(logits, label, weight, pos_weight)
-            output = P.ReduceSum()(out, reduce_indexes)
+            output = reduce_op(out, reduce_indexes)
         else:
             raise RuntimeError("For {} vmap, the attribute of reduction must in "
                                "('none', 'mean', 'sum'), but got {}."