
    [Thq                     @   S r SSKrSSKrSSKJrJrJr  SSKJrJ	r	J
r
  SSKJr  SSKJrJr  \	R                   \	R"                  \	R$                  \	R&                  \	R"                  \	R"                  \	R&                  \	R&                  \	R(                  \	R(                  \	R*                  \	R*                  \	R,                  \	R,                  \	R.                  \	R.                  \	R0                  \	R0                  \	R2                  \	R2                  0
r\	R                   \	R(                  \	R$                  \	R*                  \	R(                  \	R(                  \	R*                  \	R*                  \	R"                  \	R(                  \	R&                  \	R*                  \	R,                  \	R,                  \	R.                  \	R.                  \	R0                  \	R0                  \	R2                  \	R2                  0
r\	R                   \	R,                  \	R$                  \	R.                  \	R(                  \	R,                  \	R*                  \	R.                  \	R"                  \	R,                  \	R&                  \	R.                  \	R,                  \	R,                  \	R.                  \	R.                  \	R0                  \	R0                  \	R2                  \	R2                  0
r\	R                   \	R0                  \	R$                  \	R2                  \	R"                  \	R0                  \	R&                  \	R2                  \	R0                  \	R0                  \	R2                  \	R2                  \	R(                  \	R0                  \	R*                  \	R2                  \	R,                  \	R0                  \	R.                  \	R2                  0
rS r\R>                  " S	S
9 " S S\
5      5       r \R>                  " S	S
9 " S S\
5      5       r!\R>                  " S	S
9 " S S\
5      5       r"\R>                  " S	S
9 " S S\
5      5       r#\R>                  " S	S
9 " S S\
5      5       r$\R>                  " S	S
9 " S S\5      5       r%\R>                  " S	S
9 " S S\%5      5       r&\R>                  " S	S
9 " S S\5      5       r'\R>                  " S	S
9 " S S\
5      5       r(\R>                  " S	S
9 " S S\5      5       r)\R>                  " S	S
9 " S S \'5      5       r*\R>                  " S	S
9 " S! S"\'5      5       r+\R>                  " S	S
9 " S# S$\
5      5       r, " S% S&\RZ                  5      r.\R>                  " S	S
9 " S' S(\5      5       r/\R>                  " S	S
9 " S) S*\5      5       r0\R>                  " S	S
9 " S+ S,\5      5       r1\R>                  " S	S
9 " S- S.\5      5       r2\R>                  " S	S
9 " S/ S0\5      5       r3\R>                  " S	S
9 " S1 S2\5      5       r4 " S3 S4\5      r5\R>                  " S	S
9 " S5 S6\5      5       r6\R>                  " S	S
9 " S7 S8\5      5       r7\R>                  " S	S
9 " S9 S:\5      5       r8\R>                  " S	S
9 " S; S<\5      5       r9\R>                  " S	S
9 " S= S>\75      5       r:\R>                  " S	S
9 " S? S@\75      5       r;\R>                  " S	S
9 " SA SB\5      5       r<\R>                  " S	S
9 " SC SD\5      5       r=\R>                  " S	S
9 " SE SF\5      5       r>\R>                  " S	S
9 " SG SH\>5      5       r?\R>                  " S	S
9 " SI SJ\?5      5       r@\R>                  " S	S
9 " SK SL\>5      5       rA\R>                  " S	S
9 " SM SN\
5      5       rB\R>                  " S	S
9 " SO SP\
5      5       rC\R>                  " S	S
9 " SQ SR\
5      5       rD\R>                  " S	S
9 " SS ST\5      5       rE\R>                  " S	S
9 " SU SV\5      5       rF\R>                  " S	S
9 " SW SX\5      5       rG\R>                  " S	S
9 " SY SZ\5      5       rH\R>                  " S	S
9 " S[ S\\
5      5       rI\R>                  " S	S
9 " S] S^\
5      5       rJS_S`.Sa\
4Sb jjrKSa\
4Sc jrLSa\
4Sd jrMSa\
4Se jrNSa\
4Sf jrOg)ga  
This module provides Source classes that track the origins of values in PyTorch Dynamo.
Sources represent where values come from (e.g. local variables, globals, attributes) and
are used for guard generation and code reconstruction during compilation.

The module includes specialized sources for:
- Local variables and synthetic locals
- Global variables and constants
- Object attributes and method calls
- NN module specialization (specialized vs unspecialized)
- Random values and tensor properties
- Default argument handling
- FSDP (Fully Sharded Data Parallel) modules

Sources play a key role in Dynamo's guard system by tracking value origins for
guard generation, and in code reconstruction by providing methods to rebuild
the code needed to recreate values.
    N)AnyOptionalUnion)ChainedSourceGuardSourceSource   )utils)create_call_functioncreate_instructionc                     [        U [        5      (       a  g U R                  5       [        R                  :X  a  g g! [
         a     gf = fNTF)
isinstanceConstantSourceguard_sourcer   CONSTANTNotImplementedErrorsources    L/var/www/auris/envauris/lib/python3.13/site-packages/torch/_dynamo/source.pyis_constant_sourcer   _   sR    &.)) K$8$88 9
   s   "= 
A
	A
T)frozenc                   h    \ rS rSr% \\S'   Sr\\S'   Sr\	\
\      \S'   Sr\\S'   S rS	 rS
 rSrg)LocalSourcek   
local_nameFis_inputNdynamismis_derefed_cell_contentsc                     U R                   (       a  UR                  U R                  5        g UR                  UR	                  U R                  5      5        g N)r   
load_derefr   append_outputcreate_loadselfcodegens     r   reconstructLocalSource.reconstruct{   s:    ((t/!!'"5"5doo"FG    c                 "    [         R                  $ r!   )r   LOCALr&   s    r   r   LocalSource.guard_source   s       r*   c                 4    S[        U R                  5       S3$ )NzL[])reprr   r-   s    r   nameLocalSource.name   s    D)*!,,r*    )__name__
__module____qualname____firstlineno__str__annotations__r   boolr   r   	frozensetr   r(   r   r2   __static_attributes__r4   r*   r   r   r   k   sE    O Hd *.Hhy~&-
 &+d*H!-r*   r   c                   2    \ rS rSr% \\S'   S rS rS rSr	g)SyntheticLocalSource   r   c                 X    UR                  UR                  U R                  5      5        g r!   )r#   r$   r   r%   s     r   r(    SyntheticLocalSource.reconstruct   s    g11$//BCr*   c                 "    [         R                  $ r!   )r   SYNTHETIC_LOCALr-   s    r   r   !SyntheticLocalSource.guard_source   s    ***r*   c                 $    SU R                   < S3$ )NzSYNTHETIC_LOCAL[r0   )r   r-   s    r   r2   SyntheticLocalSource.name   s    !$//!4A66r*   r4   N
r5   r6   r7   r8   r9   r:   r(   r   r2   r=   r4   r*   r   r?   r?      s    OD+7r*   r?   c                   2    \ rS rSr% \\S'   S rS rS rSr	g)RandomValueSource   random_call_indexc                 "    [         R                  $ r!   )r   RANDOM_VALUEr-   s    r   r   RandomValueSource.guard_source   s    '''r*   c                    UR                  UR                  UR                  R                  R                  5      5        UR                  UR                  U R                  5      5        UR                  [        S5      5        g NBINARY_SUBSCR)r#   r$   txoutputrandom_values_varcreate_load_constrL   r   r%   s     r   r(   RandomValueSource.reconstruct   s^    g11'**2C2C2U2UVWg778N8NOP0ABr*   c                      SU R                    3$ )Nrandom_value_)rL   r-   s    r   r2   RandomValueSource.name   s    t55677r*   r4   N)
r5   r6   r7   r8   intr:   r   r(   r2   r=   r4   r*   r   rJ   rJ      s    (C
8r*   rJ   c                   2    \ rS rSr% \\S'   S rS rS rSr	g)GlobalSource   global_namec                 V    UR                  UR                  U R                  SS95        g NTaddr#   create_load_globalr_   r%   s     r   r(   GlobalSource.reconstruct   s&    g889I9It8TUr*   c                 "    [         R                  $ r!   r   GLOBALr-   s    r   r   GlobalSource.guard_source       !!!r*   c                 4    S[        U R                  5       S3$ )NG[r0   r1   r_   r-   s    r   r2   GlobalSource.name   s    D))*+1--r*   r4   NrH   r4   r*   r   r]   r]      s    V".r*   r]   c                   2    \ rS rSr% \\S'   S rS rS rSr	g)GlobalWeakRefSource   r_   c                 j   ^ ^ TR                  UU 4S j5        TR                  [        SS5      5        g )Nc                  V   > T R                  T R                  TR                  SS95      $ ra   rd   r'   r&   s   r   <lambda>1GlobalWeakRefSource.reconstruct.<locals>.<lambda>   s)    G))**4+;+;*Fr*   r   Fadd_push_nullextend_outputr   r%   s   ``r   r(   GlobalWeakRefSource.reconstruct   s.    	

 	21e<=r*   c                 "    [         R                  $ r!   rh   r-   s    r   r    GlobalWeakRefSource.guard_source   rk   r*   c                 4    S[        U R                  5       S3$ )Nrm   z]()rn   r-   s    r   r2   GlobalWeakRefSource.name   s    D))*+3//r*   r4   NrH   r4   r*   r   rq   rq      s    >"0r*   rq   c                   &    \ rS rSrS rS rS rSrg)WeakRefCallSource   c                 j   ^ ^ TR                  UU 4S j5        TR                  [        SS5      5        g )Nc                  (   > T " TR                   5      $ r!   baseru   s   r   rv   /WeakRefCallSource.reconstruct.<locals>.<lambda>   s    gdii&8r*   r   Frx   r%   s   ``r   r(   WeakRefCallSource.reconstruct   s(    8921e<=r*   c                 6    U R                   R                  5       $ r!   r   r   r-   s    r   r   WeakRefCallSource.guard_source       yy%%''r*   c                 <    U R                   R                  5        S3$ )Nz()r   r2   r-   s    r   r2   WeakRefCallSource.name   s    )).."#2&&r*   r4   Nr5   r6   r7   r8   r(   r   r2   r=   r4   r*   r   r   r      s    >('r*   r   c                       \ rS rSrSrg)CallFunctionNoArgsSource   r4   Nr5   r6   r7   r8   r=   r4   r*   r   r   r          r*   r   c                   8    \ rS rSr% \\S'   S rS rS rS r	Sr
g)	
AttrSource   memberc                 :   U R                   (       d   S5       eSU R                  ;   as  U R                  R                  S5      n[        R	                  U S[        U R                   SR                  US S 5      5      5        [        R	                  U SUS   5        g g )Nz9Can't construct an AttrSource without a valid base source.r   r   )r   r   splitobject__setattr__r   join)r&   member_partss     r   __post_init__AttrSource.__post_init__   s    yyUUUy$++;;,,S1LfjCHH\#2=N4OP tX|B/?@ r*   c                 |    U" U R                   5        UR                  UR                  U R                  5      5        g r!   r   rz   create_load_attrsr   r%   s     r   r(   AttrSource.reconstruct   +    		g77DEr*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   AttrSource.guard_source   r   r*   c                     U R                   R                  5       (       d,  SU R                  R                  5        SU R                   < S3$ U R                  R                  5        SU R                    3$ )Nzgetattr(, )r   )r   isidentifierr   r2   r-   s    r   r2   AttrSource.name   s[    {{''))diinn./r$++BB)).."#1T[[M22r*   r4   N)r5   r6   r7   r8   r9   r:   r   r(   r   r2   r=   r4   r*   r   r   r      s    KAF(3r*   r   c                   *    \ rS rSr% Sr\\S'   S rSrg)LocalCellSource   zw
Conceptually, this class is `LocalSource` for cell objects implicitly
generated by Python (e.g., captured variables).
r   c                 X    UR                  UR                  U R                  5      5        g r!   )r#   create_load_closurer   r%   s     r   r(   LocalCellSource.reconstruct   s!     	g99$//JKr*   r4   N)	r5   r6   r7   r8   __doc__r9   r:   r(   r=   r4   r*   r   r   r      s    
 OLr*   r   c                   6    \ rS rSr% Sr\\S'   S rS rS r	Sr
g)	
GradSourcei
  gradr   c                 |    U" U R                   5        UR                  UR                  U R                  5      5        g r!   r   r%   s     r   r(   GradSource.reconstruct  r   r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   GradSource.guard_source  r   r*   c                 T    U R                   R                  5        SU R                   3$ )Nr   )r   r2   r   r-   s    r   r2   GradSource.name  s#    )).."#1T[[M22r*   r4   N)r5   r6   r7   r8   r   r9   r:   r(   r   r2   r=   r4   r*   r   r   r   
  s    FCF(3r*   r   c                       \ rS rSrS rSrg)ParamBufferSourcei  c                 D    [         U R                  R                  5          $ r!   #_GUARD_SOURCE_SPECIALIZED_NN_MODULEr   r   r-   s    r   r   ParamBufferSource.guard_source      24993I3I3KLLr*   r4   Nr5   r6   r7   r8   r   r=   r4   r*   r   r   r     s    Mr*   r   c                       \ rS rSrSrg)UnspecializedParamBufferSourcei   r4   Nr   r4   r*   r   r   r      r   r*   r   c                   B    \ rS rSr% Sr\\   \S'   S rS r	S r
S rSrg)	EphemeralSourcei.  Ndescc                 "    [         R                  $ r!   )r   	EPHEMERALr-   s    r   r   EphemeralSource.guard_source2      $$$r*   c                 L    SU R                   b  SU R                   -    S3$ S S3$ )Nz
<ephemeralz:  >)r   r-   s    r   r2   EphemeralSource.name5  s1    		0ED499,NaPP2NaPPr*   c                     [         er!   r   r&   fns     r   
make_guardEphemeralSource.make_guard8      !!r*   c                     gNTr4   r-   s    r   is_ephemeralEphemeralSource.is_ephemeral;      r*   r4   )r5   r6   r7   r8   r   r   r9   r:   r   r2   r   r   r=   r4   r*   r   r   r   .  s%    D(3-%Q"r*   r   c                   &    \ rS rSrSrSrSrS rSrg)TensorPropertyi?  r   r	      c                 |    U [         R                  L a  gU [         R                  L a  gU [         R                  L a  gg )Nsizestridestorage_offset)r   SIZESTRIDESTORAGE_OFFSETr-   s    r   method_nameTensorProperty.method_nameD  s:    >&&&^***^222# 3r*   r4   N)	r5   r6   r7   r8   r   r   r   r   r=   r4   r*   r   r   r   ?  s    DFN$r*   r   c                   L    \ rS rSr% \\S'   Sr\\   \S'   S r	S r
S rS rS	rg)
TensorPropertySourceiM  propNidxc                     U R                   c   eU R                  [        R                  L a  U R                  b   eg U R                  c   eg r!   )r   r   r   r   r   r-   s    r   r   "TensorPropertySource.__post_init__R  sD    yy$$$9955588###88'''r*   c                   ^ ^ TR                  UU 4S j5        T" T R                  5        T R                  b*  TR                  TR	                  T R                  5      5        TR                  [        T R                  b  SOSS5      5        g )Nc                  z   > T R                  [        R                  STR                  R	                  5        35      $ )Ncall_)load_import_fromr
   r5   r   r   ru   s   r   rv   2TensorPropertySource.reconstruct.<locals>.<lambda>[  s.    G,,%		(=(=(?'@ Ar*   r   r	   F)ry   r   r   r#   rV   rz   r   r%   s   ``r   r(    TensorPropertySource.reconstructY  sn    	

 			88!!'";";DHH"EF dhh&:5I	
r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   !TensorPropertySource.guard_sourceg  r   r*   c                    U R                   [        R                  L a*  U R                  R	                  5        SU R
                   S3$ U R                   [        R                  L a*  U R                  R	                  5        SU R
                   S3$ U R                   [        R                  L a,  U R
                  b   eU R                  R	                  5        S3$ [        SU R                    35      e)Nz.size()[r0   z
.stride()[z.storage_offset()z
unhandled )	r   r   r   r   r2   r   r   r   AssertionErrorr-   s    r   r2   TensorPropertySource.namej  s    99+++iinn&'xz;;YY.///iinn&'z$((1==YY.77788###iinn&''899 :dii[!9::r*   r4   )r5   r6   r7   r8   r   r:   r   r   r[   r   r(   r   r2   r=   r4   r*   r   r   r   M  s*    
C#(
(	;r*   r   c                   8    \ rS rSr% \\S'   S rS rS rS r	Sr
g)	IndexedSourceiv  r   c                 "    U R                   c   eg r!   r   r-   s    r   r   IndexedSource.__post_init__z      yy$$$r*   c                     [         er!   r   r%   s     r   r(   IndexedSource.reconstruct}  r   r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   IndexedSource.guard_source  r   r*   c                 X    SU R                    SU R                  R                  5        S3$ )N(r   r   )r   r   r2   r-   s    r   r2   IndexedSource.name  s'    488*Btyy~~/022r*   r4   N)r5   r6   r7   r8   r[   r:   r   r(   r   r2   r=   r4   r*   r   r   r   v  s    	H%"(3r*   r   c                   ,    \ rS rSrS rS rS rS rSrg)NegateSourcei  c                 "    U R                   c   eg r!   r   r-   s    r   r   NegateSource.__post_init__  r  r*   c                     [         er!   r   r%   s     r   r(   NegateSource.reconstruct  r   r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   NegateSource.guard_source  r   r*   c                 <    U R                   R                  5        S3$ )Nz
.__neg__()r   r-   s    r   r2   NegateSource.name  s    )).."#:..r*   r4   N	r5   r6   r7   r8   r   r(   r   r2   r=   r4   r*   r   r  r    s    %"(/r*   r  c                   ,    \ rS rSrS rS rS rS rSrg)ConvertIntSourcei  c                 "    U R                   c   eg r!   r   r-   s    r   r   ConvertIntSource.__post_init__  r  r*   c                 (    U" U R                   5        g r!   r   r%   s     r   r(   ConvertIntSource.reconstruct      		r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   ConvertIntSource.guard_source  r   r*   c                 >    SU R                   R                  5        S3$ )Nz!cast_symbool_to_symint_guardless(r   r   r-   s    r   r2   ConvertIntSource.name  s    2499>>3C2DAFFr*   r4   Nr  r4   r*   r   r  r    s    %(Gr*   r  c                   ,    \ rS rSrS rS rS rS rSrg)FlattenScriptObjectSourcei  c                 "    U R                   c   eg r!   r   r-   s    r   r   'FlattenScriptObjectSource.__post_init__  r  r*   c                 (    U" U R                   5        g r!   r   r%   s     r   r(   %FlattenScriptObjectSource.reconstruct  r  r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   &FlattenScriptObjectSource.guard_source  r   r*   c                 <    U R                   R                  5        S3$ )Nz.__obj_flatten__()r   r-   s    r   r2   FlattenScriptObjectSource.name  s    )).."##566r*   r4   Nr  r4   r*   r   r!  r!    s    %(7r*   r!  c                   ,    \ rS rSrS rS rS rS rSrg)ScriptObjectQualifiedNameSourcei  c                 "    U R                   c   eg r!   r   r-   s    r   r   -ScriptObjectQualifiedNameSource.__post_init__  r  r*   c                 (    U" U R                   5        g r!   r   r%   s     r   r(   +ScriptObjectQualifiedNameSource.reconstruct  r  r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   ,ScriptObjectQualifiedNameSource.guard_source  r   r*   c                 <    U R                   R                  5        S3$ )Nz._type().qualified_name()r   r-   s    r   r2   $ScriptObjectQualifiedNameSource.name  s    )).."##<==r*   r4   Nr  r4   r*   r   r+  r+    s    %(>r*   r+  c                   &    \ rS rSrS rS rS rSrg)AttrProxySourcei  c                 (    U" U R                   5        g r!   r   r%   s     r   r(   AttrProxySource.reconstruct  r  r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   AttrProxySource.guard_source  r   r*   c                 <    U R                   R                  5        S3$ )Nz.get_base()r   r-   s    r   r2   AttrProxySource.name  s    )).."#;//r*   r4   Nr   r4   r*   r   r5  r5    s    (0r*   r5  c                       \ rS rSr% \\\4   \S'   Sr\	\S'   \
R                  " SSSS9r\\S'   \
R                  " SSSS9r\\S'   S rS	 rS
 rS rSrg)DefaultsSourcei  idx_keyFis_kw)initr1   comparefield_namec           
      h   U R                   (       d   S5       eU R                  (       a  [        U R                  [        5      (       d   e[
        R                  U SS5        [
        R                  U SU R                   R                  5        SU R                   SU R                   S35        g [        U R                  [        5      (       d   e[
        R                  U SS5        [
        R                  U SU R                   R                  5        SU R                   S	U R                   S
35        g )Nz]Base must be a valid source in order to properly track and guard this Defaults to its origin.rB  __kwdefaults__rC  r   z['z']__defaults__[r0   )
r   r?  r   r>  r9   r   r   r2   rB  r[   r-   s    r   r   DefaultsSource.__post_init__  s    yy 	
k	
y ::dllC0000tW.>?g$)).."2!31TZZL4<<.PRS dllC0000tWn=g$)).."2!31TZZL$,,qQr*   c                    U" U R                   5        UR                  UR                  U R                  5      5        UR	                  UR                  U R                  5      5        UR	                  [        S5      5        g rQ   )r   rz   r   rB  r#   rV   r>  r   r%   s     r   r(   DefaultsSource.reconstruct  s\    		g77

CDg77EF0ABr*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   DefaultsSource.guard_source  r   r*   c                     U R                   $ r!   )rC  r-   s    r   r2   DefaultsSource.name  s    zzr*   r4   N)r5   r6   r7   r8   r   r[   r9   r:   r?  r;   dataclassesrB  rC  r   r(   r   r2   r=   r4   r*   r   r=  r=    sa    38_E4""E5IE3I""E5IE3I"C(r*   r=  c                   Z   ^  \ rS rSr% \\S'   Sr\\S'   U 4S jrS r	S r
S rS	 rS
rU =r$ )GetItemSourcei  indexFindex_is_slicec                    > U R                   c   e[        U R                  [        5      (       a9  [        TU ]  SU R                  R                  5       5        [        TU ]  SS5        g g )NrR  rS  T)r   r   rR  slicesuperr   
__reduce__)r&   	__class__s    r   r   GetItemSource.__post_init__  sU    yy$$$djj%((G)>)>)@AG 0$7 )r*   c                 0   U" U R                   5        U R                  (       a/  UR                  UR                  U R	                  5       5      5        O*UR                  UR                  U R
                  5      5        UR                  [        S5      5        g rQ   )r   rS  r#   rV   unpack_slicerR  r   r%   s     r   r(   GetItemSource.reconstruct   sj    		!!'";";D<M<M<O"PQ!!'";";DJJ"GH0ABr*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   GetItemSource.guard_source  r   r*   c                 N    U R                   (       d   eU R                  u  pU" U6 $ r!   )rS  rR  )r&   slice_class
slice_argss      r   r[  GetItemSource.unpack_slice  s(    """""&**J''r*   c                    [        U R                  [        5      (       a   eU R                  (       a/  U R                  R                  5        SU R                  5       < S3$ U R                  R                  5        SU R                  < S3$ )NrG  r0   )r   rR  r   rS  r   r2   r[  r-   s    r   r2   GetItemSource.name  sp     djj&1111iinn&'q):):)<(?qAAiinn&'qa88r*   r4   )r5   r6   r7   r8   r   r:   rS  r;   r   r(   r   r[  r2   r=   __classcell__)rX  s   @r   rQ  rQ    s1    J ND 8C((
9 9r*   rQ  c                   8    \ rS rSr% \\S'   S rS rS rS r	Sr
g)	ConstDictKeySourcei  rR  c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   ConstDictKeySource.guard_source  r   r*   c                    ^ TR                  U4S j5        T" U R                  5        TR                  TR                  U R                  5      5        TR                  [        SS5      5        g )Nc                  D   > T R                  [        R                  S5      $ )Ndict_keys_getitemr   r
   r5   r'   s   r   rv   0ConstDictKeySource.reconstruct.<locals>.<lambda>$  s    G,,U^^=PQr*   r   Fry   r   r#   rV   rR  rz   r   r%   s    `r   r(   ConstDictKeySource.reconstruct"  sU    Q	
 			g77

CD21e<=r*   c                 Z    SU R                   R                  5        SU R                  < S3$ )Nzlist(dict.keys(z))[r0   r   r2   rR  r-   s    r   r2   ConstDictKeySource.name*  s'     !1 2#djj^1EEr*   c                     gr   r4   r-   s    r   is_dict_keyConstDictKeySource.is_dict_key.  r   r*   r4   N)r5   r6   r7   r8   r   r:   r   r(   r2   rv  r=   r4   r*   r   rg  rg    s    J(>Fr*   rg  c                   8    \ rS rSr% \\S'   S rS rS rS r	Sr
g)	DictGetItemSourcei3  rR  c                     SSK Jn  [        U R                  [        5      (       d#  UR                  U R                  5      (       d   eg g )Nr	   )ConstantVariable)	variablesr{  r   rR  rg  
is_literal)r&   r{  s     r   r   DictGetItemSource.__post_init__:  sC    /JJ*
 
((44	5 54
r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   DictGetItemSource.guard_sourceA  r   r*   c                 B  ^ TR                  U4S j5        T" U R                  5        [        U R                  [        5      (       a  T" U R                  5        O*TR                  TR                  U R                  5      5        TR                  [        SS5      5        g )Nc                  D   > T R                  [        R                  S5      $ )Ndict_getitemrm  rn  s   r   rv   /DictGetItemSource.reconstruct.<locals>.<lambda>I      G,,U^^^Lr*   r   F)	ry   r   r   rR  r   r#   rV   rz   r   r%   s    `r   r(   DictGetItemSource.reconstructD  sv     	L	

 			 djj&))DJJ!!'";";DJJ"GH21e<=r*   c                    [        U R                  [        5      (       a9  SU R                  R	                  5        SU R                  R	                  5        S3$ U R                  R	                  5        SU R                  < S3$ )Nzdict.__getitem__(r   r   rG  r0   )r   rR  rg  r   r2   r-   s    r   r2   DictGetItemSource.nameW  sd    djj"455&tyy~~'7&84::??;L:MQOOiinn&'qa88r*   r4   N)r5   r6   r7   r8   r   r:   r   r   r(   r2   r=   r4   r*   r   ry  ry  3  s    
 J5(>&9r*   ry  c                   $    \ rS rSrSrS rS rSrg)ListGetItemSourcei^  zQ
Same as GetItemSource with reconstruct and name overridden to be list specific.
c                   ^ TR                  U4S j5        T" U R                  5        U R                  (       a  [        S5      eTR	                  TR                  U R                  5      5        TR                  [        SS5      5        g )Nc                  D   > T R                  [        R                  S5      $ )Nlist_getitemrm  rn  s   r   rv   /ListGetItemSource.reconstruct.<locals>.<lambda>j  r  r*   >List[slice] is a temporary object and should not have a sourcer   F)	ry   r   rS  RuntimeErrorr#   rV   rR  rz   r   r%   s    `r   r(   ListGetItemSource.reconstructd  st    
 	L	

 			 P  !!'";";DJJ"GH21e<=r*   c                     [        U R                  [        5      (       a   eU R                  (       a  [	        S5      eSU R
                  R                  5        SU R                  < S3$ )Nr  zlist.__getitem__(r   r   )r   rR  r   rS  r  r   r2   r-   s    r   r2   ListGetItemSource.namez  s[     djj&1111P  'tyy~~'7&84::.JJr*   r4   N)r5   r6   r7   r8   r   r(   r2   r=   r4   r*   r   r  r  ^  s    >,
Kr*   r  c                        \ rS rSrS rS rSrg)TupleIteratorGetItemSourcei  c                    ^ TR                  U4S j5        T" U R                  5        TR                  TR                  U R                  5      5        TR                  [        SS5      5        g )Nc                  D   > T R                  [        R                  S5      $ )Ntuple_iterator_getitemrm  rn  s   r   rv   8TupleIteratorGetItemSource.reconstruct.<locals>.<lambda>  s    G,,U^^=UVr*   r   Frp  r%   s    `r   r(   &TupleIteratorGetItemSource.reconstruct  sU    V	
 			g77

CD21e<=r*   c                 Z    SU R                   R                  5        SU R                  < S3$ )Nz___tuple_iterator_getitem(r   r   rs  r-   s    r   r2   TupleIteratorGetItemSource.name  s'    +DIINN,<+=R

~QOOr*   r4   N)r5   r6   r7   r8   r(   r2   r=   r4   r*   r   r  r    s    >Pr*   r  c                   ,    \ rS rSrS rS rS rS rSrg)
TypeSourcei  c                 "    U R                   c   eg r!   r   r-   s    r   r   TypeSource.__post_init__  r  r*   c                    ^ TR                  U4S j5        T" U R                  5        TR                  [        SS5      5        g )Nc                  (   > T R                  SS5      $ )Nbuiltinstyper   rn  s   r   rv   (TypeSource.reconstruct.<locals>.<lambda>  s    g&>&>z6&Rr*   r	   Fry   r   rz   r   r%   s    `r   r(   TypeSource.reconstruct  s4    RS		21e<=r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   TypeSource.guard_source  r   r*   c                 >    SU R                   R                  5        S3$ )Nztype(r   r   r-   s    r   r2   TypeSource.name  s    tyy~~'(**r*   r4   Nr  r4   r*   r   r  r    s    %>
(+r*   r  c                   &    \ rS rSrS rS rS rSrg)OptimizerSourcei  c                 (    U" U R                   5        g r!   r   r%   s     r   r(   OptimizerSource.reconstruct  r  r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   OptimizerSource.guard_source  r   r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r2   OptimizerSource.name      yy~~r*   r4   Nr   r4   r*   r   r  r    s    ( r*   r  c                   &    \ rS rSrS rS rS rSrg)NNModuleSourcei  c                 (    U" U R                   5        g r!   r   r%   s     r   r(   NNModuleSource.reconstruct  r  r*   c                 D    [         U R                  R                  5          $ r!   r   r-   s    r   r   NNModuleSource.guard_source  r   r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r2   NNModuleSource.name  r  r*   r4   Nr   r4   r*   r   r  r    s    M r*   r  c                       \ rS rSrS rSrg)UnspecializedNNModuleSourcei  c                 D    [         U R                  R                  5          $ r!   )%_GUARD_SOURCE_UNSPECIALIZED_NN_MODULEr   r   r-   s    r   r   (UnspecializedNNModuleSource.guard_source  s    4TYY5K5K5MNNr*   r4   Nr   r4   r*   r   r  r    s    Or*   r  c                       \ rS rSrS rSrg)"UnspecializedBuiltinNNModuleSourcei  c                 D    [         U R                  R                  5          $ r!   )-_GUARD_SOURCE_UNSPECIALIZED_BUILTIN_NN_MODULEr   r   r-   s    r   r   /UnspecializedBuiltinNNModuleSource.guard_source  s    <TYY=S=S=UVVr*   r4   Nr   r4   r*   r   r  r    s    Wr*   r  c                       \ rS rSrS rSrg)FSDPNNModuleSourcei  c                 D    [         U R                  R                  5          $ r!   )_GUARD_SOURCE_FSDP_MODULEr   r   r-   s    r   r   FSDPNNModuleSource.guard_source  s    ()?)?)ABBr*   r4   Nr   r4   r*   r   r  r    s    Cr*   r  c                        \ rS rSrS rS rSrg)GlobalStateSourcei  c                     gNr   r4   r-   s    r   r2   GlobalStateSource.name      r*   c                 "    [         R                  $ r!   rh   r-   s    r   r   GlobalStateSource.guard_source  rk   r*   r4   Nr5   r6   r7   r8   r2   r   r=   r4   r*   r   r  r    s    "r*   r  c                   8    \ rS rSr% \\S'   S rS rS rS r	Sr
g)	TorchFunctionModeStackSourcei  indc                 *    SU R                  5        S3$ )Nz$___get_torch_function_mode_stack_at(r   )
_get_indexr-   s    r   r2   !TorchFunctionModeStackSource.name  s    5doo6G5HJJr*   c                 D    SSK Jn  UR                  U R                  5      $ )Nr	   )TorchFunctionModeStackVariable)variables.torch_functionr  get_mode_indexr  )r&   r  s     r   r  'TorchFunctionModeStackSource._get_index  s    L-<<TXXFFr*   c                    ^ TR                  U4S j5        TR                  TR                  U R                  5       5      /5        TR                  [	        SS5      5        g )Nc                  D   > T R                  [        R                  S5      $ )N get_torch_function_mode_stack_atrm  rn  s   r   rv   :TorchFunctionModeStackSource.reconstruct.<locals>.<lambda>  s    G,, Br*   r	   F)ry   rz   rV   r  r   r%   s    `r   r(   (TorchFunctionModeStackSource.reconstruct  sQ    	

 	w889JKLM21e<=r*   c                 "    [         R                  $ r!   rh   r-   s    r   r   )TorchFunctionModeStackSource.guard_source  rk   r*   r4   N)r5   r6   r7   r8   r[   r:   r2   r  r(   r   r=   r4   r*   r   r  r    s    	HKG
>"r*   r  c                   8    \ rS rSr% \\S'   S rS rS rS r	Sr
g)	r   i  source_namec                 V    UR                  UR                  U R                  SS95        g )NFrb   )r#   re   r  r%   s     r   r(   ConstantSource.reconstruct  s&    g889I9Iu8UVr*   c                 "    [         R                  $ r!   )r   r   r-   s    r   r   ConstantSource.guard_source  s    ###r*   c                     U R                   $ r!   )r  r-   s    r   r2   ConstantSource.name  s    r*   c                     [         er!   r   r   s     r   r   ConstantSource.make_guard  r   r*   r4   N)r5   r6   r7   r8   r9   r:   r(   r   r2   r   r=   r4   r*   r   r   r     s    W$ "r*   r   c                   .    \ rS rSrS\4S jrS rS rSrg)NumpyTensorSourcei  returnc                 >    SU R                   R                  5        S3$ )Nz___from_numpy(r   r   r-   s    r   r2   NumpyTensorSource.name  s    		 0133r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   NumpyTensorSource.guard_source  r   r*   c                    ^ TR                  U4S j5        T" U R                  5        TR                  [        SS5      5        g )Nc                  (   > T R                  SS5      $ )Ntorch	as_tensorr  rn  s   r   rv   /NumpyTensorSource.reconstruct.<locals>.<lambda>  s    g&>&>w&Tr*   r	   Fr  r%   s    `r   r(   NumpyTensorSource.reconstruct  s4    TU		21e<=r*   r4   N)	r5   r6   r7   r8   r9   r2   r   r(   r=   r4   r*   r   r  r    s    4c 4(>r*   r  c                   (    \ rS rSrS\4S jrS rSrg)SubclassAttrListSourcei  r  c                 <    U R                   R                  5        S3$ )Nz.__tensor_flatten__()[0]r   r-   s    r   r2   SubclassAttrListSource.name  s    )).."##;<<r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   #SubclassAttrListSource.guard_source  r   r*   r4   Nr5   r6   r7   r8   r9   r2   r   r=   r4   r*   r   r   r     s    =c =(r*   r   c                   (    \ rS rSrS\4S jrS rSrg)FloatTensorSourcei  r  c                 >    SU R                   R                  5        S3$ )Nz___as_tensor(r   r   r-   s    r   r2   FloatTensorSource.name  s    tyy~~/022r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   FloatTensorSource.guard_source!  r   r*   r4   Nr  r4   r*   r   r  r    s    3c 3(r*   r  c                   (    \ rS rSrS\4S jrS rSrg)CallMethodItemSourcei%  r  c                 <    U R                   R                  5        S3$ )Nz.item()r   r-   s    r   r2   CallMethodItemSource.name'  s    )).."#7++r*   c                 6    U R                   R                  5       $ r!   r   r-   s    r   r   !CallMethodItemSource.guard_source*  r   r*   r4   Nr  r4   r*   r   r  r  %  s    ,c ,(r*   r  c                        \ rS rSrS rS rSrg)ShapeEnvSourcei1  c                     gr  r4   r-   s    r   r2   ShapeEnvSource.name3  r  r*   c                 "    [         R                  $ r!   )r   	SHAPE_ENVr-   s    r   r   ShapeEnvSource.guard_source6  r   r*   r4   Nr  r4   r*   r   r  r  1  s    %r*   r  c                        \ rS rSrS rS rSrg)BackwardStateSourcei:  c                     gr  r4   r-   s    r   r2   BackwardStateSource.name<  r  r*   c                 "    [         R                  $ r!   )r   BACKWARD_STATEr-   s    r   r    BackwardStateSource.guard_source?  s    )))r*   r4   Nr  r4   r*   r   r  r  :  s    *r*   r  Fonly_allow_inputr   c                    [        U [        5      (       a  [        U R                  US9$ [        U [        5      (       d  gU(       a  U R
                  (       d  gg)Nr   FT)r   r   is_from_local_sourcer   r   r   )r   r!  s     r   r#  r#  C  s@    &-((#FKKBRSSfk**r*   c                     [        U [        5      (       a  g[        U [        5      (       a  [        U R                  5      $ gr   )r   r   r   )is_from_unspecialized_param_buffer_sourcer   r   s    r   r%  r%  M  s2    &899&-((8EEr*   c                     [        U [        5      (       a  g[        U [        5      (       a  [        U R                  5      $ gr   )r   r!  r   $is_from_flatten_script_object_sourcer   r   s    r   r'  r'  U  s2    &344	FM	*	*3FKK@@r*   c                     [        U [        5      (       a  g[        U [        5      (       a  [        U R                  5      $ gr   )r   r  r   is_from_optimizer_sourcer   r   s    r   r)  r)  ]  s1    &/**&-(('44r*   c                    [        U [        5      (       a  g[        U [        5      (       a:  [        U R                  [        5      (       a  U R                  R
                  S:X  a  g[        U [        5      (       a:  [        U R                  [        5      (       a  U R                  R
                  S:X  a  g[        U [        5      (       a  [        U R                  5      $ g)NTrE  rF  F)	r   r=  ry  r   r   r   rQ  r   is_from_defaultsr   s    r   r+  r+  g  s    &.)) 	6,--v{{J//KK"22 	6=))v{{J//KK.0&-((,,r*   )Pr   rO  enumtypingr   r   r   torch._guardsr   r   r   r   r
   bytecode_transformationr   r   r,   LOCAL_SPECIALIZED_NN_MODULEri   GLOBAL_SPECIALIZED_NN_MODULELOCAL_UNSPECIALIZED_NN_MODULEGLOBAL_UNSPECIALIZED_NN_MODULE%LOCAL_UNSPECIALIZED_BUILTIN_NN_MODULE&GLOBAL_UNSPECIALIZED_BUILTIN_NN_MODULELOCAL_FSDP_MODULEGLOBAL_FSDP_MODULEr   r  r  r  r   	dataclassr   r?   rJ   r]   rq   r   r   r   r   r   r   r   r   Enumr   r   r   r  r  r!  r+  r5  r=  rQ  rg  ry  r  r  r  r  r  r  r  r  r  r  r   r  r   r  r  r  r  r#  r%  r'  r)  r+  r4   r*   r   <module>r:     sf	  &   ' ' < <  M {>>@@++[-T-T,,k.V.V--{/X/X..0Z0Z55{7h7h668j8j!!;#@#@""K$B$B' #  {@@BB--{/X/X..0Z0Z++[-V-V,,k.X.X55{7h7h668j8j!!;#@#@""K$B$B) %" {HHJJ--{/`/`..0b0b++[-^-^,,k.`.`55{7h7h668j8j!!;#@#@""K$B$B1 - {4466++[-J-J,,k.L.L!!;#@#@""K$B$B--{/L/L..0N0N55{7T7T668V8V 	 d#-& - $-8 d#
76 
7 $
7 d#8 8 $8 d#
.6 
. $
. d#0& 0 $0$ d#	' 	' $	' d#	0 	 $	 d#3 3 $32 d#Lf L $L, d#3 3 $3 d#M
 M $M d#	Z 	 $	 d#f  $ $TYY $ d#%;= %; $%;P d#3M 3 $3  d#/= / $/ d#G} G $G d#7 7 $7 d#>m > $>0m 0 d#!] ! $!H d##9M #9 $#9L d#  $. d#'9 '9 $'9T d#%K %K $%KP d#
P 
P $
P d#+ + $+  d# m   $  d# ]   $  d#O. O $O
 d#W)D W $W
 d#C C $C
 d#" " $" d#"6 " $"0 d#"V " $"  d#
> 
> $
> d#(] ( $( d#( ( $( d#(= ( $( d#%V % $% d#*& * $* >C  f  V V r*   