[Typing][B-20] Add type annotations for python/paddle/distribution/lognormal.py (#65843)

enkilee · web-flow · commit 274d39c47988 · 2024-07-12T14:33:22.000+08:00
diff --git a/python/paddle/distribution/lognormal.py b/python/paddle/distribution/lognormal.py
@@ -11,12 +11,42 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence, Union
 
 import paddle
 from paddle.distribution.normal import Normal
 from paddle.distribution.transform import ExpTransform
 from paddle.distribution.transformed_distribution import TransformedDistribution
 
+if TYPE_CHECKING:
+    import numpy as np
+    import numpy.typing as npt
+    from typing_extensions import TypeAlias
+
+    from paddle import Tensor
+    from paddle._typing import NestedSequence
+
+    _LognormalLocBase: TypeAlias = Union[float, complex]
+    _LognormalLocNDArray: TypeAlias = Union[
+        np.float32, np.float64, np.complex64, np.complex128
+    ]
+    _LognormalLoc: TypeAlias = Union[
+        _LognormalLocBase,
+        Sequence[_LognormalLocBase],
+        NestedSequence[_LognormalLocBase],
+        npt.NDArray[_LognormalLocNDArray],
+        Tensor,
+    ]
+    _LognormalScale: TypeAlias = Union[
+        float,
+        Sequence[float],
+        NestedSequence[float],
+        npt.NDArray[Union[np.float32, np.float64]],
+        Tensor,
+    ]
+
 
 class LogNormal(TransformedDistribution):
     r"""The LogNormal distribution with location `loc` and `scale` parameters.
@@ -43,7 +73,7 @@ class LogNormal(TransformedDistribution):
     * :math:`scale = \sigma`: is the stddevs of the underlying Normal distribution.
 
     Args:
-        loc(int|float|list|tuple|numpy.ndarray|Tensor): The means of the underlying Normal distribution.
+        loc(int|float|complex|list|tuple|numpy.ndarray|Tensor): The means of the underlying Normal distribution.The data type is float32, float64, complex64 and complex128.
         scale(int|float|list|tuple|numpy.ndarray|Tensor): The stddevs of the underlying Normal distribution.
 
     Examples:
@@ -88,15 +118,17 @@ class LogNormal(TransformedDistribution):
             Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
                 [0.34939718])
     """
+    loc: Tensor
+    scale: Tensor
 
-    def __init__(self, loc, scale):
+    def __init__(self, loc: _LognormalLoc, scale: _LognormalScale) -> None:
         self._base = Normal(loc=loc, scale=scale)
         self.loc = self._base.loc
         self.scale = self._base.scale
         super().__init__(self._base, [ExpTransform()])
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         """Mean of lognormal distribution.
 
         Returns:
@@ -105,7 +137,7 @@ def mean(self):
         return paddle.exp(self._base.mean + self._base.variance / 2)
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         """Variance of lognormal distribution.
 
         Returns:
@@ -115,7 +147,7 @@ def variance(self):
             2 * self._base.mean + self._base.variance
         )
 
-    def entropy(self):
+    def entropy(self) -> Tensor:
         r"""Shannon entropy in nats.
 
         The entropy is
@@ -135,7 +167,7 @@ def entropy(self):
         """
         return self._base.entropy() + self._base.mean
 
-    def probs(self, value):
+    def probs(self, value: Tensor) -> Tensor:
         """Probability density/mass function.
 
         Args:
@@ -147,7 +179,7 @@ def probs(self, value):
         """
         return paddle.exp(self.log_prob(value))
 
-    def kl_divergence(self, other):
+    def kl_divergence(self, other: LogNormal) -> Tensor:
         r"""The KL-divergence between two lognormal distributions.
 
         The probability density function (pdf) is
