chore: import upstream snapshot with attribution

python/paddle/distribution/kl.py

@@ -0,0 +1,305 @@
+# 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 functools
+import warnings
+from typing import TYPE_CHECKING, TypeVar
+
+import paddle
+from paddle.distribution.bernoulli import Bernoulli
+from paddle.distribution.beta import Beta
+from paddle.distribution.binomial import Binomial
+from paddle.distribution.categorical import Categorical
+from paddle.distribution.cauchy import Cauchy
+from paddle.distribution.continuous_bernoulli import ContinuousBernoulli
+from paddle.distribution.dirichlet import Dirichlet
+from paddle.distribution.distribution import Distribution
+from paddle.distribution.exponential import Exponential
+from paddle.distribution.exponential_family import ExponentialFamily
+from paddle.distribution.gamma import Gamma
+from paddle.distribution.geometric import Geometric
+from paddle.distribution.laplace import Laplace
+from paddle.distribution.lognormal import LogNormal
+from paddle.distribution.multivariate_normal import MultivariateNormal
+from paddle.distribution.normal import Normal
+from paddle.distribution.poisson import Poisson
+from paddle.distribution.uniform import Uniform
+from paddle.framework import in_dynamic_mode
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle import Tensor
+
+    _T = TypeVar('_T')
+
+__all__ = ["register_kl", "kl_divergence"]
+
+_REGISTER_TABLE = {}
+
+
+def kl_divergence(p: Distribution, q: Distribution) -> Tensor:
+    r"""
+    Kullback-Leibler divergence between distribution p and q.
+
+    .. math::
+
+        KL(p||q) = \int p(x)log\frac{p(x)}{q(x)} \mathrm{d}x
+
+    Args:
+        p (Distribution): ``Distribution`` object. Inherits from the Distribution Base class.
+        q (Distribution): ``Distribution`` object. Inherits from the Distribution Base class.
+
+    Returns:
+        Tensor, Batchwise KL-divergence between distribution p and q.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> p = paddle.distribution.Beta(alpha=0.5, beta=0.5)
+            >>> q = paddle.distribution.Beta(alpha=0.3, beta=0.7)
+
+            >>> print(paddle.distribution.kl_divergence(p, q))
+            Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
+                0.21193528)
+
+    """
+    return _dispatch(type(p), type(q))(p, q)
+
+
+def register_kl(
+    cls_p: type[Distribution], cls_q: type[Distribution]
+) -> Callable[[_T], _T]:
+    """Decorator for register a KL divergence implementation function.
+
+    The ``kl_divergence(p, q)`` function will search concrete implementation
+    functions registered by ``register_kl``, according to multi-dispatch pattern.
+    If an implementation function is found, it will return the result, otherwise,
+    it will raise ``NotImplementError`` exception. Users can register
+    implementation function by the decorator.
+
+    Args:
+        cls_p (type[Distribution]): The Distribution type of Instance p. Subclass derived from ``Distribution``.
+        cls_q (type[Distribution]): The Distribution type of Instance q. Subclass derived from ``Distribution``.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> @paddle.distribution.register_kl(paddle.distribution.Beta, paddle.distribution.Beta)
+            >>> def kl_beta_beta():
+            ...     pass  # insert implementation here
+    """
+    if not issubclass(cls_p, Distribution) or not issubclass(
+        cls_q, Distribution
+    ):
+        raise TypeError('cls_p and cls_q must be subclass of Distribution')
+
+    def decorator(f):
+        _REGISTER_TABLE[cls_p, cls_q] = f
+        return f
+
+    return decorator
+
+
+def _dispatch(cls_p, cls_q):
+    """Multiple dispatch into concrete implement function."""
+
+    # find all matched super class pair of p and q
+    matches = [
+        (super_p, super_q)
+        for super_p, super_q in _REGISTER_TABLE
+        if issubclass(cls_p, super_p) and issubclass(cls_q, super_q)
+    ]
+    if not matches:
+        raise NotImplementedError
+
+    left_p, left_q = min(_Compare(*m) for m in matches).classes
+    right_p, right_q = min(_Compare(*reversed(m)) for m in matches).classes
+
+    if _REGISTER_TABLE[left_p, left_q] is not _REGISTER_TABLE[right_p, right_q]:
+        warnings.warn(
+            f'Ambiguous kl_divergence({cls_p.__name__}, {cls_q.__name__}). Please register_kl({left_p.__name__}, {right_q.__name__})',
+            RuntimeWarning,
+        )
+
+    return _REGISTER_TABLE[left_p, left_q]
+
+
+@functools.total_ordering
+class _Compare:
+    def __init__(self, *classes):
+        self.classes = classes
+
+    def __eq__(self, other):
+        return self.classes == other.classes
+
+    def __le__(self, other):
+        for cls_x, cls_y in zip(self.classes, other.classes):
+            if not issubclass(cls_x, cls_y):
+                return False
+            if cls_x is not cls_y:
+                break
+        return True
+
+
+@register_kl(Bernoulli, Bernoulli)
+def _kl_bernoulli_bernoulli(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(Beta, Beta)
+def _kl_beta_beta(p, q):
+    return (
+        (q.alpha.lgamma() + q.beta.lgamma() + (p.alpha + p.beta).lgamma())
+        - (p.alpha.lgamma() + p.beta.lgamma() + (q.alpha + q.beta).lgamma())
+        + ((p.alpha - q.alpha) * p.alpha.digamma())
+        + ((p.beta - q.beta) * p.beta.digamma())
+        + (
+            ((q.alpha + q.beta) - (p.alpha + p.beta))
+            * (p.alpha + p.beta).digamma()
+        )
+    )
+
+
+@register_kl(Binomial, Binomial)
+def _kl_binomial_binomial(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(Dirichlet, Dirichlet)
+def _kl_dirichlet_dirichlet(p, q):
+    return (
+        (p.concentration.sum(-1).lgamma() - q.concentration.sum(-1).lgamma())
+        - ((p.concentration.lgamma() - q.concentration.lgamma()).sum(-1))
+        + (
+            (
+                (p.concentration - q.concentration)
+                * (
+                    p.concentration.digamma()
+                    - p.concentration.sum(-1).digamma().unsqueeze(-1)
+                )
+            ).sum(-1)
+        )
+    )
+
+
+@register_kl(Categorical, Categorical)
+def _kl_categorical_categorical(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(Cauchy, Cauchy)
+def _kl_cauchy_cauchy(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(ContinuousBernoulli, ContinuousBernoulli)
+def _kl_continuousbernoulli_continuousbernoulli(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(Normal, Normal)
+def _kl_normal_normal(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(MultivariateNormal, MultivariateNormal)
+def _kl_mvn_mvn(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(Uniform, Uniform)
+def _kl_uniform_uniform(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(Laplace, Laplace)
+def _kl_laplace_laplace(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(Geometric, Geometric)
+def _kl_geometric_geometric(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(ExponentialFamily, ExponentialFamily)
+def _kl_expfamily_expfamily(p, q):
+    """Compute kl-divergence using `Bregman divergences <https://www.lix.polytechnique.fr/~nielsen/EntropyEF-ICIP2010.pdf>`_"""
+    if not type(p) == type(q):
+        raise NotImplementedError
+
+    p_natural_params = []
+    for param in p._natural_parameters:
+        param = param.detach()
+        param.stop_gradient = False
+        p_natural_params.append(param)
+
+    q_natural_params = q._natural_parameters
+
+    p_log_norm = p._log_normalizer(*p_natural_params)
+
+    try:
+        if in_dynamic_mode():
+            p_grads = paddle.grad(
+                p_log_norm, p_natural_params, create_graph=True
+            )
+        else:
+            p_grads = paddle.static.gradients(p_log_norm, p_natural_params)
+    except RuntimeError as e:
+        raise TypeError(
+            "Can't compute kl_divergence({cls_p}, {cls_q}) use bregman divergence. Please register_kl({cls_p}, {cls_q}).".format(
+                cls_p=type(p).__name__, cls_q=type(q).__name__
+            )
+        ) from e
+
+    kl = q._log_normalizer(*q_natural_params) - p_log_norm
+    for p_param, q_param, p_grad in zip(
+        p_natural_params, q_natural_params, p_grads
+    ):
+        term = (q_param - p_param) * p_grad
+        kl -= _sum_rightmost(term, len(q.event_shape))
+
+    return kl
+
+
+@register_kl(Exponential, Exponential)
+def _kl_exponential_exponential(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(Gamma, Gamma)
+def _kl_gamma_gamma(p, q):
+    return p.kl_divergence(q)
+
+
+@register_kl(LogNormal, LogNormal)
+def _kl_lognormal_lognormal(p, q):
+    return p._base.kl_divergence(q._base)
+
+
+@register_kl(Poisson, Poisson)
+def _kl_poisson_poisson(p, q):
+    return p.kl_divergence(q)
+
+
+def _sum_rightmost(value, n):
+    return value.sum(list(range(-n, 0))) if n > 0 else value