chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,553 @@
|
||||
# Copyright (c) 2021 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
|
||||
|
||||
import math
|
||||
from collections.abc import Iterable, Sequence
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
import paddle
|
||||
from paddle.base.data_feeder import check_type, convert_dtype
|
||||
from paddle.base.framework import Variable
|
||||
from paddle.distribution import constraint, distribution
|
||||
from paddle.framework import in_dynamic_mode
|
||||
from paddle.tensor import random
|
||||
from paddle.utils.decorator_utils import param_one_alias
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import TypeAlias
|
||||
|
||||
from paddle import Tensor, dtype
|
||||
from paddle._typing import NestedSequence
|
||||
|
||||
_NormalLocBase: TypeAlias = float | complex
|
||||
_NormalLocNDArray: TypeAlias = (
|
||||
np.float32 | np.float64 | np.complex64 | np.complex128
|
||||
)
|
||||
_NormalLoc: TypeAlias = (
|
||||
_NormalLocBase
|
||||
| Sequence[_NormalLocBase]
|
||||
| NestedSequence[_NormalLocBase]
|
||||
| npt.NDArray[_NormalLocNDArray]
|
||||
| Tensor
|
||||
)
|
||||
_NormalScale: TypeAlias = (
|
||||
float
|
||||
| Sequence[float]
|
||||
| NestedSequence[float]
|
||||
| npt.NDArray[np.float32 | np.float64]
|
||||
| Tensor
|
||||
)
|
||||
|
||||
|
||||
class Normal(distribution.Distribution):
|
||||
r"""The Normal distribution with location `loc` and `scale` parameters.
|
||||
|
||||
Mathematical details
|
||||
|
||||
If 'loc' is real number, the probability density function (pdf) is
|
||||
|
||||
.. math::
|
||||
|
||||
pdf(x; \mu, \sigma) = \frac{1}{Z}e^{\frac {-0.5 (x - \mu)^2} {\sigma^2} }
|
||||
|
||||
.. math::
|
||||
|
||||
Z = (2 \pi \sigma^2)^{0.5}
|
||||
|
||||
If 'loc' is complex number, the probability density function (pdf) is
|
||||
|
||||
.. math::
|
||||
|
||||
pdf(x; \mu, \sigma) = \frac{1}{Z}e^{\frac {-(x - \mu)^2} {\sigma^2} }
|
||||
|
||||
.. math::
|
||||
|
||||
Z = \pi \sigma^2
|
||||
|
||||
In the above equations:
|
||||
|
||||
* :math:`loc = \mu`: is the mean.
|
||||
* :math:`scale = \sigma`: is the std.
|
||||
* :math:`Z`: is the normalization constant.
|
||||
|
||||
Args:
|
||||
loc(int|float|complex|list|tuple|numpy.ndarray|Tensor): The mean of normal distribution.The data type is float32, float64, complex64 and complex128.
|
||||
scale(int|float|list|tuple|numpy.ndarray|Tensor): The std of normal distribution.The data type is float32 and float64.
|
||||
validate_args(bool|None, optional): Whether to validate input arguments. Default is None.
|
||||
name(str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.distribution import Normal
|
||||
|
||||
>>> # Define a single scalar Normal distribution.
|
||||
>>> dist = Normal(loc=0.0, scale=3.0)
|
||||
>>> # Define a batch of two scalar valued Normals.
|
||||
>>> # The first has mean 1 and standard deviation 11, the second 2 and 22.
|
||||
>>> dist = Normal(loc=[1.0, 2.0], scale=[11.0, 22.0])
|
||||
>>> # Get 3 samples, returning a 3 x 2 tensor.
|
||||
>>> dist.sample([3])
|
||||
|
||||
>>> # Define a batch of two scalar valued Normals.
|
||||
>>> # Both have mean 1, but different standard deviations.
|
||||
>>> dist = Normal(loc=1.0, scale=[11.0, 22.0])
|
||||
|
||||
>>> # Complete example
|
||||
>>> value_tensor = paddle.to_tensor([0.8], dtype="float32")
|
||||
|
||||
>>> normal_a = Normal([0.0], [1.0])
|
||||
>>> normal_b = Normal([0.5], [2.0])
|
||||
>>> sample = normal_a.sample([2])
|
||||
>>> # a random tensor created by normal distribution with shape: [2, 1]
|
||||
>>> entropy = normal_a.entropy()
|
||||
>>> print(entropy)
|
||||
Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
|
||||
[1.41893852])
|
||||
>>> lp = normal_a.log_prob(value_tensor)
|
||||
>>> print(lp)
|
||||
Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
|
||||
[-1.23893857])
|
||||
>>> p = normal_a.probs(value_tensor)
|
||||
>>> print(p)
|
||||
Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
|
||||
[0.28969154])
|
||||
>>> kl = normal_a.kl_divergence(normal_b)
|
||||
>>> print(kl)
|
||||
Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
|
||||
[0.34939718])
|
||||
"""
|
||||
|
||||
loc: Tensor
|
||||
scale: Tensor
|
||||
name: str
|
||||
dtype: dtype
|
||||
|
||||
arg_constraints = {
|
||||
"loc": constraint.real,
|
||||
"scale": constraint.positive,
|
||||
}
|
||||
support = constraint.real
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
loc: _NormalLoc,
|
||||
scale: _NormalScale,
|
||||
validate_args: bool | None = None,
|
||||
name: str | None = None,
|
||||
) -> None:
|
||||
if not in_dynamic_mode():
|
||||
check_type(
|
||||
loc,
|
||||
'loc',
|
||||
(
|
||||
int,
|
||||
float,
|
||||
complex,
|
||||
np.ndarray,
|
||||
Variable,
|
||||
paddle.pir.Value,
|
||||
list,
|
||||
tuple,
|
||||
),
|
||||
'Normal',
|
||||
)
|
||||
check_type(
|
||||
scale,
|
||||
'scale',
|
||||
(
|
||||
int,
|
||||
float,
|
||||
np.ndarray,
|
||||
Variable,
|
||||
paddle.pir.Value,
|
||||
list,
|
||||
tuple,
|
||||
),
|
||||
'Normal',
|
||||
)
|
||||
|
||||
self.all_arg_is_float = False
|
||||
self.name = name if name is not None else 'Normal'
|
||||
self.dtype = 'float32'
|
||||
self._complex_gaussian = False
|
||||
|
||||
if isinstance(loc, int):
|
||||
loc = float(loc)
|
||||
if isinstance(scale, int):
|
||||
scale = float(scale)
|
||||
|
||||
if isinstance(loc, (tuple, list)):
|
||||
loc = np.array(loc)
|
||||
if loc.dtype == np.float64:
|
||||
loc = loc.astype('float32')
|
||||
if loc.dtype == np.complex128:
|
||||
loc = loc.astype('complex64')
|
||||
|
||||
if isinstance(scale, (tuple, list)):
|
||||
scale = np.array(scale, dtype=np.float32)
|
||||
|
||||
if (
|
||||
isinstance(loc, complex)
|
||||
or (
|
||||
isinstance(loc, np.ndarray)
|
||||
and loc.dtype in [np.complex64, np.complex128]
|
||||
)
|
||||
or (self._validate_args(loc) and loc.is_complex())
|
||||
):
|
||||
self._complex_gaussian = True
|
||||
if isinstance(loc, complex) and isinstance(scale, float):
|
||||
self.all_arg_is_float = True
|
||||
|
||||
if isinstance(loc, np.ndarray):
|
||||
real_dtype = (
|
||||
'float32' if loc.dtype == np.complex64 else 'float64'
|
||||
)
|
||||
imag_dtype = (
|
||||
'float32' if loc.dtype == np.complex64 else 'float64'
|
||||
)
|
||||
real = paddle.to_tensor(loc.real, real_dtype)
|
||||
imag = paddle.to_tensor(loc.imag, imag_dtype)
|
||||
self.loc = paddle.complex(real, imag)
|
||||
elif isinstance(loc, complex):
|
||||
real = paddle.to_tensor(loc.real, dtype='float32')
|
||||
imag = paddle.to_tensor(loc.imag, dtype='float32')
|
||||
self.loc = paddle.complex(real, imag)
|
||||
else:
|
||||
self.loc = loc
|
||||
|
||||
if isinstance(scale, np.ndarray):
|
||||
self.scale = paddle.to_tensor(scale, dtype=scale.dtype)
|
||||
elif isinstance(scale, float):
|
||||
self.scale = paddle.to_tensor(scale, dtype='float32')
|
||||
else:
|
||||
self.scale = scale
|
||||
|
||||
self.dtype = convert_dtype(self.loc.dtype)
|
||||
else:
|
||||
if self._validate_args(loc, scale):
|
||||
self.loc = loc
|
||||
self.scale = scale
|
||||
self.dtype = convert_dtype(loc.dtype)
|
||||
else:
|
||||
if isinstance(loc, float) and isinstance(scale, float):
|
||||
self.all_arg_is_float = True
|
||||
if isinstance(loc, np.ndarray) and str(loc.dtype) in [
|
||||
'float32',
|
||||
'float64',
|
||||
]:
|
||||
self.dtype = loc.dtype
|
||||
elif isinstance(scale, np.ndarray) and str(scale.dtype) in [
|
||||
'float32',
|
||||
'float64',
|
||||
]:
|
||||
self.dtype = scale.dtype
|
||||
self.loc, self.scale = self._to_tensor(loc, scale)
|
||||
if self.dtype != convert_dtype(self.loc.dtype):
|
||||
self.loc = paddle.cast(self.loc, dtype=self.dtype)
|
||||
self.scale = paddle.cast(self.scale, dtype=self.dtype)
|
||||
super().__init__(self.loc.shape, validate_args=validate_args)
|
||||
if in_dynamic_mode() and self._validate_args_enabled:
|
||||
self._validate_parameters()
|
||||
|
||||
def _validate_parameters(self) -> None:
|
||||
for param, value in (("loc", self.loc), ("scale", self.scale)):
|
||||
constraint_ = self.arg_constraints[param]
|
||||
valid = constraint_.check(value)
|
||||
if not bool(valid.all()):
|
||||
raise ValueError(
|
||||
f"Expected parameter {param} "
|
||||
f"({type(value).__name__} of shape {tuple(value.shape)}) "
|
||||
f"of distribution {self!r} "
|
||||
f"to satisfy the constraint {constraint_!r}, "
|
||||
f"but found invalid values:\n{value}"
|
||||
)
|
||||
|
||||
@property
|
||||
def mean(self) -> Tensor:
|
||||
"""Mean of normal distribution.
|
||||
|
||||
Returns:
|
||||
Tensor: mean value.
|
||||
"""
|
||||
return self.loc
|
||||
|
||||
@property
|
||||
def variance(self) -> Tensor:
|
||||
"""Variance of normal distribution.
|
||||
|
||||
Returns:
|
||||
Tensor: variance value.
|
||||
"""
|
||||
return self.scale.pow(2)
|
||||
|
||||
@param_one_alias(["shape", "sample_shape"])
|
||||
def sample(self, shape: Sequence[int] = [], seed: int = 0) -> Tensor:
|
||||
"""Generate samples of the specified shape.
|
||||
|
||||
Args:
|
||||
shape (Sequence[int], optional): Shape of the generated samples.
|
||||
Alias: ``sample_shape``.
|
||||
seed (int): Python integer number.
|
||||
|
||||
Returns:
|
||||
Tensor, A tensor with prepended dimensions shape.The data type is float32.
|
||||
|
||||
"""
|
||||
if not isinstance(shape, Iterable):
|
||||
raise TypeError('sample shape must be Iterable object.')
|
||||
|
||||
if not in_dynamic_mode():
|
||||
check_type(seed, 'seed', (int), 'sample')
|
||||
|
||||
shape = list(shape)
|
||||
batch_shape = list((self.loc + self.scale).shape)
|
||||
name = self.name + '_sample'
|
||||
if -1 in batch_shape:
|
||||
output_shape = shape + batch_shape
|
||||
fill_shape = list(batch_shape + shape)
|
||||
fill_shape[0] = paddle.shape(self.loc + self.scale)[0].item()
|
||||
zero_tmp = paddle.full(fill_shape, 0.0, self.dtype)
|
||||
zero_tmp_reshape = paddle.reshape(zero_tmp, output_shape)
|
||||
|
||||
zero_tmp_shape = paddle.shape(zero_tmp_reshape)
|
||||
normal_random_tmp = random.gaussian(
|
||||
zero_tmp_shape,
|
||||
mean=(0.0 + 0.0j) if self._complex_gaussian else 0.0,
|
||||
std=1.0,
|
||||
seed=seed,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
output = normal_random_tmp * (zero_tmp_reshape + self.scale)
|
||||
output = paddle.add(output, self.loc, name=name)
|
||||
return output
|
||||
else:
|
||||
output_shape = shape + batch_shape
|
||||
output = random.gaussian(
|
||||
output_shape,
|
||||
mean=(0.0 + 0.0j) if self._complex_gaussian else 0.0,
|
||||
std=1.0,
|
||||
seed=seed,
|
||||
dtype=self.dtype,
|
||||
) * (paddle.zeros(output_shape, dtype=self.dtype) + self.scale)
|
||||
output = paddle.add(output, self.loc, name=name)
|
||||
if self.all_arg_is_float:
|
||||
return paddle.reshape(output, shape, name=name)
|
||||
else:
|
||||
return output
|
||||
|
||||
@param_one_alias(["shape", "sample_shape"])
|
||||
def rsample(self, shape: Sequence[int] = []) -> Tensor:
|
||||
"""Generate reparameterized samples of the specified shape.
|
||||
|
||||
Args:
|
||||
shape (Sequence[int], optional): Shape of the generated samples.
|
||||
Alias: ``sample_shape``.
|
||||
|
||||
Returns:
|
||||
Tensor: A tensor with prepended dimensions shape.The data type is float32.
|
||||
|
||||
"""
|
||||
if not isinstance(shape, Iterable):
|
||||
raise TypeError('sample shape must be Iterable object.')
|
||||
|
||||
shape = self._extend_shape(tuple(shape))
|
||||
eps = paddle.normal(
|
||||
mean=(0.0 + 0.0j) if self._complex_gaussian else 0.0, shape=shape
|
||||
)
|
||||
return self.loc + eps * self.scale
|
||||
|
||||
def entropy(self) -> Tensor:
|
||||
r"""Shannon entropy in nats.
|
||||
|
||||
If non-complex, the entropy is
|
||||
|
||||
.. math::
|
||||
|
||||
entropy(\sigma) = 0.5 \log (2 \pi e \sigma^2)
|
||||
|
||||
If complex gaussian, the entropy is
|
||||
|
||||
.. math::
|
||||
|
||||
entropy(\sigma) = \log (\pi e \sigma^2) + 1
|
||||
|
||||
In the above equation:
|
||||
|
||||
* :math:`scale = \sigma`: is the std.
|
||||
|
||||
Returns:
|
||||
Tensor, Shannon entropy of normal distribution.The data type is float32.
|
||||
|
||||
"""
|
||||
name = self.name + '_entropy'
|
||||
batch_shape = list((self.loc + self.scale).shape)
|
||||
|
||||
if self._complex_gaussian:
|
||||
if -1 in batch_shape:
|
||||
fill_shape = list(batch_shape)
|
||||
fill_shape[0] = paddle.shape(self.loc + self.scale)[0].item()
|
||||
fill_dtype = self.scale.dtype
|
||||
zero_tmp = paddle.full(fill_shape, 0.0, fill_dtype)
|
||||
else:
|
||||
zero_tmp = paddle.full(batch_shape, 0.0, self.scale.dtype)
|
||||
return paddle.add(
|
||||
1.0 + zero_tmp,
|
||||
math.log(math.pi) + 2.0 * paddle.log(self.scale + zero_tmp),
|
||||
name=name,
|
||||
)
|
||||
else:
|
||||
if -1 in batch_shape:
|
||||
fill_shape = list(batch_shape)
|
||||
fill_shape[0] = paddle.shape(self.loc + self.scale)[0].item()
|
||||
fill_dtype = (self.loc + self.scale).dtype
|
||||
zero_tmp = paddle.full(fill_shape, 0.0, fill_dtype)
|
||||
else:
|
||||
zero_tmp = paddle.full(batch_shape, 0.0, self.dtype)
|
||||
return paddle.add(
|
||||
0.5 + zero_tmp,
|
||||
0.5 * math.log(2 * math.pi) + paddle.log(self.scale + zero_tmp),
|
||||
name=name,
|
||||
)
|
||||
|
||||
def log_prob(self, value: Tensor) -> Tensor:
|
||||
"""Log probability density/mass function.
|
||||
|
||||
Args:
|
||||
value (Tensor): The input tensor.
|
||||
|
||||
Returns:
|
||||
Tensor: log probability.The data type is same with :attr:`value` .
|
||||
|
||||
"""
|
||||
name = self.name + '_log_prob'
|
||||
value = self._check_values_dtype_in_probs(self.loc, value)
|
||||
if in_dynamic_mode() and self._validate_args_enabled:
|
||||
self._validate_sample(value)
|
||||
|
||||
var = self.scale * self.scale
|
||||
log_scale = paddle.log(self.scale)
|
||||
if self._complex_gaussian:
|
||||
return paddle.subtract(
|
||||
-1.0 * ((value - self.loc).conj() * (value - self.loc)) / (var),
|
||||
2.0 * log_scale + math.log(math.pi),
|
||||
name=name,
|
||||
)
|
||||
else:
|
||||
return paddle.subtract(
|
||||
-1.0 * ((value - self.loc) * (value - self.loc)) / (2.0 * var),
|
||||
log_scale + math.log(math.sqrt(2.0 * math.pi)),
|
||||
name=name,
|
||||
)
|
||||
|
||||
def probs(self, value: Tensor) -> Tensor:
|
||||
"""Probability density/mass function.
|
||||
|
||||
Args:
|
||||
value (Tensor): The input tensor.
|
||||
|
||||
Returns:
|
||||
Tensor, probability. The data type is same with :attr:`value` .
|
||||
|
||||
"""
|
||||
name = self.name + '_probs'
|
||||
value = self._check_values_dtype_in_probs(self.loc, value)
|
||||
|
||||
var = self.scale * self.scale
|
||||
if self._complex_gaussian:
|
||||
return paddle.divide(
|
||||
paddle.exp(
|
||||
-1.0
|
||||
* ((value - self.loc).conj() * (value - self.loc))
|
||||
/ (var)
|
||||
),
|
||||
(math.pi * var),
|
||||
name=name,
|
||||
)
|
||||
else:
|
||||
return paddle.divide(
|
||||
paddle.exp(
|
||||
-1.0
|
||||
* ((value - self.loc) * (value - self.loc))
|
||||
/ (2.0 * var)
|
||||
),
|
||||
(math.sqrt(2 * math.pi) * self.scale),
|
||||
name=name,
|
||||
)
|
||||
|
||||
def kl_divergence(self, other: Normal) -> Tensor:
|
||||
r"""The KL-divergence between two normal distributions.
|
||||
|
||||
If non-complex, the KL-divergence is
|
||||
|
||||
.. math::
|
||||
|
||||
KL\_divergence(\mu_0, \sigma_0; \mu_1, \sigma_1) = 0.5 (ratio^2 + (\frac{diff}{\sigma_1})^2 - 1 - 2 \ln {ratio})
|
||||
|
||||
If complex gaussian:
|
||||
|
||||
.. math::
|
||||
|
||||
KL\_divergence(\mu_0, \sigma_0; \mu_1, \sigma_1) = ratio^2 + (\frac{diff}{\sigma_1})^2 - 1 - 2 \ln {ratio}
|
||||
|
||||
.. math::
|
||||
|
||||
ratio = \frac{\sigma_0}{\sigma_1}
|
||||
|
||||
.. math::
|
||||
|
||||
diff = \mu_1 - \mu_0
|
||||
|
||||
In the above equation:
|
||||
|
||||
* :math:`loc = \mu_0`: is the mean of current Normal distribution.
|
||||
* :math:`scale = \sigma_0`: is the std of current Normal distribution.
|
||||
* :math:`loc = \mu_1`: is the mean of other Normal distribution.
|
||||
* :math:`scale = \sigma_1`: is the std of other Normal distribution.
|
||||
* :math:`ratio`: is the ratio of scales.
|
||||
* :math:`diff`: is the difference between means.
|
||||
|
||||
Args:
|
||||
other (Normal): instance of Normal.
|
||||
|
||||
Returns:
|
||||
Tensor, kl-divergence between two normal distributions.The data type is float32.
|
||||
|
||||
"""
|
||||
if not in_dynamic_mode():
|
||||
check_type(other, 'other', Normal, 'kl_divergence')
|
||||
|
||||
if self._complex_gaussian != other._complex_gaussian:
|
||||
raise ValueError(
|
||||
"The kl divergence must be computed between two distributions in the same number field."
|
||||
)
|
||||
name = self.name + '_kl_divergence'
|
||||
var_ratio = self.scale / other.scale
|
||||
var_ratio = var_ratio * var_ratio
|
||||
t1 = (self.loc - other.loc) / other.scale
|
||||
if self._complex_gaussian:
|
||||
t1 = t1.conj() * t1
|
||||
return var_ratio + t1 - 1.0 - paddle.log(var_ratio)
|
||||
else:
|
||||
t1 = t1 * t1
|
||||
return paddle.add(
|
||||
0.5 * var_ratio,
|
||||
0.5 * (t1 - 1.0 - paddle.log(var_ratio)),
|
||||
name=name,
|
||||
)
|
||||
Reference in New Issue
Block a user