[Typing][B-28] Add type annotations for python/paddle/distribution/uniform.py (#65660)

NKNaN · SigureMo · web-flow · commit 0fc0aa5d92fd · 2024-07-04T15:20:52.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/distribution/uniform.py b/python/paddle/distribution/uniform.py
@@ -11,8 +11,12 @@
 # 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
 
 import numpy as np
+import numpy.typing as npt
 
 import paddle
 from paddle import _C_ops
@@ -22,6 +26,9 @@
 from paddle.framework import in_dynamic_mode
 from paddle.tensor import random
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 
 class Uniform(distribution.Distribution):
     r"""Uniform distribution with `low` and `high` parameters.
@@ -99,8 +106,25 @@ class Uniform(distribution.Distribution):
             Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
                 [0.50000000])
     """
-
-    def __init__(self, low, high, name=None):
+    low: Tensor
+    high: Tensor
+
+    def __init__(
+        self,
+        low: (
+            float
+            | Sequence[float]
+            | npt.NDArray[np.float32 | np.float64]
+            | Tensor
+        ),
+        high: (
+            float
+            | Sequence[float]
+            | npt.NDArray[np.float32 | np.float64]
+            | Tensor
+        ),
+        name: str | None = None,
+    ) -> None:
         if not in_dynamic_mode():
             check_type(
                 low,
@@ -165,7 +189,7 @@ def __init__(self, low, high, name=None):
 
         super().__init__(self.low.shape)
 
-    def sample(self, shape, seed=0):
+    def sample(self, shape: list[int], seed: int = 0) -> Tensor:
         """Generate samples of the specified shape.
 
         Args:
@@ -218,7 +242,7 @@ def sample(self, shape, seed=0):
             else:
                 return output
 
-    def log_prob(self, value):
+    def log_prob(self, value: Tensor) -> Tensor:
         """Log probability density/mass function.
 
         Args:
@@ -247,7 +271,7 @@ def log_prob(self, value):
                 paddle.log(lb * ub), paddle.log(self.high - self.low), name=name
             )
 
-    def probs(self, value):
+    def probs(self, value: Tensor) -> Tensor:
         """Probability density/mass function.
 
         Args:
@@ -272,7 +296,7 @@ def probs(self, value):
             ub = paddle.cast(ub_bool, dtype=value.dtype)
             return paddle.divide((lb * ub), (self.high - self.low), name=name)
 
-    def entropy(self):
+    def entropy(self) -> Tensor:
         r"""Shannon entropy in nats.
 
         The entropy is
