[Typing][B-10] Add type annotations for python/paddle/distribution/distribution.py (#65769)

megemini · web-flow · commit 3477632332e7 · 2024-07-07T14:12:54.000+08:00
diff --git a/python/paddle/distribution/distribution.py b/python/paddle/distribution/distribution.py
@@ -12,13 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO: define the distribution functions
-# __all__ = ['Categorical',
-#            'MultivariateNormalDiag',
-#            'Normal',
-#            'Uniform']
+from __future__ import annotations
 
 import warnings
+from typing import TYPE_CHECKING, Sequence
 
 import numpy as np
 
@@ -31,6 +28,12 @@
     in_pir_mode,
 )
 
+if TYPE_CHECKING:
+    from typing_extensions import TypeGuard
+
+    from paddle import Tensor
+    from paddle._typing import NestedNumbericSequence, TensorLike
+
 
 class Distribution:
     """
@@ -46,7 +49,9 @@ class Distribution:
             multivariate distribution, the event shape is [n].
     """
 
-    def __init__(self, batch_shape=(), event_shape=()):
+    def __init__(
+        self, batch_shape: Sequence[int] = (), event_shape: Sequence[int] = ()
+    ) -> None:
         self._batch_shape = (
             batch_shape
             if isinstance(batch_shape, tuple)
@@ -61,7 +66,7 @@ def __init__(self, batch_shape=(), event_shape=()):
         super().__init__()
 
     @property
-    def batch_shape(self):
+    def batch_shape(self) -> Sequence[int]:
         """Returns batch shape of distribution
 
         Returns:
@@ -70,7 +75,7 @@ def batch_shape(self):
         return self._batch_shape
 
     @property
-    def event_shape(self):
+    def event_shape(self) -> Sequence[int]:
         """Returns event shape of distribution
 
         Returns:
@@ -79,44 +84,44 @@ def event_shape(self):
         return self._event_shape
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         """Mean of distribution"""
         raise NotImplementedError
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         """Variance of distribution"""
         raise NotImplementedError
 
-    def sample(self, shape=()):
+    def sample(self, shape: Sequence[int] = ()) -> Tensor:
         """Sampling from the distribution."""
         raise NotImplementedError
 
-    def rsample(self, shape=()):
+    def rsample(self, shape: Sequence[int] = ()) -> Tensor:
         """reparameterized sample"""
         raise NotImplementedError
 
-    def entropy(self):
+    def entropy(self) -> Tensor:
         """The entropy of the distribution."""
         raise NotImplementedError
 
-    def kl_divergence(self, other):
+    def kl_divergence(self, other: Distribution) -> Tensor:
         """The KL-divergence between self distributions and other."""
         raise NotImplementedError
 
-    def prob(self, value):
+    def prob(self, value: Tensor) -> Tensor:
         """Probability density/mass function evaluated at value.
 
         Args:
             value (Tensor): value which will be evaluated
         """
         return self.log_prob(value).exp()
 
-    def log_prob(self, value):
+    def log_prob(self, value: Tensor) -> Tensor:
         """Log probability density/mass function."""
         raise NotImplementedError
 
-    def probs(self, value):
+    def probs(self, value: Tensor) -> Tensor:
         """Probability density/mass function.
 
         Note:
@@ -126,11 +131,11 @@ def probs(self, value):
         """
         raise NotImplementedError
 
-    def _extend_shape(self, sample_shape):
+    def _extend_shape(self, sample_shape: Sequence[int] | Tensor) -> Tensor:
         """compute shape of the sample
 
         Args:
-            sample_shape (Tensor): sample shape
+            sample_shape (Sequence[int]|Tensor): sample shape
 
         Returns:
             Tensor: generated sample data shape
@@ -141,7 +146,9 @@ def _extend_shape(self, sample_shape):
             + tuple(self._event_shape)
         )
 
-    def _validate_args(self, *args):
+    def _validate_args(
+        self, *args: TensorLike | NestedNumbericSequence
+    ) -> TypeGuard[Tensor]:
         """
         Argument validation for distribution args
         Args:
@@ -164,7 +171,9 @@ def _validate_args(self, *args):
 
         return is_variable
 
-    def _to_tensor(self, *args):
+    def _to_tensor(
+        self, *args: TensorLike | NestedNumbericSequence
+    ) -> tuple[Tensor, ...]:
         """
         Argument convert args to Tensor
 
@@ -221,7 +230,9 @@ def _to_tensor(self, *args):
 
         return tuple(variable_args)
 
-    def _check_values_dtype_in_probs(self, param, value):
+    def _check_values_dtype_in_probs(
+        self, param: Tensor, value: Tensor
+    ) -> Tensor:
         """
         Log_prob and probs methods have input ``value``, if value's dtype is different from param,
         convert value's dtype to be consistent with param's dtype.
@@ -265,7 +276,9 @@ def _check_values_dtype_in_probs(self, param, value):
             return paddle.cast(value, dtype=param.dtype)
         return value
 
-    def _probs_to_logits(self, probs, is_binary=False):
+    def _probs_to_logits(
+        self, probs: float | Tensor, is_binary: bool = False
+    ) -> Tensor:
         r"""
         Converts probabilities into logits. For the binary, probs denotes the
         probability of occurrence of the event indexed by `1`. For the
@@ -278,7 +291,9 @@ def _probs_to_logits(self, probs, is_binary=False):
             else paddle.log(probs)
         )
 
-    def _logits_to_probs(self, logits, is_binary=False):
+    def _logits_to_probs(
+        self, logits: float | Tensor, is_binary: bool = False
+    ) -> Tensor:
         r"""
         Converts logits into probabilities. For the binary, each value denotes
         log odds, whereas for the multi-dimensional case, the values along the
