[Typing][B-21] Add type annotations for python/paddle/distribution/multinomial.py (#65844)

enkilee · SigureMo · web-flow · commit 18839fe285b2 · 2024-07-10T11:22:11.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/distribution/multinomial.py b/python/paddle/distribution/multinomial.py
@@ -11,12 +11,17 @@
 # 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 collections.abc import Iterable
+from typing import TYPE_CHECKING
 
 import paddle
 from paddle.distribution import categorical, distribution
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 
 class Multinomial(distribution.Distribution):
     r"""
@@ -65,8 +70,10 @@ class Multinomial(distribution.Distribution):
               [0., 6., 4.],
               [3., 3., 4.]]])
     """
+    total_count: int
+    probs: Tensor
 
-    def __init__(self, total_count, probs):
+    def __init__(self, total_count: int, probs: Tensor) -> None:
         if not isinstance(total_count, int) or total_count < 1:
             raise ValueError(
                 'input parameter total_count must be int type and grater than zero.'
@@ -86,7 +93,7 @@ def __init__(self, total_count, probs):
         super().__init__(probs.shape[:-1], probs.shape[-1:])
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         """mean of multinomial distribution.
 
         Returns:
@@ -95,15 +102,15 @@ def mean(self):
         return self.probs * self.total_count
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         """variance of multinomial distribution.
 
         Returns:
             Tensor: variance value.
         """
         return self.total_count * self.probs * (1 - self.probs)
 
-    def prob(self, value):
+    def prob(self, value: Tensor) -> Tensor:
         """probability mass function evaluated at value.
 
         Args:
@@ -114,7 +121,7 @@ def prob(self, value):
         """
         return paddle.exp(self.log_prob(value))
 
-    def log_prob(self, value):
+    def log_prob(self, value: Tensor) -> Tensor:
         """probability mass function evaluated at value.
 
         Args:
@@ -142,7 +149,7 @@ def log_prob(self, value):
             + (value * logits).sum(-1)
         )
 
-    def sample(self, shape=()):
+    def sample(self, shape: Iterable[int] = ()) -> Tensor:
         """draw sample data from multinomial distribution
 
         Args:
@@ -163,7 +170,7 @@ def sample(self, shape=()):
             .sum(0)
         )
 
-    def entropy(self):
+    def entropy(self) -> Tensor:
         """entropy of multinomial distribution
 
         Returns:
@@ -182,7 +189,7 @@ def entropy(self):
             (binomial_pmf * paddle.lgamma(support + 1)).sum([0, -1])
         )
 
-    def _binomial_logpmf(self, count, value):
+    def _binomial_logpmf(self, count: Tensor, value: Tensor) -> Tensor:
         logits = self._probs_to_logits(self.probs, is_binary=True)
 
         factor_n = paddle.lgamma(count + 1)
