Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Typing][B-15] Add type annotations for python/paddle/distribution/gumbel.py #65774

Merged
merged 1 commit into from
Jul 6, 2024
Merged

[Typing][B-15] Add type annotations for python/paddle/distribution/gumbel.py #65774

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions python/paddle/distribution/gumbel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,22 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import math
import numbers
from typing import TYPE_CHECKING, Sequence

import numpy as np

import paddle
from paddle.base import framework
from paddle.distribution.transformed_distribution import TransformedDistribution

if TYPE_CHECKING:
from paddle import Tensor
from paddle.distribution import Transform, Uniform


class Gumbel(TransformedDistribution):
r"""The Gumbel distribution with location `loc` and `scale` parameters.
Expand Down Expand Up @@ -82,7 +89,12 @@ class Gumbel(TransformedDistribution):
[1.57721567])
"""

def __init__(self, loc, scale):
loc: Tensor
scale: Tensor
base_dist: Uniform
transforms: tuple[Transform, ...]

def __init__(self, loc: float | Tensor, scale: float | Tensor) -> None:
if not isinstance(
loc, (numbers.Real, framework.Variable, paddle.pir.Value)
):
Expand Down Expand Up @@ -118,7 +130,7 @@ def __init__(self, loc, scale):
super().__init__(self.base_dist, self.transforms)

@property
def mean(self):
def mean(self) -> Tensor:
r"""Mean of distribution

The mean is
Expand All @@ -140,7 +152,7 @@ def mean(self):
return self.loc + self.scale * np.euler_gamma

@property
def variance(self):
def variance(self) -> Tensor:
r"""Variance of distribution.

The variance is
Expand All @@ -166,7 +178,7 @@ def variance(self):
return paddle.pow(self.scale, 2) * temp / 6

@property
def stddev(self):
def stddev(self) -> Tensor:
r"""Standard deviation of distribution

The standard deviation is
Expand All @@ -183,7 +195,7 @@ def stddev(self):
"""
return paddle.sqrt(self.variance)

def prob(self, value):
def prob(self, value: Tensor) -> Tensor:
"""Probability density/mass function

Args:
Expand All @@ -199,7 +211,7 @@ def prob(self, value):

return paddle.exp(y - paddle.exp(y)) / self.scale.astype(y.dtype)

def log_prob(self, value):
def log_prob(self, value: Tensor) -> Tensor:
"""Log probability density/mass function.

Args:
Expand All @@ -211,7 +223,7 @@ def log_prob(self, value):
"""
return paddle.log(self.prob(value))

def cdf(self, value):
def cdf(self, value: Tensor) -> Tensor:
"""Cumulative distribution function.
Args:
value (Tensor): value to be evaluated.
Expand All @@ -227,7 +239,7 @@ def cdf(self, value):
)
)

def entropy(self):
def entropy(self) -> Tensor:
"""Entropy of Gumbel distribution.

Returns:
Expand All @@ -236,7 +248,7 @@ def entropy(self):
"""
return paddle.log(self.scale) + 1 + np.euler_gamma

def sample(self, shape):
def sample(self, shape: Sequence[int]) -> Tensor:
"""Sample from ``Gumbel``.

Args:
Expand All @@ -249,7 +261,7 @@ def sample(self, shape):
with paddle.no_grad():
return self.rsample(shape)

def rsample(self, shape):
def rsample(self, shape: Sequence[int]) -> Tensor:
"""reparameterized sample
Args:
shape (Sequence[int]): 1D `int32`. Shape of the generated samples.
Expand Down