Ë
    óÍ:j®  ã                   óz   — d dl mZm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
mZmZmZ dgZd„ Z G d	„ de«      Zy)
é    )ÚOptionalÚUnionN)ÚTensor)Úconstraints)ÚDistribution)Úbroadcast_allÚlazy_propertyÚlogits_to_probsÚprobs_to_logitsÚBinomialc                 óX   — | j                  d¬«      | z   | j                  d¬«      z
  dz  S )Nr   )Úmin©Úmaxé   )Úclamp)Úxs    úq/home/mcse/projects/srt_converter/srt-converter-venv/lib/python3.12/site-packages/torch/distributions/binomial.pyÚ_clamp_by_zeror      s+   € à�G‰G˜ˆG‹N˜QÑ §¡¨Q £Ñ/°1Ñ4Ð4ó    c                   óÄ  ‡ — e Zd ZdZej
                  ej                  ej                  dœZdZ		 	 	 	 dde
eef   dee   dee   dee   d	df
ˆ fd
„Zdˆ fd„	Zd„ Z ej$                  dd¬«      d„ «       Zed	efd„«       Zed	efd„«       Zed	efd„«       Zed	efd„«       Zed	efd„«       Zed	ej8                  fd„«       Z ej8                  «       fd„Zd„ Zd„ Z dd„Z!ˆ xZ"S )r   añ  
    Creates a Binomial distribution parameterized by :attr:`total_count` and
    either :attr:`probs` or :attr:`logits` (but not both). :attr:`total_count` must be
    broadcastable with :attr:`probs`/:attr:`logits`.

    Example::

        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
        >>> m = Binomial(100, torch.tensor([0 , .2, .8, 1]))
        >>> x = m.sample()
        tensor([   0.,   22.,   71.,  100.])

        >>> m = Binomial(torch.tensor([[5.], [10.]]), torch.tensor([0.5, 0.8]))
        >>> x = m.sample()
        tensor([[ 4.,  5.],
                [ 7.,  6.]])

    Args:
        total_count (int or Tensor): number of Bernoulli trials
        probs (Tensor): Event probabilities
        logits (Tensor): Event log-odds
    )Útotal_countÚprobsÚlogitsTNr   r   r   Úvalidate_argsÚreturnc                 óÜ  •— |d u |d u k(  rt        d«      ‚|�Dt        ||«      \  | _        | _        | j                  j	                  | j                  «      | _        nG|€J ‚t        ||«      \  | _        | _        | j                  j	                  | j
                  «      | _        |�| j                  n| j
                  | _        | j                  j                  «       }t        ‰| �%  ||¬«       y )Nz;Either `probs` or `logits` must be specified, but not both.©r   )
Ú
ValueErrorr   r   r   Útype_asr   Ú_paramÚsizeÚsuperÚ__init__)Úselfr   r   r   r   Úbatch_shapeÚ	__class__s         €r   r$   zBinomial.__init__7   så   ø€ ð �TˆM˜v¨˜~Ò.ÜØMóð ð Ðô ˜k¨5Ó1ñØÔ Ø”
à#×/Ñ/×7Ñ7¸¿
¹
ÓCˆDÕàÐ%Ð%Ð%ô ˜k¨6Ó2ñØÔ Ø”à#×/Ñ/×7Ñ7¸¿¹ÓDˆDÔà$)Ð$5�d—j’j¸4¿;¹;ˆŒØ—k‘k×&Ñ&Ó(ˆÜ‰Ñ˜°MÐÕBr   c                 óæ  •— | j                  t        |«      }t        j                  |«      }| j                  j                  |«      |_        d| j                  v r1| j                  j                  |«      |_        |j                  |_        d| j                  v r1| j                  j                  |«      |_	        |j                  |_        t        t        |�/  |d¬«       | j                  |_        |S )Nr   r   Fr   )Ú_get_checked_instancer   ÚtorchÚSizer   ÚexpandÚ__dict__r   r!   r   r#   r$   Ú_validate_args)r%   r&   Ú	_instanceÚnewr'   s       €r   r,   zBinomial.expandT   s¼   ø€ Ø×(Ñ(¬°9Ó=ˆÜ—j‘j Ó-ˆØ×*Ñ*×1Ñ1°+Ó>ˆŒØ�d—m‘mÑ#ØŸ
™
×)Ñ)¨+Ó6ˆCŒIØŸ™ˆCŒJØ�t—}‘}Ñ$ØŸ™×+Ñ+¨KÓ8ˆCŒJØŸ™ˆCŒJÜŒh˜Ñ% kÀÐ%ÔGØ!×0Ñ0ˆÔØˆ
r   c                 ó:   —  | j                   j                  |i |¤ŽS ©N)r!   r0   )r%   ÚargsÚkwargss      r   Ú_newzBinomial._newb   s   € Øˆt�{‰{�‰ Ð/¨Ñ/Ð/r   r   )Úis_discreteÚ	event_dimc                 óB   — t        j                  d| j                  «      S )Nr   )r   Úinteger_intervalr   ©r%   s    r   ÚsupportzBinomial.supporte   s   € ä×+Ñ+¨A¨t×/?Ñ/?Ó@Ð@r   c                 ó4   — | j                   | j                  z  S r2   ©r   r   r:   s    r   ÚmeanzBinomial.meani   s   € à×Ñ $§*¡*Ñ,Ð,r   c                 óŠ   — | j                   dz   | j                  z  j                  «       j                  | j                   ¬«      S )Né   r   )r   r   Úfloorr   r:   s    r   ÚmodezBinomial.modem   s9   € à×!Ñ! AÑ%¨¯©Ñ3×:Ñ:Ó<×BÑBÀt×GWÑGWÐBÓXÐXr   c                 óT   — | j                   | j                  z  d| j                  z
  z  S ©Nr@   r=   r:   s    r   ÚvariancezBinomial.varianceq   s$   € à×Ñ $§*¡*Ñ,°°D·J±J±Ñ?Ð?r   c                 ó0   — t        | j                  d¬«      S ©NT)Ú	is_binary)r   r   r:   s    r   r   zBinomial.logitsu   s   € ä˜tŸz™z°TÔ:Ð:r   c                 ó0   — t        | j                  d¬«      S rG   )r
   r   r:   s    r   r   zBinomial.probsy   s   € ä˜tŸ{™{°dÔ;Ð;r   c                 ó6   — | j                   j                  «       S r2   )r!   r"   r:   s    r   Úparam_shapezBinomial.param_shape}   s   € à�{‰{×ÑÓ!Ð!r   c                 ó  — | j                  |«      }t        j                  «       5  t        j                  | j                  j                  |«      | j                  j                  |«      «      cd d d «       S # 1 sw Y   y xY wr2   )Ú_extended_shaper*   Úno_gradÚbinomialr   r,   r   )r%   Úsample_shapeÚshapes      r   ÚsamplezBinomial.sample�   sa   € Ø×$Ñ$ \Ó2ˆÜ�]‰]‹_ñ 	Ü—>‘>Ø× Ñ ×'Ñ'¨Ó.°·
±
×0AÑ0AÀ%Ó0Hó÷	÷ 	ò 	ús   ¦AA8Á8Bc           	      ó  — | j                   r| j                  |«       t        j                  | j                  dz   «      }t        j                  |dz   «      }t        j                  | j                  |z
  dz   «      }| j                  t        | j                  «      z  | j                  t        j                  t        j                  t        j                  | j                  «       «      «      z  z   |z
  }|| j                  z  |z
  |z
  |z
  S rD   )
r.   Ú_validate_sampler*   Úlgammar   r   r   Úlog1pÚexpÚabs)r%   ÚvalueÚlog_factorial_nÚlog_factorial_kÚlog_factorial_nmkÚnormalize_terms         r   Úlog_probzBinomial.log_probˆ   sâ   € Ø×ÒØ×!Ñ! %Ô(ÜŸ,™, t×'7Ñ'7¸!Ñ';Ó<ˆÜŸ,™, u¨q¡yÓ1ˆÜ!ŸL™L¨×)9Ñ)9¸EÑ)AÀAÑ)EÓFÐð ×Ñœ~¨d¯k©kÓ:Ñ:Ø×Ñ¤§¡¬U¯Y©Y¼¿	¹	À$Ç+Á+Ó8NÐ7NÓ-OÓ!PÑPñQàñð 	ð �D—K‘KÑ /Ñ1Ð4EÑEÈÑVð	
r   c                 ó(  — t        | j                  j                  «       «      }| j                  j                  «       |k(  st	        d«      ‚| j                  | j                  d«      «      }t        j                  |«      |z  j                  d«       S )Nz5Inhomogeneous total count not supported by `entropy`.Fr   )
Úintr   r   r   ÚNotImplementedErrorr^   Úenumerate_supportr*   rW   Úsum)r%   r   r^   s      r   ÚentropyzBinomial.entropyœ   s   € Ü˜$×*Ñ*×.Ñ.Ó0Ó1ˆØ×Ñ×#Ñ#Ó%¨Ò4Ü%ØGóð ð —=‘= ×!7Ñ!7¸Ó!>Ó?ˆÜ—‘˜8Ó$ xÑ/×4Ñ4°QÓ7Ð7Ð7r   c                 ó¶  — t        | j                  j                  «       «      }| j                  j                  «       |k(  st	        d«      ‚t        j                  d|z   | j                  j                  | j                  j                  ¬«      }|j                  ddt        | j                  «      z  z   «      }|r|j                  d| j                  z   «      }|S )Nz?Inhomogeneous total count not supported by `enumerate_support`.r@   )ÚdtypeÚdevice)éÿÿÿÿ)r@   )r`   r   r   r   ra   r*   Úaranger!   rf   rg   ÚviewÚlenÚ_batch_shaper,   )r%   r,   r   Úvaluess       r   rb   zBinomial.enumerate_support¦   s·   € Ü˜$×*Ñ*×.Ñ.Ó0Ó1ˆØ×Ñ×#Ñ#Ó%¨Ò4Ü%ØQóð ô —‘Ø�‰O 4§;¡;×#4Ñ#4¸T¿[¹[×=OÑ=Oô
ˆð —‘˜U T¬C°×0AÑ0AÓ,BÑ%BÑBÓCˆÙØ—]‘] 5¨4×+<Ñ+<Ñ#<Ó=ˆFØˆr   )r@   NNNr2   )T)#Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   Únonnegative_integerÚunit_intervalÚrealÚarg_constraintsÚhas_enumerate_supportr   r   r`   r   Úboolr$   r,   r5   Údependent_propertyr;   Úpropertyr>   rB   rE   r	   r   r   r*   r+   rK   rR   r^   rd   rb   Ú__classcell__)r'   s   @r   r   r      s‘  ø„ ñð0 #×6Ñ6Ø×*Ñ*Ø×"Ñ"ñ€Oð
 !Ðð +,Ø"&Ø#'Ø(,ñCà˜6 3˜;Ñ'ðCð ˜ÑðCð ˜Ñ ð	Cð
   ‘~ðCð 
õCõ:ò0ð $€[×#Ñ#°ÀÔBñAó CðAð ð-�fò -ó ð-ð ðY�fò Yó ðYð ð@˜&ò @ó ð@ð ð;˜ò ;ó ð;ð ð<�vò <ó ð<ð ð"˜UŸZ™Zò "ó ð"ð #- %§*¡*£,ó ò
ò(8÷r   )Útypingr   r   r*   r   Útorch.distributionsr   Ú torch.distributions.distributionr   Útorch.distributions.utilsr   r	   r
   r   Ú__all__r   r   © r   r   ú<module>r�      s:   ðç "ã Ý Ý +Ý 9÷ó ð ˆ,€ò5ô
Zˆ|õ Zr   