diff --git a/python/paddle/distribution/uniform.py b/python/paddle/distribution/uniform.py index b9b4cf1e33480..cefbeef9c6043 100644 --- 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