o Zh@stdZddlZddlZddlZddlmZmZddlmZddl m Z m Z ddl m Z mZddlmZmZdd lmZmZdd lmZdd lmZmZmZmZd d lmZmZd dlm Z erdddl!m"Z"ddZ#GdddeZ$GdddeZ%Gddde$Z&Gddde$Z'Gddde$Z(Gddde(Z)Gddde(Z*Gd d!d!eZ+Gd"d#d#e+Z,Gd$d%d%e+Z-dS)&a Dictionary-related variable tracking classes for PyTorch Dynamo. This module implements variable tracking for different types of dictionary-like objects: - Regular Python dictionaries (dict) - Ordered dictionaries (collections.OrderedDict) - Default dictionaries (collections.defaultdict) - Dictionary views (keys and values) - Sets and frozensets (implemented internally using dictionaries) These classes are responsible for tracking dictionary operations during graph compilation, maintaining proper guards for dictionary mutations and key existence checks. They handle dictionary creation, modification, key/value access, and view operations while ensuring correct behavior in the compiled code through appropriate guard installation. The implementation uses a special _HashableTracker wrapper to handle dictionary keys while preserving proper aliasing semantics. Sets are implemented as dictionaries with None values for efficiency and code reuse. N)Optional TYPE_CHECKING)is_fake) polyfills variables)create_call_functioncreate_instruction)raise_observed_exception unimplemented) GuardBuilder install_guard)is_from_local_source)cmp_name_to_op_mapping dict_keys dict_valuesspecialize_symnode)ValueMutationNewVariableTrackerConstantVariable)InstructionTranslatorcCst|tjr|s|rdSt|tjr!|jj dduSt|tj r1t dd|j DSt|tj tjtjtjtjjtjtjtjjtjtjtjtjtjtjtjfS)NT example_valuecss|]}t|VqdSN) is_hashable.0erL/var/www/auris/lib/python3.10/site-packages/torch/_dynamo/variables/dicts.py Cszis_hashable..) isinstancerLazyVariableTracker is_realizedrTensorVariableas_proxynodemetaget TupleVariableallitemsBuiltinVariableZSymNodeVariablerZ EnumVariableZ user_definedZUserDefinedClassVariableUserFunctionVariableZSkipFunctionVariablemiscZ NumpyVariableNNModuleVariableUnspecializedNNModuleVariableZMethodWrapperVariableZTorchInGraphFunctionVariableZTypingVariableZFunctoolsPartialVariableWeakRefVariablexrrr r1s<   rcs(eZdZdhejZGdddZefdeeefddffdd Zd d Zd d Z d dZ ddZ ddZ de fddZddZddZddZddZdddefd d!Zdddefd"d#Zdefd$d%Zdefd&d'Zd(d)Zd*d+Z , - . /  0d9fd1d2 Zd3d4Zd5d6Zfd7d8ZZS):ConstDictVariableuser_clsc@sLeZdZdZdddZeddZdd Zed d Z d dde fd dZ dS)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 returnNcCs2t|}t|stdt|d|||_dS)NzDict key of type z. Key: )rrr typevt)selfr9rrr __init__hs z+ConstDictVariable._HashableTracker.__init__cst|jtjr|js|jr|jSt|jtjr(|jj j d}|St|jtj rAt j tfdd|jjD}|St|jtjrL|jjSt|jtjrW|jjSt|jtjrc|jSt|jtjrtt j |jjjS|j}|S)Nrc3s|]}|jVqdSr)underlying_valuerHashablerr r!}zFConstDictVariable._HashableTracker.underlying_value..)r"r9rr#r$roriginal_valuer%r&r'r(r*r5_HashableTrackertupler,r0valuer1r.Z get_functionr2Z referent_vtr<as_python_constant)r:r4rr=r r<qs2     z3ConstDictVariable._HashableTracker.underlying_valuecC t|jSr)hashr<r:rrr __hash__s z+ConstDictVariable._HashableTracker.__hash__csht|t|kr dSt|tr(tjt|t|ko'tfddt||DSt|r0||uS||kS)NFc3s |] \}}||VqdSr)_eq_impl)ruvr=rr r!s z>ConstDictVariable._HashableTracker._eq_impl..) r8r"rBr5rAlenr+zipr)abrr=r rIs z+ConstDictVariable._HashableTracker._eq_implothercCsNtj}t||st|sJt|t||r ||j|jS||j|Sr)r5rAr"rZ is_literalr8rIr<)r:rPr>rrr __eq__s z)ConstDictVariable._HashableTracker.__eq__r7N) __name__ __module__ __qualname____doc__r;propertyr<rH staticmethodrIboolrQrrrr rAas   rAr,r7Nc sd|vr |dd|vr|dtjd i|tjtfdd|Ds-Jfddfdd|D|_t|j |_ | |_ ||_ dS) Noriginal_itemsshould_reconstruct_allc3s,|]\}}t|tfot|tVqdSr)r"rrr4rKr=rr r!s  z-ConstDictVariable.__init__..cst|r|S|Sr)r")keyr=rr make_hashablesz1ConstDictVariable.__init__..make_hashablecsi|] \}}||qSrrr\)r^rr sz.ConstDictVariable.__init__..r) popsuperr;r5rAr+r,rsourcer[copyrZr6)r:r,r6kwargs __class__)r>r^r r;s      zConstDictVariable.__init__cCdd|jDS)NcS i|] \}}|j|qSrr9r&rkrKrrr r_s z.ConstDictVariable.as_proxy..r,rGrrr r&zConstDictVariable.as_proxycCs"dddd|jDdS)N{, css,|]\}}|jd|VqdS)z: Nr9 debug_reprrjrrr r!s z/ConstDictVariable.debug_repr..})joinr,rGrrr rqs zConstDictVariable.debug_reprcCrg)NcSrhrr9rDrjrrr r_z8ConstDictVariable.as_python_constant..rlrGrrr rDsz$ConstDictVariable.as_python_constantcCs|dd|jDS)NcSsi|] \}}|j|qSrrtrjrrr r_sz=ConstDictVariable.keys_as_python_constant..)install_dict_keys_match_guardr,rGrrr keys_as_python_constantz)ConstDictVariable.keys_as_python_constantcC|jSr)r6rGrrr python_typeszConstDictVariable.python_typecCsBt|tsJtj}t|o |||jvo t|j||tj Sr)r"rr5rArr,rDeletedVariable)r:r9r>rrr __contains__s zConstDictVariable.__contains__cCstdd|jDS)NcSsg|] }t|tjs|qSr)r"rr{rr4rrr s  z)ConstDictVariable.len..)rLr,valuesrGrrr rLs zConstDictVariable.lencs&jrdStfddjDS)NTc3s*|]\}}j|j|VqdSr) is_new_itemrZr)r9)rr]rCrGrr r!s  z2ConstDictVariable.has_new_items..)r[anyr,rGrrGr has_new_itemss  zConstDictVariable.has_new_itemscCs<|r|r|rt|t|kSt|t|kSr)r$idrealize)r:rCrPrrr rszConstDictVariable.is_new_itemcs|jtjurfddd}|jD]!\}}|j|j}|||s*|j r7|j||d7}q|jtjurO t d|dgt dddS t d|ddS)NctdgS)N OrderedDict) extend_outputcreate_load_python_module collectionscreate_load_attrrcodegenrr  z/ConstDictVariable.reconstruct..rr BUILD_MAPargF)r6rr add_push_nullr,rZr)r9rr[rr r append_output)r:rnum_argsr]rCitemrrr reconstructs(      zConstDictVariable.reconstructtxrrcCs(t|}||jvrtt||j|Sr)r5rAr,r KeyErrorr:rrr]rrr 'getitem_const_raise_exception_if_absent%s    z9ConstDictVariable.getitem_const_raise_exception_if_absentcCs.t|}||jvrtd|j|j|S)Nzdict KeyError: )r5rAr,r rCrrrr getitem_const-s   zConstDictVariable.getitem_constcCs"t|}||jvr dS|j|Sr)r5rAr,)r:rr]rrr maybe_getitem_const3s   z%ConstDictVariable.maybe_getitem_constcCsX||vsJt|}t|j|}t|j|}t|tj r*| dSdSr) r5rArBr,keysindexrZr"rr#r)r:rr]rZoriginal_key_vtrrr realize_key_vt9s    z ConstDictVariable.realize_key_vtcCs|jr t|tjdSdSr)rbr make_guardr ZDICT_KEYS_MATCHrGrrr rvBsz/ConstDictVariable.install_dict_keys_match_guardcCsddlm}|js dS|jj|rdS|d|v}|djdur.rrrc)r,Z mutation_typerb__len__ __setitem__r __delitem__T)r`r))rrr`clearupdatecs i|] \}}||qSr)createrj)rr>rr r_ruz1ConstDictVariable.call_method..)r) __getattr__r| setdefault move_to_end)*rr-rr*r5rArrdictZ call_dictrrZmutationr,rrrLrrvrbZguard_on_key_orderaddnameDictKeysVariableDictValuesVariableclonercr is_mutabler[rrr rr`rr"Zcall_custom_dictra call_methodrr)r:rrrrdr-Z arg_hashableZ temp_dict_vthas_argZ has_kwargsZdict_vtrrCr4r]valre)rr>r*r rms         &                          zConstDictVariable.call_methodcCs|dd|jDS)NcSg|]}|jqSrrr}rrr r~z9ConstDictVariable.unpack_var_sequence..)rvr,rr:rrrr unpack_var_sequencerxz%ConstDictVariable.unpack_var_sequencecCs@|jtur||jjvrtdStdStd|jddS)NTFz hasattr on z is not supported)r6r__dict__rrr )r:rrrrr call_obj_hasattr s    z"ConstDictVariable.call_obj_hasattrc s|tjdi|SNr)rvrar)r:rdrerr rszConstDictVariable.clonerrrdrr7r)rSrTrUrZ_nonvar_fieldsrArr;r&rqrDrwrzrYr|rLrrrrrrrrvrrrrr __classcell__rrrer r5[sVK "   !  ' r5csXeZdZdeddffdd ZddZdd Zd ed d ee d fdd fd dZ Z S)MappingProxyVariabledv_dictr7Nc s*tjdi|t|tsJ||_dSr)rar;r"r5rr:rrdrerr r;s zMappingProxyVariable.__init__cCs |j|Sr)rrrrrr rs z(MappingProxyVariable.unpack_var_sequencecs>|jrtdfdd|jtdddS)NzgCan't reconstruct an existing mapping variable because the connection to the original dict will be lostcr)NMappingProxyType)rrtypesrrrrr r)rz2MappingProxyVariable.reconstruct..rF)rbr rrrrr:rrrr r!s  z MappingProxyVariable.reconstructrrrdcCs,|jr |jjr td|j||||S)NaA dict has been modified while we have an existing mappingproxy object. A mapping proxy object, as the name suggest, proxies a mapping object (usually a dict). If the original dict object mutates, it is reflected in the proxy object as well. For an existing proxy object, we do not know the original dict it points to. Therefore, for correctness we graph break when there is dict mutation and we are trying to access a proxy object.)rbrrZhas_existing_dict_mutationr rrr:rrrrdrrr r3s  z MappingProxyVariable.call_method) rSrTrUr5r;rrlistrstrrrrrrer rs rc@seZdZddZddZdS)NNModuleHooksDictVariablecCdSrrrGrrr rvIz7NNModuleHooksDictVariable.install_dict_keys_match_guardcCrrrr:rrrrr rLrz5NNModuleHooksDictVariable.install_dict_contains_guardN)rSrTrUrvrrrrr rGs rcs^eZdZddfdd ZfddZfddZed d Z   dfdd ZZ S)DefaultDictVariableNr7c s.tj||fi||tjusJ||_dSr)rar;r defaultdictdefault_factory)r:r,r6rrdrerr r;Qs zDefaultDictVariable.__init__cs$|jtttfvr |js dStS)NF)rrrBrr,rais_python_constantrGrerr rVs z&DefaultDictVariable.is_python_constantcsd|jdtdS)Nz defaultdict(ro))rrqrarGrerr rq]szDefaultDictVariable.debug_reprcCs,t|tjr|jttttfvSt|tjj Sr) r"rr-fnrrBrsetZ functionsZBaseUserFunctionVariablerrrr is_supported_argbs z$DefaultDictVariable.is_supported_argrrrdrrcs|dkr=t|dks J|d|vr|||dS|jdur&t|d|j|gi}t|d|d|f||St||||S)Nrrrr)rLrrr call_functionrar)r:rrrrdZ default_varrerr ris  zDefaultDictVariable.call_methodrrRr) rSrTrUr;rrqrXrrrrrrer rPs   rcseZdZdZdeeddffdd ZddZed d Z e d d Z d dZ ddZ ddZddZdeedeeefddffdd ZdddefddZdd Zd!d"ZZS)# SetVariablez-We model a sets as dictonary with None valuesr,r7Nc s(t|t}tj|fi|dSr)rfromkeysr_default_valuerar;r:r,rdrerr r;szSetVariable.__init__cC,|jsdSdddd|jDdS)Nzset()rn,cs|]}|jVqdSrrprrkrrr r!r?z)SetVariable.debug_repr..rrr,rsrrGrrr rq"zSetVariable.debug_reprcCst|jSr)rr,rrGrrr set_itemsszSetVariable.set_itemscCs tdSr)rrrrrr rs zSetVariable._default_valuecCdd|jDS)NcSh|]}|jqSrrirrrr z'SetVariable.as_proxy..rrGrrr r&zSetVariable.as_proxycCtSr)rrGrrr rzrzSetVariable.python_typecCr)NcSrrrtrrrr rrz1SetVariable.as_python_constant..rrGrrr rDrzSetVariable.as_python_constantcCs2|dd|jD|tdt|jddS)NcSrrrr}rrr r~rz+SetVariable.reconstruct.. BUILD_SETr)foreachrrr rLrrrr rszSetVariable.reconstructrrdrcsj|dkr|rJt|dksJd}|dtf}n|dkr:|r$J|r(J|jj}t|||f||S|dkrY|rBJt|dksJJt t j  |||dgiS|dkrx|raJt|dksiJt t j  |||dgiS|dkr|rJt|dksJt t j |||dgiS|d kr|rJt|dksJt t j |||dgiS|d krt|dkr|r|rJt|dksJt t j |||dgiS|d kr|rJt|dksJ|d|vrtd t|d||S|d kr,|r Jt|dksJ|d|vr&t|d||StjddSt||||S)Nrrrrr` isdisjoint intersectionunion differencerremovezkey does not existdiscard)rC)rLrrrr`r9rarrr.rZset_isdisjointrZset_intersectionZ set_unionZset_differencerZ set_updater rr)r:rrrrdresultrerr rsx        zSetVariable.call_methodrrrcCstd)NzIllegal to getitem on a set) RuntimeError)r:rrrrr rszSetVariable.getitem_constcCrrrrGrrr rvz)SetVariable.install_dict_keys_match_guardcCrrrrrrr rrz'SetVariable.install_dict_contains_guard)rSrTrUrVrrr;rqrWrrXrr&rzrDrrrrrrvrrrrrer rs4   Arcs|eZdZdeeddffdd ZddZedd Zd d Z d d Z ddZ deede e efddffdd ZZS)FrozensetVariabler,r7Nc tj|fi|dSrrar;rrerr r;zFrozensetVariable.__init__cCr)Nz frozenset()rnrcsrrrprrrr r!r?z/FrozensetVariable.debug_repr..rrrrGrrr rqrzFrozensetVariable.debug_reprcCs |jSr)r,rrGrrr r zFrozensetVariable.set_itemscCrr) frozensetrGrrr rz rzFrozensetVariable.python_typecCr)NcSrrrtrrrr rrz7FrozensetVariable.as_python_constant..rrGrrr rDrz$FrozensetVariable.as_python_constantcs<dd|jDfddtdddS)NcSrrrr}rrr r~rz1FrozensetVariable.reconstruct..csdgS)Nr)rZcreate_load_globalrrrr rsz/FrozensetVariable.reconstruct..rF)rrrrrrrrr rs  zFrozensetVariable.reconstructrrdrc*|dvr td|dt||||S)Nrr`rrrrIllegal call_method z on a frozensetrrarrrerr rzFrozensetVariable.call_method)rSrTrUrrr;rqrWrrzrDrrrrrrrrer rs&   rcsteZdZdeeddffdd ZddZedd Zd d Z d d Z deede e efddffdd Z ZS)DictKeySetVariabler,r7Nc rrrrrerr r;)rzDictKeySetVariable.__init__cCr)Nz dict_keys([])z dict_keys([rcsrrrprrrr r!6r?z0DictKeySetVariable.debug_repr..z])rrGrrr rq0szDictKeySetVariable.debug_reprcCryrrlrGrrr r:szDictKeySetVariable.set_itemscCrrrrGrrr rz>rzDictKeySetVariable.python_typecCstdd|jDdS)NcSrrrtrrrr rCrz8DictKeySetVariable.as_python_constant..)rrrrrGrrr rDAs z%DictKeySetVariable.as_python_constantrrdrcr)Nrrz on a dict_keysrrrerr rFrzDictKeySetVariable.call_method)rSrTrUrrr;rqrWrrzrDrrrrrrrer r(s$  rcseZdZUdZdZeeed<deddffdd Z e dd Z e d d Z d d Z ddZdeddeedfddffdd ZZS)DictViewVariablezu Models _PyDictViewObject This is an "abstract" class. Subclasses will override kv and the items method Nkvrr7c s8tjdi||jdvsJt|tsJ||_dS)N)rrr)rar;r r"r5rrrerr r;[s zDictViewVariable.__init__cCst|jj|jSr)getattrrr,r rGrrr view_itemsaszDictViewVariable.view_itemscCstr)NotImplementedErrorrGrrr view_items_vteszDictViewVariable.view_items_vtcs fddfddjDS)Ncsjdkr|jS|S)Nr)r r9r3rGrr unwraplrmz4DictViewVariable.unpack_var_sequence..unwrapcsg|]}|qSrrr})rrr r~osz8DictViewVariable.unpack_var_sequence..r rr)r:rr rks z$DictViewVariable.unpack_var_sequencecCs$||j||j|ddS)Nr)rZ load_methodr rrrrr rqs  zDictViewVariable.reconstructrrrdcs,|dkr |j||||St||||S)Nr)rrrarrerr rvszDictViewVariable.call_method)rSrTrUrVr rr__annotations__r5r;rWr rrrrrrrrrrer r Rs"    r csZeZdZdZeddZeddZddZded d e e d fd d ffd d Z Z S)rrcCrEr)rr rGrrr rrzDictKeysVariable.set_itemscCr)NcSrrrr}rrr r~rz2DictKeysVariable.view_items_vt..rrGrrr rszDictKeysVariable.view_items_vtcCrrr rGrrr rzrzDictKeysVariable.python_typerrrdr7csl|dkr |j||||S|tvr-t|dttfsttStt||j |dj St ||||S)Nr|r) rrrr"rrrrNotImplementedrrarrerr rs zDictKeysVariable.call_method) rSrTrUr rWrrrzrrrrrrrrer rs   rc@s$eZdZdZeddZddZdS)rrcCrEr)rr rGrrr rrz DictValuesVariable.view_items_vtcCrr)rrGrrr rzrzDictValuesVariable.python_typeN)rSrTrUr rWrrzrrrr rs   r).rVrrrtypingrrZtorch._subclasses.fake_tensorrrrrZbytecode_transformationrr excr r Zguardsr r rbrutilsrrrrbaserrZconstantrZtorch._dynamo.symbolic_convertrrr5rrrrrrr rrrrrr s<    *?0 5t/*0!