U
    Mh                     @  s*  d dl mZ d dlZd dlZd dlZd dlZd dlZd dlZd dlm	Z	m
Z
mZmZmZmZmZmZmZ d dlmZ d dlZd dlmZm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"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z* dd	l+m,Z, e-e.Z/d
ddgZ0edej1eZ2G dd de3Z4dd Z5dd Z6dddddZ7dddddZ8ee9e:ej1f Z;ee<ef Z=ee;e=f Z>e	ej1gej1f Z?e	ej1ej1gej1f Z@e	egef ZAe	eegef ZBee?eAf ZCee@eBf ZDejEddG dd
 d
ee2 ZFG dd  d ZGG d!d deGZHd&d"d#d
d$d%dZIdS )'    )annotationsN)	CallableDictGenericOptionaloverloadSupportsFloatTYPE_CHECKINGTypeVarUnion)	TypeGuard)BooleanBooleanAtom)dtype_to_type   )_keep_floatFloatTrueDivFloorDiv
IntTrueDivOpaqueUnaryFn_expOpaqueUnaryFn_logOpaqueUnaryFn_sqrtPowByNaturalRoundDecimal
RoundToIntsafe_powToFloatTruncToFloat
TruncToInt)sympy_interpValueRangesValueRangeAnalysisbound_sympy_Tc                   @  s   e Zd ZdS )ValueRangeErrorN)__name__
__module____qualname__ r(   r(   Q/var/www/html/venv/lib/python3.8/site-packages/torch/utils/_sympy/value_ranges.pyr$   6   s   r$   c                 C  s   t | tr| rtjS tjS t | tr.t| S t | trbt	| rX| dkrPtj
S tj
 S t| S t | tjr| js|t| | tjkst| S t | tr| S tdt|  d|  d S )Nr   znot simple sympy type z: )
isinstanceboolsympytruefalseintIntegerfloatmathisinfooFloatExprZ	is_numberAssertionErrornanr   type)er(   r(   r)   simple_sympify<   s    






r;   c                 C  sL   t | tjr$t |tjst| |kS t | tr8t |ts<t| oD|  S d S N)r*   r,   r6   r7   SympyBooleanlowerupperr(   r(   r)   sympy_generic_leU   s
    rA   zValueRanges[_T]z$TypeGuard[ValueRanges[SympyBoolean]])vrreturnc                 C  s   | j S r<   is_boolrB   r(   r(   r)   
vr_is_bool_   s    rG   z"TypeGuard[ValueRanges[sympy.Expr]]c                 C  s   | j  S r<   rD   rF   r(   r(   r)   
vr_is_exprc   s    rH   T)frozenc                   @  s  e Zd ZU er,eej Zee Z	e
ee	f Zded< ded< ded< ded< ded< dd	d
dZedddddddZedddddddZddddddZdd	ddZdddddZdd Zd d	dd Zedddd!d"d#Zedddd!d$d#Zd%d%d%d!d&d#Zedddd!d'd(Zedddd!d)d(Zd%d%d%d!d*d(Zdd	d+d,Zedd	d-d.Zedd	d/d0Zeed1d2d3d4d5Zeed6d7d3d8d5Zed9d%d3d:d5Zed1d;d2d<d=d>Zeed1d;d2d<d?d@Zeed6dAd7d<dBd@Zed9dCd%d<dDd@Zed1d;d2d<dEdFZed1d;d2d<dGdHZeed1d1dId2dJdKdLZeed6d6dMd7dJdNdLZed9d9dOd%dJdPdLZe dQdR Z!dSS )Tr    r#   r?   r@   r+   rE   is_intis_floatstr)rC   c                 C  s   d| j  d| j dS )NzVR[z, ]r>   selfr(   r(   r)   __repr__   s    zValueRanges.__repr__zValueRanges[sympy.Expr]ExprInNone)rO   r?   r@   rC   c                 C  s   d S r<   r(   rO   r?   r@   r(   r(   r)   __init__   s    zValueRanges.__init__zValueRanges[SympyBoolean]BoolInc                 C  s   d S r<   r(   rS   r(   r(   r)   rT      s    AllIn)r?   r@   rC   c              
   C  s&  t |}t |}z$t||s2td| d| dW n8 tk
rl } ztd| d| |W 5 d }~X Y nX t| d| t| d| t| dt|t | jrt|tst	||ft| d	| j ot|t
jpt|t
j t| d
| j o| j  | js"| js"| js"t	||fd S )NzInvalid ranges [:rM   zCould not compare z <= r?   r@   rE   rJ   rK   )r;   rA   r$   	TypeErrorobject__setattr__r*   r=   rE   r7   r,   r0   rJ   rK   )rO   r?   r@   r:   r(   r(   r)   rT      s(    
(	c                 C  s2   t | r| S | t kr t S td|  d S )Nznot bool like )rG   r    unknownunknown_boolr7   rN   r(   r(   r)   boolify   s
    zValueRanges.boolify)xrC   c                 C  s   t || S r<   )r    wrapissubset)rO   r^   r(   r(   r)   __contains__   s    zValueRanges.__contains__c                 C  s   t |j| jot | j|jS r<   )rA   r?   r@   rO   otherr(   r(   r)   r`      s     zValueRanges.issubsetc                 C  s   | |@ S )z1Given two ValueRanges, returns their intersectionr(   rb   r(   r(   r)   tighten   s    zValueRanges.tighten)rO   rc   rC   c                 C  s   d S r<   r(   rb   r(   r(   r)   __and__   s    zValueRanges.__and__c                 C  s   d S r<   r(   rb   r(   r(   r)   re      s    AllVRc                 C  s   |t  kr| S | t  kr |S | j|jks8t| |f| j|jksPt| |f| j|jksht| |f| jrt t| j|jt	| j
|j
S t t| j|jt| j
|j
S d S r<   )r    r[   rE   r7   rJ   rK   r,   Orr?   Andr@   MaxMinrb   r(   r(   r)   re      s       c                 C  s   d S r<   r(   rb   r(   r(   r)   __or__   s    zValueRanges.__or__c                 C  s   d S r<   r(   rb   r(   r(   r)   rk      s    c                 C  s~   t  | |fkrt  S | j|jks0t| |f| jrXt t| j|jt| j|jS t t	| j|jt
| j|jS d S r<   )r    r[   rE   r7   r,   rh   r?   rg   r@   rj   ri   rb   r(   r(   r)   rk      s      c                 C  s   | j | jkS r<   r>   rN   r(   r(   r)   is_singleton	  s    zValueRanges.is_singletonc                   C  s   t tj tjS r<   r    r,   r4   r(   r(   r(   r)   r[     s    zValueRanges.unknownc                   C  s   t tjtjS r<   r    r,   r.   r-   r(   r(   r(   r)   r\     s    zValueRanges.unknown_boolzUnion[ExprIn, ExprVR]ExprVR)argrC   c                 C  s   d S r<   r(   rp   r(   r(   r)   r_     s    zValueRanges.wrapzUnion[BoolIn, BoolVR]BoolVRc                 C  s   d S r<   r(   rq   r(   r(   r)   r_     s    zUnion[AllIn, AllVR]c                 C  s4   t | tr| S t | tr*t| r*t S t| | S r<   )r*   r    r1   r2   isnanr[   rq   r(   r(   r)   r_      s
    
ExprFn)r^   fnrC   c                 C  s    t | } t || j|| jS )z#Increasing: x <= y => f(x) <= f(y).r    r_   r?   r@   r^   ru   r(   r(   r)   increasing_map)  s    
zValueRanges.increasing_mapc                 C  s   d S r<   r(   rw   r(   r(   r)   decreasing_map/  s    zValueRanges.decreasing_mapBoolFnc                 C  s   d S r<   r(   rw   r(   r(   r)   ry   4  s    AllFnc                 C  s    t | } t || j|| jS )z#Decreasing: x <= y => f(x) >= f(y).)r    r_   r@   r?   rw   r(   r(   r)   ry   9  s    
c                 C  s4   t | } || j}|| j}t t||t||S )zIt's increasing or decreasing.)r    r_   r?   r@   minmax)r^   ru   lur(   r(   r)   monotone_map@  s    


zValueRanges.monotone_mapc                 C  s>   t | } d| kr.t dt|| j|| jS t | |S dS )z$Fn is convex and has a minimum at 0.r   N)r    r_   r}   r?   r@   r   rw   r(   r(   r)   convex_min_zero_mapH  s    
zValueRanges.convex_min_zero_mapExprFn2)r^   yru   rC   c                 C  s   d S r<   r(   r^   r   ru   r(   r(   r)   coordinatewise_increasing_mapQ  s    z)ValueRanges.coordinatewise_increasing_mapBoolFn2c                 C  s   d S r<   r(   r   r(   r(   r)   r   X  s    AllFn2c                 C  s4   t | t | } }t || j|j|| j|jS )z
        It's increasing on each coordinate.

        Mathematically:
        For every 1 <= i <= n and x_i <= y_i we have that
        f(x1, .., xn) <= f(x1, , yi, ..., xn)
        rv   r   r(   r(   r)   r   _  s
    c                   sR   |  ||  | }} fddt|j|jg|j|jgD }tt|t|S )z1It's increasing or decreasing on each coordinate.c                   s   g | ]\}} ||qS r(   r(   ).0abru   r(   r)   
<listcomp>t  s   z;ValueRanges.coordinatewise_monotone_map.<locals>.<listcomp>)r_   	itertoolsproductr?   r@   r    r|   r}   )clsr^   r   ru   Zproductsr(   r   r)   coordinatewise_monotone_mapp  s
    
z'ValueRanges.coordinatewise_monotone_mapN)"r%   r&   r'   r	   r    r,   r6   ro   r=   rr   r   rf   __annotations__rP   r   rT   r]   ra   r`   rd   re   rk   rl   staticmethodr[   r\   r_   rx   ry   r   r   r   classmethodr   r(   r(   r(   r)   r    r   s   

/c                   @  s^  e Zd ZdZedd Zeded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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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d/d0 Zed1d2 Zed3d4 Zed5d6 Zed7d8 Z ed9d: Z!ed;d< Z"ed=d> Z#ed?d@ Z$edAdB Z%edCdD Z&edEdF Z'edGdH Z(edIdJ Z)edKdL Z*edMdN Z+edOdP Z,edQdR Z-edSdT Z.edUdV Z/edWdX Z0edYdZ Z1ed[d\ Z2ed]d^ Z3ed_d` Z4edadb Z5edcdd Z6dS )fSymPyValueRangeAnalysisz
    It gives bounds on a SymPy operator given bounds on its arguments
    See the function `bound_sympy` for a function that applies this logic to a full SymPy expression
    c                 C  s   t | tr|  st| j} t | tttf}|sHt | tt	j
t	jfsHtt | trt| r|tjkrnt S |jr|t S ttj d tjS |rt|}|| } n<|tjkrt | tstn"|jr| jr| jstn
| jstt| }|S )Nr   )r*   r    rl   r7   r?   r/   r1   r+   r   r,   r0   Numberr   r2   rs   torchr\   is_floating_pointr[   sysmaxsizer   	is_finiteZis_real
is_integerr_   )valuedtypeZ	is_pythontype_rr(   r(   r)   constant  s0    
 




z SymPyValueRangeAnalysis.constantNc                 C  s   |t jkrt| tS t S r<   )r   Zfloat64r    rx   r   r[   )r   r   	src_dtyper(   r(   r)   to_dtype  s    
z SymPyValueRangeAnalysis.to_dtypec                 C  s   t | tS r<   )r    rx   r   )r   r   r(   r(   r)   trunc_to_int  s    z$SymPyValueRangeAnalysis.trunc_to_intc                 C  s*   t | } |  } | jstt | tjS r<   )r    r_   r]   rE   r7   ry   r,   Not)r   r(   r(   r)   not_  s    

zSymPyValueRangeAnalysis.not_c                 C  s   t | |tjS r<   )r    r   r,   rg   r   r   r(   r(   r)   or_  s    zSymPyValueRangeAnalysis.or_c                 C  s   t | |tjS r<   )r    r   r,   rh   r   r(   r(   r)   and_  s    zSymPyValueRangeAnalysis.and_c                 C  sn   t | } t |}|  r<| r<| j|jkr<t tjS | j|jksT|j| jkr`t tjS t tjtjS r<   )r    r_   rl   r?   r,   r-   r@   r.   r   r(   r(   r)   eq  s    

zSymPyValueRangeAnalysis.eqc                 C  s   |  | ||S r<   )r   r   r   r   r   r(   r(   r)   ne  s    zSymPyValueRangeAnalysis.nec                 C  s~   t |}t |}|j|jks$t|jr<| | ||S |j|jk rTt tj	S |j|jkrlt tj
S t tj
tj	S d S r<   )r    r_   rE   r7   r   r   r@   r?   r,   r-   r.   r   r(   r(   r)   lt  s    

zSymPyValueRangeAnalysis.ltc                 C  s   |  ||S r<   )r   r   r(   r(   r)   gt  s    zSymPyValueRangeAnalysis.gtc                 C  s   |  | ||S r<   )r   r   r   r(   r(   r)   le  s    zSymPyValueRangeAnalysis.lec                 C  s   |  | ||S r<   )r   r   r   r(   r(   r)   ge  s    zSymPyValueRangeAnalysis.gec                 C  s   t | |ttjS r<   )r    r   r   operatoraddr   r(   r(   r)   r     s
      zSymPyValueRangeAnalysis.addc                 C  sP   t |}t |}|j|jks$t|jr6| ||S dd }t ||t|S )Nc                 S  s$   | dkr| S |dkr|S | | S d S Nr   r(   r   r(   r(   r)   safe_mul  s
    z-SymPyValueRangeAnalysis.mul.<locals>.safe_mul)r    r_   rE   r7   r   r   r   )r   r   r   r   r(   r(   r)   mul  s    

	zSymPyValueRangeAnalysis.mulc                 C  sf   t | } t |}d|ksHtj | ks2tj| krPtj |ksHtj|krPt  S t | |ttS d S r   )r    r_   r,   r4   r[   r   r   r   r   r(   r(   r)   int_truediv  s"    



  z#SymPyValueRangeAnalysis.int_truedivc                 C  sf   t | } t |}d|ksHtj | ks2tj| krPtj |ksHtj|krPt  S t | |ttS d S r   )r    r_   r,   r4   r[   r   r   r   r   r(   r(   r)   truediv  s"    



  zSymPyValueRangeAnalysis.truedivc                 C  sb   t | } t |}d|ksHtj | ksHtj| ksHtj |ksHtj|krPt  S t | |tS d S r   )r    r_   r,   r4   r[   r   r   r   r(   r(   r)   floordiv  s    



z SymPyValueRangeAnalysis.floordivc                   s   t |}t |}dd  dd }d|kr4t  S | rt|j||j||jkrtt | fddS |jdk rt  d dS |jdkrt dd S t d |j}t	d |j}t ||S n| |jd }t | |S d S )	Nc                 S  s$   t | t | }| dk r |d9 }|S )Nr   )abs)r   r   retr(   r(   r)   c_mod/  s    z*SymPyValueRangeAnalysis.mod.<locals>.c_modc                 S  s   | | }|j rt|S |S r<   )r   r,   r0   )r   r   r^   r(   r(   r)   c_div5  s    z*SymPyValueRangeAnalysis.mod.<locals>.c_divr   c                   s
    | S r<   r(   )r   r   Zy_valr(   r)   <lambda>A      z-SymPyValueRangeAnalysis.mod.<locals>.<lambda>r   )
r    r_   r[   rl   r   r?   r@   rx   r}   r|   )r   r^   r   r   r?   r@   r(   r   r)   mod)  s&    




zSymPyValueRangeAnalysis.modc                 C  s   |  | |||S r<   )r   r   )r   r   r   cr(   r(   r)   modular_indexingR  s    z(SymPyValueRangeAnalysis.modular_indexingc                 G  s   t  S r<   r    r[   )r   argsr(   r(   r)   &is_non_overlapping_and_dense_indicatorV  s    z>SymPyValueRangeAnalysis.is_non_overlapping_and_dense_indicatorc                   s   t |}t   | r8  r8t t|j jS |jdkr`t | t dtjd @ tS   r jd dkrt 	| fddS t 
| fddS n,t|j|j }t t| j t| jS d S )Nr   r      c                   s   t |  jS r<   r   r?   r^   r   r(   r)   r   l  r   z8SymPyValueRangeAnalysis.pow_by_natural.<locals>.<lambda>c                   s   t |  jS r<   r   r   r   r(   r)   r   p  r   )r    r_   rl   r   r?   r   r   r   r   r   rx   r}   r@   )r   r   r   Zmax_baser(   r   r)   pow_by_naturalZ  s,    


   
 
z&SymPyValueRangeAnalysis.pow_by_naturalc                 C  s   t  S r<   r   r   r(   r(   r)   pow{  s    zSymPyValueRangeAnalysis.powc                 C  s.   t | } d| krt  S t | dd S dS )zENeeded as it's used in pow, but it won't appear on a SymPy expressionr   c                 S  s
   t d| S )N      ?)r   r   r(   r(   r)   r     r   z4SymPyValueRangeAnalysis.reciprocal.<locals>.<lambda>N)r    r_   r[   ry   r   r(   r(   r)   
reciprocal  s    
z"SymPyValueRangeAnalysis.reciprocalc                 C  s   t | tS r<   )r    r   r   r   r(   r(   r)   r     s    zSymPyValueRangeAnalysis.absc                 C  s   t | tS r<   )r    rx   r   r   r(   r(   r)   exp  s    zSymPyValueRangeAnalysis.expc                 C  s(   t | } | jdkrt  S t | tS r   )r    r_   r?   r[   rx   r   r   r(   r(   r)   log  s    

zSymPyValueRangeAnalysis.logc                 C  s   |  ||tjS r<   )
min_or_maxr,   rj   r   r(   r(   r)   minimum  s    zSymPyValueRangeAnalysis.minimumc                 C  s   |  ||tjS r<   )r   r,   ri   r   r(   r(   r)   maximum  s    zSymPyValueRangeAnalysis.maximumc                 C  s"   t | } t |}t | ||S r<   )r    r_   r   )r   r   ru   r(   r(   r)   r     s    

z"SymPyValueRangeAnalysis.min_or_maxc                 C  s   t |tjjjjS r<   )r    rx   r,   	functions
elementaryintegersfloorr   r^   r   r(   r(   r)   floor_to_int  s    z$SymPyValueRangeAnalysis.floor_to_intc                 C  s   t |tjjjjS r<   )r    rx   r,   r   r   r   ceilingr   r(   r(   r)   ceil_to_int  s     
z#SymPyValueRangeAnalysis.ceil_to_intc                 C  s   t |ttjjjjS r<   )r    rx   r   r,   r   r   r   r   r   r^   r(   r(   r)   r     s     zSymPyValueRangeAnalysis.floorc                 C  s   t |ttjjjjS r<   )r    rx   r   r,   r   r   r   r   r   r(   r(   r)   ceil  s     zSymPyValueRangeAnalysis.ceilc                   s.      st S  j  fdd}t||S )Nc                   s
   t |  S r<   )r   )numberndigitsr(   r)   r     r   z7SymPyValueRangeAnalysis.round_decimal.<locals>.<lambda>)rl   r    r[   r?   rx   )r   r   r   ru   r(   r   r)   round_decimal  s
    z%SymPyValueRangeAnalysis.round_decimalc                 C  s   t |tS r<   )r    rx   r   )r   r   r   r(   r(   r)   round_to_int  s    z$SymPyValueRangeAnalysis.round_to_intc                 C  s(   t | } | jdk rt  S t | tS r   )r    r_   r?   r[   rx   r   r   r(   r(   r)   sqrt  s    

zSymPyValueRangeAnalysis.sqrtc                 C  s   t |}t |}|  } |j|jks<t  ||fks<t|jrdt t|j|jt	|j
|j
S t t|j|jt|j
|j
S d S r<   )r    r_   r]   rE   r[   r7   r,   rh   r?   rg   r@   rj   ri   )r   r   r   r(   r(   r)   where  s    

 "zSymPyValueRangeAnalysis.wherec                 C  s   |  }| |fS r<   )r]   r   r(   r(   r)   expr_cond_pair  s    z&SymPyValueRangeAnalysis.expr_cond_pairc                  G  s6   d }| D ](\}}t j|kr|d kr(|}q||B }q|S r<   )r,   r-   )rangesZ
init_rangeZ
expr_rangeZ
cond_ranger(   r(   r)   	piecewise(  s    

z!SymPyValueRangeAnalysis.piecewisec                 C  s
   t ddS Ng      r   r    r   r(   r(   r)   cos3  s    zSymPyValueRangeAnalysis.cosc                 C  s   t dtjS )Ng        rm   r   r(   r(   r)   cosh:  s    zSymPyValueRangeAnalysis.coshc                 C  s
   t ddS r   r   r   r(   r(   r)   sinF  s    zSymPyValueRangeAnalysis.sinc                 C  s   t tj tjS r<   rm   r   r(   r(   r)   sinhL  s    zSymPyValueRangeAnalysis.sinhc                 C  s   t tj tjS r<   rm   r   r(   r(   r)   tanQ  s    zSymPyValueRangeAnalysis.tanc                 C  s   t tj tjS r<   rm   r   r(   r(   r)   tanhU  s    zSymPyValueRangeAnalysis.tanhc                 C  s   t tj tjS r<   rm   r   r(   r(   r)   asinZ  s    zSymPyValueRangeAnalysis.asinc                 C  s   t tj tjS r<   rm   r   r(   r(   r)   acosd  s    zSymPyValueRangeAnalysis.acosc                 C  s   t tj tjS r<   rm   r   r(   r(   r)   atann  s    zSymPyValueRangeAnalysis.atanc                 C  s   t | tS r<   )r    rx   r   r   r(   r(   r)   truncs  s    zSymPyValueRangeAnalysis.trunc)N)7r%   r&   r'   __doc__r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r(   r(   r(   r)   r   {  s   
%




	










(


 
,















	








	
	
r   c                   @  s   e Zd Zdd Zedd Zedd Zddd	d
dZd"ddZdd Z	e
dd Zed#dddddZedd Zedd Ze
dd Ze
dd Zd d! ZdS )$r!   c                 C  s&   d| _ d}|D ]}t| || j qd S )Nr!   )xorlogical_and
logical_orZlogical_not)namesetattrbool_handler)rO   Zboolean_operatorsopr(   r(   r)   rT   y  s    zValueRangeAnalysis.__init__c                  O  s   t tjtjS r<   rn   r   kwargsr(   r(   r)   r     s    zValueRangeAnalysis.bool_handlerc                  O  s   t  S r<   r   r   r(   r(   r)   default_handler  s    z"ValueRangeAnalysis.default_handlerrL   
sympy.Expr)r   indexc                 C  s   t  S r<   r   )rO   r   r   r(   r(   r)   load  s    zValueRangeAnalysis.loadNc                 C  s   d S r<   r(   )rO   r   r   r   moder(   r(   r)   store  s    zValueRangeAnalysis.storec                 C  s   t  S r<   r   )rO   r   r   r   Zreduction_typer   r   r(   r(   r)   	reduction  s    zValueRangeAnalysis.reductionc                 C  s   t |tst| ||S r<   )r*   r    r7   r   )r   r   r   r(   r(   r)   
index_expr  s    zValueRangeAnalysis.index_exprztorch.dtypezOptional[torch.dtype])r   r   c                 C  s   t | } |tjkrN|  r,t | jdkS d| kr@t tjS t tjtjS dd }| j	r|  r| jrndnd}t |||S t |d||d|S nt || j||| j
|S d S )Nr   c                 S  s:   |j rt| S zt| W S  tk
r4   |  Y S X d S r<   )r   r,   r5   r0   rX   )r^   r   r(   r(   r)   cast  s    
z)ValueRangeAnalysis.to_dtype.<locals>.castr   )r    r_   r   r+   rl   r?   r,   r-   r.   rE   r@   )r^   r   r   r  valr(   r(   r)   r     s    

zValueRangeAnalysis.to_dtypec                 C  s   t | dd S )Nc                 S  s
   t | dS )Nr   )r   r   r(   r(   r)   r     r   z+ValueRangeAnalysis.square.<locals>.<lambda>)r    r   r   r(   r(   r)   square  s    zValueRangeAnalysis.squarec                 C  s   t | tjS r<   )r    ry   r   negr   r(   r(   r)   r	    s    zValueRangeAnalysis.negc                 C  s&   |  ||}|t kr|S | |S r<   )r   r    r[   r   )r   r   r   r^   r(   r(   r)   truncdiv  s    zValueRangeAnalysis.truncdivc                 C  s   |  || |S r<   )r   r	  r   r(   r(   r)   sub  s    zValueRangeAnalysis.subc                 C  s   t d| | jS )Nzunhandled ValueRange op %s)r   debugr   )rO   r   r(   r(   r)   __getattr__  s    zValueRangeAnalysis.__getattr__)N)N)r%   r&   r'   rT   r   r   r   r  r  r  r   r  r   r  r	  r
  r  r  r(   r(   r(   r)   r!   x  s*   



 



r   z)Optional[Dict[sympy.Symbol, ValueRanges]])exprr   rC   c                 C  s   t d| | t| tjr$t| S |p*i }tjj	
 }|rR|jjrR|jjj|}| j|  }|ri }|D ]D}|jr|jrd}q|jrd}qtj }ntj }t|tj||< ql||}tt|| S )Nzbound_sympy(%s, %s)r   r   )r   r  r*   r,   r   r    r_   r   Z_guardsZTracingContextZtry_getZ	fake_modeZ	shape_envZvar_to_rangeZfree_symbolskeysr   Zis_positiveZis_nonnegativer2   infr   r   )r  r   contextZunbounded_varsZunbounded_rangessr?   r(   r(   r)   r"     s*    

)N)J
__future__r   dataclassesr   loggingr2   r   r   typingr   r   r   r   r   r   r	   r
   r   typing_extensionsr   r,   Zsympy.logic.boolalgr   r=   r   r   Ztorch._prims_commonr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   interpr   	getLoggerr%   r   __all__r6   r#   RuntimeErrorr$   r;   rA   rG   rH   r/   r1   rQ   r+   rU   rV   rt   r   rz   r   r{   r   	dataclassr    r   r!   r"   r(   r(   r(   r)   <module>   sV   ,@



  
    c 