o
    ߗiv                     @   s  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	 d dl
mZ ddlmZmZ ddlmZmZ ddlmZ ddlmZmZ dd	lmZmZ dd
lmZmZmZ ddlmZm Z m!Z!m"Z" ddl#m$Z$m%Z% ddl&m'Z' e	rxd dl(m)Z) dd Z*G dd de%Z+G dd de+Z,G dd de+Z-G dd de-Z.G dd de%Z/G dd de/Z0G dd de/Z1d e2fd!d"Z3d e2fd#d$Z4d%d&d'e5d d(fd)d*Z6G d+d, d,e+Z7e8dd-d. Z9G d/d0 d0e%Z:G d1d2 d2e%Z;dS )3    N)DictListOptionalTYPE_CHECKING)is_fake   )	polyfills	variables)create_call_functioncreate_instruction)	skip_code)raise_observed_exceptionunimplemented)GuardBuilderinstall_guard)
AttrSourceGetItemSourceis_from_local_source)	dict_keysdict_valuesistypespecialize_symnode   )ValueMutationNewVariableTracker)ConstantVariable)InstructionTranslatorc                 C   s   t | tjr|  jjdd uS t | tjr!tdd | j	D S t | tj
tjtjtjtjjtjtjtjjtjtjtjtjtjtjfS )Nexample_valuec                 s   s    | ]}t |V  qd S N)is_hashable.0e r#   [/home/ubuntu/transcripts/venv/lib/python3.10/site-packages/torch/_dynamo/variables/dicts.py	<genexpr>'   s    zis_hashable.<locals>.<genexpr>)
isinstancer	   TensorVariableas_proxynodemetagetTupleVariableallitemsBuiltinVariableSymNodeVariabler   EnumVariableuser_definedUserDefinedClassVariableUserFunctionVariableSkipFunctionVariablemiscNumpyVariableNNModuleVariableUnspecializedNNModuleVariableMethodWrapperVariableTorchInGraphFunctionVariableTypingVariableFunctoolsPartialVariablexr#   r#   r$   r       s,   r   c                       s   e Zd ZdhejZG dd dZefdeeef ddf fddZd	d
 Z	dd Z
dd Zdd Zdd ZdefddZdd Zdd Zdd Zdd Zdddefd d!Zdddefd"d#Zdefd$d%Z	&	'	(	)		*d1 fd+d,Zd-d. Zd/d0 Z  ZS )2ConstDictVariableuser_clsc                   @   sL   e Zd ZdZdddZedd Zdd	 Zed
d Z	dd de
fddZdS )z"ConstDictVariable._HashableTrackeraB  
        Auxiliary opaque internal class that wraps a VariableTracker and makes it hashable
        This should not be seen or touched by anything outside of ConstDictVariable and its children
        Note that it's also fine to put VTs into dictionaries and sets, but doing so does not take into account aliasing
        returnNc                 C   s2   t |}t|stdt| d|  || _d S )NzDict key of type z. Key: )r   r   r   typevt)selfrD   r#   r#   r$   __init__K   s   
z+ConstDictVariable._HashableTracker.__init__c                    s   t | jtjr| j jjd }|S t | jtjr+tj	 t
 fdd| jjD }|S t | jtjr6| jjS t | jtjrA| jjS t | jtjrM| j S | j }|S )Nr   c                 3   s    | ]} |j V  qd S r   )underlying_valuer    Hashabler#   r$   r%   Z       zFConstDictVariable._HashableTracker.underlying_value.<locals>.<genexpr>)r&   rD   r	   r'   r(   r)   r*   r,   r@   _HashableTrackertupler.   r8   moduler9   valuer4   get_functionas_python_constant)rE   r?   r#   rH   r$   rG   T   s   	

z3ConstDictVariable._HashableTracker.underlying_valuec                 C   
   t | jS r   )hashrG   rE   r#   r#   r$   __hash__e   s   
z+ConstDictVariable._HashableTracker.__hash__c                    sh   t | t |kr
dS t| tr(tj t| t|ko't fddt| |D S t| r0| |u S | |kS )NFc                 3   s     | ]\}}  ||V  qd S r   )_eq_impl)r!   uvrH   r#   r$   r%   o   s    
z>ConstDictVariable._HashableTracker._eq_impl.<locals>.<genexpr>)	rC   r&   rL   r@   rK   lenr-   zipr   )abr#   rH   r$   rU   h   s   
z+ConstDictVariable._HashableTracker._eq_implotherc                 C   sN   t j}t||st|sJ t|t||r || j|jS || j|S r   )r@   rK   r&   r   
is_literalrC   rU   rG   )rE   r\   rI   r#   r#   r$   __eq__w   s   
z)ConstDictVariable._HashableTracker.__eq__rB   N)__name__
__module____qualname____doc__rF   propertyrG   rT   staticmethodrU   boolr^   r#   r#   r#   r$   rK   D   s    
	

rK   r.   rB   Nc                    s   d|v r	| d d|v r| d t jd	i | tj t fdd| D s-J  fddfdd| D | _t| j | _	|
 | _|| _d S )
Noriginal_itemsshould_reconstruct_allc                 3   s,    | ]\}}t |t fot |tV  qd S r   )r&   r   r!   r?   rW   rH   r#   r$   r%      s    
z-ConstDictVariable.__init__.<locals>.<genexpr>c                    s   t |  r| S  | S r   )r&   )keyrH   r#   r$   make_hashable   s   z1ConstDictVariable.__init__.<locals>.make_hashablec                    s   i | ]	\}} ||qS r#   r#   ri   )rk   r#   r$   
<dictcomp>   s    z.ConstDictVariable.__init__.<locals>.<dictcomp>r#   )popsuperrF   r@   rK   r-   r.   r   sourcerh   copyrg   rA   )rE   r.   rA   kwargs	__class__)rI   rk   r$   rF      s   



zConstDictVariable.__init__c                 C      dd | j   D S )Nc                 S       i | ]\}}|j  | qS r#   rD   r(   r!   krW   r#   r#   r$   rl      s     z.ConstDictVariable.as_proxy.<locals>.<dictcomp>r.   rS   r#   r#   r$   r(         zConstDictVariable.as_proxyc                 C   s"   dd dd | j D  d S )N{, c                 s   s,    | ]\}}|j   d |  V  qdS )z: NrD   
debug_reprrw   r#   r#   r$   r%      s    
z/ConstDictVariable.debug_repr.<locals>.<genexpr>})joinr.   rS   r#   r#   r$   r~      s   
zConstDictVariable.debug_reprc                 C   rt   )Nc                 S   ru   r#   rD   rP   rw   r#   r#   r$   rl      s    z8ConstDictVariable.as_python_constant.<locals>.<dictcomp>ry   rS   r#   r#   r$   rP      s   z$ConstDictVariable.as_python_constantc                 C   rt   )Nc                 S   s   i | ]
\}}|j  |qS r#   r   rw   r#   r#   r$   rl          z=ConstDictVariable.keys_as_python_constant.<locals>.<dictcomp>ry   rS   r#   r#   r$   keys_as_python_constant   rz   z)ConstDictVariable.keys_as_python_constantc                 C   s   | j S r   )rA   rS   r#   r#   r$   python_type   s   zConstDictVariable.python_typec                 C   sB   t |tsJ tj}t|o ||| jv o t | j|| tj S r   )r&   r   r@   rK   r   r.   r	   DeletedVariable)rE   rD   rI   r#   r#   r$   __contains__   s   zConstDictVariable.__contains__c                 C   s   t dd | j D S )Nc                 S   s   g | ]
}t |tjs|qS r#   )r&   r	   r   r!   r?   r#   r#   r$   
<listcomp>   s    
z)ConstDictVariable.len.<locals>.<listcomp>)rX   r.   valuesrS   r#   r#   r$   rX      s
   zConstDictVariable.lenc                    s&    j rdS t fdd j D S )NTc                 3   s*    | ]\}}   j|j|V  qd S r   )is_new_itemrg   r+   rD   )r!   rj   rN   rS   r#   r$   r%      s
    
z2ConstDictVariable.has_new_items.<locals>.<genexpr>)rh   anyr.   rS   r#   rS   r$   has_new_items   s
   zConstDictVariable.has_new_itemsc                 C   s<   |r|  r|  rt| t| kS t|t|kS r   )is_realizedidrealize)rE   rN   r\   r#   r#   r$   r      s   zConstDictVariable.is_new_itemc                    s   | j tju r  fdd d}| j D ]!\}}| j|j}| ||s*| j	r7 |j  | |d7 }q| j tju rO 
td|dgtdd d S  td|d d S )Nc                            t dgS )NOrderedDict)extend_outputcreate_load_python_modulecollectionscreate_load_attrr#   codegenr#   r$   <lambda>   
    z/ConstDictVariable.reconstruct.<locals>.<lambda>r   r   	BUILD_MAPargF)rA   r   r   add_push_nullr.   rg   r+   rD   r   rh   r   r   r
   append_output)rE   r   num_argsrj   rN   itemr#   r   r$   reconstruct   s(   
	

zConstDictVariable.reconstructtxr   r   c                 C   s(   t |}|| jvrtt| | j| S r   )r@   rK   r.   r   KeyErrorrE   r   r   rj   r#   r#   r$   'getitem_const_raise_exception_if_absent   s   



z9ConstDictVariable.getitem_const_raise_exception_if_absentc                 C   s.   t |}|| jvrtd|j  | j| S )Nzdict KeyError: )r@   rK   r.   r   rN   r   r#   r#   r$   getitem_const  s   


zConstDictVariable.getitem_constc                 C   s"   t |}|| jvrd S | j| S r   )r@   rK   r.   )rE   r   rj   r#   r#   r$   maybe_getitem_const  s   


z%ConstDictVariable.maybe_getitem_constargsList[VariableTracker]rq   Dict[str, VariableTracker]r   c                    s  ddl m}m m}m}mm} tj|ot	|d }	|dkr/t
|dks'J | ||d S |dkrT|s7|r9J | jrF|jj| j  fdd| j D S |dkro| jre|jj| j  |si|rkJ t| S |d	kr| jr|jj| j  |s|rJ t| S |d
kr|s|rJ | j| j t d dS |dkr|s|rJ  t
| jS |dkr|	r|  r|st
|dksJ |jj|  |d | j|d <  d S |dkr|	r|  rd| _|jj|  | j|d   d S |dv r%t
|dv r%|d | vr%t
|dkr! d S |d S |dkrF|	rF|  rFd| _|jj|  | j|d S |dkr_d| _|jj|  | j   d S |dkr|  rt
|dko}t|d t||t j!|f}
t
|dkot
|dk}|
s|r|jj|  t
|dkrt|d tr|d }n	|"|t#|d }| j$|j  fdd| D }| j$|  d S t% &||||S |dv r|d | v r| '||d S |dkrt
|dkr |d | v S |dkrQ|	rQ|  rQ|rJ t
|dks J | (|d }|d ur.|S t
|dkr; d }n|d }|jj|  || j|d < |S t% &||||S )Nr   )r/   r   ListIteratorVariableListVariabler,   UserDefinedObjectVariabler   __getitem__r.   c                    s   g | ]\}} |j |gqS r#   rD   rw   )r,   r#   r$   r   -  s    z1ConstDictVariable.call_method.<locals>.<listcomp>keysr   rp   )r.   mutation_typero   __len____setitem__r   __delitem__T)rm   r+   )r   r   rm   clearupdatec                    s    i | ]\}}  ||qS r#   )createrw   )r   rI   r#   r$   rl   r      z1ConstDictVariable.call_method.<locals>.<dictcomp>)r+   __getattr__r   
setdefault)) r/   r   r   r   r,   r   r@   rK   r   rX   r   ro   outputguard_on_key_orderaddnamer.   DictKeys
DictValuesclonerp   r   r   
is_mutableside_effectsmutationrh   r   rm   r   r&   r	   IteratorVariablecall_custom_dictdictr   rn   call_methodr   r   )rE   r   r   r   rq   r/   r   r   r   arg_hashableis_args_supportedis_kwargs_supporteddict_vtrN   r?   rr   )r   rI   r,   r$   r     s    	

&






zConstDictVariable.call_methodc                 C   s   dd | j  D S )Nc                 S      g | ]}|j qS r#   r   r   r#   r#   r$   r         z9ConstDictVariable.unpack_var_sequence.<locals>.<listcomp>r.   r   rE   r   r#   r#   r$   unpack_var_sequence  rz   z%ConstDictVariable.unpack_var_sequencec                 C   s@   | j tu r|| j jv rtdS tdS td| j  d d S )NTFzhasattr on z is not supported)rA   r   __dict__r   r   r   rE   r   r   r#   r#   r$   call_hasattr  s
   


zConstDictVariable.call_hasattrr   r   rq   r   rB   r   )r`   ra   rb   r   _nonvar_fieldsrK   r   r   rF   r(   r~   rP   r   r   rf   r   rX   r   r   r   r   r   r   r   r   r   __classcell__r#   r#   rr   r$   r@   >   sL    A
"			!
}r@   c                       s^   e Zd Zdd fddZ fddZ fddZed	d
 Z						d fddZ  Z	S )DefaultDictVariableNrB   c                    s.   t  j||fi | |tju sJ || _d S r   )rn   rF   r   defaultdictdefault_factory)rE   r.   rA   r   rq   rr   r#   r$   rF     s   
zDefaultDictVariable.__init__c                    s$   | j tttfvr| jsdS t  S )NF)r   listrL   r   r.   rn   is_python_constantrS   rr   r#   r$   r     s   
z&DefaultDictVariable.is_python_constantc                    s   d| j   dt   dS )Nzdefaultdict(r|   ))r   r~   rn   rS   rr   r#   r$   r~     s   zDefaultDictVariable.debug_reprc                 C   s,   t | tjr| jttttfv S t | tjj	S r   )
r&   r	   r/   fnr   rL   r   set	functionsBaseUserFunctionVariabler   r#   r#   r$   is_supported_arg  s   z$DefaultDictVariable.is_supported_argr   r   rq   r   r   c                    s   |dkr=t |dksJ |d | v r| ||d S | jd u r&t|d  | j|g i }t |d|d |f| |S t ||||S )Nr   r   r   r   )rX   r   r   r   call_functionrn   r   )rE   r   r   r   rq   default_varrr   r#   r$   r     s   
zDefaultDictVariable.call_methodr   r_   r   )
r`   ra   rb   rF   r   r~   re   r   r   r   r#   r#   rr   r$   r     s    
r   c                       s   e Zd ZdZdee ddf fddZdd Zed	d
 Z	e
dd Zdd Zdd Zdd Zdd Zdee deeef ddf fddZdddefddZ  ZS )SetVariablez-We model a sets as dictonary with None valuesr.   rB   Nc                    s(   t |t }t j|fi | d S r   )r   fromkeysr   _default_valuern   rF   rE   r.   rq   rr   r#   r$   rF     s   zSetVariable.__init__c                 C   ,   | j sdS dddd | j  D  d S )Nzset()r{   ,c                 s       | ]}|j  V  qd S r   r}   r!   rx   r#   r#   r$   r%     rJ   z)SetVariable.debug_repr.<locals>.<genexpr>r   r.   r   r   rS   r#   r#   r$   r~        "zSetVariable.debug_reprc                 C   s   t | j S r   )r   r.   r   rS   r#   r#   r$   	set_items     zSetVariable.set_itemsc                   C   s
   t d S r   r   r   r#   r#   r#   r$   r     s   
zSetVariable._default_valuec                 C      dd | j D S )Nc                 S      h | ]}|j  qS r#   rv   r   r#   r#   r$   	<setcomp>      z'SetVariable.as_proxy.<locals>.<setcomp>r   rS   r#   r#   r$   r(        zSetVariable.as_proxyc                 C      t S r   )r   rS   r#   r#   r$   r        zSetVariable.python_typec                 C   r   )Nc                 S   r   r#   r   r   r#   r#   r$   r     r   z1SetVariable.as_python_constant.<locals>.<setcomp>r   rS   r#   r#   r$   rP     r   zSetVariable.as_python_constantc                 C   s2   | dd | jD  |tdt| jd d S )Nc                 S   r   r#   r   r   r#   r#   r$   r     r   z+SetVariable.reconstruct.<locals>.<listcomp>	BUILD_SETr   )foreachr   r   r   rX   rE   r   r#   r#   r$   r     s   zSetVariable.reconstructr   rq   r   c                    sj  |dkr|rJ t |dksJ d}|d t f}n|dkr:|r$J |r(J | j j}t |||f| |S |dkrY|rBJ t |dksJJ t	t
j|| |d gi S |dkrx|raJ t |dksiJ t	t
j|| |d gi S |dkr|rJ t |dksJ t	t
j|| |d gi S |d	kr|rJ t |dksJ t	t
j|| |d gi S |d
krt |dkr|  r|rJ t |dksJ t	t
j|| |d gi S |dkr|rJ t |dksJ |d | vrtd t |d||S |dkr,|rJ t |dksJ |d | v r&t |d||S tjd dS t ||||S )Nr   r   r   r   rm   
isdisjointintersectionunion
differencer   removezkey does not existdiscardrN   )rX   r   r   r   rm   rD   rn   r   r	   r4   r   set_isdisjointr   set_intersection	set_unionset_differencer   
set_updater   r   r   )rE   r   r   r   rq   resultrr   r#   r$   r     sx   


zSetVariable.call_methodr   r   r   c                 C   s   t d)NzIllegal to getitem on a set)RuntimeError)rE   r   r   r#   r#   r$   r   8  s   zSetVariable.getitem_const)r`   ra   rb   rc   r   r   rF   r~   rd   r   re   r   r(   r   rP   r   r   strr   r   r   r#   r#   rr   r$   r     s0    


Ar   c                       s|   e Zd Zdee ddf fddZdd Zedd	 Zd
d Z	dd Z
dd Zdee deeef ddf fddZ  ZS )FrozensetVariabler.   rB   Nc                    s   t  j|fi | d S r   )rn   rF   r   rr   r#   r$   rF   =  s   zFrozensetVariable.__init__c                 C   r   )Nzfrozenset()r{   r   c                 s   r   r   r}   r   r#   r#   r$   r%   H  rJ   z/FrozensetVariable.debug_repr.<locals>.<genexpr>r   r   rS   r#   r#   r$   r~   D  r   zFrozensetVariable.debug_reprc                 C   s
   | j  S r   r   rS   r#   r#   r$   r   J     
zFrozensetVariable.set_itemsc                 C   r   r   )	frozensetrS   r#   r#   r$   r   N  r   zFrozensetVariable.python_typec                 C   r   )Nc                 S   r   r#   r   r   r#   r#   r$   r   R  r   z7FrozensetVariable.as_python_constant.<locals>.<setcomp>r   rS   r#   r#   r$   rP   Q  r   z$FrozensetVariable.as_python_constantc                    s<     dd | jD    fdd  tdd d S )Nc                 S   r   r#   r   r   r#   r#   r$   r   U  r   z1FrozensetVariable.reconstruct.<locals>.<listcomp>c                      s      dgS )Nr  )r   create_load_globalr#   r   r#   r$   r   W  s    z/FrozensetVariable.reconstruct.<locals>.<lambda>r   F)r   r   r   r   r
   r   r#   r   r$   r   T  s
   
zFrozensetVariable.reconstructr   rq   r   c                    s*   |dv rt d| dt ||||S )N)r   rm   r   r  r  r   zIllegal call_method z on a frozenset)r  rn   r   rE   r   r   r   rq   rr   r#   r$   r   _  s   zFrozensetVariable.call_method)r`   ra   rb   r   r   rF   r~   rd   r   r   rP   r   r   r  r   r   r#   r#   rr   r$   r  <  s&    

r  c                       s   e Zd ZU dZdZee ed< deddf fddZ	e
dd	 Ze
d
d Zdd Zdd Zded deedf ddf fddZ  ZS )DictViewzu
    Models _PyDictViewObject

    This is an "abstract" class. Subclasses will override kv and the items method
    Nkvdv_dictrB   c                    s8   t  jdi | | jdv sJ t|tsJ || _d S )N)r   r   r#   )rn   rF   r  r&   r@   r  )rE   r  rq   rr   r#   r$   rF   t  s   
zDictView.__init__c                 C   s   t | jj| j S r   )getattrr  r.   r  rS   r#   r#   r$   
view_itemsz  s   zDictView.view_itemsc                 C      t r   NotImplementedErrorrS   r#   r#   r$   view_items_vt~  s   zDictView.view_items_vtc                    s     fddfdd j D S )Nc                    s    j dkr| jS | S )Nr   )r  rD   r>   rS   r#   r$   unwrap  rz   z,DictView.unpack_var_sequence.<locals>.unwrapc                    s   g | ]} |qS r#   r#   r   )r  r#   r$   r     s    z0DictView.unpack_var_sequence.<locals>.<listcomp>r  r   r#   )rE   r  r$   r     s   zDictView.unpack_var_sequencec                 C   s$   || j  || j |d d S )Nr   )r  load_methodr  r   r   r#   r#   r$   r     s   
zDictView.reconstructr   r   rq   c                    ,   |dkr| j ||||S t ||||S )Nr   r  r   rn   r  rr   r#   r$   r        zDictView.call_method)r`   ra   rb   rc   r  r   r  __annotations__r@   rF   rd   r  r  r   r   r   r   r   r   r#   r#   rr   r$   r  k  s"   
 


r  c                       sZ   e Zd ZdZedd Zedd Zdd Zded	 d
e	e
d	f dd	f fddZ  ZS )r   r   c                 C   rQ   r   )r   r  rS   r#   r#   r$   r     r  zDictKeys.set_itemsc                 C   r   )Nc                 S   r   r#   r   r   r#   r#   r$   r     r   z*DictKeys.view_items_vt.<locals>.<listcomp>r  rS   r#   r#   r$   r    s   zDictKeys.view_items_vtc                 C   r   r   )r   rS   r#   r#   r$   r     r   zDictKeys.python_typer   r   rq   rB   c                    r  )Nr   r  r  rr   r#   r$   r     r   zDictKeys.call_method)r`   ra   rb   r  rd   r   r  r   r   r   r  r   r   r#   r#   rr   r$   r     s    


r   c                   @   s$   e Zd ZdZedd Zdd ZdS )r   r   c                 C   rQ   r   )r   r  rS   r#   r#   r$   r    r  zDictValues.view_items_vtc                 C   r   r   )r   rS   r#   r#   r$   r     r   zDictValues.python_typeN)r`   ra   rb   r  rd   r  r   r#   r#   r#   r$   r     s
    
r   rB   c                 C   s4   t jd}|d u rt jd}|d uot| |jS )Nztransformers.file_utilsztransformers.utils.generic)sysmodulesr+   
issubclassModelOutputclsmodr#   r#   r$   _is_matching_transformers_cls  s   r)  c                 C   s    t jd}|d uot| |jS )Nzdiffusers.utils)r"  r#  r+   r$  
BaseOutputr&  r#   r#   r$   _is_matching_diffusers_cls  s   r+  r   r   r   r   c              	   C   s  |j j| r'z|j jj| |dd}tjt|tj W S  t	y&   Y nw || j
v s2t| j|r6tdS t| jtrE| jdu rEtdS | jrnz|j j| j}tt| j|tj tt||W S  t	ym   Y nw td| jj d| d| j d| j  dS )	zXShared method between DataClassVariable and CustomizedDictVariable where items are attrsT)
deleted_okNFzhasattr(r|   z)  )r   r   is_attribute_mutation	load_attrr	   r   r   r&   r   r   r.   hasattrrA   r   r   r   ro   root_txget_example_valuer   r   
make_guardr   HASATTRr   rs   r`   )rE   r   r   r
  exampler#   r#   r$   _call_hasattr_customobj  s4   "r6  c                       s   e Zd Zedd Zedd Zedd Zedd Zed	d
 Z	d fddZ
dd Zdd Z						d  fddZdddeddf fddZeZ  ZS )!CustomizedDictVariablec                 C   s   t | pt| S r   )r)  r+  r'  r#   r#   r$   is_matching_cls_hf  s   z)CustomizedDictVariable.is_matching_cls_hfc                 C   s<   t | tjr| tjur| jtjju rt| dsdS t| S )N__post_init__T)r$  r   r   rF   r0  r7  r9  r8  r#   r#   r$   is_matching_cls  s   


z&CustomizedDictVariable.is_matching_clsc                 C      |  t|S r   r;  rC   r'  objr#   r#   r$   is_matching_object  r   z)CustomizedDictVariable.is_matching_objectc                    s<  dD ]!}t ||r#t||}t|sJ d| t |dr#t|j qt|rpt|j	|i |}|
  dd  i }| |ra|j D ]\}	}
t|
trV|
jd u sZ|
d u r[qH|
||	< qHn|j} fdd| D }n%|s|dd | D }nt|d	krt|d
 tr|s|d
 j}ntd | ||fi |S )N)rF   r:  __setattr__r   zexpect callable attr __code__c                 S   s.   t | tr| S t| rt| S td d S )Nz5expect VariableTracker or ConstantVariable.is_literal)r&   r   r   r]   r   r   r>   r#   r#   r$   make_var  s   


z/CustomizedDictVariable.create.<locals>.make_varc                    s    i | ]\}}t | |qS r#   r   rw   rC  r#   r$   rl   +  r   z1CustomizedDictVariable.create.<locals>.<dictcomp>c                 S   s   i | ]
\}}t ||qS r#   r   rw   r#   r#   r$   rl   0  r   r   r   z/custom dict init with args/kwargs unimplemented)r0  r  callabler   rB  dataclassesis_dataclassinspect	signaturebindapply_defaultsr9  	argumentsr.   r&   r   rN   rX   r@   r   )r'  rA   r   rq   options	attr_namer   bound
bound_argsrx   rW   r.   r#   rD  r$   r     s8   








zCustomizedDictVariable.createc           	      C   s   t |}| |std |j|j|jdt|j}dd t	
|D }|D ]'}t||rOt||}|j|jt|j|d|}|d urOt|}|||< q(| |||jdS )Nz.custom non-hf dict subclass wrap unimplemented)r   ro   c                 S   r   r#   )r   r!   fr#   r#   r$   r   E  r   z/CustomizedDictVariable.wrap.<locals>.<listcomp>ro   )rC   r9  r   rs   r   ro   r   r   r.   rF  fieldsr0  r  r   r   r   )	r'  builderr?  rA   r.   r   rj   valvarr#   r#   r$   wrap:  s,   



zCustomizedDictVariable.wraprB   Nc                    s(   t  j||fi | | |sJ d S r   )rn   rF   r;  )rE   r.   rA   rM  rr   r#   r$   rF   R  s   zCustomizedDictVariable.__init__c                 C   r  r   r  rS   r#   r#   r$   r(   V  r   zCustomizedDictVariable.as_proxyc                    sd    j fdd} |  } |  t| }  	t
||d d S )Nc                     sJ   r fdd}   |    jg r# tdd d S d S )Nc                      s8      jddd    d    d d S )NtorchT)r   _dynamodisable)r   r  r   r#   r   r#   r$   gen_fn2h  s   zDCustomizedDictVariable.reconstruct.<locals>.gen_fn1.<locals>.gen_fn2r   F)r   r   create_load_const_uncheckedrA   r
   )r\  r   is_hf_model_outputrE   r#   r$   gen_fn1^  s   
z3CustomizedDictVariable.reconstruct.<locals>.gen_fn1F)r9  rA   r   r   r   r   rL   r   r   create_call_function_kwrX   )rE   r   r`  dr   r#   r^  r$   r   [  s   
z"CustomizedDictVariable.reconstructr   r   rq   r   r   c                    s   t | j|}| jd u rd nt| j|}t|dr)|jttjfv r)t	 
||||S |dv r>|tj||d| gt| |S |t tj|d u rPt	 
||||S td|  d S )N__objclass__)r   to_tupler   rA  r:  rS  z,custom dict: call_method unimplemented name=)r  rA   ro   r   r0  rc  r   r   r   rn   r   inline_user_function_returnr	   r4   r   r   )rE   r   r   r   rq   r   ro   rr   r#   r$   r   }  s    z"CustomizedDictVariable.call_methodr   r   r   c                    s   t |}|| v r| |d|gi S t| jr9dd t| jD }||v r9tj || s1J tj || S t	 
||S )Nr   c                 S   s   i | ]}|j |jqS r#   )r   defaultrQ  r#   r#   r$   rl     r   z6CustomizedDictVariable.var_getattr.<locals>.<dictcomp>)r   r   r   rF  rG  rA   rT  r	   r]   rn   var_getattr)rE   r   r   name_vtdefaultsrr   r#   r$   rg    s   
z"CustomizedDictVariable.var_getattrr_   r   )r`   ra   rb   re   r9  r;  classmethodr@  r   rX  rF   r(   r   r   r  rg  r6  r   r   r#   r#   rr   r$   r7    s.    



1
""r7  c                  C   s0   t jd} | d usJ ddd }|| j_d S )N transformers.configuration_utilsz4Caller is responsible for ensuring this was importedc                 S   s   t |dsdS | j|jkS )Nr   F)r0  r   )rE   r\   r#   r#   r$   !_dynamo_overriden_transformers_eq  s   
zJ_install_PretrainedConfig_patch.<locals>._dynamo_overriden_transformers_eq)r"  r#  r+   PretrainedConfigr^   )r(  rl  r#   r#   r$   _install_PretrainedConfig_patch  s   rn  c                       sf   e Zd ZdZedd Zedd Zd fdd	Zd
dde	ddfddZ
d
dde	ddfddZ  ZS )HFPretrainedConfigVariablez/
    Hack for HuggingFace PretrainedConfig
    c                 C   s.   t jd}|d uot| |j}|rt  |S )Nrk  )r"  r#  r+   r$  rm  rn  )r'  r(  is_matchr#   r#   r$   r;    s
   z*HFPretrainedConfigVariable.is_matching_clsc                 C   r<  r   r=  r>  r#   r#   r$   r@    r   z-HFPretrainedConfigVariable.is_matching_objectrB   Nc                    s.   t  jdi | || _| t|sJ d S )Nr#   )rn   rF   r?  r;  rC   )rE   r?  rq   rr   r#   r$   rF     s   z#HFPretrainedConfigVariable.__init__r   r   r   r   c              	   C   s\   zt | j|}| jot| j|}t|||W S  ty-   td| j d| d Y d S w )Nzgetattr(r|   r   )	r  r?  ro   r   r   buildAttributeErrorr   rN   )rE   r   r   
attr_valuero   r#   r#   r$   rg    s   z&HFPretrainedConfigVariable.var_getattrc                 C   s   t jt| j|S r   )r	   r   r   r0  r?  r   r#   r#   r$   r     rz   z'HFPretrainedConfigVariable.call_hasattrr_   )r`   ra   rb   rc   re   r;  rj  r@  rF   r  rg  r   r   r#   r#   rr   r$   ro    s    
	
	ro  c                   @   s   e Zd ZdZdd Zdd Zdddee d	ee	ef fd
dZ
dddefddZdddefddZ	ddddedee fddZdddefddZdS )PythonSysModulesVariablezSpecial case for sys.modules.

    Without this we will guard on the exact set of modules imported in the
    lifetime of the python program.
    c                 C   r   r   )r   rS   r#   r#   r$   r     r   z$PythonSysModulesVariable.python_typec                    s      fdd d S )Nc                      r   )Nr#  )r   r   r"  r   r#   r   r#   r$   r     r   z6PythonSysModulesVariable.reconstruct.<locals>.<lambda>)r   r   r#   r   r$   r     s   
z$PythonSysModulesVariable.reconstructr   r   r   rq   c                 C   s   |dkr| j |g|R i |S |dkr | j|g|R i |S |dkr0| j|g|R i |S td| d| d| d d S )Nr   r+   r   zsys.modules.z(*z, **r   )call_getitemcall_getcall_containsr   r  r#   r#   r$   r     s    z$PythonSysModulesVariable.call_methodrj   c                 C   s8   |  }|tjv }t| tjtj|| d ||fS )N)rj   invert)	rP   r"  r#  r   r3  	functoolspartialr   DICT_CONTAINSrE   r   rj   rx   has_keyr#   r#   r$   _contains_helper  s   
z)PythonSysModulesVariable._contains_helperc                 C   s   |  ||\}}tj|dS Nr  )r~  r   r   r|  r#   r#   r$   rw    s   z&PythonSysModulesVariable.call_containsNrf  c                 C   sR   |  ||\}}|r| jot| j|}t|tj| |S |d ur#|S tjd dS r  )	r~  ro   r   r   rq  r"  r#  r   r   )rE   r   rj   rf  rx   r}  ro   r#   r#   r$   rv    s   z!PythonSysModulesVariable.call_getc                 C   s6   |  ||\}}| jot| j|}t|tj| |S r   )r~  ro   r   r   rq  r"  r#  )rE   r   rj   rx   r}  ro   r#   r#   r$   ru  $  s   z%PythonSysModulesVariable.call_getitemr   )r`   ra   rb   rc   r   r   r   r   r   r  r   r~  rw  r   rv  ru  r#   r#   r#   r$   rt    s,    




rt  )<r   rF  ry  rH  r"  typingr   r   r   r   torch._subclasses.fake_tensorr   r   r   r	   bytecode_transformationr
   r   
eval_framer   excr   r   guardsr   r   ro   r   r   r   utilsr   r   r   r   baser   r   constantr   torch._dynamo.symbolic_convertr   r   r@   r   r   r  r  r   r   rf   r)  r+  r  r6  r7  	lru_cachern  ro  rt  r#   r#   r#   r$   <module>   sV     _5l/0
  @
%