Table of Contents
Shortcuts

# Source code for torch.distributions.kl

    import math
    import warnings
    from functools import total_ordering
    from typing import Callable, Dict, Tuple, Type

    import torch
    from torch import inf

    from .bernoulli import Bernoulli
    from .beta import Beta
    from .binomial import Binomial
    from .categorical import Categorical
    from .cauchy import Cauchy
    from .continuous_bernoulli import ContinuousBernoulli
    from .dirichlet import Dirichlet
    from .distribution import Distribution
    from .exp_family import ExponentialFamily
    from .exponential import Exponential
    from .gamma import Gamma
    from .geometric import Geometric
    from .gumbel import Gumbel
    from .half_normal import HalfNormal
    from .independent import Independent
    from .laplace import Laplace
    from .lowrank_multivariate_normal import (
        _batch_lowrank_logdet,
        _batch_lowrank_mahalanobis,
        LowRankMultivariateNormal,
    )
    from .multivariate_normal import _batch_mahalanobis, MultivariateNormal
    from .normal import Normal
    from .one_hot_categorical import OneHotCategorical
    from .pareto import Pareto
    from .poisson import Poisson
    from .transformed_distribution import TransformedDistribution
    from .uniform import Uniform
    from .utils import _sum_rightmost, euler_constant as _euler_gamma

    _KL_REGISTRY = (
        {}
    )  # Source of truth mapping a few general (type, type) pairs to functions.
    _KL_MEMOIZE: Dict[
        Tuple[Type, Type], Callable
    ] = {}  # Memoized version mapping many specific (type, type) pairs to functions.

    __all__ = ["register_kl", "kl_divergence"]

    def register_kl(type_p, type_q):
        """
        Decorator to register a pairwise function with :meth:`kl_divergence`.
        Usage::

            @register_kl(Normal, Normal)
            def kl_normal_normal(p, q):
                # insert implementation here

        Lookup returns the most specific (type,type) match ordered by subclass. If
        the match is ambiguous, a `RuntimeWarning` is raised. For example to
        resolve the ambiguous situation::

            @register_kl(BaseP, DerivedQ)
            def kl_version1(p, q): ...
            @register_kl(DerivedP, BaseQ)
            def kl_version2(p, q): ...

        you should register a third most-specific implementation, e.g.::

            register_kl(DerivedP, DerivedQ)(kl_version1)  # Break the tie.

        Args:
            type_p (type): A subclass of :class:`~torch.distributions.Distribution`.
            type_q (type): A subclass of :class:`~torch.distributions.Distribution`.
        """
        if not isinstance(type_p, type) and issubclass(type_p, Distribution):
            raise TypeError(
                f"Expected type_p to be a Distribution subclass but got {type_p}"
            )
        if not isinstance(type_q, type) and issubclass(type_q, Distribution):
            raise TypeError(
                f"Expected type_q to be a Distribution subclass but got {type_q}"
            )

        def decorator(fun):
            _KL_REGISTRY[type_p, type_q] = fun
            _KL_MEMOIZE.clear()  # reset since lookup order may have changed
            return fun

        return decorator

    @total_ordering
    class _Match:
        __slots__ = ["types"]

        def __init__(self, *types):
            self.types = types

        def __eq__(self, other):
            return self.types == other.types

        def __le__(self, other):
            for x, y in zip(self.types, other.types):
                if not issubclass(x, y):
                    return False
                if x is not y:
                    break
            return True

    def _dispatch_kl(type_p, type_q):
        """
        Find the most specific approximate match, assuming single inheritance.
        """
        matches = [
            (super_p, super_q)
            for super_p, super_q in _KL_REGISTRY
            if issubclass(type_p, super_p) and issubclass(type_q, super_q)
        ]
        if not matches:
            return NotImplemented
        # Check that the left- and right- lexicographic orders agree.
        # mypy isn't smart enough to know that _Match implements __lt__
        # see: https://github.com/python/typing/issues/760#issuecomment-710670503
        left_p, left_q = min(_Match(*m) for m in matches).types  # type: ignore[type-var]
        right_q, right_p = min(_Match(*reversed(m)) for m in matches).types  # type: ignore[type-var]
        left_fun = _KL_REGISTRY[left_p, left_q]
        right_fun = _KL_REGISTRY[right_p, right_q]
        if left_fun is not right_fun:
            warnings.warn(
                "Ambiguous kl_divergence({}, {}). Please register_kl({}, {})".format(
                    type_p.__name__, type_q.__name__, left_p.__name__, right_q.__name__
                ),
                RuntimeWarning,
            )
        return left_fun

    def _infinite_like(tensor):
        """
        Helper function for obtaining infinite KL Divergence throughout
        """
        return torch.full_like(tensor, inf)

    def _x_log_x(tensor):
        """
        Utility function for calculating x log x
        """
        return tensor * tensor.log()

    def _batch_trace_XXT(bmat):
        """
        Utility function for calculating the trace of XX^{T} with X having arbitrary trailing batch dimensions
        """
        n = bmat.size(-1)
        m = bmat.size(-2)
        flat_trace = bmat.reshape(-1, m * n).pow(2).sum(-1)
        return flat_trace.reshape(bmat.shape[:-2])

    def kl_divergence(p: Distribution, q: Distribution) -> torch.Tensor:
        r"""
        Compute Kullback-Leibler divergence :math:`KL(p \| q)` between two distributions.

        .. math::

            KL(p \| q) = \int p(x) \log\frac {p(x)} {q(x)} \,dx

        Args:
            p (Distribution): A :class:`~torch.distributions.Distribution` object.
            q (Distribution): A :class:`~torch.distributions.Distribution` object.

        Returns:
            Tensor: A batch of KL divergences of shape `batch_shape`.

        Raises:
            NotImplementedError: If the distribution types have not been registered via
                :meth:`register_kl`.
        """
        try:
            fun = _KL_MEMOIZE[type(p), type(q)]
        except KeyError:
            fun = _dispatch_kl(type(p), type(q))
            _KL_MEMOIZE[type(p), type(q)] = fun
        if fun is NotImplemented:
            raise NotImplementedError(
                f"No KL(p || q) is implemented for p type {p.__class__.__name__} and q type {q.__class__.__name__}"
            )
        return fun(p, q)

    ################################################################################
    # KL Divergence Implementations
    ################################################################################

    # Same distributions

    @register_kl(Bernoulli, Bernoulli)
    def _kl_bernoulli_bernoulli(p, q):
        t1 = p.probs * (
            torch.nn.functional.softplus(-q.logits)
            - torch.nn.functional.softplus(-p.logits)
        )
        t1[q.probs == 0] = inf
        t1[p.probs == 0] = 0
        t2 = (1 - p.probs) * (
            torch.nn.functional.softplus(q.logits) - torch.nn.functional.softplus(p.logits)
        )
        t2[q.probs == 1] = inf
        t2[p.probs == 1] = 0
        return t1 + t2

    @register_kl(Beta, Beta)
    def _kl_beta_beta(p, q):
        sum_params_p = p.concentration1 + p.concentration0
        sum_params_q = q.concentration1 + q.concentration0
        t1 = q.concentration1.lgamma() + q.concentration0.lgamma() + (sum_params_p).lgamma()
        t2 = p.concentration1.lgamma() + p.concentration0.lgamma() + (sum_params_q).lgamma()
        t3 = (p.concentration1 - q.concentration1) * torch.digamma(p.concentration1)
        t4 = (p.concentration0 - q.concentration0) * torch.digamma(p.concentration0)
        t5 = (sum_params_q - sum_params_p) * torch.digamma(sum_params_p)
        return t1 - t2 + t3 + t4 + t5

    @register_kl(Binomial, Binomial)
    def _kl_binomial_binomial(p, q):
        # from https://math.stackexchange.com/questions/2214993/
        # kullback-leibler-divergence-for-binomial-distributions-p-and-q
        if (p.total_count < q.total_count).any():
            raise NotImplementedError(
                "KL between Binomials where q.total_count > p.total_count is not implemented"
            )
        kl = p.total_count * (
            p.probs * (p.logits - q.logits) + (-p.probs).log1p() - (-q.probs).log1p()
        )
        inf_idxs = p.total_count > q.total_count
        kl[inf_idxs] = _infinite_like(kl[inf_idxs])
        return kl

    @register_kl(Categorical, Categorical)
    def _kl_categorical_categorical(p, q):
        t = p.probs * (p.logits - q.logits)
        t[(q.probs == 0).expand_as(t)] = inf
        t[(p.probs == 0).expand_as(t)] = 0
        return t.sum(-1)

    @register_kl(ContinuousBernoulli, ContinuousBernoulli)
    def _kl_continuous_bernoulli_continuous_bernoulli(p, q):
        t1 = p.mean * (p.logits - q.logits)
        t2 = p._cont_bern_log_norm() + torch.log1p(-p.probs)
        t3 = -q._cont_bern_log_norm() - torch.log1p(-q.probs)
        return t1 + t2 + t3

    @register_kl(Dirichlet, Dirichlet)
    def _kl_dirichlet_dirichlet(p, q):
        # From http://bariskurt.com/kullback-leibler-divergence-between-two-dirichlet-and-beta-distributions/
        sum_p_concentration = p.concentration.sum(-1)
        sum_q_concentration = q.concentration.sum(-1)
        t1 = sum_p_concentration.lgamma() - sum_q_concentration.lgamma()
        t2 = (p.concentration.lgamma() - q.concentration.lgamma()).sum(-1)
        t3 = p.concentration - q.concentration
        t4 = p.concentration.digamma() - sum_p_concentration.digamma().unsqueeze(-1)
        return t1 - t2 + (t3 * t4).sum(-1)

    @register_kl(Exponential, Exponential)
    def _kl_exponential_exponential(p, q):
        rate_ratio = q.rate / p.rate
        t1 = -rate_ratio.log()
        return t1 + rate_ratio - 1

    @register_kl(ExponentialFamily, ExponentialFamily)
    def _kl_expfamily_expfamily(p, q):
        if not type(p) == type(q):
            raise NotImplementedError(
                "The cross KL-divergence between different exponential families cannot \
                                be computed using Bregman divergences"
            )
        p_nparams = [np.detach().requires_grad_() for np in p._natural_params]
        q_nparams = q._natural_params
        lg_normal = p._log_normalizer(*p_nparams)
        gradients = torch.autograd.grad(lg_normal.sum(), p_nparams, create_graph=True)
        result = q._log_normalizer(*q_nparams) - lg_normal
        for pnp, qnp, g in zip(p_nparams, q_nparams, gradients):
            term = (qnp - pnp) * g
            result -= _sum_rightmost(term, len(q.event_shape))
        return result

    @register_kl(Gamma, Gamma)
    def _kl_gamma_gamma(p, q):
        t1 = q.concentration * (p.rate / q.rate).log()
        t2 = torch.lgamma(q.concentration) - torch.lgamma(p.concentration)
        t3 = (p.concentration - q.concentration) * torch.digamma(p.concentration)
        t4 = (q.rate - p.rate) * (p.concentration / p.rate)
        return t1 + t2 + t3 + t4

    @register_kl(Gumbel, Gumbel)
    def _kl_gumbel_gumbel(p, q):
        ct1 = p.scale / q.scale
        ct2 = q.loc / q.scale
        ct3 = p.loc / q.scale
        t1 = -ct1.log() - ct2 + ct3
        t2 = ct1 * _euler_gamma
        t3 = torch.exp(ct2 + (1 + ct1).lgamma() - ct3)
        return t1 + t2 + t3 - (1 + _euler_gamma)

    @register_kl(Geometric, Geometric)
    def _kl_geometric_geometric(p, q):
        return -p.entropy() - torch.log1p(-q.probs) / p.probs - q.logits

    @register_kl(HalfNormal, HalfNormal)
    def _kl_halfnormal_halfnormal(p, q):
        return _kl_normal_normal(p.base_dist, q.base_dist)

    @register_kl(Laplace, Laplace)
    def _kl_laplace_laplace(p, q):
        # From http://www.mast.queensu.ca/~communications/Papers/gil-msc11.pdf
        scale_ratio = p.scale / q.scale
        loc_abs_diff = (p.loc - q.loc).abs()
        t1 = -scale_ratio.log()
        t2 = loc_abs_diff / q.scale
        t3 = scale_ratio * torch.exp(-loc_abs_diff / p.scale)
        return t1 + t2 + t3 - 1

    @register_kl(LowRankMultivariateNormal, LowRankMultivariateNormal)
    def _kl_lowrankmultivariatenormal_lowrankmultivariatenormal(p, q):
        if p.event_shape != q.event_shape:
            raise ValueError(
                "KL-divergence between two Low Rank Multivariate Normals with\
                              different event shapes cannot be computed"
            )

        term1 = _batch_lowrank_logdet(
            q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag, q._capacitance_tril
        ) - _batch_lowrank_logdet(
            p._unbroadcasted_cov_factor, p._unbroadcasted_cov_diag, p._capacitance_tril
        )
        term3 = _batch_lowrank_mahalanobis(
            q._unbroadcasted_cov_factor,
            q._unbroadcasted_cov_diag,
            q.loc - p.loc,
            q._capacitance_tril,
        )
        # Expands term2 according to
        # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ (pW @ pW.T + pD)
        #                  = [inv(qD) - A.T @ A] @ (pD + pW @ pW.T)
        qWt_qDinv = q._unbroadcasted_cov_factor.mT / q._unbroadcasted_cov_diag.unsqueeze(-2)
        A = torch.linalg.solve_triangular(q._capacitance_tril, qWt_qDinv, upper=False)
        term21 = (p._unbroadcasted_cov_diag / q._unbroadcasted_cov_diag).sum(-1)
        term22 = _batch_trace_XXT(
            p._unbroadcasted_cov_factor * q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1)
        )
        term23 = _batch_trace_XXT(A * p._unbroadcasted_cov_diag.sqrt().unsqueeze(-2))
        term24 = _batch_trace_XXT(A.matmul(p._unbroadcasted_cov_factor))
        term2 = term21 + term22 - term23 - term24
        return 0.5 * (term1 + term2 + term3 - p.event_shape[0])

    @register_kl(MultivariateNormal, LowRankMultivariateNormal)
    def _kl_multivariatenormal_lowrankmultivariatenormal(p, q):
        if p.event_shape != q.event_shape:
            raise ValueError(
                "KL-divergence between two (Low Rank) Multivariate Normals with\
                              different event shapes cannot be computed"
            )

        term1 = _batch_lowrank_logdet(
            q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag, q._capacitance_tril
        ) - 2 * p._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1)
        term3 = _batch_lowrank_mahalanobis(
            q._unbroadcasted_cov_factor,
            q._unbroadcasted_cov_diag,
            q.loc - p.loc,
            q._capacitance_tril,
        )
        # Expands term2 according to
        # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ p_tril @ p_tril.T
        #                  = [inv(qD) - A.T @ A] @ p_tril @ p_tril.T
        qWt_qDinv = q._unbroadcasted_cov_factor.mT / q._unbroadcasted_cov_diag.unsqueeze(-2)
        A = torch.linalg.solve_triangular(q._capacitance_tril, qWt_qDinv, upper=False)
        term21 = _batch_trace_XXT(
            p._unbroadcasted_scale_tril * q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1)
        )
        term22 = _batch_trace_XXT(A.matmul(p._unbroadcasted_scale_tril))
        term2 = term21 - term22
        return 0.5 * (term1 + term2 + term3 - p.event_shape[0])

    @register_kl(LowRankMultivariateNormal, MultivariateNormal)
    def _kl_lowrankmultivariatenormal_multivariatenormal(p, q):
        if p.event_shape != q.event_shape:
            raise ValueError(
                "KL-divergence between two (Low Rank) Multivariate Normals with\
                              different event shapes cannot be computed"
            )

        term1 = 2 * q._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(
            -1
        ) - _batch_lowrank_logdet(
            p._unbroadcasted_cov_factor, p._unbroadcasted_cov_diag, p._capacitance_tril
        )
        term3 = _batch_mahalanobis(q._unbroadcasted_scale_tril, (q.loc - p.loc))
        # Expands term2 according to
        # inv(qcov) @ pcov = inv(q_tril @ q_tril.T) @ (pW @ pW.T + pD)
        combined_batch_shape = torch._C._infer_size(
            q._unbroadcasted_scale_tril.shape[:-2], p._unbroadcasted_cov_factor.shape[:-2]
        )
        n = p.event_shape[0]
        q_scale_tril = q._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
        p_cov_factor = p._unbroadcasted_cov_factor.expand(
            combined_batch_shape + (n, p.cov_factor.size(-1))
        )
        p_cov_diag = torch.diag_embed(p._unbroadcasted_cov_diag.sqrt()).expand(
            combined_batch_shape + (n, n)
        )
        term21 = _batch_trace_XXT(
            torch.linalg.solve_triangular(q_scale_tril, p_cov_factor, upper=False)
        )
        term22 = _batch_trace_XXT(
            torch.linalg.solve_triangular(q_scale_tril, p_cov_diag, upper=False)
        )
        term2 = term21 + term22
        return 0.5 * (term1 + term2 + term3 - p.event_shape[0])

    @register_kl(MultivariateNormal, MultivariateNormal)
    def _kl_multivariatenormal_multivariatenormal(p, q):
        # From https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Kullback%E2%80%93Leibler_divergence
        if p.event_shape != q.event_shape:
            raise ValueError(
                "KL-divergence between two Multivariate Normals with\
                              different event shapes cannot be computed"
            )

        half_term1 = q._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(
            -1
        ) - p._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1)
        combined_batch_shape = torch._C._infer_size(
            q._unbroadcasted_scale_tril.shape[:-2], p._unbroadcasted_scale_tril.shape[:-2]
        )
        n = p.event_shape[0]
        q_scale_tril = q._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
        p_scale_tril = p._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
        term2 = _batch_trace_XXT(
            torch.linalg.solve_triangular(q_scale_tril, p_scale_tril, upper=False)
        )
        term3 = _batch_mahalanobis(q._unbroadcasted_scale_tril, (q.loc - p.loc))
        return half_term1 + 0.5 * (term2 + term3 - n)

    @register_kl(Normal, Normal)
    def _kl_normal_normal(p, q):
        var_ratio = (p.scale / q.scale).pow(2)
        t1 = ((p.loc - q.loc) / q.scale).pow(2)
        return 0.5 * (var_ratio + t1 - 1 - var_ratio.log())

    @register_kl(OneHotCategorical, OneHotCategorical)
    def _kl_onehotcategorical_onehotcategorical(p, q):
        return _kl_categorical_categorical(p._categorical, q._categorical)

    @register_kl(Pareto, Pareto)
    def _kl_pareto_pareto(p, q):
        # From http://www.mast.queensu.ca/~communications/Papers/gil-msc11.pdf
        scale_ratio = p.scale / q.scale
        alpha_ratio = q.alpha / p.alpha
        t1 = q.alpha * scale_ratio.log()
        t2 = -alpha_ratio.log()
        result = t1 + t2 + alpha_ratio - 1
        result[p.support.lower_bound < q.support.lower_bound] = inf
        return result

    @register_kl(Poisson, Poisson)
    def _kl_poisson_poisson(p, q):
        return p.rate * (p.rate.log() - q.rate.log()) - (p.rate - q.rate)

    @register_kl(TransformedDistribution, TransformedDistribution)
    def _kl_transformed_transformed(p, q):
        if p.transforms != q.transforms:
            raise NotImplementedError
        if p.event_shape != q.event_shape:
            raise NotImplementedError
        return kl_divergence(p.base_dist, q.base_dist)

    @register_kl(Uniform, Uniform)
    def _kl_uniform_uniform(p, q):
        result = ((q.high - q.low) / (p.high - p.low)).log()
        result[(q.low > p.low) | (q.high < p.high)] = inf
        return result

    # Different distributions
    @register_kl(Bernoulli, Poisson)
    def _kl_bernoulli_poisson(p, q):
        return -p.entropy() - (p.probs * q.rate.log() - q.rate)

    @register_kl(Beta, ContinuousBernoulli)
    def _kl_beta_continuous_bernoulli(p, q):
        return (
            -p.entropy()
            - p.mean * q.logits
            - torch.log1p(-q.probs)
            - q._cont_bern_log_norm()
        )

    @register_kl(Beta, Pareto)
    def _kl_beta_infinity(p, q):
        return _infinite_like(p.concentration1)

    @register_kl(Beta, Exponential)
    def _kl_beta_exponential(p, q):
        return (
            -p.entropy()
            - q.rate.log()
            + q.rate * (p.concentration1 / (p.concentration1 + p.concentration0))
        )

    @register_kl(Beta, Gamma)
    def _kl_beta_gamma(p, q):
        t1 = -p.entropy()
        t2 = q.concentration.lgamma() - q.concentration * q.rate.log()
        t3 = (q.concentration - 1) * (
            p.concentration1.digamma() - (p.concentration1 + p.concentration0).digamma()
        )
        t4 = q.rate * p.concentration1 / (p.concentration1 + p.concentration0)
        return t1 + t2 - t3 + t4

    # TODO: Add Beta-Laplace KL Divergence

    @register_kl(Beta, Normal)
    def _kl_beta_normal(p, q):
        E_beta = p.concentration1 / (p.concentration1 + p.concentration0)
        var_normal = q.scale.pow(2)
        t1 = -p.entropy()
        t2 = 0.5 * (var_normal * 2 * math.pi).log()
        t3 = (
            E_beta * (1 - E_beta) / (p.concentration1 + p.concentration0 + 1)
            + E_beta.pow(2)
        ) * 0.5
        t4 = q.loc * E_beta
        t5 = q.loc.pow(2) * 0.5
        return t1 + t2 + (t3 - t4 + t5) / var_normal

    @register_kl(Beta, Uniform)
    def _kl_beta_uniform(p, q):
        result = -p.entropy() + (q.high - q.low).log()
        result[(q.low > p.support.lower_bound) | (q.high < p.support.upper_bound)] = inf
        return result

    # Note that the KL between a ContinuousBernoulli and Beta has no closed form

    @register_kl(ContinuousBernoulli, Pareto)
    def _kl_continuous_bernoulli_infinity(p, q):
        return _infinite_like(p.probs)

    @register_kl(ContinuousBernoulli, Exponential)
    def _kl_continuous_bernoulli_exponential(p, q):
        return -p.entropy() - torch.log(q.rate) + q.rate * p.mean

    # Note that the KL between a ContinuousBernoulli and Gamma has no closed form
    # TODO: Add ContinuousBernoulli-Laplace KL Divergence

    @register_kl(ContinuousBernoulli, Normal)
    def _kl_continuous_bernoulli_normal(p, q):
        t1 = -p.entropy()
        t2 = 0.5 * (math.log(2.0 * math.pi) + torch.square(q.loc / q.scale)) + torch.log(
            q.scale
        )
        t3 = (p.variance + torch.square(p.mean) - 2.0 * q.loc * p.mean) / (
            2.0 * torch.square(q.scale)
        )
        return t1 + t2 + t3

    @register_kl(ContinuousBernoulli, Uniform)
    def _kl_continuous_bernoulli_uniform(p, q):
        result = -p.entropy() + (q.high - q.low).log()
        return torch.where(
            torch.max(
                torch.ge(q.low, p.support.lower_bound),
                torch.le(q.high, p.support.upper_bound),
            ),
            torch.ones_like(result) * inf,
            result,
        )

    @register_kl(Exponential, Beta)
    @register_kl(Exponential, ContinuousBernoulli)
    @register_kl(Exponential, Pareto)
    @register_kl(Exponential, Uniform)
    def _kl_exponential_infinity(p, q):
        return _infinite_like(p.rate)

    @register_kl(Exponential, Gamma)
    def _kl_exponential_gamma(p, q):
        ratio = q.rate / p.rate
        t1 = -q.concentration * torch.log(ratio)
        return (
            t1
            + ratio
            + q.concentration.lgamma()
            + q.concentration * _euler_gamma
            - (1 + _euler_gamma)
        )

    @register_kl(Exponential, Gumbel)
    def _kl_exponential_gumbel(p, q):
        scale_rate_prod = p.rate * q.scale
        loc_scale_ratio = q.loc / q.scale
        t1 = scale_rate_prod.log() - 1
        t2 = torch.exp(loc_scale_ratio) * scale_rate_prod / (scale_rate_prod + 1)
        t3 = scale_rate_prod.reciprocal()
        return t1 - loc_scale_ratio + t2 + t3

    # TODO: Add Exponential-Laplace KL Divergence

    @register_kl(Exponential, Normal)
    def _kl_exponential_normal(p, q):
        var_normal = q.scale.pow(2)
        rate_sqr = p.rate.pow(2)
        t1 = 0.5 * torch.log(rate_sqr * var_normal * 2 * math.pi)
        t2 = rate_sqr.reciprocal()
        t3 = q.loc / p.rate
        t4 = q.loc.pow(2) * 0.5
        return t1 - 1 + (t2 - t3 + t4) / var_normal

    @register_kl(Gamma, Beta)
    @register_kl(Gamma, ContinuousBernoulli)
    @register_kl(Gamma, Pareto)
    @register_kl(Gamma, Uniform)
    def _kl_gamma_infinity(p, q):
        return _infinite_like(p.concentration)

    @register_kl(Gamma, Exponential)
    def _kl_gamma_exponential(p, q):
        return -p.entropy() - q.rate.log() + q.rate * p.concentration / p.rate

    @register_kl(Gamma, Gumbel)
    def _kl_gamma_gumbel(p, q):
        beta_scale_prod = p.rate * q.scale
        loc_scale_ratio = q.loc / q.scale
        t1 = (
            (p.concentration - 1) * p.concentration.digamma()
            - p.concentration.lgamma()
            - p.concentration
        )
        t2 = beta_scale_prod.log() + p.concentration / beta_scale_prod
        t3 = (
            torch.exp(loc_scale_ratio)
            * (1 + beta_scale_prod.reciprocal()).pow(-p.concentration)
            - loc_scale_ratio
        )
        return t1 + t2 + t3

    # TODO: Add Gamma-Laplace KL Divergence

    @register_kl(Gamma, Normal)
    def _kl_gamma_normal(p, q):
        var_normal = q.scale.pow(2)
        beta_sqr = p.rate.pow(2)
        t1 = (
            0.5 * torch.log(beta_sqr * var_normal * 2 * math.pi)
            - p.concentration
            - p.concentration.lgamma()
        )
        t2 = 0.5 * (p.concentration.pow(2) + p.concentration) / beta_sqr
        t3 = q.loc * p.concentration / p.rate
        t4 = 0.5 * q.loc.pow(2)
        return (
            t1
            + (p.concentration - 1) * p.concentration.digamma()
            + (t2 - t3 + t4) / var_normal
        )

    @register_kl(Gumbel, Beta)
    @register_kl(Gumbel, ContinuousBernoulli)
    @register_kl(Gumbel, Exponential)
    @register_kl(Gumbel, Gamma)
    @register_kl(Gumbel, Pareto)
    @register_kl(Gumbel, Uniform)
    def _kl_gumbel_infinity(p, q):
        return _infinite_like(p.loc)

    # TODO: Add Gumbel-Laplace KL Divergence

    @register_kl(Gumbel, Normal)
    def _kl_gumbel_normal(p, q):
        param_ratio = p.scale / q.scale
        t1 = (param_ratio / math.sqrt(2 * math.pi)).log()
        t2 = (math.pi * param_ratio * 0.5).pow(2) / 3
        t3 = ((p.loc + p.scale * _euler_gamma - q.loc) / q.scale).pow(2) * 0.5
        return -t1 + t2 + t3 - (_euler_gamma + 1)

    @register_kl(Laplace, Beta)
    @register_kl(Laplace, ContinuousBernoulli)
    @register_kl(Laplace, Exponential)
    @register_kl(Laplace, Gamma)
    @register_kl(Laplace, Pareto)
    @register_kl(Laplace, Uniform)
    def _kl_laplace_infinity(p, q):
        return _infinite_like(p.loc)

    @register_kl(Laplace, Normal)
    def _kl_laplace_normal(p, q):
        var_normal = q.scale.pow(2)
        scale_sqr_var_ratio = p.scale.pow(2) / var_normal
        t1 = 0.5 * torch.log(2 * scale_sqr_var_ratio / math.pi)
        t2 = 0.5 * p.loc.pow(2)
        t3 = p.loc * q.loc
        t4 = 0.5 * q.loc.pow(2)
        return -t1 + scale_sqr_var_ratio + (t2 - t3 + t4) / var_normal - 1

    @register_kl(Normal, Beta)
    @register_kl(Normal, ContinuousBernoulli)
    @register_kl(Normal, Exponential)
    @register_kl(Normal, Gamma)
    @register_kl(Normal, Pareto)
    @register_kl(Normal, Uniform)
    def _kl_normal_infinity(p, q):
        return _infinite_like(p.loc)

    @register_kl(Normal, Gumbel)
    def _kl_normal_gumbel(p, q):
        mean_scale_ratio = p.loc / q.scale
        var_scale_sqr_ratio = (p.scale / q.scale).pow(2)
        loc_scale_ratio = q.loc / q.scale
        t1 = var_scale_sqr_ratio.log() * 0.5
        t2 = mean_scale_ratio - loc_scale_ratio
        t3 = torch.exp(-mean_scale_ratio + 0.5 * var_scale_sqr_ratio + loc_scale_ratio)
        return -t1 + t2 + t3 - (0.5 * (1 + math.log(2 * math.pi)))

    @register_kl(Normal, Laplace)
    def _kl_normal_laplace(p, q):
        loc_diff = p.loc - q.loc
        scale_ratio = p.scale / q.scale
        loc_diff_scale_ratio = loc_diff / p.scale
        t1 = torch.log(scale_ratio)
        t2 = (
            math.sqrt(2 / math.pi) * p.scale * torch.exp(-0.5 * loc_diff_scale_ratio.pow(2))
        )
        t3 = loc_diff * torch.erf(math.sqrt(0.5) * loc_diff_scale_ratio)
        return -t1 + (t2 + t3) / q.scale - (0.5 * (1 + math.log(0.5 * math.pi)))

    @register_kl(Pareto, Beta)
    @register_kl(Pareto, ContinuousBernoulli)
    @register_kl(Pareto, Uniform)
    def _kl_pareto_infinity(p, q):
        return _infinite_like(p.scale)

    @register_kl(Pareto, Exponential)
    def _kl_pareto_exponential(p, q):
        scale_rate_prod = p.scale * q.rate
        t1 = (p.alpha / scale_rate_prod).log()
        t2 = p.alpha.reciprocal()
        t3 = p.alpha * scale_rate_prod / (p.alpha - 1)
        result = t1 - t2 + t3 - 1
        result[p.alpha <= 1] = inf
        return result

    @register_kl(Pareto, Gamma)
    def _kl_pareto_gamma(p, q):
        common_term = p.scale.log() + p.alpha.reciprocal()
        t1 = p.alpha.log() - common_term
        t2 = q.concentration.lgamma() - q.concentration * q.rate.log()
        t3 = (1 - q.concentration) * common_term
        t4 = q.rate * p.alpha * p.scale / (p.alpha - 1)
        result = t1 + t2 + t3 + t4 - 1
        result[p.alpha <= 1] = inf
        return result

    # TODO: Add Pareto-Laplace KL Divergence

    @register_kl(Pareto, Normal)
    def _kl_pareto_normal(p, q):
        var_normal = 2 * q.scale.pow(2)
        common_term = p.scale / (p.alpha - 1)
        t1 = (math.sqrt(2 * math.pi) * q.scale * p.alpha / p.scale).log()
        t2 = p.alpha.reciprocal()
        t3 = p.alpha * common_term.pow(2) / (p.alpha - 2)
        t4 = (p.alpha * common_term - q.loc).pow(2)
        result = t1 - t2 + (t3 + t4) / var_normal - 1
        result[p.alpha <= 2] = inf
        return result

    @register_kl(Poisson, Bernoulli)
    @register_kl(Poisson, Binomial)
    def _kl_poisson_infinity(p, q):
        return _infinite_like(p.rate)

    @register_kl(Uniform, Beta)
    def _kl_uniform_beta(p, q):
        common_term = p.high - p.low
        t1 = torch.log(common_term)
        t2 = (
            (q.concentration1 - 1)
            * (_x_log_x(p.high) - _x_log_x(p.low) - common_term)
            / common_term
        )
        t3 = (
            (q.concentration0 - 1)
            * (_x_log_x(1 - p.high) - _x_log_x(1 - p.low) + common_term)
            / common_term
        )
        t4 = (
            q.concentration1.lgamma()
            + q.concentration0.lgamma()
            - (q.concentration1 + q.concentration0).lgamma()
        )
        result = t3 + t4 - t1 - t2
        result[(p.high > q.support.upper_bound) | (p.low < q.support.lower_bound)] = inf
        return result

    @register_kl(Uniform, ContinuousBernoulli)
    def _kl_uniform_continuous_bernoulli(p, q):
        result = (
            -p.entropy()
            - p.mean * q.logits
            - torch.log1p(-q.probs)
            - q._cont_bern_log_norm()
        )
        return torch.where(
            torch.max(
                torch.ge(p.high, q.support.upper_bound),
                torch.le(p.low, q.support.lower_bound),
            ),
            torch.ones_like(result) * inf,
            result,
        )

    @register_kl(Uniform, Exponential)
    def _kl_uniform_exponetial(p, q):
        result = q.rate * (p.high + p.low) / 2 - ((p.high - p.low) * q.rate).log()
        result[p.low < q.support.lower_bound] = inf
        return result

    @register_kl(Uniform, Gamma)
    def _kl_uniform_gamma(p, q):
        common_term = p.high - p.low
        t1 = common_term.log()
        t2 = q.concentration.lgamma() - q.concentration * q.rate.log()
        t3 = (
            (1 - q.concentration)
            * (_x_log_x(p.high) - _x_log_x(p.low) - common_term)
            / common_term
        )
        t4 = q.rate * (p.high + p.low) / 2
        result = -t1 + t2 + t3 + t4
        result[p.low < q.support.lower_bound] = inf
        return result

    @register_kl(Uniform, Gumbel)
    def _kl_uniform_gumbel(p, q):
        common_term = q.scale / (p.high - p.low)
        high_loc_diff = (p.high - q.loc) / q.scale
        low_loc_diff = (p.low - q.loc) / q.scale
        t1 = common_term.log() + 0.5 * (high_loc_diff + low_loc_diff)
        t2 = common_term * (torch.exp(-high_loc_diff) - torch.exp(-low_loc_diff))
        return t1 - t2

    # TODO: Uniform-Laplace KL Divergence

    @register_kl(Uniform, Normal)
    def _kl_uniform_normal(p, q):
        common_term = p.high - p.low
        t1 = (math.sqrt(math.pi * 2) * q.scale / common_term).log()
        t2 = (common_term).pow(2) / 12
        t3 = ((p.high + p.low - 2 * q.loc) / 2).pow(2)
        return t1 + 0.5 * (t2 + t3) / q.scale.pow(2)

    @register_kl(Uniform, Pareto)
    def _kl_uniform_pareto(p, q):
        support_uniform = p.high - p.low
        t1 = (q.alpha * q.scale.pow(q.alpha) * (support_uniform)).log()
        t2 = (_x_log_x(p.high) - _x_log_x(p.low) - support_uniform) / support_uniform
        result = t2 * (q.alpha + 1) - t1
        result[p.low < q.support.lower_bound] = inf
        return result

    @register_kl(Independent, Independent)
    def _kl_independent_independent(p, q):
        if p.reinterpreted_batch_ndims != q.reinterpreted_batch_ndims:
            raise NotImplementedError
        result = kl_divergence(p.base_dist, q.base_dist)
        return _sum_rightmost(result, p.reinterpreted_batch_ndims)

    @register_kl(Cauchy, Cauchy)
    def _kl_cauchy_cauchy(p, q):
        # From https://arxiv.org/abs/1905.10965
        t1 = ((p.scale + q.scale).pow(2) + (p.loc - q.loc).pow(2)).log()
        t2 = (4 * p.scale * q.scale).log()
        return t1 - t2

    def _add_kl_info():
        """Appends a list of implemented KL functions to the doc for kl_divergence."""
        rows = [
            "KL divergence is currently implemented for the following distribution pairs:"
        ]
        for p, q in sorted(
            _KL_REGISTRY, key=lambda p_q: (p_q[0].__name__, p_q[1].__name__)
        ):
            rows.append(
                f"* :class:`~torch.distributions.{p.__name__}` and :class:`~torch.distributions.{q.__name__}`"
            )
        kl_info = "\n\t".join(rows)
        if kl_divergence.__doc__:
            kl_divergence.__doc__ += kl_info  # type: ignore[operator]

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.