improved the docs in the distribution classes

mindspore/nn/probability/distribution/bernoulli.py

@@ -163,7 +163,8 @@ class Bernoulli(Distribution):
     @property
     def probs(self):
         """
-        Return the probability of that the outcome is 1.
+        Return the probability of that the outcome is 1
+        after casting to self.dtype.
         """
         return self._probs
 

mindspore/nn/probability/distribution/beta.py

@@ -190,14 +190,16 @@ class Beta(Distribution):
     @property
     def concentration1(self):
         """
-        Return the concentration1, also know as the alpha of the Beta distribution.
+        Return the concentration1, also know as the alpha of the Beta distribution,
+        after casting to self.dtype.
         """
         return self._concentration1
 
     @property
     def concentration0(self):
         """
-        Return the concentration0, also know as the beta of the Beta distribution.
+        Return the concentration0, also know as the beta of the Beta distribution,
+        after casting to self.dtype.
         """
         return self._concentration0
 

mindspore/nn/probability/distribution/categorical.py

@@ -190,7 +190,7 @@ class Categorical(Distribution):
     @property
     def probs(self):
         """
-        Return the probability.
+        Return the probability after casting to self.dtype.
         """
         return self._probs
 

mindspore/nn/probability/distribution/cauchy.py

@@ -181,14 +181,14 @@ class Cauchy(Distribution):
     @property
     def loc(self):
         """
-        Return the location of the distribution.
+        Return the location of the distribution after casting to self.dtype.
         """
         return self._loc
 
     @property
     def scale(self):
         """
-        Return the scale of the distribution.
+        Return the scale of the distribution after casting to self.dtype.
         """
         return self._scale
 

mindspore/nn/probability/distribution/exponential.py

@@ -167,7 +167,7 @@ class Exponential(Distribution):
     @property
     def rate(self):
         """
-        Return `rate` of the distribution.
+        Return `rate` of the distribution after casting to self.dtype.
         """
         return self._rate
 

mindspore/nn/probability/distribution/gamma.py

@@ -188,14 +188,16 @@ class Gamma(Distribution):
     @property
     def concentration(self):
         """
-        Return the concentration, also know as the alpha of the Gamma distribution.
+        Return the concentration, also know as the alpha of the Gamma distribution,
+        after casting to self.dtype.
         """
         return self._concentration
 
     @property
     def rate(self):
         """
-        Return the rate, also know as the beta of the Gamma distribution.
+        Return the rate, also know as the beta of the Gamma distribution,
+        after casting to self.dtype.
         """
         return self._rate
 

mindspore/nn/probability/distribution/geometric.py

@@ -172,7 +172,8 @@ class Geometric(Distribution):
     @property
     def probs(self):
         """
-        Return the probability of success of the Bernoulli trail.
+        Return the probability of success of the Bernoulli trail,
+        after casting to self.dtype.
         """
         return self._probs
 

mindspore/nn/probability/distribution/gumbel.py

@@ -127,10 +127,16 @@ class Gumbel(TransformedDistribution):
 
     @property
     def loc(self):
+        """
+        Return the location of the distribution after casting to self.dtype.
+        """
         return self._loc
 
     @property
     def scale(self):
+        """
+        Return the scale of the distribution after casting to self.dtype.
+        """
         return self._scale
 
     def extend_repr(self):

mindspore/nn/probability/distribution/log_normal.py

@@ -165,27 +165,35 @@ class LogNormal(msd.TransformedDistribution):
         self.log_2pi = np.log(2 * np.pi)
 
         #ops needed for the class
+        self.dtypeop = P.DType()
         self.exp = exp_generic
         self.expm1 = P.Expm1()
         self.log = log_generic
         self.const = P.ScalarToArray()
         self.erf = P.Erf()
         self.fill = P.Fill()
+        self.greater = P.Greater()
+        self.select = P.Select()
         self.shape = P.Shape()
         self.sq = P.Square()
         self.sqrt = P.Sqrt()
         self.cast = P.Cast()
         self.squeeze = P.Squeeze(0)
-        self.zeroslike = P.ZerosLike()
 
     @property
     def loc(self):
-        """Distribution parameter for the pre-transformed mean."""
+        """
+        Distribution parameter for the pre-transformed mean
+        after casting to self.dtype.
+        """
         return self._loc
 
     @property
     def scale(self):
-        """Distribution parameter for the pre-transformed standard deviation."""
+        """
+        Distribution parameter for the pre-transformed standard deviation
+        after casting to self.dtype.
+        """
         return self._scale
 
     def _get_dist_type(self):
@@ -254,7 +262,12 @@ class LogNormal(msd.TransformedDistribution):
         """
         mean, sd = self._check_param_type(loc, scale)
         inverse_value = self.bijector("inverse", value)
-        return self.distribution("cdf", inverse_value, mean, sd)
+        cdf = self.distribution("cdf", inverse_value, mean, sd)
+
+        # to increase numerical stability, set cdf = 0 when value <= 0
+        zeros = self.fill(self.dtypeop(cdf), self.shape(cdf), 0.0)
+
+        return self.select(self.greater(value, 0.), cdf, zeros)
 
     def _log_prob(self, value, loc=None, scale=None):
         r"""

mindspore/nn/probability/distribution/logistic.py

@@ -181,14 +181,14 @@ class Logistic(Distribution):
     @property
     def loc(self):
         """
-        Return the location of the distribution.
+        Return the location of the distribution after casting to self.dtype.
         """
         return self._loc
 
     @property
     def scale(self):
         """
-        Return the scale of the distribution.
+        Return the scale of the distribution after casting to self.dtype.
         """
         return self._scale
 

mindspore/nn/probability/distribution/poisson.py

@@ -157,7 +157,7 @@ class Poisson(Distribution):
     @property
     def rate(self):
         """
-        Return `rate` of the distribution.
+        Return `rate` of the distribution after casting to self.dtype.
         """
         return self._rate
 
@@ -212,7 +212,6 @@ class Poisson(Distribution):
         """
         value = self._check_value(value, "value")
         value = self.cast(value, self.dtype)
-        value = self.floor(value)
         rate = self._check_param_type(rate)
         log_rate = self.log(rate)
         zeros = self.fill(self.dtypeop(value), self.shape(value), 0.0)
@@ -240,7 +239,6 @@ class Poisson(Distribution):
         """
         value = self._check_value(value, 'value')
         value = self.cast(value, self.dtype)
-        value = self.floor(value)
         rate = self._check_param_type(rate)
         zeros = self.fill(self.dtypeop(value), self.shape(value), 0.0)
         comp = self.less(value, zeros)

mindspore/nn/probability/distribution/uniform.py

@@ -181,14 +181,14 @@ class Uniform(Distribution):
     @property
     def low(self):
         """
-        Return the lower bound of the distribution.
+        Return the lower bound of the distribution after casting to self.dtype.
         """
         return self._low
 
     @property
     def high(self):
         """
-        Return the upper bound of the distribution.
+        Return the upper bound of the distribution after casting to self.dtype..
         """
         return self._high