chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,212 @@
|
||||
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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
|
||||
from paddle.utils.decorator_utils import param_one_alias
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import Tensor
|
||||
|
||||
|
||||
class Multinomial(distribution.Distribution):
|
||||
r"""
|
||||
Multinomial distribution parameterized by :attr:`total_count` and
|
||||
:attr:`probs`.
|
||||
|
||||
In probability theory, the multinomial distribution is a generalization of
|
||||
the binomial distribution, it models the probability of counts for each side
|
||||
of a k-sided die rolled n times. When k is 2 and n is 1, the multinomial is
|
||||
the bernoulli distribution, when k is 2 and n is grater than 1, it is the
|
||||
binomial distribution, when k is grater than 2 and n is 1, it is the
|
||||
categorical distribution.
|
||||
|
||||
The probability mass function (PMF) for multinomial is
|
||||
|
||||
.. math::
|
||||
|
||||
f(x_1, ..., x_k; n, p_1,...,p_k) = \frac{n!}{x_1!...x_k!}p_1^{x_1}...p_k^{x_k}
|
||||
|
||||
where, :math:`n` is number of trials, k is the number of categories,
|
||||
:math:`p_i` denote probability of a trial falling into each category,
|
||||
:math:`{\textstyle \sum_{i=1}^{k}p_i=1}, p_i \ge 0`, and :math:`x_i` denote
|
||||
count of each category.
|
||||
|
||||
Args:
|
||||
total_count (int): Number of trials.
|
||||
probs (Tensor): Probability of a trial falling into each category. Last
|
||||
axis of probs indexes over categories, other axes index over batches.
|
||||
Probs value should between [0, 1], and sum to 1 along last axis. If
|
||||
the value over 1, it will be normalized to sum to 1 along the last
|
||||
axis.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.seed(2023)
|
||||
>>> multinomial = paddle.distribution.Multinomial(10, paddle.to_tensor([0.2, 0.3, 0.5]))
|
||||
>>> print(multinomial.sample((2, 3)))
|
||||
Tensor(shape=[2, 3, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
|
||||
[[[1., 5., 4.],
|
||||
[0., 4., 6.],
|
||||
[1., 3., 6.]],
|
||||
[[2., 2., 6.],
|
||||
[0., 6., 4.],
|
||||
[3., 3., 4.]]])
|
||||
"""
|
||||
|
||||
total_count: int
|
||||
probs: Tensor
|
||||
|
||||
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.'
|
||||
)
|
||||
|
||||
if probs.dim() < 1:
|
||||
raise ValueError(
|
||||
'probs parameter should not be none and over one dimension'
|
||||
)
|
||||
|
||||
self.probs = probs / probs.sum(-1, keepdim=True)
|
||||
self.total_count = total_count
|
||||
self._categorical = categorical.Categorical(
|
||||
logits=self._probs_to_logits(probs)
|
||||
)
|
||||
|
||||
super().__init__(probs.shape[:-1], probs.shape[-1:])
|
||||
|
||||
@property
|
||||
def mean(self) -> Tensor:
|
||||
"""mean of multinomial distribution.
|
||||
|
||||
Returns:
|
||||
Tensor: mean value.
|
||||
"""
|
||||
return self.probs * self.total_count
|
||||
|
||||
@property
|
||||
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: Tensor) -> Tensor:
|
||||
"""probability mass function evaluated at value.
|
||||
|
||||
Args:
|
||||
value (Tensor): value to be evaluated.
|
||||
|
||||
Returns:
|
||||
Tensor: probability of value.
|
||||
"""
|
||||
return paddle.exp(self.log_prob(value))
|
||||
|
||||
def log_prob(self, value: Tensor) -> Tensor:
|
||||
"""probability mass function evaluated at value.
|
||||
|
||||
Args:
|
||||
value (Tensor): value to be evaluated.
|
||||
|
||||
Returns:
|
||||
Tensor: probability of value.
|
||||
"""
|
||||
if paddle.is_integer(value):
|
||||
value = paddle.cast(value, self.probs.dtype)
|
||||
|
||||
logits, value = paddle.broadcast_tensors(
|
||||
[paddle.log(self.probs), value]
|
||||
)
|
||||
if paddle.in_dynamic_mode():
|
||||
logits[(value == 0) & (paddle.isinf(logits))] = 0
|
||||
else:
|
||||
logits = paddle.static.setitem(
|
||||
logits, (value == 0) & (paddle.isinf(logits)), 0
|
||||
)
|
||||
|
||||
return (
|
||||
paddle.lgamma(value.sum(-1) + 1)
|
||||
- paddle.lgamma(value + 1).sum(-1)
|
||||
+ (value * logits).sum(-1)
|
||||
)
|
||||
|
||||
@param_one_alias(["shape", "sample_shape"])
|
||||
def sample(self, shape: Iterable[int] = []) -> Tensor:
|
||||
"""draw sample data from multinomial distribution
|
||||
|
||||
Args:
|
||||
sample_shape (list|tuple, optional): [description]. Defaults to [].
|
||||
"""
|
||||
if not isinstance(shape, Iterable):
|
||||
raise TypeError('sample shape must be Iterable object.')
|
||||
|
||||
samples = self._categorical.sample([self.total_count, *list(shape)])
|
||||
return (
|
||||
paddle.nn.functional.one_hot(samples, self.probs.shape[-1])
|
||||
.cast(self.probs.dtype)
|
||||
.sum(0)
|
||||
)
|
||||
|
||||
def entropy(self) -> Tensor:
|
||||
"""entropy of multinomial distribution
|
||||
|
||||
Returns:
|
||||
Tensor: entropy value
|
||||
"""
|
||||
n = paddle.full(
|
||||
shape=[], fill_value=self.total_count, dtype=self.probs.dtype
|
||||
)
|
||||
support = paddle.arange(
|
||||
self.total_count + 1, dtype=self.probs.dtype
|
||||
).reshape((-1,) + (1,) * len(self.probs.shape))[1:]
|
||||
|
||||
binomial_pmf = paddle.exp(self._binomial_logpmf(n, support))
|
||||
|
||||
return (n * self._categorical.entropy() - paddle.lgamma(n + 1)) + (
|
||||
(binomial_pmf * paddle.lgamma(support + 1)).sum([0, -1])
|
||||
)
|
||||
|
||||
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)
|
||||
factor_k = paddle.lgamma(value + 1)
|
||||
factor_nmk = paddle.lgamma(count - value + 1)
|
||||
|
||||
norm = (
|
||||
count * _clip_by_zero(logits)
|
||||
+ count * paddle.log1p(paddle.exp(-paddle.abs(logits)))
|
||||
- factor_n
|
||||
)
|
||||
|
||||
return value * logits - factor_k - factor_nmk - norm
|
||||
|
||||
|
||||
def _binomial_support(count, dtype):
|
||||
return paddle.arange(count + 1, dtype=dtype)
|
||||
|
||||
|
||||
def _clip_by_zero(x):
|
||||
# like clip(x, min=0) but grad at 0 is 0.5
|
||||
return (x.clip(min=0) + x - x.clip(max=0)) / 2
|
||||
Reference in New Issue
Block a user