o
    ^i                     @   s  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 d dlm	Z	 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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!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/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5m6Z6m7Z7m8Z8m9Z9 d dl:m;Z;m<Z<m=Z=m>Z> d dl?m@Z@mAZAmBZBmCZD d dlEmFZF d d	lGmHZHmIZI d
ZJdhZKeLejMejN ejO ZPdd ZQe jRde j!de jde jde j8de j*de jde jSde jTde jUde jde jdiZVe jRde j!de jde jde j8de j*de jde jSd e jTd!e jUd"e jd#e jd$iZWeVXe jYd%e jZd&e j[d'e j\d(e j]d)i eWXe jYd*e jZd+e j[d,e j\d-e j]d.i eVXe j^d/i G d0d1 d1e_Z`G d2d3 d3e`ZaG d4d5 d5eaZbG d6d7 d7e`Zcd8d9 Zdd:d; Zed<d= Zfd>d? Zgd@dA ZhdadCdDZidEdF ZjG dGdH dHZkdIdJ ZldbdKdLZmdMZndcdNdOZodcdPdQZpdRdS ZqdTdU ZrdVdW ZsG dXdY dYekZtG dZd[ d[ekZuG d\d] d]ekZvev Zweu Zxet Zyddd_d`ZzdS )e    N)
namedtuple)dedent)_jit_internal),ApplyAssertAssign	Attribute	AugAssignBinOpBreakClassDefConstContinueDeclDefDeleteDictCompDictLiteralDotsEmptyTypeAnnotationExprStmtFalseLiteralForIdentIfListCompListLiteralNoneLiteralParamPassPropertyRaiseReturnSelect	SliceExprStarredStmtStringLiteral	Subscript	TernaryIfTrueLiteralTupleLiteralUnaryOpVarWhileWithWithItem)_is_drop_fnFunctionModifiersis_static_fnshould_drop)get_source_lines_and_filemake_source_context	parse_def	ParsedDef)DATACLASS_MAGIC_METHODS)get_qualified_namemonkeytype_trace__jitprintc                 C   s   |  tp| tv S N)
startswith_reserved_prefix_reserved_namesname rD   M/var/www/html/RAG/RAG_venv/lib/python3.10/site-packages/torch/jit/frontend.pyis_reserved_nameT   s   rF   zfunction definitionsz	for loopszdel statementszclass definitionszwith statementszraise statements
assertionszimport statementszglobal variableszbreak statementszcontinue statementsdeffordelclasswithraiseassertimportfromglobalbreakcontinuezasync function definitionszasync for loopszasync with statementsz
try blocksznonlocal variablesz	async defz	async forz
async withtrynonlocalzannotated assignmentsc                   @   s   e Zd Zdd Zdd ZdS )FrontendErrorc                 C   s    || _ || _tj| j | _d S r>   )source_rangemsgtorch_CErrorReporterror_report)selfrW   rX   rD   rD   rE   __init__   s   zFrontendError.__init__c                 C   s   | j | j   S r>   )rX   r\   whatlstrip)r]   rD   rD   rE   __str__      zFrontendError.__str__N)__name__
__module____qualname__r^   ra   rD   rD   rD   rE   rV      s    rV   c                   @      e Zd ZdS )NotSupportedErrorNrc   rd   re   rD   rD   rD   rE   rg          rg   c                       s   e Zd Zd fdd	Z  ZS )UnsupportedNodeError c           	         sl   t |}tt|d}||j|j|j| }t||j}| d|r(|d nd d}t	 
|| d S )N rk   zaren't supported)typelennode_start_tokensget
make_rangelineno
col_offsetpretty_node_namesrc   superr^   )	r]   ctxoffending_nodereason	node_type	range_lenrW   feature_namerX   	__class__rD   rE   r^      s   zUnsupportedNodeError.__init__)rk   )rc   rd   re   r^   __classcell__rD   rD   r|   rE   rj      s    rj   c                   @   rf   )FrontendTypeErrorNrh   rD   rD   rD   rE   r      ri   r   c                    s    fdd|D }t |S )Nc                       g | ]}t  |qS rD   )build_withitem).0irv   rD   rE   
<listcomp>       z#build_withitems.<locals>.<listcomp>)list)rv   itemsrD   r   rE   build_withitems   s   r   c                    s     fdd|D }t td |S )Nc                    r   rD   )
build_stmt)r   sr   rD   rE   r      r   zbuild_stmts.<locals>.<listcomp>)r   filter)rv   stmtsrD   r   rE   build_stmts   s   r   c              	   C   s   t j| dd d}t| dg }g }|D ]I}|d |vr\t|d js\t|d jd|d  d|d	}|d jrGt|d jd|d  d
|d	nd}|t|	 t
|	 |d || q|S )aj  
    Get a list of Property objects representing the properties of a class.

    Args:
        cls:  The class to get properties of.
        self_name: The name of the class that the properties should belong to.
    Returns:
        A list of Property objects corresponding to the properties of cls. Property
        here refers to the subclass of TreeView.
    c                 S   s
   t | tS r>   )
isinstancepropertymrD   rD   rE   <lambda>   s   
 z&get_class_properties.<locals>.<lambda>	predicate__jit_unused_properties__r      ___getter)	self_name_setterN)inspect
getmembersgetattrr4   fgetget_jit_deffsetappendr    ranger   )clsr   propsunused_properties
propertiespropgettersetterrD   rD   rE   get_class_properties   s"   "r   c                    sT   g   fdd}|j D ]}t|tjr|tj| qt|tjr'|tj| q S )Nc                    s,   z  | | W d S  ty   Y d S w r>   )r   rg   )builderentryassignsrv   rD   rE   maybe_build_assign   s
   z-get_class_assigns.<locals>.maybe_build_assign)bodyr   astr   StmtBuilderbuild_Assign	AnnAssignbuild_AnnAssign)rv   cls_astr   r   rD   r   rE   get_class_assigns   s   
r   c                    sV  t j  fddd} fddt tjj \}}}d|}t|}t	
|}|jd }	t|	t	js8J t rrdd	 |	jD }
t|D ](\}\}}t|}|rq||
vrq| }||f||< t |}tj||j qIfd
d|D }t }t|ddd t|ddd  }t||||d}t||	}t||	|||S )ap  Get definitions for each method within the current class independently.

    Args:
        cls: The class to get definition of.
        self_name: The name of the class that the properties should belong to.

    Returns:
        torch._C._jit_tree_views.ClassDef: A representation of the class,
            the methods in the class and their definition as a tree.
    c                    s8   t | s
t | ot | j o| j jv ot|  S r>   )r   ismethod
isfunctionr3   rc   __dict__r1   r   r   rD   rE   r     s    

z#get_jit_class_def.<locals>.<lambda>r   c                    s   t | ot| dd  kS )N__self__)r   r   r   )fnr   rD   rE   is_classmethod  s   z)get_jit_class_def.<locals>.is_classmethodrk   r   c                 S   s(   h | ]}t |tjr|jtv r|jqS rD   )r   r   FunctionDefrC   r9   r   methodrD   rD   rE   	<setcomp>  s    

z$get_jit_class_def.<locals>.<setcomp>c                    s$   g | ]\}}t || |d qS ))r   r   )r   )r   rC   obj)r   r   rD   rE   r   -  s    z%get_jit_class_def.<locals>.<listcomp>
r   F)r   r   r5   rY   rZ   r[   
call_stackjoinr   r   parser   r   r   dataclassesis_dataclass	enumerater9   rp   r   r   loadercachesourcer   rn   splitr6   r   build_class_def)r   r   methodssourcelinesfile_linenofilenamer   
dedent_srcpy_ast	class_ast	overridesr   rC   _synthesizer_fn
parsed_deffuncmethod_defsr   leading_whitespace_lenrv   r   rD   )r   r   r   rE   get_jit_class_def   sJ   










r   Fc                 C   sX  t | ts	t| n| }tjj|j}|jj	d }|r6|j
j
d j}t| d| j	d }|j	d| t| rtd}	t|	j	dksOt |	j	d tjs[td|j d|j |	j	d }
|
j	|_	d |j
_|j
_|j
j
|j
j D ]
}|
j
j
d j|_qtt| rd|_d|_tjj }d}trt | tst| }||}t |j!|||||dS )	af  
    Build a JIT AST (TreeView) from the given function.

    Args:
        fn: A function object to compile or a pre-parsed ParsedDef object
        def_name: The name to give to the resulting AST object. This is not
            always the same as `fn.__name__`, for example:
                def _forward(self):
                    ...
                forward = _forward
            In this case, the `__name__` attribute of the function object is "_forward",
            but we want the result AST to have the name "forward".
        self_name: If this function is a method, what the type name of `self` is.
    r    = zLdef unused_fn(self: Any):
	raise RuntimeError("Cannot call @unused methods")r   z&Expected a single top-level function: :N)r   pdt_arg_types)"r   
_ParsedDefr7   rY   jitannotationsget_type_liner   r   r   argsargr   insertr4   rn   r   RuntimeErrorr   r   kwargvararg
kwonlyargs
annotationr1   returnstype_comment_script_get_type_trace_dbr;   r:   get_args_types	build_defrv   )r   def_namer   r   r   	type_linefn_defarg_nameassign_stmtunused_fn_def
unused_defr   type_trace_dbr   qualnamerD   rD   rE   r   >  sL   

r   c                 C   sz   t | jd jtjr;| jd jj}t |tjr;|j}|j}|dkr;t |tjr;|jdkr;t |jtj	r;|jj
dkr;dS dS )Nr   _IgnoreContextManagerr   rY   TF)r   r   context_exprr   Callr   r   attrvalueNameid)stmtfunction	attr_name
attr_valuerD   rD   rE   #is_torch_jit_ignore_context_manager  s   
r	  c                   @   s   e Zd Zdd ZdS )Builderc                 C   s0   t | d|jj d }|d u rt|||||S )Nbuild_)r   r}   rc   rj   )r]   rv   noder   rD   rD   rE   __call__  s   

zBuilder.__call__N)rc   rd   re   r  rD   rD   rD   rE   r
    s    r
  c                 C   s:   |  |j|j|jtd }tt||dd |D ||S )NrK   c                 S   s   g | ]}t |qS rD   )r&   r   rD   rD   rE   r     s    z#build_class_def.<locals>.<listcomp>)rq   rr   rs   rn   r   r   )rv   py_defr   r   r   r   rrD   rD   rE   r     s   r   c                 C   s   |j }| |j|j|jtd }t| |j||}d }	t|dd d ur)t| |j	}	t
|||	}
|d u}|d urEtj|}tj|
||}
tt|||
t| |S )NrH   r   )r   rq   rr   rs   rn   build_param_listr   r   
build_exprr   r   rY   rZ   parse_type_commentmerge_type_from_type_commentr   r   r   )rv   r  r   r   r   r   r   r  
param_listreturn_typedecl	is_methodtype_comment_declrD   rD   rE   r     s    r   zfCompiled functions can't take variable number of arguments or use keyword-only arguments with defaultsc           
         s
  |j d ur|j } |j|jd |jt|j }t|t|jd ur<|j} |j|jd |jt|j }t|tt|j	dkrY|j	D ]}|d urXt
 | }t|tqFfdd|jD }fdd|jD } fdd|D }	|	 fdd|D 7 }	|	S )Nr   r   c                    0   g | ]}| rt  |j r |j nd fqS r>   boolr   r   r   r   rD   rE   r         z$build_param_list.<locals>.<listcomp>c                    r  r>   r  r  r  rD   rE   r     r  c              	      "   g | ]\}}t  |d |dqS )F
kwarg_onlypdt_arg_typebuild_paramr   r   arg_typerv   r   rD   rE   r         c              	      r  )Tr   r#  r%  r'  rD   rE   r     r(  )r   rq   rr   rs   rn   r   rg   _vararg_kwarg_errr   kw_defaultsr  r   r   r   )
rv   py_argsr   r   expr	ctx_ranger   arg_and_typesarg_and_types_kwonlyargsresultrD   )rv   r   r   rE   r    s>   







	
r  c                 C   s   |j }| |j|j|jt| }t|dd d ur t| |j}n|r*tt	||}n|d ur:|dkr:tt	||}nt
|}t|t	|||S )Nr   r]   )r   rq   rr   rs   rn   r   r  r   r-   r   r   r   )rv   py_argr   r!  r"  rC   r  annotation_exprrD   rD   rE   r$    s   r$  c                    sb  t dddg t dddg fdd}dd }d	d
 }dd }||jd jj\}}d|| | }d| }	|	dddd |D  d 7 }	||\}
}|	|
d 7 }	t|	jd }|j|_t|jd }|j| dt	| d}t

t }t|| || t |< ||}d| d|| d }t|dkr|d | }n|}t|jd }|S )N	InputTyperC   ann
OutputTypec                    sf   g }g }| D ](}|j }|jj}|d\}}|dkr"| || |dkr.||| q||fS )Nr   inpout)r   r  r   r   )r   inputsoutputsr   var_namevar_annvar_decl_typer3  r5  rD   rE   process_ins_outs  s   z6build_ignore_context_manager.<locals>.process_ins_outsc                 S   s    t dd| j}| d|j S )Nz[^a-zA-Z0-9_]r   )resubr   rr   )rv   r  r   rD   rD   rE   create_unique_name_ext  s   z<build_ignore_context_manager.<locals>.create_unique_name_extc                 S   s   d}d}t | dkr|d7 }t | dkr"d| d j }|| d j7 }t | dkrFd}|dd	d
d | D  d 7 }|d	dd | D 7 }||fS )Nrk   zreturn r   z -> Noner   z -> z	 -> tuple[, c                 S      g | ]}|j qS rD   )r4  r   varrD   rD   rE   r   '      zObuild_ignore_context_manager.<locals>.build_return_ann_stmt.<locals>.<listcomp>]c                 S   rD  rD   rB   rE  rD   rD   rE   r   (  rG  )rn   r4  rC   r   )r9  return_type_annreturn_statement_strrD   rD   rE   build_return_ann_stmt  s    z;build_ignore_context_manager.<locals>.build_return_ann_stmtc                 S   s   d dd | D S )NrC  c                 S   rD  rD   rB   r  rD   rD   rE   r   ,  rG  zDbuild_ignore_context_manager.<locals>.build_args.<locals>.<listcomp>)r   )r   rD   rD   rE   
build_args+  rb   z0build_ignore_context_manager.<locals>.build_argsr   func_ignore_z
def (rC  c                 S   s   g | ]
}|j d  |j qS )z :)rC   r4  rE  rD   rD   rE   r   4  s    z0build_ignore_context_manager.<locals>.<listcomp>)z: passz# Backward compat: These used to be imported into the outer global scope so some
# code may still expect them.
from typing import List, Dict, Tuple

@torch.jit.ignore
r   ztorch.jit.frontend.r   )r   r   r   keywordsr   r   r   r   r   unparsecopyglobalsexecrn   )rv   r  r>  rA  rK  rL  r8  r9  ignore_function_nameignore_function_str
return_annreturn_stmtignore_functionignore_func_strgassign_str_lhsassign_str_rhs
assign_str
assign_astrD   r=  rE   build_ignore_context_manager  s@   
r`  c                 C   s*   | du ri S t | }dd |j D S )a   
    Get a dictionary of default arguments for a function.

    Args:
        fn: Callable - The function to inspect for default arguments.
    Returns:
        (Dict[str, Any]): mapping argument names to their default values if
        :attr:`fn` is not None, else empty dictionary.
    Nc                 S   s&   i | ]\}}|j tjjur||j qS rD   )defaultr   	Parameterempty)r   kvrD   rD   rE   
<dictcomp>p  s
    z$get_default_args.<locals>.<dictcomp>)r   	signature
parametersr   )r   rg  rD   rD   rE   get_default_argsa  s   

ri  c                    s(   t j  fddd}dd |D }|S )aC  
    Get default arguments for all methods in a class (except for static methods).

    Args:
        cls: type - The class type to inspect for default arguments.
    Returns:
        A Dict[str, Dict[str, Any]] which maps each method name to a Dict[str, Any]
        that maps each argument name to its default value.
    c                    s.   t | s
t | ot | j o| j jv S r>   )r   r   r   r3   rc   r   r   r   rD   rE   r     s    z,get_default_args_for_class.<locals>.<lambda>r   c                 S   s   i | ]	\}}|t |qS rD   )ri  )r   method_namemethod_implrD   rD   rE   rf    s    z.get_default_args_for_class.<locals>.<dictcomp>)r   r   )r   r   defaultsrD   r   rE   get_default_args_for_classw  s   
	rm  c                   @   s   e Zd Zedd ZdS )WithItemBuilderc                 C   sZ   |j j}|j j}|tttj  }|j}| |||}t	|t
| |j |r*t
| |S d S r>   )r   rr   rs   rn   rt   r   r/   optional_varsrq   r0   r  )rv   itemrr   startendop_varsr  rD   rD   rE   r     s   
zWithItemBuilder.build_withitemN)rc   rd   re   staticmethodr   rD   rD   rD   rE   rn    s    rn  c                   @   s  e Zd Zejdejdejdejdejdej	dej
dejdejd	ejd
ejdi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 d,S )-r   +-*/%|&^<<>>**c                 C   s$   |j }|jjdkrd S tt| |S )NStr)r  r}   rc   r   r  )rv   r  r  rD   rD   rE   
build_Expr  s   zStmtBuilder.build_Exprc                    s*   t  |j} fdd|jD }t||S )Nc                    r   rD   r  )r   xr   rD   rE   r     r   z,StmtBuilder.build_Assign.<locals>.<listcomp>)r  r  targetsr   )rv   r  rhslhsrD   r   rE   r     s   
zStmtBuilder.build_Assignc                 C   s   |j d u rt| |ddt|jtjkrR|jj jdkrR| jdkrR|j}|t	d|jj
  }t|jdr?|t	d|jj 7 }| |j||}td| j d	| t| |j }t| |j}t| |j}t|g||S )
Nzwithout assigned value)rx   r]   r^   zself.r  z: zKType annotations on instance attributes must be declared in __init__, not 'z': )r  rj   rm   targetr   r   r  funcnamers   rn   r  hasattrr   rq   rr   
ValueErrorr  r   )rv   r  rq  rr  srr  r  the_typerD   rD   rE   r     s*   

zStmtBuilder.build_AnnAssignc                    s6     |j|j|jtd }t| fdd|jD S )NrJ   c                    r   rD   r  )r   r  r   rD   rE   r     r   z,StmtBuilder.build_Delete.<locals>.<listcomp>)rq   rr   rs   rn   r   r  rv   r  r  rD   r   rE   build_Delete  s   zStmtBuilder.build_Deletec                 C   s>   |  |j|j|jtd }t||jd u rd S t| |jS )Nreturn)rq   rr   rs   rn   r"   r  r  r  rD   rD   rE   build_Return  s   "zStmtBuilder.build_Returnc                 C   s2   |  |j|j|jtd }t| |j}t||S )NrM   )rq   rr   rs   rn   r  excr!   )rv   r  r  r,  rD   rD   rE   build_Raise  s   
zStmtBuilder.build_Raisec                 C   sN   |  |j|j|jtd }t| |j}|jd urt| |jnd }t|||S )NrN   )rq   rr   rs   rn   r  testrX   r   )rv   r  r  r  rX   rD   rD   rE   build_Assert  s   zStmtBuilder.build_Assertc                 C   sf   t | |j}t | |j}t|j}|tjv rtj| }ntt| |	 j
dddd|j t|||S )N=)r   )offsetsz*unsupported kind of augmented assignment: )r  r  r  rm   opr   augassign_maprg   find_beforer   rq  rc   r	   )rv   r  r  r  r  op_tokenrD   rD   rE   build_AugAssign  s   

zStmtBuilder.build_AugAssignc                 C   sH   |j rtd d| |j|j|jtd }t|t| |jt	| |j
S )Nz-else branches of while loops aren't supportedwhile)orelserg   rq   rr   rs   rn   r.   r  r  r   r   r  rD   rD   rE   build_While  s   zStmtBuilder.build_Whilec                 C   sV   |  |j|j|jtd }|jrt|dt|t| |jgt| |j	gt
| |jS )NrI   z+else branches of for loops aren't supported)rq   rr   rs   rn   r  rg   r   r  r  iterr   r   r  rD   rD   rE   	build_For  s   

zStmtBuilder.build_Forc                 C   sB   |  |j|j|jtd }t|t| |jt| |jt| |j	S )Nif)
rq   rr   rs   rn   r   r  r  r   r   r  r  rD   rD   rE   build_If  s   


zStmtBuilder.build_Ifc                    sZ     |j|j|jtd }|jrt|d fdd|jD }ttt	t
|d|g S )Nr=   z?print statements with non-default destinations aren't supportedc                    r   rD   r  )r   valr   rD   rE   r   0  r   z+StmtBuilder.build_Print.<locals>.<listcomp>)rq   rr   rs   rn   destrg   valuesr   r   r-   r   )rv   r  r  r   rD   r   rE   build_Print)  s   zStmtBuilder.build_Printc                 C   $   |  |j|j|jtd }t|S )Npass)rq   rr   rs   rn   r   r  rD   rD   rE   
build_Pass3     zStmtBuilder.build_Passc                 C   r  )NrR   )rq   rr   rs   rn   r   r  rD   rD   rE   build_Break8  r  zStmtBuilder.build_Breakc                 C   r  )NrS   )rq   rr   rs   rn   r   r  rD   rD   rE   build_Continue=  s   zStmtBuilder.build_Continuec                 C   sT   |  |j|j|jtd }t|rt| |}t| |S t|t| |j	t
| |jS )NrL   )rq   rr   rs   rn   r	  r`  r   r/   r   r   r   r   )rv   r  r  r_  rD   rD   rE   
build_WithD  s
   

zStmtBuilder.build_WithN)!rc   rd   re   r   AddSubMultDivModBitOrBitAndBitXorLShiftRShiftPowr  rt  r  r   r   r  r  r  r  r  r  r  r  r  r  r  r  r  rD   rD   rD   rE   r     sZ    
	










	
	


r   c                   @   s  e Zd Zejdejdejdejdejdej	dej
dejdejd	ejd
ejdejdiZdeej< ejdejdejdiZejdejdiZejdejdejdejdejdejdejdej dej!dej"di
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$d0d1 Z/e$d2d3 Z0e$d4d5 Z1e$d6d7 Z2e$d8d9 Z3e$d:d; Z4e$d<d= Z5e$d>d? Z6e$d@dA Z7e$dBdC Z8e$dDdE Z9e$dFdG Z:dHS )IExprBuilderru  rv  rw  rx  r  ry  z//r{  r|  rz  r}  r~  @not~andorz==z!=z<=<z>=>iszis notinznot inc                    s   t | |j}| jd  fdd}| jd }||tjv r+|d7 }||tjv s |t|j	 }| 
||}t|t||j	S )Nzutf-8c                    s   t  |  S r>   )chr)indexr   rD   rE   get_char  s   z-ExprBuilder.build_Attribute.<locals>.get_charr   )r  r  r   encoder   rr  string
whitespacern   r  make_raw_ranger#   r   )rv   r,  baser  	start_posend_pos
name_rangerD   r  rE   build_Attributex  s   zExprBuilder.build_Attributec                    s   t  |j} fdd|jD }t|dr(|jr(t  |j}|t| |g7 }g }|jD ] }t  |j}|j	s?t
| d|tt| |j	| q-t|||S )Nc                    r   rD   r  )r   r1  r   rD   rE   r     r   z*ExprBuilder.build_Call.<locals>.<listcomp>starargsz&keyword-arg expansion is not supported)r  r   r   r  r  r%   r   rP  r  r   rg   r   r   r   r   )rv   r,  r   r   stararg_exprkwargskwkw_exprrD   r   rE   
build_Call  s   
zExprBuilder.build_Callc                 C   s    |  |j|j|jd }t|S )N   )rq   rr   rs   r   rv   r,  r  rD   rD   rE   build_Ellipsis  s   zExprBuilder.build_Ellipsisc                 C   s   |  |j|j|jt|j }|jtrt|dt |jdkr%t|S |jdkr.t	|S |jdkr7t
|S |jdkr@t|S tt||jS )Nz=names of variables used in JIT-ed functions can't start with TrueFalseNoneEllipsis)rq   rr   rs   rn   r  r?   r@   rg   r*   r   r   r   r-   r   r  rD   rD   rE   
build_Name  s"   



zExprBuilder.build_Namec                 C   s|   |  |j|j|jtt|j }|jdu rt|S |jdu r#t|S |jd u r,t|S |jt	kr5t
|S tdt|j )NTFz!Name constant value unsupported: )rq   rr   rs   rn   strr  r*   r   r   r  r   r  r  rD   rD   rE   build_NameConstant  s   



zExprBuilder.build_NameConstantc                 C   s   t | |j}t | |j}t|j}|tjkr*| js*| |	 j
|	 j}t|dtj|}|d u rH| |	 j
|	 j}t|d|j t|||S )NzDivision of ints in TorchScript uses Python 3 true division semantics. Please put `from __future__ import division` at the top of your filezunsupported binary operator: )r  leftrightrm   r  r   r  uses_true_divisionr  r   rr  rq  rV   r  	binop_maprp   rg   rc   r
   )rv   r,  r  r  r  	err_ranger  rD   rD   rE   build_BinOp  s    

zExprBuilder.build_BinOpc                 C   sf   t | |j}t|j}tj|}|d u rt| d|j	 | 
|j|j|jt| }t|||S )Nzunsupported unary operator: )r  operandrm   r  r  unop_maprp   rg   r   rc   rq   rr   rs   rn   r,   )rv   r,  sub_exprr  r  r  rD   rD   rE   build_UnaryOp  s   
zExprBuilder.build_UnaryOpc                    s   t |jdk rtdtt |j  fdd|jD }t|j}tj|}|d u rC 	|d 
 j|d 
 j}t|d|j |d }|dd  D ]}t|||}qM|S )N   z.expected at least 2 values in BoolOp, but got c                    r   rD   r  )r   r  r   rD   rE   r     r   z,ExprBuilder.build_BoolOp.<locals>.<listcomp>r   r   zunsupported boolean operator: )rn   r  AssertionErrorr  rm   r  r  
boolop_maprp   r  r   rr  rq  rg   rc   r
   )rv   r,  	sub_exprsr  r  r  r  r  rD   r   rE   build_BoolOp  s$   

zExprBuilder.build_BoolOpc                 C   s$   t t| |jt| |jt| |jS r>   )r)   r  r  r   r  rv   r,  rD   rD   rE   build_IfExp  s
   


zExprBuilder.build_IfExpc                    s    fdd|j gt|j D }d }t||j|dd  D ]L\}}}t|}tj|} 	|
 j|
 j}	|d u rDt|	d|j |tjkrVtd||}
t|	d|
}nt|||}|d u rc|}qtd||}q|S )Nc                    r   rD   r  r   er   rD   rE   r     r   z-ExprBuilder.build_Compare.<locals>.<listcomp>r   z!unsupported comparison operator: r  r  r  )r  r   comparatorszipopsrm   r  	cmpop_maprp   r  r   rr  rq  rg   rc   r   NotInr
   r,   )rv   r,  operandsr0  r  op_r  r  r  r  in_exprcmp_exprrD   r   rE   build_Compare  s$     

zExprBuilder.build_Comparec           	         s  dd dd   fdd}t |j}t|j}|tju r_t|jjtjrTfdd|jjjD }|sO	|j
|jjj|jjjd	 }t|g }|| t||S t|t |jjgS |tju rot|||jgS |tju r~t||||jS |tju rg }|jjD ]}t|tjr||| q|t | q|s	|j
|jj|jjd	 }t|g }|| t||S t|t |jgS )
Nc                 S   s`   |j d urt| |j nd }|jd urt| |jnd }|jd ur%t| |jnd }t| |||S r>   )lowerr  upperstepr$   r   )rv   r  
slice_exprr  r  r   rD   rD   rE   build_SliceExpr#  s   


z4ExprBuilder.build_Subscript.<locals>.build_SliceExprc                 S   s(   t |jtjrt| dt| |jS )Nz9slicing multiple dimensions with tuples not supported yet)r   r  r   Tuplerg   r   r  )rv   r  
index_exprrD   rD   rE   build_Index5  s   z0ExprBuilder.build_Subscript.<locals>.build_Indexc                    s   g }|j D ]B}t|}|tju r| | || q|tju r)|| || q|tju r=|jtu r=|t	|
  qt|
 d| d|S )Nz!slicing multiple dimensions with z not supported)dimsrm   r   Indexr   SliceConstantr  r  r   r   rg   )rv   r  extslicer  r,  sub_type)r  r  rD   rE   build_ExtSlice=  s   



z3ExprBuilder.build_Subscript.<locals>.build_ExtSlicec                    r   rD   r  )r   r  r   rD   rE   r   T  s    
z/ExprBuilder.build_Subscript.<locals>.<listcomp>r  )r  r  rm   slicer   r  r   r  eltsrq   rr   rs   r+   r   r(   r  ExtSlice)	rv   r,  r  r  r  indicesr  tupr  rD   )r  r  rv   rE   build_Subscript!  sL   











zExprBuilder.build_Subscriptc                    .   t  |j|j|jd  fdd|jD S )Nr   c                    r   rD   r  r  r   rD   rE   r     r   z*ExprBuilder.build_List.<locals>.<listcomp>)r   rq   rr   rs   r  r  rD   r   rE   
build_List}     zExprBuilder.build_Listc                    r  )Nr   c                    r   rD   r  r  r   rD   rE   r     r   z+ExprBuilder.build_Tuple.<locals>.<listcomp>)r+   rq   rr   rs   r  r  rD   r   rE   build_Tuple  r  zExprBuilder.build_Tuplec                    s^     |j|j|jd }|jr|jd st|dt| fdd|jD  fdd|jD S )Nr   r   z1Dict expansion (e.g. `{**dict}`) is not supportedc                    r   rD   r  r  r   rD   rE   r     r   z*ExprBuilder.build_Dict.<locals>.<listcomp>c                    r   rD   r  r  r   rD   rE   r     r   )rq   rr   rs   keysrg   r   r  )rv   r,  r   rD   r   rE   
build_Dict  s   zExprBuilder.build_Dictc                 C   s0   t |j}| |j|j|jt| }t||S r>   )r  r  rq   rr   rs   rn   r   rv   r,  r  r  rD   rD   rE   	build_Num  s   

zExprBuilder.build_Numc                 C   s   |j }|d u st|trt| |S t|tttfr t| |S t|t	r+t
| |S t|ttr8t| |S | |j|j|jtt	| }t|d)Nz Unknown Constant expression type)r  r   r  r  r  intfloatcomplexr  r  	build_Strrm   r  r  rq   rr   rs   rn   rV   )rv   r,  r  error_rangerD   rD   rE   build_Constant  s   

zExprBuilder.build_Constantc                 C   s4   t |j}| |j|j|jt| d }t||S Nr   )r  r  rq   rr   rs   rn   r'   r  rD   rD   rE   r    s
   

zExprBuilder.build_Strc                 C   s   d}g }|j D ]F}| |j|j|jd }t|tjr=|jdkr%t|d|j	d ur/t|d|d7 }|
t| |j qt|tjrI||j7 }qt|d| |j|j|jd }ttt||t|d|g S )	Nrk   r   r  z%Don't support conversion in JoinedStrz%Don't support formatting in JoinedStrz{}zUnsupported value in JoinedStrformat)r  rq   rr   rs   r   r   FormattedValue
conversionrg   format_specr   r  r  r	  r   r#   r'   r   )rv   r,  r   r   r  r  rD   rD   rE   build_JoinedStr  s    





zExprBuilder.build_JoinedStrc                 C   s   |  |j|j|j}t|jdkrt|dt|jd jdkr%t|dt| |j}t| |jd j	}t| |jd j
}t||||S Nr   z.Only a single generator is currently supportedr   z'Comprehension ifs are not supported yet)rq   rr   rs   rn   
generatorsrg   ifsr  eltr  r  r   )rv   r  r  elt_exprtarget_expr	iter_exprrD   rD   rE   build_ListComp  s   

zExprBuilder.build_ListCompc                 C   s   t | |S r>   )r  r.  )rv   r  rD   rD   rE   build_GeneratorExp  s   zExprBuilder.build_GeneratorExpc                 C   s   |  |j|j|j}t|jdkrt|dt|jd jdkr%t|dt| |j}t| |j	}t| |jd j
}t| |jd j}t|||||S r'  )rq   rr   rs   rn   r(  rg   r)  r  keyr  r  r  r   )rv   r  r  key_expr
value_exprr,  r-  rD   rD   rE   build_DictComp  s   

zExprBuilder.build_DictCompc                 C   s*   |  |j|j|jd }t|t| |jS r!  )rq   rr   rs   r%   r  r  r  rD   rD   rE   build_Starred  s   zExprBuilder.build_StarredN);rc   rd   re   r   r  r  r  r  r  r  FloorDivr  r  r  r  r  r  MatMultNotUSubInvertr  AndOrr  EqNotEqLtELtGtEGtIsIsNotInr  r  rt  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r   r  r&  r.  r/  r3  r4  rD   rD   rD   rE   r  N  s    











[









r  r   r   c                 C   s8   | j d | |}| ||d  |t| |d  S )Nr   r   )r   rindexr  rn   )rv   possubstrr  new_posrD   rD   rE   r    s   $r  )NF)NNr>   )rE  ){r   rR  r   r   r?  r  collectionsr   textwrapr   rY   torch.jit.annotationsr   torch._C._jit_tree_viewsr   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/   r0   torch._jit_internalr1   r2   r3   r4   torch._sourcesr5   r6   r7   r8   r   torch.jit._dataclass_implsr9   torch.jit._monkeytype_configr:   r;   r@   rA   setascii_lowercaseascii_uppercasedigits_identifier_charsrF   r   Import
ImportFromGlobalrt   ro   updateAsyncFunctionDefAsyncFor	AsyncWithTryNonlocalr   	ExceptionrV   rg   rj   r   r   r   r   r   r   r   r	  r
  r   r   r)  r  r$  r`  ri  rm  rn  r   r  r  r   r   r  rD   rD   rD   rE   <module>   s   .

"
JC
	

4^ +   -