EiQ<SrSSKrSSKrSSKrSSKrSSKJrJr SSKJ r J r J r J r SSK r SSKJr SSKJr SSKJr SSKJr S S KJr S S KJr S S KJr "S S5r"SS5r"SS5r\"5rg)a_ This module provides the public comptime interface to TorchDynamo, enabling users to execute arbitrary Python code during symbolic evaluation of their programs. The comptime interface allows inspection and modification of TorchDynamo's compilation process while it is running. This can be useful for: - Debugging compilation issues - Inspecting intermediate state - Adding custom guards or graph breaks - Analyzing symbolic shapes and values Example usage: import torch from torch._dynamo.comptime import comptime def my_model(x): # Print the compile-time known information about x comptime.print(x) # Print the current FX graph being constructed comptime.print_graph() # Force a value to be treated as static if comptime(lambda ctx: ctx.get_local("x").is_dynamic()): comptime.force_static(x) # Add a manual graph break comptime.graph_break() Note: While this API provides significant flexibility, it intentionally avoids exposing internal implementation details of TorchDynamo to maintain compatibility across versions. N)CallableSequence)AnyOptionalTextIOUnion)InstructionTranslatorBase)VariableTracker) FakeTensor) free_symbols unimplemented) CellVariable)SymNodeVariablec\rSrSrSrS\SS4SjrS\\\\44Sjr S\ 4Sjr S\\ \ R44S jrSS \\S\\\ R44S jjrS\4S jrS\4S jrS\ 4SjrS\ 4SjrSSjrS\4SjrS\4SjrSrg) ComptimeVar7a A ComptimeVar represents a Python value, at some particular point in time, in the Python code we are symbolically evaluating with torchdynamo. This must be distinguished from a runtime value, as at compile-time there are some properties of the variable we do not know (for example, if the ComptimeVar represents a Tensor, we only know metadata about the tensor; we do NOT know what the actual data in the Tensor is.) vreturnNcXlgN_ComptimeVar__variable)selfrs S/home/mande/repo/quber/.venv/lib/python3.13/site-packages/torch/_dynamo/comptime.py__init__ComptimeVar.__init__Bsc6URR5$)a Returns an fx.Proxy (or tuple/list of fx.Proxy) representing this variable in the FX graph we are assembling to pass to the user compiler. This method only works for variables we actually track in the FX graph, aka Tensors (and ints, if you are compiling with dynamic shapes). In particular, if you have a list or tuple of tensors, you will get a list/tuple of proxies (not a single proxy representing the entire list/tuple). )ras_proxyrs rr!ComptimeVar.as_proxyEs''))rc6URR5$)z+ Returns True if as_proxy() would succeed. )ris_proxyr"s rr%ComptimeVar.is_proxySs''))rcdURR5RRS$)a@ Returns a "fake" value (either a FakeTensor or a SymInt) representing the variable in question. This only works for variables that denote Tensor or int. You can use this to query metadata; e.g., v.as_fake().size(0) will tell you the compile-time known size of the tensor. WARNING: Do NOT mutate the returned tensor. example_value)rr!nodemetar"s ras_fakeComptimeVar.as_fakeYs('')..33ODDrdimc@UR5RU5$)zw Returns the size of the tensor (if dim is None) or the size at the dimension dim. The returned size may be a SymInt. )r+size)rr-s rr/ComptimeVar.sizees ||~""3''rc6URR5$)zL Returns what type(v) would have returned for the variable at compile time. )r python_typer"s rr2ComptimeVar.python_typels **,,rc6URR5$)ay Returns the Python value this variable would have, but only if it is completely known at compile-time (e.g., it is constant). WARNING: Do NOT mutate the returned constant. The returned constant may or may not correspond to the actual value this variable may take on at runtime; for example, if the variable in question is a constant list, we may return a copy of that list. )ras_python_constantr"s rr5ComptimeVar.as_python_constantss1133rc6URR5$)z3 Returns True if as_python_constant would succeed. )ris_python_constantr"s rr8ComptimeVar.is_python_constants1133rc[UR[5(a*[URR5n[ U5$g)NF) isinstancerrr sym_numbool)rfss r is_dynamicComptimeVar.is_dynamics3 doo 7 7doo556B8Orc[UR[5(aURR5 gURR 5(ag[ SURS[ UR5S35e)zG Forces that a value is static, inducing a guard on its specific value z cannot force z (z) staticN)r;rr evaluate_exprr8AssertionErrortyper"s r force_staticComptimeVar.force_staticsi doo 7 7 OO ) ) + __ / / 1 1  043H2IR rcUR$)z Returns the internal data structure VariableTracker that Dynamo uses to represent variables at compile time. There are no BC guarantees on this API and WE RESERVE THE RIGHT TO BREAK YOUR CODE if you rely on it. rr"s r1_i_will_not_complain_if_bc_breaks_VariableTracker=ComptimeVar._i_will_not_complain_if_bc_breaks_VariableTrackersrc6URR5$r)r debug_reprr"s r__repr__ComptimeVar.__repr__s))++r) __variablerrN)__name__ __module__ __qualname____firstlineno____doc__r rrrr!r=r%r torchSymIntr+rintr/rDr2rr5r8r?rErHstrrL__static_attributes__rrrr7s/d *%/1J JK **$* Ez5<<78 E( (sELL7H1I(-T- 4C 44D4 D ?,#,rrc\rSrSrSrS\SS4SjrSS.S \S \S\ 4S jjr S$S \SS4S jjr S\ RR4SjrS\ SS4SjrSSS.S\S\\SS4SjjrS%SjrS \S\4SjrSS.S\S\\SS4SjjrSSS.S\\S \SS4SjjrSSS.S\\S \SS4SjjrSSS.S\\S \SS4SjjrSSS.S\\S \SS4SjjrSS.S\\SS4SjjrS\4S jrS!\\\ 4SS4S"jr!S#r"g)&ComptimeContextz This context class provides access to a public API for Dynamo's internals. If there is something here you would find useful that is missing, please file a feature request at https://github.com/pytorch/pytorch/ txrNcXlgr_ComptimeContext__tx)rr^s rrComptimeContext.__init__s rr stacklevelnamerdcURU5nURUn[U[5(a.[ UR R RU55$[ U5$)z< Retrieve the compile-time known information about a local. )_ComptimeContext__get_txsymbolic_localsr;rroutput side_effects load_cell)rrerdr^vars r get_localComptimeContext.get_locals\]]: &  & c< ( (ryy55??DE E3rmsgc$[SUSUS3/S9 g)z Manually trigger a graph break zComptimeContext graph breakz>## Y ##rT)verbosefilerrc[URRRR SUS9R US9 g)zg Print the partially constructed FX graph that would be passed to the user compiler after compilation. r)rrN)printrariry python_codesrc)rrrs r print_graphComptimeContext.print_graphs8  II   " " . .vw . G K KRV rc@[URR5$r)r\raparentr"s rrComptimeContext.parentstyy//00rcZURn[U5HnURnM U$r)raranger)rrdr^_s r__get_txComptimeContext.__get_txs( YYz"AB# rrc*[[U5US9 g)Nr)rrepr)rr{rs rrComptimeContext.prints d3id#r)rrdcURU5n[[R"URUR UR RS9R5US9 g)z Print the current series of opcodes being executed (not including parent frames), including where you are in the particular opcode stream. )current_offsetrN)rgrdisBytecodef_code instructionsinstruction_pointeroffset)rrrdr^s r print_disasComptimeContext.print_disassQ]]: & LL !r/E/EFMM ce  rcURU5nURHn[SUR53US9 M g)a Print the current Python value stack. Note that this is NOT the same as the traceback; use print_bt() to print that. Note that at stacklevel=0, this will typically be empty, as comptime cannot currently be used in an expression context where there would be intermediates on the stack. If you would find this useful, please file a bug at https://github.com/pytorch/pytorch/ NB: Stack grows downwards in our print z- rN)rgstackrrK)rrrdr^ss rprint_value_stack!ComptimeContext.print_value_stacks8]]: &A Bq||~&'d 3rcURU5nURR5H"upE[USUR 53US9 M$ g)z Print all of the locals available in the current context. By default this view is very limited; you can get more information about any individual local using get_local(). z = rN)rgrhitemsrrK)rrrdr^krs r print_localsComptimeContext.print_localssI]]: &&&,,.DA QCs1<<>*+$ 7/rc ,/nURU5nUb1URUR55 [USS5nUbM1[ SR [ RR[U55R55US9 g)ao Print the user code backtrace, starting at the beginning of the frame Dynamo started evaluating. Note that this MAY NOT go all the way to the torch.compile invocation, as we may have done a graph break and are compiling an intermediate frame as the starting point. If you think the other behavior would be better, file a bug at https://github.com/pytorch/pytorch/ Nrr) rgappend frame_summarygetattrrjoin traceback StackSummary from_listreversedformat)rrrdrr^s rprint_btComptimeContext.print_bts} ]]: &n LL))+ ,Xt,Bn  GGI**44Xe_ELLN O rc[SRS[URRR 555US9 g)z Print the currently installed guards for the Dynamo context. This does NOT include guards associated with variables that may or may not be installed in the future if those variables are used.  c3:# UHn[U5v M g7fr)r).0guards r /ComptimeContext.print_guards..<sT4S5e &4SsrN)rrsortedrariguards)rrs r print_guardsComptimeContext.print_guards2s6  IITF499;K;K;R;R4ST T rcUR$)z Returns the internal data structure InstructionTranslator that Dynamo uses to track state of symbolic evaluation. There are no BC guarantees on this API and WE RESERVE THE RIGHT TO BREAK YOUR CODE if you rely on it. r`r"s r7_i_will_not_complain_if_bc_breaks_InstructionTranslatorGComptimeContext._i_will_not_complain_if_bc_breaks_InstructionTranslator@syyrsecc0[R"U5 gr)timesleep)rrs rrComptimeContext.sleepKs  3r)__tx)zComptimeContext.graph_break)rr\)#rPrQrRrSrTr rrXrWrrmrvrUfxGraphryr}r=rrrrrrgrrrrrrrrfloatrrYrZrrr\r\s 49:  c  #  k    s  t  &uxx~~&   "&    -5f-=    133;?$$x'7$4$+/! ' _Comptime.Ssrfn fallback_fnrcU"5 g)zJfn gets called at compile time in TorchDynamo, calls fallback_fn otherwiseNrZ)rrs r__call___Comptime.__call__Ps  rNc[S5 g)Nc"UR5$r)rvctxs rr'_Comptime.graph_break..\ S__.rcomptimerZrrrv_Comptime.graph_breakZ ./rec(^[SU4Sj5 g)NcBURURS55$)Nr)rrmrs rr!_Comptime.print..`sSYYs}}S'9:rc>[T5$r)rrsrrr`s E!Hrrrs`rr_Comptime.print^s:.drrrrZrrr_Comptime.print_graphbrrrrcrdc[S5 g)Nc`URURS5R5S-S9$Nrdr rc)rrmr5rs rr'_Comptime.print_disas..is,==6IIKaO(rrrcs rr_Comptime.print_disasf  rc[S5 g)Nc`URURS5R5S-S9$rrrmr5rs rr-_Comptime.print_value_stack..q.--==6IIKaO.rrrcs rr_Comptime.print_value_stacknrrc[S5 U$)Nc`URURS5R5S-S9$rrrs rr8_Comptime.print_value_stack_and_return..|rrr)rrds rprint_value_stack_and_return&_Comptime.print_value_stack_and_returnys  rc[S5 g)Nc`URURS5R5S-S9$r)rrmr5rs rr(_Comptime.print_locals..s.((==6IIKaO)rrrcs rr_Comptime.print_localsrrc[S5 g)Nc`URURS5R5S-S9$r)rrmr5rs rr$_Comptime.print_bt..s, ==6IIKaO%rrrcs rr_Comptime.print_btrrc[S5 g)Nc"UR5$r)rrs rr(_Comptime.print_guards..s S--/rrrZrrr_Comptime.print_guardss /0rr{c[S5 g)NcBURURS55$Nr{)r}rmrs rr)_Comptime.assert_static..sS..s}}U/CDrrr{s rr}_Comptime.assert_statics DErc[S5 g)Nc@URS5R5$r )rmrErs rr(_Comptime.force_static..sS]]51>>@rrr s rrE_Comptime.force_statics @Arc4S[SS4Sjn[U5 g)a Like pdb breakpoint(), but drop into pdb whenever this line of code is compiled by dynamo. Use it by putting this in your model code:: from torch._dynamo.comptime import comptime comptime.breakpoint() And then, inside pdb, you can access 'ctx' to query things about the compilation context:: (Pdb) !ctx.print_bt() (Pdb) !ctx.print_locals() (Pdb) p ctx.get_local("attention").as_fake() inner_ctxrNcNUR5n[R"5 gr)rbuiltins breakpoint)rrs rinner#_Comptime.breakpoint..inners""$C    !r)r\r)rs rr_Comptime.breakpoints& "_ " " rrc[S5 g)Nc^URURS5R55$)Nr)rrmr5rs rr!_Comptime.sleep..sSYYs}}U';'N'N'PQrr)rs rr_Comptime.sleeps QRrrZrO)rPrQrRrS staticmethodrr\rrrvrrrWrrrrrrr}rErrrrrYrZrrrrOs*6 o&+ ,b#g& 00NNNN00)* 3 t  /0  T  BCC*+ C   &'  D  11F3F4FFB#B$BB0S5e$SSSrr) rTrrrrcollections.abcrrtypingrrrrrUtorch._dynamo.symbolic_convertr torch._dynamo.variables.baser torch._subclasses.fake_tensorr %torch.fx.experimental.symbolic_shapesr excr variablesrvariables.tensorrrr\rrrZrrr'sl"H .// D84>#-l,l,bddNjSjSZ ;r