o I h@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.0erQ/var/www/vscode/kcb/lib/python3.10/site-packages/torch/_dynamo/variables/dicts.py Cszis_hashable..) isinstancerLazyVariableTracker is_realizedrTensorVariableas_proxynodemetaget TupleVariableallitemsBuiltinVariableSymNodeVariabler EnumVariable user_definedUserDefinedClassVariableUserFunctionVariableSkipFunctionVariablemisc NumpyVariableNNModuleVariableUnspecializedNNModuleVariableMethodWrapperVariableTorchInGraphFunctionVariableTypingVariableFunctoolsPartialVariableWeakRefVariablexrrr 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)selfrCrrr __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"rCrr#r$roriginal_valuer%r&r'r(r*r?_HashableTrackertupler,r6valuer7r2 get_functionr< referent_vtrFas_python_constant)rDr>rrGr rFqs2     z3ConstDictVariable._HashableTracker.underlying_valuecC t|jSr)hashrFrDrrr __hash__s z+ConstDictVariable._HashableTracker.__hash__csht|t|kr dSt|tr(tjt|t|ko'tfddt||DSt|r0||uS||kS)NFc3s |] \}}||VqdSr)_eq_impl)ruvrGrr r!s z>ConstDictVariable._HashableTracker._eq_impl..) rBr"rLr?rKlenr+zipr)abrrGr rUs z+ConstDictVariable._HashableTracker._eq_implothercCsNtj}t||st|sJt|t||r ||j|jS||j|Sr)r?rKr"r is_literalrBrUrF)rDr\rHrrr __eq__s z)ConstDictVariable._HashableTracker.__eq__rAN) __name__ __module__ __qualname____doc__rEpropertyrFrT staticmethodrUboolr^rrrr rKas   rKr,rANc sd|vr |dd|vr|dtjd i|tjtfdd|Ds-Jfddfdd|D|_t|j |_ | |_ ||_ dS) Noriginal_itemsshould_reconstruct_allc3s,|]\}}t|tfot|tVqdSr)r"rrr>rWrGrr r!s  z-ConstDictVariable.__init__..cst|r|S|Sr)r")keyrGrr make_hashablesz1ConstDictVariable.__init__..make_hashablecsi|] \}}||qSrrri)rkrr sz.ConstDictVariable.__init__..r) popsuperrEr?rKr+r,rsourcerhcopyrgr@)rDr,r@kwargs __class__)rHrkr rEs      zConstDictVariable.__init__cCdd|jDS)NcS i|] \}}|j|qSrrCr&rkrWrrr rls z.ConstDictVariable.as_proxy..r,rSrrr r&zConstDictVariable.as_proxycCs"dddd|jDdS)N{, css,|]\}}|jd|VqdS)z: NrC debug_reprrwrrr r!s z/ConstDictVariable.debug_repr..})joinr,rSrrr r~s zConstDictVariable.debug_reprcCrt)NcSrurrCrPrwrrr rlz8ConstDictVariable.as_python_constant..ryrSrrr rPsz$ConstDictVariable.as_python_constantcCs|dd|jDS)NcSsi|] \}}|j|qSrrrwrrr rlsz=ConstDictVariable.keys_as_python_constant..)install_dict_keys_match_guardr,rSrrr keys_as_python_constantz)ConstDictVariable.keys_as_python_constantcC|jSr)r@rSrrr python_typeszConstDictVariable.python_typecCsBt|tsJtj}t|o |||jvo t|j||tj Sr)r"rr?rKrr,rDeletedVariable)rDrCrHrrr __contains__s zConstDictVariable.__contains__cCstdd|jDS)NcSsg|] }t|tjs|qSr)r"rrrr>rrr s  z)ConstDictVariable.len..)rXr,valuesrSrrr rXs zConstDictVariable.lencs&jrdStfddjDS)NTc3s*|]\}}j|j|VqdSr) is_new_itemrgr)rC)rrjrMrSrr r!s  z2ConstDictVariable.has_new_items..)rhanyr,rSrrSr has_new_itemss  zConstDictVariable.has_new_itemscCs<|r|r|rt|t|kSt|t|kSr)r$idrealize)rDrMr\rrr 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)r@rr add_push_nullr,rgr)rCrrhrr r append_output)rDrnum_argsrjrMitemrrr reconstructs(      zConstDictVariable.reconstructtxrrcCs(t|}||jvrtt||j|Sr)r?rKr,r KeyErrorrDrrrjrrr 'getitem_const_raise_exception_if_absent%s    z9ConstDictVariable.getitem_const_raise_exception_if_absentcCs.t|}||jvrtd|j|j|S)Nzdict KeyError: )r?rKr,r rMrrrr getitem_const-s   zConstDictVariable.getitem_constcCs"t|}||jvr dS|j|Sr)r?rKr,)rDrrjrrr maybe_getitem_const3s   z%ConstDictVariable.maybe_getitem_constcCsX||vsJt|}t|j|}t|j|}t|tj r*| dSdSr) r?rKrLr,keysindexrgr"rr#r)rDrrjroriginal_key_vtrrr realize_key_vt9s    z ConstDictVariable.realize_key_vtcCs|jr t|tjdSdSr)ror make_guardr DICT_KEYS_MATCHrSrrr rBsz/ConstDictVariable.install_dict_keys_match_guardcCsddlm}|js dS|jj|rdS|d|v}|djdur.rrrp)r, mutation_typero__len__ __setitem__r __delitem__T)rmr))rrrmclearupdatecs i|] \}}||qSr)createrw)rrHrr rlrz1ConstDictVariable.call_method..)r) __getattr__r setdefault move_to_end)*rr-rr*r?rKrrdict call_dictrrmutationr,rrrXrrroguard_on_key_orderaddnameDictKeysVariableDictValuesVariableclonerpr is_mutablerhrrr rrmrr"call_custom_dictrn call_methodrr)rDrrrrqr- arg_hashable temp_dict_vthas_arg has_kwargsdict_vtrrMr>rjvalrr)rrHr*r rms         &                          zConstDictVariable.call_methodcCs|dd|jDS)NcSg|]}|jqSrrrrrr rz9ConstDictVariable.unpack_var_sequence..)rr,rrDrrrr unpack_var_sequencerz%ConstDictVariable.unpack_var_sequencecCs@|jtur||jjvrtdStdStd|jddS)NTFz hasattr on z is not supported)r@r__dict__rrr )rDrrrrr call_obj_hasattr s    z"ConstDictVariable.call_obj_hasattrc s|tjdi|SNr)rrnr)rDrqrrrr rszConstDictVariable.clonerrrqrrAr)r`rarbr_nonvar_fieldsrKrrEr&r~rPrrrfrrXrrrrrrrrrrrrr __classcell__rrrrr r?[sVK "   !  ' r?csXeZdZdeddffdd ZddZdd Zd ed d ee d fdd fd dZ Z S)MappingProxyVariabledv_dictrANc s*tjdi|t|tsJ||_dSr)rnrEr"r?rrDrrqrrrr rEs 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)ror rrrrrDrrrr r!s  z MappingProxyVariable.reconstructrrrqcCs,|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.)rorrhas_existing_dict_mutationr rrrDrrrrqrrr r3s  z MappingProxyVariable.call_method) r`rarbr?rErrlistrstrrrrrrrr rs rc@seZdZddZddZdS)NNModuleHooksDictVariablecCdSrrrSrrr rIz7NNModuleHooksDictVariable.install_dict_keys_match_guardcCrrrrDrrrrr rLrz5NNModuleHooksDictVariable.install_dict_contains_guardN)r`rarbrrrrrr rGs rcs^eZdZddfdd ZfddZfddZed d Z   dfdd ZZ S)DefaultDictVariableNrAc s.tj||fi||tjusJ||_dSr)rnrEr defaultdictdefault_factory)rDr,r@rrqrrrr rEQs zDefaultDictVariable.__init__cs$|jtttfvr |js dStS)NF)rrrLrr,rnis_python_constantrSrrrr rVs z&DefaultDictVariable.is_python_constantcsd|jdtdS)Nz defaultdict(r|))rr~rnrSrrrr r~]szDefaultDictVariable.debug_reprcCs,t|tjr|jttttfvSt|tjj Sr) r"rr-fnrrLrset functionsBaseUserFunctionVariablerrrr is_supported_argbs z$DefaultDictVariable.is_supported_argrrrqrrcs|dkr=t|dks J|d|vr|||dS|jdur&t|d|j|gi}t|d|d|f||St||||S)Nrrrr)rXrrr call_functionrnr)rDrrrrq default_varrrrr ris  zDefaultDictVariable.call_methodrr_r) r`rarbrErr~rerrrrrrrr 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,rANc s(t|t}tj|fi|dSr)rfromkeysr_default_valuernrErDr,rqrrrr rEszSetVariable.__init__cC,|jsdSdddd|jDdS)Nzset()r{,cs|]}|jVqdSrr}rrxrrr r!rIz)SetVariable.debug_repr..rr,rrrSrrr r~"zSetVariable.debug_reprcCst|jSr)rr,rrSrrr set_itemsszSetVariable.set_itemscCs tdSr)rrrrrr rs zSetVariable._default_valuecCdd|jDS)NcSh|]}|jqSrrvrrrr z'SetVariable.as_proxy..r rSrrr r&zSetVariable.as_proxycCtSr)rrSrrr rrzSetVariable.python_typecCr )NcSr rrrrrr r r z1SetVariable.as_python_constant..rrSrrr rPrzSetVariable.as_python_constantcCs2|dd|jD|tdt|jddS)NcSrrrrrrr rrz+SetVariable.reconstruct.. BUILD_SETr)foreachr rr rXrrrr rszSetVariable.reconstructrrqrcsj|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)Nrrrrrm isdisjoint intersectionunion differencerremovezkey does not existdiscard)rM)rXrrr rmrCrnrrr2rset_isdisjointrset_intersection set_unionset_differencer set_updater rr)rDrrrrqresultrrrr rsx        zSetVariable.call_methodrrrcCstd)NzIllegal to getitem on a set) RuntimeError)rDrrrrr rszSetVariable.getitem_constcCrrrrSrrr rz)SetVariable.install_dict_keys_match_guardcCrrrrrrr rr z'SetVariable.install_dict_contains_guard)r`rarbrcrrrEr~rdr rerr&rrPrrrrrrrrrrrrr rs4   Arcs|eZdZdeeddffdd ZddZedd Zd d Z d d Z ddZ deede e efddffdd ZZS)FrozensetVariabler,rANc tj|fi|dSrrnrErrrrr rEzFrozensetVariable.__init__cCr)Nz frozenset()r{rcsrrr}rrrr r!rIz/FrozensetVariable.debug_repr..rrrSrrr r~rzFrozensetVariable.debug_reprcCs |jSr)r,rrSrrr r  zFrozensetVariable.set_itemscCrr) frozensetrSrrr r rzFrozensetVariable.python_typecCr )NcSr rrrrrr r r z7FrozensetVariable.as_python_constant..rrSrrr rPrz$FrozensetVariable.as_python_constantcs<dd|jDfddtdddS)NcSrrrrrrr rrz1FrozensetVariable.reconstruct..csdgS)Nr&)rcreate_load_globalrrrr rsz/FrozensetVariable.reconstruct..rF)rr rrrrrrr rs  zFrozensetVariable.reconstructrrqrc*|dvr td|dt||||S)NrrmrrrrIllegal call_method z on a frozensetrrnrrrrrr rzFrozensetVariable.call_method)r`rarbrrrEr~rdr rrPrrrrrrrrrr r!s&   r!csteZdZdeeddffdd ZddZedd Zd d Z d d Z deede e efddffdd Z ZS)DictKeySetVariabler,rANc r"rr#rrrrr rE)r$zDictKeySetVariable.__init__cCr)Nz dict_keys([])z dict_keys([rcsrrr}rrrr r!6rIz0DictKeySetVariable.debug_repr..z])rrSrrr r~0szDictKeySetVariable.debug_reprcCrrryrSrrr r :szDictKeySetVariable.set_itemscCrrrrSrrr r>rzDictKeySetVariable.python_typecCstdd|jDdS)NcSr rrrrrr r Cr z8DictKeySetVariable.as_python_constant..)rrr rrSrrr rPAs z%DictKeySetVariable.as_python_constantrrqrcr()Nr)r*z on a dict_keysr+rrrrr rFr,zDictKeySetVariable.call_method)r`rarbrrrEr~rdr rrPrrrrrrrrr r-(s$  r-cseZdZUdZdZeeed<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 NkvrrAc s8tjdi||jdvsJt|tsJ||_dS)N)rrr)rnrEr0r"r?rrrrrr rE[s zDictViewVariable.__init__cCst|jj|jSr)getattrrr,r0rSrrr view_itemsaszDictViewVariable.view_itemscCstr)NotImplementedErrorrSrrr view_items_vteszDictViewVariable.view_items_vtcs fddfddjDS)Ncsjdkr|jS|S)Nr)r0rCr=rSrr unwraplrzz4DictViewVariable.unpack_var_sequence..unwrapcsg|]}|qSrrr)r5rr rosz8DictViewVariable.unpack_var_sequence..r2rr)rDr5r rks z$DictViewVariable.unpack_var_sequencecCs$||j||j|ddS)Nr)r load_methodr0rrrrr rqs  zDictViewVariable.reconstructrrrqcs,|dkr |j||||St||||S)Nr)rrrnrrrrr rvszDictViewVariable.call_method)r`rarbrcr0rr__annotations__r?rErdr2r4rrrrrrrrrrr r/Rs"    r/csZeZdZdZeddZeddZddZded d e e d fd d ffd d Z Z S)rrcCrQr)rr2rSrrr r r%zDictKeysVariable.set_itemscCr )NcSrrrrrrr rrz2DictKeysVariable.view_items_vt..r6rSrrr r4szDictKeysVariable.view_items_vtcCrrr.rSrrr rrzDictKeysVariable.python_typerrrqrAcsl|dkr |j||||S|tvr-t|dttfsttStt||j |dj St ||||S)Nrr) rrrr"rrrrNotImplementedr rnrrrrr rs zDictKeysVariable.call_method) r`rarbr0rdr r4rrrrrrrrrrr rs   rc@s$eZdZdZeddZddZdS)rrcCrQr)rr2rSrrr r4r%z DictValuesVariable.view_items_vtcCrr)rrSrrr rrzDictValuesVariable.python_typeN)r`rarbr0rdr4rrrrr rs   r).rcrrrtypingrrtorch._subclasses.fake_tensorrrrrbytecode_transformationrr excr r guardsr r rorutilsrrrrbaserrconstantrtorch._dynamo.symbolic_convertrrr?rrrrr!r-r/rrrrrr s<    *?0 5t/*0!