U
    ©Ê¹j¬  ã                   @  s  d Z ddlmZ ddlZddlmZ ddlmZ ddlZddl	m
Z
mZ ddlmZmZmZ d	d
ddddddddddgZdZdd„ Zed(dd	„ƒZed)dd
„ƒZedd„ ƒZeZeZed*dd„ƒZedd„ ƒZed+dd„ƒZeed d!œd"d„ƒƒZed,d#d„ƒZeed-d d%d&œd'd„ƒƒZdS ).z§Wrapper to mimic (parts of) np.random API surface.

NumPy has strict guarantees on reproducibility etc; here we don't give any.

Q: default dtype is float64 in numpy

é    )ÚannotationsN)Úsqrt)ÚOptionalé   )Ú_dtypes_implÚ_util)Úarray_or_scalarÚ	ArrayLikeÚ
normalizerÚseedÚrandom_sampleÚsampleÚrandomÚrandÚrandnÚnormalÚchoiceÚrandintÚshuffleÚuniformÚUSE_NUMPY_RANDOMFc                   s   t  ˆ ¡‡ fdd„ƒ}|S )Nc                    sP   t dkrˆ | |ŽS t dkr<ddlm} t|ˆ jƒ}|| |ŽS tdt › d�ƒ‚d S )NFTr   )r   zUSE_NUMPY_RANDOM=z not understood.)r   Únumpyr   ÚgetattrÚ__name__Ú
ValueError)ÚargsÚkwdsÚnrÚf©Úfunc© úR/var/www/html/bwave/python/venv/lib/python3.8/site-packages/torch/_numpy/random.pyÚinner(   s    

zdeco_stream.<locals>.inner)Ú	functoolsÚwraps)r    r#   r!   r   r"   Údeco_stream'   s    r&   c                 C  s   | d k	rt j | ¡ d S ©N)Útorchr   Zmanual_seed)r   r!   r!   r"   r   7   s    c                 C  s8   | d krd} t  ¡ j}tj| |d� ¡ }t|| d kd�S ©Nr!   ©Údtype©Zreturn_scalar©r   Údefault_dtypesÚfloat_dtyper(   ÚemptyZuniform_r   ©Úsizer+   Úvaluesr!   r!   r"   r   =   s
    
c                  G  s   t | ƒS r'   )r   ©r2   r!   r!   r"   r   F   s    ç        ç      ð?c                 C  s<   |d krd}t  ¡ j}tj||d� | |¡}t||d kd�S r)   r-   )ÚlowÚhighr2   r+   r3   r!   r!   r"   r   O   s
    
c                  G  s(   t  ¡ j}tj| |d�}t|| d kd�S )Nr*   r,   )r   r.   r/   r(   r   r   r1   r!   r!   r"   r   X   s    
c                 C  s<   |d krd}t  ¡ j}tj||d� | |¡}t||d kd�S r)   )r   r.   r/   r(   r0   Znormal_r   )ÚlocÚscaler2   r+   r3   r!   r!   r"   r   _   s
    
r	   )Úxc                 C  s&   t  | jd ¡}| | }|  |¡ d S )Nr   )r(   ZrandpermÚshapeZcopy_)r;   ÚpermZxpr!   r!   r"   r   h   s    c                 C  sT   |d krd}t |ttfƒs |f}|d kr2d|  } }tj| ||d�}t|t|d kd�S )Nr!   r   r4   r,   )Ú
isinstanceÚtupleÚlistr(   r   r   Úint)r7   r8   r2   r3   r!   r!   r"   r   p   s    
TzOptional[ArrayLike])ÚaÚpc           	      C  sÄ   |   ¡ dkrt | ¡} |d kr$d}n&t |¡rFd}|D ]}||9 }q6n|}|d krft | ¡| jd  }tt |j	¡j
ƒ}t| ¡ d ƒ|kr”tdƒ‚tj|||d�}t |¡r¸| |¡}| | }|S )Nr   r   r6   zprobabilities do not sum to 1.)Úreplacement)Znumelr(   Zaranger   Zis_sequenceZ	ones_liker<   r   Zfinfor+   ÚepsÚabsÚsumr   ZmultinomialZreshape)	rB   r2   ÚreplacerC   Znum_elÚelZatolÚindicesZsamplesr!   r!   r"   r   |   s&    



)N)N)r5   r6   N)r5   r6   N)NN)NTN)Ú__doc__Ú
__future__r   r$   Úmathr   Útypingr   r(   Ú r   r   Z_normalizationsr   r	   r
   Ú__all__r   r&   r   r   r   r   r   r   r   r   r   r   r   r!   r!   r!   r"   Ú<module>   sX   ô

