[Typing][B-24] Add type annotations for python/paddle/distribution/poisson.py (#65852)

enkilee · SigureMo · web-flow · commit a88f4f7fbaf3 · 2024-07-14T12:27:10.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/distribution/poisson.py b/python/paddle/distribution/poisson.py
@@ -11,12 +11,19 @@
 # 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 Sequence
+from typing import TYPE_CHECKING
 
 import paddle
+from paddle.base.data_feeder import convert_dtype
 from paddle.distribution import distribution
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing.dtype_like import _DTypeLiteral
+
 
 class Poisson(distribution.Distribution):
     r"""
@@ -72,7 +79,10 @@ class Poisson(distribution.Distribution):
              [0.06825157  , 1.53426421  ]])
     """
 
-    def __init__(self, rate):
+    rate: Tensor
+    dtype: _DTypeLiteral
+
+    def __init__(self, rate: float | Tensor) -> None:
         self.dtype = paddle.get_default_dtype()
         self.rate = self._to_tensor(rate)
 
@@ -82,7 +92,7 @@ def __init__(self, rate):
             batch_shape = self.rate.shape
         super().__init__(batch_shape)
 
-    def _to_tensor(self, rate):
+    def _to_tensor(self, rate: float | Tensor) -> Tensor:
         """Convert the input parameters into tensors.
 
         Returns:
@@ -92,11 +102,11 @@ def _to_tensor(self, rate):
         if isinstance(rate, (float, int)):
             rate = paddle.to_tensor([rate], dtype=self.dtype)
         else:
-            self.dtype = rate.dtype
+            self.dtype = convert_dtype(rate.dtype)
         return rate
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         """Mean of poisson distribution.
 
         Returns:
@@ -105,15 +115,15 @@ def mean(self):
         return self.rate
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         """Variance of poisson distribution.
 
         Returns:
             Tensor: variance value.
         """
         return self.rate
 
-    def sample(self, shape=()):
+    def sample(self, shape: Sequence[int] = ()) -> Tensor:
         """Generate poisson samples of the specified shape. The final shape would be ``shape+batch_shape`` .
 
         Args:
@@ -133,7 +143,7 @@ def sample(self, shape=()):
         with paddle.no_grad():
             return paddle.poisson(output_rate)
 
-    def entropy(self):
+    def entropy(self) -> Tensor:
         r"""Shannon entropy in nats.
 
         The entropy is
@@ -162,7 +172,7 @@ def entropy(self):
         )
         return paddle.multiply(proposed, mask)
 
-    def _enumerate_bounded_support(self, rate):
+    def _enumerate_bounded_support(self, rate: float | Tensor) -> Tensor:
         """Generate a bounded approximation of the support. Approximately view Poisson r.v. as a
         Normal r.v. with mu = rate and sigma = sqrt(rate). Then by 30-sigma rule, generate a bounded
         approximation of the support.
@@ -203,7 +213,7 @@ def false_func():
             values = paddle.arange(0, upper, dtype=self.dtype)
             return values
 
-    def log_prob(self, value):
+    def log_prob(self, value: Tensor) -> Tensor:
         """Log probability density/mass function.
 
         Args:
@@ -223,7 +233,7 @@ def log_prob(self, value):
             neginf=-eps,
         )
 
-    def prob(self, value):
+    def prob(self, value: Tensor) -> Tensor:
         """Probability density/mass function.
 
         Args:
@@ -234,7 +244,7 @@ def prob(self, value):
         """
         return paddle.exp(self.log_prob(value))
 
-    def kl_divergence(self, other):
+    def kl_divergence(self, other: Poisson) -> Tensor:
         r"""The KL-divergence between two poisson distributions with the same `batch_shape`.
 
         The probability density function (pdf) is
