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

mindspore/nn/optim/optimizer.py

@@ -26,12 +26,13 @@ from mindspore.nn.cell import Cell
 from mindspore.nn.layer.container import CellList
 from mindspore.common.parameter import Parameter, ParameterTuple
 from mindspore.common.initializer import initializer
-from mindspore.common.tensor import Tensor, RowTensor
+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, _enable_distributed_mindrt
+from mindspore.parallel._ps_context import _is_ps_mode
 from mindspore.context import ParallelMode
 from mindspore import context
 from mindspore.nn.learning_rate_schedule import LearningRateSchedule
@@ -180,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"""
@@ -252,7 +254,7 @@ class Optimizer(Cell):
         chunk_tensors = Tensor._get_flattened_tensors(parameters)  # pylint: disable=W0212
         if not chunk_tensors:
             # Failed to get chunk tensors.
-            logger.warning("Parameters are not properly falttened, fallback to not flattened parameters.")
+            logger.warning("Parameters are not properly flattened, fallback to not flattened parameters.")
             return parameters
         # Convert chunk tensors to parameters.
         self._use_flattened_params = True
@@ -350,7 +352,7 @@ class Optimizer(Cell):
         """
         Whether use distributed optimizers.
         """
-        return _is_ps_mode() and _enable_distributed_mindrt()
+        return _is_ps_mode()
 
     def flatten_gradients(self, gradients):
         """
@@ -545,7 +547,7 @@ class Optimizer(Cell):
             parameters = self._parameters_base_check(group_param['params'], "group `params`")
             for index, param in enumerate(parameters):
                 if not isinstance(param, Parameter):
-                    raise TypeError(f"For 'Optimizer', the elemeter in group parameters must be Parameter type, "
+                    raise TypeError(f"For 'Optimizer', the element in group parameters must be Parameter type, "
                                     f"but got {type(param)} at index {index}.")
 
     def _parse_group_params(self, parameters, learning_rate):
@@ -717,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()))
@@ -837,7 +840,7 @@ def _tensor_apply_decay_with_sparse(weight_decay, if_apply, weight, gradient):
         indices = gradient.indices
         values = op_add((op_gather(weight, indices, 0) * F.cast(weight_decay, F.dtype(weight)), gradient.values))
         shape = gradient.dense_shape
-        return RowTensor(indices, values, shape)
+        return RowTensorInner(indices, values, shape)
     return gradient
 
 
@@ -863,7 +866,7 @@ def _tensor_apply_grad_centralization_with_sparse(if_apply, gradient):
             if grad_shape[1] % 16 != 0:
                 return gradient
             values = op_gc(gradient.values, axis)
-            return RowTensor(indices, values, shape)
+            return RowTensorInner(indices, values, shape)
     return gradient
 
 
@@ -903,7 +906,13 @@ def tensor_grad_scale_with_tensor(scale, grad):
 @_grad_scale.register("Tensor", "RowTensor")
 def tensor_grad_scale_with_sparse(scale, grad):
     """Get grad with scale."""
-    return RowTensor(grad.indices, grad.values * F.cast(scale, F.dtype(grad.values)), grad.dense_shape)
+    return RowTensorInner(grad.indices, grad.values * F.cast(scale, F.dtype(grad.values)), grad.dense_shape)
+
+
+@_grad_scale.register("Tensor", "MapTensor")
+def tensor_grad_scale_with_map_tensor(scale, grad):
+    """Get grad with scale."""
+    return grad
 
 
 @_indices_deduplicate.register("RowTensor")
@@ -915,7 +924,7 @@ def rowtensor_deduplicate_indices_slices(grad):
     unique_indices, index_position = P.Unique()(indices)
     summed_values = P.UnsortedSegmentSum()(values, index_position, P.TensorShape()(unique_indices)[0])
 
-    return RowTensor(unique_indices, summed_values, grad.dense_shape)
+    return RowTensorInner(unique_indices, summed_values, grad.dense_shape)
 
 
 @_indices_deduplicate.register("Tensor")

mindspore/nn/optim/proximal_ada_grad.py

@@ -18,8 +18,8 @@ from __future__ import absolute_import
 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 ms_function
-from mindspore._checkparam import Validator as validator
+from mindspore.common.api import jit
+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``
+        ``Ascend`` ``GPU``
 
     Examples:
         >>> import mindspore as ms
@@ -197,7 +197,7 @@ class ProximalAdagrad(Optimizer):
         self.opt = P.ApplyProximalAdagrad(use_locking=use_locking)
         self.sparse_opt = P.SparseApplyProximalAdagrad(use_locking=use_locking)
 
-    @ms_function
+    @jit
     def construct(self, grads):
         params = self._parameters
         accum = self.accum
@@ -231,7 +231,8 @@ class ProximalAdagrad(Optimizer):
                              "but got {}.".format(value))
 
         if value == 'CPU':
-            self.sparse_opt = P.FusedSparseProximalAdagrad(self.use_locking).add_prim_attr("primitive_target", "CPU")
+            self.sparse_opt = P.FusedSparseProximalAdagrad(self.use_locking)
+            self.sparse_opt.set_device("CPU")
         else:
             self.sparse_opt = P.SparseApplyProximalAdagrad(self.use_locking)
 

mindspore/nn/optim/rmsprop.py

@@ -16,8 +16,8 @@
 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.common.api import ms_function
+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
 
@@ -225,7 +225,7 @@ class RMSProp(Optimizer):
         self.epsilon = epsilon
         self.decay = decay
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         params = self._parameters
         gradients = self.flatten_gradients(gradients)

mindspore/nn/optim/rprop.py

@@ -18,9 +18,8 @@ from __future__ import absolute_import
 from mindspore import ops
 from mindspore.ops import operations as P
 import mindspore.common.dtype as mstype
-from mindspore.common.api import ms_function
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore.common.api import jit
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 
@@ -30,19 +29,19 @@ class Rprop(Optimizer):
     Implements Resilient backpropagation.
 
     Further information about this implementation can be found at  `A Direct Adaptive Method for Faster Backpropagation
-    Learning: The RPROP Algorithm <https://ieeexplore.ieee.org/document/298623>`_ .
+    Learning: The RPROP Algorithm <https://ieeexplore.ieee.org/document/298623>`_.
 
     The updating formulas are as follows:
 
     .. 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
@@ -50,8 +49,12 @@ class Rprop(Optimizer):
 
     Note:
         If parameters are not grouped, the `weight_decay` in optimizer will be applied on the parameters without 'beta'
-        or 'gamma' in their names. Users can group parameters to change the strategy of decaying weight. When parameters
-        are grouped, each group can set `weight_decay`. If not, the `weight_decay` in optimizer will be applied.
+        or 'gamma' in their names.
+
+        Users can group parameters to change the strategy of decaying weight.
+
+        When parameters are grouped, each group can set `weight_decay`. If not, the `weight_decay` in optimizer will be
+        applied.
 
     Args:
         params (Union[list[Parameter], list[dict]]): Must be list of `Parameter` or list of `dict`. When the
@@ -80,7 +83,7 @@ class Rprop(Optimizer):
               If `order_params` in the keys, other keys will be ignored and the element of 'order_params' must be in
               one group of `params`.
 
-        learning_rate (Union[float, int, Tensor, Iterable, LearningRateSchedule]):
+        learning_rate (Union[float, int, Tensor, Iterable, LearningRateSchedule]): Learning_rate. Default: 0.1.
 
             - float: The fixed learning rate value. Must be equal to or greater than 0.
 
@@ -95,8 +98,9 @@ class Rprop(Optimizer):
               LearningRateSchedule with step as the input to get the learning rate of current step.
 
         etas (tuple[float, float]): The factor of multiplicative increasing or
-            descreasing(etaminus, etaplus).
+            descreasing(etaminus, etaplus). Default: (0.5, 1.2).
         step_sizes(tuple[float, float]): The allowed minimal and maximal step size(min_step_sizes, max_step_size).
+            Default: (1e-6, 50.).
         weight_decay (Union[float, int, Cell]): Weight decay (L2 penalty). Default: 0.0.
 
             - float: The fixed weight decay value. Must be equal to or greater than 0.
@@ -170,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]))
@@ -191,7 +195,7 @@ class Rprop(Optimizer):
         self.select = P.Select()
         self.ones_like = P.OnesLike()
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         gradients = self.flatten_gradients(gradients)
         gradients = self.decay_weight(gradients)

mindspore/nn/optim/sgd.py

@@ -18,17 +18,17 @@ from __future__ import absolute_import
 from mindspore.ops import functional as F, composite as C, operations as P
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
-from mindspore.common.api import ms_function
+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,13 +186,20 @@ 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')
         self.stat = self._parameters.clone(prefix="stat", init='ones')
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         params = self._parameters
         accum = self.accum
@@ -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

@@ -21,10 +21,11 @@ from mindspore.ops import functional as F, composite as C, operations as P
 from mindspore.common.initializer import initializer
 from mindspore.common.parameter import Parameter, ParameterTuple
 from mindspore.common.tensor import Tensor
+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
@@ -87,8 +88,8 @@ def _clip_grad(clip_type, clip_value, grad):
         return grad
     dt = F.dtype(grad)
     if clip_type == 0:
-        new_grad = C.clip_by_value(grad, F.cast(F.tuple_to_array((-clip_value,)), dt),
-                                   F.cast(F.tuple_to_array((clip_value,)), dt))
+        new_grad = ops.clip_by_value(grad, F.cast(F.tuple_to_array((-clip_value,)), dt),
+                                     F.cast(F.tuple_to_array((clip_value,)), dt))
     else:
         new_grad = nn.ClipByNorm()(grad, F.cast(F.tuple_to_array((clip_value,)), dt))
     return new_grad
@@ -210,7 +211,7 @@ def find_net_layertype_recur(net, layertype_map):
             if subcell.gamma.requires_grad:
                 layertype_map.append(BatchNorm)
         elif isinstance(subcell, (nn.Conv2d, nn.Dense, nn.Embedding, nn.Conv2dTranspose, nn.Conv1d, nn.Conv1dTranspose,
-                                  nn.BatchNorm1d, nn.GroupNorm, nn.GlobalBatchNorm)):
+                                  nn.BatchNorm1d, nn.GroupNorm)):
             if isinstance(subcell, (nn.Dense, nn.Conv2d)):
                 get_layer_type_for_dense_and_conv(subcell, prefix, layertype_map)
             else:
@@ -253,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::
@@ -270,11 +266,11 @@ def thor(net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0
             \otimes\left(G_{i}^{(k)}+\lambda I\right)^{-1}\right) \nabla_{w_{i}} J^{(k)}
         \end{array}
 
-    :math:`a_{i-1}` represents the input of i-th layer,and which is the activations of previous layer,
-    :math:`D_{s_i}` represents the derivative of the loss function of the output of the i-th layer,
-    :math:`I` represents the identity matrix,
-    :math:`\lambda` represents :math:`damping`, :math:`g_i` represents gradients of the i-th layer,
-    :math:`\otimes` represents Kronecker product, :math:`\gamma` represents 'learning rate'
+    :math:`a_{i-1}` represents the input of i-th layer,and which is the activations of previous layer.
+    :math:`D_{s_i}` represents the derivative of the loss function of the output of the i-th layer.
+    :math:`I` represents the identity matrix.
+    :math:`\lambda` represents :math:`damping`, :math:`g_i` represents gradients of the i-th layer.
+    :math:`\otimes` represents Kronecker product, :math:`\gamma` represents 'learning rate'.
 
     Note:
         When a parameter group is separated, 'weight_decay' of each group is applied to the corresponding parameter.
@@ -313,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`.
@@ -337,6 +333,12 @@ def thor(net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0
         ``Ascend`` ``GPU``
 
     Examples:
+        .. 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/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
         >>> from mindspore import nn
@@ -966,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

@@ -16,9 +16,10 @@
 from mindspore import context
 from mindspore.nn.cell import Cell
 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
@@ -27,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.
@@ -96,8 +97,8 @@ class Bijector(Cell):
         self.dtype_base = P.DType()
         self.shape_base = P.Shape()
         self.fill_base = P.Fill()
-        self.sametypeshape_base = P.SameTypeShape()
-        self.issubclass_base = P.IsSubClass()
+        self.sametypeshape_base = inner.SameTypeShape()
+        self.issubclass_base = inner.IsSubClass()
 
     @property
     def name(self):

mindspore/nn/probability/bijector/invert.py

@@ -13,13 +13,15 @@
 # limitations under the License.
 # ============================================================================
 """Invert Bijector"""
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from .bijector import Bijector
 
 
 class Invert(Bijector):
     r"""
-    Invert Bijector. Compute the inverse function of the input bijector.
+    Invert Bijector. Compute the inverse function of the input bijector. If the function of the forward mapping,
+    namely the input of `bijector` below, is :math:`Y = g(X)`,
+    then the function of corresponding inverse mapping Bijector is :math:`Y = h(X) = g^{-1}(X)`.
 
     Args:
         bijector (Bijector): Base 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/__init__.py

@@ -31,6 +31,9 @@ from .log_normal import LogNormal
 from .normal import Normal
 from .poisson import Poisson
 from .uniform import Uniform
+from .half_normal import HalfNormal
+from .laplace import Laplace
+from .student_t import StudentT
 
 __all__ = ['Distribution',
            'TransformedDistribution',
@@ -47,4 +50,7 @@ __all__ = ['Distribution',
            'Normal',
            'Poisson',
            'Uniform',
+           'HalfNormal',
+           'Laplace',
+           'StudentT',
            ]

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

@@ -15,6 +15,7 @@
 """Utility functions to help distribution class."""
 import numpy as np
 from mindspore.ops import operations as P
+from mindspore.ops.operations import _inner_ops as inner
 from mindspore.common import dtype as mstype
 
 
@@ -26,7 +27,7 @@ def exp_generic(input_x):
     exp = P.Exp()
     cast = P.Cast()
     dtype = P.DType()
-    checktype = P.IsSubClass()
+    checktype = inner.IsSubClass()
 
     if not checktype(dtype(input_x), mstype.float_):
         input_x = cast(input_x, mstype.float32)
@@ -48,7 +49,7 @@ def log_generic(input_x):
     dtype = P.DType()
     shape = P.Shape()
     select = P.Select()
-    checktype = P.IsSubClass()
+    checktype = inner.IsSubClass()
 
     if not checktype(dtype(input_x), mstype.float_):
         input_x = cast(input_x, mstype.float32)