
    9i|                   j'   % S r SSKJr  SSKrSSKrSSKrSSKrSSKrSSKrSSK	r	SSK
r
SSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSK r SSKJ!r!J"r"  SSKJ#r#J$r$  SSKJ%r%  SSKJ&r&  SSKJ'r'J(r(  SS	KJ)r)J*r*J+r+J,r,J-r-J.r.J/r/J0r0J1r1  SS
K2J3r3J4r4J5r5J6r6J7r7  SSK8r8SSK9r8SSK:r8SSK;J<s  J=r>  SSK8J?r?  SSK@JArAJBrBJCrCJDrD  SSKEJFrF  SSKGJHrHJIrI  SSKJJKrKJLrLJMrM  SSKNJOrO  SSKPJQrQJRrRJSrSJTrTJUrU  SSKVJWrWJXrX  SSKYJZrZ  SSK[J\r\  SSK]J^r^J_r_  SSK`Jara  SSKbJcrc  \R                  (       a0  SSKeJfrfJgrgJhrhJiriJjrjJkrkJlrlJmrmJnrn  SSKoJprp  SSKqJrrrJsrs  SSKtJuru  SSKvJwrw   SSKxry SSK{r8SSK|J}r~  SSKJJr  SSK{Jr  SS KJr  \y(       ay  \y\yGR                  \yGR                  \yGR
                  4rS!\S"'   \y\~\yGR                  \~GR                  \yGR                  \~GR                  \yGR
                  \~GR
                  0rOS#r0 rSS$KJrJrJr  \0" S%5      r\0" S&5      r\4" S'5      r\8GR"                  GR$                  GR&                  r\8GR"                  GR$                  GR*                  r\8GR"                  GR$                  GR.                  r\GR2                  " \RB                  5      rS(\S)'   0 rS*\S+'   S,rS-rS.\ S/3r\GR>                  " \5      r0 rS0\S1'   \GR2                  " \5      rS2\S3'   \GRJ                  " 5       r " S4 S5\
GRN                  5      r " S6 S75      r      GS[S8 jrSqGS\S9 jrGS\S: jrSqS;\S<'   GS]S= jrGS^S> jrSqGS_S? jrGS`S@ jrGS\SA jrSB\SC'   SD\SE'   GSaSF jrGSbSG jr " SH SI\
GRN                  5      r " SJ SK5      r\GRr                  " 5       r\$        GSc                   GSdSM jj5       r\/GSeGSfSN jj5       r\/ GSe     GSgSO jj5       r GSh     GSiSP jjr\GRz                  GS\SQ j5       r\8GR~                  \8GR                  \8GRF                  4\8GR                  \8GR                  \8GR                  4\8GR                  \8GR                  \8GR                  4\8GR                  \8GR                  4\8GR                  \8GR                  4\8GR                  \8GR                  4\8GR                  \8GR                  \8GR                  4\8GR                  \8GR                  \8GR                  4\8GR                  \8GR                  \8GR                  4\8GR                  \8GR                  40
r " SR SS5      r\" 5       rGSjST jrGS\SU jrGSjSV jrGSjSW jrGSkSX jrGSlSY jrGSmSZ jrGSnS[ jrGSoS\ jrGSpS] jr " S^ S_5      r\/GSqS` j5       r\/      GSrSa j5       r\/GSsSb j5       rGStSc jr\GR                  Sd:  aH  \GR                  \GR                  \Rh                  \R`                  \GR                  \GR                  4rGSuSe jrGSuSf jrGSuSg jr\/    GSvSh j5       r\/    GSwSi j5       r    GSuSj jr\1\GR                  \GR                  \GR                  \GR                  4   rSk\Sl'       GSxSm jr    GSySn jr\GR                  \GR                  \GR                  \GR                  \GR                  \GR                  So.rSpSqSrSsStSuSo.r    GSzSv jGr GS{Sw jGrGS|Sx jGrGS}Sy jGrGS~Sz jGrGSS{ jGr\RL                  " S|5      GSS} j5       GrGSGSS~ jjGrGSS jGrGSS jGr	GSS jGr
SGr\GR                   " S S5      5       GrSGr\GR                  " G\S9GqS\S'   GSS jGrGSS jGrGSS jGr            GSS jGr\H" G\S9Gr\I" G\S9GrGSS jGrGS\S jGrGSS jGr " S S5      GrSGqS\S'   GSS jGrGSS jGr\$  GS       GSS jj5       Gr\GR                   " S S5      5       Gr " S S\5      Gr G\ " 5       G\ Gl!        GSS jGr"SS.     GSS jjGr#\/    GSS j5       Gr$\/GSS j5       Gr$GSS jGr$GSS jGr%\$GSS j5       Gr&    GSS jGr'GSeGSS jjGr(GSS jGr)GS~S jGr*GSS jGr+\RL                  " S5      GSS j5       Gr,GSS jGr- GS       GSS jjGr.GSS jGr/\&" S5      GSS j5       Gr0\\G\1\G\2G\3G\4" S5      G\5GRl                  G\7GRl                  \RN                  \8GRp                  \8GRr                  \8GRt                  \8GRv                  \8GRx                  \8GRz                  \8GR"                  GR|                  GR~                  \8GR                  GR                  1GrBS\S'   \_" 5       (       a0  SSGKCGrCG\BGR                  G\CGR                  GRr                  5         GSS jGrF\GR                  GSS j5       GrHGSS jGrIGSS jGrJGSS jGrKGSS jGrLGSS jGrMGSS jGrN      GSS jGrOGSS jGrPG\4" 0 GR                  5       5      GrRS\S'   G\4" 0 GR                  5       5      GrTS\S'   G\4" 0 GR                  5       5      GrVS\S'   G\4" \"" 5       GR                  5       5      GrWS\S'   G\4" G\X" S#5      5      GrYS\S'   G\4" G\X" G\Z" S5      5      5      Gr[S\S'   G\YGR                  Gr]G\^GR                  Gr`G\aGR                  Grb\GR                  " G\aGR                  GR                  5       \"GR                  GR                  5       5       V s1 s H  n G\e" U 5      (       d  M  U iM     sn GrfG\gGR                  GR                  5        V s1 s H  n G\e" U 5      (       d  M  U iM     sn GrhG\iGR                  GR                  5        V s1 s H  n G\e" U 5      (       d  M  U iM     sn GrjG\kGR                  GrlG\kGR                  GR                  5        V s1 s H  n G\e" U 5      (       d  M  U iM     sn GrmG\nGR                  GR                  5        V s1 s H  n G\e" U 5      (       d  M  U iM     sn GroG\nGR                  GrqG\2GR                  GR                  5        V s1 s H  n G\e" U 5      (       d  M  U iM     sn Grr\0" S5      Grs\0" S5      GrtGSS jGruGSS jGrvGSS jGrwGSS jGrxGSS jGryGSS jGrzG\{Gr|GSS jGr}GSS jGr~G\aGR                  GrGSS jGrGSS jGrGSS jGrGSS jGrGSS jGrGSS jGrGSS jGr GSe         GSS jjGr    GSS jGrGSS jGrGSS jGrGSS jGrSGrSSGKGJGr  GSS jGrGSS jGrSLS.       GSS jjGrGSS jGr      GSS jGrGSS jGrSSLSSLSSLSL\GR*                  SLSL4
                         GSS jjGrGSS jGr\RH                  GSS j5       Gr\" 5       Gr\GR2                  " G\n5      GrS\S'   / GrS\S'   \" 5       Gr\GR                  GSS j5       GrGSS jGrGSGSS jjGrGSS jGr        GSS jGr GSe       GSS jjGr\GRr                  " 5       GrGSS jGr\$GSS j5       Gr            GSS jGrGSS jGrGSS jGrGSS jGrGSS jGrGSS jGrGSuS jGrGSS jGrGSS jGr GSe     GSS jjGr " GS  GS\
GRN                  5      GrGSGS jGr        GSGS jGrGSGS jGr            GSGS jGrGSGS/GrGSGS	/Gr/ GS
QGrG\G\-   G\-   GrGSGS jGr   GS         GSGS jjGr   GS         GSGS jjGrGSGS jGrGSGS jGr " GS GS\-\\4   5      GrGSGS jGr " GS GS5      Gr " GS GS\-\\4   5      GrGSGS jGr        GSGS jGrGS~GS jGrGS~GS jGrGSGS jGrGSGS jGrGSGS jGrGSGS jGr\GR                   " GS GS 5      5       GrGSGS! jGrGSGS" jGrGSGS# jGrGSuGS$ jGrGSuGS% jGrGSGS& jGrGSuGS' jGrGSGS( jGr      GSGS) jGrGSuGS* jGrGSGS+ jGr\RH                   GS       GSGS, jj5       GrGSGS- jGrGSGS. jGr " GS/ GS0\8GR"                  GR$                  5      Gr        GSGS1 jGrGSGS2 jGrGSGS3 jGr " GS4 GS55      GrG\g" 5       GrGS6\GS7'   GSGSGS8 jjGrGSGS9 jGr\RH                  GSGS: j5       GrGSGS; jGrGSGS< jGrGSGS= jGrGSGS> jGrGS\GS? jGrGSGS@ jGrGSGSA jGrGSGSB jGr\8GR                  GSGSC j5       Gr\8GR                  GSGSD j5       Gr\8GR                  GSGSE j5       GrGSGSF jGr0 GrGSG\GSH'   GSGSI jGrGSGSJ jGr " GSK GSL5      Gr " GSM GSN\5      GrGSGSO jGrGSPGrS\GSQ'   GSGSR jGr\$GSGSS j5       Gr\$GSGST j5       GrGSGSU jGrSLGqGSGSV jGrGSGSW jGr\8GR                  GSGSX j5       Gr\8GR                  GSGSY j5       GrGSGSZ jGrg! \z a    Sry GNf = f! \ a     GN]f = fs  sn f s  sn f s  sn f s  sn f s  sn f s  sn f (   a  
Utility functions and classes used throughout the TorchDynamo system.

This module contains a collection of helper utilities used by various parts of Dynamo for:
- Performance metrics collection and reporting
- Compilation timing and debugging
- Graph manipulation and tensor operations
- Runtime guards and checks
- Common data structure operations
- Testing and development tools

This is an internal module that provides shared functionality used across the Dynamo codebase.
    )annotationsN)CounterOrderedDict)AbstractContextManagercontextmanager)is_dataclass)	lru_cache)CodeTypeMethodWrapperType)	AnyCallablecastClassVarGenericOptionaloverloadTypeVarUnion)Literal	ParamSpec	TypeAlias	TypeGuardTypeIs)fx)_instruction_counter_len_torch_function_stack_pop_torch_function_stack_push_on_torch_function_stack)enable_python_dispatcher)MetricsContextRuntimeMetricsContext)	CompileIdSourceTracingContext)is_sparse_compressed)justknobs_checklog_chromium_event_internallog_compilation_eventrecord_chromium_event_internalsignpost_event)_format_graph_codelazy_format_graph_code)_WaitCounter)LazyModuleMixin)
has_tritonhas_triton_package)RemovableHandle   )_get_flat_args)		Container	Generator	ItemsViewIterableIteratorKeysViewMappingSequence
ValuesView)ExecutionRecord)InstructionTranslatorInstructionTranslatorBase)VariableTracker)DeviceLikeType)detect_fake_mode
LazyStringconfigztuple[types.ModuleType, ...]NP_SUPPORTED_MODULES )
FakeTensoris_fakemaybe_get_fake_modeTR_Pz*collections.defaultdict[str, Counter[str]]countersdict[str, Any]optimus_scuba_logzAhttps://pytorch.org/docs/main/torch.compiler_troubleshooting.htmlz;https://pytorch.org/docs/main/torch.compiler_nn_module.htmlzSee z& for more information and limitations.zdict[str, list[float]]compilation_time_metricsdict[str, float]cumulative_time_spent_nsc                       \ rS rSrSrSrSrSrg)ReInplaceTrigger   r2         rH   N)__name__
__module____qualname____firstlineno__AUTO_FUNC_V1AUTO_FUNC_V2
TRITON_OPS__static_attributes__rH       S/var/www/html/land-doc-ocr/venv/lib/python3.13/site-packages/torch/_dynamo/utils.pyrV   rV      s    LLJrb   rV   c                      \ rS rSr% \R
                  " \5      rS\S'   \	SS j5       r
\	SS j5       r\	SS j5       r\	SS j5       r\	SS j5       r\	SS	 j5       rS
rg)ReinplaceCounters   z!collections.defaultdict[str, int]_valuesc                Z    US:w  a%  U R                   SUR                   3==   U-  ss'   g g Nr   missed_bytes_)rg   name)clstriggerbytess      rc   add_missed_bytes"ReinplaceCounters.add_missed_bytes   s-    A:KK-~675@7 rb   c                F    US:w  a  U R                   SU 3==   U-  ss'   g g Nr   missed_tensors_)rg   )rl   rm   counts      rc   add_missed_opportunities*ReinplaceCounters.add_missed_opportunities   s)    A:KK/'34=4 rb   c                8    U R                   R                  5         g N)rg   clearrl   s    rc   ry   ReinplaceCounters.clear   s    rb   c                f    Sn[          H$  nXR                  R                  SU 3S5      -  nM&     U$ rr   )rV   rg   getrl   sumrm   s      rc   get_total_missed"ReinplaceCounters.get_total_missed   s4    'G;;??_WI#>BBC (
rb   c                z    Sn[          H.  nXR                  R                  SUR                   3S5      -  nM0     U$ ri   )rV   rg   r}   rk   r~   s      rc   get_total_missed_bytes(ReinplaceCounters.get_total_missed_bytes   s8    'G;;??]7<<.#A1EEC (
rb   c                V    U R                   (       a  [        SSU R                   5        g g )Ninductorreinplace_counters)rg   r*   rz   s    rc   logReinplaceCounters.log   s!     ;;:';S[[I rb   rH   N)rm   rV   rn   intreturnNone)rm   rV   rt   r   r   r   r   r   r   r   )rZ   r[   r\   r]   collectionsdefaultdictr   rg   __annotations__classmethodro   ru   ry   r   r   r   ra   rH   rb   rc   re   re      s    1<1H1H1MG.M A A
 > >       J Jrb   re   c                     SS K nUR                  XS9$ ! [         a1    SR                  S [        R                  " U/U 5       5       5      s $ f = f)Nr   headers
c              3  `   #    U  H$  nS R                  [        [        U5      5      v   M&     g7f), N)joinmapstr).0rows     rc   	<genexpr>tabulate.<locals>.<genexpr>   s&      
0PDIIc#sm$$0Ps   ,.)tabulateImportErrorr   	itertoolschain)rowsr   r   s      rc   r   r      sX    
   77 
yy 
09	40P
 
 	

s    8AAc                     [         S-   q g Nr2   )
curr_framerH   rb   rc   increment_framer      s    aJrb   c                 X    [         R                  5         [        R                  5         Sqg Nr   )rT   ry   rR   r   rH   rb   rc   reset_frame_countr      s    ""$""$Jrb   !Optional[list[Callable[[], str]]]_recompile_user_contextsc                @    [         c  / q [         R                  U 5        g)aR  
Register a hook to be called when a recompile is triggered. The hook
should return a string describing user contexts that are not available
to the compiler, such as the current training epoch. This is useful for
debugging and data analysis for recompile. For data retention purposes,
the user context string is capped at 256 characters.
N)r   append)hooks    rc   (register_hook_for_recompile_user_contextr   
  s      '#% ##D)rb   c                     [         $ rx   )r   rH   rb   rc   #get_hook_for_recompile_user_contextr         ##rb   c                    [         U -  q g rx   )op_count)cnts    rc   increment_op_countr     s    OHrb   c                     0 n [         R                  5        H  u  pUS-  X'   M     U R                  SS5      U R                  SS5      -   U S'   U $ )N    eAentire_frame_compiler   entire_backward_compiletotal_wall_time)rT   itemsr}   )total_by_keyphasetimings      rc   calculate_time_spentr   &  sc    L1779$sl : '3&6&6'2A6'7L"# rb   c            	         [        5       n SnU R                  5        H  u  p#U SU S[        US5       3nM     [        U5        g )NzTIMING: :   )r   r   roundprint)r   outkeyvalues       rc   print_time_reportr   6  sJ    ')L
C"((*
Qse1U5!_-. + 
#Jrb   r    _METRICS_CONTEXTr!   _RUNTIME_METRICS_CONTEXTc                     [         $ rx   )r   rH   rb   rc   get_metrics_contextr   Y  s    rb   c                     [         $ rx   )r   rH   rb   rc   get_runtime_metrics_contextr   ]  r   rb   c                  $    \ rS rSrSrSrSrSrSrg)CompileEventLogLevelia  z
Enum that loosely corresponds with a "log level" of a given event.

CHROMIUM_EVENT: Logs only to tlparse.
COMPILE_EVENT: Logs to tlparse + PT2 Compile Events
COMPILATION_METRIC: Logs to tlparse, PT2 Compile Events, and dynamo_compile
r2   rX   rY   rH   N)	rZ   r[   r\   r]   __doc__CHROMIUMPT2_COMPILECOMPILATION_METRICra   rH   rb   rc   r   r   a  s     HKrb   r   c                     \ rS rSrSr\S\R                  4         SS jj5       r\ S         SS jj5       r	\ S       SS jj5       r
\          SS j5       r\S\R                  4       SS	 jj5       r\          SS
 j5       r\\R                  4       SS jj5       r\SS j5       r\SS j5       r\SSS jj5       r\ S       SS jj5       r\SS j5       r\SS j5       rSrg) CompileEventLoggerio  a  
Helper class for representing adding metadata(i.e. columns) to various compile events.
Use CompileEventLogger to add event data to:
- Chromium events
- PT2 Compile Events
- CompilationMetrics

This should be used in conjunction with dynamo_timed() and metrics contexts, which create
timed spans and events. CompileEventLogger uses three log levels (described in CompileEventLogLevel),
where each log level logs to all sources below it in the hierarchy.

Example usages:
- I want to log to an existing chromium event within dynamo timed:
with dynamo_timed("my_event"):
    CompileEventLogger.chromium("my_event", foo=bar)

- I want to log my event to both chromium + pt2_compile_events:
with dynamo_timed("my_event", log_pt2_compile_event=True):
    CompileEventLogger.pt2_compile("my_event", foo=bar)

- I want to add information to dynamo events and dynamo_compile
    CompileEventLogger.compilation_metric(foo=bar)
Nc                    Uc  [         R                  " 5       n[        5       nU[        R                  :X  a  SnO"U[        R
                  :X  a  SnO[        S5      eUR                  XX5        g )NFTzjCannot log instant event at COMPILATION_METRIC level. Please choose one of CHROMIUM_EVENT or COMPILE_EVENT)timetime_nsget_chromium_event_loggerr   r   r   RuntimeErrorlog_instant_event)
event_namemetadatar   	log_levelchromium_loglog_pt2_compile_events         rc   r   $CompileEventLogger.log_instant_event  sk     ?llnG02,555$)!.:::$(!|  	&&	
rb   c                "   [        5       nUR                  5       nU[        R                  :X  a  UR                  " U 40 UD6  gU[        R
                  :X  a5  UR                  5       nX;  a  [        SU 5      eUR                  " U 40 UD6  gU[        R                  :X  d   eUR                  5       nX:w  a  [        S5      e[        5       nUR                  5       (       d  [        S5      eUR                  X25        UR                  " U 40 UD6  g)z
Centralized API for adding data to various events
Log an event to a toplevel "dynamo" event or metrics context
depending on log level.
zError: specified log level PT2_COMPILE, but the event %s is not logged to pt2_compile_events. Make sure the event is active and you passed log_pt2_compile_event=True to dynamo_timedzLog level is COMPILATION_METRIC, but event_name isn't the toplevel event. CompilationMetrics must be logged to the toplevel event. Consider using `log_toplevel_event_data` directly.z[No metrics context is in progress. Please only call this function within a metrics context.N)r   get_pt2_compile_substackr   r   add_event_datar   r   r   get_outermost_eventr   in_progressupdate)r   r   	overwriter   r   pt2_compile_substack	top_eventmetrics_contexts           rc   add_dataCompileEventLogger.add_data  s    12+DDF,555''
?h?.:::#/#H#H#J 5"A 	  ''
?h? 4 G GGGG$88:I&"B  23O"..00"q 
 ""87''
?h?rb   c                    [        5       R                  5       nUc  [        S5      e[        R                  " X0U40 UD6  g)z3
Syntactic sugar for logging to the toplevel event
NzWNo toplevel event active. Please only call this function within a dynamo_timed context.)r   r   r   r   r   )r   r   r   r   s       rc   add_toplevelCompileEventLogger.add_toplevel  sB     ./CCE	i  	##I)PxPrb   c                   [        5       nU[        R                  :X  d  U[        R                  :X  a  UR	                  XU5        gU[        R
                  :X  d   eUR                  5       nX:w  a  [        S5      e[        5       nUR                  5       (       d  [        S5      eUR	                  X#5        UR	                  XU5        g)z*
Increments an existing field, or adds it
zLog level is COMPILATION_METRIC, but event_name isn't the toplevel event. CompilationMetrics must be logged to the toplevel event. Consider using `increment_toplevel` directly.hNo metrics context is in progress. Please only call this function within a metrics context/dynamo_timed.N)
r   r   r   r   	incrementr   r   r   r   r   r   r   r   r   r   r   r   s          rc   r   CompileEventLogger.increment  s     12-6660<<<"":E: 4 G GGGG$88:I&"} 
 23O"..00"~  %%c1"":E:rb   r2   c                    [        5       nUR                  5       nUc  [        S5      e[        R	                  XBX5        g)zH
Increments a value on the toplevel metric. By default, logs to metric.
N_No toplevel event active. Please only call this function within a metrics context/dynamo_timed.)r   r   r   r   r   r   r   r   r   r   s        rc   increment_toplevel%CompileEventLogger.increment_toplevel  sB     12 446	q  	$$Y3Frb   c                   [        5       nU[        R                  :X  d  U[        R                  :X  a  UR	                  XU5        gU[        R
                  :X  d   eUR                  5       nX:w  a  [        S5      e[        5       nUR                  5       (       d  [        S5      eUR	                  X#5        UR	                  XU5        g)z\
Add metadata <value> to a set of values with key <key>. Creates a set if it doesn't exist.
zLog level is COMPILATION_METRIC, but event_name isn't the toplevel event. CompilationMetrics must be logged to the toplevel event. Consider using `add_to_set_metric` directly.r   N)
r   r   r   r   
add_to_setr   r   r   r   r   r   s          rc   r  CompileEventLogger.add_to_set
  s     12-6660<<<##JU; 4 G GGGG$88:I&"| 
 23O"..00"~  &&s2##JU;rb   c                    [        5       nUR                  5       nUc  [        S5      e[        R	                  XBX5        g)z
Same as add to set, just does it automatically to the toplevel event instead of having to explicitly name it.
Defaults to COMPILATION_METRIC log level.
Nr   )r   r   r   r   r  r   s        rc   add_to_set_toplevel&CompileEventLogger.add_to_set_toplevel)  sB     12 446	q  	%%iCGrb   c                T    [         R                  " U [        R                  4SS0UD6  g)z
Add <metadata> to <event_name> in chromium. Each key/value of metadata will appear in the chromium trace.
<event_name> should be the name of a timed event span passed to `dynamo_timed`.
r   FN)r   r   r   r   r   r   s     rc   chromiumCompileEventLogger.chromium=  s.     	##,55	
AF	
JR	
rb   c                T    [         R                  " U [        R                  4SS0UD6  g)aS  
Add <metadata> to <event_name> in chromium and PT2 Compile Events.
Each key/value of metadata will appear in the chromium trace. Each kwarg name becomes
a column in PT2 Compile Events, with the corresponding kwarg value.
<event_name> should be the name of a timed event span passed to `dynamo_timed`,
with log_to_pt2_compile_events=True.
r   FN)r   r   r   r   r  s     rc   pt2_compileCompileEventLogger.pt2_compileG  s.     	##,88	
DI	
MU	
rb   c                P    [         R                  " [        R                  U 40 UD6  g)a  
Add <metadata> to the CompilationMetrics context. Also logs to PT2 Compile Events
and chromium.
Each key/value of metadata will appear in the chromium trace. Each kwarg name becomes
a column in PT2 Compile Events and Dynamo Compile, with the corresponding kwarg value.
N)r   r   r   r   )r   r   s     rc   compilation_metric%CompileEventLogger.compilation_metricT  s%     	'' 33Y	
BJ	
rb   c                N    [         R                  XU[        R                  5        g)z
Log an instant event to chromium logs with name <event_name> at time <time_ns>. The `args` field in
Perfetto will point to metadata. <time_ns> should be a value obtained from time.time_ns().
N)r   r   r   r   )r   r   r   s      rc   instantCompileEventLogger.instant`  s      	,,'+?+H+H	
rb   c                ^    [        5       (       d  g[        5       nUR                  " U 40 UD6  g)z
Adds to an existing pt2_compile event, but silently returns if the event doesn't exist
or ChromiumEventLogger is not initialized.
This function is syntactic sugar for chromium_event_logger().try_add_event_data.
N)chromium_event_log_activer   try_add_event_data)r   r   r   s      rc   try_add_pt2_compile&CompileEventLogger.try_add_pt2_compilel  s+     )**02''
?h?rb   c                t    [        5       (       d  g[        5       nUR                  5       (       d  gU " U0 UD6  g)zz
Special function that quietly runs a given method, returning if CHROMIUM_EVENT_LOG is None or metrics context is not set
N)r  r   r   )	method_fnargskwargsr   s       rc   try_CompileEventLogger.try_x  s7    
 )**-/**,,4"6"rb   rH   )
r   r   r   rP   r   Optional[int]r   r   r   r   F)
r   r   r   r   r   boolr   objectr   r   )r   r   r   r   r   r!  r   r   )
r   r   r   r   r   r   r   r   r   r   )r   r   r   r   r   r   r   r   )
r   r   r   r   r   r   r   r   r   r   )r   r   r   r   r   r   r   r   )r   r   r   r!  r   r   )r   r   r   r!  r   r   rx   )r   r   r   rP   r   r  r   r   )r  zCallable[_P, Any]r  _P.argsr  	_P.kwargsr   r   )rZ   r[   r\   r]   r   staticmethodr   r   r   r   r   r   r   r   r  r  r  r  r  r  r  r  ra   rH   rb   rc   r   r   o  s   0  "&*>*G*G	

 
 
 (	

 

 
*   +@+@'+@ +@ 	+@
 
+@ +@Z ;@Q'Q48QNTQ	Q Q ;;$8;?B;KN;	; ;<  *>*Q*QGGG (G 
	G G  <<$8<?B<KN<	< <<  +?*Q*QHHH (H 
	H H& 
 
 

 

 	
 	
 LP	
	
#1	
<I	
		
 	
 	@ 	@ 	# 	#rb   r   Fc	           
   #  8	  #    U(       a  Un	U n
OU n	Sn
U [         ;  a	  / [         U '   0 nU(       a  UR                  U5        U
(       a  UR                  SU
05        Ub  UR                  SU05        [        5       n[        R                  " 5       nUR                  XXU5        [        R                  R                  U  S35      /nU(       a6  U(       a  UOU nUR                  [        SU 35      R                  5       5        [        R                  R                  R                  5       SLnU(       a  UR                  S5      (       d   e[!        ["        S5      (       d  S["        l        ["        =R$                  S	-  sl        ["        R$                  S	:X  aY  UR                  [        S
5      R                  5       5        U(       d*  SnUR                  [        U5      R                  5       5         [&        R(                  " 5        nU H  nUR+                  U5        M     Sv   SSS5        [        R                  " 5       nUU-
  n[         U    R                  US-  5        UR-                  U	U0 XU5        U(       a  [.        U	==   U-  ss'   ["        =R$                  S	-  sl        ["        R$                  S:H  nUS-  nU(       aM  [1        5       nUR3                  5       (       a-  UR5                  UU5        U(       a  UR5                  SU5        ggg[7        5       nUR5                  UU5        U(       a  USU(       + S.nUR5                  SUU5        ggg! , (       d  f       GN:= f! [        R                  " 5       nUU-
  n[         U    R                  US-  5        UR-                  U	U0 XU5        U(       a  [.        U	==   U-  ss'   ["        =R$                  S	-  sl        ["        R$                  S:H  nUS-  nU(       aM  [1        5       nUR3                  5       (       a-  UR5                  UU5        U(       a  UR5                  SU5        f f f [7        5       nUR5                  UU5        U(       a  USU(       + S.nUR5                  SUU5        f f f = f7f)a8  
dynamo_timed is a context manager
By wrapping a function in dynamo_timed, we can get a few things:

1) Optionally log timings to pt2_compile_events.
2) Optionally log timings to CompilationMetrics (dynamo_compile).
3) Optionally log chromium events.
4) Optionally increment a WaitCounter.
5) Store a record in compilation_time_metrics
   For example:

    def _foo(...):
        with dynamo_timed("_foo"):
            ...

    Would show up as an entry in our timing dict:
    OrderedDict([('_foo', [0.083690, 0.23949, 3.1425e-05])])
    This is extremely useful for granular debugging.

Although it is tempting to use dynamo_timed as a decorator, please do not.
In its decorator form it makes cProfile traces less useful as dynamo_timed
suddenly becomes a bottleneck for lots of function calls (as only one parent
pointer is recorded).

Params:
- key: key into compile_time_metrics. If phase_name is not provided, this is
  also the event name used for pt2_compile_events logs and chromium events.
- phase_name: Optional override for the event name.
- log_pt2_compile_event: Whether to log a pt2 compile event internally.
- metadata: Extra metadata to put in pt2_compile_events.
- dynamo_compile_column_us: If provided, updates the specified CompilationMetrics
  field to be logged to dyname_compile column. We expect all columns to be _us;
  therefore, the field name must end with "_us".
- compile_id: In the typical case, this parameter should not be needed. Use to
  supply the compile_id for those cases where we want to log a compile_id where
  it's not naturally available, e.g., for runtime autotuning.
- is_backward: Specify forward/backward directly when not available in a
  CompileContext, e.g., during runtime autotuning.
  that support it.
- log_waitcounter: If set, we'll log a waitcounter of the form "pytorch.dynamo_timed.{key}"
Nfn_nameis_backwardz (dynamo_timed)zpytorch.wait_counter._usdepthr   r2   z#pytorch.wait_counter.dynamo_compilez.pytorch.wait_counter.compile_runtime_overheadsr     duration_usT)
compile_id
is_runtime
is_forward)rR   r   r   r   r   log_event_starttorchprofilerrecord_functionr   r-   guard_guardsCompileContextcurrent_compile_idendswithhasattr_dynamo_timed_tlsr)  
contextlib	ExitStackenter_contextlog_event_endrT   r   r   r   r   )r   
phase_namer   r   dynamo_compile_column_usr,  r'  log_waitcounterwaitcounter_name_overrider   r&  event_metadatar   start_nscx_mgrswc_nameis_compile_time
runtime_wcstackcxend_nstime_spent_nsis_outer_eventr+  r   runtime_contextextras                              rc   dynamo_timedrO    s"    l 


**(* %Nh'y'23}k:;(A(CL||~H  nZ
 	&&#o'>?!G /H+c|&;G9$EFLLNOmm22EEGtSO'007777 ('22&'#1$ ""a'NN<(MNTTVW"M
|J7==?@(Q!!#u##B'  $
 ) %,,]S-@A""HZ	
 $ %Z0MA0 ##q(#.449N'4/K"5"7"..00#--.FT%'11-M & 1
 #>"?))*BKP!&0&**5oE
 $--m[%P "- $ $#
 ) %,,]S-@A""HZ	
 $ %Z0MA0 ##q(#.449N'4/K"5"7"..00#--.FT%'11-M & 1
 #>"?))*BKP!&0&**5oE
 $--m[%P "- $s8   G/R2M+ M&M+ .D+R
M(#M+ +D,RRc                    g rx   rH   repr	aggregates     rc   compile_timesrT    s    ILrb   c                    g rx   rH   rQ  s     rc   rT  rT    s     &)rb   c           	     N  ^ [         4SU4S jjjnU S:X  a5  [         Vs/ s H  nX2" [        U   S S94PM     nnSnU[        USS9-  nU$ U S:X  aJ  [        R                  5        Vs/ s H  nU" US	 S9PM     nn[	        [        R                  5       5      nX4$ g
s  snf s  snf )a  
Get metrics about torchdynamo frontend/backend compilation times.

Accumulates information from functions tagged with `dynamo_timed`.

repr='str' returns a printable string for user interaction, and 'csv'
returns headers, rows which can be logged for output

aggregate causes values from multiple compilations (e.g. split graphs)
to be accumulated into one value.  If false, expect more than one value
per metric.
r   c                h   > T(       a  U" [        U 5      5      $ SR                  [        X5      5      $ )Nr   )r   r   r   )valuesitem_fnrS  s     rc   fmt_fncompile_times.<locals>.fmt_fn4  s)    3v;''yyW-..rb   c                
    U S $ )Nz.4frH   xs    rc   <lambda>compile_times.<locals>.<lambda>;  s
    3jrb   )rY  z!TorchDynamo compilation metrics:
)FunctionzRuntimes (s)r   csvc                
    U S $ )Nz.6frH   r]  s    rc   r_  r`  C  s
    1S'
rb   N)rX  zlist[float]rY  zCallable[[float], str]r   r   )r   rR   r   rX  listkeys)	rR  rS  rZ  kr   r   vrX  r   s	    `       rc   rT  rT  $  s      GJ / /
 u} .
- /2<PQR- 	 
 3x&BCC
	 .446
6 1236 	 
 /4467

s   B(B"c                 >    [         R                  [        SSS95        g )Nr   TrQ  )r   inforT  rH   rb   rc   dump_compile_timesrj  K  s    HH]67rb   c                  6    \ rS rSrSSS jjrS	S jrS
S jrSrg)DuplicateWarningCheckeri^  c                0    Xl         U R                  5         g rx   )maxsizereset)selfrn  s     rc   __init__ DuplicateWarningChecker.__init___  s    

rb   c                "    [        5       U l        g rx   )r   setrp  s    rc   ro  DuplicateWarningChecker.resetc  s    *5-rb   c                f   XR                   ;   a2  U R                   R                  USS9  [        R                  (       d  g gS U R                   U'   [	        U R                   5      U R
                  :  a>  U R                   R                  SS9  [	        U R                   5      U R
                  :  a  M>  g)NT)lastF)rt  move_to_endrF   verboselenrn  popitemrp  r   s     rc   addDuplicateWarningChecker.addf  s    ((?HH  4 0>> "  !DHHSMdhh-$,,.  e , dhh-$,,.rb   )rn  rt  N)   )rn  r   r   r   r   )r   z!Union[str, tuple[object, object]]r   r   )rZ   r[   r\   r]   rq  ro  r~  ra   rH   rb   rc   rl  rl  ^  s    8	rb   rl  c                     [         R                  R                  SS5      S:H  n U (       a
  [        5       $ [        R
                  " 5       $ )NTORCH_COMPILE_DEBUG01)osenvironr}   add_file_handlerr:  r;  )compile_debugs    rc   setup_compile_debugr  u  s7    JJNN#8#>#EM!!!!rb   c                 ,    [         R                  5         g rx   )graph_break_dup_warning_checkerro  rH   rb   rc   reset_graph_break_dup_checkerr  ~  s    #))+rb   c                   ^^ [         R                  R                  [        5       S5      n [         R                  " U SS9  [
        R                  " [         R                  R                  U S5      5      m[
        R                  " S5      mTR                  T5        [        R                  " 5       nUR                  UU4S j5        U$ )NtorchdynamoTexist_okz	debug.logztorch._dynamoc                 &   > TR                  T 5      $ rx   removeHandlerlog_file_handlerloggers   rc   r_  "add_file_handler.<locals>.<lambda>  s    v334DErb   )r  pathr   get_debug_dirmakedirsloggingFileHandler	getLogger
addHandlerr:  r;  callback)log_path	exitstackr  r  s     @@rc   r  r    s    ww||MO];HKK4(**277<<++NO/F
&'$$&IEFrb   c                 N  ^^ [         R                  " 5       n [        R                  b|  [        R
                  " [        R                  5      m[        R                  R                  R                  5        H*  mTR                  T5        U R                  UU4S j5        M,     U $ U $ )Nc                 &   > TR                  T 5      $ rx   r  r  s   rc   r_   setup_log_file.<locals>.<lambda>  s    v';';<L'Mrb   )r:  r;  rF   log_file_namer  r  r0  _logging	_internalget_loggersr  r  )r  r  r  s    @@rc   setup_log_filer    s}    $$&I'"..v/C/CDnn..::<F./MN = rb   c                |    [        5        SUR                   S[        U 5      R                   SUR                   S3$ )Nz/error_recordings/_z.rec)r  co_nametyperZ   co_firstlineno)exccodes     rc   gen_record_file_namer    sB    o aS	""#1T%8%8$9? ?rb   c                    [         R                  R                  U 5      (       a  [        R	                  SU 5        g [         R
                  " [         R                  R                  U 5      SS9  [        U S5       nUR                  U5        S S S 5        g ! , (       d  f       g = f! [         a    [        R                  SU 5         g f = f)Nz9Unable to write execution record %s; file already exists.Tr  wbz#Unable to write execution record %s)r  r  existsr   warningr  dirnameopendump	Exception	exception)filenameexec_recordfs      rc   write_record_to_filer    s    
G77>>(##KKKX KK1DAh%  # &%% G;XFGs4   :B' >B' ;BB' 
B$ B' $B' ' C
	C
c                ^    SnU R                    H  nSUR                  ;   d  M  US-  nM     U$ )Nr   callr2   )nodesop)gcns      rc   count_callsr    s0    	AWWQTT>FA  Hrb   c                    U $ rx   rH   r]  s    rc   identityr    s    Hrb   c                R     [        U 5        g! [         a     g[         a     gf = fNTF)hash	TypeError
ValueErrorr]  s    rc   hashabler    s.    Q  s    
&	&&c                     g rx   rH   r  r  s     rc   nothingr    s    rb   c                  b    \ rS rSrSrSS jrSS jrSSS jjrSS jrSS jr	SS	 jr
SS
 jrSrg)ExactWeakKeyDictionaryi  z\Similar to weakref.WeakKeyDictionary, but use `is`/`id` rather than `==` to compare equalityc                     0 U l         0 U l        g rx   rX  refsru  s    rc   rq  ExactWeakKeyDictionary.__init__  s    &(;=	rb   c                2    U R                   [        U5         $ rx   )rX  idr}  s     rc   __getitem__"ExactWeakKeyDictionary.__getitem__  s    {{2c7##rb   Nc                L    U R                   R                  [        U5      U5      $ rx   )rX  r}   r  )rp  r   defaults      rc   r}   ExactWeakKeyDictionary.get  s    {{r#w00rb   c                2    [        U5      U R                  ;   $ rx   )r  rX  r}  s     rc   __contains__#ExactWeakKeyDictionary.__contains__  s    #w$++%%rb   c                   ^ ^ [        U5      mTT R                  ;  a)  [        R                  " UUU 4S j5      T R                  T'   UT R                  T'   g )Nc                &   > TR                  T5      $ rx   )
_remove_id)refidxrp  s    rc   r_  4ExactWeakKeyDictionary.__setitem__.<locals>.<lambda>  s    $//#:Nrb   )r  r  weakrefr  rX  )rp  r   r   r  s   `  @rc   __setitem__"ExactWeakKeyDictionary.__setitem__  s?    gdii$[[.NODIIcN Crb   c                v    XR                   ;   a  U R                   U	 XR                  ;   a  U R                  U	 g g rx   r  )rp  r  s     rc   r  !ExactWeakKeyDictionary._remove_id  s2    ++C ))		# rb   c                l    U R                   R                  5         U R                  R                  5         g rx   )r  ry   rX  ru  s    rc   ry   ExactWeakKeyDictionary.clear  s     		rb   )r  rX  r   )r   r   r   r   rx   )r   r   r  r   r   r   )r   r   r   r   )r   r   r   r   r   r   r  r   r   r   )rZ   r[   r\   r]   r   rq  r  r}   r  r  r  ry   ra   rH   rb   rc   r  r    s(    f>$1&!rb   r  c                    g rx   rH   objallowed_typess     rc   istyper    s    >Arb   c                    g rx   rH   r  s     rc   r  r    s     rb   c                    g rx   rH   r  s     rc   r  r    s    @Crb   c                x    [        U[        [        [        45      (       a  [	        U 5      U;   $ [	        U 5      UL $ )zisinstance() without subclasses)
isinstancetuplerd  rt  r  r  s     rc   r  r    s4    -%s!344CyM))9%%rb   rY      c                    [         R                  S:  a  [        U [        5      (       a  g[        U [        R
                  5      =(       d    U [        R                  L $ )Nr  T)sysversion_infor  _builtin_final_typing_classestyping_Finalr   r   s    rc   	is_typingr    sA     7"z%9V'W'WeV]]+Fu/FFrb   c                "   [         (       d  g[        U [         R                  [         R                  [         R                  [         R
                  [         R                  [         R                  [         R                  [         R                  45      $ NF)
npr  int8int16int32int64uint8uint16uint32uint64r  s    rc   is_numpy_int_typer    sW    2GGHHHHHHHHIIIIII		
 rb   c                    [         (       d  g[        U [         R                  [         R                  [         R                  45      $ r  )r  r  float16float32float64r  s    rc   is_numpy_float_typer  -  s4    2JJJJJJ	
 rb   c                    g rx   rH   r  s    rc   is_lru_cache_wrapped_functionr  ;  s     25rb   c                    g rx   rH   r  s    rc   r  r  A  s     47rb   c                    [        U [        R                  5      =(       a     [        [        R
                  " U S5      5      $ )N__wrapped__)r  	functools_lru_cache_wrapperis_functioninspectgetattr_staticr  s    rc   r  r  G  s6     eY99: {um4@ rb   r   
_FuncTypesc                    [        U 5      =(       d>    [        U [        R                  R                  [        R                  R
                  45      $ rx   )r  r  r0  _opsOpOverloadPacket
OpOverloadr  s    rc   is_function_or_wrapperr$  W  s;     u 

++UZZ-B-BC" rb   c                    [        U [        R                  [        R                  [        R                  [        R
                  45      $ rx   )r  typesFunctionTypeBuiltinFunctionTypeMethodDescriptorTypeWrapperDescriptorTyper  s    rc   r  r  _  s=     %%&&''		
 rb   )__eq____ne____lt____le____gt____ge__z==z!=<z<=>z>=c                    [        U [        R                  [        R                  [        R                  [        R
                  [        R                  45      $ rx   )r  r&  GetSetDescriptorTyper)  r*  MemberDescriptorTyper   r  s    rc   is_wrapper_or_member_descriptorr6    sH      &&&&''&&##	
 rb   c                    [        U 5      S   $ r   ) unwrap_with_attr_name_if_wrapper)fns    rc   unwrap_if_wrapperr:    s    +B/22rb   c                    [        U 5      (       a:  [        R                  " U SS5      (       a  [        R                  " U SU 5      n SnX4$ S nX4$ )N_torchdynamo_inlineF)r  r  r  )r9  	attr_names     rc   r8  r8    sU     2711"6KUSS##B(=rB)	 = 	=rb   c                N    [         (       d  g[        U [         R                  5      $ r  )r  r  ndarrayr  s    rc   is_numpy_ndarrayr@    s    2%$$rb   c                    [         R                  [         R                  R                  /[        R
                  Q7nU[         R                  R                  4-   n[        X5      $ )zCheck of obj is a tensor)	r0  Tensornn	ParameterrF   traceable_tensor_subclasses_subclassesrI   r  )r  tensor_lists     rc   istensorrH    sU     	% 
	+	+%K
 !2!2!=!= ??K###rb   c                "    [        U [        5      $ rx   )r  r.   )mods    rc   is_lazy_modulerK    s    c?++rb   r  c                     [        U 6   g rx   )r   r  s    rc   
print_oncerN    s	    	4Lrb   c                   ^ U mSU4S jjnUR                   b  [        UR                   5      S:X  d   eUR                   S   $ )zNSome black magic to create a cell object that usually only exists in a closurec                    > T $ rx   rH   r]  s   rc   r  make_cell.<locals>.f  s    rb   r2   r   )r   r   )__closure__r{  )valr  r^  s     @rc   	make_cellrT    s>    A ==$Q]]);q)@@@==rb   c                >    [        S U  5       5      nUR                  5        VVs0 s H  u  p4X4R                  5       _M     nnnX%4$ s  snnf ! [         aC  nSSKJn  SSKJn  U" SSU" U 6  SU" [        UR                  5       5      6  3S/ US	9   S nAg S nAff = f)
Nc              3  @   #    U  H  oR                  5       v   M     g 7frx   )as_proxyr   args     rc   r   $proxy_args_kwargs.<locals>.<genexpr>  s     :Tc<<>>T   r2   unimplemented_v2)typestrz&Failed to convert args/kwargs to proxyzcall_function args: r   z7Missing `as_proxy()` implementation for some arg/kwarg.gb_typecontextexplanationhintsfrom_exc)
r  r   rW  NotImplementedErrorr  r]  variables.baser^  rd  rX  )	r  r  
proxy_argsr   rY  proxy_kwargser]  r^  s	            rc   proxy_args_kwargsrj    s    
:T::
<BLLNKN\\^+NK'' L 

)+<*7D>*:!GT&--/EZ<[;\]Q	
	

s(   %A A	A 	A 
B9BBc                (    U c  S $ [        U S-  5      $ Nr*  r   rg  s    rc   	to_int_msro    s    94/#a$h-/rb   c                (    U c  S $ [        U S-  5      $ )Ni@B rm  rn  s    rc   	to_int_usrq    s    944#a)m"44rb   rY   c                     \ rS rSr% SrS\S'   SrS\S'   SrS\S'   SrS\S'   Sr	S\S	'   Sr
S\S
'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS\S'   SrS \S!'   SrS \S"'   Sr S\S#'   Sr!S$\S%'   Sr"S\S&'   Sr#S\S''   Sr$S$\S('   Sr%S$\S)'   Sr&S$\S*'   Sr'S\S+'   Sr(S$\S,'   Sr)S\S-'   Sr*S\S.'   Sr+S\S/'   Sr,S\S0'   Sr-S\S1'   Sr.S\S2'   Sr/S\S3'   Sr0S\S4'   Sr1S\S5'   Sr2S\S6'   Sr3S\S7'   Sr4S\S8'   Sr5S\S9'   Sr6S\S:'   Sr7S\S;'   Sr8S\S<'   Sr9S\S='   Sr:S\S>'   Sr;S\S?'   Sr<S\S@'   Sr=S\SA'   Sr>S\SB'   \?r@SC\SD'   SrAS\SE'   SrBS\SF'   SrCS\SG'   SrDS\SH'   SrES\SI'   SrFS\SJ'   SrGS\SK'   SrHS\SL'   SrIS\SM'   SrJSN\SO'   SrKS\SP'   SQrLS$\SR'   SrMS\SS'   SrNS$\ST'   SrOS$\SU'   SrPS\SV'   SrQS\SW'   SrRS\SX'   SrSS\SY'   SrTS\SZ'   SrUS\S['   SrVS\S\'   SrWS\S]'   SrXS\S^'   SrYS\S_'   SrZS\S`'   Sr[S\Sa'   Sr\S\Sb'   Sr]S\Sc'   SQr^S$\Sd'   \_SgSe j5       r`Sfrag)hCompilationMetricsi  NOptional[str]r,  	frame_keyr  co_filenamer  r  
cache_sizeaccumulated_cache_sizeguard_countshape_env_guard_countgraph_op_countgraph_node_countgraph_input_countOptional[float]
start_timeentire_frame_compile_time_sbackend_compile_time_sinductor_compile_time_scode_gen_time_s	fail_typefail_reasonfail_user_frame_filenamefail_user_frame_linenozOptional[set[str]]non_compliant_opscompliant_custom_opsrestart_reasonsdynamo_time_before_restart_szOptional[list[str]]stack_traceexception_stack_tracegraph_node_shapesOptional[bool]has_guarded_coderemote_cache_time_saved_sstructured_logging_overhead_sconfig_suppress_errors config_inline_inbuilt_nn_modulesspecialize_floatdynamo_configr.  num_triton_bundles!remote_fx_graph_cache_get_time_ms!remote_fx_graph_cache_put_time_msstart_time_usr+  !dynamo_cumulative_compile_time_us'aot_autograd_cumulative_compile_time_us#inductor_cumulative_compile_time_us,inductor_code_gen_cumulative_compile_time_ustriton_compile_time_usruntime_cudagraphify_time_usruntime_triton_autotune_time_us%dynamo_compile_time_before_restart_us distributed_ephemeral_timeout_usstructured_logging_overhead_us!remote_fx_graph_cache_get_time_us!remote_fx_graph_cache_put_time_us#backward_cumulative_compile_time_usend_time_uspre_grad_pass_time_uspost_grad_pass_time_usjoint_graph_pass_time_usr   log_format_versioninductor_configremote_cache_version"inductor_fx_remote_cache_hit_count#inductor_fx_remote_cache_miss_count%inductor_fx_remote_cache_backend_type!inductor_fx_remote_cache_hit_keys"inductor_fx_remote_cache_miss_keyscuda_versiontriton_versionzOptional[dict[str, bool]]feature_usagecompile_time_autotune_time_usFr-  
gc_time_ustensorify_float_attempttensorify_float_successtensorify_float_failureguard_latency_usrecompile_reasonnum_graph_breakstriton_kernel_compile_times_usir_countcudagraph_skip_reasonpython_version!pgo_put_remote_code_state_time_us!pgo_get_remote_code_state_time_usparam_numelparam_bytesparam_countrecompile_user_contexts#inline_inbuilt_nn_modules_candidatec                   SS jnSS jnSS jnSS jnU" UR                  S5      5      U" UR                  S5      5      U" UR                  S5      5      U" UR                  S5      5      U" UR                  S	5      5      U" UR                  S
5      5      U" UR                  S5      5      U" UR                  S5      5      U" UR                  S5      5      S.	n0 UEUEnU" UR                  S5      5      US'   U" UR                  S5      5      US'   U" UR                  S5      5      US'   UR                  S5      nU(       a  [        U5      OSUS'   U " S0 UD6$ )z
Factory method to create a CompilationMetrics from a dict of fields.
Includes the logic to add legacy fields and any pre-processing, e.g.,
we transform some fields to comma-separated strings for scuba logging.
c                    U b  U S-  $ S $ )Ng    .ArH   metrics    rc   us_to_s*CompilationMetrics.create.<locals>.us_to_s`  s    #)#56C<?4?rb   c                    U b  U S-  $ S $ rl  rH   r  s    rc   us_to_ms+CompilationMetrics.create.<locals>.us_to_msc  s    %+%76T>ATArb   c                   ^ SS jmU c  g [        U [        [        45      (       d  gSR                  U4S j[	        U 5       5       5      $ )Nc                :     [        U 5      $ ! [         a     gf = fN	<unknown>)r   r  )items    rc   safe_strFCompilationMetrics.create.<locals>.collection_to_str.<locals>.safe_strg  s#    't9$  '&'s   
 
r  ,c              3  4   >#    U  H  nT" U5      v   M     g 7frx   rH   )r   r  r  s     rc   r   GCompilationMetrics.create.<locals>.collection_to_str.<locals>.<genexpr>s  s     F~tHTNN~   )r  r   r   r   )r  rt  rd  r   sorted)r  r  s    @rc   collection_to_str4CompilationMetrics.create.<locals>.collection_to_strf  s?    ' ~fsDk22"88Fvf~FFFrb   c                j    U c  g  [         R                  " [        U 5      5      $ ! [         a     gf = fr  )jsondumpsrd  r  r  s    rc   collection_to_json_str9CompilationMetrics.create.<locals>.collection_to_json_stru  s4    ~#zz$v,// #"#s   % 
22r  r  r  r  r  r  r  r  r  )	r  r  r  r  r  r  r  r  r  r  r  r  r,  N)r  r  r   r~  )r  r  r   r  )r  Optional[Any]r   rt  rH   )r}   r   )	rl   metricsr  r  r  r  legacy_metricsall_metricsr,  s	            rc   createCompilationMetrics.createX  s   	@	B	G	# "'++o">?+2?@, '.EF' (/AB(  'JK  *1>?* 2:?@2 2:?@2 .5<=./
8 4373 <MOO?@<
78 =NOO@A=
89 9OOO<=9
45 !__\2
7AC
OtL!![!!rb   rH   )r  rP   r   rs  )brZ   r[   r\   r]   r,  r   ru  r  rv  r  rw  rx  ry  rz  r{  r|  r}  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r.  r  r  r  r  r+  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  LOG_FORMAT_VERSIONr  r  r  r  r  r  r  r  r  r  r  r  r-  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r   r  ra   rH   rb   rc   rs  rs    s    $J$#I}#!G]!!%K%$(NM( $J$,0M0!%K%+/=/$(NM(&*m*'+}+"&J&377.2O2/3_3'+O_+#I}#!%K%.2m2,0M0,0)0/3,3*.O'.48 /8'+K$+15.5'+}+ (,n+15559!?9-1N17;$n;'+n+#'M='!%J%(,,7;%};7;%};#'M='!%K%7;%};=A+]A9='=BF0-F,0M026 -659#]9;?)=?6:$m:48"M87;%};7;%};9='=!%K%+/=/,0M0.2m200%)O])*.-.8<&<9='=;?)=?7;%};8<&<"&L-&$(NM(/3M,337!=7!&J& $J$.2^2.2^226/6(,o,&*m*&*m*48"M8"Hm"+/=/$(NM(7;%};7;%}; "&K% "&K% "&K%26/6:?'?Q" Q"rb   rs  @   maxlenz%collections.deque[CompilationMetrics]_compilation_metricsc                   [        5       nUR                  5       nU(       d  gUR                  UU R                  U R                  U R
                  U R                  U R                  U R                  U R                  U R                  U R                  U R                  U R                  U R                  U R                  U R                   U R"                  U R$                  b  ['        U R$                  5      OSU R(                  b  ['        U R(                  5      OSU R*                  b  ['        U R*                  5      OSU R,                  U R.                  U R0                  S9  g)a  
These are the common fields in CompilationMetrics that existed before
metrics_context, and aren't set by MetricsContext.set(). We add the subset
of them that make sense in `dynamo`/toplevel events in PT2 Compile Events
directly.

If you're tempted to add to this list, consider using CompileEventLogger.compilation_metric()
instead, which will automatically also add it to tlparse and PT2 Compile Events.
TODO: Get rid of this function and replace it with CompileEventLogger directly instead.
N)r   ru  r  rv  r  rw  rx  ry  rz  r{  r|  r}  r  r  r  r  r  r  r  r  r  r  )r   r   r   ru  r  rv  r  rw  rx  ry  rz  r{  r|  r}  r  r  r  r  r  rd  r  r  r  r  r  )r  event_loggerr   s      rc   #add_compilation_metrics_to_chromiumr    s     -.L113J++		MM''<< 77MM55''++--++MM!"!;!; 77 *+)<)<)HD$$%d -.,B,B,ND''(TX ()'8'8'DD""#$%&%C%C++oo;   rb   c                 l    SS jn U " [         R                  " 5       5      n[        R                  " USS9$ )Nc           	         1 SknU R                  5        VVs0 s H/  u  p#X!;  d  M  U[        U[        5      (       a  [        U5      OU_M1     snn$ s  snnf )N>   rz  base_dirrepro_afterrepro_levelTYPE_CHECKINGr  debug_dir_rootrepro_tolerancerepro_forward_only_custom_ops_profile_save_config_ignorerepro_ignore_non_fpignore_logger_methodslog_compilation_metricssame_two_models_use_fp64rE  reorderable_logging_functionsnontraceable_tensor_subclasses)_autograd_backward_strict_mode_banned_ops+inject_BUILD_SET_unimplemented_TESTING_ONLY)r   r  rt  r  )d	blocklistr   r   s       rc   clean_for_json6_get_dynamo_config_for_logging.<locals>.clean_for_json  sX    
	2  ggi
'
# DC*UC"8"8eC'
 	
 
s
   A'AT)	sort_keys)r  rP   r   rP   )rF   get_config_copyr  r  )r	  config_dicts     rc   _get_dynamo_config_for_loggingr    s-    
< !!7!7!9:K::kT22rb   c                 P    " S S[         R                  5      n [        5       nSnSn[        R                  R
                  (       a)   [        R                  R
                  R                  5       nUb   UR                  5        HO  u  pE[        U[        5      (       d  UR                  U5        [        U[        5      (       d  MB  [        U5      X4'   MQ     U H  nX4	 M     [         R                  " UU SSS9nU$ U$ ! [        [        4 a    Sn Nf = f! [         a    Sn U$ f = f)zF
Method to parse and scrub uninteresting configs from inductor config
c                  ,   ^  \ rS rSrSU 4S jjrSrU =r$ )A_scrubbed_inductor_config_for_logging.<locals>.TypeSafeSerializeri  c                D   >  [         TU ]  U5      $ ! [         a     gf = f)NzValue is not JSON serializable)superr  r  )rp  o	__class__s     rc   r  I_scrubbed_inductor_config_for_logging.<locals>.TypeSafeSerializer.default  s)    8wq)) 878s    
rH   )r  r   r   r   )rZ   r[   r\   r]   r  ra   __classcell__r  s   @rc   TypeSafeSerializerr    s    	8 	8rb   r  Nz!Inductor Config cannot be pickledT)rl   skipkeysr  z(Inductor Config is not JSON serializable)r  JSONEncoderrt  r0  	_inductorrF   r  r  AttributeErrorr   r  r   r~  rd  r  r  )r  keys_to_scrubinductor_conf_strinductor_config_copyr   rS  s         rc   %_scrubbed_inductor_config_for_loggingr!    s&   8T-- 8 "eM	D#(??#9#9#I#I#K  '	K0668!#s++!%%c*c3''04S	(- 9 %(- % !%

$&	! 1 >* 	D C	D*  	K J	Ks+   (C= 4AD 3D =DDD%$D%c                   ^ [         R                  R                  R                  5       (       a   SSKJn  UnSnOS nS nUR                  S5      nU(       d(  [         R                  R                  R                  5       nUU S-  US-  U(       a  UR                  OS U(       a  [        U5      OS [        [         R                  R                  5       5      [!        5       ["        R$                  ["        R&                  [)        5       [         R*                  R,                  [/        5       (       a  [0        R2                  OSUU[4        R*                  S.n	[6        R9                  0 U	EUE5      m[:        R=                  T5        Sn
TR>                  S	L a  S
U
-   n
TR@                  SL a  U
S-   n
[         R                  RC                  U
U4S jS	US9  [E        T5        ["        RF                  (       a  [I        T5        g g ! [         a    S nS n GNf = f)Nr   )REMOTE_CACHE_VERSION_ManifoldCacher,  r*   )r,  r  r  r  r  r  r  r  r  r  r  r  r  r  r  compilation_metricsFbwd_T_runtimec            	        > [         R                  " T5      R                  5        V Vs0 s H(  u  pU [        U[        5      (       a  [        U5      OU_M*     snn $ s  snn f rx   )dataclassesasdictr   r  rt  rd  )rf  rg  r&  s     rc   r_  ,record_compilation_metrics.<locals>.<lambda>q  sQ    #**+>?EEG
G *Q,,tAw!3G
 
s   /A)record_logging_overheadr,  )%r0  r  utils should_use_remote_fx_graph_cachetorch._inductor.fb.remote_cacher#  ModuleNotFoundErrorr}   r4  r5  r6  r\   r   rq  r  get_structured_logging_overheadr  rF   suppress_errorsinline_inbuilt_nn_modulesr!  versioncudar/   triton__version__r  rs  r  r  r   r.  r-  trace_structuredr  r  r(   )start_time_nsend_time_nsr  exc_type	exc_valuer#  r  r  r,  common_metricsrk   r&  s              @rc   record_compilation_metricsr?  8  s    ==??	9L#7 4D1
 15-# \*J]]11DDF
 !&$."d*.6X**D)2s9~*3NN::<+
 89"("8"8,2,L,L@B**0:&,," 41V++#N( -334Q~4Q4QR 34 D%%.}%%-j 	NN##	
 !&  $ $ ((;< %%12 & # 	9#' 481	9s   
G; ;HH)on_exitc                    [        [        5      U :  a)  [        R                  5         [        [        5      U :  a  M)  [        R                  " [        U S9nUqg )Nr  )r{  r  popleftr   deque)new_size	new_deques     rc   set_compilation_metrics_limitrF    sE    
"
#h
.$$& "
#h
.!!"6xHI$rb   c                 ,    [         R                  5         g rx   )r  ry   rH   rb   rc   clear_compilation_metricsrH    s     rb   c                      [        [        5      $ rx   )rd  r  rH   rb   rc   get_compilation_metricsrJ    s    $%%rb   c                  0   \ rS rSrSrSS jrSS jrSS jrSS jrSS jr	SS jr
      SS	 jrSS
 jr        SS jr  S           SS jjrSS jr S             SS jjr S         SS jjr  S         S S jjrSrg)!ChromiumEventLoggeri  a  Logs chromium events to structured logs. tlparse will concatenate these into a perfetto UI link.

See https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#heading=h.yr4qxyxotyw for
a specification of the Chromium Event JSON format.
c                    [        U R                  S5      (       a  U R                  R                  $ / U R                  l        U R                  R                  $ )zE
The main event stack, with every chromium event.
Logged to tlparse.
rH  )r8  tlsrH  ru  s    rc   	get_stackChromiumEventLogger.get_stack  s<    
 488W%%88>>!DHHN88>>!rb   c                >    U R                  5       nU(       a  US   $ S$ )z^
Get the outermost event name (i.e. the longest running event)
or None if the stack is empty.
r   N)rO  )rp  rH  s     rc   r   'ChromiumEventLogger.get_outermost_event  s!    
   uQx*d*rb   c                    [        U R                  S5      (       a  U R                  R                  $ / U R                  l        U R                  R                  $ )zY
A smaller subset of the main stack that gets used to log
PT2 Compile Events internally.
r   )r8  rN  r   ru  s    rc   r   ,ChromiumEventLogger.get_pt2_compile_substack  sB    
 48834488000,.DHH)88000rb   c                    [        U R                  S5      (       d  0 U R                  l        U R                  R                  $ )N
event_data)r8  rN  rV  ru  s    rc   get_event_data"ChromiumEventLogger.get_event_data  s/    txx.."$DHHxx"""rb   c                @   [         R                  " 5       U l        SSKJn  UR
                  (       a*  UR
                   S[        R                  " 5        3U l        O#[        [        R                  " 5       5      U l        [        R                  SU R                  5        g )Nr2   rE   -z*ChromiumEventLogger initialized with id %s)	threadinglocalrN  r%  rF   pt2_compile_id_prefixuuiduuid4id_r   r   ri  )rp  rF   s     rc   rq  ChromiumEventLogger.__init__  sd    ??$ '' 667qGDH4::<(DH 	=txxHrb   c                R    XR                  5       ;  a  gU R                  " U40 UD6  g)zT
Same as add_event_data, but will silently not log if the event isn't in the stack.
N)rO  r   )rp  r   r  s      rc   r  &ChromiumEventLogger.try_add_event_data  s'     ^^--J1&1rb   c                    XR                  5       ;  a)  [        S[        U5       SU R                  5        S35      eU R                  5       nX;  a  0 X1'   X1   R	                  U5        g)zb
Adds additional metadata info to an in-progress event
This metadata is recorded in the END event
Event  not in q. Cannot add metadata to events that aren't in progress. Please make sure the event has started and hasn't ended.N)rO  r   rR  rW  r   )rp  r   r  rV  s       rc   r   "ChromiumEventLogger.add_event_data  su     ^^--j)*(4>>3C2D EK K 
 ((*
'%'J"%%f-rb   c                    XR                  5       ;  a)  [        S[        U5       SU R                  5        S35      eU R                  5       nX;  a  0 XA'   X$U   ;  a  SXA   U'   XA   U==   U-  ss'   g)z;
Increment an integer event data field by the given amount
re  rf  rg  r   N)rO  r   rR  rW  rp  r   r   r   rV  s        rc   r   ChromiumEventLogger.increment  s     ^^--j)*(4>>3C2D EK K  ((*
'%'J",,*+J"3's#u,#rb   c                   XR                  5       ;  a)  [        S[        U5       SU R                  5        S35      eU R                  5       nX;  a  0 XA'   X$U   ;  a  [	        5       XA   U'   XA   U   R                  U5        g)zB
Add a value to a set within a event_name's metadata if it exists
re  rf  rg  N)rO  r   rR  rW  rt  r~  rj  s        rc   r  ChromiumEventLogger.add_to_set  s     ^^--j)*(4>>3C2D EK K 
 ((*
'%'J",,*-%J"3's#''.rb   Nc                \   U=(       d(    [         R                  R                  R                  5       n[	        U5      US'   U R                  UUSU5        U R                  5       R                  U5        U R                  " U40 UD6  U(       a   U R                  5       R                  U5        gg)aQ  
Logs the start of a single event.
:param str event_name Name of event to appear in trace
:param time_ns Timestamp in nanoseconds
:param metadata: Any extra metadata associated with this event
:param log_pt2_compile_event: If True, log to pt2_compile_events
:param compile_id: Explicit compile_id (rather than using the current context)
r,  BN)
r0  r4  r5  r6  r   _log_timed_eventrO  r   r   r   )rp  r   r   r   r   r,  s         rc   r/  #ChromiumEventLogger.log_event_start  s       T5==#?#?#R#R#T
!$Z		
 	
+J3(3 ))+22:> !rb   c                    U R                  5       nU R                  5       nUR                  5         UR                  5         U R                  5       nUR                  5         g rx   )rO  r   ry   rW  )rp  rH  substackrV  s       rc   ro  ChromiumEventLogger.reset9  sK      002((*
rb   c                4  ^ U=(       d(    [         R                  R                  R                  5       n[	        U5      US'   U R                  5       nTU;   a	  UT   nUT	 O0 nUR                  U5        U R                  TUSU5      n	SU4S jjn
U R                  5       nTU;  a  [        R                  S5        gU
" U5        U(       a?  U R                  5       nU
" U5        [        XU R                  U5        UR                  5         UR                  5         g)a  
Logs the end of a single event. This function should only be
called after log_event_start with the same event_name.
:param event_name: Name of event to appear in trace
:param time_ns: Timestamp in nanoseconds
:param metadata: Any extra metadata associated with this event
:param start_time_ns: The start time timestamp in nanoseconds
:param log_pt_compile_event: If True, log to pt2_compile_events
:param compile_id: Explicit compile_id (rather than using the current context)
r,  Ec                z   > TU S   :w  a1  [         R                  S5        U R                  5         TU S   :w  a  M0  g g )Nz>ChromiumEventLogger: Detected overlapping events, fixing stack)r   r  pop)rH  r   s    rc   	pop_stack4ChromiumEventLogger.log_event_end.<locals>.pop_stackj  s9    b	) T 		 b	)rb   z7ChromiumEventLogger: Start event not in stack, ignoringN)rH  	list[str]r   r   )r0  r4  r5  r6  r   rW  r   rp  rO  r   r  r   r'   r`  ry  )rp  r   r   r   r:  r   r,  all_event_datarB  eventrz  event_stackr   s    `           rc   r=  !ChromiumEventLogger.log_event_endC  s   &  T5==#?#?#R#R#T
!$Z ,,.'+J7Nz*Nh'%%	
	 nn& [( KKQR+ #'#@#@#B *+'TXX} !$$& 	rb   c                   ^ UUS-  UUSSSS.m[         R                  R                  SU4S jSSS9  [        T5        T$ )	zQ
Logs a timed event in chromium format. See log_event_start, log_event_end, etc.
r*  rO  r   )rk   tsr  phcattidpidchromium_eventc                    > T $ rx   rH   r~  s   rc   r_  6ChromiumEventLogger._log_timed_event.<locals>.<lambda>      urb   F
payload_fnsuppress_contextexpect_trace_id)r0  r  r9  r)   )rp  r   r   r   r   r~  s        @rc   rp  $ChromiumEventLogger._log_timed_event  sY     D.!	
 	''$"!	 	( 	
 	'u-rb   c           	     :  ^ Uc  0 n[        [        R                  R                  R	                  5       5      nXSS'   UUS-  USSSSSS.m[        R
                  R                  S	U4S
 jSSS9  U(       a'  [        TU R                  5       U R                  U5        gg)a*  
Log an instant event with no associated duration.
:param str event_name: Name of event to appear in trace
:param int time_ns Timestamp in nanoseconds
:param Optional[Dict[str, Any]] metadata: Any extra metadata associated with this event
:param str cname optional color for the arrow in the trace
Nr,  r*  irO  r   p)rk   r  r  r  r  r  r  sr  c                    > T $ rx   rH   r  s   rc   r_  7ChromiumEventLogger.log_instant_event.<locals>.<lambda>  r  rb   FTr  )
r   r0  r4  r5  r6  r  r9  r'   r   r`  )rp  r   r   r   r   r,  r~  s         @rc   r   %ChromiumEventLogger.log_instant_event  s     H55HHJK
!+D.!

 	''$" 	 	( 	
 !'t446' !rb   )r`  rN  )r   r|  r   rt  )r   rP   r   )r   r   r  r   r   r   )r   r   r   r   r   r   r   r   )r   r   r   r   r   r   r   r   FN)r   r   r   r   r   rP   r   r   r,  Optional[CompileId]r   r   rx   )r   r   r   r   r   rP   r:  r   r   r   r,  r  r   r   )
r   r   r   r   r   r   r   Optional[dict[str, Any]]r   rP   r  )
r   r   r   r   r   r  r   r   r   r   )rZ   r[   r\   r]   r   rO  r   r   rW  rq  r  r   r   r  r/  ro  r=  rp  r   ra   rH   rb   rc   rL  rL    s   	"+	1#
I2.. . 
	.(-$// / 	/
 
/6 ',*.?? ? !	?
  $? (? 
?<" +/GG G !	G
 G  $G (G 
G\ .2  	
 + 
B .2&+(( ( +	(  $( 
( (rb   rL  zOptional[ChromiumEventLogger]CHROMIUM_EVENT_LOGc                 0    [         c
  [        5       q [         $ rx   )r  rL  rH   rb   rc   r   r     s    !02rb   c                     [         S L$ rx   )r  rH   rb   rc   r  r    s    T))rb   c           	   #    #    [        5       n[        R                  " 5       nUR                  U U0 U5         Sv   UR	                  U [        R                  " 5       0 UU5        U(       a  UR                  5         gg! UR	                  U [        R                  " 5       0 UU5        U(       a  UR                  5         f f = f7f)z
Context manager that creates a chromium start and end event. Chromium event
logging is integrated with dynamo_timed, so you probably want to use that
instead. Use this context manager only if you want to avoid dynamo_timed.
N)r   r   r   r/  r=  ro  )r   reset_event_log_on_exitr   chromium_event_logchromium_start_times        rc   chromium_event_timedr    s      34,,.&&
	'((LLN!	
 #$$& # 	((LLN!	
 #$$& #s   4CA< AC<AB>>Cc                  L    \ rS rSr% SrS\S'   S\S'   SS jr\SS j5       rS	r	g
)CleanupHooki  z,Remove a global variable when hook is calledrP   scoper   rk   c                r    [         b  [         =R                  S-  sl        U R                  U R                  	 g r   )CleanupManagerrt   r  rk   rp  r  s     rc   __call__CleanupHook.__call__  s)    %  A% JJtyy!rb   c                `    X;  d   e[         =R                  S-  sl        X U'   [        X5      $ r   )r  rt   r  )r  rk   rS  s      rc   r  CleanupHook.create  s1       !d5''rb   rH   Nr  r   r   r   )r  rP   rk   r   rS  r   r   r  )
rZ   r[   r\   r]   r   r   r  r$  r  ra   rH   rb   rc   r  r    s)    6
I" ( (rb   r  c                  <   ^  \ rS rSr% SrS\S'   SU 4S jjrSrU =r$ )r  i  r   zClassVar[CleanupManager]instancec                ^   > U R                   U    H
  nU" 5         M     [        TU ]	  U5        g rx   )rX  r  r  )rp  r  r   r  s      rc   r  CleanupManager._remove_id  s(    KK$DF %3rb   rH   r  )	rZ   r[   r\   r]   rt   r   r  ra   r  r  s   @rc   r  r    s    E&&   rb   r  c                    U R                  5       R                  U R                  5      nU R                  (       a,  U R                  b  U R                  R                  5       Ul        U$ )z!Clone the tensor and its gradient)clonerequires_grad_requires_gradis_leafgrad)r^  ys     rc   clone_tensorr  &  sB    		  1AyyQVV'Hrb   dtypec          	     :  ^ [        U 5      (       a  U $ SU4S jjn[        R                  " 5          U R                  R                  S:X  a  U" U 5      sSSS5        $ U R
                  [        R                  L a`  [        R                  " U" U R                  5       5      U" U R                  5       5      U R                  U R                  5       S9sSSS5        $ [        U 5      (       a  U R
                  [        R                  [        R                  1;   a!  U R                  5       nU R!                  5       nO U R#                  5       nU R%                  5       n[        R&                  " U" U5      U" U5      U" U R)                  5       5      U R                  U R
                  S9sSSS5        $ [+        S [-        U R/                  5       U R1                  5       5       5       5      nU R2                  (       a  [        R4                  " US-   4U 5      nO6[        R6                  " US-   T=(       d    U R8                  U R                  S9nU R;                  5       UR;                  5       -
  S-  U R=                  5       -  nUR?                  U R/                  5       U R1                  5       U5         URA                  U RC                  5       5        U RD                  (       a  URG                  U RH                  5        U RD                  (       a&  U RJ                  b  [M        U RJ                  TS	9Ul%        [Q        U S
5      (       a  U RR                  RU                  5       Ul)        UsSSS5        $ ! [N         a    U" U 5      s sSSS5        $ f = f! , (       d  f       g= f)zcopy while preserving stridesc                Z  > [         R                  " U 5      nU R                  (       a  UR                  U R                  5        U R                  (       a&  U R
                  b  [        U R
                  TS9Ul        [        U S5      (       a  U R                  R                  5       Ul        U$ )Nr  _dynamo_dynamic_indices)
r0  r  r  r  r  r  clone_inputr8  r  copy)r^  r  r  s     rc   torch_clone clone_input.<locals>.torch_clone7  su    KKN99Q__-99+ u5AF1/00()(A(A(F(F(HA%rb   xlaN)is_coalesced)layoutc              3  6   #    U  H  u  pUS -
  U-  v   M     g7f)r2   NrH   )r   shapestrides      rc   r   clone_input.<locals>.<genexpr>]  s      
6O]UUQY& 6Os       )r  devicer  r  r^  torch.Tensorr   r  )+rJ   r0  no_gradr  r  r  
sparse_coosparse_coo_tensor_indicesrg   r  r  r%   
sparse_csr
sparse_bsrcrow_indicescol_indicesccol_indicesrow_indicessparse_compressed_tensorrX  r   zipsizer  is_quantizedempty_quantizedemptyr  data_ptrelement_sizeas_strided_copy_r  r  r  r  r  r  r   r8  r  r  )r^  r  r  compressed_indicesplain_indicesneeded_sizeresultcache_line_offsets    `      rc   r  r  .  s   
 qzz 
88==E!q> 
 88u'''**AJJL)AIIK(^^-	 
 "!$$xxE,,e.>.>??%&^^%5" !%&^^%5" !11./M*AHHJ'xx) 
8  
69!&&(AHHJ6O
 
 >>**K",<+>BF[[b (8F ZZ\FOO--3^^ 	1668QXXZ1BC
	"LL#yy%%aoo6yyQVV/)!&&> 1/00-.-F-F-K-K-MF*m 
^  	" q>!g 
^	"_ 
sE   "NA3NC
N.C?N.BM+01N+N	=NN		N
Nc                    g rx   rH   example_inputss    rc   clone_inputsr  z  s     rb   c                    g rx   rH   r  s    rc   r  r    s    :=rb   c                   [        U 5      [        L a  [        U 5      nUR                  5        He  u  p#[        U[        5      (       a  [        U5      X'   M)  [        U[        R                  5      (       d   [        U5      5       e[        U5      X'   Mg     U$ [        U 5      n[        [        U5      5       H5  n[        X   [        R                  5      (       d  M&  [        X   5      X'   M7     U$ rx   )r  dictr   r  r  r  r0  rB  r  rd  ranger{  )r  resr   r   r  s        rc   r  r    s    Nt#>"))+JC%'''.!%66CUC6&u- & 

~
C3s8_cfell++ (CF  Jrb   c                     U R                  5         g ! [         a8  nSSKJn  [        R
                  " SS[        U 5      5      nU" SU 35      UeS nAff = f)Nr2   )	SkipFramez\(.*r%  z(torch.compile cannot be run in context: )r  r   r  r  resubrR  )rS  ri  r  functorch_subclass_names       rc   skip_frame_if_in_functorch_moder    sX    	 " #%&&"d3i"@67N6OP
	s    
A3AAc               #    #    [         R                  R                  n [         R                  R                  R
                  nU" 5          U " 5          [         R                  " [         R                  R                  5       5      n[        U5        [         R                  R                  5       (       a2  [         R                  " [         R                  R                  5       5      nS S S 5        S S S 5         S v   [         R                  R                  R                  5          [         R                  R                  W5        [         R                  R                  5       (       a  [         R                  R                  W5        S S S 5        g ! , (       d  f       N= f! , (       d  f       N= f! , (       d  f       g = f! [         R                  R                  R                  5          [         R                  R                  W5        [         R                  R                  5       (       a  [         R                  R                  W5        S S S 5        f ! , (       d  f       f = f= f7frx   )r0  _C_DisableFuncTorchr.  _python_dispatch_disable_current_modesr  randomget_rng_stater  r6  is_availableset_rng_state)disable_functorchdisable_current_modes	rng_statecuda_rng_states       rc   preserve_rng_stater    s`    22!KK88OO		 "3"5KK : : <=	'	2::""$$"[[)A)A)CDN	 #6	 
9[[))@@BLL&&y1zz&&((

((8 CB #6"5	 	  CBU[[))@@BLL&&y1zz&&((

((8 CBBs   AI&FBF#F+I&4F> 8(I& A"F-	I&
F	F
F*&I&-
F;7I&>)I#'A"I		I#
I I##I&c                
   [        U [        R                  R                  R                  [        R                  R
                  R                  [        R                  R                  [        R                  R                  45      $ rx   )	r  r0  jit_traceTopLevelTracedModule_scriptRecursiveScriptModuleScriptFunctionScriptModule)model0s    rc   is_jit_modelr    sY     II11II33II$$II""		
 rb   c                R   [        U 5      (       a  U $  [        R                  R                  X5      $ ! [         ai     [        R                  R                  U 5      s $ ! [         a7    U(       a  [        R                  S5         O[        R                  S5          g f = f g f = f)Nz	jit errorz0Both torch.jit.trace and torch.jit.script failed)	r  r0  r  tracer  scriptr   r  error)modelr  rz  s      rc   torchscriptr    s    E	Nyyu55 N	N99##E** 	Nk*		LM	N* Ns>   3 
B&AB&&B B&B B&B  B&%B&c                \     [         R                  " U 5      $ ! [        [        4 a     g f = frx   )r  getfiler  OSErrorr  s    rc   r  r    s-    s##w s    ++c                *    [        [        U 5      5      $ )zLTest if an object is a namedtuple or a torch.return_types.* quasi-namedtuple)is_namedtuple_clsr  r  s    rc   is_namedtupler    s    T#Y''rb   c                  ^   [        T [        5      (       a  [        T SS5      nUS;   a  g[        [        T SS5      [        5      (       ae  [	        [        T SS5      5      (       aJ  [        S T R
                   5       5      nU[        4:X  a  gU(       a  [        U 4S jU 5       5      (       a  gg	! [         a     g	f = f)
zhTest if an object is a namedtuple or a (torch.return_types|torch.autograd.forward_ad).* quasi-namedtupler[   N)torch.return_typesztorch.autograd.forward_adT_fields_makec              3  :   #    U  H  o[         Ld  M  Uv   M     g 7frx   )r   )r   ts     rc   r   $is_namedtuple_cls.<locals>.<genexpr>  s     KA7:Jaas   	c              3  x   >#    U  H/  n[        U5      =(       a    TR                  UR                  L v   M1     g 7frx   )r  __new__)r   r!  rl   s     rc   r   r"    s;      ! #	 *!, 5KK19945 #s   7:F)
issubclassr  getattrr  callable	__bases__anyr  )rl   modulebasess   `  rc   r  r    s    c5!!S,5FLL'#y$7??HWd+E E KKKUH$S ! #!       s   (B2 A"B2 !B2 2
B?>B?tuple[str, ...]c                8   U [         L a  g[        U [        5      (       d   e[        U S5      (       a  U R                  $ [
        R                   " S S5      5       nU R                  S:X  d   eU " [        U[        U R                  5      5      5      n0 n[        U5       HA  nUS   S:w  d  M  [        [        X$5      U5      (       d  M*  [        X$5      R                  X4'   MC     [        U5      U R                  :X  d   e[        [!        X3R"                  S95      $ )	zIGet the fields of a namedtuple or a torch.return_types.* quasi-namedtuple)startstopstepr  c                       \ rS rSr% S\S'   Srg)!namedtuple_fields.<locals>.Markeri	  r   indexrH   NrZ   r[   r\   r]   r   ra   rH   rb   rc   Markerr2  	  s    
rb   r5  r  r   r  )r   )slicer%  r  r8  r  r*  	dataclassr[   r   r  n_fieldsdirr  r&  r3  r{  r  r}   )rl   r5  r  fieldsrk   s        rc   namedtuple_fieldsr;  	  s     e|(c5!!!!sI{{   >>1111
c&%-.
/CFC7c>j);VDD"3-33FL  v;#,,&&&JJ/00rb   c           	     ^  ^^^ [         R                  " 5          [         R                  " [         R                  R	                  5       5      m[         R
                  R                  5       (       a2  [         R                  " [         R
                  R	                  5       5      m[        R                  " U R                  5       U R                  5       5       Vs/ s H%  nXR                  [         R                  " U5      4PM'     snmS S S 5        SUUU4S jjnU$ s  snf ! , (       d  f       N= f)Nc                   > [         R                  " 5          [         R                  R                  T5        [         R                  R                  5       (       a  [         R                  R                  T5        T H)  u  pnU R                  U:w  d  M  U R                  U5        M+     S S S 5        g ! , (       d  f       g = frx   )r0  r  r  r  r6  r  _versionr  )paramr5  original_valuer  r  saved_states      rc   restore"checkpoint_params.<locals>.restore+	  st    ]]_LL&&y1zz&&((

((82=.>>W,KK/ 3>	 __s   A8B1B11
B?r   )r0  r  r  r  r  r6  r  r   r   
parametersbuffersr>  )gmr?  rB  r  r  rA  s      @@@rc   checkpoint_paramsrG  !	  s    	KK : : <=	::""$$"[[)A)A)CDN #"**,G
G NNEKK$67G
	 
0 0 N
	 
s   B=D,DDD
D,c                   [         R                  R                  5       (       a  [         R                  R                  nO[        nU" 5         [
        R                  " 5         [         R                  " S5        [        R                  " 5       n[        U5       H  nU " U6 nU" 5         M     [        R                  " 5       nWXt-
  4$ )Ni9  )r0  r6  r  synchronizer  gccollectmanual_seedr   perf_counterr  )r  r  timesrI  t0r  r  t1s           rc   timedrQ  7	  s     zz  jj,,MJJL	d				B5\'  
			B27?rb   c           	     l    [        S [        R                  " XR                  S5      5       5       5      $ )Nc              3  8   #    U  H  oR                   v   M     g 7frx   )is_cudar   r^  s     rc   r    check_is_cuda.<locals>.<genexpr>K	  s     W"VQyy"V   T)allr   r   rD  )rF  r  s     rc   check_is_cudarY  J	  s&    W)//.--PTBU"VWWWrb   r  c                    U S:  d   e[        U 5       Vs/ s H  nSU 3PM
     nn[        USS  US S -   5      n[        SSR                  U5       SSR                  U5       S35      nSU  S	3Ul        U$ s  snf )
Nr2   rg  rx  zlambda r  z: ()rot__helper)r  reversedevalr   rZ   )r  r  varsrotatedr9  s        rc   rot_n_helperrb  N	  s    q5L5"1X&XasGXD&tBCy49,-G	's388G+<*=Q?	@B7#BKI	 's   A8z	set[type]common_constant_typesc                B   [        U [        [        45      (       a  [        [	        [
        U 5      5      $ [        U [        R                  [        [        R                  [        R                  [        R                  45      =(       d    [        U [         ["        1-  5      $ rx   )r  r  	frozensetrX  r   is_safe_constantr  enumEnumr  r0  Sizer   _GenericAliasr&  GenericAliasrc  r6  rn  s    rc   rf  rf  {	  s{    a%#$$3'+,,	IIJJ  	
	  
	'
rb   c                 
    SS1$ Nr   r2   rH   rH   rb   rc   common_constantsrn  	  s    
 	
		 rb   c                    [        U [        R                  [        R                  45      =(       aG    [        U R                  [        R
                  R                  R                  R                  5      (       + $ rx   )	r  r0  SymBoolSymIntnodenestedr  
nested_intNestedIntNoder  s    rc   is_torch_symrv  	  sP    eemmU\\:; J

ELL**55CCE A rb   c                   SSK Jn  [        R                  " 5       R                  (       + =(       Ga     UR                  5       R                  5       (       + =(       d    UR                  5       R                  5       =(       a    [        R                  (       + =(       d    UR                  5       R                  5       =(       a    [        R                  (       + =(       dt    UR                  5       R                  5       =(       a    [        R                  (       + =(       d3    U" U5      =(       d$    [        S5      (       + =(       a    U [        5       ;   $ )Nr2   )is_from_defaultsz5pytorch/dynamo:enable_unspecialize_zero_one_plain_int)sourcerx  r$   r}   #force_unspec_int_unbacked_size_likeguard_sourceis_localis_specialized_nn_modulerF   allow_unspec_int_on_nn_module"is_unspecialized_builtin_nn_moduleis_unspecialized_nn_moduler&   rn  )r   ry  rx  s      rc   is_int_specialization_caser  	  s    (!!#GGG  !**,, 	

 !::< 9888	
 !DDF 9888	
 !<<> 9888	
  F#!	
(   WXX ,)++/rb   c                   SSK JnJnJn  [	        X5      (       a  U R                  5       (       d  U R                  5       nU R                  5       n[        U5      =(       dD    [        R                  (       + =(       a(    [        U5      [        L =(       a    [        XT5      (       + nU(       d  U $ [	        X5      (       a  UR                  U R                  5       5      $ U $ )Nr2   )ConstantVariableLazyVariableTrackerSymNodeVariable)	variablesr  r  r  r  is_realizedoriginal_sourcer@  rv  rF   specialize_intr  r   r  r  evaluate_expr)rY  r  r  r  ry  r   is_symnode_vts          rc   specialize_symnoder  	  s    QQ #++COO4E4E $$&""$$U+ 
%%% >Us">.u== 	 J#''&&s'8'8':;;Jrb   c                h    SSK Jn  [        U 5      n [        X5      (       a  U R	                  5       $ U $ )Nr2   r  )r  r  r  r  as_python_constant)rY  r  s     rc   guard_if_dynr  	  s.    +
S
!C#((%%''Jrb   c                j    [        S [        R                  " XR                  5       5       5       5      $ )Nc              3  @   #    U  H  oR                  5       v   M     g 7frx   )is_python_constantrU  s     rc   r   &check_constant_args.<locals>.<genexpr>	  s     V/U!##%%/Ur[  )rX  r   r   rX  r  s     rc   check_constant_argsr  	  s"    Vyt]]_/UVVVrb   c                    SSK Jn  SSKJn  Sn[        R
                  " XR                  5       5       H,  n[        XS5      (       a  US-  nM  [        XR5      (       a  M,    g   US:  $ )Nr2   r  UnspecializedPythonVariabler   F)variables.constantr  variables.tensorr  r   r   rX  r  )r  r  r  r  unspec_countr^  s         rc   check_unspec_python_argsr  	  sV    4=L__T==?3a55ALA00	 4
 !rb   c                    SSK Jn  [        R                  " XR	                  5       5       H,  nUR                  5       (       a  M  [        X25      (       a  M,    g   g)Nr2   r  FT)r  r  r   r   rX  r  r  )r  r  r  r^  s       rc   check_unspec_or_constant_argsr  	  sA    
 >__T==?3$$&&*Q*T*T 4 rb   c                ~   ^ SSK Jm  [        U4S j[        R                  " XR                  5       5       5       5      $ )Nr2   )NumpyNdarrayVariablec              3  <   >#    U  H  n[        UT5      v   M     g 7frx   )r  )r   r^  r  s     rc   r   +check_numpy_ndarray_args.<locals>.<genexpr>
  s#      7A 	1*++7s   )r  r  r)  r   r   rX  )r  r  r  s     @rc   check_numpy_ndarray_argsr  	  s/    6 }}7  rb   ztype[KeysView[Any]]	dict_keysztype[ValuesView[Any]]dict_valuesztype[ItemsView[Any, Any]]
dict_itemsodict_valuesztype[Iterator[Any]]tuple_iteratorrange_iteratorKVc                Z    [        U [        5      (       d   e[        R                  U 5      $ rx   )r  r  re  )r  s    rc   builtin_dict_keysr  %
  s#    a99Q<rb   c                   [        U [        5      (       d   e[        U [        [        45      (       a  U R	                  5       $ [        U [        5      (       a<  [        R
                  " U 5       Vs/ s H  o[        R                  " X5      4PM     sn$ [        R                  U 5       Vs/ s H  o[        R                  X5      4PM     sn$ s  snf s  snf rx   )r  r  r  r   r   re  r  )r  rf  s     rc   get_items_from_dictr  +
  s    c4    cD+&''yy{	C	%	%>I>N>Ns>ST>SK++C34>STT7;yy~F~!D$$S,-~FF UFs   0!C+ Cc                n    [        U 5      n[        R                  R                  R	                  U5        U$ rx   )
object_newr0  rC  Modulerq  )rl   r  s     rc   nn_module_newr  6
  s&    
S/C	HHOOS!Jrb   c                N    [         R                  " [        R                  U S5      $ r   )r  reduceoperatormul)its    rc   productr  <
  s    HLL"a00rb   c                :    U R                  5       u  nu  p4X4U-      $ rx   )
__reduce__)r  r3  r  r  r.  s        rc   tuple_iterator_getitemr  @
  s!    }}Avu}rb   c                h    [         R                  R                  [        R                  5      " U 5      $ rx   )r0  _dynamodisabler*  r:  rz   s    rc   dataclass_fieldsr  E
  s"    ==  !3!34S99rb   c                    U R                  5       u  nu  p#UR                  U=(       d    SUR                  -  -   nUR                  nUR                  nXEU4$ r   )r  r.  r0  r/  )
range_iterr  	range_obj	maybe_idxr.  r/  r0  s          rc   normalize_range_iterr  L
  sQ    !+!6!6!8A|	 OOy~A??E>>D>>Drb   c                $    U R                  U5      $ rx   )as_subclass)r!  rl   s     rc   to_subclassr  Z
  s    ==rb   c                    [         n[        U [        5      (       a  [        n[        [        R
                  " UR                  U 5      XS-   5      5      $ r   )r  r  r   nextr   islicere  )r  r  
dict_classs      rc   dict_keys_getitemr  a
  s>    J![!! 
	  !3QA>??rb   c                    [        U 5      U   $ rx   rd  )r  r  s     rc   set_getitemr  i
  s    71:rb   c                t    U R                   R                  nU R                  nU(       a  SOSnU SU SU 3nU$ )NLGz["z"].)r  rZ   rk   )r   r\  rk   rS  r  
local_names         rc   	enum_reprr  n
  sA     ??##D
**CCcE7"TF#cU+Jrb   c                   XR                   S'   [        R                  " 5       R                  nUc   eUR                  n[
        R                  R                  R                  R                  X15      =n(       a  X@R                   S'   g g )Nexample_valueunbacked_bindings)
metar$   r}   	fake_mode	shape_envr0  r   experimentalsymbolic_shapescompute_unbacked_bindings)rr  r  r  r  symbol_to_paths        rc   set_example_valuer  x
  s~     "/IIo""$..I   ##I 88  00JJ
	
 	

 *8		%&	
rb   c                    U R                  5       R                  R                  R                  S5      n[	        U5      (       d*  SSKJn  SSKJn  U" S[        U5      S/ UR                  QS9  U$ )Nr  r2   graph_break_hintsr\  z:Cannot check Tensor object identity without its fake valuez/TensorVariable is missing a fake example_value.r`  ra  rb  rc  )rW  rr  r  r}   rJ   r%  r  r  r]  r   
DYNAMO_BUG)vtfake_tensorr  r]  s       rc   _get_fake_tensorr  
  sb    ++-$$))--o>K;')P$I1%001		
 rb   c                f    U R                  U5      u  p#n[        SX2-
  XDS:  a  SOS-
  -   U-  5      $ )Nr   r2   rx  )indicesmax)r  seq_lenr.  r/  r0  s        rc   slice_lengthr  
  s9    		'*Eq4<4q1b#ABtKLLrb   c           	     J    SSK Jn  SSKJn  U" [        U U" SU S35      /S9  g )Nr   )raise_observed_exceptionr  zwrong number of arguments for z() callrM  )torch._dynamo.excr  torch._dynamo.variablesr  r  )txrk   r  r  s       rc   raise_args_mismatchr  
  s+    :8
!?vWMNOrb   c                D  ^ SSK JnJnJn  TR	                  5       (       a&  [        U4S jU  5       5      nUR                  U5      $ SnU(       a  [        TU5      (       a  Sn[        T5      mS n	U  H  n
U(       a5  [        X5      (       a#  T[        U
5      L a  UR                  S5      s  $ M=  M?  U" [        R                  5      R                  X*T/0 5      nU	c  Un	Mn  U" [        R                  5      R                  X+U	/0 5      n	M     U	c  UR                  S5      n	U	$ )Nr2   )BuiltinVariabler  TensorVariablec              3     >#    U  H=  nUR                  5       =(       a!    UR                  5       TR                  5       :H  v   M?     g 7frx   )r  r  )r   r^  searchs     rc   r    iter_contains.<locals>.<genexpr>
  sI      
    " F$$&&*C*C*EEFs   AAFT)r  r  r  r  r  r)  r  r  r  r  eqcall_functionor_)r   r  r  check_tensor_identityr  r  r  found_constmust_check_tensor_idfoundr^  checks    `          rc   iter_containsr  
  s    ML  "" 
 
 

  &&{33 FN!C!C#!&)'+E!,,-a00+22488 1 - $HKK0>>rv;PRSE}'5CC  } ''.Lrb   c                t    [        U [        R                  [        R                  R                  [
        45      $ )z4Returns whether it indexes dictionaries using its id)r  r0  rB  rC  r  r   )rf  s    rc   	key_is_idr  
  s%     a%,,9JKLLrb   c                    U R                  5        Vs/ s H   n[        U5      (       a  [        U5      OUPM"     sn$ s  snf rx   )re  r  r  )r   rf  s     rc   	key_to_idr  
  s1    27**,?,QYq\\BqEq(,???s   '=c                 ^ SSK Jn  [        U [        [        45      (       ah  SR                  U4S jU  5       5      n[        U [        5      (       a  SU S3$ [        U [        5      (       d   e[        U 5      S:X  a  SU S3$ SU S	3$ [        U [        R                  5      (       a  [        U TS
9R                  SS5      $ U" U 5      (       a  U R                  $ [        U [        5      (       a  SS jnU" U 5      $ U < $ )Nr2   )is_builtin_callabler  c              3  8   >#    U  H  n[        UTS 9v   M     g7fr\  N
const_reprr   r  r\  s     rc   r   const_repr.<locals>.<genexpr>
  s     D!Qj%8!   [](z,)r[  r  'r%  c                z    U R                   nUR                  nUS:X  a  UR                  $ US-   UR                  -   $ )Nbuiltins.)r  r[   r\   )r  klassr*  s      rc   fullnameconst_repr.<locals>.fullname
  s?    KKE%%F#)))C<%"4"444rb   )r  r   r   r   )trace_rulesr  r  rd  r  r   r{  rg  rh  r  replacerZ   r  )r^  r\  r  
elems_reprr  s    `   rc   r  r  
  s    0!dE]##XXD!DD
azl!$$a''''1v{:,b)):,a((	Atyy	!	! %(00b99	Q		zz	At			5 {rb   c               J   ^ SR                  U4S jU  5       5      nSU-   S-   $ )Nr  c              3  8   >#    U  H  n[        UTS 9v   M     g7fr
  r  r  s     rc   r   !dict_keys_repr.<locals>.<genexpr>  s     GJq
1E2Jr  r  r  )r   )
const_keysr\  keys_strs    ` rc   dict_keys_reprr#     s&    xxGJGGH>Crb   
__dict_key)UnsupportedFakeTensorExceptionc                P    U S[        U5       SU R                  R                   3$ )Nr  _c)r  outputr,  )r  rootr  s      rc   get_safe_global_namer*    s*     V1RWIR		 4 4566rb   c                $    U H
  nX;   d  M
    g   gr  rH   )r  
containers	containers      rc   is_inr.    s    	   rb   )requires_suffixc                   U(       d  [        U /UQ76 (       d  U $ [        R                  " 5        H  nU  SU 3n[        U/UQ76 (       a  M  Us  $    [        S5      e)za
Return a name that starts with `prefix` and is not in any of the
`containers` (e.g., map, set).
r  unreachable)r.  r   rt   AssertionError)prefixr/  r,  r  	candidates        rc   get_unique_name_wrtr5    s^     5#=*#=#=__hasO	Y,,, 
 
''rb   c           	          U " 5       $ ! [          aR  nSSKJn  SUR                   S3n[        R                  U5        U" S[        UR                  5      U/ US9   S nAg S nAff = f)Nr2   r\  zEncountered exception (z!) during fake tensor propagation.z!Fake tensor propagation exceptionr_  )r%  r  r]  reasonr   r  r   )r9  ri  r]  msgs       rc   wrap_fake_exceptionr9  .  s_    
t) 
)'z1RSC7M	

s   	 
A%AA  A%c                   ^  [         R                  R                  R                  U5         [	        U 4S j5      sS S S 5        $ ! , (       d  f       g = f)Nc                 0   > [         R                  " T 5      $ rx   )r  deepcopyr  s   rc   r_  )deepcopy_to_fake_tensor.<locals>.<lambda>C  s    4==+=rb   )r0  rF  r  FakeCopyModer9  )r  r  s   ` rc   deepcopy_to_fake_tensorr?  ?  s4     
			&	&	3	3I	>"#=> 
?	>	>s   A
Ac                    [         R                  " [         R                  " [         R                  " X-
  5      5      5      $ )z#
Calculate root mean squared error
)r0  sqrtmeansquare)r  r  s     rc   rmserD  F  s'     ::ejjci!89::rb   g-C6?Tc                L  ^ ^^^^^^^^^	^
^^ Tc  T m[        T [        [        [        R                  [
        R                  R                  [
        R                  45      (       a  [        T[        [        [        R                  45      (       d   S[        T 5       S[        T5       35       e[        T 5      [        T5      :w  a	  T	" S5        g[        T 5      [        T5      :H  =(       a(    [        UUUUUU	UUU
4	S j[        T TT5       5       5      $ [        T 5      R                  S:X  a2  [        T R                  TR                  TR                  TTTTTTT	T
TS9$ [        T [         5      (       a  [        T[         5      (       d   e[#        T R%                  5       5      [#        TR%                  5       5      :X  d;   S	[#        T R%                  5       5       S
[#        TR%                  5       5       35       e['        T R%                  5       5       H0  n[        T U   TU   TU   TTTTTTT	T
TS9(       a  M'  T	" SU5          g   g[        T ["        5      (       aO  [        T["        5      (       d   e[#        T 5      [#        T5      :X  d   S[#        T 5       S
[#        T5       35       eg[        T [
        R(                  [*        45      (       Ga  [        T [
        R,                  R.                  5      (       a   e[        T[
        R,                  R.                  5      (       a   eS&S jmU4S jT TT4 5       u  m mmT R0                  (       a3  TR0                  (       d   eT R3                  5       m TR3                  5       m[        T[
        R(                  5      (       d   S[        T 5       S[        T5       35       eT(       a  T R4                  TR4                  :w  a  T	" ST R4                  TR4                  5        gT R4                  [
        R6                  :X  ag  T(       a  g[
        R8                  " T R;                  [
        R<                  S9TR;                  [
        R<                  S9TTTS9nU(       d  T	" S5        U$ T(       a  T R?                  5       R;                  [
        R@                  5      m TR?                  5       R;                  [
        R@                  5      m[
        R8                  " T TTTSS9(       a  g[
        R                  RB                  RE                  T TSSS9nUS:  a@  [F        RI                  SURK                  5       RM                  5       RO                  5       5        [7        US:  5      $ T(       d  T R;                  TR4                  5      m [
        R8                  " T TTTTS9(       a  gTR4                  [
        RP                  :X  Ga  SnTRS                  5       RU                  5       (       d{  TRS                  5       RU                  5       (       dX  T RS                  5       RU                  5       (       a5  [
        R8                  " TR;                  TR4                  S9TUUTS9(       a  g[W        TT 5      RO                  5       n[X        RR                  " U5      (       a  [F        RI                  S5        [W        TT5      RO                  5       nS'UUU UUU
4S jjnU" 5       nUUU-  TS-  -   :*  nU(       dr  T(       ak  [X        RR                  " U5      (       aP  [X        RR                  " U5      (       a5  [
        RZ                  R\                  R^                  R`                  (       d  SnU(       d'  T	" SUUTRc                  5       TR4                  UTT
5        U$ T(       a  gT	" ST5        g[        T [d        [f        [        S5      [6        [
        Rh                  45      (       a*  T(       a  gT T:H  nU(       d  T	" S [        T 5      T T5        U$ [k        T 5      (       d  [m        T 5      (       al  T(       a0  [k        T5      (       d   [m        T5      (       d  T RO                  5       m [        T 5      [        T5      L =(       a    T T:H  nU(       d
  T	" S!T T5        U$ [o        T 5      (       aX  [        T 5      [        T5      L =(       a;    [        [
        Rp                  " T 5      [
        Rp                  " T5      TTTTTTTT	T
S"9$ [        T 5      R                  S#;   aP  [        T 5      [        T5      L d   e[        UUUUUU	U UUUU
4S$ jT Rr                  R%                  5        5       5      $ [u        S%[        T 5      R                   35      e)(z-Check correctness to see if ref and res matchNztype mismatch r   zLength mismatchFc              3  R   >	#    U  H  u  pn[        UUUTTTTT
TT	TTS 9v   M     g7f)	log_error(use_larger_multiplier_for_smaller_tensorforce_max_multiplierN)same)r   aibi	fp64_reficos_similarity	equal_nanexact_dtyperJ  ignore_non_fprH  relax_numpy_equalitytolrI  s       rc   r   same.<locals>.<genexpr>g  sJ      ,
 &=!	 $#9a%9 &=s   $'QuestionAnsweringModelOutputrG  zkeys mismatch z == )	rO  rT  rP  rQ  rS  rR  rH  rI  rJ  zAccuracy failed for key name %sTzelements mismatch c                p    [        U [        R                  5      (       a  U $ [        R                  " U 5      $ rx   )r  r0  rB  tensorr!  s    rc   	to_tensorsame.<locals>.to_tensor  s%    "1ell331HaHrb   c              3  4   >#    U  H  nT" U5      v   M     g 7frx   rH   )r   rS  rZ  s     rc   r   rU    s     M8Linn8Lr  zdtype mismatch %s, %sr  )atolrtolrP  z+Accuracy failed: uint8 tensor did not matchr   gư>)dimepsgGz?zSimilarity score=%sg{Gz?z=Found nan in reference. Consider running in higher precision.c                   > T(       a  gTR                   [        R                  [        R                  4;   a  SOSn T(       a  TR	                  5       S::  a  Sn U $ T(       a  TR	                  5       S::  a  Sn U $ TR	                  5       S:  d?  TR
                  S:X  a)  TR                  S	   TR                  S
   s=:X  a  S:X  d  O  TS:  a  Sn U $ )N      $@g      @g       @
   i  g       @r*     rx  r2   g{Gz?)r  r0  r  bfloat16numelndimr  )
multiplierrJ  fp64_refr  r  rT  rI  s    rc   get_multipliersame.<locals>.get_multiplier  s     ,#  #yyU]]ENN,KKQT  @ (B.%)
 &% B (C/%(
 &% !(4/HHMciimsyy}.QPQ.Q(?
 &)
%%rb   rb  zRMSE (res-fp64): %.5f, (ref-fp64): %.5f and shape=%s. res.dtype: %s, multiplier: %f, tol: %f, use_larger_multiplier_for_smaller_tensor: %dz+Accuracy failed: allclose not within tol=%szAccuracy failed (%s): %s != %sz!Accuracy failed (numpy): %s != %srO  rT  rP  rQ  rS  rR  rH  rI  )MaskedLMOutputSeq2SeqLMOutput!CausalLMOutputWithCrossAttentionsLongformerMaskedLMOutput	InstancesSquashedNormalBoxesNormalTanhTransformFooVariablec              3     >#    U  H6  n[        [        TU5      [        T
U5      [        TU5      TTTTT	TTTS 9v   M8     g7f)rm  N)rK  r&  )r   r   rO  rP  rQ  rj  rR  rH  r  rS  r  rT  rI  s     rc   r   rU  i  sW      
 + S!S!#&-#'%9+#9a +s   >Azunsupported type: )r!  r   r   r  )r   float);r  rd  r  r   rC  r0  rC  ParameterListri  r  r{  rX  r  rZ   rK  lossr  rt  re  r  rB  rz  rF  rI   	is_sparseto_denser  r   allclosetor  flattenr  
functionalcosine_similarityr   r  detachcpur  r  isnanr)  rD  mathr  rF   cppinject_relu_bug_TESTING_ONLYr  r   r   r  r  r  r@  	as_tensor__dict__r   )r  r  rj  rO  rT  rP  rQ  rS  rR  rH  rI  rJ  rf  rscore	loose_tol	ref_error	res_errorrk  ri  passes_testrZ  s   ````````````         @rc   rK  rK  M  s    dE;,,ehh.D.DejjQ  #e[->->?@@ 	
T#YKqc4	
@ s8s3x'(3x3s8# 
 ,
 ,
 &)c8%<,
 )
 	
" 
c		=	= HHHHMM 5]!5
 	
 
C		#t$$$$388:#chhj/1 	
S_-T#chhj/1BC	
1 
#AFFQK#1' +)="/'=e)=  ;Q?% $& 	C		#s####3x3s8#R'9#c(4Cz%RR#	C%,,.	/	/c5#4#4#?#?@@@@c5#4#4#?#?@@@@	I Nc88LMS(==== =,,.C,,.C#u||,,VtCyk4PS9+.VV,yyCII%1399ciiHyyEJJ& NNFFF-FFF-' KL++-""5==1C++-""5==1C~~c3SsdK HH''99#st9TEt|15<<>3E3E3G3L3L3NO&&ffSYY' ~~c3SsiP ~~. %	 (,,..IIKOO--		)) #))4&&"+   3/446	::i((KKW !3/446	!& !&F ,-
'J,BS4Z,OP#!

9--

9-- "OO2266SS"&K"I!!
		"@
 #"CSI	C#sDJellC	D	D3J6S	3L	3		#6s#;#;c""&9#&>&>((*C#Y$s)#5#*93D	#		S	T#Y& 
DOOC OOC )#!5'5]-
 	
 
c		   
 CyDI%%% 
 
 ||((*
 
 	
" /S	0B0B/CDEErb   c                ~    U R                   R                  S5      S   nSU R                   SU SU R                   S3$ )N/rx  r  z' (r   r[  )rv  splitr  r  )r  short_filenames     rc   format_func_infor  }  sD    %%++C04Nt||nC/q1D1D0EQGGrb   c               #    #    [         R                  n [        R                  [         l        [         R                  n[        R                  [         l         S v   U [         l        U[         l        g ! U [         l        U[         l        f = f7frx   )rF   recompile_limitr  rn  accumulated_recompile_limit)priorprior_acc_limits     rc   disable_cache_limitr    s^     ""E [[F88O),F&=!&-<* "'-<*s   ABA1 B1B		Bz'collections.defaultdict[Any, list[Any]]guard_failuresz3list[torch._dynamo.output_graph.GraphCompileReason]graph_break_reasonsc                    S[         R                   R                  5       R                  S5      -   S-   [        [        R
                  " 5       5      -   n[        R                  R                  X5      $ )Nrun_z%Y_%m_%d_%H_%M_%S_%fz-pid_)datetimenowstrftimer   r  getpidr  r   )root_dirdir_names     rc   _get_debug_dirr    sd     	





!
*
*+A
B	C 	 biik
		  77<<++rb   c                 8    [         R                  n [        U 5      $ rx   )rF   r  r  )
debug_roots    rc   r  r    s    &&J*%%rb   c                    SU R                   ;   a,  [        U R                   S   5      (       a  U R                   S   $ U(       a/  SSKJn  SSKJn  U" S[        U 5      SU  S3/ UR                  QS	9  g g )
Nr  r   r\  r2   r  z Missing FakeTensor example valuez,`FakeTensor` example value was required for z but not available.r  )r  rJ   r  r]  r%  r  r   r  )rr  requiredr]  r  s       rc   extract_fake_example_valuer    sl    $))#		/0J(K(Kyy))	6'6IFtfL_`1%001		
 rb   c                :    [        U 5      UR                  L d   eU $ rx   )rK   r  )ri  r  s     rc   ensure_graph_faker    s    q!R\\111Hrb   c                j   ^ ^ SUU 4S jjn[         R                  R                  R                  X5      $ )Nc                  > U R                   S:X  a  SU R                  ;  a  [        U TT5      $ U R                   S:X  a  SU R                  ;  av  U R                  TR                  R
                  ;   d   eTR                  R
                  U R                     n[        U[        R                  R                  5      (       d   eU$ U R                  S   nT(       d+  [        U[        R                  5      (       a  [        UT5      $ U$ )Nr  r  get_attr)r  r  get_fake_valuetargetr(  
nn_modulesr  r0  r   GraphModulerB  r  )r  rF  r   allow_non_graph_faker  s      rc   visit)get_fake_values_from_nodes.<locals>.visit  s    44?"aff'D "!R)=>>TTZO166$A88ryy33333%%ahh/Bb%(("6"67777Iff_%#
3(E(E$S"--
rb   )r  torch.fx.Noder   r   )r0  r   rr  map_arg)r  r  r  r  s   ` ` rc   get_fake_values_from_nodesr    s(     " 88==  ..rb   c           	       ^ ^^^^ SSK Jn  SSKJnJnJnJnJn  T R                  n	ST R                  ;   a,  [        T R                  S   5      (       a  T R                  S   $ [        TT R                  T R                  4U5      u  mm[        R                  R                   R"                  (       d)  [        R                  R                   R$                  (       aP  [        T['        T 0 5      U5      n
U
 Vs0 s H+  n[        U5      (       d  M  [)        U5      UR*                  _M-     nnO/ n
0 nSmTR,                  nUc   eU	S:X  aZ  [/        T5      S:  aK  [1        TS   [        R2                  R4                  5      (       a  [7        TS   U5      4[9        TSS 5      -   mU	S:X  ab  TR:                  R<                  T R>                     m[A        T5      (       a#  [C        TS	5      (       a  TRE                  TT5        [7        TU5      mT RF                  S
;   d  [I        S T 5       5      (       a  [9        S T 5       5      m U   [K        5          [M        UUUU U4S j5      nSSS5        SSS5        U(       d=  [        R                  " [        R                  [        R                  " [        TS(9W5      n[        R                  R                   R"                  (       d)  [        R                  R                   R$                  (       a'  TR:                  R                  R                  T U
U5        W$ s  snf ! , (       d  f       N= f! , (       d  f       N= f! U a    e [N         Gay  nUnURP                  b  URP                  n[1        U[        RR                  RT                  RV                  5      (       a  URX                  [        RZ                  R\                  R^                  R`                  [        RZ                  R\                  Rb                  R`                  4;   a  S/nOS/nU" S[e        URX                  5      SURX                   S3US9  GO^[1        U[        RR                  RT                  Rf                  5      (       a  [        R                  R                   Rh                  (       d-  U" S[e        URX                  5      SURX                   S3S/S9  GOU" S[e        URX                  5      SURX                   S3S/S9  GO[1        U[        RR                  RT                  Rj                  5      (       a  URX                  n	Sn[1        U	[        Rl                  Rn                  5      (       a\  [        Rp                  Rs                  U	Rt                  RF                  U	Rt                  Rv                  5      nUb  Uu  nnSU SU SU S3nU" SS URX                   3SU S!3/S9  O[1        U[        Rx                  Rz                  R|                  R~                  5      (       a  U" UR                  [e        U5      S"S#9e[1        UU5      (       a!  U" UR                  UR                  S   5      Ue[1        U[        5      (       a(  S$[e        U5      ;   a  U" S%S&T R>                   S'U 3S/ S9  U" [e        U5      5      R                  UR                  5      SeSnAff = f))ab  
Run the computation represented by `node` using fake tensors and return the result.

allow_non_graph_fake: whether to allow the return result to be:
    1. non-fake or 2. fake that is not created by this instance of Dynamo.
    If `True`, you must be prepared to deal with such return values, ideally
    by further wrapping them as this graph's fakes.
r   )ValueRangeErrorr2   )TorchRuntimeErrorr]  Unsupported	UserErrorUserErrorTyper  Ncall_methodcall_module_initialize_hook)interpolate
is_integerwrapped_gradientc              3  B   #    U  H  n[        U[        5      v   M     g 7frx   )r  complex)r   as     rc   r   !get_fake_value.<locals>.<genexpr>!  s      M(,1
1gs   c              3     #    U  HH  n[        U[        R                  5      (       a"  UR                  R                  b  [        U5      OUv   MJ     g 7frx   )r  r0  SymFloatrr  hintrz  rX  s     rc   r   r  %  sG      
  c5>>22sxx}}7P c
 s   AAc                 6   > [        TR                  TT TT5      $ rx   )run_noder(  )r  r  nnmodulerr  r  s   rc   r_   get_fake_value.<locals>.<lambda>1  s    D$Irb   zkEnable tracing of data-dependent output operators with `torch._dynamo.config.capture_scalar_outputs = True`zConsider wrapping the operator into a PyTorch-understood custom operator (see https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html)zData dependent operatorz
Operator `zP` has a non-Tensor output whose value is dependent on the data of Tensor inputs.r  zDynamic shape operatorz.`'s output shape depends on input Tensor data.zmEnable tracing of dynamic shape operators with `torch._dynamo.config.capture_dynamic_output_shape_ops = True`z'Dynamic shape operator (no meta kernel)zA` does not have a meta kernel that supports dynamic output shapesz!Please report an issue to PyTorchr%  z:It's possible that the support was implemented in module `z` and you may need to `import z`(z), otherwise z3Operator does not support running with fake tensorszunsupported operator: z}see https://docs.google.com/document/d/1GgvOe7C8_NVOMLOCwDaYV1mXXyHMXY7ExoewHqooxrs/edit#heading=h.64r4npvq0w0 for how to fixconstrain_as_size_example)	case_nameargumentz&TypeError when making fake tensor callz
TypeError z: )r  )Ltorch.utils._sympy.value_rangesr  r  r  r]  r  r  r  r  r  rJ   r  r  r  r0  r  rF   use_graph_deduplicationtrack_nodes_for_deduplicationr3   r  r>  r  r{  r  rC  r  r?  r  r(  r  r  rK  r8  _infer_parametersrk   r)  r   r9  r   	__cause__rF  r  DataDependentOutputExceptionfuncopsatenr  r  _local_scalar_denser   DynamicOutputShapeException capture_dynamic_output_shape_opsUnsupportedOperatorExceptionr!  r#  r  _dispatch_pystub_schemaoverload_namer   r  r  GuardOnDataDependentSymNodeCONSTRAINT_VIOLATIONr  with_traceback__traceback__pytreetree_map_onlyrB  r  partialr  region_trackertrack_node_mutations)rr  r  r  r  r  r]  r  r  r  r  flat_args_kwargsrY  id_to_initial_versionr  ret_valri  causerc  import_suggestionmaybe_pystubr*  ctxr  r  r  r  s   ``                     @@@rc   r  r    s    @  
B $))#		/0J(K(Kyy))-
TYY$&:LD&
 	44====5tR(*>
 .>!
-=c!BsGS\\!-= 	 !
  "HI   	]s4y1}DGUXX__1U1U'Q;=d12hO	]99''4(##:L(M(M
 &&x6 +8Y?yyEE M(,M J J  
 
 
hR02)IIG 3YR    LL)++,="Ew

 	44====
		  55!	
 NC!
R 32YY   aR ;;"KKE5$$00MM
 

 zz		##++		22:: K` 1EJJ( 5I I 5$$00LL
 
 ==''HH 4

O",UZZL8f gY	 !E

O",UZZL8y z;	 5$$00MM
 
 B ""ejj3344$xx88JJOORZZ%=%=   +".KFC##)(*H Q5/ &
 M0=() *& &		 588((88TT
 
 22E
5 
 //M>>q	JPQQy))jCJ.F@$T[[ME7;	  A'66qGTQCaRs[   5L3L3M M	L8$M	,M 8
M	M		
MM M [$+M4[[$c                 $    [        [        SS 5      $ )Nr   )r&  _current_noderH   rb   rc   get_current_noder    s    ='400rb   c              #  v   #    [        5       nU [        l         S v   U[        l        g ! U[        l        f = f7frx   )r  r  r   )rr  olds     rc   set_current_noder    s-     

CM"!cs   9) 969c           
       ^^^^
 TR                   m
[        T5         SUUUU
4S jjnSSKJn   T
S:X  a  TR                  " T0 TD6sSSS5        $ T
S:X  a^  [        TS   TR                  5      (       d  SSKJn  U" S	S
U" S5      / S9  [        TS   TR                  5      " TSS 0 TD6sSSS5        $ T
S:X  a  Uc   eU" T0 TD6sSSS5        $ T
S:X  a.  U R                  R                  TR                  5      sSSS5        $ T
S:X  a*  STR                  ;   d   eTR                  S   sSSS5        $  SSS5        [%        T
5      e! [        [        4 a:  nSSKJn  / n	[        U[        5      (       a  S/n	U" SS
U" U5      U	US9   SnANXSnAfU a    e [         a0  n[        U" U5      5      R!                  UR"                  5      UeSnAff = f! , (       d  f       N= f)a  
Runs a given node, with the given args and kwargs.

Behavior is dictated by a node's op.

run_node is useful for extracting real values out of nodes.
See get_real_value for more info on common usage.

Note: The tracer arg is only used for 'get_attr' ops
Note: The nnmodule arg is only used for 'call_module' ops

Nodes that are not call_function, call_method, call_module, or get_attr will
raise an AssertionError.
c           	     N   > ST STR                    ST ST S3	[        U 5      -   $ )Nz0Dynamo failed to run FX node with fake tensors: r   z(*z, **z): got )r  rR  )ri  r  r  rr  r  s    rc   make_error_message$run_node.<locals>.make_error_message  s=    B2$a}TVW[V\\`ag`hhopq'rb   r2   )r  r  Nr  r   r\  z/Missing attribute when running call_method noder%  zattribute not definedr  r  r  placeholderr  z;If the op is a PyTorch op, please file an issue to PyTorch.zGNotImplementedError/UnsupportedFakeTensorException when running FX noder_  )ri  r   r   r   )r  r  r  r  r  r8  r]  r&  output_graphget_submoduler  re  r%  r  r  r   r  r  r2  )tracerrr  r  r  r  r  r  r]  ri  rc  r  s    ```      @rc   r  r    s   " 
B	$		 	 	%-	_${{D3F3 
 	 }$tAw445$ Q "$67N$O 	 tAw4d12hI&I/ 
 	0 }$+++005 
 	6 z!**88E9 
 	: }$&$))333yy1? 
 	: %; 
 r 
1 $%CD 	-E!011Q a.q1  	 	1!45DD	i 
 	sY   GEAE4E*E&E4GG0FGG +GGG
G!c                  ^ SSK Jn  TR                  nX;   a  X0   $ U R                  n[        R
                  R                  R                  U R                  U R                  4U4S j5      u  pVUS:X  a)  SU R                  ;   a  U R                  S   R                  $ US:X  aS  TR                  R                  U R                     n[        U5      (       d  [         R"                  " U5      nOU" U0 UD6  OSn [%        TXXg5      nXU '   U$ ! [&         a0  n	U" [)        U	5      5      R+                  U	R,                  5      SeSn	A	ff = f)z
Run the actual computation represented by `node` and return the result.
This will execute any dependent nodes in the graph as well.
r2   )r  c                   > [        U T5      $ rx   )get_real_value)r  r  s    rc   r_   get_real_value.<locals>.<lambda>  s    .F+rb   r   graphargr  N)r  r  real_value_cacher  r0  r   rr  r  r  r  r  exampler  r  r  rK  r  r<  r  r   r   r  r  )
rr  r  r  cacher  r  r  	nn_module
real_valueri  s
    `        rc   r  r  
  s&   
 '##E}{	B88==((	DKK +LD
 
]zTYY6yy$,,,	]''224;;?	i((i0I t&v&	Rfd&D
 d   RA'66qGTQRs   :D 
E+EEc                
  ^ SSK JmJn  SU4S jjnU R                  5        H%  u  p4U" U5      (       d  M   SU SU" U5       35       e   U R	                  5        H%  u  p5U" U5      (       d  M   SU SU" U5       35       e   g )Nr   )FakeTensorConfigrJ   c                n   > TR                   (       a#  SS KnSUR                  " U R                  5       3$ g)Nr   z"FAKE TENSOR CREATION TRACEBACK: 
 zNEnable TORCH_FAKE_TENSOR_DEBUG=1 to get creation stack traces on fake tensors.)debug	tracebackformat_list_debug_trace)r!  r  r  s     rc   stack_or_hint7assert_no_fake_params_or_buffers.<locals>.stack_or_hint4  s/    !!89N9Nq~~9^8_``crb   zUnexpected fake buffer r   zUnexpected fake param )r!  r   r   r   )torch._subclasses.fake_tensorr  rJ   named_buffersnamed_parameters)rF  rJ   r  rk   bufferr?  r  s         @rc    assert_no_fake_params_or_buffersr  1  s    Gd ((*6?? 	
%dV1]6-B,CD	
" + **,5>> 	
$TF!M%,@+AB	
! -rb   c                8    U R                    SU R                   3$ )z1
Returns the fully qualified name of the object.
r  )r[   r\   r  s    rc   fqnr  F  s!     nnQs//011rb   c                \    [         R                  R                  R                  (       a  U $ U$ rx   )r0  r  rF   assume_static_by_default)count1count2s     rc   ifdynstaticdefaultr"  M  s    }}44rb   c           
     P   [        [        R                  " [        R                  R	                  [        [        U R                  5      5      5      5       HN  nUR                  S5      (       d  M  US   S:w  d  M&  [        R                  " U R                   SUSS  35        MP     g)z8
Ensure all the files in a given submodule are imported
z.pyr   r  r  N)r  r  listdirr  r  r   r   __file__r7  	importlibimport_modulerZ   )rJ  r  s     rc   import_submoduler)  T  st     2::bggood36M&NOPU##s(:##s||nAhsm_$EF Qrb   c                *    [        [        U 5      5      $ rx   )class_has_getattributer  r  s    rc   object_has_getattributer,  ]  s    !$u+..rb   c                8    [         R                  U S5      nX#U'   g )Nr  )r!  __getattribute__)r  rk   r   r  s       rc    object_setattr_ignore_descriptorr/  a  s    Z0AdGrb   c                     [        [        R                  " U S5      [        R                  5      (       a  g g! [
         a     gf = f)Nr.  TF)r  r  r  r&  r'  r  rz   s    rc   r+  r+  g  sT    ""3(:;
 
 	
   s   49 
AAc                     [         R                  " [        U 5      S5      nU(       a)  U[        R
                  R                  R                  L a  S nU$ ! [         a    S n N@f = f)N__getattr__)r  r  r  r  r0  rC  r  r2  )r   ignore_nn_module_getattr
getattr_fns      rc   get_custom_getattrr5  s  sX    ++DKG
  J%((//2M2M$M
  
s    A A#"A#c                       \ rS rSrSrSrSrSrg)TensorStaticReasoni  rX   rd  r   rH   N)rZ   r[   r\   r]   	PARAMETER
NOT_TENSORNN_MODULE_PROPERTYra   rH   rb   rc   r7  r7    s    IJrb   r7  c                    U [         R                  :X  a  gU [         R                  :X  a  gU [         R                  :X  a  g[	        SU  35      e)Nz>mark_dynamic on parameter, parameters are always static today.z2mark_dynamic on a non tensor, how did this happen?z4tensor is static because it is nn module associated.zIllegal reason )r7  r8  r9  r:  r2  )r7  s    rc   tensor_static_reason_to_messager<    sJ    #---O#...C#666E
?6(3
44rb   c                   SSK Jn  UR                  5       R                  5       (       d#  UR                  5       R	                  5       (       a'  [
        R                  (       a  S[        R                  4$ [        U 5      [        R                  R                  L d  U" U5      (       a'  [
        R                  (       a  S[        R                  4$ U(       d  S[        R                  4$ g)a   
Given a tensor, source, and is_tensor flag, determine if a shape should be static.

Args:
tensor - the real tensor to evaluate, parameters force a static shape.
is_tensor - internal dynamo check, essentially "is_tensor": target_cls is TensorVariable,
tensors not in a TensorVariable for whatever reason are forced static.

Returns a tuple, where the first element is the bool of whether or not this tensor should have a static shape.
The second element is a TensorStaticReason, useful for passing to tensor_static_reason_to_message if needed.
r2   ))is_from_unspecialized_param_buffer_sourceTr  )ry  r>  r{  r}  r  rF   &force_nn_module_property_static_shapesr7  r:  r  r0  rC  rD  force_parameter_static_shapesr8  r9  )rX  	is_tensortensor_sourcer>  s       rc   tensor_always_has_static_shaperC    s      B 	""$==??%%'JJLL

7
7':::: 	V***4]CC

.
.'1111'2222rb   c                .   ^ ^ SU U4S jjn[        U5      $ )Nc                   >  SSK J n   TR                  R
                   Vs/ s H;  oR                  UR                  UR                  UR                  UR                  /PM=     nnU " U/ SQS9n[        TTR                  R                  R                  U5      $ ! [         a    S[        [        TT5      5      -   s $ f = fs  snf )Nr   )r   zkTabulate module missing, please install tabulate to log the graph in tabular format, logging code instead:
)opcoderk   r  r  r  r   )r   r   r   r,   graphr  r  rk   r  r  r  r+   forward__code__rv  )r   r  
node_specs	graph_strr&  rF  s       rc   inner(lazy_format_graph_tabular.<locals>.inner  s    	) AC
@N1TT166188QVVQXX6 	 
  N
	 "'2::+>+>+J+JIVV  	~,Wb9:;	
s   B AC"B>=B>r   r   rC   )r&  rF  rL  s   `` rc   lazy_format_graph_tabularrO    s    W W" erb   c                h    U  SU SU SU S[         R                  " U5      R                  5        S3
$ )Nr   z line z 
r   )disBytecode)r3  rk   r  line_nor  s        rc   format_bytecoderT    s=     XQtfAhZvgYc#,,t:L:P:P:R9SSUVVrb   _forward_pre_hooks_forward_hooks_backward_pre_hooks_backward_hooks)_state_dict_pre_hooks_state_dict_hooks_load_state_dict_pre_hooks_load_state_dict_post_hooksc                     [        [        [        R                  R                  R
                  R                  5      =(       d7    [        [        R                  R                  R
                  R                  5      5      $ rx   )r   r{  r0  rC  modulesr*  _global_backward_hooks_global_backward_pre_hooksrH   rb   rc   nn_module_has_global_hooksra    sS     EHH##::; 	Cuxx&&AAB rb   c                ~   / nU(       + =(       a    U(       + =(       a    U(       + nU(       d  U(       a  UR                  [        5        U(       d  U(       a  UR                  [        5        U(       a  UR                  [        5        / nU H-  n[	        X/ 5      nU H  n	X   n
UR                  U
5        M     M/     U$ )z
Sometimes its useful to differentiate between types of hooks such as forward/backward/pre
hooks executed during module.__call__, and state_dict hooks which are executed separately.
)extendforward_hook_namesbackward_hook_namesstate_dict_hook_namesr&  r   )rJ  check_forward_hookscheck_backward_hookscheck_state_dict_hookshook_dicts_to_checkcheck_all_hooks	all_hookshook_dict_namehooks	hook_namer   s              rc   nn_module_get_all_hooksrp    s      	'$$	'&& 
 o""#56""#67""#89I-R0I#DT"  . rb   c                0    [        U UUUS9n[        U5      $ )zD
Helper function to check if a module has any hooks attached to it.
)rg  rh  ri  )rp  r   )rJ  rg  rh  ri  rn  s        rc   nnmodule_has_hooksrr    s&     $/15	E ;rb   c                T   [        U 5      (       a  U $ [        U [        R                  5      (       a  [	        U R
                  5      $ [        U [        R                  5      (       a  U R                  SS9$ [        U [        [        45      (       a  [        U 5      " S U  5       5      $ U $ )z0Convert tensor and tnp.ndarray to numpy.ndarray.T)forcec              3  8   #    U  H  n[        U5      v   M     g 7frx   )to_numpy_helperr   r  s     rc   r   "to_numpy_helper.<locals>.<genexpr>       A5C?3//5rW  )rJ   r  tnpr?  rv  rX  r0  rB  numpyr  rd  r  r  s    rc   rv  rv    s}    u~~%%%u||,,	E5<<	(	({{{&&	EE4=	)	)E{A5AAArb   c                >   [         c   e[        U [         R                  5      (       a  [        R                  " U 5      $ [        U [
        R                  5      (       a  U R                  $ [        U [        [        45      (       a  [        U 5      " S U  5       5      $ U $ )zeConvert tnp.ndarray to tensor, leave other types intact. If a list/tuple, loop through it to convert.c              3  8   #    U  H  n[        U5      v   M     g 7frx   )numpy_to_tensorrw  s     rc   r   "numpy_to_tensor.<locals>.<genexpr>*  ry  rW  )
r  r  r?  r0  r  rz  rX  r  rd  r  r  s    rc   r~  r~  "  sr    >>%$$u%%%%%||	EE4=	)	)E{A5AAArb   c                  2    \ rS rSrSS jrSS jrS	S jrSrg)
numpy_to_tensor_wrapperi/  c                L    Xl         SU R                   R                  -   U l        g Nwrapped_r  rZ   )rp  r  s     rc   rq   numpy_to_tensor_wrapper.__init__0  s    "TVV__4rb   c                6    SU R                   R                   S3$ )Nz<Wrapped function <original >>r  ru  s    rc   __repr__ numpy_to_tensor_wrapper.__repr__4  s    -dffoo->bAArb   c                <    U R                   " U0 UD6n[        U5      $ rx   )r  r~  rp  r  r  r   s       rc   r   numpy_to_tensor_wrapper.__call__7  s     ffd%f%s##rb   )rZ   r  N)r  Callable[_P, R]r   r   rN  r  r"  r  r#  r   r   )rZ   r[   r\   r]   rq  r  r  ra   rH   rb   rc   r  r  /  s    5B$rb   r  c                   [        U [        R                  5      (       a  [        X5      n[	        U5      $ [        U [
        R                  5      (       a+  [        [        R                  " U 5      U5      n[	        U5      $ g rx   )r  rz  r?  r&  r~  r0  rB  )r  rk   r   s      rc   numpy_attr_wrapperr  <  s^    #s{{##c s##	C	&	&ckk#&-s## 
'rb   c                  6    \ rS rSrSrSS jrS	S jrS
S jrSrg)numpy_method_wrapperiE  zgConvert obj from torch.Tensor to tnp.ndarray and call method. Then convert result back to torch.Tensor.c                8    Xl         SU R                   -   U l        g r  )methodrZ   )rp  r  s     rc   rq  numpy_method_wrapper.__init__H  s    "T[[0rb   c                "    SU R                    S3$ )Nz<Wrapped method <original r  r  ru  s    rc   r  numpy_method_wrapper.__repr__L  s    +DKK=;;rb   c                    US   n[        U[        R                  5      (       a  [        R                  " U5      n[        X0R                  5      nU" USS  0 UD6n[        U5      $ rm  )r  r0  rB  rz  r?  r&  r  r~  )rp  r  r  r  method_callabler   s         rc   r  numpy_method_wrapper.__call__O  sX    1gc5<<((++c"C!#{{3tABx262s##rb   )rZ   r  N)r  r   r   r   rN  )r  r   r  r   r   r   	rZ   r[   r\   r]   r   rq  r  r  ra   rH   rb   rc   r  r  E  s    q1<$rb   r  c                  6    \ rS rSrSrSS jrS	S jrS
S jrSrg)numpy_operator_wrapperiX  zQImplements dunder methods for tnp.ndarray via functions from the operator libraryc                8    Xl         SUR                   3U l        g r  )r  rZ   )rp  r  s     rc   rq  numpy_operator_wrapper.__init__[  s    "2;;-0rb   c                "    SU R                    S3$ )Nz<Wrapped operator <original r  )rZ   ru  s    rc   r  numpy_operator_wrapper.__repr___  s    -dmm_B??rb   c                Z    U(       a   eS U 5       nU R                   " U6 n[        U5      $ )Nc              3     #    U  H<  n[        U[        R                  5      (       a  [        R                  " U5      OUv   M>     g 7frx   )r  r0  rB  rz  r?  rX  s     rc   r   2numpy_operator_wrapper.__call__.<locals>.<genexpr>e  s2      
RV3
3 = =CKK3FRVs   AA)r  r~  r  s       rc   r  numpy_operator_wrapper.__call__b  s3    z
RV
 ggtns##rb   )rZ   r  N)r  Callable[..., Any]r   r   rN  r  r  rH   rb   rc   r  r  X  s    [1@$rb   r  c                r   [        U [        5      (       d  U $ U R                  (       Ga'  / nU R                  5        H}  n[        U[        R
                  5      (       aJ  UR                  UR                  R                  R                  UR                  R                  5      5        Ml  UR                  U5        M     / nU R                  5        H}  n[        U[        R
                  5      (       aJ  UR                  UR                  R                  R                  UR                  R                  5      5        Ml  UR                  U5        M     O U R                  5       nU R                  5       n[        R                  " UUU R                  U R                  U R                  S9nUR!                  5         U$ )N)r  r  r  )r  rI   _has_symbolic_sizes_stridesr  r0  rq  r   rr  r  	size_hintexprr  empty_stridedr  r  r  zero_)r^  r  r  r  r  s        rc   defaker  l  s,   a$$ 	$$$A!U\\**AFF,,66qvv{{CDA	 
 A!U\\**aff..88EFa 	  vvxggxxoo	A GGIHrb   c                    U " U0 UD6$ rx   rH   )r9  r  r  s      rc   8_disable_side_effect_safety_checks_for_current_subtracerr    s     tvrb   c                L    SS K nXR                  R                  R                  L $ r   )torch.utils.checkpointr.  
checkpoint)r  r0  s     rc   is_utils_checkpointr    s    !++((3333rb   c                    SSK Jn  XL $ )Nr   )invoke_subgraph_placeholder)'torch._higher_order_ops.invoke_subgraphr  )r  r  s     rc   is_invoke_subgraphr    s    S--rb   c                 d    SSK Jn  UR                  " [        R                  R
                  40 U D6$ )Nr2    TorchHigherOrderOperatorVariable)variables.higher_order_opsr  maker0  _higher_order_opsinvoke_subgraph)optionsr  s     rc   build_invoke_subgraph_variabler    s/    L+00//
 rb   c                     SS K Js  Jn  SSKJn  UR
                  n[        R                  R                  R                  (       a  UR                  nUR                  " U40 U D6$ )Nr   r2   r  )torch._higher_order_ops.wrapr  wrapr  r  tag_activation_checkpointr0  
_functorchrF   functionalize_rng_opswrap_activation_checkpointr  )r  higher_order_opsr  activation_checkpoint_ops       rc   build_checkpoint_variabler    s\    ;;L
 	22  44#3#N#N +00 
 rb   c                    SSK Jn  [        R                  " U 5      R                  nU" 5       nUS:X  a   U$ US;   a  U(       a  [        5       nU$ SnU$ )Nr2   )is_dynamo_supportedr  )r6  xpumtiaF)
eval_framer  r0  r  r  r/   )device_typer  r  compile_supporteds       rc   is_compile_supportedr    s`    /<<$))D+-u}
 	 
(	(->&L  "rb   c                \    U R                  S5      n[        USU R                  SSS95      $ )z
Convert byte offset `offset` of `str` into character offset.
Byte offset is used for 3.11+ instruction column data.
Takes things like unicode characters into consideration.

Unchanged from CPython implementation.
zutf-8Nr  )errors)encoder{  decode)r   offsetas_utf8s      rc   _fix_offsetr    s4     jj!Gww&&wy&ABBrb   c                  >    \ rS rSr% S\S'   S\S'   S\S'   S\S'   Srg)	_Anchorsi  r   left_end_linenoleft_end_offsetright_start_linenoright_start_offsetrH   Nr4  rH   rb   rc   r  r    s     rb   r  c                  ^^ [         R                  S:  d   eSSKn UR                  SU -   S-   5      n[        UR                  5      S:w  a  gU R                  S5      mSU4S jjnSU4S	 jjmSUU4S
 jjnSUU4S jjnUR                  S   n[        XaR                  5      (       Ga  UR                  n[        XqR                  5      (       Ga  [        [        UR                  R                  5      S-
  nUR                  R                   c   eU" XR                  R                   5      n	T" X5      u  pTU   U	   =n
R#                  5       (       d  U
S;   aB  U
S;   a  U" X5      u  pO
U" X5      u  pTU   U	   =n
R#                  5       (       a  M:  U
S;   a  MB  U	S-   nU[        TU   5      :  a(  TU   U   =n
R#                  5       (       d  U
S;  a  US-  n[%        XX5      $ [        XqR&                  5      (       a  [        [        UR                  R                  5      S-
  nUR                  R                   c   eU" XR                  R                   5      nT" X5      u  pTU   U   S:w  a  U" X5      u  pTU   U   S:w  a  M  [        [        UR                  5      S-
  nUR                   c   eU" XR                   5      n[%        XX5      $ [        XqR(                  5      (       a  [        [        UR*                  R                  5      S-
  nUR*                  R                   c   eU" XR*                  R                   5      nT" X5      u  pTU   U   S:w  a  U" X5      u  pTU   U   S:w  a  M  [        [        UR                  5      S-
  nUR                   c   eU" XR                   5      n[%        XX5      $ g! [         a     gf = f)z
Given source code `segment` corresponding to a bytecode
instruction, determine:
    - for binary ops, the location of the binary op
    - for indexing, the location of the brackets.
`segment` is expected to be a valid Python expression
)rY      r   Nz(
z
)r2   r   c                "   > [        TU    U5      $ rx   )r  )linenor  liness     rc   	normalize-_extract_anchors_from_expr.<locals>.normalize  s    5=&11rb   c                   > U [        T5      :  a<  U[        TU    5      :  a*  SnU S-  n U [        T5      :  a  U[        TU    5      :  a  M*  U [        T5      :  a  U[        TU    5      :  d   eX4$ rm  r{  )r  colr  s     rc   next_valid_char3_extract_anchors_from_expr.<locals>.next_valid_char  sw    s5z!cSv-?&?CaKF s5z!cSv-?&? E
"sSv-?'???{rb   c                n   > US-  nT" X5      u  pU [        T5      :  a  U[        TU    5      :  d   eX4$ r   r  r  r  r  r  s     rc   r   -_extract_anchors_from_expr.<locals>.increment
  sB    q%f2E
"sSv-?'???{rb   c                r   > SnU S-  n T" X5      u  pU [        T5      :  a  U[        TU    5      :  d   eX4$ rm  r  r  s     rc   nextline,_extract_anchors_from_expr.<locals>.nextline  sG    !%f2E
"sSv-?'???{rb   rX   z)\#z\#r  r  )r  r   r  r   r   r   )r  r   r  r   r   ztuple[int, int])r  r  astparseSyntaxErrorr{  bodyr  r  Exprr   BinOpr   r   left
end_linenoend_col_offsetisspacer  	SubscriptCallr  )segmentr  treer  r   r  	statementr  
cur_linenocur_colch	right_colleft_linenoleft_colright_linenor  r  s                  @@rc   _extract_anchors_from_exprr
    s    w&&& yy501 499~MM$E2
   		!I)XX&&dII&& c499#7#781<J99++777
II,D,DEG"1*"FJ z*7332<<>>",;*2:*G'J*3J*H'J	 z*7332<<>>",  !ICj 122$Z0;;DDFFeOQ	 JGGmm,, sDJJ$9$9:Q>K::,,888 jj.G.GHH$3K$J!K$X.#5(1+(H% $X.#5  T__59L&&222!,0C0CDIK<KKhh''
 sDII$8$89A=K99++777 ii.F.FGH$3K$J!K$X.#5(1+(H% $X.#5  T__59L&&222!,0C0CDIK<KKU  s   O 
OOc           	     d   [         R                  S:  Ga  [        R                  " U R                  UR
                  R                  U R                  UR
                  R                  UR
                  R                  UR
                  R                  S9n[        R                  " U/5      S   nSR                  UR                  5       SS 5      n[        UR
                  R                  UR
                  R                  S-   5       Vs/ s H2  n[        R                   " U R                  U5      R#                  5       PM4     nn[$        R&                  " SR                  U5      5      R                  5       n[)        US   5      [)        US   5      -
  nUS   SU n[$        R*                  " [$        R&                  " U5      U5      nU$ UR
                  c   eUR
                  R                  c  g[        R                   " U R                  UR
                  R                  5      R#                  5       n	UR
                  R                  c  U	$ UR
                  R                  b  UR
                  R                  c  U	$ [-        XR
                  R                  5      n
SnSn/ nUR
                  R                  UR
                  R                  :X  aA  [-        XR
                  R                  5      nXU nUR/                  SU
-  S	X-
  -  -   5        GOXS S-   nUR/                  SU
-  S	[)        U	5      U
-
  -  -   5        [        R                   " U R                  UR
                  R                  5      R#                  5       n[-        XR
                  R                  5      n[        UR
                  R                  S-   UR
                  R                  5       H  n[        R                   " U R                  U5      R#                  5       nXS-   -  n[)        U5      [)        UR1                  5       5      -
  nUR/                  SU-  S	[)        U5      U-
  -  -   5        M     XSU -  n[)        U5      [)        UR1                  5       5      -
  nUR/                  SU-  S	UU-
  -  -   5        Sn [3        U5      nUc$  U Vs/ s H  nUR7                  S	S
5      PM     nnGO=U Vs/ s H  n[9        U5      PM     nnUR:                  S:X  a  U=R<                  U
-  sl        UR>                  S:X  a  U=R@                  U
-  sl         [        [)        U5      5       H  n[        [)        UU   5      5       H  nUUR:                  :  a  M  UUR:                  :X  a  UUR<                  :  a  M7  UUR>                  :X  a  UUR@                  :  a  MY  UUR>                  :  a  Mk  UU   U   S	:X  d  My  S
UU   U'   M     M     U Vs/ s H  nSR                  U5      PM     nnSn[        [)        U5      5       HY  nU[        R                   " U R                  UR
                  R                  U-   5      R#                  5       S-   -  nX=U   S-   -  nM[     U$ s  snf ! [4         a     GNf = fs  snf s  snf s  snf )a  
Python 3.11+ only. Returns lines of source code (from code object `code`)
corresponding to `inst`'s location data, and underlines relevant code to `inst`.

Example: CALL on `g`:
f(g(
  ^^
    h(x)))
    ^^^^^

We need our own implementation in < 3.13 since `format_frame_summary` in
Python's `traceback` module doesn't handle multi-line expressions
(and their anchor extraction code is not completely correct).
)rY      )r  colno	end_colnor   r   r2   Nr%  r   ~^)!r  r  r  FrameSummaryrv  	positionsr  r  r  
col_offsetr  r  r   
splitlinesr  	linecachegetlinerstriptextwrapdedentr{  indentr  r   lstripr
  r2  r  rd  r  r  r  r  )r  instframe_summaryr  r  
orig_linesorig_lines_dedent
indent_lenr  
first_linestart_offset
end_offsetr   markers	last_lineline
num_spacesanchorsmarkermutable_markersr  r  s                         rc   get_instruction_source_311r+  a  s    7"!..NN!!LL~~00..++nn33
 &&7:6,,.qr23   5 5t~~7P7PST7TU
U d..7>>@U 	 
 %OODIIj,ABMMOA'#.?.B*CC
A{
+!8&A>>%%%~~$ ""4#3#3T^^5J5JKRRTJ~~  (~~  (DNN,I,I,Q z>>+D+DEL JGG ~~  DNN$9$99 ^^-J-JK
*5s\)C:3L,MMN]+d2s\)C3z?\3Q,RRS%%dnn77

&( 	 !NN,I,IJ
DNN11A5t~~7P7PQF$$T%5%5v>EEGDd{"GTS%77JNN3+cSY5K.LLM R 	[j))^c)*:*:*<&==
sZ'#j1H*IIJ"&G,W5
 :AB'6>>#s+'B HO+OwVDLw+O ""a'##|3#%%*&&,6& CL)FS!89:G333W444w?V?V9Vg888w999G666"6*3/3636OF+C0 ; *" 2AAv2776?AF3w< d..0E0E0IJQQS	
 	!*t## ! MI
r  
 C ,P4 Bs*   9Z)Z ;Z#Z(;Z-
Z Z c                \    [        U [        R                  5      (       a  [        U SS 5      $ g )N_dynamo_static_input_type)r  r0  rB  r&  rY  s    rc   get_static_address_typer.    s&    !U\\""q5t<<rb   c                   [         R                  R                  R                  [         R                  R                  [         R
                  [         R                  R
                  4n[         R                  R                  R                  [         R                  R                  [         R                  [         R                  R                  4nU / UQUQ7;   $ rx   )	r0  r  r5   	get_statedefault_generatorr  r6  	set_stater  )r   getterssetterss      rc   is_rng_state_getter_or_setterr5    s     	$$))

  G 	$$))

  	G (g((((rb   c                    [        U [        R                  5      =(       aF    U R                  S:H  =(       a0    U R                  R
                  [        R                  R                  L $ )N__get__)	r  r&  r   rZ   __self____objclass__r0  r  _TensorBaser  s    rc   is_tensor_base_attr_getterr;    sK    5%112 	@NNi'	@NN''588+?+??rb   c                     [         R                  " [        R                  U 5      n[	        U5      [
        R                  L $ ! [         a     gf = fr  )r  r  r0  rB  r  r&  r4  r  )rk   attrs     rc   is_tensor_getset_descriptorr>    sD    %%ellD9DzU7777 s   ?A 
AAc                    [        U S5      $ )N__torch_function__)r8  r  s    rc   is_torch_function_objectrA    s    5.//rb   c                   SSK Jn  SSKJn  U R	                  5       (       d0  [        U S5      (       a  [        U R                  5       S5      (       ap  S n[        X5      (       a  [        U R                  SS 5      nO'[        X5      (       a  [        U R                  SS 5      nUS [        R                  R                  4;  $ g)Nr   )UserDefinedObjectVariable)TensorWithTFOverrideVariable
peek_valuer@  F)r  rC  &torch._dynamo.variables.torch_functionrD  r  r8  rE  r  r&  
class_typer   r0  r  _disabled_torch_function_impl)r  rC  rD  r  s       rc   has_torch_functionrI    s     BS 
~~L!!gbmmo?S&T&Tb772==*>ED66288%94@DD%(("H"HIIIrb   c                    S nS n[         R                  R                  R                  5       =n(       a*  XR                  ;   a  UR                  U    nUR
                  nUR                  U SX#S9$ )NF)static_shapessymbolic_contextry  )r0  r4  r$   try_gettensor_to_contextrB  from_tensor)r!  r  rL  ry  tracing_contexts        rc   to_fake_tensorrQ  /  su     F--66>>@@@111.@@C%33F  	1A !  rb   c                   [        U 5      (       + =(       ao    [        U 5      (       + =(       aX    [        U 5      =(       aF    [        U S5      =(       a3    [        U R                  S5      =(       a    U R                  R
                  $ )N__dataclass_params__frozen)r,  r+  r   r8  rS  rT  r  s    rc   is_frozen_dataclassrU  ?  sn    #E** 	.&u--	.	. E12	. E..9		.
 &&--rb   c                n    U H   n[        X5      (       d  M  [        X5      s  $    [        U  SU 35      e)zP
Return the first available attribute or throw an exception if none is present.
z% does not has any of the attributes: )r8  r&  r2  )r  attrsr=  s      rc   get_first_attrrX  J  s=     33%%  C5 EeWM
NNrb   c              #     ^^#    U (       d  S v   g SUU4S jjn[         R                  R                  R                  U5       nUv   S S S 5        g ! , (       d  f       g = f7f)Nc                :   > SS jn[         R                  " XTTS9$ )Nc                    [         R                  R                  R                  S   S==   S-  ss'   [         R                  R                  X5      $ )Ncompiled_autogradcompilesr2   )r0  r  r.  rO   r  compile)gm_example_inputs_s     rc   inner_compilerKmaybe_enable_compiled_autograd.<locals>.compiler_fn.<locals>.inner_compiler^  s>    ##,,-@A*MQRRM..sDDrb   )backend	fullgraphdynamic)r_  r   r`  r   r   r   )r0  r^  )rF  ra  re  rd  s     rc   compiler_fn3maybe_enable_compiled_autograd.<locals>.compiler_fn]  s#    E ==i rb   )rF  r   r   r   )r0  r  r\  _enable)should_enablerd  re  rf  r  s    ``  rc   maybe_enable_compiled_autogradrj  U  sE      	 	 ]],,44[ASI BAAs   ?A"A	A"
AA"c                 @     " S S[         5      n [        U " 5       5      $ )Nc                      \ rS rSrSrg)*invalid_removeable_handle.<locals>.Invalidil  rH   N)rZ   r[   r\   r]   ra   rH   rb   rc   Invalidrm  l  s    rb   rn  )r  r1   )rn  s    rc   invalid_removeable_handlero  j  s    $  79%%rb   c                   [        U [        R                  R                  5      (       d  U $ [        U [        R                  R
                  5      (       a  U $ U R                  R                  U R                  5      nU R                  Ul        U$ rx   )	r  r0  rC  r  r   r  r  r$  r  )rJ  proxys     rc   nn_module_proxyrr  v  s`    c588??++
#uxx++,,
MM!!#--0E\\ENLrb   c                  B   ^  \ rS rSr      SU 4S jjrSS jrSrU =r$ )	GmWrapperi  c                :   > [         TU ]  5         Xl        X l        g rx   )r  rq  rF  unflatten_fn)rp  rF  rv  r  s      rc   rq  GmWrapper.__init__  s     	(rb   c                T    [        U5      nU R                  " U R                  U5      6 $ rx   )rd  rF  rv  r  s     rc   rH  GmWrapper.forward  s%    t*ww))$/00rb   )rF  rv  )rF  torch.fx.GraphModulerv  zCallable[[list[Any]], Any]r   r   r  r   r   r   )rZ   r[   r\   r]   rq  rH  ra   r  r  s   @rc   rt  rt    s*    )&)6P)	)1 1rb   rt  c                  ^	^
^^ [        U R                  R                  5       VVs/ s H<  u  p4UR                  S:X  d  M  UR                  R                  SS5      (       d  M:  UPM>     snnm[        R                  R                  R                  (       aW  TS/:X  d   e[        US   [        5      (       d   e[        US   5      m	S	S jmS
U	4S jjnU" [        X5      T" U5      5      m
O][        R                  " U5      u  pg[         R"                  " [        R$                  US9nU" [        X5      U5      m
[        R&                  mS	U
UU4S jjnU$ s  snnf )z
Mutate inputs so that they are flat and wrap gm such that it
accepts those inputs.  This is needed for graphs that take
bumpy inputs.
r   	steal_argFr   c                *    U S   [        U SS  5      -   $ rm  r  rM  s    rc   
flatten_fn(flatten_graph_inputs.<locals>.flatten_fn  s    7T$qr(^++rb   c                   > U S T /U TS  Q7$ rx   rH   )	flat_argsboxed_inputs_counts    rc   rv  *flatten_graph_inputs.<locals>.unflatten_fn  s#    112TY?Q?R5STTrb   )treespecc                 Z   > T" U 5      nT H  nX   R                  5         M     T" U5      $ rx   )ry   )r  r  r  compiled_fnr  inputs_idx_to_clears      rc   wrapper%flatten_graph_inputs.<locals>.wrapper  s1    t$	 %AGMMO % 9%%rb   r{  )r  r   r   r   )	enumeraterG  r  r  r  r}   r0  r  r\  in_compiled_autograd_regionr  rd  r{  rt  r  tree_flattenr  r  tree_unflattenarg_tree_leaves)rF  inputs
compile_gmr  rr  rv  flat_inputsspecr  r  r  r  r  s            @@@@rc   flatten_graph_inputsr    s$    !00GA77m# 	
(,		k5(I 	
0 }}&&BB #qc)))&)T**** ^	,	U !2!<j>PQ #//7 (()>)>N 2!<kJ++
& & NMs   EE$Ec                    [        U [        R                  R                  5      (       a  [	        U S5      (       d  / $ U R
                  R                  S/ 5      $ )Nr  locals_to_steal)r  r0  r   r  r8  r  r}   )maybe_gms    rc   get_locals_to_stealr    sB    h 4 455WXv=V=V	==.33rb   c                     XR                   S'   g )Nr  r  )rF  r  s     rc   set_locals_to_stealr    s    !0GGrb   c                  (    \ rS rSrSS jrSS jrSrg)Liti  c                    Xl         g rx   r  )rp  r  s     rc   rq  Lit.__init__  s    rb   c                    U R                   $ rx   r  ru  s    rc   r  Lit.__repr__  s    vvrb   r  N)r  r   r   r   rN  )rZ   r[   r\   r]   rq  r  ra   rH   rb   rc   r  r    s    rb   r  zset[str]warn_once_cachec                r    U [         ;   a  g [         R                  U 5        [        R                  " XS-   S9  g )Nr2   )
stacklevel)r  r~  warningswarn)r8  r  s     rc   	warn_oncer    s-     oMM#q.1rb   c                R    [         R                  " S5      nUR                  SU 5      $ )Nz\x1B[@-_][0-?]*[ -/]*[@-~]r%  )r  r^  r  )textansi_escapes     rc   strip_color_from_stringr    s"    **:;K??2t$$rb   c               #    #     [         R                  R                  R                  S5      n S v   [         R                  R                  R                  U 5        g ! [         R                  R                  R                  W 5        f = f7f)NT)r0  r  	_autograd _saved_tensors_hooks_set_tracing)r  s    rc   +_disable_saved_tensors_hooks_during_tracingr    sX     C""CCDI;;EB;;EBs   B	-A *B	+BB	c                     [         R                  R                  R                  =(       a    [         R                  " 5       (       + $ rx   )r0  r  rF   freezingis_grad_enabledrH   rb   rc   is_parameter_freezingr    s)    ??!!**J53H3H3J/JJrb   c                 f    [        [        5       5       V s/ s H  n [        U 5      PM     sn $ s  sn f rx   )r  r    get_torch_function_mode_stack_at)r  s    rc   get_torch_function_mode_stackr    s3    5:;T;V5W5W(+5W  s   .c                l    U [        5       :  a  U S:  d   e[        R                  R                  U 5      $ r   )r   r0  r  _get_function_stack_at)inds    rc   r  r    s/    *,,9988**3//rb   c                t    [        [        5       5       H  n[        5         M     U  H  n[        U5        M     g rx   )r  r   r   r   )rH  r  modes      rc   set_torch_function_mode_stackr    s/    ,./!# 0 %d+ rb   c                 L    [        [        5       5       H  n [        5         M     g rx   )r  r   r   )r  s    rc   clear_torch_function_mode_stackr    s    ,./!# 0rb   c                     [        5         g rx   )
breakpointrM  s    rc   _breakpoint_for_c_dynamor    s    Lrb   c                    U R                   n[        R                  " U5      n[        UR                  5      S:w  a  SSKJn  U" S5      eUR                  U R                  :w  a  SSKJn  U" S5      eg )NrX   r2   )InternalTorchDynamoErrorz[Tensor subclass method __metadata_guard__ must take exactly two subclass metadata argumentsz?Tensor subclass method __metadata_guard__ must be a classmethod)	__metadata_guard__r  	signaturer{  rD  r  r  r8  r  )r   r9  sigr  s       rc   verify_guard_fn_signaturer    sk    		!	!B


B
C
3>>a1&i
 	
 
{{eoo%1&M
 	
 &rb   c                   U R                   [        R                   [        R                   4;   =(       a    U R                  [        R                  [        R                  4;   =(       aa    U R                  [        R                  [        R                  4;   =(       a-    U R
                  [        R
                  [        R
                  4;   $ rx   )r   r  r   rX  re  __iter__)user_clss    rc   #does_not_override_dict_iter_methodsr  "  s    4::{'8'899 	GOO[-?-?@@	GMMdii)9)9::	G $--1E1E!FF	rb   c                $    U R                  U5      $ rx   )r  r^  r  s     rc   	call_sizer  .  s    66!9rb   c                $    U R                  U5      $ rx   )r  r  s     rc   call_strider  3  s    88A;rb   c                "    U R                  5       $ rx   )storage_offsetr]  s    rc   call_storage_offsetr  8  s    rb   c                    SS/nU Vs0 s H8  o"U R                   ;   d  M  U[        R                  " U R                   U   5      _M:     nnU$ s  snf )Nr-  tag)r  r  )r!  KEYS_TO_COPYr   tensor_dicts       rc   _extract_tensor_dictr  ?  s]    #L 4@3?C!**CT'TYYqzz#''<   	s
   A(Az(dict[int, weakref.ReferenceType[object]]user_obj_id_to_weakrefc                6    [         U    " 5       nUc   S5       eU$ )NzUser object is no longer alive)r  )obj_idr  s     rc   get_user_object_from_idr  R  s$    
 
(
*C?<<<?Jrb   c                T    [        U 5      n[        R                  " U 5      [        U'   g rx   )r  r  r  r  )r  r  s     rc   store_user_object_weakrefr  X  s    WF%,[[%56"rb   c                      \ rS rSr% SrS\S'   SrS\S'   Sr\SS j5       r	\SS j5       r
\SS	 j5       r\SS
 j5       r\\SS j5       5       rSrg)CompileTimeInstructionCounteri]  r   r   _counterrx  _idc                    U R                   S-   U l         U R                   S:X  a  [        R                  " 5       U l        g g r   )_depthr   r.  r  rz   s    rc   r.  #CompileTimeInstructionCounter.startb  s1    ZZ!^
::?*002CG rb   c                    U R                   S-
  U l         U R                   S:X  a;  U =R                  [        R                  " U R                  5      -  sl        SU l        g g )Nr2   r   rx  )r  r  r   endr  rz   s    rc   r  !CompileTimeInstructionCounter.endh  sF    ZZ!^
::?LL044SWW==LCG rb   c                    SU l         g r   r  rz   s    rc   ry   #CompileTimeInstructionCounter.clearo  s	    rb   c                    U R                   $ rx   r  rz   s    rc   r   #CompileTimeInstructionCounter.values  s    ||rb   c              #    #     [         R                  (       a  U R                  5         S v   [         R                  (       a  U R                  5         g g ! [         R                  (       a  U R                  5         f f = f7frx   )rF   %record_compile_time_instruction_countr.  r  rz   s    rc   record$CompileTimeInstructionCounter.recordw  sM     	;;		;;	 <v;;	 <s   A?)A 'A?(A<<A?rH   Nr   r   r   Generator[None, None, None])rZ   r[   r\   r]   r  r   r  r  r   r.  r  ry   r   r   r  ra   rH   rb   rc   r  r  ]  s    HcCMF3 3
         rb   r  c                  ,   ^  \ rS rSrSU 4S jjrSrU =r$ )CompileCounterInti  c                4   > [        [        TU ]	  U5      5      $ rx   )r  r  __add__)rp  otherr  s     rc   r  CompileCounterInt.__add__  s     !788rb   rH   )r  r   r   r  )rZ   r[   r\   r]   r  ra   r  r  s   @rc   r  r    s    9 9rb   r  c                t    [        5       R                  5       (       a  [        5       R                  SX5        gg)zE
Records whether we are using a feature
Generally a feature is a JK.
r  N)r   r   set_key_value)featureusages     rc   set_feature_user    s.     ((**++OWL +rb   )ddp_optimizerpython_reducer'python_reducer_without_compiled_forwardno_optimization_ddp_optimization_modec                     [         R                  n [        U [        5      (       a  U (       a  SOSnO0[        U [        5      (       a  U nO[        S[        U 5      < 35      eU[        ;   d   SU< 35       eU$ )Nr   r  z;Invalid dynamo config optimize_ddp type type(optimize_ddp)=z.Invalid dynamo config optimize_ddp value mode=)rF   optimize_ddpr  r   r   r  r  r  )r  r  s     rc   get_optimize_ddp_moder    s|    &&L,%%".4E	L#	&	&JtL7I6KL
 	
 )) 
9D7;) Krb   c               #  B  #    [         R                  =(       a    [        R                  " 5       n U (       aB  [        R                  " S5         [        R
                  " 5          Sv   SSS5        SSS5        gSv   g! , (       d  f       N= f! , (       d  f       g= f7f)aV  
Disables torch.inference_mode for the compilation (still on at runtime).
This simplifies the compile stack where we can assume that inference_mode
will always be off.

Since inference_mode is equivalent to no_grad + some optimizations (version
counts etc), we turn on no_grad here. The other optimizations are not
relevant to torch.compile.
FN)rF   "fake_tensor_disable_inference_moder0  is_inference_mode_enabledinference_moder  )is_inference_mode_ons    rc   maybe_disable_inference_moder    sd      	11We6U6U6W    'MMO  ('
 		 O ('s6   A
BB"A='B/B=
B	B
BBc               #     #    [         R                  (       a6  [        R                  R                  R                  5          Sv   SSS5        gSv   g! , (       d  f       g= f7f)z
Turns off tracking of inference_mode for fake tensor propagation. With this
context manager, when a real tensor is converted to fake tensor, the fake
tensor looses its inference-ness.
N)rF   r	  r0  rF  
meta_utils$disable_inference_mode_for_fake_proprH   rb   rc   *maybe_disable_inference_mode_for_fake_propr    sB      00))NNP QP 	 QPs   >A$ AA$
A!A$c                b    U S L =(       d%    SU R                   ;   =(       d    SU R                   ;   $ )Nr  rS  r  )rr  s    rc   is_node_meta_validr    s(    4<M?dii7M5DII;MMrb   c                     [         $ rx   _error_on_graph_breakrH   rb   rc   _get_error_on_graph_breakr    s      rb   c                    U q g rx   r  r  s    rc   _set_error_on_graph_breakr    s    !rb   c                     [         R                  R                  R                  (       a)  [         R                  R
                  R                  S5      O[        R                  " 5       n U R                  5         U $ )NzPregraph bytecode)
r0  autogradr1  _is_profiler_enabledr  	_profiler_RecordFunctionFastr:  nullcontext	__enter__cms    rc   record_pregraph_bytecode_enterr#    sS     >>""77 	../BC##% 
 LLNIrb   c                *    U R                  S S S 5        g rx   )__exit__r!  s    rc   record_pregraph_bytecode_exitr&    s    KKdD!rb   c                 0    SSK Jn   U R                  " 5       $ )Nr   r$   )torch._guardsr$   get_traced_coder(  s    rc   r*  r*    s    ,))++rb   )r   z-Union[list[tuple[str, Any]], list[list[Any]]]r   z!Union[tuple[str, ...], list[str]]r   r   r   )r   zCallable[[], str]r   r   )r   r   )r   r   r   r   )r   rS   )r   r    )r   r!   )NFNNNNFN)r   r   r>  rt  r   r   r   zOptional[dict[str, object]]r?  rt  r,  r  r'  r  r@  r   rA  rt  r   Generator[Any, None, None]r  )rR  zLiteral['str']rS  r   r   r   )rR  zLiteral['csv']rS  r   r   ztuple[list[str], list[object]])r   F)rR  r   rS  r   r   z-Union[str, None, tuple[list[str], list[str]]])r   zcontextlib.ExitStack)r  r  r  r
   r   r   )r  r   r  r=   r   r   )r  zfx.Graphr   r   )r^  rL   r   rL   )r^  r   r   r   )r  r   r  r   r   r   )r  r!  r  ztype[T]r   	TypeIs[T])r  r!  r  z)tuple[type[list[T]], type[tuple[T, ...]]]r   r,  )r  r!  r  zIterable[type]r   r   )r  r!  r  r   r   r   )r   r   r   r   )r   zCallable[..., T]r   z*TypeGuard[functools._lru_cache_wrapper[T]])r   r   r   z,TypeGuard[functools._lru_cache_wrapper[Any]])r   r   r   zMTypeIs[Union[_FuncTypes, torch._ops.OpOverloadPacket, torch._ops.OpOverload]])r   r   r   zTypeIs[_FuncTypes])r   r   r   zTypeIs[Union[types.GetSetDescriptorType, types.MethodDescriptorType, types.WrapperDescriptorType, types.MemberDescriptorType, types.MethodWrapperType]])r9  r   r   r   )r9  r   r   ztuple[Any, Optional[str]])r   r   r   zTypeGuard[np.ndarray])r  r   r   r   )rJ  r   r   r   r  rx   )rS  r   r   ztypes.CellType)r  r   r  r   r   z&tuple[tuple[Any, ...], dict[str, Any]])rg  r~  r   r  )r  rs  r   r   r  )r:  r   r;  r   r  rP   r<  zOptional[type[BaseException]]r=  zOptional[BaseException]r   r   )rD  r   r   r   )r   zlist[CompilationMetrics])r   rL  )r   r   )FF)r   r   r  r   r   r   r   r+  r  )r^  r  r  zOptional[torch.dtype]r   r  )r  z"dict[str, Union[T, tuple[T, ...]]]r   zdict[str, list[T]])r  zSequence[T]r   zlist[T])r  r   r   r   )rS  r  r   r   r  )r  r   r   zTypeIs[Union[torch.jit._trace.TopLevelTracedModule, torch.jit._script.RecursiveScriptModule, torch.jit.ScriptFunction[Any, Any], torch.jit.ScriptModule]])r  r   r  r   rz  r   r   r   )r  r   r   rt  )rl   r   r   r   )rl   r  r   r,  )rF  rz  r   zCallable[[], None])r2   )r  r   r  Iterable[Any]rN  r   r   ztuple[Any, float])rF  rz  r  r-  r   r   )r  r   r   r  )rg  r   r   r   )r   zset[int])r   r   r   z-TypeGuard[Union[torch.SymBool, torch.SymInt]])r   r   ry  r   r   r   )rY  r   r   r   )r  r-  r  zMapping[Any, Any]r   r   )r  
dict[K, V]r   zKeysView[K])r  r.  r   z!Iterable[tuple[K, Union[V, Any]]])rl   r   r   r   )r  zIterable[T]r   r   )r  r   r3  r   r   r   )r  r   r   ztuple[int, int, int])r!  r   rl   r  r   r   )r  zdict[Any, Any]r  r   r   r   )r  zset[T]r  r   r   rL   )r   r   r\  r   r   r   )rr  r  r  r   r   r   )r  r@   r   r   )r  r6  r  r   r   r   )r  r?   rk   r   r   r   )
r   r-  r  r   r  r>   r  r   r   r   )rf  r   r   z?TypeIs[Union[torch.Tensor, torch.nn.Module, MethodWrapperType]])r   r   r   	list[Any])r^  r   r\  r   r   r   )r!  r   r\  r   r   r   )r  r?   r)  r   r  r   r   r   )r  rL   r,  zContainer[T]r   r   )r3  r   r,  r   r/  r   r   r   )r9  zCallable[[], Any]r   r   )r  r   r  ,torch._subclasses.fake_tensor.FakeTensorModer   r   )r  r  r  r  r   r  )r  r   r  r   rj  r   rO  r   rT  rz  rP  r   rQ  r   rS  r   rR  r   rH  zCallable[..., None]rI  r   rJ  r   r   r   )r  r
   r   r   )r  r   r   r   rN  )T)rr  r  r  r   r   r   )ri  r   r  r?   r   r   )r  r?   r  r   r  r   r   r   )rr  r  r  r?   r  r   r   r   )r   Optional[torch.fx.Node])rr  r  r   r  )r  r   rr  r  r  r   r  r   r  r   r   r   )rr  r  r  r   r   r   )rF  rz  r   r   )r  r   r   r   )r   r   r!  r   r   r   )rJ  ztypes.ModuleTyper   r   )r  r   rk   r   r   r   r   r   )rl   r  r   r   )r   r   r3  r   r   r  )r7  r7  r   r   )rX  zUnion[torch.Tensor, Any]rA  r   rB  r#   r   z)tuple[bool, Optional[TensorStaticReason]])r&  r   rF  rz  r   r   )r3  r   rk   r   r  r   rS  r   r  r   r   r   )FFF)
rJ  torch.nn.Modulerg  r   rh  r   ri  r   r   r/  )
rJ  r2  rg  r   rh  r   ri  r   r   r   )r   r   r   r   )r  r   rk   r   r   r   )r^  r   r   r   )r9  r  r  r"  r  r#  r   rM   )r  r   r   r   )r  rA   r   r   )r   r   r  r   r   r   )r   r   r   zOptional[_Anchors])r  ztypes.CodeTyper  zdis.Instructionr   r   )r!  r   r   r   )rk   r   r   r   )r  r@   r   r   )r!  r  r  r0  r   r   )r  r   rW  r   r   r   )TT)ri  r   rd  r   re  r   r   r+  )r   r1   )rJ  r   r   r   )rF  rz  r  r   r  zCallable[[Any, Any], Any]r   r  )r  r   r   r/  )rF  rz  r  r/  r   r   )r8  r   r  r   r   r   )r  r   r   r   )r   r/  )r  r   r   r   )rH  r/  r   r   )r   r   r   r   )r  r   r   r   )r^  r   r  r   r   r   )r^  r   r   r   )r!  r  r   rP   )r  r   r   r   )r  r!  r   r   )r  r   r  r   r   r   )rr  r1  r   r   )r   r   r   r   )r   AbstractContextManager[None])r"  r3  r   r   )r   zOptional[list[CodeType]](   r   
__future__r   atexitr   r:  r  r*  r  rQ  rg  r  rJ  r'  r  r   r  r  r  r  r  r  r  r  r  r[  r   r  r&  r   r^  r  r  r   r   r   r   r   r	   r
   r   r   r   r   r   r   r   r   r   r   typing_extensionsr   r   r   r   r   r0  torch._functorch.config%torch.fx.experimental.symbolic_shapestorch.utils._pytreer.  _pytreer  r   torch._Cr   r   r   r   torch._dispatch.pythonr   torch._dynamo.metrics_contextr    r!   r)  r"   r#   r$   torch._subclasses.meta_utilsr%   torch._utils_internalr&   r'   r(   r)   r*   torch.fx._utilsr+   r,   torch.monitorr-   torch.nn.modules.lazyr.   torch.utils._tritonr/   r0   torch.utils.hooksr1   graph_utilsr3   r  collections.abcr4   r5   r6   r7   r8   r9   r:   r;   r<   torch._dynamo.replay_recordr=   torch._dynamo.symbolic_convertr>   r?   torch._dynamo.variables.baser@   torch._prims_commonrA   r{  r  r1  torch._loggingtorch._numpy_numpyrz  rB   rD   r%  rF   fftlinalgr  rG   r   NP_TO_TNP_MODULEr  rI   rJ   rK   r   rL   rM   rN   rC  r  r2  unpatched_nn_module_getattrr  unpatched_nn_module_call
_call_implunpatched_nn_module_call_implr   rO   rQ   troubleshooting_urlnnmodule_doc_urlnnmodule_doc_url_msgr  rZ   r   rR   rz  rT   rt   timer_counterrh  rV   re   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r\  r9  rO  rT  registerrj  FloatTensorr  DoubleTensorr  double
HalfTensorr  halfBFloat16Tensorrf  
ByteTensorr  
CharTensorr  
LongTensorr
  long	IntTensorr	  r   ShortTensorr  short
BoolTensorr   tensortype_to_dtyperl  r  r  r  r  r  r  r  r  r  r  r  r  r  r  ParamSpecArgsParamSpecKwargsTypeVarTupleTypeAliasTyper  r  r  r  r  r'  r(  r)  r*  r  r$  r  r  neltlegtgecmp_name_to_op_mappingcmp_name_to_op_str_mappingr6  r:  r8  r@  rH  rK  rN  rT  rj  ro  rq  r  r7  rs  !DEFAULT_COMPILATION_METRICS_LIMITrC  r  r  r  r!  r?  r   r   rF  rH  rJ  rL  r  r   r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r;  rG  rQ  rY  rb  r  r   rn   r  Ellipsisr  NotImplementedr  r  memory_formatr  finfoiinfo	attention
SDPBackendr6  _CudaDevicePropertiesrc  r7  r~  languagerf  r  rn  rv  r  r  r  r  r  r  r  re  r  rX  r  r   r  r  iterr  r  r  __length_hint__tuple_iterator_lenr!  r$  r  r  dict_newr   r  r'  dict_methodsrt  set_methodsre  frozenset_methodsr  	tuple_newtuple_methodsrd  list_methodsr  list_getitemstr_methodsr  r  r  r  r  r  r  r  r  	iter_nextr  r  dict_getitemr  r  r  r  r  r  r  r  r  r  r  r#  GLOBAL_KEY_PREFIXtorch._subclassesr%  r*  r.  r5  r9  r?  rD  r  rK  r  r  orig_code_mapr  r  seen_code_mapr  r  r  r  r  r  r  r  r  r  r  r  r  r"  r)  r,  r/  r+  r5  r7  r<  rC  rO  rT  rd  re  rf  all_hook_namesra  rp  rr  rv  r~  r  r  r  r  r  r  r  r  r  r  r  r  r  r
  r+  r.  r5  r;  r>  rA  rI  rQ  rU  rX  rj  ro  rr  rt  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  _disable_dynamor  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r#  r&  r*  r  s   0rc   <module>r     s   #       
   	         	 	 
          , = $  -
 
 
 O N   , $ $   < O ; ; =  G & 1 > - ' 

 
 
 < =2	.) 
FFIIII	>
: 	
 FFCGGIIszzIIszz	
  "VV
 CLCLt_#hhoo99  88??33  % : : 7B7N7N8
4  %' > &G  Q .//UV ! 46 0 5
 .9-D-DU-K * K!tyy %J %JP

7
.
 	
 
  ?C ; B*$  > !  / /$499 S# S#l OO%   !%"',0.2&*"&!/3NQ	NQ NQ  	NQ
 *NQ ,NQ $NQ  NQ NQ  -NQ  NQ NQb 
 L 
 L 
,1)
)%))#) 
) */$
$"&$2$N 8 8
 
u{{3	5	u}}ejj1	5>>+	u{{n	uzzm	u{{EJJ/	OOekk599-	U[[1	uzzm  ( #:"; ",
	?
G	 B 
 A 
 A 
	 I 

 
 C 
 C& w 	%!G& 
55/5 
5
 
7717 
7
	 				!
I R kkkkkkkkkkkk   83%$, T  
$05   v" v" v"r %' ! ?J>O>O,? ; 
-` 3F-`M3M3M3 M3 ,	M3
 'M3 
M3b "*DE 09ST %!
&s sl	 59 1 8*
  %*"'''!'  '  	' '@ ( ( (( +   )*  8<II4IIX 
6 

 
 = 
 =&
 9 9"* $(
< Q1 12. => -69&X 2  		J	NN	LL	KK		LL	KK	KK	HH!!	JJ$$'$ y , foo334$  >2W



!2
	
 "&bggi	 0%)"))+%6" 6(,RXXZ(8
% 8&*;=+?+?+A&B# B&*48n# 4&*4a>&:# :#33 ^^
<< //$--"6"6"8+:N:N:U:U:WXX X
 %(LL$7$7$9N$9&Xf=Mv$9N"++2244v8HF4  MM	&+nn&;&;&=R&=F&AQ&=R%)]]%9%9%;P%;6x?O%;P$'LL$7$7$9N$9&Xf=Mv$9NCLCLG1
: 	 @
8(M
 #(	&&& 	&  	&
 	&RM
MDM@B 
 !  = <7 <A((!(48(($
"?	?E??;  !&%(YY5:!&mF	mF	mF mF 	mF
 
mF mF mF mF mF #mF /3mF mF 
mF`	H
 
= 
= '( ;F:Q:QRV:W7 W LN H M '( , ,&
 $
/!/*-/EI//4 "'I
I!I I 		IX !1 " "LL$L,/L9<LHKLL^$N
*2G/	 27

*.

 5$  /	D*WWW&)W47W?BWW +,<= ,.?@   $&99<QQ !&!&#(		  !	
  F !&!&#(		  !	
 
 $

$gb!en 
$$$ $&$WRU^ $(@ '3<4.&$	C   {|AH)"0< L O AE$(:> (&
1 
1..&).7P..b41   E !	2 	2% C CK0
,$
"      
  DF @ E6
# #L9 9
M+  "  . 
 
N
  !"
   " ", ,yS  	B>  		rK
 O
 SP Os   AM; B*AN
 'AN=AN&AN<AN%AN!;AN!3AN&	AN&2AN+AN+ AN0AN0M;ANNANN
ANNAN