Vh & ddlZddlZddlZddlZddlZddlZddlZddlmZddl m Z ddlm Z ddl m Z mZmZmZmZmZddlmZddlmZddlmZmZmZdd lmZdd lmZdd lm Z m!Z!dd l"m#Z#m$Z$dd l%m&Z&er ddl'Z'ddl(m)Z)ddl*Z*ddl+m,cm-Z.ddl/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5m6Z6m7Z7m8Z8m9Z9m:Z:m;Z;mZ>ddl?m@Z@ddlAmBZBddlmCZCddlDmEZEmFZFddlGmHZHddlImJZJddlKmLZLddlMmNZNddlOmPZPmQZQmRZRmSZSmTZTmUZUmVZVmWZWmXZXmYZYmZZZm[Z[m\Z\gdZ]ee*jjgeeLfZ`ejGddZbejGddZcdZdd eeed!e.jd"eeed#e.jd$egf d%Zhe*jjje*jjje*jjje*jjje*jjje*jjje*jjje*jjje*jjje*jjje*jjje*jjjg Zwe*jjje*jjje*jjje*jjje*jjje*jjje*jjje*jjje*jjjg Ze d@d&Ze d'Zd(ee*j j efd$eee*j j efd)ffd*ZdAd+Zd,ee*j j efd-ee*j j efd.eefd/Zd0e*jjd1eSd$ee*jjeSffd2Zd0e*jjfd3Zd0e*jjd4eecfd5Zd,ee*j j efd-ee*j j efd.eed6egfd7ZGd8d9Zd:Z dBd0e*jjd;dZd?Zy)CN) namedtuple)Iterator)contextmanager)AnyCallablefinalOptional TYPE_CHECKINGUnion)autograd_not_implemented)FakeScriptObject)_deregister_op_impl_is_op_registered_to_fake_ruleregister_op_impl)FakeTensorMode)#first_call_function_nn_module_stack)_PyTreeCodeGen _PyTreeInfo)immutable_dictimmutable_list)insert_deferred_runtime_asserts) ValueRanges)_collect_all_valid_cia_ops_collect_and_set_constant_attrs_collect_param_buffer_metadata_detect_fake_mode_from_gm_fakify_params_buffers_get_decomp_for_cia_is_preservable_cia_op_name_hoo_subgraph_placeholders7_override_graph_signature_for_temp_registered_constants/_overwrite_signature_for_non_persistent_buffers)_populate_param_buffer_metadata_to_new_gm_register_constants_as_buffers_rename_without_collisions_special_op_to_preserve_ciaplaceholder_naming_pass)Verifierdetect_fake_modeunset_fake_temporarily) is_equivalentreorder_kwargsCustomDecompTable) compatibility) PassResult) PassManager) ArgumentSpecConstantArgumentCustomObjArgumentExportGraphSignature InputKind InputSpec OutputKind OutputSpecSymBoolArgumentSymFloatArgumentSymIntArgumentTensorArgument TokenArgument)ExportedProgramModuleCallEntryModuleCallSignaturedefault_decompositionsceZdZUeeed<eeed<ejed<ejed<dZe ee ed<dZ y)rDinputsoutputsin_specout_specNforward_arg_namesc|jD]-}|j|jk(s|j|_/|jD]-}|j|jk(s|j|_/yN)rGnamerH)self original_nodenew_nodeios M/home/dcms/DCMS/lib/python3.12/site-packages/torch/export/exported_program.pyreplace_all_uses_withz)ModuleCallSignature.replace_all_uses_withfsd 'Avv+++! ' 'Avv+++! ') __name__ __module__ __qualname__listr5__annotations__pytreeTreeSpecrKr strrUrVrTrDrD^sE   ,  __oo-1xS *1'rVrDc,eZdZUeed<dZeeed<y)rCfqnN signature)rWrXrYr^r[rbr rDr_rVrTrCrCos H/3Ix+,3rVrCcBtjfd}|S)NcTt5|i|cdddS#1swYyxYwrMr+)argskwargsfns rTwrapperz/_disable_prexisiting_fake_mode..wrappervs+ # % 't&v& ' ' 's') functoolswraps)rgrhs` rT_disable_prexisiting_fake_moderkus%__R'' NrV spec1_type spec1_context spec2_type spec2_contextreturnc|| ||uxr||k(St|ttfrt|ttfr||k(St|ttfrt|ttfr||k(S||uxr||k(S)zeTreat containers and their immutable variants as the same type. Otherwise compare as normal. ) issubclassdictrrZr)rlrmrnros rT_fx_collection_equivalence_fnrt~sZ/Z'JM],JJ*t^45*T>*; --*t^45*T>*; --  # F (FFrVc#Ki}t}|jD]z\}}|jj||<|j |t D]3}||jvs|j |t|d5t|}tjjj|jvr/|jtjjj=|r9|j tjjj|d}t|s&t|tj ||t"D](} | |jvs|j | |*} d|D]_} | jj%| jj'|| | j(j%t+| ay#|D]_} | jj%| jj'|| | j(j%t+| awxYww)NT)deferred_errorcR|d}|d=|5||i|cdddS#1swYyxYw)Noriginal_callabler_)fake_tensor_modeoprerforig_cia_callables rT$_force_dispatch_to_orig_cia_callablezQ_override_composite_implicit_decomp.._force_dispatch_to_orig_cia_callables< &': ; *+! :($9&9 : : :s&)rx)setitems py_kernelscopyadd(_AUTOGRAD_ALIAS_BACKEND_KEYS_TO_OVERRIDEpy_implr rtorch_C DispatchKeyCompositeImplicitAutogradrrripartial_BACKEND_KEYS_TO_OVERRIDEclearupdate_dispatch_cacher) cia_ops_to_callablesafe saved_tables patched_ops op_overloaddecomp_callableoverride_dispatch_keyr{r|keyrzs rT#_override_composite_implicit_decomprs L%K(;(A(A(C><$ _$/$:$:$?$?$A [! $%M  !$K,B,BB: ##$9:,[N 0 < 88   9 9[=S=S S&&uxx';';'U'UV  OK   4 4 N N O  : .k: ) [ )!!8&7 - <@$  $B MM   ! MM b!1 2    $ $ &  #  $+ $B MM   ! MM b!1 2    $ $ &  #  $s-A)I0,DI0I0H"A%I0A&I--I0c# Kttjjjj t tjjjjt id5ddddy#1swYyxYww)NF)r)rropsatento dtype_layoutr&dtyper_rVrT!_override_decomp_aten_to_variantsrsg - IINN   * *,G IINN   # #%@    sA2B4B9 BB B decomp_table.cJtt}i}t|jD]_}||vr||||<|j |||=$t |s0|j }|jdrJd||||<a|D] }t||< ||fS)NrzThis should be a custom op) r}rrZkeysremoverrN startswithr&)rall_preservable_cia_opsrrzop_nameks rT,_split_decomp_table_to_cia_and_python_decompr's""<">?<$$&'7$ ( (&22&6  # # * *2 .R $B 'ggiG))&1 O3O O1&22&6  #97@%=!<A=  ,,rVctS)z This is the default decomposition table which contains decomposition of all ATEN operators to core aten opset. Use this API together with :func:`run_decompositions()` r/r_rVrTrErESs  rV cia_to_decomppython_decomp_tablejoint_loss_indexcddlm}ddlm}ddlm}m}m} m} m } m } ddl m } d}|||s|j}it|jdt|j!d}dd lm}||it|jdt|j!d}t'|j(}|t+| d }|j,}|j.j0j2j4}|j6t8t:fvrt=j>t:d|g}tAtC||jD||j._|jGtI|jJ|jL|tO||jP|jJjR}tU||}tW|}dd l,m-}m.}g}|j.j^D]}|j`d k(stc|jddtfrwd}|jddjh'|jL|jddjj}n#|jddjhjl}|jo||jo|jdd|5tq5ts|5||jt5t=jv||jD}||g|d|djyi|5\}} }!}"}#||| |!||"|dd| }$|$jLj{D%&cic]\}%}&|%tc|&t|r|#|&n|&c}&}%|$_&|$j~}'|$j}(t|(|}(t||'|(t|jJ|(}(||'|(|"})t|'|(|| |!||)ddddddddddddddd| '| |'| |'(t|'|(\}'}(|j{D]\}*}+|*|vs |+|jP|*<|j{D]h\}*}+|*|jPvr&tc|+tjjrJ|*|jPvsI|*|vsN|jPj|*j|'|(|jPfS|j(j.j^Dcgc]}|j`d k(s|},}|,Dcgc]}|jdd}-}|j(j!D*.cgc]\}*}.|* }/}*}.|/D]}*t|j(|*t|-}|tjn|}|rtjn|}0| 5|5ts|5|05||j(|-||d nd||nd\}'}1|'j.jddddddddddddd}2'j.j^Dcgc]}|j`d k(s|}3}t9|'j.j^djd}4t|3t|,k(sJt|,|3D]\}5}6|5jjx|6_5|6_R|3D7cic]}7|7jj|7jj}8}7|'j.j^D]8}|j`d k(rt|8|jj|jj|_5:t|'ddlUmV}9mW}:ddlXmY};tjjjs[d}}*tc|4|>tjjs.|*|4|>_59t|jJjD>?cic]C\}>}?tc|?jts$|?jjj|3|>jjE}@}>}?t|jJjD>?cgc]E\}>}?t|?j|2|?j|3|>|?j|?jG}A}>}?g}Bt1jD>*cic]\}>}*|*|> }C}>}*t9|1jj}D|D|4dt|DDcgc]}|jjc}k(sJ|1jj{D]\}E}FC|F}>|jJj|>}G|GjtjtjfvsJGjtjk(rtjntj}HGjtjk(r GjnGjjj}IBjotHtE|It|jJjD]\}>}?Bjot|>|k(rtjn |?j|2|?j|4tD|>z@j|?j|?j||1jJ|1jj}Jt|1jt|jJjk(sJt|jJjD>?cic]/\}>}?tc|?jtr|1j|>|?1}K}>}?t|4tBdD]\}>}J|jj}LK|L}?|?jtjk(rtj}M|?j}InH|?jtjk(rtj}ML}Intd|?jBjotMt|jjIt|3t|,k(sJtAB}(t|,|3D]\}N}Otc|NjddtjsNjddOjdd<Oj|(jvsOj|(jvs}Njdj{D]\}P}Q|QOjd|P<|'|(|jPfScc}&}%w#1swY xYw#1swY xYw#1swY xYw#1swY xYw#1swY xYwcc}wcc}wcc}.}*w#1swY#xYw#1swY(xYw#1swY-xYw#1swY2xYwcc}wcc}7w#1swYxYwcc}?}>wcc}?}>wcc}*}>wcc}wcc}?}>w)Nr)_materialize_and_lift_constants)aot_export_module)2_disable_custom_triton_op_functional_decomposition_export_to_aten_ir_ignore_backend_decomps_verify_nn_module_stack_verify_placeholder_names_verify_stack_trace)ShapeEnvc>|duxs|jjduSrM)graph_signaturebackward_signature)eprs rT_is_joint_ir_decompzO_decompose_and_get_gm_with_new_signature_constants.._is_joint_ir_decompps) D ( A!!44D@ rVF)remove_duplicate)!unwrap_tensor_subclass_parametersT) shape_envexport)&_enable_graph_inputs_of_type_nn_module_fakify_script_objects placeholdervalr4)r_check_autograd_state_prettify_placeholder_namesdecompose_custom_triton_ops)decompositions trace_jointoutput_loss_indexct|tr|St|trt|jSt|trt |jSt|t rt |jSt|t rt |jStdt|)NrNzType of old_arg not supported: ) isinstancer6r@rNr?r>r= RuntimeErrortype)old_argnew_phs rT update_argzF_decompose_and_get_gm_with_new_signature_constants..update_argbs g/ 0N  0!v{{3 3  0!v{{3 3 !1 2#5 5  1" 4 4  &DII==++77AA ==u -udXJ?H+       ((:(:BLL#N%C  2 2 I I% ! 5YD&DS&I#  IIOO &># L+-F$ :/5H 'V&&(;'#<+-D ('!#'{E P P P P d #B!"&9:"C ##  0557 (GD!11&' d# (.335 ,GD!2==(%a););<<<r}}$77MM%%d+ ,& 55..44=8P/??d5!?I?-/__-J-J-LM'$MM!'&'!+I,5,= &&(9I '  ?( ! " 'I '/R0 '*, '0 OO . 0 <%$4$@ d  O $$& ' ' ' ' N*,T477m;STTrxx~~&r*//2K  C(8$9 99 9.0@A2&,kk1 fm2,<-@!AB GAttyy)F=DyyI///!77i222!88$';DII;%GHH   " 2  &  C(8$9 99 9.l""24DE %((--. =#+==#7HMM%  OO2GG G"5"G"GG ++- %1#$ a  % % "BMM 11g29E E P P P P P P P P `@M ' ' ' ' ' ' ' ':U=0  $ "XWJ s* y ' x>3x1Ax$ 5x "x"A+x x$ x1x>%y 5y yy y"3z6 zy5 =y( y5zzzzz!0&z&Az3A z9/ z?3{4{ x x!x$ $x. )x11x; 6x>>y y  y(y2 -y55y? :zz zz&z0rr0c|j|j5|jjD]}|jdk(st j |j\}}|D]}t|tjjs(|j|jvsA|jdk(sQ|jtjj j"j$k(s|j'|jd|jj)||j+ddd||fS#1swY||fSxYw)zQ Removes redundant copy_ node that was introduced due to mutated buffer. output call_functionr4N)_set_replace_hookget_replace_hookrrrzr\ tree_flattenrerrr r rNbuffers_to_mutaterrrrdefaultrU erase_noder)rr0r&rer4outs rTrrs  1BBD EHHNN 5Dww(" --dii8a 5C"3 6HH(;(M(MM FFo5 # eiinn.A.A.I.I I55chhqkBHH//4 5 5  " ""! " ""s+)E A EE/E?;E;AEE!c |j|j5|jD]}t|tj j s)i}i}t|jjD]}|jdk(r|jtjk(r|j\}}||d|} | |vrc|j|| |D].} | j | j j||| 0|jj#|||| <| ||<|j$||< dddy#1swYyxYw)NrL.)rMrNmodulesrrr  GraphModulerZrrrzroperatorgetitemrerUrbrRrN) rrmodule_call_graphrnode_idgetitemsr&rFidxnew_identrys rT _common_getitem_elimination_passr`'sG  o>>@ A.jjl .Ffehh&:&:;02G13HV\\//0 .77o-$++AQAQ2Q"&))KFC '0#7F)228F3CD%6"E$: % E E$((6*:!""  //5+/((. $(IIGDM! .  ....sCE/AEEold_module_call_graphctj|}i}|jjD]=}|jj dgx}s"|j ||dj <?|D][}|j}|g|j|jD]-}|j |j |j |_/]|S)N from_noder) rdeepcopyrrrrrNrbrGrH) rranew_module_call_graph provenancer&historyr_rbxs rT_get_updated_module_call_graphriDs!MM*?@"$J5iimmK4 47 4+/99Jwr{'' (5 '4OO   89##8i&7&78 4A^^AFFAFF3AF 4 4 ! rVrc Ft|||||\}}}t||j}|jj |j jt ||j} t||j||| ||j|j} | S)Nrrrr)rootrrrrange_constraintsrZrr) rIrirZrrr_get_updated_range_constraintsrmrBrrr) rrrrrrr0rrenew_range_constraintsexported_programs rT_decompose_exported_programrq\s ; #/)$?    ;  GGNN2??''(:   ' hh+//((,,  rVc eZdZdZ dDdddeej jee e ffdejjde dee eejej jffdd d eed eeee d fee e ffd eee eejeej*j,ffdeeeefdZeeddZej:eddZeeddZej:eddZeeddZej:eddZeeddZ e j:eddZ edde!ej jfdZ"edde!ee ej jffdZ#edde!ejfdZ$edde!ee ejffdZ%eeddZ&e&j:edd Z&eedd!Z'e'j:edd"Z'eedd#Z(e(j:edd$Z(eedd%Z)e)j:edd&Z)eedde fd'Z*e*j:edd(Z*eedde fd)Z+e+j:edd*Z+eedd+Z,e,j:edd,Z,eedd-Z-e-j:edd.Z-eedd/Z.e.j:edd0Z.d1Z/d2e d3e de fd4Z0d2e d3e de fd5Z1d6Z2de fd7Z3dej jfd8Z4d9Z5e6 dEd:eeejnjpe9fd;e:ddfd<Z;d=eZ=d?Z>edd@Z?e@dAZAddddB dFdCZBy)GrBa Package of a program from :func:`export`. It contains an :class:`torch.fx.Graph` that represents Tensor computation, a state_dict containing tensor values of all lifted parameters and buffers, and various metadata. You can call an ExportedProgram like the original callable traced by :func:`export` with the same calling convention. To perform transformations on the graph, use ``.module`` property to access an :class:`torch.fx.GraphModule`. You can then use `FX transformation `_ to rewrite the graph. Afterwards, you can simply use :func:`export` again to construct a correct ExportedProgram. N) verifiersrlrrrrmdict[sympy.Symbol, Any]rZr.rrsc $tjjj|_t |||_t|tjjr/|j jj|jt|j ||||_ ||_ ||_|J||_||_|xsi|_| xst$g} t'd| DsJ| |_|j+y)Nc3<K|]}t|tywrM)rrr().0rHs rT z+ExportedProgram.__init__..s>q:a*>s)rr rCodeGenr_create_graph_module_for_export _graph_modulerrWrrr`_graph_signature _state_dict_range_constraints_module_call_graph_example_inputs _constantsr(all _verifiersvalidate) rOrlrrrrmrZrrrss rT__init__zExportedProgram.__init__s //1I>>>># rVF)is_backward_compatiblec|jSrM)r{rOs rTrzExportedProgram.graph_modules!!!rVctd)Nz7Unable to set ExportedProgram's graph_module attribute.rrOvalues rTrzExportedProgram.graph_modulesTUUrVc.|jjSrM)rrrs rTrzExportedProgram.graphs  &&&rVctd)Nz0Unable to set ExportedProgram's graph attribute.rrs rTrzExportedProgram.graphsMNNrVc|jSrM)r|rs rTrzExportedProgram.graph_signatures$$$rVctd)Nz:Unable to set ExportedProgram's graph_signature attribute.rrs rTrzExportedProgram.graph_signaturesWXXrVc|jSrM)r}rs rTrzExportedProgram.state_dictsrVctd)Nz5Unable to set ExportedProgram's state_dict attribute.rrs rTrzExportedProgram.state_dictsRSSrVrpc#DK|jD] \}}| yw)zH Returns an iterator over original module's parameters. N)r)rOr4params rT parameterszExportedProgram.parameterss) --/ HAuK  c#hK|jjD]}||j|fyw)z Returns an iterator over original module parameters, yielding both the name of the parameter as well as the parameter itself. N)rrr)rO param_names rTrz ExportedProgram.named_parameterss7 ..99 :Jdooj99 9 :s02c#DK|jD] \}}| yw)zC Returns an iterator over original module buffers. N)r)rOr4bufs rTbufferszExportedProgram.bufferss) ((* FAsI rc#Kt|jj}|jjD]-}||vr||j|f||j |f/yw)z Returns an iterator over original module buffers, yielding both the name of the buffer as well as the buffer itself. N)r}rrrrr)rOr buffer_names rTrzExportedProgram.named_bufferssm "%T%9%9%P%P!Q//77 @K44!4>>+#>>>!4??;#???  @sA'A)c|jSrM)r~rs rTrmz!ExportedProgram.range_constraints&&&rVctd)Nz.s<1!A$Trying to flatten user inputs with exported input tree spec: z- but actually got inputs with tree spec of: F)rrIrr-rtrrrrr9rrrr rrrrCONSTANT_TENSOR CUSTOM_OBJr)rOrerfrIrradditional_inputsinput_s rT_graph_module_flat_inputsz)ExportedProgram._graph_module_flat_inputssf..((#'#A#A$#O =  } 7$A( Q)@ /# **66 HF{{i222##  !$$-&,,T^^FMM-JK%,,T__V]]-KL))$$!"(( )FG# H$""34!9,,rVctd)Nz\Unable to call ExportedProgram directly. You should use `exported_program.module()` instead.r)rOrerfs rT__call__zExportedProgram.__call__s B  rVc :ddlmcm}|j||\}}|jj |j j}|j j}t|t|z} |d| } |j j} | 'tt| j} |d| }|| d} tj||jj } |j j&Dcgc]!}|j(t*j,k(r|#}}t/| D]\}}|j j0|j(t2j4k(r(j6J||j8j6<dj(t2j:k(r@j6Jtfdt/|D}||j=|t?dj(|S|S#t $rDtj"|\}} |j%d|jj d| wxYwcc}w#|j j&Dcgc]!}|j(t*j,k(r|#ncc}w}}t/| D]\}}|j j0|j(t2j4k(r(j6J||j8j6<dj(t2j:k(r@j6Jtfdt/|D}||j=|t?dj(wxYw)zProcess potential mutations to the input. Because self.graph_module is functional, so mutations has to be written back after execution of graph_module. rNz@Trying to flatten user outputs with exported output tree spec: z. but actually got outputs with tree spec of: c3nK|],\}}|jjjk(r|.ywrM)rrNr)rwrRr< output_specs rTrxzDExportedProgram._postprocess_graph_module_outputs.. s3% '4#xx}} 0B0BB%s25zUnexpected kind: ) torch._export.error_exporterrorrrrJrrPrrassertion_dep_tokennextiterrr\r ExceptionrO InternalErrorrrr9rr rr;rrrrcopy_r)rOresr orig_kwargsrrr4buffer_mutationuser_input_mutation num_mutatedmutated_valuesrassertion_dep_token_indexrr<r rRrindexrs @rT!_postprocess_graph_module_outputsz1ExportedProgram._postprocess_graph_module_outputss ,+55iM 1 >> " " ."22DDO"&"6"6"L"L o.5H1IIK +.N#'"6"6"J"J ".,06I6N6N6P1Q,R)445kl#C U++C1H1HI!% 4 4 @ @yyI$8$88  !*. 9UHAu"&"6"6"C"CA"FK"'':+E+EE*11===>C (:(:;$))Z-K-KK*11=== $%+4[+A%! "%(..u5,/@AQAQ@R-STTU s ; #)#6#6s#; =))W~~../0E$o'  $ 4 4 @ @yyI$8$88  !*. 9UHAu"&"6"6"C"CA"FK"'':+E+EE*11===>C (:(:;$))Z-K-KK*11=== $%+4[+A%! "%(..u5,/@AQAQ@R-STTUs,*H&I(A I%%I--N&J-,C.Nc|jjddjdd}d|d|jd|jd}|S)NF) print_outputcolored z zExportedProgram: z Graph signature: z Range constraints: )rprint_readablereplacerrm)rOrstrings rT__str__zExportedProgram.__str__sp((778 '$ !  .! $ 4 456""&"8"8!9 =   rVcddlm}||}ddtfd}ddtfd}tj|||_tj|||_|S)z_ Returns a self contained GraphModule with all the parameters/buffers inlined. r4)&_unlift_exported_program_lifted_statesmodectd)Nz%Calling train() is not supported yet.NotImplementedErrorrOrs rT_trainz&ExportedProgram.module.._train*s%&MN NrVctd)Nz$Calling eval() is not supported yet.rrs rT_evalz%ExportedProgram.module.._eval-s%&LM MrVT)rrbooltypes MethodTypetraineval)rOrrrrs rTrzExportedProgram.module"s] D7= Ot O Nd N''7 &&uf5  rVctdt|jjDt |jjS)Nc3`K|]&\}}|jtjk(r|(ywrM)rr9r)rwrRss rTrxz=ExportedProgram._num_lifted_params_buffers..6s. Aq66Y111 s,.)rr r|rrrs rT_num_lifted_params_buffersz*ExportedProgram._num_lifted_params_buffers4sC %d&;&;&G&GH  %%11 2   rVrrc| tn t|}t|tr|j }t |\}}t |||d|S)a Run a set of decompositions on the exported program and returns a new exported program. By default we will run the Core ATen decompositions to get operators in the `Core ATen Operator Set `_. For now, we do not decompose joint graphs. Args: decomp_table: An optional argument that specifies decomp behaviour for Aten ops (1) If None, we decompose to core aten decompositions (2) If empty, we don't decompose any operator Some examples: If you don't want to decompose anything .. code-block:: python ep = torch.export.export(model, ...) ep = ep.run_decompositions(decomp_table={}) If you want to get a core aten operator set except for certain operator, you can do following: .. code-block:: python ep = torch.export.export(model, ...) decomp_table = torch.export.default_decompositions() decomp_table[your_op] = your_custom_decomp ep = ep.run_decompositions(decomp_table=decomp_table) Nrk)rErsrr0 materializerrq)rOrr _decomp_tablerrs rTrun_decompositionsz"ExportedProgram.run_decompositions>skP)5(< " $$|BT  m%6 7)557M" 9 G  + ' 3!(C   rVpassesc Ttt|}ddlm}|5||j}ddd |jn |j}|J||jur|j s|Sdt dtjjdt fd}t||j||j||jt||jt!j"|j$|j&|j(|j* }|jj,j/|jj,|jj,j/|jj,|S#1swYjxYw)Nr)r old_signaturenew_gmrpc g}t|jjD]\}}|jdk7rn|t |j ksJd|j |}t |jttfr |jn$t|j|j}|jt|j||j|j t#|jjd}|jdk(sJg}t|j$dD]\}}|t |j&ksJd|j&|} t | jttfr | jn$t| j|j}|jt)| j|| jt+||} | S)zO Update the graph signature's user_input/user_outputs. rz-Number of inputs changed after transformationrrKrz.Number of outputs changed after transformationr)r rrrzrrrrr6r7rrNrr:rrr rZrerr<r8) rr new_input_specsrRr&old_input_specr output_nodenew_output_specsold_output_spec new_signatures rT_get_updated_graph_signaturezKExportedProgram._transform_do_not_use.._get_updated_graph_signatures!O$V\\%7%78 477m+3!--CBC"/!:!:1!="&**-=?P,Q#&&2n001$))<  &&&++&--&11  0v||11226K>>X- --! $[%5%5a%89 43!..DCD#0"<"@Q-R$''3o112499= !''33S/:P:PQ  1+:JM! rV rlrrrrmrZrrrs)r3rZrrrmodifiedr8rr rWrBrrrrnrmrrdrrrrsrr)rOrpmrrtransformed_gmrtransformed_eps rT_transform_do_not_usez%ExportedProgram._transform_do_not_useso f & @ $ & (T&&'C (-0_))$BSBS))) T.. .s||K7 !/7 !HH((7 !"7 !r) &&8$$n<&&#mmD,C,CD..nnnn  ##((//0A0A0F0FG##((//0@0@0E0EFi ( (s FF'cRddlm}|jjDcgc]}|jdk(s|}}t ||j jDcgc]$\}}|jtjk(r|&}}}||||jycc}wcc}}w)Nr)"_check_input_constraints_for_graphr) torch._export.utilsrrrrzrrrrr9rrm)rOrrr1 placeholdersrinput_placeholderss rTrz(ExportedProgram._check_input_constraintssJ#'::#3#3Maqtt}7LM ML$*>*>*J*JK 1vv---   +  3T5K5K N sBB)B#c$|jyrM) _validaters rTrzExportedProgram.validates  rVct|jdkDsJd|jD]}|j|y)Nrz0ExportedProgram must have at least one verifier.)rrscheck)rOrHs rTrzExportedProgram._validatesE  ! # > = > # A CIIdO rV)rrrsc $t||j|||n |jtj|j tj|j |j||n |j|| S|j S)Nr) rBrrrrdrmrrrrs)rOrrrrrss rT_updatezExportedProgram._updates$$+%/%;z"mmD,B,BC"mmD,C,CD..#,#8idnn#,#8i  ?Cnn  rV)NN)NF)rprB)CrWrXrY__doc__r rrModulersr^rr Graphr8rrrZrCr rr r ScriptObjectrr(rpropertyr1rsetterrrrrrrrrrmrZrrrrrsrrrrrrrrrrk_ops OperatorBaserrrPassTyperrrrrr"r_rVrTrBrBs .LP %59%EHHOOT#s(^34%xx~~%. % eELL%((2D2D$DEEF % 5 % 0%!uS#XS#X'F!GH% eELL*:EHHr?r@rA__all__r rWr+ dataclassrDrCrkrContextrrtrr AutogradCPU AutogradCUDA AutogradMeta AutogradXLA AutogradLazy AutogradIPU AutogradXPU AutogradMPS AutogradHPUAutogradPrivateUse1AutogradPrivateUse2AutogradPrivateUse3rCPUCUDAMetaXLALazyIPUXPUMPSHPUrrrrsr)r*rrrEintrIrr`rZrirqrBrrnrzr_rVrTris  "$%GGB? 9?6IJ ; $$",*@B716:"  UXX))*HZ,@@ A  ' ' '  444 GG>>GG>> G  G: HH$$ HH%% HH%% HH$$ HH%% HH$$ HH$$ HH$$ HH$$ HH,, HH,, HH,, ,(, HH HH HH HH HH HH HH HH HH X$X$v  )-uzz..89)- 4 ''1 2C 78)-Xo2 //9:o2ejj55x?@ o2 sm o2d # #3G# 588  !5 56#2. .:! !0!01 //9:1ejj55x?@ 1 sm 1 "& 1h    D $$BF >.rV