a
    MSic                     @   sT   d dl Z d dlmZ d dlmZ d dlmZmZmZm	Z	 dd Z
G dd deZdS )	    N)constraints)Distribution)broadcast_allprobs_to_logitslazy_propertylogits_to_probsc                 C   s    | j dd|  | j dd d S )Nr   )minmax   )clamp)x r   X/var/www/html/django/DPS/env/lib/python3.9/site-packages/torch/distributions/binomial.py_clamp_by_zero   s    r   c                       s   e Zd ZdZejejejdZdZ	d$ fdd	Z
d% fdd		Zd
d Zejddddd Zedd Zedd Zedd Zedd Zedd Zedd Ze fddZdd Zd d! Zd&d"d#Z  ZS )'Binomiala  
    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::

        >>> 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logitsT   Nc                    s   |d u |d u krt d|d urDt||\| _| _| j| j| _n"t||\| _| _| j| j| _|d urt| jn| j| _| j }tt	| j
||d d S )Nz;Either `probs` or `logits` must be specified, but not both.validate_args)
ValueErrorr   r   r   type_asr   _paramsizesuperr   __init__)selfr   r   r   r   batch_shape	__class__r   r   r   '   s    
zBinomial.__init__c                    s   |  t|}t|}| j||_d| jv rD| j||_|j|_d| jv rd| j	||_	|j	|_t
t|j|dd | j|_|S )Nr   r   Fr   )_get_checked_instancer   torchSizer   expand__dict__r   r   r   r   r   _validate_args)r   r   	_instancenewr    r   r   r%   5   s    


zBinomial.expandc                 O   s   | j j|i |S N)r   r)   )r   argskwargsr   r   r   _newC   s    zBinomial._newr   )is_discrete	event_dimc                 C   s   t d| jS )Nr   )r   integer_intervalr   r   r   r   r   supportF   s    zBinomial.supportc                 C   s   | j | j S r*   r   r   r1   r   r   r   meanJ   s    zBinomial.meanc                 C   s   | j d | j  j| j dS )Nr   r	   )r   r   floorr   r1   r   r   r   modeN   s    zBinomial.modec                 C   s   | j | j d| j  S Nr   r3   r1   r   r   r   varianceR   s    zBinomial.variancec                 C   s   t | jddS NT)	is_binary)r   r   r1   r   r   r   r   V   s    zBinomial.logitsc                 C   s   t | jddS r9   )r   r   r1   r   r   r   r   Z   s    zBinomial.probsc                 C   s
   | j  S r*   )r   r   r1   r   r   r   param_shape^   s    zBinomial.param_shapec                 C   sR   |  |}t , t| j|| j|W  d    S 1 sD0    Y  d S r*   )_extended_shaper#   no_gradbinomialr   r%   r   )r   sample_shapeshaper   r   r   sampleb   s    

zBinomial.samplec              	   C   s   | j r| | t| jd }t|d }t| j| d }| jt| j | jttt	| j   | }|| j | | | S r7   )
r'   _validate_sampler#   lgammar   r   r   log1pexpabs)r   valueZlog_factorial_nZlog_factorial_kZlog_factorial_nmkZnormalize_termr   r   r   log_probg   s    
zBinomial.log_probc                 C   sJ   t | j }| j |ks$td| | d}t|| 	d S )Nz5Inhomogeneous total count not supported by `entropy`.Fr   )
intr   r
   r   NotImplementedErrorrH   enumerate_supportr#   rE   sum)r   r   rH   r   r   r   entropyw   s
    zBinomial.entropyc                 C   sp   t | j }| j |ks$tdtjd| | jj| jj	d}|
ddt| j  }|rl|d| j }|S )Nz?Inhomogeneous total count not supported by `enumerate_support`.r   )dtypedevice))r   )rI   r   r
   r   rJ   r#   aranger   rN   rO   viewlen_batch_shaper%   )r   r%   r   valuesr   r   r   rK      s    zBinomial.enumerate_support)r   NNN)N)T)__name__
__module____qualname____doc__r   nonnegative_integerunit_intervalrealarg_constraintshas_enumerate_supportr   r%   r-   dependent_propertyr2   propertyr4   r6   r8   r   r   r   r;   r#   r$   rA   rH   rM   rK   __classcell__r   r   r    r   r      s6   






r   )r#   Ztorch.distributionsr    torch.distributions.distributionr   torch.distributions.utilsr   r   r   r   r   r   r   r   r   r   <module>   s
   