
    Bj{                       d Z ddlZddlZddlZddlZddlZddlZddlZddlZddl	Z	ddl
Z
ddlZddlZddlZddlmZmZ ddlmZ ddlmZ ddlmZmZmZmZmZ ddlZddlmZ ddlmc m Z! ddl"m#Z# ddl$m%Z% dd	l&m'Z' d
dl(m)Z)m*Z*m+Z+m,Z, d
dl-m.Z.m/Z/m0Z0 d
dl1m2Z2 d
dl3m4Z4m5Z5m6Z6 d
dl7m8Z8m9Z9 d
dl:m;Z; d
dl<m=Z=m>Z>m?Z?m@Z@mAZAmBZB d
dlmCZCmDZDmEZEmFZFmGZGmHZHmIZI ddlJmKZKmLZLmMZM ddlNmOZO ddlPmQZQmRZR ddlSmTZTmUZUmVZV erddlWmXZX ddlYmZZZ  G d deM          Z[ G d deM          Z\ G d deM          Z] G d  d!eM          Z^ G d" d#eM          Z_ G d$ d%e_          Z` G d& d'eM          Za G d( d)eM          Zb G d* d+eM          Zcd,ed-ed.ef         fd/Zd G d0 d1eM          Zeejf         G d2 d3                      Zg G d4 d5eV          Zh G d6 d7eV          Zi G d8 d9eM          Zj G d: d;eM          Zk G d< d=eM          Zl G d> d?eM          Zm G d@ dAeM          Zn G dB dCeM          Zo ejp        dD          d-eqej        ejr        f         fdE            Zs ejp        dD          d-eqejr        ej        f         fdF            Zt G dG dHeM          Zu G dI dJeM          Zv G dK dLeM          Zw G dM dNeM          Zx G dO dPeM          Zy G dQ dReM          Zz G dS dTeM          Z{ G dU dVeM          Z| G dW dXeM          Z} G dY dZe}          Z~ G d[ d\e}          Zej        e}ej        e}ej        eiZ G d] d^eM          Z G d_ d`eM          Z G da dbeM          ZdS )ca  
This module contains miscellaneous variable tracker implementations for various Python types
and features used in Dynamo's symbolic execution. These classes help track and propagate
information about different kinds of variables during graph capture.

Key classes include:
- SuperVariable: Handles super() calls and method resolution
- ExceptionVariable: Tracks exception objects
- RandomVariable: Manages random number generators
- GetAttrVariable: Tracks attribute access
- MethodWrapperVariable: Handles method wrappers
- PythonModuleVariable: Tracks Python modules
- NumpyVariable: Handles numpy functions and types
- StringFormatVariable: Manages string formatting
- DebuggingVariable: Handles print and logging
    N)CallableSequence)Random)BuiltinFunctionType)AnyLiteralTYPE_CHECKING	TypeGuardUnion)MutationType)TupleVariable)Source   )configgraph_break_hintstrace_rules	variables)create_call_functioncreate_call_function_excreate_instruction)$do_not_convert_to_tracable_parameter)raise_observed_exceptionraise_type_errorunimplemented)GuardBuilderinstall_guard)unpatched_nn_module_init)
AttrSourceGenericAttrSourceGetItemSourceTypeMROSource
TypeSourceWeakRefCallSource)check_unspec_or_constant_argscmp_name_to_op_mappingidentityis_tensor_base_attr_getteristypeproxy_args_kwargsraise_args_mismatch   )#AsPythonConstantNotImplementedErrorNO_SUCH_SUBOBJVariableTracker)ConstantVariable)NestedUserFunctionVariableUserFunctionVariable)call_random_fnis_standard_setattrUserDefinedObjectVariable)	PyCodegen)InstructionTranslatorc            
            e Zd ZeZh ej        Z	 ddededz  deddf fdZde	fdZ
d fd
Zdddedeeedz  f         f fdZdddedefdZdddedee         deeef         def
dZ xZS )SuperVariableNtypevarobjvarkwargsreturnc                 V     t                      j        di | || _        || _        d S N )super__init__r9   r:   )selfr9   r:   r;   	__class__s       _/var/www/html/Carbon-Document/venv/lib/python3.11/site-packages/torch/_dynamo/variables/misc.pyrA   zSuperVariable.__init__Y   s8     	""6""" 
     c                     t           j        S N)builtinsr@   rB   s    rD   python_typezSuperVariable.python_typej   s
    ~rE   codegenr5   c                                         fd            | j                   | j        5 | j                                       t	          dd                     d S                     t	          dd                     d S )Nc                  H     t          j        t                              S rG   )r   BuiltinVariabler@   )rC   rK   s   rD   <lambda>z+SuperVariable.reconstruct.<locals>.<lambda>n   s    ggi.G.N.N&O&O rE   r   Fr+   )add_push_nullr9   r:   extend_outputr   rB   rK   rC   s    `rD   reconstructzSuperVariable.reconstructm   s    OOOOOPPP;"GDK   !!"6q%"@"@AAAAA!!"6q%"@"@AAAAArE   txr6   namec                    | j         s(t          ddd| d| j         d| j          ddg           | j         J | j                                        }| j                                         }| j         j        rt          | j         j                  nd }t          |t                    r| j         j	        }| j         j        }d }|j
        }	 |                    |          d	z   }n0# t          $ r# t          t          ||          |          d fcY S w xY wt          |t!          |                    D ]f}	||	         j                            |t&                    x}
r<|
t&          ur3|r+t)          t+          t-          |          |	          |          }|
|fc S gt          d
dd| d| j         d| j          dddg           d S )Nz1-arg super not implemented z"Dynamo failed to trace attribute `z$` accessed via `super()` (for type `z` and object `z4`) because one-argument of super() is not supported.z-Use two-argument super(type, object_or_type).gb_typecontextexplanationhintsr+   zUnable to resolve super getattrz8`) because the resolved attribute type is not supported.z0Ensure the attribute exists in the parent class.z(Check the arguments passed to `super()`.)r:   r   r9   as_python_constantrJ   sourcer"   
issubclasstypevalue__mro__index
ValueErrorgetattrr@   rangelen__dict__getr-   r   r    r!   )rB   rT   rU   search_typetype_to_usetype_to_use_sourcer^   
search_mrostart_indexrc   resolved_getattrrC   s              rD   _resolved_getattr_and_sourcez*SuperVariable._resolved_getattr_and_sourcev   sr    { 
	5D D D,0LD DHLD D D D	 	 	 	 {&&&l5577 k--//.2k.@JJt{)***d 	 k4(( 	4++K!%!3 (
	H$**;77!;KK 	H 	H 	H 5k::DAA4GGGG	H ;J88 
	4 
	4E#-e#4#=#A#A$#W#WW 4#>99) !+)-8J*K*KUSS " " ,V33335DT D D(,D DDHKD D D C:
	
 
	
 
	
 
	
 
	
 
	
s   
C# #*DDc                 N   |                      ||          \  }}t          j                            |          st	          | |t          |                    S |r,t          |                    t          j	                             t          j        
                    ||          S )Npy_typer^   )rp   r   r/   
is_literalGetAttrVariabler`   r   
make_guardr   CONSTANT_MATCHcreate)rB   rT   rU   ra   r^   s        rD   var_getattrzSuperVariable.var_getattr   s     99"dCCv)44U;; 	D"4tE{{CCCC 	J&++L,GHHIII)00v0FFFrE   argsc                    |                      ||          \  }}| j        J |t          j        u rt	          t
                    S |t          j        j        j        u r| j        }ddl	m
} t          |t          j                  r}t          |j        |          rh|sf|sdt                      5  t!          j        |t$          |          }	|	                    || j        g|z   |          cd d d            S # 1 swxY w Y   nt)          dd|  d| d| d| d| dg t*          j        	           n| j        j        rt1          |d
          r|j        dk    rt          j                            |          rq|j        }
t1          |
d          r!|
j        dk    rt!          j        ||
          }n!|J |j        }t          j        |
|          }|                    |d||          S t          |t@                    rSt          |j!        tD          j#                  r4t!          j        ||j!        |d          }	|	                    |||          S t          |tH                    rt          |j!        tD          j#                  rt          | j        t          j                  r| j        }nGd }| j        j        rtK          | j        j                  }t!          j        || j        j&        |          }|J t!          j        ||j!        tO          |d          d          }	|	                    ||g||          S t          |tD          j#                  r8t!          j        |||d          }	|	                    || j        g|z   |          S t          |tD          j(                  r6t          j)        |j!        | j        |                              |||          S tU          |          r0t          | j        t                    r | j        j+        |g|R i |S |t          j,        u r|d         }	 |-                                }n2# t\          $ r%}t)          dd|  d| ddg|           Y d }~nd }~ww xY w|j/        j	        0                    | j                  s/t)          dd|  d| d| j         ddgt*          j1        	           t          |td                    sJ |j/        j	        3                    | j        |t          j4                               t          j5        6                    d           S t          | j        t          j                  rH| j        j7        <| j        j8        0|| j        j8        v r"| j        j7                            ||||          S |t          j9        u r|d         j:        }|j/        j	        ;                    | j        |          rX|j/        j	        <                    | j        |d          }t          |t          j4                  rt{          t|          |           |S d }	 t          9                    | j        j:        |          }n%# t|          $ r t{          t|          |           Y nw xY wd }| j        j        t          | j        j        |          }t!          j        |||          S |t          j@        jA        u r|d         }i }|d         jB        }|d         jB        B                                D ]#\  }}|jC        -                                }|||<   $|jD        jE        }d|jD        _E        	 |                    |||          ||jD        _E        S # ||jD        _E        w xY wt          |tD          jF                  rM|t          jH                    v r8t!          j        |||d          }|                    || j        g|z   |          S t)          dd|  d| d| d | d!d"g	           d S )#Nr   )AttributeMutationNewrt   z#Unsupported super().__init__() callcall_method  z0Dynamo encountered a super().__init__() call on zK that resolved to a `torch.nn.Module.__init__()` call that we cannot trace.rX   __name____new__
__module__rH   Tr^   realize__func__r   z7Non-constant attribute given to `super().__delattr__()`zbDynamo requires the attribute name passed to `super().__delattr__(...)` to be a constant (string).zEEnsure the attribute name is a string literal or a constant variable.)rY   rZ   r[   r\   from_exczFAttempted super().__delattr__() on an object without mutation trackingzmDynamo needs to track mutations on an object before `super().__delattr__` can be used on it. But the object (z3) doesn't have attribute mutation tracking enabled.z<Ensure the object is tracked by Dynamo's side effect system.)
deleted_ok   F)r   zFAttempted to call a super() attribute that is not a function or methodz4Dynamo does not know how to trace the call `super().z()` because `super().z(` is not a function or method attribute.zMEnsure the attribute accessed via `super()` is a standard method or function.)Irp   r:   objectrA   LambdaVariabler&   torchnnModuleside_effectsr}   
isinstancer   r4   mutation_typer   r.   buildr   call_functionr   r   	DIFFICULTr^   hasattrr   UserDefinedClassVariableis_supported_new_method__self__r   membercall_methodstaticmethodr   typesFunctionTypeclassmethodr"   
value_typer   
MethodTypeUserMethodVariabler3   method_setattr_standard__delattr__r]   NotImplementedErroroutputis_attribute_mutation
DYNAMO_BUGstr
store_attrDeletedVariabler/   ry   _base_vt_base_methods__getattribute__ra   has_pending_mutation_of_attr	load_attrr   AttributeErrorr   _C_disabled_torch_function_implitemsvtsymbolic_torch_function_statetorch_function_subclass_enabledMethodDescriptorTyper   get_tensor_method)rB   rT   rU   r{   r;   inner_fnr^   r:   r}   fn_vtuser_clsuser_cls_vtuser_cls_sourcecls_variable
cls_sourceattrexc	attr_nameresult
attr_valueattr_sourcefunc	tf_kwargstf_argshash_key_vtvalue_vtkey_strtx_oldfn_vars                                rD   r   zSuperVariable.call_method   s	     <<RFF&{&&& v&&!(+++111[F;;;;;; 69#FGGv35IJJ  "(
 :;; Q Q+14V  E !..rDK=43GPP	Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q AH4HH$HHHHHH!1 !1 !1 !1 9-78     Kb	J*--b	J !Y..2JJ8TT /  (Hx.. 83F*3T3T/>/DR/R/R)))"(-'@_   **2y$GGG,// R	JJu15
 5
 R	J $)H%fd  E &&r4888+.. K	J:u14
 4
 K	J $+y'IJJ   ${ "
;% @!+DK,>!?!?J.4K*   
 %%%#)!!&*55	  E &&rL+@4+@&III%"455 l	J#)"hvtTTTE&&rDK=4+?HHH%"233 i	J/!4;v  mBf--. !** e	JzK20
 0
 e	J 74;6rKDKKKFKKK+++7D..00& 
 
 
U8488$88!L ` !	 	 	 	 	 	 	 	 	
 9)??LL d8488$88!(#{!( !( !(
 W*5    dC(((((I"--T9#<#>#>   -44T:::t{I$GHH?	J$0)5DK555;'33BdFKKK000 QIy%BBY   /99Kt :   fi&?@@ A,^R@@@J= $44T[5F	RR

! = = =(<<<<<= K{!-/0BINN"(ZEEE???
 7DI1gmG)-a)<)<)>)> . .%X%.;;==%-	'""5UFOTB,L))"gyAA  0PP 0P     x!;<<	JK9;;;; %*2xNNNF''T[MD,@&III'0400$00,, ,37, , , `
	
 
	
 
	
 
	
 
	
 
	
sH   +<C44C8;C85Q
 

Q9Q44Q9?%Y% %ZZ]1 1]?rG   rK   r5   r<   N)r   r   __qualname__r@   _cpython_typer.   _nonvar_fieldsr   rA   r`   rJ   rS   r   tupler   rp   rz   listdictr   __classcell__rC   s   @rD   r8   r8   Q   s       M		'N *.    $& 	
 
     "T    B B B B B BB
)B
14B
	sJ%%	&B
 B
 B
 B
 B
 B
HG5 GS G_ G G G G X
#X
 X
 ?#	X

 S/)*X
 
X
 X
 X
 X
 X
 X
 X
 X
rE   r8   c                   Z     e Zd Zdej        deddf fdZdefdZddd	e	de
f fd
Z xZS )FrameSummaryVariableframe_summaryr;   r<   Nc                 H     t                      j        di | || _        d S r>   )r@   rA   r   )rB   r   r;   rC   s      rD   rA   zFrameSummaryVariable.__init__  s.    ""6"""*rE   c                     t           j        S rG   )	tracebackFrameSummaryrI   s    rD   rJ   z FrameSummaryVariable.python_type  s    %%rE   rT   r6   rU   c                 p   |dk    rt          j        || j        j                  S |dk    rt          j        || j        j                  S |dk    rt          j        || j        j                  S |dk    rt          j        || j        j                  S t                                          ||          S )NlinenofilenamerU   line)	r.   r   r   r   r   rU   r   r@   rz   rB   rT   rU   rC   s      rD   rz   z FrameSummaryVariable.var_getattr  s    8"(T-?-FGGGZ"(T-?-HIIIV^^"(T-?-DEEEV^^"(T-?-DEEEww""2t,,,rE   )r   r   r   r   r   r   rA   r`   rJ   r   r.   rz   r   r   s   @rD   r   r     s        +i&< + +PT + + + + + +&T & & & &	-5 	-S 	-_ 	- 	- 	- 	- 	- 	- 	- 	- 	- 	-rE   r   c            
       d    e Zd Zdeded ef         deddf fdZede	j
        ded ef         dd fd            Zededefd	            Zdee	j
        ez           fd
ZdedefdZdeej                 fdZdddededefdZdddedef fdZdddedee         deeef         def
 fdZ xZS )TracebackVariabler   tb_nextr;   r<   Nc                 ^     t                      j        di | || _        |J || _        d S r>   )r@   rA   r   r   )rB   r   r   r;   rC   s       rD   rA   zTracebackVariable.__init__  sA     	""6"""*"""rE   c                 6     | t          |          |          S )N)r   )r   )clsr   r   s      rD   from_frame_summaryz$TracebackVariable.from_frame_summary  s"     s'66HHHHrE   objc                 T    t          | t                    p|                                 S rG   )r(   r   is_constant_noner   s    rD   is_valid_tracebackz$TracebackVariable.is_valid_traceback  s$    c,--G1E1E1G1GGrE   c                     t          | j        t                    r| j        gS | j        g| j                                        z   S rG   )r(   r   r/   r   
extract_tbrI   s    rD   r   zTracebackVariable.extract_tb  sB    $, 011 	(&''"#dl&=&=&?&???rE   tbc                 x    | }t          |t                    r"||u rdS |j        }t          |t                    "dS )NTF)r(   r   r   )rB   r   curr_tbs      rD   has_reference_cyclez%TracebackVariable.has_reference_cycle  sK    8<W/00 	&"}}toG W/00 	& urE   c                     t           j        S rG   )r   TracebackTyperI   s    rD   rJ   zTracebackVariable.python_type  s    ""rE   rT   r6   name_varvalc                    |                                 }|dk    r|                     |          st          t          |           t	          |t
          t          f          sJ |                     |          s*t          |t
                    r*|                    |           rt          t          |           || _
        t          j                            d           S )Nr   )r]   r   r   	TypeErrorr   r   r/   r   r(   rd   r   r   ry   rB   rT   r   r   rU   s        rD   call_setattrzTracebackVariable.call_setattr  s     **,,9**3// 8(B777c$57G#HIIIII'',, 9s-..9363J3J43P3P9 )R888DL)00666rE   rU   c                    |dk    r| j         S |dk    r| j                            |d          S |dk    r| j        S |dk    r"t          d|  ddg t          j        	           t                                          ||          S )
Nr   	tb_linenor   r   tb_lastiz traceback.tb_lasti not supportedz accessing 'tb_lasti'zNDynamo does not support accessing the tb_lasti attribute of traceback objects.rX   )r   r   rz   r   r   SUPPORTABLEr@   r   s      rD   rz   zTracebackVariable.var_getattr  s    9<[  %11"h???_$$%%Z:666l6)56	    ww""2t,,,rE   r{   c                     |dk    rt          j        || |d         u           S |dk    r | j        |g|R  S t                                          ||||          S )N__eq__r   __setattr__)r.   r   r   r@   r   )rB   rT   rU   r{   r;   rC   s        rD   r   zTracebackVariable.call_method  so     8"(TT!W_===]""$4$R/$////ww""2tT6:::rE   )r   r   r   r   r   r/   r   rA   r   r   r   r   r   r.   boolr   r   r   r   r`   r   r   rJ   r   r   rz   r   r   r   r   s   @rD   r   r     s3       + *,<<= 	
 
     $ I -I *,<<=I 
	I I I [I H HD H H H \H@D!7:N!NO @ @ @ @
o $    #T%"56 # # # #7#7 "7 	7
 
7 7 7 7$-5 -S -_ - - - - - - ;#; ; ?#	;
 S/)*; 
; ; ; ; ; ; ; ; ; ;rE   r   c                   >    e Zd ZeZ	 	 	 ddedee         dee	ef         dz  de
dz  dedz  ddf fdZd	eddfd
ZddZdefdZdddededefdZddde	dee         dee	ef         def
 fdZddde	def fdZde	fdZeZedede	fd            Zde	fdZ xZS )ExceptionVariableNexc_typer{   init_kwargsr^   r   r<   c                    t                                          ||           || _        || _        |r$t	          d|  d| dg t
          j                   t          j        d           | _	        t          j        d           | _
        t          j        d          | _        t          j        d           | _        d | _        d S )N)r^   r   z,Keyword args passed to exception constructorz with kwargs zRDynamo does not know how to handle keyword args passed to an exception constructorrX   F)r@   rA   r  r{   r   r   r   r/   ry   __context__	__cause____suppress_context____traceback__python_stack)rB   r  r{   r  r^   r   rC   s         rD   rA   zExceptionVariable.__init__  s     	mDDD 	 	F;;k;;p6)56	    -=,CD,I,I +;*A$*G*G5E5LU5S5S!.>.Ed.K.K <@rE   rZ   c                     || _         d S rG   )r	  )rB   rZ   s     rD   set_contextzExceptionVariable.set_context@  s    "rE   rK   r5   c                 "                          fd                                j                                       t	           j                  d           dt
          dd f fd} |d            |d            |d           d S )	Nc                  D                          dj        j                  S )NrH   )load_import_fromr  r   rK   rB   s   rD   rO   z/ExceptionVariable.reconstruct.<locals>.<lambda>E  s    G,,Z9OPP rE   FrU   r<   c                 4   t          |           }t          |t                    r|j        dv s
J |            d S                                   |                                                   d                                         |            d S )N)TFNr   )re   r(   r/   ra   dup_toprQ   rot_nr   )rU   r   rK   rB   s     rD   codegen_attrz3ExceptionVariable.reconstruct.<locals>.codegen_attrJ  s    4&&Dd,-- )z%8888$88888!!!%%gmmA&6&6777""4(((((rE   r	  r
  r  )rP   foreachr{   r   rg   r   )rB   rK   r  s   `` rD   rS   zExceptionVariable.reconstructC  s    PPPPP	
 	
 	
 		"""c$)nne444	)s 	)t 	) 	) 	) 	) 	) 	) 	) 	]###[!!!+,,,,,rE   c                     | j         S rG   )r  rI   s    rD   rJ   zExceptionVariable.python_typeX  s
    }rE   rT   r6   r   r   c                    |                                 }|dk    ri|                                s>t          |t          j        t          j        t          j        f          sJ | d            |                     |           n)|dk    r|                                s<t          |t          j        t          j        t          j        t          j        f          r,|| _	        t          j
                            d          | _        nt          |d           n|dk    r/|                    dd          r|| _        nrt          |d           na|dk    r2t                              |          st          |d	           || _        n)t%          d
d|  d| d| dg t&          j                   t          j
                            d           S )Nr	  z! is not a valid exception contextr
  Tz9exception cause must be None or derive from BaseExceptionr  Fr  z0__traceback__ must be a traceback object or Nonez4Unsupported attribute assignment on Exception objectzcall_setattr r   z/Dynamo does not support setting the attribute 'z{' on tracked exception objects. Only `__context__`, `__cause__`, `__suppress_context__`, and `__traceback__` are supported.rX   )r]   r   r   r   r  !UserDefinedExceptionClassVariable"UserDefinedExceptionObjectVariabler  rN   r
  r/   ry   r  r   is_constant_matchr   r   r  r   r   r   r   s        rD   r   zExceptionVariable.call_setattr[  s)    **,,=  '')) 9Z/?@. . 9 9 8889 9  S!!!![  ##%% -/?@	* *  "%,5,F,M,Md,S,S)) S    +++$$T511 ,/)) S    _$$$77<< Y %WXXX!$DN555t55ZZ Z Z 7)56    )00666rE   rU   r;   c                     |dk    r | j         |g|R  S |dk    r0|\  }|                      |t          j        |d          |           | S t                                          ||||          S )Nr  with_tracebackr  )r   r.   r   r@   r   )rB   rT   rU   r{   r;   r   rC   s         rD   r   zExceptionVariable.call_method  s     =  $4$R/$////%%%DRb/"7O"L"LbQQQK77&&r4v>>>rE   c                    |dk    rt          j        || j                  S |dk    r| j        S |dk    r| j        S |dk    r| j        S |dk    r| j        S |dk    rCt          j        |t          | j                  | j	        ot          | j	        d                    S t                                          ||          S )NrC   r	  r
  r  r  r{   rt   )r.   r   r  r	  r
  r  r  r   r{   r^   r   r@   rz   r   s      rD   rz   zExceptionVariable.var_getattr  s    ;"(T];;;]""##[  >!+++,,_$$%%V^^"(di  {Fz$+v'F'F   
 ww""2t,,,rE   c                 0    | j         j         d| j         dS )N())rC   r   r  rI   s    rD   __str__zExceptionVariable.__str__  s     .)<<DM<<<<rE   argc                     	 t          |                                           S # t          $ r |                                 cY S w xY wrG   )reprr]   	Exception
debug_reprr%  s    rD   _debug_format_argz#ExceptionVariable._debug_format_arg  sR    	$..00111 	$ 	$ 	$>>#####	$s    # AAc                      d                      fd j        D                       }                                  d| dS )N, c              3   B   K   | ]}                     |          V  d S rG   r+  .0r%  rB   s     rD   	<genexpr>z/ExceptionVariable.debug_repr.<locals>.<genexpr>  s1      JJ//44JJJJJJrE   r"  r#  )joinr{   python_type_name)rB   r{   s   ` rD   r)  zExceptionVariable.debug_repr  sN    yyJJJJ	JJJJJ''))33D3333rE   )NNNr   )r   r   r   BaseExceptionr   r   r   r.   r   r   r   r   rA   r  rS   r`   rJ   r   r   rz   r$  __repr__r   r+  r)  r   r   s   @rD   r  r    s=       !M :> $-1@ @@ ?#@ #./$6	@
 @ $d*@ 
@ @ @ @ @ @B#? #t # # # #- - - -*T    67#67 "67 	67
 
67 67 67 67p?#? ? ?#	?
 S/)*? 
? ? ? ? ? ? -5 -S -_ - - - - - -&= = = = = H$ $3 $ $ $ \$4C 4 4 4 4 4 4 4 4rE   r  c                       e Zd ZdZdS )UnknownVariablez
    It could be anything!
    Nr   r   r   __doc__r?   rE   rD   r8  r8    s           rE   r8  c                   l     e Zd ZdZddedz  deddf fdZddd	ee         de	eef         defd
Z
 xZS )DelayGraphBreakVariablez^
    Used to insert a dummy variable in the stack to do the graph break at CALL_FUNCTION.
    Nmsgr;   r<   c                 H     t                      j        di | || _        d S r>   )r@   rA   r=  )rB   r=  r;   rC   s      rD   rA   z DelayGraphBreakVariable.__init__  s+    ""6"""rE   rT   r6   r{   c                 |    | j         dn| j         j        }t          dd| j          d| d| j         g            d S )NrW   z#Unsupported function call (delayed)zsource: z@Dynamo determined that a graph break should occur when calling `z`. Reason: rX   )r^   rU   r   r=  )rB   rT   r{   r;   rU   s        rD   r   z%DelayGraphBreakVariable.call_function  sn     [(rrdk.>9,t{,,9!9 9.2h9 9	
 	
 	
 	
 	
 	
rE   rG   )r   r   r   r:  r   r   rA   r   r.   r   r   r   r   s   @rD   r<  r<    s          C$J        
#
 '
 S/)*	

 

 
 
 
 
 
 
 
rE   r<  c                   b    e Zd ZdZddZddd	edefd
Zdddee         de	eef         defdZ
dS )ComptimeVariableza
    This variable is special, it lets you execute arbitrary code at
    Dynamo compile time
    rK   r5   r<   Nc                      t          d          )Nzcomptime is special form)r   rB   rK   s     rD   rS   zComptimeVariable.reconstruct  s    !"<===rE   rT   r6   rU   c                     ddl m } | j        J t          j        |t	          ||          t          | j        |                    S )Nr   )comptimert   )rE  r^   r.   r   re   r   )rB   rT   rU   rE  s       rD   rz   zComptimeVariable.var_getattr  sY    ''''''{&&&$$''
4;0M0M
 
 
 	
rE   r{   r;   c           
         ddl m} |st          |          dk    r2t          |ddt          |           dt          |           d           |d         }t	          |t
                    r' |                                 ||                     nt	          |t                    r|                                }|j	        rt          |d|j                    t          j        ||j        |j                                        |j        rt%          |j        j                  nd d	          } | ||                     nt)          d
t+          |                     t,          j                            d           S )Nr   )ComptimeContextz
comptime()zat most 2 args and 0 kwargsz
 args and  kwargsr   zOcomptime function must not have free variables, but these variables were free: r?   z"unsupported argument to comptime: )rE  rG  rg   r*   r   r1   get_functionr0   get_codeclosurer   co_freevarsr   r   	f_globalsfn_namer]   defaultsr   r   RuntimeErrorr`   r   r/   ry   )rB   rT   r{   r;   rG  fncoder   s           rD   r   zComptimeVariable.call_function  s    	/.....  	SYY]]-t99<<F<<<	   !Wb.// 	PBOOoob112222677 	P;;==Dz  xfjfvxx   %
--//,.KAbk'(((T  D D$$%%%%NDHHNNOOO)00666rE   r   )r   r   r   r:  rS   r   r.   rz   r   r   r   r?   rE   rD   rA  rA    s         
> > > >
5 
S 
_ 
 
 
 
-7#-7 '-7 S/)*	-7
 
-7 -7 -7 -7 -7 -7rE   rA  c                   v     e Zd ZU ej        Zedz  ed<   dZe	dz  ed<   	 ddedz  de
ddf fdZdefdZ xZS )	CellVariableNpre_existing_contents
local_namer;   r<   c                 H     t                      j        di | || _        d S r>   )r@   rA   rU  )rB   rU  r;   rC   s      rD   rA   zCellVariable.__init__8  s0     	""6"""%:"""rE   c                     t           j        S rG   )r   CellTyperI   s    rD   rJ   zCellVariable.python_type>  s
    ~rE   rG   )r   r   r   r   rY  r   r.   __annotations__rV  r   r   rA   r`   rJ   r   r   s   @rD   rT  rT  (  s         NM +T1111 "Jd
!!! ?C; ;%4t%;;NQ;	; ; ; ; ; ;T        rE   rT  c                   (     e Zd Zdeddf fdZ xZS )NewGlobalVariabler;   r<   Nc                 :     t                      j        di | d S r>   r@   rA   rB   r;   rC   s     rD   rA   zNewGlobalVariable.__init__C  &    ""6"""""rE   )r   r   r   r   rA   r   r   s   @rD   r\  r\  B  sK        # # # # # # # # # # # #rE   r\  fn_clsr<   .c                 V     dt           dt           dt           f fd}t          |_        |S )Nr{   r;   r<   c                       j         | i |S rG   )apply)r{   r;   ra  s     rD   trampoline_autograd_applyzDproduce_trampoline_autograd_apply.<locals>.trampoline_autograd_applyH  s    v|T,V,,,rE   )r   !produce_trampoline_autograd_apply_origin)ra  re  s   ` rD   rf  rf  G  sH    - - - - - - - - - )J%$$rE   c            
       $    e Zd ZdZdhej        Zdededdf fdZdefdZ	dd	d
e
e         deeef         defdZdd	d
e
e         deeef         defdZdd	d
ee         deeef         dd fdZdd	ded
e
e         deeef         def
dZ xZS )AutogradFunctionVariablez-represents a torch.autograd.Function subclassra  r;   r<   Nc                 H     t                      j        di | || _        d S r>   )r@   rA   ra  )rB   ra  r;   rC   s      rD   rA   z!AutogradFunctionVariable.__init__X  s+    ""6"""rE   c                     t           S rG   r`   rI   s    rD   rJ   z$AutogradFunctionVariable.python_type\      rE   rT   r6   r{   c           	      n   ddt           dd ffd}t          j        |||f           rt          j                    r| j        }ddlm} ddlm} | j	        j
        } || j	        j                  }	|	r" || j	        j
        | j	        j                  }d }| j	        j        }
|
t          j        j        j        ur*t          dd	|  d
| d
| dddgt           j                   | j	        j        }|t          j        j        j        ur)t          dd	|  d
| d
| ddgt           j                   ddlm} |9|	s7t+                              | j	        j                  | j	        j                  }|ot+          |d          } ||| j	        j        ||                              ||          }| j        r|	rt+          | j        d          }t7          |                    t:          j                             t+          | j        d          }t7          |                    t:          j                             |S | j        rt+          | j        d          }nd }| j	        j
        }t>                               ||          }|g|}tC          |tD          j#                  rutI          j%        |          }tM          |          dz
  tM          |j'                  k    r
|dd          }t          j(        ||d          }|                    ||          S tC          |tD          j)                  rHtU          j+        |j,        tU          j-        | j	                  |                              ||          S t          dd	|  d
| d
| dt]          |          j         ddg           d S )NFr   r<   c                     |                                  r| j        durdt          | t          j                  r|                               rdd S d S d S )NFT)	is_tensorrequires_gradr   r   NNModuleVariableis_training)r   rq  rT   s    rD   visitz2AutogradFunctionVariable.call_apply.<locals>.visitg  sl    ||~~ )#500$(M"i899 )>>"%% )$(MMM) )) )rE   r   )#autograd_function_forward_rewritten)_is_setup_context_definedzUnsupported custom vjpzcall_apply r   zgDynamo does not support tracing `torch.autograd.Function` subclasses that define a custom `vjp` method.z+Remove the custom `vjp` method if possible.z.Use standard `backward` instead if applicable.rX   zUnsupported custom jvpzgDynamo does not support tracing `torch.autograd.Function` subclasses that define a custom `jvp` method.z+Remove the custom `jvp` method if possible.r+   )AutogradFunctionApplyVariablerd  )r   rt   forwardsetup_contextTr   z=Non-function or method in subclass of torch.autograd.FunctionzDynamo requires the `forward` attribute of a `torch.autograd.Function` subclass to be a standard Python function or method. Found type `z
` instead.zPEnsure the `forward` method is defined as a regular function or instance method.)/r.   rt  r   is_grad_enabledr^   "torch._functorch.autograd_functionru  torch.autograd.functionrv  ra  rx  ry  vjpautogradFunctionr   r   r   jvphigher_order_opsrw  r   import_sourcer   r   backwardr   r   rw   r   CLOSURE_MATCHAutogradFunctionContextVariablery   r   r   r   inspect	signaturerg   
parametersr   r   r   r   r   r   r`   )rB   rT   r{   r;   rt  r^   ru  rv  
forward_fnis_setup_ctx_definedvjp_fnjvp_fnrw  apply_sourcer   fwd_srcsetup_ctx_srcrQ  ctxsigr   rq  s    `                   @rD   
call_applyz#AutogradFunctionVariable.call_apply_  s    	)o 	)$ 	) 	) 	) 	) 	) 	) 	) 	edF^444 F	U244 F	[F      JIIIII,J#<#<T[=V#W#W # 	 A@K')B 
 [_FU^48884@$@@@@@@!- FH +6    [_FU^48884@$@@@@@@!- F*6
 
 
 
 HGGGGG~&:~#$$T[%;<<dk>R  "Hj&H&H&HL//$#	  
 mBf--  { T3 T$T[)<<g001KLLMMM *4; H Hm66|7QRRSSSJ; 	Y77FFF[ -44RvFF|d|b%,-- 	#B''C4yy1}CN 3 333ABBx#)"bNNNE&&r4888E,-- 	/24;??   mBf--	. W<d<<T<<F<<Q37883DQ Q Q3
 
 
 
 
 
rE   c                 P   | j         j        }t          |d         j                  t          j        j        j        u sJ t          |t          j
                  sJ | j        J t          | j        d          }t          j        |||d          }|                    |||          S )Nr   r  Tr   )ra  r  r`   ra   r   _dynamoexternal_utilsFakeBackwardCFunctionr   r   r   r^   r   r.   r   r   )rB   rT   r{   r;   rQ  	fn_sourcer   s          rD   call_backwardz&AutogradFunctionVariable.call_backward  s     [!DGM""em&B&XXXXX"e011111{&&&t{J77	%b"YMMM""2tV444rE   c                 *    t          | j                  S rG   )ri  ra  rB   rT   r{   r;   s       rD   r   z&AutogradFunctionVariable.call_function  s     (444rE   rU   c           
         ddl m} |dk    rqt          j        | j                  rAt          | j                  } || |j        j        d|gt          ||          R            S | 	                    |||          S |dk    r| 
                    |||          S | j        t          | j        |          nd }	 t          j        | j        |          }n# t          $ r d }Y nw xY wt!          |t"                    r|                    | j                  }	t          j        |	          }
|
J |,|
                    |	|                              |||          S  |
|	                              |||          S t!          |t,                    r1t/          j        |j        | |                              |||          S t5          dd	|  d
| d| ddg           d S )Nr+   )wrap_fx_proxyrd  r   )rT   proxyr  rt   z$Unsupported autograd.Function methodr~   r   ,Dynamo does not support calling the method `z` directly on the `torch.autograd.Function` instance. Supported methods include `apply`, `backward`, static methods, and class methods.znEnsure the method is decorated with `@staticmethod` or `@classmethod` if it's meant to be called on the class.rX   )builderr  r   is_callable_allowedra  rf  r   create_proxyr)   r  r  r^   r   r  getattr_staticr   r   r   __get__lookupcreate_with_sourcer   r   r   r   r   r   )rB   rT   rU   r{   r;   r  re  r^   r   r   traceds              rD   r   z$AutogradFunctionVariable.call_method  sr    	+*****7??.t{;; 9,MK- -) %}0")0'1 +488      r4888Z%%b$7776:k6MZT222SWF,T[$??!    #|,, {{4;//$+D11)))% 11$v1FFTTf  "6$<<55b$GGGC--  3L$v  -D&112 B8488$88!9!9 !9 !9
U     s   :C C$#C$)r   r   r   r:  r.   r   r   rA   r`   rJ   r   r   r   r  r  r   r   r   r   r   s   @rD   ri  ri  P  s       77 			'N
s c d      T    {#{ ?#{ S/)*	{
 
{ { { {z5#5 ?#5 S/)*	5
 
5 5 5 55#5 '5 S/)*	5
 
$5 5 5 5@#@ @ ?#	@
 S/)*@ 
@ @ @ @ @ @ @ @rE   ri  c                   H    e Zd ZU  ej        e          Zee         ed<   dS )SavedTensorBox)default_factorytensorsN)	r   r   r   dataclassesfieldr   r  r.   rZ  r?   rE   rD   r  r  6  s7         %6[%6t%L%L%LGT/"LLLLLrE   r  c                   F    e Zd ZdZdddhej        Z	 	 	 	 	 ddededz  dededz  d	e	ed
f         dz  dedz  deddf fdZ
e	 	 ddddee         dz  deeef         dz  defd            ZdefdZdddedee         deeef         def
 fdZdddedef fdZ xZS )r  zZ
    Tracks an autograd.Function() context using mutation tracking in side_effects.py
    r  	inferencesaved_tensorsNFra   r   needs_input_grad.non_differentiabler;   r<   c                 x     t                      j        d||d| || _        || _        || _        || _        d S N)ra   r   r?   )r@   rA   r  r  r  r  )	rB   ra   r   r  r  r  r  r;   rC   s	           rD   rA   z(AutogradFunctionContextVariable.__init__G  sN     	FuFFvFFF"* 0"4rE   rT   r6   r{   c           
          d }|r|st          d |D                       }| j        j                            d t          j        j        j        t          j	        t          dt                      |          i           }|S )Nc              3   L   K   | ]}|                                 o|j        V   d S rG   )rp  rq  r1  xs     rD   r2  z9AutogradFunctionContextVariable.create.<locals>.<genexpr>`  s1      $U$U1Q[[]]%Fq$U$U$U$U$U$UrE   T)r  r  r  )r   r   r   track_object_newr   r~  functionFunctionCtx	functoolspartialr  r  )rT   r{   r;   r  outs        rD   ry   z&AutogradFunctionContextVariable.createW  s       	V 	V$$U$UPT$U$U$UUUi$55N#//,..!1	   

 

 
rE   c                 b    | j         "t          dd|  dg t          j                   | j         S )Nzproxy not setz	as_proxy zMDynamo requires the autograd.Function context to be initialized with a proxy.rX   )r  r   r   r   rI   s    rD   as_proxyz(AutogradFunctionContextVariable.as_proxyn  sJ    :'*D**25)45    zrE   rU   c                 *   |dk    r$t                                          ||||          S |dk    r^|r"t          ||dt          |           d           t	          |i           d         | _        t          j                            d           S |dk    r)t          dd|  d	| d
| dg t          j                   | j        t          dd|  d	| ddg           | j        J | j        s9|s| j        st          |d           |j        j                            | |           t          | j        j                  dk    rg | j        _        |D ]!}| j        j                            |           "t          j                            d           S )Nr  mark_non_differentiablez0 kwargsrH  r   save_for_backwardz,Unsupported autograd.Function context methodr~   r   r  z` on `autograd.Function` context objects. Supported methods are `__setattr__`, `save_for_backward` and `mark_non_differentiable`.rX   z9Unsupported autograd.Function context `save_for_backward`zjDynamo requires the `saved_tensors` attribute to be initialized on the `autograd.Function` context object.zEnsure that the `saved_tensors` attribute is properly initialized before calling `save_for_backward`. `save_for_backward` only supported on a newly constructed `torch.autograd.function.FunctionCtx`.z>save_for_backward() requires a source and no keyword arguments)r@   r   r*   rg   r)   r  r   r/   ry   r   r   r   r  r  r^   r   r   r   track_save_for_backwardr  append)rB   rT   rU   r{   r;   r%  rC   s         rD   r   z+AutogradFunctionContextVariable.call_method{  s    =  77&&r4v>>>... S#BjS[[:Q:Q:QRRR&7b&A&A!&DD#-44T:::&&&F4t44d44-- - - 7)56    %S4t44d44Ow
 
 
 
 !---~ 	G T[  X   I"::4FFF t!)**Q..)+D& 	3 	3C&--c2222)00666rE   c                     dv rt           fd          S dk    r2 j        +t          j        t	           j        j                            S dk    rg j        $t          j                             j                  S  j	        r5t           j	        d          }t          j         j        j        |          S t                                                    S )N)r  r  c                  N                         t          |           |          S rG   )r   r   )r{   r;   rU   rB   rT   s     rD   rO   z=AutogradFunctionContextVariable.var_getattr.<locals>.<lambda>  s!    (8(8T4::v(V(V rE   r  r  )r   r  r   r   r   r  r  r/   ry   r^   r   r.   r   ra   r@   rz   rB   rT   rU   r^   rC   s   ``` rD   rz   z+AutogradFunctionContextVariable.var_getattr  s    CCC!VVVVVV   ?""t'9'E*40B0J+K+KLLL%%%$0 1889NOOO{ V#DK1CDD&,R1LfUUUww""2t,,,rE   )NFNNNNN)r   r   r   r:  r4   r   r   r`   r  r   rA   r   r   r.   r   r   ry   r  r   r   rz   r   r   s   @rD   r  r  ;  s        
 	 
#	1	N #'$(48)-5 55 4K5 	5
 Tz5  c	*T15  $J5 5 
5 5 5 5 5 5   2648 #'$. S/)*T1 
	   \,#    37#37 37 ?#	37
 S/)*37 
37 37 37 37 37 37j-5 -S -_ - - - - - - - - - -rE   r  c            
            e Zd ZdZ	 ddej        j        deej        j                 dz  deddf fdZ	dd	d
e
dee         dee
ef         def
dZ xZS )AutogradEngineVariablez;
    Represents a torch._C._ImperativeEngine instance.
    Nra   r   r;   r<   c                 @     t                      j        d||d| d S r  r^  )rB   ra   r   r;   rC   s       rD   rA   zAutogradEngineVariable.__init__  s/     	FuFFvFFFFFrE   rT   r6   rU   r{   c                    |dk    rt           j        j        j        r}|j        s|j        s
J d            t          j        |t           j        j        j	        j
                  }|                    ||j        j                                        g|R |          S t          dd|  d| dg            d S t          dd|  d| d	| d
g            d S )Nqueue_callbackzXqueue_callback() is only supported when Compiled Autograd is enabled with fullgraph=Truez7Unsupported torch._C._ImperativeEngine.queue_callback()r~   r   zYqueue_callback() is only supported when Compiled Autograd is enabled with fullgraph=True.rX   z-Unsupported torch._C._ImperativeEngine methodzgDynamo only supports the `queue_callback` method on a torch._C._ImperativeEngine instance, but found: `z`.)r   r  compiled_autogradin_compiled_autograd_region	one_grapherror_on_graph_breakr.   r   r  FakeCompiledAutogradEnginer  r   r   r   get_ca_final_callbacks_varr   )rB   rT   rU   r{   r;   r   s         rD   r   z"AutogradEngineVariable.call_method  sF    ###}.J | r'>  n > (-M0KZ  **Y+FFHHP4PP   U8488$88!H      G4t44d44RIMR R R     rE   rG   )r   r   r   r:  r   r   _ImperativeEnginer`   r   rA   r   r   r.   r   r   r   r   s   @rD   r  r    s          ?CG Gx)G 34t;G 	G
 
G G G G G G%#% % ?#	%
 S/)*% 
% % % % % % % %rE   r  c                   |     e Zd Zdedef         deddf fdZdefdZdd	d
e	e         de
eef         defdZ xZS )r   rQ  .r;   r<   Nc                 H     t                      j        di | || _        d S r>   )r@   rA   rQ  )rB   rQ  r;   rC   s      rD   rA   zLambdaVariable.__init__  s+    ""6"""rE   c                     t           j        S rG   )r   r   rI   s    rD   rJ   zLambdaVariable.python_type  s    !!rE   rT   r6   r{   c                      | j         |i |S rG   )rQ  r  s       rD   r   zLambdaVariable.call_function   s     tw''''rE   )r   r   r   r   r.   r   rA   r`   rJ   r   r   r   r   r   r   s   @rD   r   r     s        8C$89 S T      "T " " " "(#( '( S/)*	(
 
( ( ( ( ( ( ( (rE   r   c                       e Zd Zddhej        Z	 ddedededz  deddf
 fdZdef fdZ	defd	Z
ed
ej        j        dedefd            ZdefdZdefdZdddedefdZddZdddee         deeef         defdZ xZS )rv   rU   rs   Nr   r;   r<   c                      t                      j        di | t          |t                    sJ t          |t                    sJ || _        || _        || _        d S r>   )r@   rA   r   r.   r   r   rU   rs   )rB   r   rU   rs   r;   rC   s        rD   rA   zGetAttrVariable.__init__  sg     	""6"""#/////$$$$$$	rE   c                 `    | j         | j         S t                                                      S rG   )rs   r@   rJ   rB   rC   s    rD   rJ   zGetAttrVariable.python_type  s)    <#<77&&(((rE   c                 @    | j         j         d| j         d| j         dS Nr"  r-  r#  )rC   r   r   rU   rI   s    rD   r6  zGetAttrVariable.__repr__$  s*    .)DDDHDD	DDDDrE   
base_proxyr   c                 "    t          | |          S rG   )re   )r  r   s     rD   create_getattr_proxyz$GetAttrVariable.create_getattr_proxy'  s    z4(((rE   c                 p    t                               | j                                        | j                  S rG   )rv   r  r   r  rU   rI   s    rD   r  zGetAttrVariable.as_proxy+  s(    33DH4E4E4G4GSSSrE   c                     | j                                         }	 t          || j                  S # t          $ r t          |  d          d w xY w)Nz is not a constant)r   r]   re   rU   r   r   )rB   constants     rD   r]   z"GetAttrVariable.as_python_constant.  se    8..00	M8TY/// 	M 	M 	M%&A&A&ABBL	Ms	   0 ArT   r6   c                 >   t          | j        t          j                  st          |j                            | j        j                  }| j        |j	        vrt          t          j        || j                  }||j	        vrt          t          j        ||          S rG   )r   r   r   rr  r   r   get_submodule
module_keyrU   rh   r  r  )rB   rT   rU   step1step2s        rD   const_getattrzGetAttrVariable.const_getattr5  s    $(I$>?? 	&%%	''(;<<9EN**%%&udi88u~%%%%%eT222rE   rK   r5   c                      || j                    |                    |                    | j                             d S rG   )r   rQ   create_load_attrsrU   rC  s     rD   rS   zGetAttrVariable.reconstruct@  s>    g77	BBCCCCCrE   r{   c                 `    | j                             || j        t          |          |          S rG   )r   r   rU   r   r  s       rD   r   zGetAttrVariable.call_functionD  s(     x##B	4::vFFFrE   rG   r   )r   r   r   r.   r   r   r`   r   rA   rJ   r6  r   r   fxProxyr  r  r]   r  rS   r   r   r   r   r   s   @rD   rv   rv   	  s        
	'N  $	   	
  
     )T ) ) ) ) ) )E# E E E E ) )s )s ) ) ) \)T# T T T TMC M M M M	3 7 	3s 	3s 	3 	3 	3 	3D D D DG#G 'G S/)*	G
 
G G G G G G G GrE   rv   c                        e Zd Zdej        deddf fdZdej        fdZddd	ee	         de
ee	f         de	f fd
Zded         fdZdej        fdZded         fdZdefdZdedefdZ xZS )MethodWrapperVariablemethod_wrapperr;   r<   Nc                 H     t                      j        di | || _        d S r>   )r@   rA   r  )rB   r  r;   rC   s      rD   rA   zMethodWrapperVariable.__init__N  s.    ""6""",rE   c                     | j         S rG   r  rI   s    rD   get_real_python_backed_valuez2MethodWrapperVariable.get_real_python_backed_valueR      ""rE   rT   r6   r{   c           	      (   t          | j                  rt          |d         t          j                  rat          |          dk    rt          |          dk    st          |d           |d                             || j        j        j	                  S | j        j        }| j        j	        }|dk    rWt          |          j        }|t          j        u r3t          j        |t                                        |||g||          S net           j        dk    r|t          j        d         u r|dk    rddlm} t          |          dk    rt|sr	 |                    ||                     |d                                                             S # t.          $ r t1          t.          |           Y nt2          $ r Y nw xY wt5          d	d
|  d| d| dddgt6          j                   n|t          j        d         u r|dk    s|t          j        d         u rS|dk    rL|t          j        d         u rdnd}t          |          dk    r|s	 |                     |d                                                   }	|d         j        }
|
||d                                         }t=          j        t          |          |          }||u r<|dk    rtA          |
          }
ntC          |
|          }
t          j        ||	|
          S ddlm} |                    ||	          S # t2          $ r Y nw xY wt5          d| dd
|  d| d| d| dd| dd| dgt6          j                   tE                      #                    |||          S )Nr   r+   z2tensor attribute getter takes exactly one argumentrA   )r      rZ  r  )SourcelessBuilderz9unsupported type.__dict__['__annotations__'].__get__ callzcall_function , args: 
, kwargs: z{`torch.compile` only supports calling type.__dict__['__annotations__'].__get__ on a single constant argument (i.e. a type).zpMake sure your call to type.__dict__['__annotations__'] only has one positional argument (no keyword arguments).zMake sure the argument to type.__dict__['__annotations__'] is a constant (i.e. type). For example, `object`, `int`, `MyCustomClass`.rX   rb   rh   zunsupported type.__dict__['z'].__get__ callz5`torch.compile` only supports calling type.__dict__['z7'].__get__ on a single constant argument (i.e. a type).z&Make sure your call to type.__dict__['zC'].__get__ only has one positional argument (no keyword arguments).z)Make sure the argument to type.__dict__['zT'].__get__ is a constant (i.e. type). For example, `object`, `int`, `MyCustomClass`.)$r'   r  r   r   TensorVariablerg   r   rz   r   r   r`   rA   r   r.   r   r   sysversion_inforh   r  r  ry   r]   r   r   r,   r   r   r   r^   r  r  r!   r   r@   r   )rB   rT   r{   r;   self_objwrapper_namefn_objr  r   ra   r^   cls_valstatic_descrC   s                rD   r   z#MethodWrapperVariable.call_functionU  s|    &d&9:: 	RzGY-@
 @
 	R IINNs6{{a'7'7 L   7&&r4+>+G+PQQQ &/*3 :%%(^^,F((&,R88DD %%   ) '' DM*;<<<	))2222224yyA~~f~,33D//Q0J0J0L0LMM   & A A A,^R@@@@@:   D SOOOtOOvOO?FR
 '2     $-	222|y7P7Pj111li6O6O &y)AAA		z  4yyA~~f~? //Q0J0J0L0LMME "!W^F)"&q'"<"<">">&-&<T']]I&V&V&(22(I55)6v)>)>)3FI)F)F#2#8UF#K#KK::::::,33B>>>% ;   D( PiPPPOOOtOOvOO?T] ? ? ?FY F F FR	 R R R
 '2    ww$$Rv666s+   A F G 4	G ?G -L' '
L43L4Tc                     dS NTr?   rI   s    rD   is_python_constantz(MethodWrapperVariable.is_python_constant      trE   c                     | j         S rG   r  rI   s    rD   r]   z(MethodWrapperVariable.as_python_constant  r  rE   c                     dS r  r?   rI   s    rD   is_python_hashablez(MethodWrapperVariable.is_python_hashable  r  rE   c                 D    t          |                                           S rG   hashr]   rI   s    rD   get_python_hashz%MethodWrapperVariable.get_python_hash      D++--...rE   otherc                     t          |t                    o)|                                 |                                k    S rG   r   r.   r]   rB   r  s     rD   is_python_equalz%MethodWrapperVariable.is_python_equal  :    uo.. H''))U-E-E-G-GG	
rE   )r   r   r   r   MethodWrapperTyper   rA   r  r   r.   r   r   r   r   r  r]   r  intr  r   r  r  r   r   s   @rD   r  r  M  sk       -u'> -# -RV - - - - - -#e.E # # # #o7#o7 'o7 S/)*	o7
 
o7 o7 o7 o7 o7 o7bGDM    #E$; # # # #GDM    / / / / /
V 
 
 
 
 
 
 
 
 
rE   r  c                        e Zd Zdej        deddf fdZdej        fdZddd	ede	f fd
Z
ded         fdZdej        fdZ xZS )GetSetDescriptorVariabledescr;   r<   Nc                 H     t                      j        di | || _        d S r>   )r@   rA   r  )rB   r  r;   rC   s      rD   rA   z!GetSetDescriptorVariable.__init__  s+    ""6"""			rE   c                     | j         S rG   r  rI   s    rD   r  z5GetSetDescriptorVariable.get_real_python_backed_value  
    yrE   rT   r6   rU   c                 ^   |dk    r<| j         r5t          | j         d          }t          j        || j        j        |          S |dv rE| j         ot          | j         |          }t          j        |t          | j        |          |          S t                                          ||          S )Nr  )__objclass__r   )	r^   r   r.   r   r  r  re   r@   rz   r  s       rD   rz   z$GetSetDescriptorVariable.var_getattr  s    9Y77F"(TY->GGG111[BZT%B%BF"(WTY-E-EvNNN77&&r4000rE   Tc                     dS r  r?   rI   s    rD   r  z+GetSetDescriptorVariable.is_python_constant  r  rE   c                     | j         S rG   r  rI   s    rD   r]   z+GetSetDescriptorVariable.as_python_constant  r  rE   )r   r   r   r   GetSetDescriptorTyper   rA   r  r   r.   rz   r   r  r]   r   r   s   @rD   r  r    s        U7 3 4      e.H    15 1S 1_ 1 1 1 1 1 1GDM    E$>        rE   r  c                        e Zd Zej        Zddhej        Zdej        deddf fdZ	de
ej                 fdZdej        fdZdej        fd	Zdefd
ZdddeddfdZdddedefdZ xZS )PythonModuleVariablera   is_torchr;   r<   Nc                      t                      j        di | || _        | j        t          u p| j        j                            d          | _        d S )Nztorch.r?   )r@   rA   ra   r   r   
startswithr'  rB   ra   r;   rC   s      rD   rA   zPythonModuleVariable.__init__  sQ    ""6"""

e+Wtz/B/M/Mh/W/WrE   c                     t           j        S rG   )r   
ModuleTyperI   s    rD   rJ   z PythonModuleVariable.python_type  s    rE   c                     | j         S rG   ra   rI   s    rD   r]   z'PythonModuleVariable.as_python_constant  
    zrE   c                     | j         S rG   r.  rI   s    rD   r  z1PythonModuleVariable.get_real_python_backed_value  r/  rE   c                     d| j          dS )NzPythonModuleVariable(r#  r.  rI   s    rD   r6  zPythonModuleVariable.__repr__
  s    4tz4444rE   rT   r6   rU   r/   c                 V    t          | j        |          }t          j        ||          S rG   )r   ra   r.   r   rB   rT   rU   r   s       rD   call_obj_hasattrz%PythonModuleVariable.call_obj_hasattr  s(     T**$R000rE   c                    |j         j                            | |          r |j         j                            | |          S d }| j        s|| j        j        vr<	 t          | j        |          }n7# t          $ r t          t          |           Y nw xY w| j        j        |         }| j
        ot          | j
        |          }t          j        |||          S rG   )r   r   r   r   r'  ra   rh   re   r   r   r^   r   r.   r   )rB   rT   rU   r   r^   s        rD   rz   z PythonModuleVariable.var_getattr  s    9!>>tTJJ 	@9)33D$???
= 	3D
(;;;=$TZ66

! = = =(<<<<<= ,T2J>DK!>!>$RV<<<s   A/ /BB)r   r   r   r   r,  r   r.   r   r   rA   r`   rJ   r]   r  r   r6  r4  rz   r   r   s   @rD   r&  r&    sR       $M 	 
	'NXe. X# X$ X X X X X X
 T%"23        E$4    e.>    5# 5 5 5 51)1141	1 1 1 1=5 =S =_ = = = = = = = =rE   r&  c            
            e Zd Zdededdf fdZdddedefd	Zddd
edee         de	eef         def
dZ
ddd
edefdZdefdZdefdZd fdZded         fdZdefdZdedefdZ xZS )TypingVariablera   r;   r<   Nc                 H     t                      j        di | || _        d S r>   r@   rA   ra   r*  s      rD   rA   zTypingVariable.__init__%  +    ""6"""


rE   rT   r6   keyc                 ^    | j         |                                         }t          |          S rG   )ra   r]   r7  )rB   rT   r;  
new_typings       rD   mp_subscript_implz TypingVariable.mp_subscript_impl)  s*     Z 6 6 8 89
j)))rE   rU   r{   c                 n   |dk    ret          |          dk    rR|sPt          |d         t                    o| j        |d         j        k    }t          j                            |          S t          dd| j         d| d| d| d	| d
| j         dd| d| j         dgt          j	                   d S )Nr  r+   r   z,unsupported method call on `typing` variableztyping variable: z, method name: r  r  z.`torch.compile` does not support method call `z` on `typing` variable f.zAvoid calling the z method on rX   )
rg   r(   r7  ra   r   r/   ry   r   r   r   )rB   rT   rU   r{   r;   r   s         rD   r   zTypingVariable.call_method3  s     84yyA~~f~Q88XTZ4PQ7==X 188@@@Bi
ii4iiQUiiagiitttgkgqtttCTCCdjCCC".		
 	
 	
 	
 	
 	
rE   c           
         ddl m}m} |t          v r7t	          j        | |t          t          | j        |                              S |j	        j
                            | |          r |j	        j
                            | |          S t          | j        |          }| j        r*t          | j        |          }  |||          |          S |                    ||          S )Nr+   )r  VariableBuilderrr   )r  r  rB  r%   r   rv   r`   re   ra   r   r   r   r   r^   r   ry   )rB   rT   rU   r  rB  ra   r   s          rD   rz   zTypingVariable.var_getattrH  s    ????????))),dDT)B)B$C$C    9!>>tTJJ 	@9)33D$???
D)); 	7$T[$77K3??2{33E:::$++B666rE   c                     | j         S rG   r.  rI   s    rD   r]   z!TypingVariable.as_python_constantZ  r/  rE   c                     | j         S rG   r.  rI   s    rD   r  z+TypingVariable.get_real_python_backed_value]  r/  rE   rK   r5   c                     t          | j        t          j                  s!t	                                          |          S |                    |                    | j                             d S rG   )r   ra   r   GenericAliasr@   rS   append_outputcreate_load_constrR   s     rD   rS   zTypingVariable.reconstruct`  s^    $*e&899 	077&&w///0 	g77
CCDDDDDrE   Tc                     dS r  r?   rI   s    rD   r  z!TypingVariable.is_python_hashable|  r  rE   c                 D    t          |                                           S rG   r  rI   s    rD   r  zTypingVariable.get_python_hash  r  rE   r  c                     t          |t                    o)|                                 |                                k    S rG   r  r  s     rD   r  zTypingVariable.is_python_equal  r  rE   r   )r   r   r   r   rA   r.   r>  r   r   r   r   rz   r]   r  rS   r   r  r  r  r   r  r  r   r   s   @rD   r7  r7  $  s       c S T      *#* * 
	* * * *
#
 
 ?#	

 S/)*
 

 
 
 
*75 7S 7_ 7 7 7 7$C    c    E E E E E E8GDM    / / / / /
V 
 
 
 
 
 
 
 
 
rE   r7  )maxsizec                      ddl m}  i }|                                 D ]K\  }}|j                                        D ],\  }}t	          |          rt          ||d          x}r|||<   -L|S )zd
    This generates a mapping from numpy modules to their torch._numpy
    modules equivalents.
    r   )NP_TO_TNP_MODULEN)utilsrN  r   rh   callablere   )rN  np_fn_to_tnp_fnnp_modtnp_modrN  tnp_fnnp_fns          rD   get_np_to_tnp_maprV    s     )(((((O+1133 4 4&/5577 	4 	4OGV 4 $FGT:::5 4-3OE*	4 rE   c                  Z    t                      } d |                                 D             S )z
    This is just the reverse mapping of get_np_to_tnp_map() - mapping from
    torch._numpy modules to numpy equivalents.
    c                     i | ]\  }}||	S r?   r?   r1  kvs      rD   
<dictcomp>z%get_tnp_to_np_map.<locals>.<dictcomp>  s    '''TQAq'''rE   )rV  r   )ms    rD   get_tnp_to_np_mapr^    s+     	A''QWWYY''''rE   c            
       \    e Zd ZdZej        fZdededdf fdZdefdZ	e
dej        defd	            Ze
dej        defd
            Zdddee         deeef         defdZdddedee         deeef         def
dZdefdZdef fdZded         fdZdefdZdedefdZ xZ S )NumpyVariablezz
    Wrapper around `numpy.*`. Currently, is able to trace a small subset of numpy functions as well as numpy dtypes.
    ra   r;   r<   Nc                 H     t                      j        di | || _        d S r>   r9  r*  s      rD   rA   zNumpyVariable.__init__  r:  rE   c                     | j         S rG   r.  rI   s    rD   r  z*NumpyVariable.get_real_python_backed_value  r/  rE   rQ  c                     |j                             d          }t          |          dk    r|d d         ddgk    sJ || j        v S Nr@  r   r   _numpy)r   splitrg   constant_fold_functionsr   rQ  mods      rD   can_constant_fold_throughz'NumpyVariable.can_constant_fold_through  sQ    m!!#&&3xx1}}RaRWh,?!?!?!??S000rE   c                     |j                             d          }t          |          dk    r|d d         ddgk    sJ t                              |          S rd  )r   rf  rg   np_constant_collections_mapri   rh  s      rD    get_constant_collection_for_funcz.NumpyVariable.get_constant_collection_for_func  sW    m!!#&&3xx1}}RaRWh,?!?!?!??*..r222rE   rT   r6   r{   c                    t           j        s+t          dd| j         d| d| d| j         ddg           d	d
lm} ddlm} t                      	                    | j                  }|:t          dd| j         d| d| d| d	d| j         dg t          j                   |J |                     |          x}	 	  | | j        d |D             i d |                                D                       S # t          $ r> t          dd| j         d| d| d| d	d| j         dg t          j                   Y d S w xY w|j        dk    rEt           j        r9t          dd| j         d| d| d| d	d| j         ddd| j         dg           |                    |j        ||          \  }}|                     |          rct+          ||          rSt-          j        | |                                 d |D             i d |                                D                       S  |j        j        d  ||          gt7          ||          R  }|                    ||          S )!Nz?attempted to trace numpy function with config.trace_numpy=Falsenumpy function: r  r  z"Attempted to trace numpy function z; while `torch._dynamo.config.trace_numpy` was set to False.zHSet `torch._dynamo.config.trace_numpy` to True to trace numpy functions.rX   r   )numpy_to_tensor_wrapperr+   )NumpyNdarrayVariablez8attempted to trace numpy function unsupported by PyTorchz  (corresponding torch function: r#  z Can't find numpy numpy function z in torch._numpy.c                 6    g | ]}|                                 S r?   r]   r  s     rD   
<listcomp>z/NumpyVariable.call_function.<locals>.<listcomp>  $    ???Q!..00???rE   c                 >    i | ]\  }}||                                 S r?   rs  rY  s      rD   r\  z/NumpyVariable.call_function.<locals>.<dictcomp>  *    PPPA1a2244PPPrE   zTnumpy function that produces a const collection type encountered non-const argumentsznumpy function zs that produces a const collection type (e.g. np.dtype, np.iinfo/np.finfo) received arguments that are not constant.ztorch._numpy.randomzXattempted to trace torch._numpy.random function with config.use_numpy_random_stream=TruezAttempted to trace zD when `torch._dynamo.config.use_numpy_random_stream` is set to True.z<Set `torch._dynamo.config.use_numpy_random_stream` to False.zAvoid calling r@  c                 6    g | ]}|                                 S r?   rs  r  s     rD   rt  z/NumpyVariable.call_function.<locals>.<listcomp>  ru  rE   c                 >    i | ]\  }}||                                 S r?   rs  rY  s      rD   r\  z/NumpyVariable.call_function.<locals>.<dictcomp>  rw  rE   r   )r   trace_numpyr   ra   rO  rp  tensorrq  rV  ri   r   r   rm  r   r,   
USER_ERRORr   use_numpy_random_stream
patch_argsr   rj  r$   r.   r   r]   r   r  r)   ry   )	rB   rT   r{   r;   rp  rq  r   collection_variable_typr  s	            rD   r   zNumpyVariable.call_function  s    ! 		YW4:WWtWWvWWM M M M _    	433333000000 ""&&tz22<R~4:~~t~~v~~w{~~~\tz\\\&2	    '+'L'LT'R'RR#..DJ??$???PPPPP    7 
 
 
r Ctz  C  C4  C  CSY  C  C{  C  C  C!@$* !@ !@ !@*5	 	 	 	 	 	 	
 #8882 9 v Ctz  C  C4  C  CSY  C  C{  C  C  C!&dj !& !& !& W6666	 	 	 	 0::4=$PVWWLD&--d33 
-dF;;
 ',-D++--??$???PPPPP    +BI*''-- #400  E
 (..r5999s   ;C? ?AEErU   c           	      `    t          dd| j         d| d| dg t          j                   d S )Nz/attempted to trace numpy.* function as a methodro  r  r  z6Tracing numpy.* functions as methods is not supported.rX   )r   ra   r   r   )rB   rT   rU   r{   r;   s        rD   r   zNumpyVariable.call_method  sX     	EStzSS4SS6SSP",		
 	
 	
 	
 	
 	
rE   c                     | j         S rG   r.  rI   s    rD   r]   z NumpyVariable.as_python_constant.  r/  rE   c                     t           j        rLt          | j        t          j                  r| j        S t          | j        t                    r| j        j        S t                      	                                S rG   )
r   rz  r   ra   enumEnumMetar`   r   r@   r  r  s    rD   r  zNumpyVariable.as_proxy1  sc     		+$*dm44 "z!$*d++ + z**ww!!!rE   Tc                     dS r  r?   rI   s    rD   r  z NumpyVariable.is_python_hashable?  r  rE   c                 D    t          |                                           S rG   r  rI   s    rD   r  zNumpyVariable.get_python_hashB  r  rE   r  c                     t          |t                    o)|                                 |                                k    S rG   r  r  s     rD   r  zNumpyVariable.is_python_equalE  r  rE   )!r   r   r   r:  tnp
issubdtyperg  r   rA   r  r   r   r   r  rj  rm  r   r.   r   r   r   r   r   r   r]   r  r   r  r  r  r   r  r   r   s   @rD   r`  r`    s(          #~/c S T      c     15+= 1$ 1 1 1 [1
 3%2D 3 3 3 3 [3
Z:#Z: 'Z: S/)*	Z:
 
Z: Z: Z: Z:x
#
 
 ?#	

 S/)*
 

 
 
 
 $7    "# " " " " " "GDM    / / / / /
V 
 
 
 
 
 
 
 
 
rE   r`  c                   <     e Zd Zdeddf fdZdefdZd	dZ xZS )
NullVariabler;   r<   Nc                 :     t                      j        di | d S r>   r^  r_  s     rD   rA   zNullVariable.__init__N  r`  rE   c                     dS )Nr  r?   rI   s    rD   r6  zNullVariable.__repr__Q  s    ~rE   rK   r5   c                     t           j        dk     rt          dddg t          j                   |                    t          d                     d S )N)r      z0cannot reconstruct NullVariable in Python < 3.11rW   zdAttempted to generate PUSH_NULL instruction in Python < 3.11; where this instruction does not exist.rX   	PUSH_NULL)r  r  r   r   r   rG  r   rC  s     rD   rS   zNullVariable.reconstructT  sg    g%%J9&1    	0==>>>>>rE   r   )	r   r   r   r   rA   r   r6  rS   r   r   s   @rD   r  r  M  sz        # # # # # # # ##    ? ? ? ? ? ? ? ?rE   r  c                       e Zd ZdZdS )r   z"Marker used to implement delattr()Nr9  r?   rE   rD   r   r   b  s        ,,,,rE   r   c            
            e Zd ZdZdhej        ZdefdZede	de
e         dee	ef         defd            Zde	de
e         dee	ef         dedd	f
 fd
Zde	fdZededefd            Zde	fdZddZ xZS )StringFormatVariablez[
    Represents a call to str.format(), we delay calling format until after the graph.
    format_stringr<   c                     t           S rG   )r   rI   s    rD   rJ   z StringFormatVariable.python_typem  s    
rE   sym_args
sym_kwargsc           	      l   t          d t          j        ||                                          D                       rPt          j                             |j        d |D             i d |                                D                       S  | |t          |          t          |                    S )Nc              3   >   K   | ]}|                                 V  d S rG   )r  r  s     rD   r2  z.StringFormatVariable.create.<locals>.<genexpr>w  sB       
 
   ""
 
 
 
 
 
rE   c                 6    g | ]}|                                 S r?   rs  )r1  r[  s     rD   rt  z/StringFormatVariable.create.<locals>.<listcomp>}  s$    ???a**,,???rE   c                 >    i | ]\  }}||                                 S r?   rs  rY  s      rD   r\  z/StringFormatVariable.create.<locals>.<dictcomp>~  s*    PPPTQq!..00PPPrE   )all	itertoolschainvaluesr   r/   ry   formatr   r   r   )r   r  r  r  s       rD   ry   zStringFormatVariable.createp  s      
 
_Xz/@/@/B/BCC
 
 
 
 
 		 -44$$??h???PPZ=M=M=O=OPPP    s=$x..$z2B2BCCCrE   r;   Nc                      t                      j        di | t          |t                    sJ || _        || _        || _        d S r>   )r@   rA   r   r   r  r  r  )rB   r  r  r  r;   rC   s        rD   rA   zStringFormatVariable.__init__  sR     	""6"""------* $rE   c                 P    | j         j         d| j        d| j        d| j        dS r  )rC   r   r  r  r  rI   s    rD   r6  zStringFormatVariable.__repr__  s7    .)kkD,>kkDMkkW[WfkkkkrE   r%  c                 t    	 |                                  S # t          $ r |                                 cY S w xY wrG   )r]   r(  r)  r*  s    rD   r+  z&StringFormatVariable._debug_format_arg  sJ    	$))+++ 	$ 	$ 	$>>#####	$s    77c                      	   j         j         fd j        D             i  fd j                                        D             }n# t
          $ r t                     cY S w xY wt          |          S )Nc                 :    g | ]}                     |          S r?   r/  r0  s     rD   rt  z3StringFormatVariable.debug_repr.<locals>.<listcomp>  s'    GGG#$((--GGGrE   c                 B    i | ]\  }}|                     |          S r?   r/  )r1  r;  ra   rB   s      rD   r\  z3StringFormatVariable.debug_repr.<locals>.<dictcomp>  s=       "U //66  rE   )r  r  r  r  r   r(  r'  )rB   rendereds   ` rD   r)  zStringFormatVariable.debug_repr  s    		0t)0GGGGGGG   &*o&;&;&=&=   HH  	 	 	::	H~~s   AA
 
A&%A&rK   r5   c                 d                          fdd            t          j        t           j                                       d  j                                        D             } t          j        |                                         t          dd                     d S )Nc                                                                 j                                       d          g          S )Nr  )rQ   rH  r  create_load_attrr  s   rD   rO   z2StringFormatVariable.reconstruct.<locals>.<lambda>  sA    G))--d.@AA,,X66  rE   T)call_function_exc                 T    i | ]%\  }}t           j                            |          |&S r?   r   r/   ry   rY  s      rD   r\  z4StringFormatVariable.reconstruct.<locals>.<dictcomp>  s<     
 
 
8<1I&--a00!
 
 
rE   F)
rP   r   r   r   r  r  r   ConstDictVariablerQ   r   )rB   rK   r;   s   `` rD   rS   z StringFormatVariable.reconstruct  s         " 	 	
 	
 	
 		'T](;(;<<===
 
@D@U@U@W@W
 
 
 		+F334445dEBBCCCCCrE   r   )r   r   r   r:  r.   r   r`   rJ   r   r   r   r   ry   r   rA   r6  r   r   r+  r)  rS   r   r   s   @rD   r  r  f  s         &G(FGNT     DD ?+D o-.	D
 
D D D [D$%% ?+% o-.	%
 % 
% % % % % %l# l l l l $ $6 $ $ $ \$C    D D D D D D D DrE   r  c                   T     e Zd ZeZdededdf fdZdefdZdee         fdZ	 xZ
S )ObjectVariablera   r;   r<   Nc                 H     t                      j        di | || _        d S r>   r9  r*  s      rD   rA   zObjectVariable.__init__  r:  rE   c                     | j         S rG   r.  rI   s    rD   r  z+ObjectVariable.get_real_python_backed_value  r/  rE   c                     t           S rG   )r   rI   s    rD   rJ   zObjectVariable.python_type  s    rE   )r   r   r   r   r   r   rA   r  r`   rJ   r   r   s   @rD   r  r    s        Mf        f    T&\        rE   r  c            	           e Zd ZdZdededdf fdZdefdZedede	e
j        e
j        z           fd	            Zd
ddee         deeef         defdZddZededee         deeef         defd            Z xZS )DebuggingVariablez
    Represents a call to a debugging function like print(), or something
    registered to config.reorderable_logging_functions.
    ra   r;   r<   Nc                 H     t                      j        di | || _        d S r>   r9  r*  s      rD   rA   zDebuggingVariable.__init__  r:  rE   c                 *    t          | j                  S rG   r`   ra   rI   s    rD   rJ   zDebuggingVariable.python_type      DJrE   r   c                     t          |           o=t          | t          j        t          j        f          o| t
          j        j        j        v S rG   )	rP  r   r   r   r   r   r  r   reorderable_logging_functionsr   s    rD   is_reorderable_logging_functionz1DebuggingVariable.is_reorderable_logging_function  sF    
 SMM J3!3U5N OPPJu}+II	
rE   rT   r6   r{   c           	      @   |j         rt          j        d           S |                     | j        ||          s+t          dd| j         d| d| dd| j         dg           |j                            | t          |          f           t          j        d           S )	NzJattempted to reorder a debugging function that can't actually be reorderedzfn: r  r  zl`torch.compile` can only reorder functions where the arguments are Tensors, constants, or string formatters.z#Avoid calling the logging function z" with args that are not supported.rX   )	exportr/   ry   can_reorder_logsra   r   debug_localsr  r   r  s       rD   r   zDebuggingVariable.call_function  s     9 	1#*4000$$TZv>> 		dKtzKK4KK6KK@ i$*hhh    	d4jj1222&t,,,rE   rK   r5   c                 H    | j         J | j                             |          S rG   )r^   rS   rC  s     rD   rS   zDebuggingVariable.reconstruct  s&    {&&&{&&w///rE   rQ  c                     t           j        t           j        t          f}t	          j        ||g          }|D ]}t          ||          s dS dS )zm
        Run some additional checks for what sort of function calls can we
        actually reorder.
        FT)r   r  r/   r  pytreetree_leavesr   )rQ  r{   r;   allowed_input_types	flat_argsr%  s         rD   r  z"DebuggingVariable.can_reorder_logs  sf     $& 
 &f~66	 	 	Cc#677 uu trE   r   )r   r   r   r:  r   rA   r`   rJ   r   r
   r   r   r   r  r   r.   r   r   r   rS   r  r  r   r   s   @rD   r  r    s]        
c S T       T         

	5%(AA	B
 
 
 \
-#- '- S/)*	-
 
- - - -00 0 0 0 S  tCH~ RV    \    rE   r  c                   |     e Zd ZdZdededdf fdZdefdZdefdZd	d
de	e
         deee
f         de
fdZ xZS )IgnoredFunctionVariablezL
    Represents a call to an arbitrary function that should be ignored.
    ra   r;   r<   Nc                 H     t                      j        di | || _        d S r>   r9  r*  s      rD   rA   z IgnoredFunctionVariable.__init__  r:  rE   c                 *    t          | j                  S rG   r  rI   s    rD   rJ   z#IgnoredFunctionVariable.python_type  r  rE   c                     | j         S rG   r.  rI   s    rD   r  z4IgnoredFunctionVariable.get_real_python_backed_value  r/  rE   rT   r6   r{   c                 @    t           j                            d           S rG   r  r  s       rD   r   z%IgnoredFunctionVariable.call_function   s     )00666rE   )r   r   r   r:  r   rA   r`   rJ   r  r   r.   r   r   r   r   r   s   @rD   r  r    s         c S T       T        c    7#7 '7 S/)*	7
 
7 7 7 7 7 7 7 7rE   r  c            
            e Zd ZdZdej        deddf fdZdefdZ	dej        fdZ
d	d
dedee         deeef         def
dZ xZS )LoggingLoggerVariablez:
    Represents a call to any logging.Logger methods.
    ra   r;   r<   Nc                 H     t                      j        di | || _        d S r>   r9  r*  s      rD   rA   zLoggingLoggerVariable.__init__.  r:  rE   c                 *    t          | j                  S rG   r  rI   s    rD   rJ   z!LoggingLoggerVariable.python_type2  r  rE   c                     | j         S rG   r.  rI   s    rD   r  z2LoggingLoggerVariable.get_real_python_backed_value5  r/  rE   rT   r6   rU   r{   c                 d   |j         rt          j                            d           S t	          | j        |d           }t	          |dd           }t          j        j        j	        }||v s||v rt          j                            d           S t          dd| j         d| d| d| ddg	           d S )
Nr   z8logging.Logger method not supported for non-export caseszmethod: r@  r  r  z>logging.Logger methods are not supported for non-export cases.zJAdd the logging method to `torch._dynamo.config.ignore_logging_functions`.rX   )r  r   r/   ry   re   ra   r   r  r   ignore_logging_functionsr   )rB   rT   rU   r{   r;   methodr  
ignore_sets           rD   r   z!LoggingLoggerVariable.call_method8  s     9 	;-44T:::T4006:t44 ])B
Z8z#9#9-44T:::NRtzRRDRR$RR&RRX\		
 	
 	
 	
 	
 	
rE   )r   r   r   r:  loggingLoggerr   rA   r`   rJ   r  r   r   r.   r   r   r   r   s   @rD   r  r  )  s         gn         T        gn    
#
 
 ?#	

 S/)*
 

 
 
 
 
 
 
 
rE   r  c            
           e Zd ZdZ	 ddlmZmZm	Z
 n%# e$ r  eddi           Z eddi           ZY nw xY wdededd	f fd
Zedefd            ZdefdZdddedee         deeef         def
dZdddedefdZ xZS )ConstantLikeVariablez8self.value is a compile-time constant, but not a literalr   )dtypefloatinggenericinvalid_typer?   ra   r;   r<   Nc                 H     t                      j        di | || _        d S r>   r9  r*  s      rD   rA   zConstantLikeVariable.__init__e  r:  rE   c                 p    t          | j                  }|j        dk    r|j        S |j         d|j         S )z4Dynamically compute the prefix from the value's typerH   r@  )r`   ra   r   r   )rB   ts     rD   _error_prefixz"ConstantLikeVariable._error_prefixi  sB      <:%%>!,11111rE   c                     | j         S rG   r.  rI   s    rD   r]   z'ConstantLikeVariable.as_python_constantt  r/  rE   rT   r6   rU   r{   c                 Z   g i }}	 d |D             }d |                                 D             }nA# t          $ r4 t          d| j         d| d| d| dd| j         d| d	d
g           Y nw xY w t	          | j        |          |i |}t          j                            |          rt          j
        ||          S t          |t          j                  rt          |          S t          d| j         d| d| d| d| 	d| j         d| d| dg t          j                   d S )Nc                 6    g | ]}|                                 S r?   rs  r  s     rD   rt  z4ConstantLikeVariable.call_method.<locals>.<listcomp>  s$    :::Q))++:::rE   c                 >    i | ]\  }}||                                 S r?   rs  rY  s      rD   r\  z4ConstantLikeVariable.call_method.<locals>.<dictcomp>  s*    LLLTQq!..00LLLrE   z0constant-like method call with non-constant argsr@  z(*z, **r#  zAttempted to call z with non-constant args.zFEnsure that the args to the method call are constant (int, str, etc.).rX   z6constant-like method call with unsupported return typez) returned z, got unsupported return value )r   r   r   r  re   ra   r   r/   ru   r.   r   r   reMatchr  r   r   )rB   rT   rU   r{   r;   cargsckwargsr   s           rD   r   z ConstantLikeVariable.call_methodw  s    Rw	::T:::ELLV\\^^LLLGG" 	 	 	J-KKKKKK&KKKd1Cdddddd\	     	 +T**E=W==%0088 	5"(V444fbh'' 	0'///L)YYDYYDYYfYYQWYYpT-?pp$ppgmppp".		
 	
 	
 	
 	
 	
s   *1 ;A/.A/c                    t          | j        |          }t          || j                  rt	          |          }t          || j                  rt          |          S t          |t                    r$t          || j	                  rt          |          S t          j                            |          rt          j        ||          S t!          | |t          |                    S )Nrr   )re   ra   r   np_floatingfloatnp_dtypeNumpyDTypeVariabler`   r_   
np_genericr`  r   r/   ru   r.   r   rv   r3  s       rD   rz   z ConstantLikeVariable.var_getattr  s    T**fd.// 	#6]]Ffdm,, 	.%f---fd## 	)
64?(K(K 	) (((%0088 	5"(V444tT4<<@@@@rE   )r   r   r   r:  numpyr  r  r  r  r  r  ImportErrorr`   r   rA   propertyr   r  r]   r   r.   r   r   rz   r   r   s   @rD   r  r  V  s       BB
0	
 	
 	
 	
 	
 	
 	
 	
 	
 	
 	

  0 0 0d>2r224B//	0c S T       2s 2 2 2 X2C    %
#%
 %
 ?#	%

 S/)*%
 
%
 %
 %
 %
NA5 AS A_ A A A A A A A As   
 66r  c                   ,     e Zd ZdZdeddf fdZ xZS )TorchVersionVariableztorch.__version__r;   r<   Nc                     |                     dt          j                   |d         t          j        u sJ  t                      j        di | d S )Nra   r?   )
setdefaultr   __version__r@   rA   r_  s     rD   rA   zTorchVersionVariable.__init__  sW    '5#4555g%"33333""6"""""rE   )r   r   r   r  r   rA   r   r   s   @rD   r  r    sP        'M# # # # # # # # # # # #rE   r  c                       e Zd ZdefdZdS )r  r<   c                 $    | j         j        j        S )a0  Similar to how numpy dtype descriptors (e.g. np.float32 ) are handled by NumpyVariable:

        np.dtype() objects are serialized as strings, torch._numpy wrappers will normalize to the torch dtype.
        This also handles unsupported things nicely (i.e. structured arrays and object arrays).
        )ra   r`   r   rI   s    rD   r  zNumpyDTypeVariable.as_proxy  s     z''rE   N)r   r   r   r   r  r?   rE   rD   r  r    s/        (# ( ( ( ( ( (rE   r  c                   l     e Zd ZdZdeddf fdZdefdZddd	ee	         de
ee	f         dd
fdZ xZS )RandomClassVariablezrandom.Randomr;   r<   Nc                 :     t                      j        di | d S r>   r^  r_  s     rD   rA   zRandomClassVariable.__init__  r`  rE   c                     t           S rG   rl  rI   s    rD   rJ   zRandomClassVariable.python_type  rm  rE   rT   r6   r{   RandomVariablec                 D   t          |          dk    s|r%t          dd| d| dg t          j                   t          |          dk    rt          j                            d           n|d         }t          |t          j        	                                          S )	Nr+   z'random.Random() with improper argumentszargs: r  z=random.Random() with > 1 arg or with kwargs is not supported.rX   r   )seedr   )
rg   r   r   r|  r   r/   ry   r  baseValueMutationNew)rB   rT   r{   r;   r   s        rD   r   z!RandomClassVariable.call_function  s     t99q==F=A99999[&1	    ;>d))q..y)00666dSTgY^%D%D%F%F
 
 
 	
rE   )r   r   r   r:  r   rA   r`   rJ   r   r.   r   r   r   r   r   s   @rD   r  r    s        # # # # # # # #T    
#
 '
 S/)*	

 

 
 
 
 
 
 
 
rE   r  c            
           e Zd ZdZej        Zdhej        Zh dZ		 	 ddej        dz  dedz  de
ddf fd	Zdeej                 fd
Zdej        fdZededefd            Zedeeeedf         edz  f         ddfd            Zedeeeedf         edz  f         defd            Zededeeeedf         edz  f         fd            Zdddedee         deeef         def
 fdZddZ xZS )r  zrandom.Random()

    Implemented by wrapping a VariableTracker around a random.Random object.
    The supported methods for the random.Random object cannot be overridden.
    Assumes that random objects behave the same given a set seed or state.
    random>   r  randintuniform	randrangeNrandr   r;   r<   c                 Z    t                      j        di | |]|                     |          sJ t          j                    | _        | j                            |                                           d S ||                                nd }t          j        |          | _        d S r>   )r@   rA   is_supported_random_objr  r   setstategetstater]   )rB   r  r   r;   rC   s       rD   rA   zRandomVariable.__init__  s     	""6"""//55555 -//DKK  11111040@4**,,,dD ---DKKKrE   c                     t           j        S rG   )r  r   rI   s    rD   rJ   zRandomVariable.python_type		  s
    }rE   c                     | j         S rG   )r  rI   s    rD   r]   z!RandomVariable.as_python_constant	  s
    {rE   r   c                    t          |           t          j        urdS t          j        t
          j        d          D ]}t          | |          s dS t          | |          }t          j
        |          r5|t          t          j        |                              |           k    r dS nt          |dd           t          t          j        |          ur dS dS )NF)r   r  r  r   T)r`   r  r   r  r  r  _supported_fn_namesr   re   r  	isbuiltinr  )r   rU   meths      rD   r
  z&RandomVariable.is_supported_random_obj	  s    99FM))5O.0P
 
 	! 	!D 3%% uu3%%D && !76=$77??DDDD 55 E 4T22'&-QU:V:VVV 55 WtrE   state.c                 @   t          |           t          u sJ t          | d                   t          u sJ t          | d                   t          u sJ t          d | d         D                       sJ | d         t          | d                   t          u sJ d S d S )Nr   r+   c              3   B   K   | ]}t          |          t          u V  d S rG   )r`   r  r  s     rD   r2  z-RandomVariable.check_state.<locals>.<genexpr>'	  s,      44a477c>444444rE   r   )r`   r   r  r  r  r  s    rD   check_statezRandomVariable.check_state"	  s    E{{e####E!H~~$$$$E!H~~&&&&445844444444Qx4a>>U#:#:#::#:#:rE   c                 4   t                               |            t          j        t          j                            | d                   t          j        d | d         D                       t          j                            | d                   g          S )Nr   c                 L    g | ]!}t           j                            |          "S r?   r  r  s     rD   rt  z-RandomVariable.wrap_state.<locals>.<listcomp>1	  s)    LLLaY/66q99LLLrE   r+   r   )r  r  r   r   r/   ry   r  s    rD   
wrap_statezRandomVariable.wrap_state*	  s    ""5)))&*11%(;;'LL58LLL  *11%(;;
 
 	
rE   c                 b    |                                  }t                              |           |S rG   )r]   r  r  )r  	state_objs     rD   unwrap_statezRandomVariable.unwrap_state7	  s/     ,,..	""9---rE   rT   r6   rU   r{   c                    dk    rv|j         j                            |             | j        j        d |D             i d |                                D              t          j                            d           S dk    r,| 	                    | j        
                                          S dk    rq|j         j                            |            | j                            |                     |d                              t          j                            d           S | j        v r|j         j                            |            | j        
                                dt          dt          d	t          ffd
} t          | j                  d |D             i d |                                D              t!          ||||          S t#                                          |||          S )Nr   c                 6    g | ]}|                                 S r?   rs  r  s     rD   rt  z.RandomVariable.call_method.<locals>.<listcomp>I	  $    777Q!&&((777rE   c                 >    i | ]\  }}||                                 S r?   rs  )r1  r;  r   s      rD   r\  z.RandomVariable.call_method.<locals>.<dictcomp>J	  s*    PPPXS#3..00PPPrE   r  r  r   r{   r;   r<   c                      t          j                    }|                                t          |          | i |S rG   )r  r   r  re   )r{   r;   rrU   r  s      rD   call_random_methz4RandomVariable.call_method.<locals>.call_random_methW	  s?    MOO

5!!!'wq$''8888rE   c                 6    g | ]}|                                 S r?   rs  r  s     rD   rt  z.RandomVariable.call_method.<locals>.<listcomp>_	  r   rE   c                 >    i | ]\  }}||                                 S r?   rs  rY  s      rD   r\  z.RandomVariable.call_method.<locals>.<dictcomp>`	  s*    HHHA1a**,,HHHrE   )r   r   mutationr  r   r   r   r/   ry   r  r  r  r  r  r   re   r2   r@   r   )rB   rT   rU   r{   r;   r$  r  rC   s     `   @rD   r   zRandomVariable.call_method?	  s    6>>I"++D111DK77$777PPPPP   -44T:::Z??4;#7#7#9#9:::ZI"++D111K  !2!247!;!;<<<-44T:::T---I"++D111K((**E9 9s 9s 9 9 9 9 9 9 9 'GDK&&77$777HHHHH  
 ""&6fEEEww""2tT6:::rE   rK   r5   c                 t                        fd                               dd                                                                d            |                     | j                                                                                 dd                                            d S )Nc                                                                 t                                         d          g          S )Nr   )rQ   create_load_python_moduler  r  rK   s   rD   rO   z,RandomVariable.reconstruct.<locals>.<lambda>h	  s>    G))55f==,,X66  rE   r   Fr  r+   T)rP   r   r  r   r  r  r  pop_toprC  s    `rD   rS   zRandomVariable.reconstructf	  s       	
 	
 	
 	a''' 	*%%% 4 4 6 677888a&&&rE   r  r   )r   r   r   r:  r  r   r   r.   r   r  r   rA   r`   rJ   r]   r   r  r
  r   r  r  r  r   r  r  r   r   r   r   rS   r   r   s   @rD   r  r    sk         MM 			'N
   &*'+. .md". $. 	.
 
. . . . . .T&-0    FM     V     \$ ;5eCHout|!CD ; ; ; ; \; 

%U38_edl BC 

 

 

 

 \

 	sE#s(OUT\1	2   \%;#%; %; ?#	%;
 S/)*%; 
%; %; %; %; %; %;N       rE   r  c                        e Zd ZdefdZedddej        e         de	dz  dedd f
d	            Z
d
edededdf fdZdddee         deeef         defdZddZdefdZdefdZdedefdZ xZS )WeakRefVariabler<   c                     t           j        S rG   )weakrefrefrI   s    rD   rJ   zWeakRefVariable.python_typez	  s
    {rE   rT   r6   weakref_valuer^   Noptionsc                     |J |j         }|ot          |d          }t          j        | ||          } |            }|ot	          |          }t          j        | ||          }||d<   t          ||fi |S )N__callback__r^   )r5  r   r.   r   r#   r.  )	rT   r2  r^   r3  callbackcallback_sourcecallback_vtreferentreferent_vts	            rD   r   zWeakRefVariable.build}	  s     !!! - GZ%G%G%+B/JJ =??5-f55%+B&AA"{KCC7CCCrE   r:  r8  c                 V     t                      j        di | || _        || _        d S r>   )r@   rA   r:  r8  )rB   r:  r8  r3  rC   s       rD   rA   zWeakRefVariable.__init__	  s8     	##7###&&rE   r{   r;   c                     | j         S rG   )r:  r  s       rD   r   zWeakRefVariable.call_function	  s     rE   rK   r5   c                                          fd            | j                    | j                                       t	          dd                     d S )Nc                  0                          dd          S )Nr0  r1  )r  r+  s   rD   rO   z-WeakRefVariable.reconstruct.<locals>.<lambda>	  s    g&>&>y%&P&P rE   r   F)rP   r:  r8  rQ   r   rC  s    `rD   rS   zWeakRefVariable.reconstruct	  sl    PPPPQQQ !!! !!!21e<<=====rE   c                 4    | j                                         S rG   )r:  r  rI   s    rD   r  z"WeakRefVariable.is_python_hashable	  s    22444rE   c                 4    | j                                         S rG   )r:  r  rI   s    rD   r  zWeakRefVariable.get_python_hash	  s    //111rE   r  c                 n    t          |t                    sdS | j                            |j                  S )NF)r   r.  r:  r  r  s     rD   r  zWeakRefVariable.is_python_equal	  s3    %11 	5//0ABBBrE   r   )r   r   r   r`   rJ   r   r0  ReferenceTyper   r   r   r.   rA   r   r   r   r   rS   r  r  r  r  r   r  r   r   s   @rD   r.  r.  y	  s       T     D#D,S1D D 	D
 
D D D \D '*'9H'UX'	' ' ' ' ' ' #  '  S/)*	 
 
       > > > >5D 5 5 5 52 2 2 2 2CV C C C C C C C C CrE   r.  )r:  rH   r  r  r  r  r  r  r  r  r  r   r   r0  collections.abcr   r   r   r   typingr   r   r	   r
   r   torch._Cr   torch._numpyre  r  torch.utils._pytreerO  _pytreer  torch._dynamo.variables.baser   torch._dynamo.variables.listsr   torch._guardsr   rW   r   r   r   r   bytecode_transformationr   r   r   create_parameter_opr   r   r   r   r   guardsr   r   mutation_guardr   r^   r   r   r    r!   r"   r#   r$   r%   r&   r'   r(   r)   r*   r  r,   r-   r.   r  r/   	functionsr0   r1   user_definedr2   r3   r4   torch._dynamo.codegenr5   torch._dynamo.symbolic_convertr6   r8   r   r   r  r8  r<  rA  rT  r\  rf  ri  	dataclassr  r  r  r   rv   r  r  r&  r7  	lru_cacher   r   rV  r^  r`  r  r   r  r  r  r  r  r  r  r  finfoiinfor  rl  r  r  r.  r?   rE   rD   <module>rX     s_
   "                  				 



       . . . . . . . .       % % % % % % @ @ @ @ @ @ @ @ @ @ @ @ @ @        $ $ $ $ $ $ $ $ $ 5 5 5 5 5 5 7 7 7 7 7 7             @ @ @ @ @ @ @ @ @ @ @ @         
 G F F F F F K K K K K K K K K K 0 0 0 0 0 0 0 0 5 5 5 5 5 5                                 W V V V V V V V V V & & & & & & G G G G G G G G X X X X X X X X X X  E//////DDDDDDQ
 Q
 Q
 Q
 Q
O Q
 Q
 Q
h
- - - - -? - - -(^; ^; ^; ^; ^; ^; ^; ^;Bj4 j4 j4 j4 j4 j4 j4 j4Z    o   
 
 
 
 
o 
 
 
2?7 ?7 ?7 ?7 ?7 ?7 ?7 ?7D    ?   4# # # # # # # #
%c %hsCx6H % % % %c c c c c c c cL M M M M M M M MD- D- D- D- D-&? D- D- D-N2 2 2 2 26 2 2 2j( ( ( ( (_ ( ( ($AG AG AG AG AGo AG AG AGHI
 I
 I
 I
 I
O I
 I
 I
X       2/= /= /= /= /=? /= /= /=db
 b
 b
 b
 b
_ b
 b
 b
J Q4 95;M MN     ( Q(4 2E4M MN ( ( (  (a
 a
 a
 a
 a
O a
 a
 a
J? ? ? ? ?? ? ? ?*- - - - -o - - -PD PD PD PD PD? PD PD PDf    _    E E E E E E E EP7 7 7 7 7o 7 7 70*
 *
 *
 *
 *
O *
 *
 *
ZSA SA SA SA SA? SA SA SAl# # # # #/ # # #( ( ( ( (- ( ( ( I#I#I! 
 
 
 
 
/ 
 
 
<R R R R R_ R R Rj5C 5C 5C 5C 5Co 5C 5C 5C 5C 5CrE   