jXUddlZddlZddlZddlmZddlmZddlmZm Z m Z m Z ddl m Z ddlmZddlmZgd Ziae eefed <d ae d Gd deeZe d GddZe d edZe d de efdZe d dZe d dZe d de efdZe d defdZe d eddZ e d de eeffdZ!e d dedefdZ"dS)N)contextmanager)Enum)AnyDictListOptional) compatibility)Graph)Node) preserve_node_metahas_preserved_node_metaset_stack_traceset_grad_fn_seq_nrreset_grad_fn_seq_nr format_stackset_current_metaget_current_meta NodeSourceNodeSourceActionget_graph_provenance_json current_metaF)is_backward_compatibleceZdZdZdZdS)rcreatereplaceN)__name__ __module__ __qualname__CREATEREPLACECC:\Projects\Historical\venv\Lib\site-packages\torch/fx/traceback.pyrr s FGGGr#rceZdZUdZGddZeed<eded<eded<eded< dd ee dededfd Z e d efdZ e d efdZ e d efdZdZddZd efdZd S)rz NodeSource is a data structure that contains the provenance information of a node. If node `a` is created from node `b`, then `a.meta["from_node"]` may contain NodeSource(b). c"eZdZdededefdZdS)NodeSource.NodeInfonametargetgraph_idc0||_||_||_dSNr(r)r*)selfr(r)r*s r$__init__zNodeSource.NodeInfo.__init__.sDI DK$DMMMr#N)rrrstrintr/r"r#r$NodeInfor'-s= % %c %S % % % % % %r#r2 pass_nameraction from_node node_infoNnodec0||_||_|rw||jt |jt |j|_d|j vrtj |j dng|_ dSd|_g|_ dS)Nr-r5) r3r4r2r(r0r)idgraphr6metacopydeepcopyr5)r.r8r3r4s r$r/zNodeSource.__init__8s #  !]]Ys4;'7'7"TZ..+DN $)++ di 4555 NNN "DNDNNNr#returnc,|jr |jjndSNr7)r6r(r.s r$r(zNodeSource.nameMs&*nt~$$B>r#c,|jr |jjndS)N)r6r*rBs r$r*zNodeSource.graph_idUs*..@t~&&b@r#c*|Sr,)print_readablerBs r$__repr__zNodeSource.__repr__Ys""$$$r#rc |dkrdSd}|d|zdzd|jd|jd|jd|jd zz }|jD]}|||d zz }|S) N r7 z(name=z , pass_name=z , action=z , graph_id=z) r )r(r3r4r*r5rG)r.indentresultitems r$rGzNodeSource.print_readable\s A::2 &L1 qtyqqdnqqt{qq_c_lqqq r N 6 6D d))&1*55 5FF r#cd|j|j|j|j|jd|jDdS)Nc6g|]}|Sr"to_dict).0r8s r$ z&NodeSource.to_dict..ps DDDT$,,..DDDr#r(r)r*r3r4r5rVrBs r$rSzNodeSource.to_dicths?Ik kDDT^DDD    r#)r7N)r)rrr__doc__r2r0__annotations__rrr r/propertyr(r)r1r*rHrGdictrSr"r#r$rr&s %%%%%%%% NNN ' ((((L!!!! #### /3   tn  +,     *=c===X=????X?A#AAAXA%%%             r#rc#rKt}t} dadV|a|adS#|a|awxYw)NT)should_preserve_node_metarr=)saved_should_preserve_node_metasaved_current_metas r$r r ts_ '@#%**,,*$(! $C!) %D!) ))))s06stackcZtr!|r!d|td<dSdSdS)Nr7 stack_trace)r\joinr)r_s r$rrsA!5U5&(ggenn ]###5555r#ctrOtdg|gztd<tdddztd<dSdS)Ngrad_fn_seq_nr in_grad_fnrr r\rget)seq_nrs r$rrsj!K)5)9)9:JB)O)O S * %&&2%5%5lA%F%F%J \""" KKr#ctrgtdd}|dksJ|dkrtd=td=dS|dz td<tdddtd<dSdS)Nrerr rdrErf) current_levels r$rrs !Q$((q99 q    A  \*-...)6):L &-9:J-KCRC-PL) * * *QQr#r?ctrtddgStjtjddS)Nrar7rE)r\rrg traceback format_list extract_stackr"r#r$rrsI E  3344$Y%<%>%>ss%CDDDr#ctSr,)r\r"r#r$rrs $$r#r7c#Ktr[|jrTt} |jat ||t jgtd<dV|adS#|awxYwdVdS)Nr5)r\r<rr=rrr )r8r3 saved_metas r$rrs!TY!  &9>>++L4,<,CDD)L % EEE%LLL:L % % % % s AAA"ctSr,)rr"r#r$rrs r#r;ci}|jD]7}|jdkr*d|jvrd|jdDng||j<8t j|S)zh Given an fx.Graph, return a json string that contains the provenance information of each node. call_functionr5c6g|]}|Sr"rR)rTsources r$rUz-get_graph_provenance_json..s"GGGf!!GGGr#)nodesopr<r(jsondumps)r;provenance_tracking_jsonr8s r$rrsx "  7o % %$)++HG +0FGGGG %TY / :. / //r#)r7)#r=ryrl contextlibrenumrtypingrrrr_compatibilityr r;r r8r __all__rr0rXr\rrr rrrrboolrrrrr"r#r$rsW %%%%%%,,,,,,,,,,,,))))))    " d38n!!!!e,,,sD-, e,,,J J J J J J J -,J Ze,,, * *-, *e,,,549555-,5e,,,KK-,Ke,,, Q Q-, Qe,,,Ed3iEEE-,Ee,,,%%%%-,%e,,,-,*e,,,$sCx.-,e,,, 0U 0s 0 0 0-, 0 0 0r#