Ë
    ˆ\;j|   ã                   óX  — d dl Z d dlZd dlZd dlmZ d dlmZ d dlmZ d dl	m
Z
 d dlmZ d dlmZ d dlmZ d d	lmZ d d
lmZ d dlmZ d dlmZ d dlmZ d dlmZ ddgZi Zd„ Zd„ Z d„ Z!e jD                   G d„ d«      «       Z# e ee«      d„ «       Z$ e ee«      d„ «       Z% e ee«      d„ «       Z& e ee«      d„ «       Z' e e
e
«      d„ «       Z( e ee«      d„ «       Z) e ee«      d„ «       Z* e ee«      d„ «       Z+ e ee«      d„ «       Z, e ee«      d„ «       Z- e ee«      d „ «       Z.d!„ Z/y)"é    N)Ú	Bernoulli)ÚBeta)ÚCategorical)ÚCauchy)Ú	Dirichlet)ÚDistribution)ÚExponentialFamily)Ú	Geometric)ÚLaplace)Ú	LogNormal)ÚNormal)ÚUniform)Úin_dynamic_modeÚregister_klÚkl_divergencec                 óL   —  t        t        | «      t        |«      «      | |«      S )a-  
    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:: python

            >>> 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)

    )Ú	_dispatchÚtype©ÚpÚqs     ú_G:\00. PROJECTS\API\Inventory\templateJSON\kerjaOCR\Lib\site-packages\paddle/distribution/kl.pyr   r   %   s"   € ð: 'Œ9”T˜!“Wœd 1›gÓ& q¨!Ó,Ð,ó    c                 ól   ‡ ‡— t        ‰ t        «      rt        ‰t        «      st        d«      ‚ˆ ˆfd„}|S )a�  Decorator for register a KL divergence implemention function.

    The ``kl_divergence(p, q)`` function will search concrete implemention
    functions registered by ``register_kl``, according to multi-dispatch pattern.
    If an implemention function is found, it will return the result, otherwise,
    it will raise ``NotImplementError`` exception. Users can register
    implemention function by the decorator.

    Args:
        cls_p (Distribution): The Distribution type of Instance p. Subclass derived from ``Distribution``.
        cls_q (Distribution): The Distribution type of Instance q. Subclass derived from ``Distribution``.

    Examples:
        .. code-block:: python

            >>> import paddle

            >>> @paddle.distribution.register_kl(paddle.distribution.Beta, paddle.distribution.Beta)
            >>> def kl_beta_beta():
            ...     pass # insert implementation here
    z0cls_p and cls_q must be subclass of Distributionc                 ó   •— | t         ‰‰f<   | S ©N)Ú_REGISTER_TABLE)ÚfÚcls_pÚcls_qs    €€r   Ú	decoratorzregister_kl.<locals>.decorator`   s   ø€ Ø()Œ˜˜u˜Ñ%Øˆr   )Ú
issubclassr   Ú	TypeError)r   r    r!   s   `` r   r   r   E   s6   ù€ ô, �eœ\Ô*´*ØŒ|ô3ô ÐJÓKÐKõð Ðr   c                 óè  — t         D ��cg c]!  \  }}t        | |«      rt        ||«      r||f‘Œ# }}}|st        ‚t        d„ |D «       «      j                  \  }}t        d„ |D «       «      j                  \  }}t         ||f   t         ||f   urTt        j                  dj                  | j                  |j                  |j                  |j                  «      t        «       t         ||f   S c c}}w )z3Multiple dispatch into concrete implement function.c              3   ó,   K  — | ]  }t        |Ž –— Œ y ­wr   )Ú_Compare©Ú.0Úms     r   Ú	<genexpr>z_dispatch.<locals>.<genexpr>s   s   è ø€ Ð6©v¨!œ 1œ©vùs   ‚c              3   ó>   K  — | ]  }t        t        |«      Ž –— Œ y ­wr   )r&   Úreversedr'   s     r   r*   z_dispatch.<locals>.<genexpr>t   s   è ø€ ÐB¹6°aœ8¤X¨a£[Ô1¹6ùs   ‚z;Ambiguous kl_divergence({}, {}). Please register_kl({}, {}))
r   r"   ÚNotImplementedErrorÚminÚclassesÚwarningsÚwarnÚformatÚ__name__ÚRuntimeWarning)	r   r    Úsuper_pÚsuper_qÚmatchsÚleft_pÚleft_qÚright_pÚright_qs	            r   r   r   g   sï   € õ !0ôá /ÑˆG�WÜ�e˜WÔ%¬*°U¸GÔ*Dð 
�'ÒØ /ð ñ ñ
 Ü!Ð!äÑ6©vÓ6Ó6×>Ñ>�N€FˆFÜÑB¹6ÓBÓB×JÑJÑ€GˆWä�v˜v�~Ñ&¬o¸gÀwÐ>NÑ.OÑOÜ�‰ØI×PÑPØ—‘Ø—‘Ø—‘Ø× Ñ ó	ô ô	
ô ˜6 6˜>Ñ*Ð*ùó-s   Š&C.c                   ó   — e Zd Zd„ Zd„ Zd„ Zy)r&   c                 ó   — || _         y r   ©r/   )Úselfr/   s     r   Ú__init__z_Compare.__init__†   s	   € Øˆ�r   c                 ó4   — | j                   |j                   k(  S r   r>   )r?   Úothers     r   Ú__eq__z_Compare.__eq__‰   s   € Ø�|‰|˜uŸ}™}Ñ,Ð,r   c                 ó|   — t        | j                  |j                  «      D ]  \  }}t        ||«      s y||usŒ y y)NFT)Úzipr/   r"   )r?   rB   Úcls_xÚcls_ys       r   Ú__le__z_Compare.__le__Œ   s?   € Ü §¡¨e¯m©mÖ<‰LˆE�5Ü˜e UÔ+ÙØ˜EÒ!ØØð =ð
 r   N)r3   Ú
__module__Ú__qualname__r@   rC   rH   © r   r   r&   r&   „   s   „ òò-ór   r&   c                 ó$   — | j                  |«      S r   ©r   r   s     r   Ú_kl_bernoulli_bernoullirN   •   ó   € à�?‰?˜1ÓÐr   c                 ó  — |j                   j                  «       |j                  j                  «       z   | j                   | j                  z   j                  «       z   | j                   j                  «       | j                  j                  «       z   |j                   |j                  z   j                  «       z   z
  | j                   |j                   z
  | j                   j                  «       z  z   | j                  |j                  z
  | j                  j                  «       z  z   |j                   |j                  z   | j                   | j                  z   z
  | j                   | j                  z   j                  «       z  z   S r   )ÚalphaÚlgammaÚbetaÚdigammar   s     r   Ú_kl_beta_betarU   š   s  € ð 
�‰�‰Ó	˜AŸF™FŸM™M›OÑ	+¨q¯w©w¸¿¹Ñ/?×.GÑ.GÓ.IÑ	IØ�7‰7�>‰>Ó˜aŸf™fŸm™m›oÑ-°·±¸1¿6¹6Ñ1A×0IÑ0IÓ0KÑKñ	Mà�G‰G�a—g‘gÑ §¡§¡Ó!2Ñ2ñ	4ð �F‰F�Q—V‘V‰O˜qŸv™vŸ~™~Ó/Ñ/ñ	1ð
 �g‰g˜Ÿ™Ñ 1§7¡7¨Q¯V©VÑ#3Ñ4Ø�w‰w˜Ÿ™Ñ×(Ñ(Ó*ñ+ñ	
ð	r   c                 ó,  — | j                   j                  d«      j                  «       |j                   j                  d«      j                  «       z
  | j                   j                  «       |j                   j                  «       z
  j                  d«      z
  | j                   |j                   z
  | j                   j                  «       | j                   j                  d«      j                  «       j	                  d«      z
  z  j                  d«      z   S )Néÿÿÿÿ)ÚconcentrationÚsumrR   rT   Ú	unsqueezer   s     r   Ú_kl_dirichlet_dirichletr[   ¨   s×   € ð 
�‰×	Ñ	˜RÓ	 ×	'Ñ	'Ó	)¨A¯O©O×,?Ñ,?ÀÓ,C×,JÑ,JÓ,LÑ	LØ�O‰O×"Ñ"Ó$ q§¡×'=Ñ'=Ó'?Ñ?×DÑDÀRÓHñ	Jð —‘ 1§?¡?Ñ2à—O‘O×+Ñ+Ó-Ø—o‘o×)Ñ)¨"Ó-×5Ñ5Ó7×AÑAÀ"ÓEñFñ÷
 ‰c�"‹gñ
	
ðr   c                 ó$   — | j                  |«      S r   rM   r   s     r   Ú_kl_categorical_categoricalr]   ¹   rO   r   c                 ó$   — | j                  |«      S r   rM   r   s     r   Ú_kl_cauchy_cauchyr_   ¾   rO   r   c                 ó$   — | j                  |«      S r   rM   r   s     r   Ú_kl_normal_normalra   Ã   rO   r   c                 ó$   — | j                  |«      S r   rM   r   s     r   Ú_kl_uniform_uniformrc   È   rO   r   c                 ó$   — | j                  |«      S r   rM   r   s     r   Ú_kl_laplace_laplacere   Í   rO   r   c                 ó$   — | j                  |«      S r   rM   r   s     r   Ú_kl_geometric_geometricrg   Ò   rO   r   c           	      óÀ  — t        | «      t        |«      k(  st        ‚g }| j                  D ]*  }|j                  «       }d|_        |j                  |«       Œ, |j                  } | j                  |Ž }	 t        «       rt        j                  ||d¬«      }n t        j                  j                  ||«      } |j                  |Ž |z
  }t!        |||«      D ]0  \  }	}
}|
|	z
  |z  }|t#        |t%        |j&                  «      «      z  }Œ2 |S # t        $ rH}t        dj                  t        | «      j                  t        |«      j                  ¬«      «      |‚d}~ww xY w)zuCompute kl-divergence using `Bregman divergences <https://www.lix.polytechnique.fr/~nielsen/EntropyEF-ICIP2010.pdf>`_FT)Úcreate_graphzlCann't compute kl_divergence({cls_p}, {cls_q}) use bregman divergence. Please register_kl({cls_p}, {cls_q}).)r   r    N)r   r-   Ú_natural_parametersÚdetachÚstop_gradientÚappendÚ_log_normalizerr   ÚpaddleÚgradÚstaticÚ	gradientsÚRuntimeErrorr#   r2   r3   rE   Ú_sum_rightmostÚlenÚevent_shape)r   r   Úp_natural_paramsÚparamÚq_natural_paramsÚ
p_log_normÚp_gradsÚeÚklÚp_paramÚq_paramÚp_gradÚterms                r   Ú_kl_expfamily_expfamilyr‚   ×   sm  € ô �‹7”d˜1“gÒÜ!Ð!àÐØ×&Ô&ˆØ—‘“ˆØ#ˆÔØ×Ñ Õ&ð 'ð
 ×,Ñ,Ðà"�×"Ñ"Ð$4Ð5€JðÜÔÜ—k‘kØÐ,¸4ô‰Gô —m‘m×-Ñ-¨jÐ:JÓKˆGð 
ˆ×	Ñ	Ð,Ð	-°
Ñ	:€BÜ$'ØÐ*¨Gö%Ñ ˆ�˜&ð ˜'Ñ! VÑ+ˆØ
Œn˜T¤3 q§}¡}Ó#5Ó6Ñ6‰ð	%ð €Iøô ò ÜØz÷  Bñ  BÜ˜1“g×&Ñ&¬d°1«g×.>Ñ.>ð Bó ó
ð ð		ûðús   Á5AD Ä	EÄAEÅEc                 óL   — | j                   j                  |j                   «      S r   )Ú_baser   r   s     r   Ú_kl_lognormal_lognormalr…   ÿ   s   € à�7‰7× Ñ  §¡Ó)Ð)r   c                 óZ   — |dkD  r%| j                  t        t        | d«      «      «      S | S )Nr   )rY   ÚlistÚrange)ÚvalueÚns     r   rt   rt     s)   € Ø,-°ªEˆ5�9‰9”Tœ%   A›,Ó'Ó(Ð<°uÐ<r   )0Ú	functoolsr0   ro   Úpaddle.distribution.bernoullir   Úpaddle.distribution.betar   Úpaddle.distribution.categoricalr   Úpaddle.distribution.cauchyr   Úpaddle.distribution.dirichletr   Ú paddle.distribution.distributionr   Ú&paddle.distribution.exponential_familyr	   Úpaddle.distribution.geometricr
   Úpaddle.distribution.laplacer   Úpaddle.distribution.lognormalr   Úpaddle.distribution.normalr   Úpaddle.distribution.uniformr   Úpaddle.frameworkr   Ú__all__r   r   r   r   Útotal_orderingr&   rN   rU   r[   r]   r_   ra   rc   re   rg   r‚   r…   rt   rK   r   r   Ú<module>r›      s§  ðó Û ã Ý 3Ý )Ý 7Ý -Ý 3Ý 9Ý DÝ 3Ý /Ý 3Ý -Ý /Ý ,à˜/Ð
*€à€ò-ò@òD+ð: ×Ñ÷ð ó ðñ  ˆY˜	Ó"ñó #ðñ ˆT�4Óñ
ó ð
ñ ˆY˜	Ó"ñó #ðñ  ˆ[˜+Ó&ñó 'ðñ ˆV�VÓñó ðñ ˆV�VÓñó ðñ ˆW�gÓñó ðñ ˆW�gÓñó ðñ ˆY˜	Ó"ñó #ðñ ÐÐ 1Ó2ñ$ó 3ð$ñN ˆY˜	Ó"ñ*ó #ð*ó=r   