[Th6SSKrSSKrSSKrSSKrSSKrSSKJr SSKJr SSK J r J r SSK r SSK r SSKJr SSKJr SSKJs Jr SSK Jr SSKJr SSKJr SS KJrJrJr SS KJ r SS K!J"r"J#r#J$r$ \RJ"\&5r'S r(\S 5r)\S\RTS\ 4Sj5r+S Sjr,Sr-\S\RTS\ 4Sj5r."SS\R^5r0\S\RTS\ 4Sj5r1\S5r2Sr3\ RhRjr5\5Rl\5Rn\5Rp\5Rr\5Rt\5Rv\5Rx\5Rz\5R|\5R~\5R\5R\5R\5RR\5RR\5R\5R\5R\5R\5R\5R\5R1rM\"\M5rM\S5rNS\ \ \R44SjrPSrQSqRSrSSrTS!SjrUg)"N)contextmanager)partial)CallableUnion)SymInt)get_decompositions) bind_symbols) aot_function aot_modulemake_boxed_compiler)strip_overloads)default_partition draw_graph#min_cut_rematerialization_partitioncURRS[RRR S9H,n[RRR UlM. UR5 U$)N call_functionoptarget) graph find_nodestorchopsaten_to_copytor recompile)fx_gnodes R/var/www/auris/envauris/lib/python3.13/site-packages/torch/_functorch/compilers.py _canonicalizer"$s[ %% 599>>#:#:&iinn''  NN Kc## [RRS5nSv [RRU5 g![RRU5 f=f7f)NF)r_C_jit_set_autocast_mode)old_jit_autocast_flags r!_disable_jit_autocastr(-sH!HH;;EB?  ''(=>''(=>s A+A A+!A((A+rreturncp[5 [U5 URRS[R R RS9Htn[UR5S:XdM[UR5S:XdM9SUR;dMK[R R RUl Mv URRH]n0nURR5H4upE[U[R 5(a UR"nXSU'M6 X2l M_ URR%5 UR'5 [R(R+U5n[R,R/UR5 [R(R1UR355n[R(R5U5n[7SU55(dU"U6 SSS5 U$!,(df  W$=f)z Compiles the :attr:`fx_g` with Torchscript compiler. .. warning:: This API is experimental and likely to change. Args: fx_g(fx.GraphModule): The input Fx graph module to be compiled. Returns: Torch scripted model. rrr dtypec3j# UH)n[U[RR5v M+ g7fN) isinstancer _subclasses FakeTensor).0ts r! ts_compile..`s&M1:a!2!2!=!=>>s13N)r(rrrrrrrlenargskwargsrrnodesitemsr.devicetypelintrjitscriptr%_jit_pass_remove_mutationfreezeevaloptimize_for_inferenceany)rinpsr new_kwargskvfs r! ts_compilerI6s  JJ))uyy~~'>'>* D499~"s4;;'71'<DKKAW#iinn//  JJ$$DJ ))+a..A !1 ,%K %   II  T " **1773 II  QVVX & II , ,Q /MMMM tH9 !: H; ! : HsA"H&1H& H&E>H&& H5cD[UR5 [XUS9 U$)N) clear_meta)printcoder)r_namerKs r!_draw_graph_compilerPes $))tj1 Kr#c0[[[US95$)NrO)r rrPrRs r!draw_graph_compilerSks w':F GGr#cU$)z Returns the :attr:`fx_g` Fx graph module as it is. This is a no-op compiler and can be used to check accuracy. .. warning:: This API is experimental and likely to change. rrNs r!noprWos  Kr#c4^\rSrSrU4SjrU4SjrSrU=r$)DebugInterpreter|cV>[UR/UQ76Ul[TU]"U6 gr-)r modulesymbol_mappingsuperrun)selfr6 __class__s r!r_DebugInterpreter.run}s%*4;;>>  Tr#c>^^ ^ ^ ^ U4Sjm U 4Sjm U 4Sjm U U 4Sjn[TT]U5nSUR;a[R"URS5upE[R"U5upg[ U5[ U5:Xd[ U5S[ U535e[ [[ U55XF5H5um p[U [R5(dM'U"XU U4Sj5 M7 U$)Nc>[U[5(dU$[R"URR R TR55nUR(dU5e[U5$r-) r.rsympyexpandr exprxreplacer] is_numberint)nirr`s r! subst_symint/DebugInterpreter.run_node..subst_symintsVb&))  RWW\\2243F3FGHA;; ! !;q6Mr#c.>[U4SjU55$)Nc34># UH nT"U5v M g7fr-rU)r1rkrms r!r3HDebugInterpreter.run_node..subst_symint_tuple..s8Cbb))Cs)tuple)nisrms r!subst_symint_tuple5DebugInterpreter.run_node..subst_symint_tuples8C88 8r#c>T"UR55S:ae[UR5HLnT"URU55URU5:wdM/T"UR U55S:dML g g)Nrr FT)numelrangendimstridesize)abidxrms r!check_significant_strides.check_significant_stridessbAGGI&* =C$QXXc]3qxx}D(59$ ) r#c 8>[U5(deURUR:Xd(U"5SURSUR35eT"UR55UR5:XdGU"5SUR5ST"UR55SUR535eT"X5nU(dGU"5SUR5ST"UR55SUR535eg)Nz:  != z aka )callabler+r{rz)nvrvdesc same_stridesrrts r!check(DebugInterpreter.run_node..checksD>> !>88rxx' NDF82bhhZtBHH:)N N'"2779-: [&BGGI;e,>rwwy,I+J$rwwykZ [:4RSTSTR3$)Nzoutput z where r])ir`sr!+DebugInterpreter.run_node..ss'$:M:M9N&Or#) r^run_nodemetapytree tree_flattenr5ziprxr.rTensor)r`nrrln_vals_n_specr_vals_r_specrrrrrmrtras` @@@@r!rDebugInterpreter.run_nodes  9  a G Q  AFF?$11!&&-@OF$11!4OFv;#f+- P#f+d3v;-/P P- s6{!3VD 2!"ell33bOPEr#r)__name__ __module__ __qualname____firstlineno__r_r__static_attributes__ __classcell__)ras@r!rYrY|s//r#rYc,[U5R$)z Returns a (slow) interpreter over the FX graph module that also checks various debugging properties (e.g., that tracing strides matched real strides.) )rYr_rVs r! debug_noprs D ! % %%r#c[U5 [RRU5n[RR UR 55nU$r-)rrr=r>r@rA)rrNrHs r!simple_ts_compilers=D A "A Hr#c"[U[5$r-)r r)rHs r!nnc_jitrs , --r#c0[UR5 U$r-)rLrMrVs r! print_compilers $)) Kr#fnc [[[[S.nURU5 [ U[ R R5(a [U40UD6$[U40UD6$)a; Wrapper function over :func:`aot_function` and :func:`aot_module` to perform memory efficient fusion. It uses the :func:`min_cut_rematerialization_partition` partitioner to perform efficient recomputation. It uses NVFuser to compile the generated forward and backward graphs. .. warning:: This API is experimental and likely to change. Args: fn (Union[Callable, nn.Module]): A Python function or a ``nn.Module`` that takes one ore more arguments. Must return one or more Tensors. **kwargs: Any other overrides you want to make to the settings Returns: Returns a ``Callable`` or ``nn.Module`` that retains the eager behavior of the original :attr:`fn`, but whose forward and backward graphs have gone through recomputation optimizations, and the graphs have been compiled with nvfuser.  fw_compiler bw_compiler partition_fndecompositions) rIrdefault_decompositionsupdater.rnnModuler r )rr7configs r!memory_efficient_fusionrsZ6"!;0 F  MM&"ehhoo&&"'''B)&))r#cURS5 [SUVs/sHo"RUR4PM snS35 SSKJn U"5R 5"U6 [X5$s snf)NfooaQ ############################################################## # To minimize FX graph, copy and paste the below and run it # ############################################################## import torch import torch.fx as fx from functorch.compile import minifier, check_nvfuser_subprocess, check_nvfuser_correctness_subprocess inps = a? inps = [torch.ones(shape, dtype=dtype, device='cuda') for (shape, dtype) in inps] from foo import FxModule mod = FxModule().cuda() with torch.jit.fuser("fuser2"): # check_nvfuser_subprocess can be replaced with check_nvfuser_correctness_subprocess minifier(fx.symbolic_trace(mod), inps, check_nvfuser_subprocess) r)FxModule) to_folderrLshaper+rrcudarI)rrDrrs r! debug_compilersmNN5  &**T''177 T*+, ( JOOt d !! +s!A0 c P/n[US5n[R"U5n/nUHn[U5S:XaUnU"[R "55nOUupWpn U [ R[ R[ R[ R[ R[ R[[1;a[ R"SSXyU S9nO[ R "XyU S9nURU5 M SSS5 U$!,(df  U$=f)zR Return a random input for the given inputs meta generated from _save_fx_default. rbr r)r+r:N)openpickleloadr5randomrandrrjint32int64booluint8floatrandintappend) input_data_pathinputsrH inputs_metarr;inputr_strider+r:s r! get_inputsr3sF ot $kk!n D4yA~V[[]+6:3WVIIKKKKJJIIKK  "MM!Q6RE!JJu&IE MM% '  %. M/ % $. Ms C=D D%c x^^^^ ^ SSKJn U 4Sjm UUUU 4Sjm U 4SjnU 4SjnU 4SjnU"UUUUU[S9$) a The forward, backward, and joint computation graph will be stored in {folder_name}/{current_name}/{current_name}_forward_{graph_index}, {folder_name}/{current_name}/{current_name}_backward_{graph_index}, and {folder_name}/{current_name}/{current_name}_joint_{graph_index} respectively. The input shape of the graphs will be stored in the .input files. These files can be loaded with pickle, and is a list of format (type, shape, stride, dtype, device). In the case of type = int or float, it is just (type,). For joint graph input, it is a nested list [[],[]] where the two inner lists have the same format. If dump_example_input is True, example_inputs will be stored in .pt file. Since each function might produce multiple graphs, the graph_index is used to distinguish difference graphs r)aot_module_simplifiedc>/n[U5S:a6[US[5(aUT"US5- nUT"US5- nU$UHn[U5[:Xd[U5[ :XaUR [U545 MFUR [U5URUR5URUR45 M U$)Nrr ) r5r.rrr;rjrrrrzr+r:)r6 input_metaargget_input_metas r!r(_save_fx_default..get_input_metads t9q=ZQ77 .a1 1J .a1 1J CCyC49#5!!49,/!!#Y 3::<CJJO r#c>[URR5S:Xa,[R[R STU[ 5 g[R"U5nURR[RRR55 UR5 T"U5n[R"TST3SS9 UR!TSTSTSUS[ 3 5 ["R$"U['TSTSTSUS[ STSUS[ S3S55 T(a8[R("UTSTSTSUS[ STSUS[ S 35 gg) Nrz!No nodes in graph {%s}_{%s}_{%s}./T)exist_okrNz.inputwbz.pt)r5rr8logloggingWARNING graph_indexcopydeepcopy set_codegenrfxCodeGenrosmakedirsrrdumprsave) gm_to_saver6 type_namegmr current_namedump_example_input folder_namers r!graph_saver_helper,_save_fx_default..graph_saver_helperss z%% &! + GG3    ]]: & UXX^^3356 #D)  {m1\N3dC m1\N!L>9+Q{m T    -qa ~Qyk;-WXYeXffghqgrrsttAAGH    JJ-qa ~Qyk;-WXYeXffghqgrrsttAADE  r#c>T"XS5 U$)NforwardrU)rfw_argsrs r!graph_saver_forward-_save_fx_default..graph_saver_forwards2 2 r#c,>T"XS5 [S- qU$)Nbackwardr )r)rbw_argsrs r!graph_saver_backward._save_fx_default..graph_saver_backwards2 3q  r#c,>T"XS5 [X5$)Njoint)r)r joint_argsrs r!graph_saver_joint+_save_fx_default..graph_saver_joints273 00r#r)functorch.compilerr) rrrrexample_inputsrrrrrrs ``` @@r!_save_fx_defaultrRsH 8 !!F 1 ! '(&-  r#c(Sq[[XU5$)aK Dump the forward, backward, and joint computation graph. Example Usage: save_fx_func = graph_dumper_aot(current_name, folder_name, dump_example_input = False) optimize_ctx = torchdynamo.optimize( save_fx_func ) with torch.enable_grad(): with optimize_ctx: result = forward_and_backward_pass(model, example_inputs) r)rrr)rrrs r!graph_dumper_aotrsK #\@R SSr#)T)F)Vrrrrr contextlibr functoolsrtypingrrrertorch.fxrtorch.nnrtorch.utils._pytreeutils_pytreerr torch._decompr%torch.fx.experimental.symbolic_shapesr aot_autogradr r r compile_utilsr partitionersrrr getLoggerrrr"r( GraphModulerIrPrSrW InterpreterrYrrrrrdetach gelu_backwardleaky_relu_backwardsigmoid_backwardthreshold_backwardhardtanh_backwardhardsigmoid_backwardhardswish_backward tanh_backward silu_backward elu_backwardcudnn_batch_normcudnn_batch_norm_backward masked_fillScalarrelu leaky_reluhardtanh hardswish hardsigmoid conj_physical is_same_sizerrrrrrrrrrUr#r!r,s[  %" $$,>GG*! ??+ R^^+ h+ + \ H bnn H  4r~~4n&BNN&(&&  . yy~~KK""HHOOMMNN-2,,BC $* h !"$*N": >YzTr#