
    קgg                         d dl 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 dlmZ dgZ G d	 de          ZdS )
    )NumberN)nan)constraints)ExponentialFamily)broadcast_alllazy_propertylogits_to_probsprobs_to_logits) binary_cross_entropy_with_logits	Bernoullic                   B    e Zd ZdZej        ej        dZej        Z	dZ
dZd fd	Zd fd	Zd Zed	             Zed
             Zed             Zed             Zed             Zed             Z ej                    fdZd Zd ZddZed             Zd Z xZS )r   a1  
    Creates a Bernoulli distribution parameterized by :attr:`probs`
    or :attr:`logits` (but not both).

    Samples are binary (0 or 1). They take the value `1` with probability `p`
    and `0` with probability `1 - p`.

    Example::

        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
        >>> m = Bernoulli(torch.tensor([0.3]))
        >>> m.sample()  # 30% chance 1; 70% chance 0
        tensor([ 0.])

    Args:
        probs (Number, Tensor): the probability of sampling `1`
        logits (Number, Tensor): the log-odds of sampling `1`
    )probslogitsTr   Nc                    |d u |d u k    rt          d          |,t          |t                    }t          |          \  | _        n+t          |t                    }t          |          \  | _        || j        n| j        | _        |rt          j                    }n| j        	                                }t                                          ||           d S )Nz;Either `probs` or `logits` must be specified, but not both.validate_args)
ValueError
isinstancer   r   r   r   _paramtorchSizesizesuper__init__)selfr   r   r   	is_scalarbatch_shape	__class__s         Y/var/www/html/ai-engine/env/lib/python3.11/site-packages/torch/distributions/bernoulli.pyr   zBernoulli.__init__,   s    TMv~..M   "5&11I)%00MTZZ"6622I*622NT[$)$5djj4; 	-*,,KK+**,,KMBBBBB    c                    |                      t          |          }t          j        |          }d| j        v r+| j                            |          |_        |j        |_        d| j        v r+| j                            |          |_        |j        |_        t          t          |          
                    |d           | j        |_        |S )Nr   r   Fr   )_get_checked_instancer   r   r   __dict__r   expandr   r   r   r   _validate_args)r   r   	_instancenewr   s       r   r$   zBernoulli.expand>   s    ((I>>j--dm##
))+66CICJt}$$++K88CJCJi&&{%&HHH!0
r    c                 &     | j         j        |i |S N)r   r'   )r   argskwargss      r   _newzBernoulli._newK   s    t{////r    c                     | j         S r)   r   r   s    r   meanzBernoulli.meanN   s
    zr    c                 r    | j         dk                        | j                   }t          || j         dk    <   |S )Ng      ?)r   tor   )r   modes     r   r3   zBernoulli.modeR   s4    
c!%%dj11"%TZ3r    c                 &    | j         d| j         z
  z  S )N   r.   r/   s    r   variancezBernoulli.varianceX   s    zQ^,,r    c                 .    t          | j        d          S NT)	is_binary)r
   r   r/   s    r   r   zBernoulli.logits\   s    tzT::::r    c                 .    t          | j        d          S r8   )r	   r   r/   s    r   r   zBernoulli.probs`   s    t{d;;;;r    c                 4    | j                                         S r)   )r   r   r/   s    r   param_shapezBernoulli.param_shaped   s    {!!!r    c                     |                      |          }t          j                    5  t          j        | j                            |                    cd d d            S # 1 swxY w Y   d S r)   )_extended_shaper   no_grad	bernoullir   r$   )r   sample_shapeshapes      r   samplezBernoulli.sampleh   s    $$\22]__ 	= 	=?4:#4#4U#;#;<<	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	=s   ,A""A&)A&c                     | j         r|                     |           t          | j        |          \  }}t	          ||d           S Nnone)	reduction)r%   _validate_sampler   r   r   )r   valuer   s      r   log_probzBernoulli.log_probm   sN     	)!!%(((%dk5990&QQQQQr    c                 :    t          | j        | j        d          S rE   )r   r   r   r/   s    r   entropyzBernoulli.entropys   s$    /Kv
 
 
 	
r    c                     t          j        d| j        j        | j        j                  }|                    ddt          | j                  z  z             }|r|                    d| j        z             }|S )N   )dtypedevice))r5   )	r   aranger   rO   rP   viewlen_batch_shaper$   )r   r$   valuess      r   enumerate_supportzBernoulli.enumerate_supportx   sn    at{'8ASTTTUTC0A,B,B%BBCC 	>]]54+<#<==Fr    c                 6    t          j        | j                  fS r)   )r   logitr   r/   s    r   _natural_paramszBernoulli._natural_params   s    DJ''))r    c                 N    t          j        t          j        |                    S r)   )r   log1pexp)r   xs     r   _log_normalizerzBernoulli._log_normalizer   s    {59Q<<(((r    )NNNr)   )T) __name__
__module____qualname____doc__r   unit_intervalrealarg_constraintsbooleansupporthas_enumerate_support_mean_carrier_measurer   r$   r,   propertyr0   r3   r6   r   r   r   r<   r   r   rC   rJ   rL   rW   rZ   r_   __classcell__)r   s   @r   r   r      s        $ !, 9[EUVVO!G C C C C C C$     0 0 0   X   X
 - - X- ; ; ]; < < ]< " " X" #-%*,, = = = =
R R R
 
 

    * * X*) ) ) ) ) ) )r    )numbersr   r   r   torch.distributionsr   torch.distributions.exp_familyr   torch.distributions.utilsr   r   r	   r
   torch.nn.functionalr   __all__r    r    r   <module>rt      s                 + + + + + + < < < < < <            A @ @ @ @ @ -p) p) p) p) p)! p) p) p) p) p)r    