o
    8ήc*                     @   s  d dl Z d dlZd dlmZ ddlmZmZmZm	Z	m
Z
mZ d dlmZ d dlmZ e
 ZejZejZejZeeejejfZejgZdd ZG d	d
 d
e	ZG dd deZ ee j!ddG dd deZ"ee j ddeej j ddG dd de Z#edkreej j$dde# eej j%dde# eej j&dde# ee j'ddG dd deZ(eej j'ddG dd de Z(ee j)ddG dd  d eZ*ee j+d!deej j+d"dG d#d$ d$eZ,eej j-d%deej j.d&deej j/d'dG d(d) d)e Z0eej j1d*deej j2d+dG d,d- d-e Z3eej j4d.dG d/d0 d0e Z5eej j6d1deej j7d2dG d3d4 d4e Z8eej j9d5dG d6d7 d7e Z:eej j;d8deej j<d9deej j=d:deej j>d;dG d<d= d=e Z?eej j@d>dG d?d@ d@e ZAeej jBdAdG dBdC dCe ZCeej jDdDdeej jEdEdeej jFdFdeej jGdGdeej jHdHdeej jIdIdee jJdJdee jKdKdee jLdLdee jMdMdee jNdNdee jGdOdee jOdPdee jPdQdG dRdS dSe ZQeej jRdTdeej jSdUdeej jTdVdeej jUdWdeej jVdXdeej jWdYdee jXdZdee jYd[dG d\d] d]e ZZeej j[d^deej j\d_deej j]d`dG dadb dbe Z^ee jBdcdG ddde deeZ_dS )f    N)types   )ConcreteTemplateAbstractTemplateAttributeTemplateCallableTemplateRegistry	signature)numpy_version)glue_typingc                    sP   t  tjrtjdfS  fdd D r!t }ttj||fS td f )Nr   c                 3   s,    | ]}t  tjott |tjV  qd S N)
isinstancer   	BaseTupleallInteger).0vshape C/tmp/pip-target-vg8gfxp4/lib/python/numba/core/typing/randomdecl.py	<genexpr>#   s    
z"normalize_shape.<locals>.<genexpr>zinvalid size type %s)r   r   r   intplenUniTuple	TypeError)r   ndimr   r   r   normalize_shape    s   

r   c                   @      e Zd ZdZdddZdS )RandomTemplatezl
    A template helper to transparently handle the typing of array-returning
    np.random.* functions.
    Nc                    sP   j dd   dv sJ j d u rfdd}|S  fdd}|S )N.r   )nprandomc                     s    | i |S r   r   argskwargs)scalar_typerr   r   typer7      z)RandomTemplate.array_typer.<locals>.typerc                     sb    dkrt djf t\}}| i |}|d ur/tt|j|dg|j|f R  S d S )Nr"   z unexpected size parameter for %rC)r   keyr   r	   r   Arrayreturn_typer$   )r$   r%   r   r   sigprefixr&   selfsizer   r   r'   ;   s   
)r*   split)r0   r&   r1   r'   r   r.   r   array_typer1   s   zRandomTemplate.array_typerr   __name__
__module____qualname____doc__r3   r   r   r   r   r   +       r   c                   @   r   )ConcreteRandomTemplatezk
    A RandomTemplate subclass using the `cases` attribute as a list of
    allowed scalar signatures.
    Nc                    s0   | j | j | j fdd}t| ||S )Nc                     s<   | r| d d u r| d d } | r| d d u s  | |S )N)resolve_overloadr#   casescontextr*   r   r   concrete_scalar_typerV   s   zAConcreteRandomTemplate.array_typer.<locals>.concrete_scalar_typer)r*   r>   r?   r   r3   )r0   r1   r@   r   r=   r   r3   Q   s
   z"ConcreteRandomTemplate.array_typerr   r4   r   r   r   r   r:   K   r9   r:   zrandom.getrandbits)
typing_keyc                   @      e Zd ZeejejgZdS )Random_getrandbitsN)r5   r6   r7   r	   r   uint64int32r>   r   r   r   r   rC   _   s    rC   zrandom.randomznp.random.randomc                   @   s    e Zd ZeejgZdd ZdS )Random_randomc                       d fdd	}|S )Nc                         |  S r   r3   r1   r0   r   r   r'   i      z$Random_random.generic.<locals>.typerr   r   r0   r'   r   rK   r   generich      zRandom_random.genericN)r5   r6   r7   r	   r   float64r>   rN   r   r   r   r   rF   c   s    rF   )r      znp.random.random_sampleznp.random.sampleznp.random.ranfzrandom.randintc                   @   s   e Zd Zdd eD ZdS )Random_randintc                 C      g | ]}t |||qS r   r	   r   tpr   r   r   
<listcomp>~       Random_randint.<listcomp>Nr5   r6   r7   
_int_typesr>   r   r   r   r   rR   |   s    rR   znp.random.randintc                   @   4   e Zd Zdd eD Zedd eD 7 Zdd ZdS )rR   c                 C      g | ]}t ||qS r   rT   rU   r   r   r   rW          rY   c                 C   rS   r   rT   rU   r   r   r   rW      rX   c                    rG   )Nc                         || |S r   rI   )lowhighr1   rK   r   r   r'         z%Random_randint.generic.<locals>.typerNNr   rM   r   rK   r   rN      rO   zRandom_randint.genericNr5   r6   r7   r[   r>   rN   r   r   r   r   rR          zrandom.randrangec                   @   s>   e Zd Zdd eD Zedd eD 7 Zedd eD 7 ZdS )Random_randrangec                 C   r]   r   rT   rU   r   r   r   rW      r^   zRandom_randrange.<listcomp>c                 C   rS   r   rT   rU   r   r   r   rW      rX   c                 C      g | ]	}t ||||qS r   rT   rU   r   r   r   rW          NrZ   r   r   r   r   rf      s    rf   zrandom.seedznp.random.seedc                   @   rB   )Random_seedN)r5   r6   r7   r	   r   voiduint32r>   r   r   r   r   ri      s    ri   znp.random.geometricznp.random.logseriesznp.random.zipfc                   @   "   e Zd Zdd eD Zdd ZdS )Numpy_geometricc                 C      g | ]}t tj|qS r   r	   r   int64rU   r   r   r   rW      rX   zNumpy_geometric.<listcomp>c                    rG   )Nc                         || S r   rI   ar1   rK   r   r   r'      r(   z&Numpy_geometric.generic.<locals>.typerr   r   rM   r   rK   r   rN      rO   zNumpy_geometric.genericNr5   r6   r7   _float_typesr>   rN   r   r   r   r   rm          rm   znp.random.binomialznp.random.negative_binomialc                   @   rl   )Numpy_negative_binomialc                 C   s   g | ]
}t tjtj|qS r   ro   rU   r   r   r   rW      s    z"Numpy_negative_binomial.<listcomp>c                    rG   )Nc                    r_   r   rI   )npr1   rK   r   r   r'      rb   z.Numpy_negative_binomial.generic.<locals>.typerr   r   rM   r   rK   r   rN      rO   zNumpy_negative_binomial.genericNrt   r   r   r   r   rw      rv   rw   znp.random.poissonc                   @   s2   e Zd Zdd eD Zeeejg7 Zdd ZdS )Numpy_poissonc                 C   rn   r   ro   rU   r   r   r   rW      rX   zNumpy_poisson.<listcomp>c                    rG   )Nc                    rq   r   rI   )lamr1   rK   r   r   r'      r(   z$Numpy_poisson.generic.<locals>.typerrc   r   rM   r   rK   r   rN      rO   zNumpy_poisson.genericN)	r5   r6   r7   ru   r>   r	   r   rp   rN   r   r   r   r   rz      s    rz   znp.random.exponentialznp.random.rayleighc                   @   r\   )Numpy_exponentialc                 C   r]   r   rT   rU   r   r   r   rW      r^   zNumpy_exponential.<listcomp>c                 C      g | ]}t |qS r   rT   rU   r   r   r   rW          c                    rG   )Nc                    rq   r   rI   )scaler1   rK   r   r   r'      r(   z(Numpy_exponential.generic.<locals>.typerrc   r   rM   r   rK   r   rN      rO   zNumpy_exponential.genericNrt   r   r   r   r   r|      s    r|   znp.random.hypergeometricc                   @   rl   )Numpy_hypergeometricc                 C   rg   r   rT   rU   r   r   r   rW      rh   zNumpy_hypergeometric.<listcomp>c                    rG   )Nc                         || ||S r   rI   )ngoodnbadnsampler1   rK   r   r   r'         z+Numpy_hypergeometric.generic.<locals>.typerr   r   rM   r   rK   r   rN      rO   zNumpy_hypergeometric.genericNrd   r   r   r   r   r          r   znp.random.laplaceznp.random.logisticznp.random.lognormalznp.random.normalc                   @   sF   e Zd Zdd eD Zedd eD 7 Zedd eD 7 Zdd ZdS )Numpy_normalc                 C   rS   r   rT   rU   r   r   r   rW      rX   zNumpy_normal.<listcomp>c                 C   r]   r   rT   rU   r   r   r   rW      r^   c                 C   r}   r   rT   rU   r   r   r   rW      r~   c                    rG   )Nc                    r_   r   rI   )locr   r1   rK   r   r   r'      rb   z#Numpy_normal.generic.<locals>.typer)NNNr   rM   r   rK   r   rN      rO   zNumpy_normal.genericNrt   r   r   r   r   r      s
    r   znp.random.gammac                   @   r\   )Numpy_gammac                 C   rS   r   rT   rU   r   r   r   rW      rX   zNumpy_gamma.<listcomp>c                 C   r]   r   rT   rU   r   r   r   rW      r^   c                    rG   )Nc                    r_   r   rI   )r   r   r1   rK   r   r   r'      rb   z"Numpy_gamma.generic.<locals>.typerrc   r   rM   r   rK   r   rN      rO   zNumpy_gamma.genericNrt   r   r   r   r   r      re   r   znp.random.triangularc                   @   rl   )Random_ternary_distributionc                 C   rg   r   rT   rU   r   r   r   rW      rh   z&Random_ternary_distribution.<listcomp>c                    rG   )Nc                    r   r   rI   )leftmoderightr1   rK   r   r   r'      r   z2Random_ternary_distribution.generic.<locals>.typerr   r   rM   r   rK   r   rN      rO   z#Random_ternary_distribution.genericNrt   r   r   r   r   r      r   r   znp.random.betaznp.random.fznp.random.gumbelznp.random.uniformznp.random.vonmisesznp.random.waldzrandom.betavariatezrandom.gammavariatezrandom.gausszrandom.lognormvariatezrandom.normalvariatezrandom.uniformzrandom.vonmisesvariatezrandom.weibullvariatec                   @   rl   )Random_binary_distributionc                 C   rS   r   rT   rU   r   r   r   rW      rX   z%Random_binary_distribution.<listcomp>c                    rG   )Nc                    r_   r   rI   )rs   br1   rK   r   r   r'     rb   z1Random_binary_distribution.generic.<locals>.typerr   r   rM   r   rK   r   rN     rO   z"Random_binary_distribution.genericNrt   r   r   r   r   r      s    r   znp.random.chisquareznp.random.paretoznp.random.powerznp.random.standard_gammaznp.random.standard_tznp.random.weibullzrandom.expovariatezrandom.paretovariatec                   @   rl   )Random_unary_distributionc                 C   r]   r   rT   rU   r   r   r   rW     r^   z$Random_unary_distribution.<listcomp>c                    rG   )Nc                    rq   r   rI   rr   rK   r   r   r'     r(   z0Random_unary_distribution.generic.<locals>.typerr   r   rM   r   rK   r   rN     rO   z!Random_unary_distribution.genericNrt   r   r   r   r   r     s    	r   znp.random.standard_cauchyznp.random.standard_normalznp.random.standard_exponentialc                   @   rl   )Random_nullary_distributionc                 C   r}   r   rT   rU   r   r   r   rW     r~   z&Random_nullary_distribution.<listcomp>c                    rG   )Nc                    rH   r   rI   rJ   rK   r   r   r'   !  rL   z2Random_nullary_distribution.generic.<locals>.typerr   r   rM   r   rK   r   rN      rO   z#Random_nullary_distribution.genericNrt   r   r   r   r   r     s    r   zrandom.triangularc                   @   s,   e Zd Zdd eD Zedd eD 7 ZdS )Random_triangularc                 C   rS   r   rT   rU   r   r   r   rW   (  rX   zRandom_triangular.<listcomp>c                 C   rg   r   rT   rU   r   r   r   rW   )  rh   N)r5   r6   r7   ru   r>   r   r   r   r   r   &  s    r   )`r"   numpyr!   
numba.corer   	templatesr   r   r   r   r   r	   numba.np.numpy_supportr
   numba.core.overload_gluer   registryregisterinferregister_globalinfer_globalregister_attrinfer_getattrsortedsetr   rp   r[   rP   ru   r   r   r:   getrandbitsrC   rF   random_samplesampleranfrandintrR   	randrangerf   seedri   	geometric	logserieszipfrm   binomialnegative_binomialrw   poissonrz   exponentialrayleighr|   hypergeometricr   laplacelogistic	lognormalnormalr   gammar   
triangularr   betafgumbeluniformvonmiseswaldbetavariategammavariategausslognormvariatenormalvariatevonmisesvariateweibullvariater   	chisquareparetopowerstandard_gamma
standard_tweibullexpovariateparetovariater   standard_cauchystandard_normalstandard_exponentialr   r   r   r   r   r   <module>   s      
		
		,	 		