g^(h UdZddlmZddlZddlZddlZddlZddlZddlZddl Z ddl Z ddl Z ddl Z ddl Z ddlZddlZddlZddlmZddlmZddlmZmZddlmZmZmZmZmZmZddlm Z ddl!Z!ddl"Z"ddl#Z"ddl$m%cm&Z'ddl(Z"dd l"m)Z)dd l*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0dd l1m2Z2dd l3m4Z4m5Z5m6Z6dd l7m8Z8ddl9m:Z:m;Z;ddlZ>m?Z?ddl#m@Z@ddlAmBZBmCZCddlDmEZEddlFmGZGmHZHmIZImJZJddlKmLZLmMZMddlNmOZOmPZPmQZQmRZRm%Z%ddlSmTZTmUZUddlVmVZVddlWmXZXmYZYmZZZm[Z[m\Z\ddl]m^Z^ddl_m`Z`ddl%maZambZberddlcmdZdddlmeZemfZfejehZie%jZkejZmGdd eZniaod!epd"<enjZrdcd#ZseGd$d%Ztetauddd&Zvd'ev_wdaxd(epd)<ded*Zyd+Zzdfd,Z{d-Z|d.Z}d/Z~d0Zeje@eeeed1hZ dgd2ZGd3d4e"jj Zd5Zd6Zd7Zd8Zdhd9ZGd:d;e jZeZd<Zej d=ZGd>d?ZGd@dAeZGdBdCeZGdDdEeZemdFdddf didGZdHZGdIdJejZdKZdLZdMZdNZdOZdPZ djdFdddFddQ dkdRZe dSd'dTZGdUdVe"jBjDZGdWdXeZdYZdZZdFdFdd[dd'dFd'dFdFdFd'd\ dld]Ze^dddFdddd^d_ZGd`daZdmdbZy)na This module implements the core frame evaluation handler for TorchDynamo's compilation system. The eval frame handler intercepts Python bytecode execution at runtime to enable dynamic compilation and optimization of PyTorch code. Key components defined here: - Frame evaluation handlers that intercept and analyze Python execution frames - Guards management for tracking dependencies and invalidating compiled code - Optimization contexts and decorators (optimize, run_once, disable, etc.) - Export functionality for saving optimized graphs - Backend compiler integrations and callback management Functions in this file are responsible for modifying the eval frame handler at RUNTIME. Therefore, all functions in this file are hot and performance-critical. Functions that only execute at compile time should be placed in torch._dynamo.convert_frame. The eval frame handler is the core mechanism that enables TorchDynamo to dynamically intercept, analyze and optimize PyTorch code during execution. It works by registering a custom frame evaluation function that gets called for every Python frame, allowing us to detect PyTorch operations and trigger compilation as needed. ) annotationsN) dataclass)Enum)dirnamejoin)AnyCallable NamedTupleOptional TYPE_CHECKINGUnion)patch)_guards) reset_codeset_code_exec_strategyset_eval_frameset_guard_error_hookset_skip_guard_eval_unsafe unsupported)enable_python_dispatcher)ConvertFrameReturn FrameActionFrameExecStrategy)unset_fake_temporarily)justknobs_checklog_export_usage) _combine_args_process_dynamic_shapes_RelaxedConstraint) GraphModule)clone_and_convert_to_metatrack_dynamism_across_examples)make_fx)ConstraintViolationError DimDynamicShapeEnvStatelessSymbolicContext)_PyTreeCodeGen _PyTreeInfo)config convert_frameexternal_utils trace_rulesutils) CompilerFnlookup_backend) code_context)CondOpArgsMismatchErrorShortenTraceback Unsupported UserError UserErrorType)Hooksinstall_generation_tagging_init)common_constant_types compile_times) fake_tensor) CacheEntryDynamoCallbackceZdZdZy)UnsetrN)__name__ __module__ __qualname__tokenV/var/www/skyplay_api_hub/venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.pyrArAus ErGrAzdict[int, CompilerFn]cached_backendsctds+tjjj d|St |S)Nz/pytorch/compiler:enable_compiler_set_eval_framezTDynamo disabled by Justknob: enable_compiler_set_eval_frame, skipping set_eval_frame)rtorch_dynamor/ warn_oncercallbacks rH_maybe_set_eval_framerP~s: L M %% b h''rGc8eZdZUdZded<dZded<dZded <y) DynamoStancedefaultstrstanceFboolskip_guard_eval_unsafeNz$Union[str, Callable[..., Any], None]backend)rBrCrDrU__annotations__rWrXrFrGrHrRrRs"FC#(D(48G 18rGrRcRddlm}|}|dur | tdt}|a|S)Nr)get_eval_frame_callbackFz1attempted to set_stance in a torch.compile region)torch._C._dynamo.eval_framer[ RuntimeError_stance)rUr[rOpriors rH _set_stancer`s8D&(Hu!5NOO EG LrGTzOptional[dict[str, list[Any]]]_EXAMPLE_INPUTScFtia|tvr gt|<t|SN)ra)keys rHget_example_inputsres* /!! 3 rGc$tjdk(r8tj&|dvr"tt tj}|Stjdk(r |dvrt |tjS|Stjdk(r |dvrt |tjS|Stjdk(rytjdk(rytjdk(r|dvr|Sd }||_|Std td ) NrS)FNeager_then_compileaot_eager_then_compile force_eagereager_on_recompileFfail_on_recompilectd)NzCDetected recompile when torch.compile stance is 'fail_on_recompile')r])argskwargss rH fail_callbackz,_callback_from_stance..fail_callbacksU rGzinvalid torch.compile stance '')r^rUrX_create_wrapped_callbackget_compiler_fn _create_delayed_compile_callback_torchdynamo_orig_callabler])rOros rH_callback_from_stancerus~~" ?? &8=+H/0PQH / / = (3HgnnM M 3 3 = (3HgnnM M = ( / / . . } $O  4< 0;G9AFGGrGcjt}tjtj|||Src)r8r,catch_errors_wrapper) compiler_fndynamismhookss rHrqrqs5 GE  - -##      rGc|jjt|jjz}t |}t |dkr$|j t|j|S)N) f_code co_filenamerTco_firstlinenorelenappendr!f_locals)framerdexample_inputss rH_get_or_add_example_inputsrsZ ,, " "S)D)D%E EC',N >Q7GH rGcfd}|S)Nc|d}t|}t|dk(rXdk(r2tttj tj Sdk(rt d}t||i|St|}|tj|jd<jj}t|||i|S)Nrr*rg)frame_exec_strategyrh aot_eagerry) rrrrrDEFAULTrrrqr"r2 get_contextr}rt) rmrnrr aot_eager_fnryrxrOrUs rH callback_fnz5_create_delayed_compile_callback..callback_fnsQ3E: ~ ! #--)(9#++[-@-@) 33.{; =/ =tNvNN1.A=E   .z:99TT >' X>OOOrGrF)rOrUrs`` rHrsrssP( rGc"tjSrc)r^rWrFrGrH!_is_skip_guard_eval_unsafe_stancers  ) ))rGctjD]}t|ds|j!tj y)Nreset)rIvalueshasattrrclear)rXs rH_reset_guarded_backend_cacher s;"))+ 7G $ MMOrGz+onnx/_internal/fx/dynamo_graph_extractor.pyct|r |j}tjjj j |S)zh Given a code object or a callable object, retrieve the cache entries stored in this code. )callable__code__rK_CrL eval_frame_debug_get_cache_entry_listcodes rHrrs5~}} 88   & & B B4 HHrGceZdZUdZded<ded<hdZdfd ZdZd Zd Z d Z e d Z e jfd Z dZdfd ZdZfdZxZS)OptimizedModulezx Wraps the original nn.Module object and later patches its forward method to optimized self.forward method. Callable[..., Any]rtzCallable[[], Any]get_compiler_config>forward__dict___forward _orig_mod dynamo_ctxrnamed_children_walkrtct|||_||_|j |jj |_yrc)super__init__rr _initializetraining)selfmodr __class__s rHrzOptimizedModule.__init__9s9 $ // rGct|jtr+|j|jj|_nt|jj t jrttj|jj st|jddr4|jtj|j|_n*|j|jj|_t|jdr#|j |_|j|_yy)N_is_fsdp_managed_moduleF_initialize_hook) isinstancerDisableContextr__call__rtypes MethodTyper.checkgetattrr- wrap_inlinerr_call_lazy_checkrs rHrzOptimizedModule._initializeAs doo~ 6??4>>+B+BCDL ..0@0@ A   dnn44 5t~~'@%H  ??>+E+Ednn+UVDL ??4>>+B+BCDL 4>>#5 6 LLDM00DL 7rGcJ|j|j|jffSrc)rrrrs rH __reduce__zOptimizedModule.__reduce__Vs ABBrGcxt|j}|jdd|jdd|S)Nrr)dictrpoprstates rH __getstate__zOptimizedModule.__getstate__Ys1T]]# )T" *d# rGc2||_|jyrc)rrrs rH __setstate__zOptimizedModule.__setstate___s  rGc.|jjSrc)rrrs rHrzOptimizedModule.trainingcs~~&&&rGcf t|d||j_y#t$rYywxYwNr)r __getattr__rrAttributeError)rvaluers rHrzOptimizedModule.traininggs3  G  ,&+DNN #   s $ 00cV|dk(r|jdSt|j|Sr)_modulesrr)rnames rHrzOptimizedModule.__getattr__ps* ; ==- -t~~t,,rGctt||rt| ||S|tj vrt| ||St |j||Src)rtyper __setattr__r_opt_mod_attributessetattrr)rrvalrs rHrzOptimizedModule.__setattr__usV 4:t $7&tS1 1 ?66 67&tS1 1t~~tS11rGc t|jdr\t|jdrFt|jjr'|jj|j|||j|i|S)Nr_infer_parameters)rrrrr)rrmrns rHrz OptimizedModule._call_lazy_check~sd DNN$6 7(;<99: NN , ,T^^T6 Jt}}d-f--rGc|jj}|t|Dcgc] }||vs| c}zScc}wrc)r__dir__r)rorig_mod_attrsattrrs rHrzOptimizedModule.__dir__sH//1"W_.! $n2LD!   ! s AA)rztorch.nn.ModulereturnNonerr)rBrCrD__doc__rYrrrrrrpropertyrsetterrrrr __classcell__rs@rHrr%s} !32** 01*C ''__- 2 .  rGrcHt|tjr t|yt |drt|j yt t |dddr t|jj yddlm }|tjd|y)zA Make sure f.__code__ is not cached to force a recompile rrNr*rz#could not determine __code__ for %s) rrCodeTyperrrrrrlogwarning)frs rHremove_from_cachersk!U^^$1 J 1:: It,j 9199%%&  91=rGcyrcrFrFrGrHnothingrsrGcyNFrFrFrGrH always_falsers rGcn|}t|dr&|j}t|sJt|dr&|S)z In case of nesting of _TorchDynamoContext calls, find the innermost function. TorchDynamo caches on fn.__code__ object, so its necessary to find the innermost function to pass on the optimize, run, disable etc. rt)rrtr)fn unaltered_fns rH innermost_fnrsB L , < =#>>  %%% , < = rGct|tsJ|rtjdStjddS)NF)assume_static_by_defaultT)automatic_dynamic_shapesr)rrVr+_make_closure_patcher)enables rHmake_set_enable_dynamicrs> fd ## # ++UKK++%*T  rGceZdZUgZded<y) DynamoTLS list[str]traced_frame_infosN)rBrCrDrrYrFrGrHrrs%' &rGrc@tjjyrc) dynamo_tlsrrrFrGrHclear_dynamo_tlsrs!!'')rGcdjtj}tj|d}d|d}t j |y)z At program exit, log all of the frames Dynamo has attempted to trace from, excluding the continuation frames generated by Dynamo.  z * z7TorchDynamo attempted to trace the following frames: [ z ]N)rrrtextwrapindentrinfo)msgs rH_log_traced_framesrsA ))J11 2C //#v &C DSE MCHHSMrGcLeZdZeeedfdddd dfdZdZdZdZxZ S) _TorchDynamoContextFNexportdynamiccompiler_configc6t |t|s|dus|J||_|_t |_||_||_||_ ||_ g|_ g|_ |t|} tjt!| | |$|jj#t%|t&ur fd} |jj#| t(j*ur!fd} |jj#| yy)NFctSrc)r)on_entersrH call_on_enterz3_TorchDynamoContext.__init__..call_on_enters  rGcv}|jtj|jdddSrc) __enter__ functoolspartial__exit__)ctxbackend_ctx_ctors rHcall_backend_ctxz6_TorchDynamoContext.__init__..call_backend_ctx s/&(  ((tT4HHrG)rrrrO_backend_ctx_ctorunsetr_ first_ctxr _dynamicr cleanup_fnsenter_exit_hooksrrI setdefaultidrrr contextlib nullcontext) rrOrrpatch_fnrr r rrXrrrs `` rHrz_TorchDynamoContext.__init__s !X%6(:JJJ(0 !138 "  .46 " x(""2g;8    ! ! ( ()@)I J 7 "   ! ! ( ( 7 :#9#9 9 I  ! ! ( ()9 : :rGctjr tdtd|_|j Dcgc] }| c}|_tt|_ tt|jycc}w)Nztorch._dynamo.optimize(...) is used with a context manager. Please refer to https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html to use torch._dynamo.optimize(...) as an annotation/decorator. ) r+raise_on_ctx_manager_usager]rr_r rrrprior_skip_guard_eval_unsaferPrurO)renters rHrz_TorchDynamoContext.__enter__st  , ,R  $D) 151F1FGEGG,F - /- ) 3DMMBC HsBc"|jtusJtdt|j|j D] }| |j j tt|jt|_yrc) r_rrrr(rrrPru)rexc_typeexc_valexc_tbcleanups rHrz_TorchDynamoContext.__exit__sozz&&&t"4#D#DE'' G I   3DJJ?@ rGc  fd}tttr?tjt j jjd<ttjjr6}t|}|j|_ t|drJ||_|St!j"rV}|j$|_t'|tjjr|j(|_|St+sJ t!j,}|t1j2r,t5dddvr|t6vrt9j:d}tdr j< n| tj>j@ tjBjDjF tIjJ  fd }|_&|_ t|drJ||_ d vr=td stOtQjRd d tTj<|S#t.$rd}Y4wxYw)NcjSrc)rrsrHrz9_TorchDynamoContext.__call__..get_compiler_config*s'' 'rGorig_graphmodulerrBr) _call_impl_wrapped_call_impl _lazy_forwardcyrcrF)argrns rH do_nothingz0_TorchDynamoContext.__call__..do_nothing`s rGrOctd}  r.tjr td |i|t |S r td j Dcgc] }| }}t t}tjjj}t t   |i|tdtjjj|t ||D] }|  t |Scc}w#t$r(}tjr|j!ddd}~wt"$r}|j%dd}~wwxYw#tdtjjj|t ||D] }| wxYw#t |wxYw)NzvDetected that you are using FX to symbolically trace a dynamo-optimized function. This is not supported at the moment.zsDetected that you are using FX to torch.jit.trace a dynamo-optimized function. This is not supported at the moment.)rr+error_on_nested_fx_tracer]rPr rrrKr _functorchget_dynamic_layer_stack_depthru)pop_dynamic_layer_stack_and_undo_to_depthr5verbosewith_tracebackr4remove_dynamo_frames)rmrnr_r)cleanupsr(saved_dynamic_layer_stack_depthr.erOr is_fx_tracingis_jit_tracingrs rH_fnz)_TorchDynamoContext.__call__.._fnks"4(E5 - ?66*` "4262Z&e,W"#&\ 261F1FGEGGG/I570,HH''EEG0&&;H&EF"t.v.#4(HH''QQ7//KL#+" "&e,KH$#;~~**40d:'=002<= #4(HH''QQ7//KL#+" "&e,sb)G  G #D$1AG D) A G $G ) E72#E E7!E22E77E::AG  G Grra torch._dynamo.optimize is called on a non function object. If this is a callable class, please wrap the relevant code into a function and optimize the wrapper function. >> class CallableClass: >> def __init__(self) -> None: >> super().__init__() >> self.relu = torch.nn.ReLU() >> >> def __call__(self, x): >> return self.relu(torch.sin(x)) >> >> def print_hello(self): >> print("Hello world") >> >> mod = CallableClass() If you want to optimize the __call__ function and other code, wrap that up in a function >> def wrapper_fn(x): >> y = mod(x) >> return y.sum() and then optimize the wrapper_fn >> opt_wrapper_fn = torch._dynamo.optimize(wrapper_fn) T)+rrr weakrefrefr2rrrrKnnModulerrtrrinspectisclassr issubclassr2r getsourcefile TypeErrorr.rrDONT_WRAP_FILESr-rrOr _is_tracingfx_symbolic_tracerCrwraps_torchdynamo_inliner]rdedentalways_optimize_code_objects) rrrrnew_modcls_objfilenamer7rErOrCrDs `` @@@rHrz_TorchDynamoContext.__call__(sE ("  b+ & B  $ $RZZ%8%8 9:L M b%((// *C%c40G25G .w(=> >>*=G 'N ??2 G#G$4$45G '588??3%)'*<*<%="N|| ,,R0H !2!22!6J+LM/ ++B/B  4 $}}H!H--00>>  7 - 7 -t#%*,&3 5666"5 = (2z*"OO  B9= ( 5 K H s8I33 JJ)rOr?rr) rBrCrDr null_contextrrrrrrs@rHr r sJ% /;/; /; /;b DprGr c<eZdZ dddddd dfdZdZxZS)OptimizeContextFN)r r r rebuild_ctxc d}t |||tj||||tj r"fd} j j| yy)Nctyrcr9rFrGrHrz*OptimizeContext.__init__..on_enters + -rG)rOrrr%rr r rcJ}tjjj|j}|j t j|jdddS)N)r ) rKrLcompiled_autograd_enablerrrrr)rxrr]rs rHcall_compiled_autogradz8OptimizeContext.__init__..call_compiled_autogradsd"...)m mm55==>  ((tT4HHrG)rr TorchPatcherrr+rar r) rrOrrr r rr]rrcrs ` ` rHrzOptimizeContext.__init__sg . -!''+    # # I  ! ! ( ()? @ $rGc|j|j|j|jf|j|j |j dfS)Nr )rrOrrr rrrs rHrzOptimizeContext.__reduce__sG NN ]]D22DNN C++==#'#7#7   rG)F)r]z>Optional[Callable[[], Union[OptimizeContext, _NullDecorator]]]rrrBrCrDrrrrs@rHr\r\s@  &A  &A &A &AP  rGr\c&eZdZdfd ZdZxZS)RunOnlyContextc.d}t|d|y)Ncrtjjjxjdz c_y)Nr*)rKrLmutation_guardGenerationTracker generationrFrGrHrz)RunOnlyContext.__init__..on_enters# MM ( ( : : E E J ErGF)rOrrr)rrrs rHrzRunOnlyContext.__init__s K %(;rGc|jdfSNrFrrs rHrzRunOnlyContext.__reduce__##rGrrfrs@rHrhrhs <$rGrhc,eZdZdfd ZdZdZxZS)rc&t|dy)NrNrn)rrs rHrzDisableContext.__init__s $'rGc tttjjr!}t |}|j |_|Stjrm}|j|_ |j|_ t|tjjr|j|_ |StsJtj fd}d|_|_|S)Nc(td} tt}tt j  |i|tdt|t|S#tdt|wxYw#t|wxYwrc)rrrrPrurO)rmrnr_r(rrs rHrEz$DisableContext.__call__.._fn=s"4(E -/I570,&&;DMM&JKMt.v."4(./KL%e,#4(./KL%e,s#1BA)B)BB BT)rrrKrHrIrrrtrJrKrrrLr2rrrS_torchdynamo_disable)rrrrWrXrEs`` rHrzDisableContext.__call__ s "  b%((// *C%c40G14G .N ??2 G $G$4$45G #G$4$45G '588??3%)'*<*<%="N||   -  -$( *,& rGc|jdfSrprrs rHrzDisableContext.__reduce__UrqrGr)rBrCrDrrrrrs@rHrrs(3j$rGrFc Nttj|||d||||S)NT)rrr r rr])r\r,rw) compile_fnrzrr r rr]s rH_optimize_catch_errorsrzYs3 **:u=)' rGcddlm}t|dr |j}nt |t r|}nd}t |}|||S)Nr*)wrap_backend_debug compiler_name)repro.after_dynamor|rr}rrTr1)rxr| compiler_strs rHrrrrmsE6{O,"00 K %"   -K k< 88rGceZdZdZy)_NullDecoratorc t|sJ|Src)r)rrs rHrz_NullDecorator.__call__{s|| rGN)rBrCrDrrFrGrHrrzsrGrctjdk\r tdtjddk(r tdy)N)z0Python 3.14+ not yet supported for torch.compilePy_GIL_DISABLEDr*z@torch.compile is not supported on Python built with GIL disabled)sys version_infor] sysconfigget_config_varrFrGrHcheck_if_dynamo_supportedrsE 7"MNN  ! !"3 4 9 N   :rGc8 ty#t$rYywxYwNTF)r ExceptionrFrGrHis_dynamo_supportedrs"!#  ctyrc)rrFrGrHcheck_if_inductor_supportedrsrGc8 ty#t$rYywxYwr)rrrFrGrHis_inductor_supportedrs"#% rcTtjrtjrJdyy)NzfDynamo configs suppress_error and fail_on_recompile_limit_hit can not both be active at the same time.)r+suppress_errorsfail_on_recompile_limit_hitrFrGrHcheck_for_incompatible_configsrs)&&6+M+Mp N+M&rGc0fd}t|giS)Nctj}|r=t|jdhk(sJd|j|dd<t iS)N fullgraphz>Only `fullgraph` kwarg override is supported for now, but got nopython)r+!compiled_autograd_kwargs_overridesetkeysoptimize)ca_kwargs_overridermrns rHr]zoptimize..rebuild_ctxsp#EE )..01k]B PQcQhQhQjPkl B"4K!@F: (((rG) _optimize)rmrnr]s`` rHrrs ) [ 24 26 22rG)rguard_export_fn guard_fail_fndisabler c ttt||}tjj d|s.t jjdddk(s tds tSt|}t|dt}|rt||||Stt!j || |||t#|d r|j%| Sd | S) a The main entrypoint of TorchDynamo. Do graph capture and call backend() to optimize extracted graphs. Args: backend: One of the two things: - Either, a function/callable taking a torch.fx.GraphModule and example_inputs and returning a python callable that runs the graph faster. One can also provide additional context for the backend, like torch.jit.fuser("fuser2"), by setting the backend_ctx_ctor attribute. See AOTAutogradMemoryEfficientFusionWithContext for the usage. - Or, a string backend name in `torch._dynamo.list_backends()` nopython: If True, graph breaks will be errors and there will be a single whole-program graph. disable: If True, turn this decorator into a no-op dynamic: If True, upfront compile as dynamic a kernel as possible. If False, disable all dynamic shapes support (always specialize). If None, automatically detect when sizes vary and generate dynamic kernels upon recompile. Example Usage:: @torch._dynamo.optimize() def toy_example(a, b): ... rrztorch._dynamo.optimizeTORCHDYNAMO_DISABLEr1zpytorch/compiler:enable_dynamor)r rzr])rzrN)r rr])rrr8rKr_log_api_usage_onceosenvirongetrrrrrrZoptimize_assertrzr,rr) r]rXrrrrr rzrs rHrrsF"$ / OE HH  !9: ::>>/ 4 ; @Ag&Gw(:LI #    "##G59 w 56  ' ' )     rGz&torch._dynamo.symbolic_convert.explaincbfd}|s|r$tjdtd||i|S|S)Nc  ddlm}|g g d g g d fd }fd}t|d|}||i|t }|dz }t d }|dd lm} | || |S) Nr*rrcLddlm}||\}|jS)Nr*)_explain_graph_detail)backends.debuggingrr)gmrr break_reasonsgraphsop_count ops_per_graphs rH"dynamo_graph_accumulating_compilerzBexplain..inner..dynamo_graph_accumulating_compilers: BBWFHm]B >B-:: rGc(j|yrc)extendguards out_guardss rHguard_export_printz2explain..inner..guard_export_print"s   f %rGF)rrrT)repr) ExplainOutputrtorch.fx.GraphModule)rrrrr<rr)rmrnrrropt_f graph_countgraph_break_count compile_timerrrrrrrs @@@@@rHinnerzexplain..inners -/#% -/ *,  $    &  ..    tv&k '!O$%0  5          rGzexplain(f, *args, **kwargs) is deprecated, use explain(f)(*args, **kwargs) instead. If you don't migrate, we may break your explain call in the future if your user defined kwargs conflict with future kwargs added to explain(f).r| stacklevel)warningswarn FutureWarning)r extra_args extra_kwargsrs` rHexplainrs@9 v\  ?   j1L11 rGcjeZdZ d dfd ZdZfdZfdZfdZxZS)FlattenInputOutputSignaturec Xt|t|t|k(sJt|D cic] \} } | ||  } } } g_t dt|D]} td| di} | | vr| | | jjd<n|t|| tjr|j|| tt t|| jDcgc])}||| vrtj ntj"+c}dgt|| jz| jjd<jj%| fd|D_|_|_ycc} } wcc}w)Nrr6rFr) dynamic_sizesconstraint_sizes)symbolic_contextc3<K|]}j|ywrc)new_args).0irs rH z7FlattenInputOutputSignature.__init__..{sX!T]]1-Xs)rrr enumeraterrange placeholdernodemetarrKTensor from_tensorr'shaper%DYNAMICSTATICr old_args_gen!matched_output_elements_positions flat_results)rm flat_args matched_input_elements_positionsrrexample_fake_inputsflat_args_dynamic_dims fake_modeixrmatched_input_elements_to_fakerr6drs` rHrz$FlattenInputOutputSignature.__init__Os )*c)n<<<%%EF* C $R( (* &*  q#i.) &A'%A3iR8C22'Ea'H e$(Z ! ell-S+4+@+@!! )A*/s9Q<3E3E/F)G + %&(),B1,E'E%/$6$6)3):):%;+/3Vc)A,:L:L6M-M *,A ,CHHMM%( MM  %/ &0Y7WX1R.(A* $+s F!7.F'cpt|j}d|jjvr0|jjd|jjd<d|jjvr0|jjd|jjd<d|jjvr0|jjd|jjd<d|jjvr0|jjd|jjd<|S)Nr tensor_dict example_valueunbacked_bindings)nextr current_noderr)rtargetrmrnr6s rHrz'FlattenInputOutputSignature.placeholders4$$% D%%** *#'#4#4#9#9%#@CHHMM% D--22 2+/+<+<+A+A-+PCHHMM- ( d//44 4-1->->-C-CO-TCHHMM/ * $"3"3"8"8 8151B1B1G1G#2CHHMM- . rGct|d}g||j}g}tt|jD]n}|j|"|j ||j|4|j|}t |ttsJ|j |pt |)||fiS)Nr) rrrrrrrtupler;routput) rrrmrndynamo_result_flatlookupnew_results_flatr const_valrs rHrz"FlattenInputOutputSignature.outputs!!W6%6 6s4,,-. 3A55a8D ''4AA!DE!--a0 !)U3H-IJJJ '' 2 3w~f'7&92>>rGc||_t| |}d|jjvr0|jjd|jjd<d|jjvr0|jjd|jjd<d|jjvr0|jjd|jjd<|jj dk7rD|jj t|jd|jj|S)Nrrrrr) rrrun_noderrop_renamerr)rn result_proxyrs rHrz$FlattenInputOutputSignature.run_nodes$w'* D%%** *,0,=,=,B,B5,IL   " "5 ) d//44 46:6G6G6L6L7L   " "? 3 $"3"3"8"8 8:>:K:K:P:P#;L   " "#6 7     8 +    % %))6<3D3D3I3IJ rGct|}d|jjvr&|jjd|jd<d|jjvr&|jjd|jd<|S)N dynamo_flat_name_to_original_fqndynamo_compile_id)r transformmoduler)r result_gmrs rHr z%FlattenInputOutputSignature.transforms}G%' -1A1A AAEAQAQ2BINN= > $++"2"2 226++2B2BCV2WINN. /rGrc)rrrz tuple[Any]r list[int]r list[Any]rr rzlist[torch.Tensor]rzlist[set[int]]rz$Optional[fake_tensor.FakeTensorMode]rr) rBrCrDrrrrr rrs@rHrrNs;?.) .).)+4 .)  .) ,5 .)0.)!/.)8.) .)`  ?(rGrc"eZdZUded<ded<y) ExportResultr graph_module_guards.GuardsSetrN)rBrCrDrYrFrGrHrrs&& rGrcg}|jjdD]}t|dsJt|dsJ|j}|jj |}|It |dkDsJd}|D]}t |dk(r|}n||jd}n]djtj|}d} t |dkDrd t |dz d } |jd || }|j||r,ttjd d j|zy)Nrr_dynamo_source_source_to_user_stacksrz, a closed over free variablerr*z(elided z more accesses)z, accessed at: aCannot export model which references tensors that are neither buffers/parameters/constants nor are direct inputs. For each tensor, if you'd like this tensor to be an explicit input, add it as a dummy argument to the top-level model definition you are exporting; if you would like its value to be embedded as an exported constant, wrap its access in a function marked with @assume_constant_result. z )graph find_nodesrrrrrrr traceback format_listrr6r7 INVALID_INPUT) r input_errorsrsource user_stacksstacksrtbextras rHcheck_signature_rewritabler"ssL &&-&8!t-...u6777$$2266v>   ;!### A1v{E    =[[]O#@AC..u56BE;!#"3{#3a#7"8H[[]O#3B4w?CC 7!:  ' ' E kk,'  (  rGc tj||\} } d} | |tjtj|\} }| | tj ddfd}|d||}|J|dt ||z| }t|||| ||| |j}dd}tt|| | |||j_ |j|S) Nc tjtjtjtjtj j gttzfd}|tjk(rdnd}|D]B}||r |tjk(r|"t|d|dt|d|dd y) Nc.t|tSrc)rr)rsupported_typess rHis_supported_typezMrewrite_signature..check_user_input_output..is_supported_typesc5#9: :rGinputrzIt looks like one of the z s with type `z;` is not supported or pytree-flattenable. Exported graphs z2s can only contain the following supported types: z. If you are using a custom class object, please register a pytree_flatten/unflatten function using `torch.utils._pytree.register_pytree_node` or `torch.export.register_dataclass`.) rKrSymIntSymFloatSymBoolr ScriptObjectlistr;r7rr6r) flat_values error_typer' value_typevr&s @rHcheck_user_input_outputz2rewrite_signature..check_user_input_outputs LL LL NN MM HH ! !  & ' ( ;!+m.I.I IWx  A$Q'!.check_optional_input_and_error!sm ++113 KD%}}G$5$5$;$;;9 IMM "I. rGcrg}i}t|D]\}}||t|<t|D]\}}t|ttr|j d1t||vr(|dk(rt dt|d|d|j |t||S)NinputszUnexpectedly found a z in the zO. Please file an issue along with a paste of the logs from TORCH_LOGS="+export")rr"rrr;rAssertionErrorr) debug_typesources candidatesmatched_elements_positionsdict_of_source_valsrrr?r=s rHproduce_matchingz+rewrite_signature..produce_matching1s:<" ( -FAs+, 3 ( - + PFAs#u%:;<*11$7C 33)259$+DI;hzlKdd +112Ebg2NO P*)rGrAoutputscdd}||}|jdt|}t|t|jkDrU|jJd|tdt|t|z Dcgc]}|jd|c}z }ndt|t|jkrC|jt|t|jxsg D]}||vrJd||t |j z }|jD]%}|jxsi} ||vr|| vrJd||Scc}w)Nc ht|jj}|Dcgc]5}|jtj j k(r |j7}}|Dcgc]6}|jtj jk(s+|j8}}td|Dd}td|Dd}td|D}|Dcic]f}|jtj jk(r=|jtj jur|j|jh}}i} |jrd|ji} |D]} | j| | j<t j||||||| Scc}wcc}wcc}w)Nc3K|]8}|jtjjk(s+|j:ywrc)kindrJr9VAR_POSITIONALrrps rHrz^rewrite_signature..argument_names..signature_to_fullargspec..is+VA166W5F5F5U5U+UV ,AAc3K|]8}|jtjjk(s+|j:ywrc)rMrJr9 VAR_KEYWORDrrOs rHrz^rewrite_signature..argument_names..signature_to_fullargspec..ms+SA166W5F5F5R5R+RSrQc3K|]]}|jtjjk(r4|jtjj ur|j_ywrc)rMrJr9POSITIONAL_OR_KEYWORDrSr:rOs rHrz^rewrite_signature..argument_names..signature_to_fullargspec..qsL66W..DDDIIW%6%6%<%<< sA#A%r)r-r7rrMrJr9rUr KEYWORD_ONLYrrrSr:return_annotation annotation FullArgSpec) sigparamsrPrm kwonlyargsvarargsvarkwdefaultskwonlydefaultsr parameters rHsignature_to_fullargspeczKrewrite_signature..argument_names..signature_to_fullargspec\s#..//12F 66W..DDDD !'!&&G4E4E4R4R*RJVVGSSE H 66W..;;;IIW%6%6%<%<< !NK$$')>)>? # C .7.B.B INN+ C&&guh NK I $s:F%(,F*F*!A+F/zMore arguments than expectedr_zMissing argument zMissing keyword only argument )rZinspect.Signature) rmrr]rr_r-rr\r`) r=rmrnrb fullargspec input_strsrunprovided_arg kwonly_argr`s rHargument_namesz)rewrite_signature..argument_names[s* X/u5 !%%kD 2 t9s;++, ,&&2 R4R R2 q#d)c*o"=>&&'q, JY[--. . #."2"2D S!5!5!;<<# V&/U3D^DT1UU/ V d6;;=)) &00 J(77=2N':+G 0 = G  5s=E)r=rd)rr)pytreetree_unflattenr7r tree_flattenINVALID_OUTPUTr-rr r(r)r_codegen recompile)r=rrrin_specrgraph_captured_inputgraph_captured_outputdynamo_traced_resultr orig_args orig_kwargsr2flat_results_tracedout_spec_tracedrHrr new_graphrir?s` @rHrewrite_signaturerys$229gFI{>I}'B'BC+1+>+>?S+T(/1M1MN **(8)1($ ! ,, ,(84-.:d |zf|_n$|r t)|}nt@jCd||jDjGD]]}t|tHjJst)djMtOjP|jR|d|d}_|r|$2sJdg"d%d&tjTjWtYd#'t[j\''Jt_|j`jG}tjbje}tg|D]o\}}tjh| |s|jk|}#jm| |d|jnd<%| |fz%"jq| |q|js&tjj}tjbju||$t@jCd$jwddnA'Jd#Jt@jCd$jwdd2r ty$"Dcgc]}#jm|}},rԈ$fd}t{5t}5#5 t|-dd0d|$ ddddddddd$J$jjd !D]e}tt+$|jtjs2#jmt+$|jd|jnd<g2r| D cgc]}|xsdD chc]w} | jt|k(r]t| tsM| jjj| jjjk7r | jyc} }!}} t|$#| | |%&| |! $t$'S#t($r }|}Yd}~d}~wwxYw#1swY xYwcc}w#t$r*}ttjt|dd}~wwxYw#1swYxYw#1swYxYw#1swYxYwcc} wcc} }w)$Nztorch._dynamo.exportzaSpecifying a decomposition_table table or tracing mode is illegal without setting aten_graph=Truez7pre_dispatch=True can only be used when aten_graph=TruecJd|y)Nz3whole graph export entails exactly one guard exportrFrs rHrz1export..inner..guard_export_print s % E % JrGcZJd|tj|fd}|S)Nz\Tried to emit a second graph during export. Tracing through 'f' must produce a single graph.c| Jt jd}t jd}tj|tj|n }t }|J|j x}r|j}|5t5|5i||}i}|jD]\}} |j| d||<tj|j|} tjj || dddddddddS#1swYxYw#1swYxYw#1swYSxYw)NF)remove_duplicateT static_shapes)rnamed_parameters named_buffersrdetect_fake_moderZ shape_envignore_fresh_unbacked_symbolsrr8rrjtree_maprKfuncfunctional_call) graph_inputsrrambient_fake_modeignore_fresh_unbackedrparams_and_buffersfake_params_buffersrrfake_graph_inputsrrrqgraph_captured_results rHresult_capturing_wrapperzhexport..inner..dynamo_normalization_capturing_compiler..result_capturing_wrapper;s(4$(((#'(>(>PU(>(V#W $U%8%8%%8%P Q // =I,,\:"")5%(444 1 ; ;;9;,5,S,S,U)&,.* ***'*&+-''9'?'?'A e4E4Q4Q!5R5+D1 )/)55|)%-2JJ,F,F24E-)'.-,/.-,s= E#E&A8D8E&E8E =EE EE)rr)rinner_example_inputsrrrrrqrs rH'dynamo_normalization_capturing_compilerzFexport..inner..dynamo_normalization_capturing_compiler)sC= n =E 002I1N1 -f, +rGrzTF)specialize_intrrr capture_dynamic_output_shape_opscapture_scalar_outputsrrr)rzr export_constraintsrrz#Summary of dimension constraints:%srzk It appears that you're trying to set a constraint on a value which we evaluated to have a static value of z0. Set TORCH_LOGS="+export" for more information.zhFailed to produce a graph during tracing as no tensor operations were found and same_signature is False.rF)rr rrzQFailed to capture a graph during tracing as no tensor operations were found.: %s) print_outputcoloredz'Failed to produce guards during tracingzDynamo captured graph: %sctjjj5tjj j |cdddS#1swYyxYwrc)rKrQrpreserve_node_meta Interpreterrun)rmrs rHgraph_with_interpreterz5export..inner..graph_with_interpretersOXX''::<B 88//6::DABBBs +AA(real)r}r~_allow_non_fake_inputsr|_allow_fake_constant cond_operandsr5get_attrr)rrr)SrrrrKrrrrrHrIrrJ signaturerjrlrr+rrr8r$rdim_constraints_opsOpOverloadPacket OpOverloadr.rsolveforced_specializationsprettify_resultsrmrr var_to_rangersympyIntegerrrr var_to_stack _subclassesFakeTensorModer&r GuardsSetr-r7rQGraphr is_tensorrrrrrr print_readabler"rrr#r3r6r7DYNAMIC_CONTROL_FLOWrTrrrrt_idr"rconstraint_rangevrlowerupperdimryr)4rmrn combined_args constraintsrrrcall_to_inspectoriginal_signature result_tracedrrrrpconstraint_violation_errorrrBrrrrkparameter_namesfx_graphrrrr trrxcrrrrrqrr_assume_static_by_default_f_specialize_floatrr{r}rrr|rrr~s4 @@@@@@rHrzexport..inners%b$7 -m^L ,#< !# $$%;<  * s :  XX X: O'1!UXX__'E!))1$..? #DH   E ,$E ,E ,R$00$@ 7!%)" : %'+ $ \\-%=%*-1#'8c4[   /O7$6"&#.E / %t 6v 6 + /0 !*%idCCP$-$=$==J%**"="=uzz?T?T!U %%o6  ! ! #%4%K%K%M "!22"*& C *.33A6<3*/*1I#1N.HH= ++002 a/1I779#8#89O9OPQ9R#STUVNNOSQII2.  &, , =! z > N#% $& !))88"*T9I!$..0 ) ))"#5#@#@#E#E#GHOxx~~'H$_5 84??9Q<0#//5D'0'<'<!! D(=(DIIe$,@9QY$ Z9 Z,ZY)&Z,.Z9 [ A<[ [ > YY YYYY!) Z 2%Z Z ZZ) $Z,,Z6 1Z99[[ zexport(f, *args, **kwargs) is deprecated, use export(f)(*args, **kwargs) instead. If you don't migrate, we may break your export call in the future if your user defined kwargs conflict with future kwargs added to export(f).r|r)rrrr)rr{r|r}r~rrrrrrrrrrrrrrs ````` ```` @@@rHr r snP3I;G B( 8a/a/F \  >   j1L11 rG)rzr rr r]ct|}t|dt}tt j ||||||||S)zF The same as `torch._dynamo.optimize(backend, nopython=True)` r)r r)r r r])rrrrZrzr,convert_frame_assert)rXrzr rr r]rs rHrr8sTg&Gw(:LI !** F7I    rGcVeZdZeej ddZedZy)rdNc &ddlm}|tjjtj_|tjj tj_|tjj tj_|tjjjjtjjj_tjjjdddl m}m}m}m}m}m}m}m}m} m} m} m} m} ||||||||| | | | | h }|D]O}|j6j9dd}d|}t;||s3t=|||t?||Qtj@jBjEDcgc]=}tGjH|r&tK|tj@jLr|?}}tj@jNtj@jPh}|D]A}||vr||jR|_)t;|d s+||jT|_*Cycc}w) Nr*)rFr) adadeltaadagradadamadamaxadamwasgdlbfgsnadamradamrmsproprpropsgd sparse_adam._fused_ _init_group)+ decoratorsrrKjittrace trace_module_get_trace_graphrQrRTracer distributions Distributionset_default_validate_args torch.optimrrrrrrrrrrrrrrBsplitrrroptimrrrJrKrL Optimizer SparseAdamLBFGSstepr)rrrrrrrrrrrrrroptimizer_modulesopt_modopt_name fused_fn_nameoptoptimizer_classesexcluded_optimizer_classess rHrzTorchPatcher.patchVs (!%))//2 !()?)?!@ %,UYY-G-G%H "07 HH $ $ + + 1 11   ''- ((BB5I    "               ) G''--c226H%hZ0Mw .]GGG]4S,T  {{++224 s# 3 8M8M(N    KK " " KK  & " % ;C00"388,sM*")#//":  ; sAJcfd}|S)NcLtjdtd|i|S)Nignoreztorch.distributed)categoryr )rfilterwarnings UserWarning)rmrnrs rHinner_fnzBTorchPatcher.suppress_torch_distributed_warnings..inner_fns+  # #;7J t&v& &rGrF)rr s` rH#suppress_torch_distributed_warningsz0TorchPatcher.suppress_torch_distributed_warningss ' rG)rBrCrD staticmethodr lru_cacherr rFrGrHrdrdUs@YH;H;TrGrdcht|ttjtjyrc)rrrSKIPrrs rH skip_coders#  0 0+2E2EFrG)rOr?)rUrRrrR)rr rc)rz)Union[types.CodeType, Callable[..., Any]]rzlist[CacheEntry])rrV)rzr8)inductor)r]z4Callable[[], Union[OptimizeContext, _NullDecorator]]rz&Union[OptimizeContext, _NullDecorator])rrr{rVr|rVr}z9Optional[dict[torch._ops.OpOverload, Callable[..., Any]]]r~rTrz6Optional[Union[dict[str, Any], tuple[Any], list[Any]]]rrVrrVrrVrrVrrVrrVrrVrzCallable[..., ExportResult])rztypes.CodeType)r __future__ratexitr#rrJloggingrrrr threadingrrrrF dataclassesrenumros.pathrrtypingrr r r r r unittest.mockrrrKtorch.fxtorch.utils._pytreer/_pytreerjtorch.utils.checkpointrr\rrrrrrtorch._dispatch.pythonrtorch._dynamo.typesrrrtorch._subclasses.fake_tensorrtorch._utils_internalrrtorch.export.dynamic_shapesrrrr torch.fx.experimental._dynamismr!r""torch.fx.experimental.proxy_tensorr#%torch.fx.experimental.symbolic_shapesr$r%r&r'torch.fx.graphr(r)rr+r,r-r.backends.registryr0r1r2excr3r4r5r6r7rzr8rkr:r;r<torch._subclassesr=r>r? getLoggerrBrExactWeakKeyDictionaryrVr$rZrArIrYrErrPrRr^r`_dynamo_forbiddenrarerurqrrsrrrM__file__rOrrHrIrrrrrrlocalrrrregisterrr r\rhrrzrrrrrrrrrrrrQ Transformerrrr"ryr rrdrrFrGrHr4s=,#   !!LL $$<RR@C !7 7GG9&;7-1g! "   ' ' [ *zzz2 )2 j $( $:$(:$@"   ( 9Z++    3" M  MEM,Mb/6F7FRk%(("6"6k\:) XEV "MQ!%*&+8=49"!}}} }  }}K}}#}} $}26}.2} !}$!%}F d   :UUprG