mindspore 2.0.0a0__cp38-cp38-win_amd64.whl → 2.0.0rc1__cp38-cp38-win_amd64.whl

mindspore/nn/optim/lars.py

@@ -18,7 +18,7 @@ from __future__ import absolute_import
 from mindspore.ops import operations as P
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.common import Tensor, Parameter, dtype as mstype
 from mindspore.common.api import jit
 from mindspore.nn.optim.optimizer import _grad_scale, Optimizer
@@ -83,7 +83,7 @@ class LARS(Optimizer):
         \end{array}
 
     :math:`w` represents the network parameters, :math:`g` represents `gradients`,
-    :math:`t` represents the current step, :math:`\delta` represents `weight_decay` in `optimizer`,
+    :math:`t` represents the current step, :math:`\lambda` represents `weight_decay` in `optimizer`,
     :math:`\gamma` represents `learning_rate` in `optimizer`, :math:`\eta` represents `coefficient`.
 
     Args:

mindspore/nn/optim/lazyadam.py

@@ -23,8 +23,7 @@ from mindspore.ops import composite as C
 from mindspore.ops import functional as F
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 from mindspore.nn.optim._dist_optimizer_registry import _register_dist_optimizer
@@ -86,6 +85,46 @@ def _run_opt_with_sparse_dist(opt, sparse_opt, push, pull, use_locking, use_nest
     return success
 
 
+@_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "MapTensor", "MapTensor", "MapTensor", "MapTensor",
+                         "Bool", "Bool", "Function", "Bool", "Function", "Bool")
+def _run_map_tensor_opt_with_sparse_dist(opt, sparse_opt, push, pull, use_locking, use_nesterov, target, beta1_power,
+                                         beta2_power, beta1, beta2, eps, lr, gradient, params, m, v,
+                                         ps_parameter, cache_enable, distributed_opt, use_flag, distributed_sparse_opt,
+                                         use_sparse_flag):
+    """Apply sparse lazy adam optimizer to the weight parameter when the gradient is sparse."""
+    success = True
+    indices, values = gradient.get_data()
+    if use_sparse_flag:
+        # PS Mode.
+        success = F.depend(success, distributed_sparse_opt(params, m, v, beta1_power, beta2_power, lr, beta1, beta2,
+                                                           eps, values, indices))
+    else:
+        # PS Cache mode.
+        op_sqrt = P.Sqrt()
+
+        m_slice = m.get(indices)
+        v_slice = v.get(indices)
+
+        next_m = m_slice * beta1 + values * (1 - beta1)
+        next_v = v_slice * beta2 + values * values * (1 - beta2)
+
+        lr_t = lr * op_sqrt(1 - beta2_power) / (1 - beta1_power)
+
+        if use_nesterov:
+            m_temp = beta1 * next_m + values * (1 - beta1)
+            param_update = m_temp / (op_sqrt(next_v) + eps)
+        else:
+            param_update = next_m / (op_sqrt(next_v) + eps)
+
+        params_need_update = params.get(indices)
+        params.put(indices, params_need_update - lr_t * param_update)
+        m.put(indices, next_m)
+        v.put(indices, next_v)
+
+    return success
+
+
 @_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
                          "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool",
                          "Function", "Bool", "Function", "Bool")
@@ -212,8 +251,8 @@ def _check_param_value(beta1, beta2, eps, weight_decay, prim_name):
     validator.check_value_type("beta2", beta2, [float], prim_name)
     validator.check_value_type("eps", eps, [float], prim_name)
     validator.check_value_type("weight_dacay", weight_decay, [float], prim_name)
-    validator.check_float_range(beta1, 0.0, 1.0, Rel.INC_NEITHER, "beta1", prim_name)
-    validator.check_float_range(beta2, 0.0, 1.0, Rel.INC_NEITHER, "beta2", prim_name)
+    validator.check_float_range(beta1, 0.0, 1.0, validator.INC_NEITHER, "beta1", prim_name)
+    validator.check_float_range(beta2, 0.0, 1.0, validator.INC_NEITHER, "beta2", prim_name)
     validator.check_positive_float(eps, "eps", prim_name)
     validator.check_non_negative_float(weight_decay, "weight_decay", prim_name)
 

mindspore/nn/optim/momentum.py

@@ -20,7 +20,7 @@ from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 from mindspore.common.api import jit
 import mindspore.common.dtype as mstype
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 from mindspore.nn.optim._dist_optimizer_registry import _register_dist_optimizer
@@ -69,19 +69,20 @@ class Momentum(Optimizer):
     learning <https://dl.acm.org/doi/10.5555/3042817.3043064>`_ for more details.
 
     .. math::
-            v_{t+1} = v_{t} \ast u + grad
+        v_{t+1} = v_{t} \ast u + grad
 
     If use_nesterov is True:
 
     .. math::
-            p_{t+1} =  p_{t} - (grad \ast lr + v_{t+1} \ast u \ast lr)
+        p_{t+1} =  p_{t} - (grad \ast lr + v_{t+1} \ast u \ast lr)
 
     If use_nesterov is False:
 
     .. math::
-            p_{t+1} = p_{t} - lr \ast v_{t+1}
+        p_{t+1} = p_{t} - lr \ast v_{t+1}
 
-    Here: where grad, lr, p, v and u denote the gradients, learning_rate, params, moments, and momentum respectively.
+    Here: where :math:`grad`, :math:`lr`, :math:`p`, :math:`v` and :math:`u` denote the gradients,
+    learning_rate, params, moments, and momentum respectively.
 
     Note:
         If parameters are not grouped, the `weight_decay` in optimizer will be applied on the network parameters without

mindspore/nn/optim/optimizer.py

@@ -29,7 +29,7 @@ from mindspore.common.initializer import initializer
 from mindspore.common import Tensor
 from mindspore.common.sparse_tensor import RowTensorInner
 import mindspore.common.dtype as mstype
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore import log as logger
 from mindspore.parallel._utils import _get_global_rank, _get_device_num, _get_parallel_mode
 from mindspore.parallel._ps_context import _is_ps_mode
@@ -181,6 +181,7 @@ class Optimizer(Cell):
             self._init_group_params(parameters, learning_rate, weight_decay, self.grad_centralization)
 
         self._init_opt_attrs(learning_rate, parameters, weight_decay)
+        self.add_flags(skip_auto_parallel_compile=True)
 
     def _init_opt_attrs(self, learning_rate, parameters, weight_decay):
         """initialize optimizer attributions"""
@@ -718,6 +719,7 @@ class Optimizer(Cell):
 
         Examples:
             >>> from mindspore import nn
+            >>> # net = LeNet5()
             >>> net = Net()
             >>> conv_params = list(filter(lambda x: 'conv' in x.name, net.trainable_params()))
             >>> no_conv_params = list(filter(lambda x: 'conv' not in x.name, net.trainable_params()))

mindspore/nn/optim/proximal_ada_grad.py

@@ -19,7 +19,7 @@ from mindspore.ops import functional as F, composite as C, operations as P
 from mindspore.common import Tensor
 import mindspore.common.dtype as mstype
 from mindspore.common.api import jit
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 
@@ -158,7 +158,7 @@ class ProximalAdagrad(Optimizer):
         ValueError: If `accum`, `l1`, `l2` or `weight_decay` is less than 0.
 
     Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
+        ``Ascend`` ``GPU``
 
     Examples:
         >>> import mindspore as ms

mindspore/nn/optim/rmsprop.py

@@ -16,7 +16,7 @@
 from __future__ import absolute_import
 
 from mindspore.ops import functional as F, composite as C, operations as P
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.common.api import jit
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register

mindspore/nn/optim/rprop.py

@@ -19,8 +19,7 @@ from mindspore import ops
 from mindspore.ops import operations as P
 import mindspore.common.dtype as mstype
 from mindspore.common.api import jit
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 
@@ -37,12 +36,12 @@ class Rprop(Optimizer):
     .. math::
         \begin{gather*}
             &\hspace{-10mm}  \textbf{if} \:   g_{t-1} g_t  > 0                                     \\
-            &\hspace{25mm}  \Delta_t \leftarrow \mathrm{min}(\Delta_{t-1} \eta_{+}, \Delta_{max}) \\
-            &\hspace{0mm}  \textbf{else if}  \:  g_{t-1} g_t < 0                                 \\
-            &\hspace{25mm}  \Delta_t \leftarrow \mathrm{max}(\Delta_{t-1} \eta_{-}, \Delta_{min}) \\
-            &\hspace{-25mm}  \textbf{else}  \:                                                      \\
-            &\hspace{-5mm}  \Delta_t \leftarrow \Delta_{t-1}                                      \\
-            &\hspace{15mm} w_{t} \leftarrow w_{t-1}- \Delta_{t} \mathrm{sign}(g_t)                \\
+            &\hspace{25mm}  \Delta_t \leftarrow \mathrm{min}(\Delta_{t-1} \eta_{+}, \Delta_{max})  \\
+            &\hspace{0mm}  \textbf{else if}  \:  g_{t-1} g_t < 0                                   \\
+            &\hspace{25mm}  \Delta_t \leftarrow \mathrm{max}(\Delta_{t-1} \eta_{-}, \Delta_{min})  \\
+            &\hspace{-25mm}  \textbf{else}  \:                                                     \\
+            &\hspace{-5mm}  \Delta_t \leftarrow \Delta_{t-1}                                       \\
+            &\hspace{15mm} w_{t} \leftarrow w_{t-1}- \Delta_{t} \mathrm{sign}(g_t)                 \\
         \end{gather*}
 
     :math:`\Delta_{min/max}` represents the min/max step size, :math:`\eta_{+/-}` represents the factors of
@@ -175,7 +174,7 @@ class Rprop(Optimizer):
             raise ValueError("For Rprop, maximal step size should not be less than minimal step size, "
                              "but got {} > {}.".format(step_sizes[0], step_sizes[1]))
 
-        validator.check_float_range(etas[0], 0.0, 1.0, Rel.INC_NEITHER, "etaminus", self.cls_name)
+        validator.check_float_range(etas[0], 0.0, 1.0, validator.INC_NEITHER, "etaminus", self.cls_name)
         validator.check_value_type("etaplus", etas[1], [float], self.cls_name)
         if etas[1] <= 1.0:
             raise ValueError("For Rprop, etaplus must be greater than 1.0, but got etaplus {}.".format(etas[1]))

mindspore/nn/optim/sgd.py

@@ -20,15 +20,15 @@ from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 from mindspore.common.api import jit
 import mindspore.common.dtype as mstype
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 
 _sgd_opt = C.MultitypeFuncGraph("sgd_opt")
 
 
-@_sgd_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
-def _tensor_run_opt_ext(opt, momentum, learning_rate, gradient, weight, accum, stat):
+@_sgd_opt.register("Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Function")
+def _tensor_run_opt_ext(momentum, learning_rate, gradient, weight, accum, stat, opt):
     """Apply sgd optimizer to the weight parameter using Tensor."""
     success = True
     success = F.depend(success, opt(weight, gradient, learning_rate, accum, momentum, stat))
@@ -76,7 +76,9 @@ class SGD(Optimizer):
             - lr: Optional. If "lr" in the keys, the value of corresponding learning rate will be used.
               If not, the `learning_rate` in optimizer will be used. Fixed and dynamic learning rate are supported.
 
-            - weight_decay: Using different `weight_decay` by grouping parameters is currently not supported.
+            - weight_decay: Optional. If "weight_decay" in the keys, the value of corresponding weight decay
+              will be used. If not, the `weight_decay` in the optimizer will be used. It should be noted that weight
+              decay must be float, dynamic weight decay is currently not supported.
 
             - grad_centralization: Optional. Must be Boolean. If "grad_centralization" is in the keys, the set value
               will be used. If not, the `grad_centralization` is False by default. This configuration only works on the
@@ -164,7 +166,7 @@ class SGD(Optimizer):
 
         if isinstance(momentum, float) and momentum < 0.0:
             raise ValueError("For 'SGD', the argument 'momentum' must be at least 0.0, "
-                             "but got {}".format(momentum))
+                             "but got {}.".format(momentum))
 
         if isinstance(dampening, int):
             dampening = float(dampening)
@@ -177,9 +179,6 @@ class SGD(Optimizer):
                              "but got 'dampening' {}".format(dampening))
         self.dampening = dampening
 
-        if isinstance(weight_decay, int):
-            weight_decay = float(weight_decay)
-
         validator.check_value_type("nesterov", nesterov, [bool], self.cls_name)
 
         if nesterov and (momentum <= 0.0 or dampening != 0.0):
@@ -187,7 +186,14 @@ class SGD(Optimizer):
                              "equal to 0.0, but got 'momentum' {}, 'dampening' {}".format(momentum, dampening))
         self.nesterov = nesterov
 
-        self.opt = P.SGD(dampening, weight_decay, nesterov)
+        if self.dynamic_weight_decay:
+            raise TypeError("For 'SGD', dynamic weight decay is currently not supported, the argument 'weight_decay' "
+                            "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, wd, nesterov) for wd in self.group_weight_decay)
+        else:
+            self.opt = tuple([P.SGD(dampening, float(weight_decay), nesterov)] * len(self._parameters))
 
         self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")
         self.accum = self._parameters.clone(prefix="accum", init='zeros')
@@ -203,9 +209,9 @@ class SGD(Optimizer):
         gradients = self.scale_grad(gradients)
         lr = self.get_lr()
         if self.is_group_lr:
-            success = self.hyper_map_reverse(F.partial(_sgd_opt, self.opt, self.momentum),
-                                             lr, gradients, params, accum, stat)
+            success = self.hyper_map_reverse(F.partial(_sgd_opt, self.momentum),
+                                             lr, gradients, params, accum, stat, self.opt)
         else:
-            success = self.hyper_map_reverse(F.partial(_sgd_opt, self.opt, self.momentum, lr),
-                                             gradients, params, accum, stat)
+            success = self.hyper_map_reverse(F.partial(_sgd_opt, self.momentum, lr),
+                                             gradients, params, accum, stat, self.opt)
         return success

mindspore/nn/optim/thor.py

@@ -25,7 +25,7 @@ import mindspore.ops as ops
 import mindspore.nn as nn
 import mindspore.common.dtype as mstype
 import mindspore.log as logger
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.parallel._utils import _get_device_num, _get_gradients_mean
 from mindspore import context
@@ -254,11 +254,6 @@ def thor(net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0
     r"""
     Updates gradients by second-order algorithm--THOR.
 
-    Trace-based Hardware-driven layer-ORiented Natural Gradient Descent Computation (THOR) algorithm is proposed in:
-
-    `THOR: Trace-based Hardware-driven layer-ORiented Natural Gradient Descent Computation
-    <https://www.aaai.org/AAAI21Papers/AAAI-6611.ChenM.pdf>`_
-
     The updating formulas are as follows,
 
     .. math::
@@ -314,9 +309,9 @@ def thor(net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0
 
         enable_clip_grad (bool): Whether to clip the gradients. Default: False
 
-        frequency(int): The update interval of A/G and $A^{-1}/G^{-1}$. When frequency equals N (N is greater than 1),
-            A/G and $A^{-1}/G^{-1}$ will be updated  every N steps, and other steps will use the stale A/G and
-            $A^{-1}/G^{-1}$ to update weights. Default: 100.
+        frequency(int): The update interval of A/G and :math:`A^{-1}/G^{-1}`. When frequency equals N
+            (N is greater than 1), A/G and :math:`A^{-1}/G^{-1}` will be updated every N steps,
+            and other steps will use the stale A/G and :math:`A^{-1}/G^{-1}` to update weights. Default: 100.
 
     Inputs:
         - **gradients** (tuple[Tensor]) - The gradients of `params`, the shape is the same as `params`.
@@ -341,8 +336,8 @@ def thor(net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0
         .. note::
             Before running the following example, you need to customize the network Net and
             dataset preparation function create_dataset. Refer to
-            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0.0-alpha/beginner/model.html>`_
-            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0.0-alpha/beginner/dataset.html>`_ .
+            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0/beginner/model.html>`_
+            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0/beginner/dataset.html>`_ .
 
         >>> import mindspore as ms
         >>> from mindspore.nn import thor
@@ -973,15 +968,15 @@ class ThorAscend(Optimizer):
             matrix_g_combine_shape = self.shape(matrix_g_inv)
             if matrix_a_inv_shape[0] == 2048 and matrix_g_combine_shape[0] == 1001:
                 matrix_a_inv = self.reshape(matrix_a_inv,
-                                            (matrix_a_inv_shape[0] / 16, 16,
-                                             matrix_a_inv_shape[0] / 16, 16))
+                                            (matrix_a_inv_shape[0] // 16, 16,
+                                             matrix_a_inv_shape[0] // 16, 16))
                 matrix_a_inv = self.transpose(matrix_a_inv, (2, 0, 1, 3))
                 matrix_g_inv = P.Pad(((0, 7), (0, 7)))(matrix_g_inv)
 
                 matrix_g_inv_shape = self.shape(matrix_g_inv)
                 matrix_g_inv = self.reshape(matrix_g_inv,
-                                            (matrix_g_inv_shape[0] / 16, 16,
-                                             matrix_g_inv_shape[0] / 16, 16))
+                                            (matrix_g_inv_shape[0] // 16, 16,
+                                             matrix_g_inv_shape[0] // 16, 16))
                 matrix_g_inv = self.transpose(matrix_g_inv, (2, 0, 1, 3))
 
         matrix_a_allreduce = matrix_a_allreduce + (matrix_a_inv,)

mindspore/nn/probability/__init__.py

@@ -20,5 +20,3 @@ The high-level components used to construct the probabilistic network.
 
 from . import bijector
 from . import distribution
-from . import infer
-from . import dpn

mindspore/nn/probability/bijector/bijector.py

@@ -19,7 +19,7 @@ from mindspore.ops import operations as P
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.common import dtype as mstype
 from mindspore.common.tensor import Tensor
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from ..distribution._utils.utils import CheckTensor, cast_to_tensor, raise_type_error
 from ..distribution import Distribution
 from ..distribution import TransformedDistribution
@@ -28,9 +28,9 @@ from ..distribution import TransformedDistribution
 class Bijector(Cell):
     """
     Bijecotr class. A bijector perform a mapping from one distribution to the other via some function.
-    If X is a random variable following the original distribution,
-    and g(x) is the mapping function,
-    then Y = g(X) is the random variable following the transformed distribution.
+    If :math:`X` is a random variable following the original distribution,
+    and :math:`g(x)` is the mapping function,
+    then :math:`Y = g(X)` is the random variable following the transformed distribution.
 
     Args:
         is_constant_jacobian (bool): Whether the Bijector has constant derivative. Default: False.

mindspore/nn/probability/bijector/invert.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """Invert Bijector"""
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from .bijector import Bijector
 
 

mindspore/nn/probability/bijector/softplus.py

@@ -41,7 +41,7 @@ class Softplus(Bijector):
         TypeError: When the dtype of the sharpness is not float.
 
     Supported Platforms:
-        ``Ascend`` ``GPU``
+        ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore
@@ -51,7 +51,7 @@ class Softplus(Bijector):
         >>>
         >>> # To initialize a Softplus bijector of sharpness 2.0.
         >>> softplus = msb.Softplus(2.0)
-        >>> # To use a ScalarAffine bijector in a network.
+        >>> # To use a Softplus bijector in a network.
         >>> value = Tensor([1, 2, 3], dtype=mindspore.float32)
         >>> ans1 = softplus.forward(value)
         >>> print(ans1.shape)

mindspore/nn/probability/bnn_layers/dense_variational.py

@@ -14,7 +14,7 @@
 # ============================================================================
 """dense_variational"""
 from mindspore.ops import operations as P
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from ...cell import Cell
 from ...layer.activation import get_activation
 from ..distribution.normal import Normal

mindspore/nn/probability/bnn_layers/layer_distribution.py

@@ -29,7 +29,7 @@ class NormalPrior(Cell):
     To initialize a normal distribution of mean 0 and standard deviation 0.1.
 
     Args:
-        dtype (:class:`mindspore.dtype`): The argument is used to define the data type of the output tensor.
+        dtype (mindspore.dtype): The argument is used to define the data type of the output tensor.
             Default: mindspore.float32.
         mean (int, float): Mean of normal distribution. Default: 0.
         std (int, float): Standard deviation of normal distribution. Default: 0.1.
@@ -55,7 +55,7 @@ class NormalPosterior(Cell):
     Args:
         name (str): Name prepended to trainable parameter.
         shape (list, tuple): Shape of the mean and standard deviation.
-        dtype (:class:`mindspore.dtype`): The argument is used to define the data type of the output tensor.
+        dtype (mindspore.dtype): The argument is used to define the data type of the output tensor.
             Default: mindspore.float32.
         loc_mean (int, float): Mean of distribution to initialize trainable parameters. Default: 0.
         loc_std (int, float): Standard deviation of distribution to initialize trainable parameters. Default: 0.1.

mindspore/nn/probability/distribution/_utils/utils.py

@@ -15,12 +15,12 @@
 """Utility functions to help distribution class."""
 import numpy as np
 from mindspore import context
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.common.tensor import Tensor
 from mindspore.common.parameter import Parameter
 from mindspore.common import dtype as mstype
 from mindspore.ops import operations as P
-from mindspore.ops.primitive import constexpr, PrimitiveWithInfer, prim_attr_register
+from mindspore.ops.primitive import constexpr, _primexpr, PrimitiveWithInfer, prim_attr_register
 import mindspore.ops as ops
 import mindspore.nn as nn
 
@@ -230,48 +230,42 @@ def probs_to_logits(probs, is_binary=False):
     return P.Log()(ps_clamped)
 
 
-@constexpr
+@constexpr(check=False)
 def raise_none_error(name):
     raise TypeError(f"the type {name} must be subclass of Tensor."
                     f" It can not be None since it is not specified during initialization.")
 
 
-@constexpr
-def raise_probs_logits_error():
-    raise TypeError(
-        "Either 'probs' or 'logits' must be specified, but not both.")
-
-
-@constexpr
+@_primexpr
 def raise_broadcast_error(shape_a, shape_b):
     raise ValueError(f"Shape {shape_a} and {shape_b} is not broadcastable.")
 
 
-@constexpr
+@constexpr(check=False)
 def raise_not_impl_error(name):
     raise ValueError(
         f"{name} function must be implemented for non-linear transformation")
 
 
-@constexpr
+@constexpr(check=False)
 def raise_not_implemented_util(func_name, obj, *args, **kwargs):
     raise NotImplementedError(
         f"{func_name} is not implemented for {obj} distribution.")
 
 
-@constexpr
+@constexpr(check=False)
 def raise_type_error(name, cur_type, required_type):
     raise TypeError(
         f"For {name} , the type must be or be subclass of {required_type}, but got {cur_type}")
 
 
-@constexpr
+@constexpr(check=False)
 def raise_not_defined(func_name, obj, *args, **kwargs):
     raise ValueError(
         f"{func_name} is undefined for {obj} distribution.")
 
 
-@constexpr
+@constexpr(check=False)
 def check_distribution_name(name, expected_name):
     if name is None:
         raise ValueError(

mindspore/nn/probability/distribution/bernoulli.py

@@ -16,16 +16,16 @@
 from mindspore.common import dtype as mstype
 from mindspore.ops import operations as P
 from mindspore.ops import composite as C
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from .distribution import Distribution
 from ._utils.utils import check_prob, check_distribution_name, clamp_probs
 from ._utils.custom_ops import exp_generic, log_generic
 
 
 class Bernoulli(Distribution):
-    """
+    r"""
     Bernoulli Distribution.
-    A Bernoulli Distribution is a discrete distribution with the range {0, 1}
+    A Bernoulli Distribution is a discrete distribution with the range :math:`\{0, 1\}`
     and the probability mass function as :math:`P(X = 0) = p, P(X = 1) = 1-p`.
 
     Args:

mindspore/nn/probability/distribution/beta.py

@@ -17,7 +17,7 @@ import numpy as np
 from mindspore.ops import operations as P
 from mindspore.ops import composite as C
 import mindspore.nn as nn
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
 from .distribution import Distribution
 from ._utils.utils import check_greater_zero, check_distribution_name

mindspore/nn/probability/distribution/categorical.py

@@ -19,20 +19,20 @@ from mindspore.ops import operations as P
 from mindspore.ops import composite as C
 from mindspore.ops.functional import stop_gradient
 from mindspore.ops.operations import _inner_ops as inner
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 import mindspore.ops as ops
 import mindspore.nn as nn
 from mindspore.common import dtype as mstype
 from .distribution import Distribution
 from ._utils.utils import check_prob, check_sum_equal_one, check_rank,\
-    check_distribution_name, raise_not_implemented_util
+    check_distribution_name
 from ._utils.custom_ops import exp_generic, log_generic, broadcast_to
 
 
 class Categorical(Distribution):
-    """
+    r"""
     Categorical distribution.
-    A Categorical Distribution is a discrete distribution with the range {1, 2, ..., k}
+    A Categorical Distribution is a discrete distribution with the range :math:`\{1, 2, ..., k\}`
     and the probability mass function as :math:`P(X = i) = p_i, i = 1, ..., k`.
 
     Args:
@@ -238,7 +238,7 @@ class Categorical(Distribution):
         """
         probs = self._check_param_type(probs)
         logits = self.log(probs)
-        return self.squeeze(-self.reduce_sum(logits * probs, -1))
+        return self.squeeze(P.Neg()(self.reduce_sum(logits * probs, -1)))
 
     def _kl_loss(self, dist, probs_b, probs=None):
         """
@@ -405,8 +405,6 @@ class Categorical(Distribution):
         Returns:
             Tensor, shape is shape(probs)[:-1] + sample_shape
         """
-        if self.device_target == 'Ascend':
-            raise_not_implemented_util('On d backend, sample', self.name)
         shape = self.checktuple(shape, 'shape')
         probs = self._check_param_type(probs)
         num_classes = self.shape(probs)[-1]

mindspore/nn/probability/distribution/cauchy.py

@@ -16,7 +16,7 @@
 import numpy as np
 from mindspore.ops import operations as P
 from mindspore.ops import composite as C
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
 from .distribution import Distribution
 from ._utils.utils import check_greater_zero, check_distribution_name, raise_not_defined
@@ -26,13 +26,13 @@ from ._utils.custom_ops import exp_generic, log_generic, log1p_generic
 class Cauchy(Distribution):
     r"""
     Cauchy distribution.
-    A Cauchy distributio is a continuous distribution with the range :math:`[0, 1]`
+    A Cauchy distributio is a continuous distribution with the range of all real numbers
     and the probability density function:
 
     .. math::
         f(x, a, b) = 1 / \pi b(1 - ((x - a)/b)^2),
 
-    where a and b are loc and scale parameter respectively.
+    where :math:`a, b` are loc and scale parameter respectively.
 
     Args:
         loc (int, float, list, numpy.ndarray, Tensor): The location of the Cauchy distribution. Default: None.

mindspore/nn/probability/distribution/distribution.py

@@ -18,7 +18,7 @@ from mindspore.ops import operations as P
 from mindspore.nn.cell import Cell
 from mindspore.ops.primitive import constexpr
 from mindspore.ops.operations import _inner_ops as inner
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from ._utils.utils import raise_none_error, cast_to_tensor, set_param_type, cast_type_for_device,\
     raise_not_implemented_util
 from ._utils.utils import CheckTuple, CheckTensor
@@ -102,7 +102,7 @@ class Distribution(Cell):
         self.device_target = context.get_context('device_target')
         self.checktuple = CheckTuple()
 
-        @constexpr
+        @constexpr(check=False)
         def _check_tensor(x, name):
             CheckTensor()(x, name)
             return x