wehub-resource-sync
112 lines
3.9 KiB
Python
112 lines
3.9 KiB
Python
# 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 typing import TYPE_CHECKING
|
|
|
|
from paddle.distribution import distribution
|
|
from paddle.utils.decorator_utils import param_one_alias
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
|
|
from paddle import Tensor
|
|
|
|
|
|
class Independent(distribution.Distribution):
|
|
r"""
|
|
Reinterprets some of the batch dimensions of a distribution as event dimensions.
|
|
|
|
This is mainly useful for changing the shape of the result of
|
|
:meth:`log_prob`.
|
|
|
|
Args:
|
|
base (Distribution): The base distribution.
|
|
reinterpreted_batch_rank (int): The number of batch dimensions to
|
|
reinterpret as event dimensions.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.distribution import independent
|
|
|
|
>>> beta = paddle.distribution.Beta(paddle.to_tensor([0.5, 0.5]), paddle.to_tensor([0.5, 0.5]))
|
|
>>> print(beta.batch_shape, beta.event_shape)
|
|
(2,) ()
|
|
>>> print(beta.log_prob(paddle.to_tensor(0.2)))
|
|
Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True,
|
|
[-0.22843921, -0.22843921])
|
|
>>> reinterpreted_beta = independent.Independent(beta, 1)
|
|
>>> print(reinterpreted_beta.batch_shape, reinterpreted_beta.event_shape)
|
|
() (2,)
|
|
>>> print(reinterpreted_beta.log_prob(paddle.to_tensor([0.2, 0.2])))
|
|
Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True,
|
|
-0.45687842)
|
|
"""
|
|
|
|
def __init__(
|
|
self, base: distribution.Distribution, reinterpreted_batch_rank: int
|
|
) -> None:
|
|
if not isinstance(base, distribution.Distribution):
|
|
raise TypeError(
|
|
f"Expected type of 'base' is Distribution, but got {type(base)}"
|
|
)
|
|
if not (0 < reinterpreted_batch_rank <= len(base.batch_shape)):
|
|
raise ValueError(
|
|
f"Expected 0 < reinterpreted_batch_rank <= {len(base.batch_shape)}, but got {reinterpreted_batch_rank}"
|
|
)
|
|
self._base = base
|
|
self._reinterpreted_batch_rank = reinterpreted_batch_rank
|
|
|
|
shape = base.batch_shape + base.event_shape
|
|
super().__init__(
|
|
batch_shape=shape[
|
|
: len(base.batch_shape) - reinterpreted_batch_rank
|
|
],
|
|
event_shape=shape[
|
|
len(base.batch_shape) - reinterpreted_batch_rank :
|
|
],
|
|
)
|
|
|
|
@property
|
|
def mean(self) -> Tensor:
|
|
return self._base.mean
|
|
|
|
@property
|
|
def variance(self) -> Tensor:
|
|
return self._base.variance
|
|
|
|
@param_one_alias(["shape", "sample_shape"])
|
|
def sample(self, shape: Sequence[int] = []) -> Tensor:
|
|
return self._base.sample(shape)
|
|
|
|
def log_prob(self, value: Tensor) -> Tensor:
|
|
return self._sum_rightmost(
|
|
self._base.log_prob(value), self._reinterpreted_batch_rank
|
|
)
|
|
|
|
def prob(self, value: Tensor) -> Tensor:
|
|
return self.log_prob(value).exp()
|
|
|
|
def entropy(self) -> Tensor:
|
|
return self._sum_rightmost(
|
|
self._base.entropy(), self._reinterpreted_batch_rank
|
|
)
|
|
|
|
def _sum_rightmost(self, value: Tensor, n: int) -> Tensor:
|
|
return value.sum(list(range(-n, 0))) if n > 0 else value
|