o ߗix@sUddlZddlZddlZddlmZddlmZddlmZm Z m Z m Z ddl m Z ddlmZddlmZgd Ziae eefed <d ae d d Gd ddeeZe d d GdddZe d d eddZe d d de efddZe d d ddZe d d ddZe d d de efddZe d d defddZe d d ed'd d!Z e d d de eeffd"d#Z!e d d d$edefd%d&Z"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_compatiblec@seZdZdZdZdS)rcreatereplaceN)__name__ __module__ __qualname__CREATEREPLACEr r P/home/ubuntu/transcripts/venv/lib/python3.10/site-packages/torch/fx/traceback.pyr src@seZdZUdZGdddZeed<eded<eded<eded< dd ee dededfd d Z e defddZ e defddZ e defddZddZdddZdefddZd 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@s"eZdZdededefddZdS)zNodeSource.NodeInfonametargetgraph_idcCs||_||_||_dSNr"r#r$)selfr"r#r$r r r!__init__.s zNodeSource.NodeInfo.__init__N)rrrstrintr(r r r r!NodeInfo-sr+ pass_nameraction from_node node_infoNnodecCsj||_||_|r-|j|jt|jt|jd|_d|j vr(t |j d|_ dSg|_ dSd|_g|_ dS)Nr&r.) r,r-r+r"r)r#idgraphr/metacopydeepcopyr.)r'r1r,r-r r r!r(8s  zNodeSource.__init__returncC|jr|jjSdSNr0)r/r"r'r r r!r"MzNodeSource.namecCr8r9)r/r#r:r r r!r#Qr;zNodeSource.targetcCr8)N)r/r$r:r r r!r$Ur;zNodeSource.graph_idcCs|Sr%)print_readabler:r r r!__repr__YszNodeSource.__repr__rc Csh|dkrdSd}|d|dd|jd|jd|jd|jd 7}|jD] }|||d 7}q&|S) N r0 z(name=z , pass_name=z , action=z , graph_id=z) r)r"r,r-r$r.r=)r'indentresultitemr r r!r=\s $ zNodeSource.print_readablecCs(|j|j|j|j|jdd|jDdS)NcSg|]}|qSr to_dict).0r1r r r! pz&NodeSource.to_dict..r"r#r$r,r-r.rKr:r r r!rGhszNodeSource.to_dict)r0N)r)rrr__doc__r+r)__annotations__rrr r(propertyr"r#r*r$r>r=dictrGr r r r!r&s2       rccs2t}t}z dadVW|a|adS|a|aw)NT)should_preserve_node_metarr5)saved_should_preserve_node_metasaved_current_metar r r!r tsr stackcCs"tr |rd|td<dSdSdS)Nr0 stack_trace)rPjoinr)rSr r r!rsrcCs6trtdg|gtd<tdddtd<dSdS)Ngrad_fn_seq_nr in_grad_fnrrrPrget)seq_nrr r r!rs   rcCs\tr,tdd}|dksJ|dkrtd=td=dS|dtd<tdddtd<dSdS)NrWrrrVr<rX) current_levelr r r!rs    rr7cCs(tr tddgSttddS)NrTr0r<)rPrrY traceback format_list extract_stackr r r r!rsrcCtSr%)rPr r r r!r r r0ccsPtr#|jr#t}z|jat||tjgtd<dVW|adS|awdVdS)Nr.)rPr4rr5rrr)r1r, saved_metar r r!rs    rcCr_r%)rr r r r!rr`rr3cCsLi}|jD]}|jdkr d|jvrdd|jdDng||j<qt|S)zh Given an fx.Graph, return a json string that contains the provenance information of each node. call_functionr.cSrEr rF)rHsourcer r r!rIrJz-get_graph_provenance_json..)nodesopr4r"jsondumps)r3provenance_tracking_jsonr1r r r!rs    r)r0)#r5rfr\ contextlibrenumrtypingrrrr_compatibilityr r3r r1r __all__rr)rMrPrrr rrrrboolr rrrr r r r!sH      M