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

mindspore/nn/probability/distribution/exponential.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
@@ -26,7 +26,7 @@ from ._utils.custom_ops import exp_generic, log_generic
 class Exponential(Distribution):
     r"""
     Exponential Distribution.
-    An Exponential distributio is a continuous distribution with the range :math:`[0, 1]`
+    An Exponential distributio is a continuous distribution with the range :math:`[0, \inf)`
     and the probability density function:
 
     .. math::

mindspore/nn/probability/distribution/gamma.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
@@ -27,14 +27,14 @@ from ._utils.custom_ops import log_generic
 class Gamma(Distribution):
     r"""
     Gamma distribution.
-    A Gamma distributio is a continuous distribution with the range :math:`[0, 1]`
+    A Gamma distributio is a continuous distribution with the range :math:`(0, \inf)`
     and the probability density function:
 
     .. math::
         f(x, \alpha, \beta) = \beta^\alpha / \Gamma(\alpha) x^{\alpha - 1} \exp(-\beta x).
 
     where :math:`G` is the Gamma function,
-    and :math:`\alpha, \beta` are the concentration and the rate of the distribution respectively.
+    and :math:`\alpha` and :math:`\beta` are the concentration and the rate of the distribution respectively.
 
     Args:
         concentration (int, float, list, numpy.ndarray, Tensor): The concentration,

mindspore/nn/probability/distribution/geometric.py

@@ -17,7 +17,7 @@ import numpy as np
 from mindspore.ops import operations as P
 from mindspore.ops.operations import _inner_ops as inner
 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_prob, check_distribution_name

mindspore/nn/probability/distribution/gumbel.py

@@ -15,7 +15,7 @@
 """Gumbel Distribution"""
 import numpy as np
 from mindspore.ops import operations as P
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
 import mindspore.nn as nn
 import mindspore.nn.probability.bijector as msb
@@ -28,13 +28,13 @@ from ._utils.custom_ops import exp_generic, log_generic
 class Gumbel(TransformedDistribution):
     r"""
     Gumbel distribution.
-    A Gumbel distributio is a continuous distribution with the range :math:`[0, 1]`
+    A Gumbel distributio is a continuous distribution with the range of all real numbers
     and the probability density function:
 
     .. math::
         f(x, a, b) = 1 / b \exp(\exp(-(x - a) / b) - x),
 
-    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 Gumbel distribution.

mindspore/nn/probability/distribution/half_normal.py

@@ -16,8 +16,9 @@
 from __future__ import absolute_import
 from __future__ import division
 import numpy as np
+from mindspore import ops
 from mindspore.ops import operations as P
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
 from mindspore.nn.probability.distribution import Distribution
 from mindspore.nn.probability.distribution._utils.utils import check_greater_zero
@@ -35,16 +36,19 @@ class HalfNormal(Distribution):
     where :math:`\mu, \sigma` are the mean and the standard deviation of the half normal distribution respectively.
 
     Args:
-        mean (int, float, list, numpy.ndarray, Tensor): The mean of the distribution. Default: None.
-        sd (int, float, list, numpy.ndarray, Tensor): The standard deviation of the distribution. Default: None.
-        seed (int): The seed used in sampling. The global seed is used if it is None. Default: None.
-        dtype (mindspore.dtype): The type of the event samples. Default: mstype.float32.
-        name (str): The name of the distribution. Default: 'HalfNormal'.
+        mean (Union[int, float, list, numpy.ndarray, Tensor], optional): The mean of the distribution.
+            If this arg is None, then the mean of the distribution will be passed in runtime. Default: None.
+        sd (Union[int, float, list, numpy.ndarray, Tensor], optional): The standard deviation of the distribution.
+            If this arg is None, then the sd of the distribution will be passed in runtime. Default: None.
+        seed (int, optional): The seed used in sampling. The global seed is used if it is None. Default: None.
+        dtype (mindspore.dtype, optional): The type of the event samples. Default: mstype.float32.
+        name (str, optional): The name of the distribution. Default: 'HalfNormal'.
 
     Note:
         - `sd` must be greater than zero.
-        - `dist_spec_args` are `mean` and `sd`.
         - `dtype` must be a float type because HalfNormal distributions are continuous.
+        - If the arg `mean` or `sd` is passed in runtime, then it will be used as the parameter value.
+          Otherwise, the value passed in the constructor will be used.
 
     Raises:
         ValueError: When sd <= 0.
@@ -104,18 +108,18 @@ class HalfNormal(Distribution):
 
         self.exp = P.Exp()
         self.cast = P.Cast()
-        self.const = np.sqrt(2. / np.pi)
+        self.const = ops.scalar_to_tensor(np.sqrt(2. / np.pi))
         self.sq = P.Square()
         self.type = dtype
 
     def _prob(self, value, mean=None, sd=None):
         r"""
-        Evaluate probability.
+        Evaluate probability of the value of the HalfNormal distribution.
 
         Args:
             value (Tensor): The value to be evaluated.
-            mean (Tensor): The mean of the distribution. Default: self._mean_value.
-            sd (Tensor): The standard deviation the distribution. Default: self._sd_value.
+            mean (Tensor, optional): The mean of the distribution. Default: self._mean_value.
+            sd (Tensor, optional): The standard deviation the distribution. Default: self._sd_value.
 
         .. math::
             P(x) = 1 / \sigma \sqrt{2\pi} \exp(-(x - \mu)^2 / 2\sigma^2)

mindspore/nn/probability/distribution/laplace.py

@@ -16,7 +16,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from mindspore.ops import operations as P
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
 from mindspore.nn.probability.distribution import Distribution
 from mindspore.nn.probability.distribution._utils.utils import check_greater_zero
@@ -25,25 +25,28 @@ from mindspore.nn.probability.distribution._utils.utils import check_greater_zer
 class Laplace(Distribution):
     r"""
     Laplace distribution.
-    A Laplace distribution is a continuous distribution with the range :math:`[-\inf, \inf)`
+    A Laplace distribution is a continuous distribution with the range :math:`(-\inf, \inf)`
     and the probability density function:
 
     .. math::
-        f(x, \mu, b) = 1 / (2. * b) * \exp(-abs(x - \mu) / b).
+        f(x, \mu, b) = 1 / (2 * b) * \exp(-abs(x - \mu) / b).
 
     where :math:`\mu, b` are the mean and the scale of the laplace distribution respectively.
 
     Args:
-        mean (int, float, list, numpy.ndarray, Tensor): The mean of the distribution. Default: None.
-        sd (int, float, list, numpy.ndarray, Tensor): The standard deviation of the distribution. Default: None.
-        seed (int): The seed used in sampling. The global seed is used if it is None. Default: None.
-        dtype (mindspore.dtype): The type of the event samples. Default: mstype.float32.
-        name (str): The name of the distribution. Default: 'Laplace'.
+        mean (Union[int, float, list, numpy.ndarray, Tensor], optional): The mean of the distribution.
+            If this arg is None, then the mean of the distribution will be passed in runtime. Default: None.
+        sd (Union[int, float, list, numpy.ndarray, Tensor], optional): The scale of the distribution.
+            If this arg is None, then the scale of the distribution will be passed in runtime. Default: None.
+        seed (int, optional): The seed used in sampling. The global seed is used if it is None. Default: None.
+        dtype (mindspore.dtype, optional): The type of the event samples. Default: mstype.float32.
+        name (str, optional): The name of the distribution. Default: 'Laplace'.
 
     Note:
         - `sd` must be greater than zero.
-        - `dist_spec_args` are `mean` and `sd`.
         - `dtype` must be a float type because Laplace distributions are continuous.
+        - If the arg `mean` or `sd` is passed in runtime, then it will be used as the parameter value.
+          Otherwise, the value passed in the constructor will be used.
 
     Raises:
         ValueError: When sd <= 0.
@@ -57,7 +60,7 @@ class Laplace(Distribution):
         >>> import mindspore.nn as nn
         >>> from mindspore.nn.probability.distribution import Laplace
         >>> from mindspore import Tensor
-        >>> # To initialize a Laplace distribution of the mean 3.0 and the standard deviation 4.0.
+        >>> # To initialize a Laplace distribution of the mean 3.0 and the scale 4.0.
         >>> n1 = Laplace(3.0, 4.0, dtype=mindspore.float32)
         >>> # A Laplace distribution can be initialized without arguments.
         >>> # In this case, `mean` and `sd` must be passed in through arguments.
@@ -107,12 +110,12 @@ class Laplace(Distribution):
 
     def _log_prob(self, value, mean=None, sd=None):
         r"""
-        Evaluate log probability.
+        Evaluate log probability of the laplace distribution.
 
         Args:
             value (Tensor): The value to be evaluated.
-            mean (Tensor): The mean of the distribution. Default: self._mean_value.
-            sd (Tensor): The standard deviation the distribution. Default: self._sd_value.
+            mean (Tensor, optional): The mean of the distribution. Default: self._mean_value.
+            sd (Tensor, optional): The scale the distribution. Default: self._sd_value.
 
         .. math::
             L(x) = -1* \abs{\frac{x - \mu}{\sigma}} - \log(2. * \sigma))

mindspore/nn/probability/distribution/logistic.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
@@ -32,7 +32,7 @@ class Logistic(Distribution):
     .. math::
         f(x, a, b) = 1 / b \exp(\exp(-(x - a) / b) - x).
 
-    where a and b are loc and scale parameter respectively.
+    where :math:`a, b` are loc and scale parameter respectively.
 
     Args:
         loc (float, list, numpy.ndarray, Tensor): The location of the Logistic distribution. Default: None.

mindspore/nn/probability/distribution/normal.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 mindspore.common import Tensor
 from .distribution import Distribution

mindspore/nn/probability/distribution/poisson.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

mindspore/nn/probability/distribution/student_t.py

@@ -18,7 +18,7 @@ from __future__ import division
 import numpy as np
 import mindspore.nn as nn
 from mindspore.ops import operations as P
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
 from mindspore.nn.probability.distribution import Distribution
 from mindspore.nn.probability.distribution._utils.utils import check_greater_zero
@@ -27,28 +27,33 @@ from mindspore.nn.probability.distribution._utils.utils import check_greater_zer
 class StudentT(Distribution):
     r"""
     StudentT distribution.
-    A StudentT distribution is a continuous distribution with the range :math:`[-\inf, \inf)`
+    A StudentT distribution is a continuous distribution with the range :math:`(-\inf, \inf)`
     and the probability density function:
 
     .. math::
         f(x, \nu, \mu, \sigma) = (1 + y^2 / \nu)^{(-0.5*(\nu + 1))} / Z
 
-    where :math:`y = (x-\mu)/\sigma`, :math:`Z = abs(\sigma)*\sqrt(\nu * \pi)*\Gamma(0.5 * \nu)/\Gamma(0.5*(\nu + 1))`,
-        :math:`\nu, \mu, \sigma` are the degrees of freedom , mean and scale of the laplace distribution respectively.
+    where :math:`y = (x - \mu)/ \sigma`,
+    :math:`Z = abs(\sigma) * \sqrt{(\nu * \pi)} * \Gamma(0.5 * \nu) / \Gamma(0.5 * (\nu + 1))`,
+    :math:`\nu, \mu, \sigma` are the degrees of freedom , mean and sd of the laplace distribution respectively.
 
     Args:
-        df (int, float, list, numpy.ndarray, Tensor): The degrees of freedom. Default: None.
-        mean (int, float, list, numpy.ndarray, Tensor): The mean of the distribution. Default: None.
-        sd (int, float, list, numpy.ndarray, Tensor): The standard deviation of the distribution. Default: None.
-        seed (int): The seed used in sampling. The global seed is used if it is None. Default: None.
-        dtype (mindspore.dtype): The type of the event samples. Default: mstype.float32.
-        name (str): The name of the distribution. Default: 'StudentT'.
+        df (Union[int, float, list, numpy.ndarray, Tensor], optional): The degrees of freedom.
+            If this arg is None, then the df of the distribution will be passed in runtime. Default: None.
+        mean (Union[int, float, list, numpy.ndarray, Tensor], optional): The mean of the distribution.
+            If this arg is None, then the df of the distribution will be passed in runtime. Default: None.
+        sd (Union[int, float, list, numpy.ndarray, Tensor], optional): The standard deviation of the distribution.
+            If this arg is None, then the sd of the distribution will be passed in runtime. Default: None.
+        seed (int, optional): The seed used in sampling. The global seed is used if it is None. Default: None.
+        dtype (mindspore.dtype, optional): The type of the event samples. Default: mstype.float32.
+        name (str, optional): The name of the distribution. Default: 'StudentT'.
 
     Note:
         - `df` must be greater than zero.
         - `sd` must be greater than zero.
-        - `dist_spec_args` are `mean` and `sd`.
         - `dtype` must be a float type because StudentT distributions are continuous.
+        - If the arg `df`, `mean` or `sd` is passed in runtime, then it will be used as the parameter value.
+          Otherwise, the value passed in the constructor will be used.
 
     Raises:
         ValueError: When df <= 0.
@@ -122,13 +127,13 @@ class StudentT(Distribution):
 
     def _log_prob(self, value, df=None, mean=None, sd=None):
         r"""
-        Evaluate log probability.
+        Evaluate log probability of the value of the StudentT distribution.
 
         Args:
             value (Tensor): The value to be evaluated.
-            df (Tensor): The degrees of freedom of the distribution. Default: self._df_value.
-            mean (Tensor): The mean of the distribution. Default: self._mean_value.
-            sd (Tensor): The standard deviation the distribution. Default: self._sd_value.
+            df (Tensor, optional): The degrees of freedom of the distribution. Default: self._df_value.
+            mean (Tensor, optional): The mean of the distribution. Default: self._mean_value.
+            sd (Tensor, optional): The standard deviation the distribution. Default: self._sd_value.
 
         .. math::
             L(x) = -0.5 * (\nu + 1.) * \log((x - \mu) / \sigma + 1.)) + \log(\sqrt(\pi * \mu * \sigma^2))

mindspore/nn/probability/distribution/transformed_distribution.py

@@ -14,7 +14,7 @@
 # ============================================================================
 """Transformed Distribution"""
 import numpy as np
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.ops import operations as P
 from mindspore.common import dtype as mstype
 import mindspore.nn as nn
@@ -28,9 +28,9 @@ class TransformedDistribution(Distribution):
     Transformed Distribution.
     This class contains a bijector and a distribution and transforms the original distribution
     to a new distribution through the operation defined by the bijector.
-    If X is an random variable following the underying distribution,
-    and g(x) is a function represented by the bijector,
-    then Y = g(X) is a random variable following the transformed distribution.
+    If :math:`X` is an random variable following the underying distribution,
+    and :math:`g(x)` is a function represented by the bijector,
+    then :math:`Y = g(X)` is a random variable following the transformed distribution.
 
     Args:
         bijector (Bijector): The transformation to perform.

mindspore/nn/probability/distribution/uniform.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, check_distribution_name
@@ -32,7 +32,7 @@ class Uniform(Distribution):
     .. math::
         f(x, a, b) = 1 / (b - a),
 
-    where a and b are the lower and upper bound respectively.
+    where :math:`a, b` are the lower and upper bound respectively.
 
     Args:
         low (int, float, list, numpy.ndarray, Tensor): The lower bound of the distribution. Default: None.

mindspore/nn/reinforcement/_tensors_queue.py

@@ -19,7 +19,7 @@ from __future__ import absolute_import
 
 from mindspore.nn.cell import Cell
 from mindspore.ops.operations import _rl_inner_ops as rl_ops
-from mindspore._checkparam import Rel, Validator
+from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
 
 
@@ -59,9 +59,9 @@ class TensorsQueue(Cell):
         """Initialize TensorsQueue"""
         super(TensorsQueue, self).__init__()
         Validator.check_subclass("dtype", dtype, mstype.number_type + (mstype.bool_,), self.cls_name)
-        Validator.check_int(size, 0, Rel.GE, "size", self.cls_name)
+        Validator.check_int(size, 0, Validator.GE, "size", self.cls_name)
         elements_num = len(shapes)
-        Validator.check_int(elements_num, 1, Rel.GE, "len(shapes)", self.cls_name)
+        Validator.check_int(elements_num, 1, Validator.GE, "len(shapes)", self.cls_name)
         self.handle_ = rl_ops.TensorsQueueCreate(dtype, shapes, size, name)()
         self.tensors_q_put = rl_ops.TensorsQueuePut(dtype, shapes)
         self.tensors_q_get = rl_ops.TensorsQueueGet(dtype, shapes)

mindspore/nn/reinforcement/tensor_array.py

@@ -19,7 +19,7 @@ from __future__ import absolute_import
 
 from mindspore.nn.cell import Cell
 from mindspore.ops.operations import _tensor_array as ta
-from mindspore._checkparam import Rel, Validator
+from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
 
 
@@ -62,7 +62,7 @@ class TensorArray(Cell):
         """Initialize TensorArray"""
         super(TensorArray, self).__init__()
         Validator.check_subclass("dtype", dtype, mstype.number_type + (mstype.bool_,), self.cls_name)
-        Validator.check_int(size, 0, Rel.GE, "size", self.cls_name)
+        Validator.check_int(size, 0, Validator.GE, "size", self.cls_name)
         self.handle_ = ta.TensorArray(dtype, element_shape, dynamic_size, size, name)()
         self.tensor_array_write = ta.TensorArrayWrite()
         self.tensor_array_read = ta.TensorArrayRead(dtype, element_shape)

mindspore/nn/sparse/sparse.py

@@ -45,7 +45,7 @@ class SparseToDense(Cell):
         TypeError: If `sparse_tensor.shape` is not a tuple.
 
     Supported Platforms:
-        ``CPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore as ms
@@ -118,7 +118,7 @@ class SparseTensorDenseMatmul(Cell):
             and shape of `dense` don't meet the parameter description.
 
     Supported Platforms:
-        ``CPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore as ms

mindspore/nn/wrap/cell_wrapper.py

@@ -23,11 +23,11 @@ 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
-from mindspore._checkparam import Validator as validator
+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.ops.primitive import constexpr
+from mindspore.ops.primitive import _primexpr
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
 from mindspore.ops import operations as P
@@ -110,7 +110,7 @@ class WithLossCell(Cell):
         super(WithLossCell, self).__init__(auto_prefix=False)
         self._backbone = backbone
         self._loss_fn = loss_fn
-        if backbone.jit_config_dict:
+        if isinstance(backbone, Cell) and backbone.jit_config_dict:
             self._jit_config_dict = backbone.jit_config_dict
 
     def construct(self, data, label):
@@ -147,7 +147,7 @@ class WithGradCell(Cell):
             output value. Default: None.
 
     Inputs:
-        - **(\*inputs)** (Tuple(Tensor)) - Tuple of input tensors with shape :math:`(N, \ldots)`.
+        - **\*inputs** (Tuple(Tensor)) - Tuple of input tensors with shape :math:`(N, \ldots)`.
 
     Outputs:
         list, a list of Tensors with identical shapes as trainable weights.
@@ -182,6 +182,8 @@ class WithGradCell(Cell):
         else:
             self.network_with_loss = WithLossCell(self.network, self.loss_fn)
         self.network_with_loss.set_train()
+        if isinstance(network, Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     def construct(self, *inputs):
         weights = self.weights
@@ -216,8 +218,8 @@ class ForwardValueAndGrad(Cell):
             the input parameter.
 
     Inputs:
-        - **(\*inputs)** (Tuple(Tensor...)) - Tuple of inputs with shape :math:`(N, \ldots)`.
-        - **(sens)** - A sensitivity (gradient with respect to output) as the input of backpropagation.
+        - **\*inputs** (Tuple(Tensor...)) - Tuple of inputs with shape :math:`(N, \ldots)`.
+        - **sens** - A sensitivity (gradient with respect to output) as the input of backpropagation.
           If network has single output, the sens is a tensor.
           If network has multiple outputs, the sens is the tuple(tensor).
 
@@ -282,6 +284,8 @@ class ForwardValueAndGrad(Cell):
         self.get_by_list = get_by_list
         self.sens_param = sens_param
         self.grad = C.GradOperation(get_all=self.get_all, get_by_list=self.get_by_list, sens_param=self.sens_param)
+        if isinstance(network, Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     def construct(self, *inputs):
         grad_inputs = inputs
@@ -309,7 +313,7 @@ class TrainOneStepCell(Cell):
         sens (numbers.Number): The scaling number to be filled as the input of backpropagation. Default value is 1.0.
 
     Inputs:
-        - **(\*inputs)** (Tuple(Tensor)) - Tuple of input tensors with shape :math:`(N, \ldots)`.
+        - **\*inputs** (Tuple(Tensor)) - Tuple of input tensors with shape :math:`(N, \ldots)`.
 
     Outputs:
         Tensor, a tensor means the loss value, the shape of which is usually :math:`()`.
@@ -375,6 +379,8 @@ class TrainOneStepCell(Cell):
                 create_group(server_group_name, group_list[current_index])
                 group = server_group_name
             self.grad_reducer = DistributedGradReducer(self.weights, self.mean, self.degree, group=group)
+        if isinstance(network, Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     def construct(self, *inputs):
         loss = self.network(*inputs)
@@ -453,18 +459,19 @@ class _VirtualDatasetCell(Cell):
         super(_VirtualDatasetCell, self).__init__(auto_prefix=False)
         self._backbone = backbone
         self._virtual_dataset = _VirtualDataset()
+        if isinstance(backbone, Cell) and backbone.jit_config_dict:
+            self._jit_config_dict = backbone.jit_config_dict
 
     def construct(self, *inputs):
         output = self._virtual_dataset(*inputs)
         return self._backbone(*output)
 
 
-@constexpr
+@_primexpr
 def _check_shape_value_on_axis_divided_by_target_value(input_shape, micro_size):
     if input_shape[0] % micro_size != 0:
         raise ValueError(f"For micro batch initialization, the 0th dimension shape of input({input_shape[0]}) must be "
                          f"divided by micro size({micro_size})")
-    return True
 
 
 class _MicroBatch(Cell):
@@ -545,6 +552,8 @@ class MicroBatchInterleaved(Cell):
             interleave_data.strided_slice.add_prim_attr("strided_slice_flag", True)
             interleave_data.strided_slice.add_prim_attr("interleave_num", interleave_num)
             self.interleave_inputs.append(interleave_data)
+        if isinstance(network, Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     def construct(self, *inputs):
         output = 0.0
@@ -583,6 +592,8 @@ class PipelineCell(Cell):
             self.micro_inputs.append(micro_input)
             self.add = P.Add().add_prim_attr("pipeline_end", i)
             self.add_list.append(self.add)
+        if isinstance(network, Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     def construct(self, *inputs):
         ret = None
@@ -611,6 +622,8 @@ class _TrainPipelineAccuStepCell(TrainOneStepCell):
         self.accu_grads = self.weights.clone(prefix="accu_grads", init="zeros")
         self.hyper_map = ops.HyperMap()
         self.opt_shard = _get_enable_parallel_optimizer()
+        if isinstance(network, Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     def construct(self, *inputs):
         weights = self.weights
@@ -652,6 +665,8 @@ class VirtualDatasetCellTriple(Cell):
         super(VirtualDatasetCellTriple, self).__init__(auto_prefix=False)
         logger.warning("WARN_DEPRECATED: The usage of VirtualDatasetCellTriple is deprecated.")
         self._backbone = backbone
+        if isinstance(backbone, Cell) and backbone.jit_config_dict:
+            self._jit_config_dict = backbone.jit_config_dict
 
     def construct(self, a, b, c):
         return self._backbone(a, b, c)
@@ -694,6 +709,8 @@ class WithEvalCell(Cell):
         self._network = network
         self._loss_fn = loss_fn
         self.add_cast_fp32 = validator.check_value_type("add_cast_fp32", add_cast_fp32, [bool], self.cls_name)
+        if isinstance(network, Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     def construct(self, data, label):
         outputs = self._network(data)
@@ -717,7 +734,7 @@ class ParameterUpdate(Cell):
         - **x** (Tensor) - A tensor whose shape and type are the same with `param`.
 
     Outputs:
-        Tensor, the input `x`.
+        Tensor, the updated value.
 
     Raises:
         KeyError: If parameter with the specified name does not exist.

mindspore/nn/wrap/grad_reducer.py

@@ -315,11 +315,11 @@ class DistributedGradReducer(Cell):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_ascend.html#preparations>`_
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_gpu.html#preparation>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
 
             This example should be run with multiple devices.