
    i%                    0   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
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mZmZmZ ddlmZmZ ddlmZ ddl m!Z!m"Z"m#Z# dd	l$m%Z%m&Z& dd
l'm(Z( ddl)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/ ddlm0Z0m1Z1m2Z2m3Z3m4Z4m5Z5m6Z6m7Z7 ddl8m9Z9 ddl:m;Z; ddl<m=Z=m>Z> ddl?m@Z@mAZAmBZB erddlCmDZD ddlEmFZF  G d d      ZG G d de9      ZH G d de9      ZI G d de9      ZJ G d deJ      ZK G d de9      ZL G d  d!e9      ZM G d" d#e9      ZNd$ ZO G d% d&e9      ZPej                   G d' d(             ZR G d) d*eB      ZS G d+ d,eB      ZT G d- d.e9      ZU G d/ d0e9      ZV G d1 d2e9      ZW G d3 d4e9      ZX G d5 d6e9      ZY G d7 d8e9      ZZ ej                  d9      d:        Z\ ej                  d9      d;        Z] G d< d=e9      Z^ G d> d?e9      Z_ G d@ dAe9      Z` G dB dCe9      Za G dD dEe9      Zb G dF dGe9      Zc G dH dIe9      Zd G dJ dKed      Ze G dL dMed      Zf G dN dOed      Zg G dP dQed      Zh G dR dSed      Ziej                  ehej                  ehej                  eiiZm G dT dUe9      Zn G dV dWe9      Zo G dX dYe9      Zpy)Za  
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)OptionalTYPE_CHECKING   )configgraph_break_hintstrace_rules	variables)create_call_functioncreate_instruction)$do_not_convert_to_tracable_parameter)raise_observed_exceptionunimplementedunimplemented_v2)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list_methodsproxy_args_kwargstuple_methods   )VariableTracker)ConstantVariable)NestedUserFunctionVariableUserFunctionVariable)call_random_fnis_standard_setattrUserDefinedObjectVariable)	PyCodegen)InstructionTranslatorc                       e Zd Zy)NO_SUCH_SUBOBJN)__name__
__module____qualname__     V/var/www/html/engine/venv/lib/python3.12/site-packages/torch/_dynamo/variables/misc.pyr,   r,   F   s    r1   r,   c                        e Zd Zh ej                  Zdd fdZd fdZd fdZdddeddfd	Z		 	 	 	 	 	 	 	 dd
Z
 xZS )SuperVariablereturnc                 @    t        |   di | || _        || _        y Nr0   )super__init__typevarobjvar)selfr:   r;   kwargs	__class__s       r2   r9   zSuperVariable.__init__O   s%    "6" 
 r1   c                     j                  fd        | j                         | j                  . | j                         j                  t	        dd             y j                  t	        dd             y )Nc                  B     t        j                  t                    S N)r	   BuiltinVariabler8   )r>   codegens   r2   <lambda>z+SuperVariable.reconstruct.<locals>.<lambda>\   s    gi.G.G.N&O r1   r   Fr!   )add_push_nullr:   r;   extend_outputr
   )r<   rC   r>   s    `r2   reconstructzSuperVariable.reconstruct[   s`    OP;;"DKK !!"6q%"@A!!"6q%"@Ar1   txr*   c                 J   | j                   sJ d       | j                  j                         }| j                   j                         }| j                   j                  rt        | j                   j                        nd }t        |t              r,| j                   j                  }| j                   j                  }d }|j                  }	 |j                  |      dz   }t        |t        |            D ]W  }	||	   j                   j#                  |t$              x}
s)|
t$        us2|rt'        t)        t+        |      |	      |      }|
|fc S  t-        ddd| d| j                   d| j                    dd	d
g       y # t        $ r t        t        ||      |      d fcY S w xY w)Nz1-arg super not implementedr!   zUnable to resolve super getattr z"Dynamo failed to trace attribute `z$` accessed via `super()` (for type `z` and object `z8`) because the resolved attribute type is not supported.z0Ensure the attribute exists in the parent class.z(Check the arguments passed to `super()`.gb_typecontextexplanationhints)r;   r:   as_python_constantpython_typesourcer   
issubclasstypevalue__mro__index
ValueErrorgetattrr8   rangelen__dict__getr,   r   r   r   r   )r<   rH   namesearch_typetype_to_usetype_to_use_sourcerR   
search_mrostart_indexrW   resolved_getattrr>   s              r2   _resolved_getattr_and_sourcez*SuperVariable._resolved_getattr_and_sourced   s   {{999{ll557 kk--/.2kk.@.@Jt{{))*d 	 k4(++++K!%!3!3 ((
	H$**;7!;K ;J8 
	4E#-e#4#=#=#A#A$#WWW#>9)!+)-8J*KUS " ,V33
	4 	5<TF C((,~^DKK= QDD C:
	
'  	H 5k:DA4GG	Hs   E> >!F"!F"r^   r"   c                    | j                  | |      \  }}t        j                  j                  |      st	        | |      S |r(t        |j                  t        j                               t        j                  j                  ||      S )NrR   )
re   r	   r#   
is_literalGetAttrVariabler   
make_guardr   CONSTANT_MATCHcreate)r<   rH   r^   rU   rR   s        r2   var_getattrzSuperVariable.var_getattr   ss     99$Ev))44U;"4..&++L,G,GHI))00v0FFr1   c                    | j                  | |      \  }}|t        j                  u rt        t              S |t
        j                  j                  j                  u r| j                  }ddl	m
} t        |t        j                        rht        |j                  |      rR|sP|sNt               5  t        j                   t"        |      j%                  || j                  g|z   |      cd d d        S t'        dd|  d| d| d| d| dg t(        j*                  	       n| j                  j,                  rt/        |d
      r|j0                  dk(  rt        j2                  j5                  |      rt|j6                  }	t/        |	d      r%|	j8                  dk(  rt        j:                  |	      }
n#|j<                  }t        j2                  |	|      }
|
j?                  |d||      S t        |t@              rVt        |jB                  tD        jF                        r2t        j                   |jB                  |      j%                  |||      S t        |tH              rt        |jB                  tD        jF                        rt        | j                  t        j2                        r| j                  }nbd }| j                  j,                  rtK        | j                  j,                        }tM        jN                  || j                  jP                  |      }t        j                   |jB                  tS        |d            j%                  ||g||      S t        |tD        jF                        r6t        j                   ||      j%                  || j                  g|z   |      S t        |tD        jT                        r=t        jV                  |jB                  | j                  |      j%                  |||      S tY        |      r:t        | j                  t              r  | j                  jZ                  |g|i |S |t        j\                  u r|d   }	 |j_                         }|jb                  j                  je                  | j                        s4t'        dd|  d| d| j                   ddgt(        jf                  	       |jb                  j                  ji                  | j                  |t        jj                                t        jl                  d       S t        | j                  t        jn                        r@|| j                  jp                  v r(| j                  jr                  j?                  ||||      S t        | j                  t        jt                        r@|| j                  jv                  v r(| j                  jx                  j?                  ||||      S t        | j                  t        jz                        r0|t|        v r(| j                  j~                  j?                  ||||      S t        | j                  t        j                        r0|t        v r(| j                  j                  j?                  ||||      S |t        j                  u r|d   j                  }|jb                  j                  j                  | j                  |      r^|jb                  j                  j                  | j                  |d      }t        |t        jj                        rt        t        |       |S 	 t        j                  | j                  j                  |      }d }| j                  j,                   t        | j                  j,                  |      }tM        jN                  ||      S |t
        j                  j                  u r|d   }i }|d   j                  }|d   j                  j                         D ]$  \  }}|j                  j_                         }|||<   & |j                  j                  }d|j                  _O        	 |j%                  |||      ||j                  _O        S t        |tD        j                        rN|t        j                         v r8tM        jN                  |||      }|j%                  || j                  g|z   |      S t'        dd|  d| d| d| dd g	       y # 1 sw Y   'xY w# t`        $ r"}t'        dd|  d| ddg|       Y d }~d }~ww xY w# t        $ r t        t        |       Y w xY w# ||j                  _O        w xY w)!Nr   )AttributeMutationNewrg   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.rK   r-   __new__r.   builtins__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.)rL   rM   rN   rO   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.T)
deleted_ok   FzFAttempted 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.)Sre   objectr9   LambdaVariabler   torchnnModuler;   side_effectsro   
isinstancer	   r(   mutation_typer   r%   r   call_functionr   r   	DIFFICULTrR   hasattrr-   UserDefinedClassVariableis_supported_new_method__self__r.   rB   membercall_methodstaticmethodrt   typesFunctionTypeclassmethodr   r"   build
value_typer   
MethodTypeUserMethodVariabler'   method_setattr_standard__delattr__rP   NotImplementedErroroutputis_attribute_mutation
DYNAMO_BUG
store_attrDeletedVariabler#   UserDefinedDictVariable_dict_methods_dict_vtUserDefinedSetVariable_set_methods_set_vtUserDefinedTupleVariabler    	_tuple_vtUserDefinedListVariabler   _list_vt__getattribute__rU   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)r<   rH   r^   argsr=   inner_fnrR   r;   ro   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                               r2   r   zSuperVariable.call_method   s;     <<T4H& v&!(++111[[F; 69#F#FGv335IJ9; F$990#mB(<fEF F
 !A*4&$qaxH!  "1!1 9-778 KK*-!!Y.22JJ8T((Hx.83F3F*3T'77A"(--'@@_ **2y$GG,/Ju115
 11!!&mBf-. +.:u114
 $++y'I'IJ  ${{ "
;;%%!+DKK,>,>!?J.44..
  11!!*VZ*HmB 5 5v>? %"4"4511mB 4f=> %"2"23//!!4;;vmBf-. !*zKK20
 74;;66rKDKFKK+++7D..0 99))??L d*4&$8!#{{m ,(!(
 W*55 II""--T9#<#<#> --d33t{{I$E$EFDKK555;;''33BdFKKt{{I$D$DEDKK444;;&&222tT6JJt{{I$F$FGM);;((44RtVLLt{{I$E$EFL(;;''33BdFKK000 QIyy%%BBY //99KKt :  fi&?&?@,^R@=#44T[[5F5F	R
 K{{!!-/0B0BIN"((ZEE???
 7DI1gmmG)-a)<)<)> .%X%..;;=%-	'". 55UUFOTB,,L))"gyA  00P x!;!;<K99;; %**2x@F''T[[MD,@&II'"4&$0v24& 9,, `
	
eF FV ' 
 U*4&$8!L ` !	 	
D " =(<=4  00PsB   -:dd *d= e d	d:d55d:=eee0rA   r5   NrC   r)   rH   r*   rH   r*   r   list[VariableTracker]r=   dict[str, VariableTracker]r5   r"   )r-   r.   r/   r"   _nonvar_fieldsr9   rG   re   strrm   r   __classcell__r>   s   @r2   r4   r4   J   sy    		'	'N
B5
nG5 GS GEV G W
#W
 &	W

 -W
 
W
r1   r4   c                   h     e Zd Zd fdZddZddZd Zdddedefd	Z fd
Z	 fdZ
d ZeZ xZS )ExceptionVariablec                     t        |   di | || _        || _        t	        d       | _        t	        d       | _        t	        d      | _        t	        d       | _        y )NFr0   )	r8   r9   exc_typer   r#   __context__	__cause____suppress_context____traceback__)r<   r   r   r=   r>   s       r2   r9   zExceptionVariable.__init__  sZ    "6" 	 ,D1 *$/$4U$;! .d3r1   c                     || _         y rA   )r   )r<   rM   s     r2   set_contextzExceptionVariable.set_context  s
    "r1   c                      j                   fd       j                   j                         j                  t	         j                        d       dt
        dd f fd} |d        |d        |d       y )	Nc                  P     j                  dj                  j                        S )Nrs   )load_import_fromr   r-   rC   r<   s   r2   rD   z/ExceptionVariable.reconstruct.<locals>.<lambda>  s    G,,Z9O9OP r1   Fr^   r5   c                     t        |       }t        |t              r|j                  dv sJ |       y j	                           |       j                  j                  d             j                  |        y )N)TFNr   )rY   r   r#   rU   dup_toprF   rot_nr   )r^   r   rC   r<   s     r2   codegen_attrz3ExceptionVariable.reconstruct.<locals>.codegen_attr  sh    4&Dd,-zz%88>$>8!%%gmmA&67""4(r1   r   r   r   )rE   foreachr   r   r[   r   )r<   rC   r   s   `` r2   rG   zExceptionVariable.reconstruct  sl    P	
 			"c$))ne4	)s 	)t 	) 	]#[!+,r1   c                     | j                   S rA   )r   r<   s    r2   rQ   zExceptionVariable.python_type  s    }}r1   rH   r*   name_varvalc                    fd}|j                         }|dk(  r| j                  |       n1|dk(  rt        |t              r|j                  Ht        |t
        j                  t
        j                  t
        j                  t
        j                  f      r"|| _
        t        j                  d      | _        n |d       n|dk(  r/t        |t              r|j                  dv r|| _        nr |d       ni|dk(  r;t        |t              r|j                  || _        n@t        d	d
|  d| ddg       n)t        dd
|  d| d| dg t        j                         t        j                  d       S )Nc                 >    t        t        t        |       g       y )N)r   )r   	TypeErrorr#   )msgrH   s    r2   raise_errorz3ExceptionVariable.call_setattr.<locals>.raise_error  s    $Y:J3:O9PQr1   r   r   Tz9exception cause must be None or derive from BaseExceptionr   )TFr   z<Set Exception object `__traceback__` attribute to not-`None`zcall_setattr rq   zwDynamo does not support setting the attribute '__traceback__' on tracked exception objects to anything other than None.zYAvoid setting '__traceback__' on exception objects within traced code, or set it to None.rK   z4Unsupported attribute assignment on Exception objectz/Dynamo does not support setting the attribute 'z{' on tracked exception objects. Only `__context__`, `__cause__`, `__suppress_context__`, and `__traceback__` are supported.)rP   r   r~   r#   rU   r	   rB   r   !UserDefinedExceptionClassVariable"UserDefinedExceptionObjectVariabler   r   r   r   r   SUPPORTABLE)r<   rH   r   r   r   r^   s    `    r2   call_setattrzExceptionVariable.call_setattr  sz   	R **,= S![ 3 01cii6GJ--//??@@	M "%,5,F,Ft,L)WX++#/0SYY-5O,/)WX_$#/0SYY5F%(" Z+D64&9!'A
 N'vQtf56 ZZ 7)556 ))$//r1   c                     |dk(  r | j                   |g| S |dk(  r"|\  }| j                  |t        d      |       | S t        |   ||||      S )N__setattr__with_tracebackr   )r   r#   r8   r   )r<   rH   r^   r   r=   tbr>   s         r2   r   zExceptionVariable.call_method  sf    = $4$$R/$//%%DRb"2?"CRHK7&r4v>>r1   c                    |dk(  r| j                   S |dk(  r| j                  S |dk(  r| j                  S |dk(  rt        j                  d       S |dk(  r+t        j
                  | j                  | j                        S t        | %  ||      S )Nr   r   r   r   r   rg   )
r   r   r   r	   r#   ListVariabler   rR   r8   rm   )r<   rH   r^   r>   s      r2   rm   zExceptionVariable.var_getattr  s    = ###[ >>!++,,,_$--d33V^))$))DKKHHw"2t,,r1   c                 N    | j                   j                   d| j                   dS )N())r>   r-   r   r   s    r2   __str__zExceptionVariable.__str__  s$    ..))*!DMM?!<<r1   r   )rM   r   r   )r-   r.   r/   r9   r   rG   rQ   r"   r   r   rm   r   __repr__r   r   s   @r2   r   r     sO    4"#-*70#70 "70 	70r?-= Hr1   r   c                       e Zd ZdZy)UnknownVariablez
    It could be anything!
    Nr-   r.   r/   __doc__r0   r1   r2   r   r   	  s    r1   r   c                   <     e Zd ZdZd fd	Z	 	 	 	 	 	 	 	 ddZ xZS )DelayGraphBreakVariablez^
    Used to insert a dummy variable in the stack to do the graph break at CALL_FUNCTION.
    c                 2    t        |   di | || _        y r7   )r8   r9   r   )r<   r   r=   r>   s      r2   r9   z DelayGraphBreakVariable.__init__  s    "6"r1   c                     t        dd| j                   d| j                  j                          d| j                   g        y )Nz#Unsupported function call (delayed)zsource: z@Dynamo determined that a graph break should occur when calling `z`. Reason: rK   )r   rR   r^   r   r<   rH   r   r=   s       r2   r   z%DelayGraphBreakVariable.call_function  sH     	9t{{m,![[--/0DHH:G	
r1   rA   r   r-   r.   r/   r   r9   r   r   r   s   @r2   r   r     s8    
#
 &
 -	

 

r1   r   c                   D    e Zd ZdZddZdddeddfdZ	 	 	 	 	 	 	 	 dd	Zy
)ComptimeVariableza
    This variable is special, it lets you execute arbitrary code at
    Dynamo compile time
    c                     t        d      )Nzcomptime is special form)r   r<   rC   s     r2   rG   zComptimeVariable.reconstruct-  s    !"<==r1   rH   r*   r^   r5   r"   c                 j    ddl m } ddlm}  |t        ||      t	        | j
                  |            S )Nr   )comptimer!   )r%   rg   )r  	functionsr%   rY   r   rR   )r<   rH   r^   r  r%   s        r2   rm   zComptimeVariable.var_getattr0  s.    ' 	4#Hd#Jt{{D,I
 	
r1   c                 p   ddl m} |rJ t        |      dk  sJ |d   }t        |t              r |j                          ||             nt        |t              r|j                         }|j                  rJ d|j                          t        j                  ||j                  |j                  j                         |j                  rt!        |j                  j"                        nd d      } | ||             nt%        dt'        |             t(        j*                  j-                  d       S )Nr   )ComptimeContextr   zOcomptime function must not have free variables, but these variables were free: r0   z"unsupported argument to comptime: )r  r
  r[   r~   r%   get_functionr$   get_codeclosureco_freevarsr   r   	f_globalsfn_namerP   defaultstupler   RuntimeErrorrT   r	   r#   rl   )r<   rH   r   r=   r
  fncoder   s           r2   r   zComptimeVariable.call_function:  s    	/ z4yA~~!Wb./BOOob1267;;=Dzz 2262B2B1CE> %%

--/,.KKbkk''(T D $%!CDH:NOO))0066r1   Nr   r   )r-   r.   r/   r   rG   r   rm   r   r0   r1   r2   r  r  '  sS    
>
5 
S 
EV 
'7#'7 &'7 -	'7
 
'7r1   r  c                   X     e Zd ZU ee   ed<   dZee   ed<   	 ddee   ddf fdZ xZ	S )CellVariablepre_existing_contentsN
local_namer5   c                 2    t        |   di | || _        y r7   )r8   r9   r  )r<   r  r=   r>   s      r2   r9   zCellVariable.__init__q  s     	"6"%:"r1   rA   )
r-   r.   r/   r   r"   __annotations__r  r   r9   r   r   s   @r2   r  r  d  sF     $O44 !%J$ BF;%-o%>;	; ;r1   r  c                         e Zd Zd fdZ xZS )NewGlobalVariablec                 $    t        |   di | y r7   r8   r9   r<   r=   r>   s     r2   r9   zNewGlobalVariable.__init__y      "6"r1   r   )r-   r.   r/   r9   r   r   s   @r2   r  r  x  s    # #r1   r  c                 (      fd}t         |_        |S )Nc                  (     j                   | i |S rA   )apply)r   r=   fn_clss     r2   trampoline_autograd_applyzDproduce_trampoline_autograd_apply.<locals>.trampoline_autograd_apply~  s    v||T,V,,r1   )!produce_trampoline_autograd_apply_origin)r%  r&  s   ` r2   r'  r'  }  s    - )J%$$r1   c                   n     e Zd ZdZdhej
                  Zd fdZd	dZd	dZd	dZ		 	 	 	 	 	 d
dZ
 xZS )AutogradFunctionVariablez-represents a torch.autograd.Function subclassr%  c                 2    t        |   di | || _        y r7   )r8   r9   r%  )r<   r%  r=   r>   s      r2   r9   z!AutogradFunctionVariable.__init__  s    "6"r1   c           	      ^   dfd}t        j                  |||f       rt        j                         rt        j
                  du rt        j                  d       ddlm	} ddl
m} | j                  j                  } || j                  j                        }|r1 || j                  j                  | j                  j                        }| 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} | j0                  }|Ct3        j5                  | j                  j6                        | j                  j8                        } ||| j                  j:                  |t3        |d            j=                  ||      }| j0                  r~t3        | j0                  d      }t?        |jA                  tB        jD                               |r>t3        | j0                  d      }t?        |jA                  tB        jD                               |S | j0                  rt3        | j0                  d      }nd }| j                  j                  }tF        jI                  ||      }|g|}tK        |tL        jN                        rftQ        jR                  |      }tU        |      dz
  tU        |jV                        k(  r|dd  }tY        jZ                  ||      j=                  ||      S tK        |tL        j\                        rPtY        j^                  |j`                  tY        jb                  | j                        |      j=                  ||      S t%        dd|  d	| d	| dte        |      j8                   ddg       y )NFc                     t        | t        j                        r| j                  durdt        | t        j                        r| j                        rdy y y )NFT)r~   r	   TensorVariablerequires_gradNNModuleVariableis_training)r   r/  rH   s    r2   visitz2AutogradFunctionVariable.call_apply.<locals>.visit  sQ    "i667##50$(M"i889>>"%$(M & :r1   zNThe config.capture_autograd_function flag is deprecated, it's now always true.r   )#autograd_function_forward_rewritten)_is_setup_context_definedzUnsupported custom vjpzcall_apply rq   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.rK   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!   )AutogradFunctionApplyVariabler$  )r   rg   forwardsetup_contextz=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.)3r"   r2  rz   is_grad_enabledr   capture_autograd_functionwarningswarn"torch._functorch.autograd_functionr3  torch.autograd.functionr4  r%  r6  r7  vjpautogradFunctionr   r   r   jvphigher_order_opsr5  rR   r   import_sourcer.   r-   backwardr   r   rj   r   FUNCTION_MATCHAutogradFunctionContextVariablerl   r~   r   r   inspect	signaturer[   _parametersr	   r%   r   r   rt   r   rT   )r<   rH   r   r=   r2  r3  r4  
forward_fnis_setup_ctx_definedvjp_fnjvp_fnr5  rR   r   fwd_srcsetup_ctx_srcr  ctxsigr/  s    `                 @r2   
call_applyz#AutogradFunctionVariable.call_apply  s   	) 	edF^4U224//58d J,,J#<T[[=V=V#W # AKK'')B)B
 [[__FU^^44888 4)$qax@!- FH +66 [[__FU^^44888 4)$qax@!- F*66
 H[[F~#$$T[[%;%;<dkk>R>R 0$$!&9	
 mBf-  {{$T[[)<g001L1LMN'$.t{{O$LM!-":":<;V;V"WXJ;;Y7FF[[  -44RvF|d|b%,,-##B'C4y1}COO 44ABx11"VDRRD&  E,,-//224;;? mBf-	. W%dV1TF!F8<33783D3D2EZQ3
r1   c                 ^   | j                   j                  }t        |d   j                        t        j
                  j                  j                  u sJ t        |t        j                        sJ t        | j                  d      }t        j                  ||      j                  |||      S )Nr   rD  rg   )r%  rD  rT   rU   rz   _dynamoexternal_utilsFakeBackwardCFunctionr~   r   r   r   rR   r	   r%   r   )r<   rH   r   r=   r  	fn_sources         r2   call_backwardz&AutogradFunctionVariable.call_backward  s    [[!!DGMM"emm&B&B&X&XXXX"e00111t{{J7	--bCQQf
 	
r1   c                 ,    t        | j                        S rA   )r*  r%  r   s       r2   r   z&AutogradFunctionVariable.call_function  s    '44r1   c           
         ddl m} |dk(  rwt        j                  | j                        rEt        | j                        } || |j                  j                  d|gt        ||             S | j                  |||      S |dk(  r| j                  |||      S | j                  t        | j                  |      nd }	 t        j                  | j                  |      }t!        |t"              r|j%                  | j                        }	|7t        j&                  |	      j)                  |	|      j+                  |||      S  t        j&                  |	      |	      j+                  |||      S t!        |t,              r3t/        j0                  |j2                  | |      j+                  |||      S t5        dd	|  d
| d| ddg       y # t        $ r d }Y w xY w)Nr!   )wrap_fx_proxyr$  r   )rH   proxyrD  rg   z$Unsupported autograd.Function methodrp   rq   ,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.rK   )builderr[  r   is_callable_allowedr%  r'  r   create_proxyr   rR  rX  rR   r   rG  getattr_staticr   r~   r   __get__lookupcreate_with_sourcer   r   r	   r   rt   r   )
r<   rH   r^   r   r=   r[  r&  rR   objr   s
             r2   r   z$AutogradFunctionVariable.call_method  s    	+7?..t{{;,MKK-) %0"))00'1 +48  r488Z%%b$776:kk6MZT2SWF,,T[[$? #|,{{4;;/%#**40++D+@&r48 4;--d3D9GGD&  C- 33LL$v-D&12 !B*4&$8!v 9!9
U) " s     G GGr   r   )rH   r*   r   r   r=   r   )r-   r.   r/   r   r"   r   r9   rR  rX  r   r   r   r   s   @r2   r*  r*    sW    7 			'	'N
||
5>#> &	>
 ->r1   r*  c                   D    e Zd ZU  ej                  e      Zee   ed<   y)SavedTensorBox)default_factorytensorsN)	r-   r.   r/   dataclassesfieldlistri  r"   r  r0   r1   r2   rg  rg  ]  s    %6[%6%6t%LGT/"Lr1   rg  c                        e Zd ZdZdddhej
                  Z	 	 	 	 	 d
	 d fdZeddd       Zd Z		 	 	 	 	 	 	 	 d fdZ
d fd	Z xZS )rF  zZ
    Tracks an autograd.Function() context using mutation tracking in side_effects.py
    r\  	inferencesaved_tensorsc                 b    t        |   d||d| || _        || _        || _        || _        y N)rU   r   r0   )r8   r9   rn  ro  needs_input_gradnon_differentiable)	r<   rU   r   rn  ro  rr  rs  r=   r>   s	           r2   r9   z(AutogradFunctionContextVariable.__init__n  s;     	FuFvF"* 0"4r1   c           
         d }|r|st        d |D              }| j                  j                  j                  d t        j
                  j                  j                  t        j                  t        dt               |      i       }|S )Nc              3   l   K   | ],  }t        |t        j                        xr |j                   . y wrA   )r~   r	   r.  r/  .0xs     r2   	<genexpr>z9AutogradFunctionContextVariable.create.<locals>.<genexpr>  s1      % 1i667KAOOK%s   24T)rn  ro  rr  )r  r   r}   track_object_newrz   r?  functionFunctionCtx	functoolspartialrF  rg  )rH   r   r=   rr  outs        r2   rl   z&AutogradFunctionContextVariable.create~  s    $ %%   ii$$55NN##///,.!1	 

 
r1   c                 v    | 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.rK   )r\  r   r   r   r   s    r2   as_proxyz(AutogradFunctionContextVariable.as_proxy  sA    ::'#D6*25)445 zzr1   c                    |dk(  rt         |   ||||      S |dk(  rCt        |      dk(  sJ t        |i       d   | _        t
        j                  j                  d       S |dk7  r)t        dd|  d| d| d	g t        j                  
       | j                  t        dd|  d| ddg
       | j                  s6| j                  r|rJ |j                  j                  j!                  | |       t        | j                  j"                        dkD  rg | j                  _        |D ]'  }| j                  j"                  j%                  |       ) t
        j                  j                  d       S )Nr   mark_non_differentiabler   save_for_backwardz,Unsupported autograd.Function context methodrp   rq   r]  z` on `autograd.Function` context objects. Supported methods are `__setattr__`, `save_for_backward` and `mark_non_differentiable`.rK   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`.)r8   r   r[   r   rs  r	   r#   rl   r   r   r   ro  rn  rR   r   r}   track_save_for_backwardri  append)r<   rH   r^   r   r=   argr>   s         r2   r   z+AutogradFunctionContextVariable.call_method  s    = 7&r4v>>..v;!###&7b&A!&DD#--44T::&&F&tfAdV46 -- 7)556 %S&tfAdV4Ow
 ~~;;v--II""::4F t!!))*Q.)+D& 	3C&&--c2	3))0066r1   c                     dv rt         fd      S dk(  r> j                  2t        j                  t	         j                  j
                              S dk(  r j                  )t        j                  j                   j                        S  j                  rAt         j                  d      }t        j                   j                  j                  |      S t         =        S )N)r  r  c                  ,    j                  | |      S rA   )r   )r   r=   r^   r<   rH   s     r2   rD   z=AutogradFunctionContextVariable.var_getattr.<locals>.<lambda>  s    (8(8T4(P r1   ro  rr  )ry   ro  r	   TupleVariablerl  ri  rr  r#   rl   rR   r   r"   r   rU   r8   rm   r<   rH   r^   rR   r>   s   ``` r2   rm   z+AutogradFunctionContextVariable.var_getattr  s    CC!P  ?"t'9'9'E**40B0B0J0J+KLL%%$$0 11889N9NOO{{#DKK1CD&,,R1L1LfUUw"2t,,r1   )NFNNNr   NNr   r   )r-   r.   r/   r   r(   r   r9   r   rl   r  r   rm   r   r   s   @r2   rF  rF  b  s    
 	 
#	1	1	N 5 
5   (	.7#.7 &	.7
 -.7 
.7`- -r1   rF  c                   B     e Zd ZdZ	 d	 d fdZ	 	 	 	 	 	 	 	 ddZ xZS )AutogradEngineVariablez;
    Represents a torch._C._ImperativeEngine instance.
    c                 *    t        |   d||d| y rq  r  )r<   rU   r   r=   r>   s       r2   r9   zAutogradEngineVariable.__init__  s     	FuFvFr1   c                    |dk(  rt         j                  j                  j                  r|j                  s|j
                  sJ d       t        j                  t         j                  j                  j                  j                  | j                        j                  ||j                  j                  j                         g||      S t!        dd|  d| dg        y t!        d	d|  d| d
| dg        y )Nqueue_callbackzXqueue_callback() is only supported when Compiled Autograd is enabled with fullgraph=Truerg   z7Unsupported torch._C._ImperativeEngine.queue_callback()rp   rq   zYqueue_callback() is only supported when Compiled Autograd is enabled with fullgraph=True.rK   z-Unsupported torch._C._ImperativeEngine methodzgDynamo only supports the `queue_callback` method on a torch._C._ImperativeEngine instance, but found: `z`.)rz   rT  compiled_autogradin_compiled_autograd_region	one_grapherror_on_graph_breakr	   r%   rU  FakeCompiledAutogradEnginer  rR   r   r   r}   get_ca_final_callbacks_varr   r<   rH   r^   r   r=   s        r2   r   z"AutogradEngineVariable.call_method  s     ##}}..JJ||r'>'> n> !55MM00KKZZ;;  -YY++FFHP4P !U*4&$8!H G&tfAdV4IIMbRr1   rA   r   r   r-   r.   r/   r   r9   r   r   r   s   @r2   r  r    sI     G
 
G### &	#
 -# 
#r1   r  c                   8     e Zd Zd fdZ	 	 	 	 	 	 	 	 ddZ xZS )ry   c                 2    t        |   di | || _        y r7   )r8   r9   r  )r<   r  r=   r>   s      r2   r9   zLambdaVariable.__init__  s    "6"r1   c                 &     | j                   |i |S rA   )r  r   s       r2   r   zLambdaVariable.call_function  s     tww'''r1   r   r   )r-   r.   r/   r9   r   r   r   s   @r2   ry   ry     s3    (#( &( -	(
 
(r1   ry   c                        e Zd Zddhej                  Zdd fdZ fdZdefdZe	de
j                  j                  fd       Zd	 Zd
 ZddZddZ	 	 	 	 	 	 	 	 ddZdddee   deeef   def fdZd Z xZS )ri   r^   py_typer5   c                     t        |   di | t        |t              sJ t        |t              sJ || _        || _        || _        y r7   )r8   r9   r~   r"   r   re  r^   r  )r<   re  r^   r  r=   r>   s        r2   r9   zGetAttrVariable.__init__'  sH    "6"#///$$$$	r1   c                 P    | j                   | j                   S t        | 	         S rA   )r  r8   rQ   r<   r>   s    r2   rQ   zGetAttrVariable.python_type/  s%    <<#<<7&((r1   c                 h    | j                   j                   d| j                   d| j                   dS Nr   z, r   )r>   r-   re  r^   r   s    r2   r   zGetAttrVariable.__repr__5  s.    ..))*!DHH:R		{!DDr1   
base_proxyc                     t        | |      S rA   )rY   )r  r   s     r2   create_getattr_proxyz$GetAttrVariable.create_getattr_proxy8  s    z4((r1   c                 r    t         j                  | j                  j                         | j                        S rA   )ri   r  re  r  r^   r   s    r2   r  zGetAttrVariable.as_proxy<  s&    33DHH4E4E4GSSr1   c                     | j                   j                         }	 t        || j                        S # t        $ r t        |  d      d w xY w)Nz is not a constant)re  rP   rY   r^   r   r   )r<   constants     r2   rP   z"GetAttrVariable.as_python_constant?  sQ    88..0	M8TYY// 	M%.@&ABL	Ms	   2 ArH   r*   c                    t        | j                  t        j                        st        |j
                  j                  | j                  j                        }| j                  |j                  vrt        t        j                  || j                        }||j                  vrt        t        j                  ||      S rA   )r~   re  r	   r0  r   r   get_submodule
module_keyr^   r\   rG  ra  )r<   rH   r^   step1step2s        r2   const_getattrzGetAttrVariable.const_getattrF  s    $((I$>$>?%%		''(;(;<99ENN*%%&&udii8u~~%%%%%eT22r1   c                 |     || j                          |j                  |j                  | j                               y rA   )re  rF   create_load_attrsr^   r  s     r2   rG   zGetAttrVariable.reconstructQ  s+    g77		BCr1   r   r=   c                 R    | j                   j                  || j                  ||      S rA   )re  r   r^   r   s       r2   r   zGetAttrVariable.call_functionU  s#     xx##B		4@@r1   c                    |dv r| j                   dk(  r|s|d   j                         rt        | j                  t        j
                  t        j                  t        j                  f      rq| j                  }|d   j                         }|j                  ||      r|j                  ||      S |dk(  rt        |      dk(  r|d   S t	        j                  d       S |dk(  r| j                   dk(  rt        |      dk(  r|d   j                         r|st        | j                  t        j
                  t        j                  t        j                  f      r[| j                  }|d   j                         }|j                  ||      rt	        j                  d      S t	        j                  d	      S |d
k(  r| j                   dk(  r|st        | j                  t        j
                        r%| j                  j                  ||d   |d   d      S t        | j                  t        j                        r| j                  j                  |       t        | A  ||||      S )N)__getitem__r]   r\   r   r]   r   r!   __contains__TF__setitem__)directly_update_dict)r^   is_python_constantr~   re  r	   r(   r0  r   rP   has_key_in_generic_dictrm   r[   r#   r   convert_to_unspecializedr8   r   )r<   rH   r^   r   r=   re  keyr>   s          r2   r   zGetAttrVariable.call_method]  s    **		Z'Q**,77..66 ((Cq',,.C**2s3r3// u}t9>7N$55d;; N"		Z'D	QQ**,77..66 ((Cq',,.C**2s3 11$77 11%88]"tyyJ'>v$((I$G$GHxx77Qat 8   $((I$>$>?11"5w"2tT6::r1   c                 j   | j                   dk(  rSt        | j                  t        j                        r/|j
                  j                  j                  | j                        rJ d| j                  _        t        j                  || j                  j                  j                  | j                        S )Nr\   T)r^   r~   re  r	   r   r   r}   has_pending_mutationban_mutationr"   r   rU   r\   rR   )r<   rH   s     r2   get_forwarded_dictz"GetAttrVariable.get_forwarded_dict  s~    II#488Y%G%GHII**??I	
J !%$$R)@)@$++NNr1   rA   r   r   r   r   )r-   r.   r/   r"   r   r9   rQ   r   r   r   rz   fxProxyr  r  rP   r  rG   r   rl  dictr   r  r   r   s   @r2   ri   ri      s     
	'	'N)E# E ) ) )TM	3DA#A &A -	A
 
AB;#B; ?#	B;
 S/)*B; 
B;HOr1   ri   c                   H     e Zd Zd fdZ	 	 	 	 	 	 	 	 d fdZd Zd Z xZS )MethodWrapperVariablec                 @    t        |   di | || _        i | _        y r7   )r8   r9   method_wrapper_builtin_fns)r<   r  r=   r>   s      r2   r9   zMethodWrapperVariable.__init__  s#    "6",r1   c                 8   t        | j                        rnt        |d   t        j                        rQt        |      dk(  rt        |      dk(  sJ |d   j                  || j                  j                  j                        S | j                  j                  }| j                  j                  }|dk(  rUt        |      j                  }|t        j                  u r.t        j                  t              j                  |||g||      S t        | =  |||      S )Nr   r!   r9   )r   r  r~   r	   r.  r[   rm   r   r-   rT   r9   rx   rB   r   r8   r   )r<   rH   r   r=   self_objwrapper_namefn_objr>   s          r2   r   z#MethodWrapperVariable.call_function  s     &d&9&9:zGY--@
 t9>c&kQ&6667&&r4+>+>+G+G+P+PQQ &&//**33 :%(^,,F( 008DDx&7$&7  w$Rv66r1   c                      yNTr0   r   s    r2   r  z(MethodWrapperVariable.is_python_constant      r1   c                     | j                   S rA   )r  r   s    r2   rP   z(MethodWrapperVariable.as_python_constant  s    """r1   r   r   )r-   r.   r/   r9   r   r  rP   r   r   s   @r2   r  r    s=    
7#7 &7 -	7
 
7>#r1   r  c                   8     e Zd Zd fdZd fdZd Zd Z xZS )GetSetDescriptorVariablec                 2    t        |   di | || _        y r7   )r8   r9   desc)r<   r  r=   r>   s      r2   r9   z!GetSetDescriptorVariable.__init__  s    "6"	r1   c                     |dk(  rM| j                   rAt        | j                   d      }t        j                  || j                  j
                  |      S t        |   ||      S )Nrb  )rR   r   r"   r   r  rb  r8   rm   r  s       r2   rm   z$GetSetDescriptorVariable.var_getattr  sR    9Y7F"((TYY->->GG7&r400r1   c                      yr  r0   r   s    r2   r  z+GetSetDescriptorVariable.is_python_constant  r  r1   c                     | j                   S rA   )r  r   s    r2   rP   z+GetSetDescriptorVariable.as_python_constant  s    yyr1   r   r   )r-   r.   r/   r9   rm   r  rP   r   r   s   @r2   r  r    s    1r1   r  c                        e Zd Zddhej                  Zdej                  ddf fdZd Zd Z	de
fdZdd	Zdd
Z xZS )PythonModuleVariablerU   is_torchr5   Nc                     t        |   di | || _        | j                  t        u xs% | j                  j                  j                  d      | _        y )Nztorch.r0   )r8   r9   rU   rz   r-   
startswithr  r<   rU   r=   r>   s      r2   r9   zPythonModuleVariable.__init__  sD    "6"


e+Wtzz/B/B/M/Mh/Wr1   c                 "    t         j                  S rA   )r   
ModuleTyper   s    r2   rQ   z PythonModuleVariable.python_type  s    r1   c                     | j                   S rA   rU   r   s    r2   rP   z'PythonModuleVariable.as_python_constant      zzr1   c                 "    d| j                    dS )NzPythonModuleVariable(r   r  r   s    r2   r   zPythonModuleVariable.__repr__  s    &tzzl!44r1   c                 l    t        | j                  |      }t        j                  j	                  |      S rA   )r   rU   r	   r#   rl   r<   rH   r^   r   s       r2   call_obj_hasattrz%PythonModuleVariable.call_obj_hasattr  s)    T*))0088r1   c                    |j                   j                  j                  | |      r&|j                   j                  j                  | |      S | j                  s|| j
                  j                  vr	 t        | j
                  |      }n| j
                  j                  |   }| j                  xr t        | j                  |      }t        j                  ||      S # t        $ r t        t        |       Y Vw xY wrA   )r   r}   r   r   r  rU   r\   rY   r   r   rR   r   r"   r   )r<   rH   r^   r   rR   s        r2   rm   z PythonModuleVariable.var_getattr  s    99!!>>tTJ99))33D$??==D

(;(;;=$TZZ6
 ,,T2J>DKK!>$$RV<< " =(<=s   2C C98C9r   )r-   r.   r/   r"   r   r   r  r9   rQ   rP   r   r   r  rm   r   r   s   @r2   r  r    s\     
	'	'NXe.. XT X
 5# 59=r1   r  c                   V     e Zd Zd	 fdZ	 	 	 	 	 	 	 	 d
dZdddefdZd ZddZ xZ	S )TypingVariablec                 2    t        |   di | || _        y r7   r8   r9   rU   r  s      r2   r9   zTypingVariable.__init__      "6"
r1   rH   r*   c                     |dk(  r9t        |      dk(  r+| j                  |d   j                            }t        |      S t	        d       y )Nr  r!   r   z+unsupported method call on typing variablel)r[   rU   rP   r  r   )r<   rH   r^   r   r=   
new_typings         r2   r   zTypingVariable.call_method  sE     = SY!^DG$>$>$@AJ!*--CDr1   r^   c                    ddl m}m} |t        v rt	        j
                  | |      S |j                  j                  j                  | |      r|j                  j                  | |      S t        | j                  |      }| j                  r%t        | j                  |      }  |||      |      S |j                  ||      S )Nr!   )SourcelessBuilderVariableBuilder)r^  r  r  r   r	   ri   r   r}   r   r   rY   rU   rR   r   rl   )r<   rH   r^   r  r  rU   r   s          r2   rm   zTypingVariable.var_getattr%  s    ?)),,T48899!!>>tTJ??,,T488

D);;$T[[$7K3?2{3E::$++B66r1   c                     | j                   S rA   r  r   s    r2   rP   z!TypingVariable.as_python_constant5  r  r1   c                 X    |j                  |j                  | j                               y rA   )append_outputcreate_load_constrU   r  s     r2   rG   zTypingVariable.reconstruct8  s!    0 	g77

CDr1   r   r   )rC   r)   r5   N)
r-   r.   r/   r9   r   r   rm   rP   rG   r   r   s   @r2   r  r    sW    E#E &	E
 -E 
E75 7S 7 Er1   r  )maxsizec                      ddl m}  i }| j                         D ]H  \  }}|j                  j                         D ]&  \  }}t	        |      st        ||d      x}s"|||<   ( J |S )zd
    This generates a mapping from numpy modules to their torch._numpy
    modules equivalents.
    r   )NP_TO_TNP_MODULEN)utilsr  r   r\   callablerY   )r  np_fn_to_tnp_fnnp_modtnp_modr  tnp_fnnp_fns          r2   get_np_to_tnp_mapr  S  sz     )O+113 4&//557 	4OGV $FGT::5:-3OE*	44 r1   c                  j    t               } | j                         D ci c]  \  }}||
 c}}S c c}}w )z
    This is just the reverse mapping of get_np_to_tnp_map() - mapping from
    torch._numpy modules to numpy equivalents.
    )r  r   )mkvs      r2   get_tnp_to_np_mapr  h  s.     	AWWY'TQAqD'''s   /c                        e Zd ZdZej
                  fZd	 fdZed        Z	ed        Z
	 	 	 	 	 	 	 	 d
dZ	 	 	 	 	 	 	 	 d
dZd Z 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.
    c                 2    t        |   di | || _        y r7   r  r  s      r2   r9   zNumpyVariable.__init__y  r  r1   c                     |j                   j                  d      }t        |      dk\  r
|d d ddgk(  sJ || j                  v S N.r   rz   _numpy)r.   splitr[   constant_fold_functionsclsr  mods      r2   can_constant_fold_throughz'NumpyVariable.can_constant_fold_through}  sJ    mm!!#&3x1}RaWh,?!???S0000r1   c                     |j                   j                  d      }t        |      dk\  r
|d d ddgk(  sJ t        j	                  |d       S r	  )r.   r  r[   np_constant_collections_mapr]   r  s      r2    get_constant_collection_for_funcz.NumpyVariable.get_constant_collection_for_func  sL    mm!!#&3x1}RaWh,?!???*..r488r1   c                 
   t         j                  st        d| j                   d       ddlm} ddlm} t               j                  | j                        }|t        d| j                   d       | j                  |      x}	 e	  | | j                  |D cg c]  }|j                          c}i |j                         D 	
ci c]  \  }	}
|	|
j                          c}
}	      S |j                  dk(  rAt         j                   r1d|j"                   d}|dt         j                    z  }t        |       |j%                  |j                  ||      \  }}| j'                  |      rt)        ||      rt*        j,                  j/                   | j                         |D cg c]  }|j                          c}i |j                         D 	
ci c]  \  }	}
|	|
j                          c}
}	      S  |j0                  j2                  d ||      gt5        ||       }|j/                  ||      S c c}w c c}
}	w # t        $ r* t        | j                  j                   d	| d
|        Y y w xY wc c}w c c}
}	w )Nznumpy.z()r   )numpy_to_tensor_wrapperr!   )NumpyNdarrayVariablezCan't find numpy function zM in torch._numpy.  Please file an issue to request support for this function.z with non-const args: rq   ztorch._numpy.randomz
delegate 'z' to NumPy itself via zconfig.use_numpy_random_stream=r   )r   trace_numpyr   rU   r  r  tensorr  r  r]   r  rP   r   r   r-   r.   use_numpy_random_streamr/   
patch_argsr  r   r	   r#   rl   r   r`  r   )r<   rH   r   r=   r  r  r   collection_variable_typrx  r  r  r   r\  s                r2   r   zNumpyVariable.call_function  so    !!F4::,b1230 "&&tzz2<,TZZL 9N N (,'L'LT'RR#
.DJJ:>?Q!..0?AGPA1a2244P  #8822"4#4#4"55KL5f6T6T5UV c"/::4==$PVWLD&--d3-dF; !1188-D++-:>?Q!..0?AGPA1a2244P  +BII**'- #40E
 (..r599M @P ' zz**++A$qQ. @Ps<   	I H92I H>#	I 'I:I?9I 0I76I7c                     t        d       y )Nnumpy)r   r  s        r2   r   zNumpyVariable.call_method  s     	gr1   c                     | j                   S rA   r  r   s    r2   rP   z NumpyVariable.as_python_constant  r  r1   c                     t         j                  r0t        | j                  t              r| j                  j
                  S t        |          S rA   )r   r  r~   rU   rT   r-   r8   r  r  s    r2   r  zNumpyVariable.as_proxy  s9    *TZZ"> ::&&&w!!r1   r   r   )r-   r.   r/   r   tnp
issubdtyper  r9   r   r  r  r   r   rP   r  r   r   s   @r2   r  r  r  s      #~~/ 1 1
 9 9
@:#@: &@: -	@:
 
@:D# &	
 - 
" "r1   r  c                   4     e Zd Zd fdZdefdZddZ xZS )NullVariabler5   c                 $    t        |   di | y r7   r  r   s     r2   r9   zNullVariable.__init__  r!  r1   c                      y)Nr$  r0   r   s    r2   r   zNullVariable.__repr__  s    r1   c                 t    t         j                  dk  rt        d       |j                  t	        d             y )N)rw      z0cannot reconstruct NullVariable in < Python 3.11	PUSH_NULL)sysversion_infor   r  r   r  s     r2   rG   zNullVariable.reconstruct  s-    g%LM0=>r1   r   r   )r-   r.   r/   r9   r   r   rG   r   r   s   @r2   r$  r$    s    ## ?r1   r$  c                       e Zd ZdZy)r   z"Marker used to implement delattr()Nr   r0   r1   r2   r   r     s    ,r1   r   c                   f     e Zd ZdZdhej
                  Zed        Zd 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_stringc                    t        d t        j                  ||j                               D              r{t        j
                  j                   |j                  |D cg c]  }|j                          c}i |j                         D ci c]  \  }}||j                          c}}      S  | |t        |      t        |            S c c}w c c}}w )Nc              3   <   K   | ]  }|j                           y wrA   )r  rv  s     r2   ry  z.StringFormatVariable.create.<locals>.<genexpr>  s!      
   "
s   )all	itertoolschainvaluesr	   r#   rl   formatrP   r   rl  r  )r  r/  sym_args
sym_kwargsr  r  s         r2   rl   zStringFormatVariable.create  s     
__Xz/@/@/BC
 
 --44$$$6>?a**,?=G=M=M=OPTQq!..00P  =$x.$z2BCC	 @Ps   CCr5   c                 r    t        |   di | t        |t              sJ || _        || _        || _        y r7   )r8   r9   r~   r   r/  r7  r8  )r<   r/  r7  r8  r=   r>   s        r2   r9   zStringFormatVariable.__init__  s9    "6"----* $r1   c                     | j                   j                   d| j                  d| j                  d| j                  dS r  )r>   r-   r/  r7  r8  r   s    r2   r   zStringFormatVariable.__repr__  sB    ..))*!D,>,>+ADMMCTTVW[WfWfViijkkr1   c                     j                   fdd        t        j                   j                                j                  j                         D ci c]%  \  }}t        j                  j                  |      |' }}} t        j                  |             j                  t        dd             y c c}}w )Nc                  z     j                   j                  j                         j                  d      g      S )Nr6  )rF   r  r/  create_load_attrr   s   r2   rD   z2StringFormatVariable.reconstruct.<locals>.<lambda>  s9    G))--d.@.@A,,X6 r1   T)call_function_exCALL_FUNCTION_EXr!   )r  )rE   r	   r  r7  r8  r   r#   rl   ConstDictVariabler  r   )r<   rC   r  r  r=   s   ``   r2   rG   z StringFormatVariable.reconstruct  s     " 	 	
 		''67@D@U@U@W
8<1I&&--a0!3
 
 		++F3401CKL	
s   *Cr   r   )r-   r.   r/   r   r"   r   r   rl   r9   r   r   rG   r   r   s   @r2   r.  r.    sJ     &G(F(FGND D%l# lMr1   r.  c                   V     e Zd ZdZd fdZed        ZddZd	dZed
d       Z	 xZ
S )DebuggingVariablez
    Represents a call to a debugging function like print(), or something
    registered to config.reorderable_logging_functions.
    c                 2    t        |   di | || _        y r7   r  r  s      r2   r9   zDebuggingVariable.__init__*  r  r1   c                     t        |       xrR t        | t        j                  t        j                  f      xr& | t
        j                  j                  j                  v S rA   )	r  r~   r   r   BuiltinFunctionTyperz   rT  r   reorderable_logging_functions)re  s    r2   is_reorderable_logging_functionz1DebuggingVariable.is_reorderable_logging_function.  sO     SM J3!3!3U5N5N OPJu}}++III	
r1   c           	          |j                   ry | j                  | j                  ||      st        d| j                   d| d| d       |j                  j                  | t        |      f       y )NzReordering debugging function z with inputs rq   z is not yet implemented.)exportcan_reorder_logsrU   r   debug_localsr  rl  r   s       r2   r   zDebuggingVariable.call_function6  sk    99$$TZZv>0 =#fAfX-EG
 	d4j12r1   c                 8    | j                   j                  |      S rA   )rR   rG   r  s     r2   rG   zDebuggingVariable.reconstructC  s    {{&&w//r1   c                     t         j                  t         j                  t        f}t	        j
                  ||g      }|D ]  }t        ||      r y y)zm
        Run some additional checks for what sort of function calls can we
        actually reorder.
        FT)r	   r.  r#   r.  pytreetree_leavesr~   )r  r   r=   allowed_input_types	flat_argsr  s         r2   rJ  z"DebuggingVariable.can_reorder_logsF  sY     $$&& 
 &&f~6	 	Cc#67	 r1   r   r   r   )r5   T)r-   r.   r/   r   r9   r   rG  r   rG   rJ  r   r   s   @r2   rB  rB  $  s?    
 
 
30  r1   rB  c                   <     e Zd ZdZd fdZ	 	 	 	 	 	 	 	 ddZ xZS )LoggingLoggerVariablez<
    Represents a call to any of logging.Logger methods
    c                 2    t        |   di | || _        y r7   r  r  s      r2   r9   zLoggingLoggerVariable.__init__`  r  r1   c                 $   |j                   ry t        | j                  |d       }t        |dd       }||hj                  t        j
                  j                  j                        rt        j                  j                  d       S t        d       y )Nrt   zLogger not supported for non-export cases. To avoid graph breaks caused by logger in compile-mode, it is recommended to disable logging by adding logging methods to config.ignore_logger_methods)rI  rY   rU   intersectionrz   rT  r   ignore_logger_methodsr	   r#   rl   r   )r<   rH   r^   r   r=   methodr{  s          r2   r   z!LoggingLoggerVariable.call_methodd  sv     99T406:t4H**5==+?+?+U+UV--44T::Y	
r1   r   r   r  r   s   @r2   rS  rS  [  s8    
#
 &	

 -
 

r1   rS  c                        e Zd ZdZd Z	 ddlmZmZ	m
Z d fdZd Zd	d
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 # e$ r  eddi       Z	 eddi       ZY Fw xY w)ConstantLikeVariablez8self.value is a compile-time constant, but not a literalr   )dtypefloatinggenericinvalid_typer0   r5   c                 2    t        |   di | || _        y r7   r  r  s      r2   r9   zConstantLikeVariable.__init__  r  r1   c                     | j                   S rA   r  r   s    r2   rP   z'ConstantLikeVariable.as_python_constant  r  r1   rH   r*   r   r=   c                 P   	 |D cg c]  }|j                          }}|j                         D ci c]  \  }}||j                          }	}} t        | j                  |      i 	}
t        j                  j                  |
      rt        j                  j                  |
      S t        |
t        j                        rt        |
      S t        | j                   d| d|
        y c c}w c c}}w # t        $ r$ t        | j                   d| d| d| d       Y w xY w)Nr
  z(*z, **r   z() -> )rP   r   r   r   _error_prefixrY   rU   r	   r#   rh   rl   r~   reMatchConstantRegexMatchVariable)r<   rH   r^   r   r=   rx  cargsr  r  ckwargsr   s              r2   r   z ConstantLikeVariable.call_method  s   	O59:Q))+:E:=C\\^LTQq!..00LGL +T*E=W=%%008--44V<<fbhh'-f55++,AdV6&BC ;L" 	OT//0$r$tF81MN	Os,   C8 C-C8 C2C8 -C8 8*D%$D%r^   c                    t        | j                  |      }t        || j                        rt	        |      }t        || j
                        rt        |      S t        |t              r!t        || j                        rt        |      S t        j                  j                  |      rt        j                  j                  |      S t        | |      S rA   )rY   rU   r~   np_floatingfloatnp_dtypeNumpyDTypeVariablerT   rS   
np_genericr  r	   r#   rh   rl   ri   r  s       r2   rm   z ConstantLikeVariable.var_getattr  s    T*fd../6]Ffdmm,%f--fd#
64??(K ((%%008--44V<<tT**r1   r   )r-   r.   r/   r   rb  r  r[  rk  r\  ri  r]  rm  ImportErrorrT   r9   rP   rl  r"   r  r   r   rm   r   r   s   @r2   rZ  rZ  y  s    B*M0	
 	
D#D ?#	D
 S/)*D 
D.+5 +S +_ +E  0>2r2B/0s   
A   AArZ  c                       e Zd ZdZy)RegexPatternVariablez
re.PatternNr-   r.   r/   rb  r0   r1   r2   rp  rp    s     Mr1   rp  c                       e Zd ZdZy)re  zre.MatchNrq  r0   r1   r2   re  re    s    Mr1   re  c                   $     e Zd ZdZd fdZ xZS )TorchVersionVariableztorch.__version__c                     |j                  dt        j                         |d   t        j                  u sJ t        |   di | y )NrU   r0   )
setdefaultrz   __version__r8   r9   r   s     r2   r9   zTorchVersionVariable.__init__  sA    '5#4#45g%"3"3333"6"r1   r   )r-   r.   r/   rb  r9   r   r   s   @r2   rt  rt    s    'M# #r1   rt  c                       e Zd ZdZy)NumpyTypeInfoVariableznp.iinfo/np.finfoNrq  r0   r1   r2   ry  ry    s    'Mr1   ry  c                       e Zd ZdZd Zy)rl  znp.dtype[...]c                 B    | 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).
        )rU   rT   r-   r   s    r2   r  zNumpyDTypeVariable.as_proxy  s     zz'''r1   N)r-   r.   r/   rb  r  r0   r1   r2   rl  rl    s    #M(r1   rl  c                   ,     e Zd ZdZd fdZddZ xZS )RandomClassVariablezrandom.Randomc                 $    t        |   di | y r7   r  r   s     r2   r9   zRandomClassVariable.__init__  r!  r1   c                    t        |      dkD  rt        d       n|rt        d       t        |      dk(  rt        j                  j	                  d       n|d   }t        |t        j                  j                               S )Nr!   zrandom.Random() with > 1 argzrandom.Random() with kwargsr   )seedr   )r[   r   r	   r#   rl   RandomVariablebaseValueMutationNew)r<   rH   r   r=   r  s        r2   r   z!RandomClassVariable.call_function  sh    t9q=8978:=d)q.y))006dSTgY^^%D%D%F
 	
r1   r   r   r  r   s   @r2   r}  r}    s    #
r1   r}  c                        e Zd ZdZdhej
                  Zh dZ	 	 ddeej                     dee   ddf fdZ
d	 Zd
 Zed        Zed        Zed        Zed        Zdd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  r5   c                 0   t        |   di | |V| j                  |      sJ t        j                         | _        | j                  j                  |j                                y ||j                         nd }t        j                  |      | _        y r7   )r8   r9   is_supported_random_objr  RandomsetstategetstaterP   )r<   r  r  r=   r>   s       r2   r9   zRandomVariable.__init__   sy     	"6"//555 --/DKKK  1040@4**,dD ---DKr1   c                 "    t         j                  S rA   )r  r  r   s    r2   rQ   zRandomVariable.python_type  s    }}r1   c                     | j                   S rA   )r  r   s    r2   rP   z!RandomVariable.as_python_constant  s    {{r1   c                    t        |       t        j                  uryt        j                  t
        j                  d      D ]  }t        | |      s yt        | |      }t        j                  |      r/|t        t        j                  |      j                  |       k7  s_ yt        |dd       t        t        j                  |      us y y)NF)r  r  r  rt   T)rT   r  r  r3  r4  r  _supported_fn_namesr   rY   rG  	isbuiltinrb  )r   r^   meths      r2   r  z&RandomVariable.is_supported_random_obj  s    9FMM)OO..0P
 	!D 3%3%D  &76==$7??DD 4T2'&--QU:VV 	! r1   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 y y )Nr   r!   c              3   >   K   | ]  }t        |      t        u   y wrA   )rT   intrv  s     r2   ry  z-RandomVariable.check_state.<locals>.<genexpr>-  s     4a47c>4s   r   )rT   r  r  r2  rj  )states    r2   check_statezRandomVariable.check_state(  s}    E{e###E!H~$$$E!H~&&&4584444Qx4a>U#:::#:r1   c                 h   t         j                  |        t        j                  t        j                  j                  | d         t        j                  | d   D cg c]!  }t        j                  j                  |      # c}      t        j                  j                  | d         g      S c c}w )Nr   r!   r   )r  r  r	   r  r#   rl   )r  rx  s     r2   
wrap_statezRandomVariable.wrap_state0  s    ""5)&&**11%(;''CH8LaY//66q9L **11%(;
 	
 Ms   &B/c                 P    | j                         }t        j                  |       |S rA   )rP   r  r  )r  	state_objs     r2   unwrap_statezRandomVariable.unwrap_state=  s$    ,,.	""9-r1   rH   r*   r   r=   c                    dk(  r|j                   j                  j                  |         | j                  j                  |D cg c]  }|j                          c}i |j                         D ci c]  \  }}||j                          c}} t        j                  j                  d       S dk(  r)| j                  | j                  j                               S dk(  rq|j                   j                  j                  |        | j                  j                  | j                  |d                t        j                  j                  d       S | j                  v r|j                   j                  j                  |        | j                  j                         fd} t        | j                        |D cg c]  }|j                          c}i |j                         D 	
ci c]  \  }	}
|	|
j                          c}
}	 t!        ||||      S t"        | I  |||      S c c}w c c}}w c c}w c c}
}	w )Nr  r  r  r   c                  r    t        j                         }|j                          t        |      | i |S rA   )r  r  r  rY   )r   r=   rr^   r  s      r2   call_random_methz4RandomVariable.call_method.<locals>.call_random_meth[  s3    MMO

5!'wq$'888r1   )r   r}   mutationr  r  rP   r   r	   r#   rl   r  r  r  r  r  rY   r&   r8   r   )r<   rH   r^   r   r=   rx  r  r   r  r  r  r  r>   s     `        @r2   r   zRandomVariable.call_methodC  s    6>II""++D1DKK267Q!&&(7AGPXS#3..00P --44T::Z??4;;#7#7#9::ZII""++D1KK  !2!247!;<--44T::T---II""++D1KK((*E9 'GDKK&267Q!&&(79?HA1a**,,H
 ""&6fEEw"2tT6::7 8P* 8Hs   H45H9H?3Ic                 6   j                  fd       j                  dd       j                          j                  d        | j	                  | j
                  j                                      j                  dd       j                          y )Nc                  n     j                   j                  t               j                  d      g      S )Nr  )rF   create_load_python_moduler  r=  rC   s   r2   rD   z,RandomVariable.reconstruct.<locals>.<lambda>l  s3    G))55f=,,X6 r1   r   Fr  r!   T)rE   r   r   r   r  r  r  pop_topr  s    `r2   rG   zRandomVariable.reconstructj  sz    	
 	a' 	*% 4 4 678a&r1   r  r   )r-   r.   r/   r   r"   r   r  r   r  r  r9   rQ   rP   r   r  r  r  r  rl  r  r   r   rG   r   r   s   @r2   r  r    s     			'	'N
 )-*..v}}%. '.
 
.  $ ; ; 

 

  
%;#%; ?#	%;
 S/)*%; 
%;Nr1   r  c                   N     e Zd Zed        Z fdZ	 	 	 	 	 	 	 	 ddZddZ xZS )WeakRefVariablec                 
   |j                  dd       }|j                  }|xr t        |d      }t        j                  | ||      } |       }|xr t        |      }t        j                  | ||      }||d<   t        ||fi |S )NrR   __callback__)r]   r  r   r"   r   r   r  )	rH   weakref_valueoptionsrR   callbackcallback_sourcecallback_vtreferentreferent_vts	            r2   r   zWeakRefVariable.build~  s    Xt, -- GZ%G%++B/J ?5-f5%++B&A"{KC7CCr1   c                 @    t        |   di | || _        || _        y r7   )r8   r9   r  r  )r<   r  r  r  r>   s       r2   r9   zWeakRefVariable.__init__  s#    #7#&&r1   c                     | j                   S rA   )r  r   s       r2   r   zWeakRefVariable.call_function  s     r1   c                     j                  fd        | j                          | j                         j                  t	        dd             y )Nc                  (     j                  dd      S )Nweakrefref)r   r  s   r2   rD   z-WeakRefVariable.reconstruct.<locals>.<lambda>  s    g&>&>y%&P r1   r   F)rE   r  r  rF   r
   r  s    `r2   rG   zWeakRefVariable.reconstruct  sD    PQ  !  !21e<=r1   r   r   )	r-   r.   r/   r   r   r9   r   rG   r   r   s   @r2   r  r  }  sL    	D 	D'
 #  &  -	 
 
 >r1   r  )qr   rj  r}  rG  r3  r  rc  r*  r   r:  typingr   r   torch._Crz   torch._numpyr  r!  torch.utils._pytreer  _pytreerN  rJ   r   r   r   r	   bytecode_transformationr
   r   create_parameter_opr   r   r   r   r   guardsr   r   mutation_guardr   rR   r   r   r   r   r   r   r   r   r   r   r   r   r   r    r  r"   r  r#   r  r$   r%   user_definedr&   r'   r(   torch._dynamo.codegenr)   torch._dynamo.symbolic_convertr*   r,   r4   r   r   r   r  r  r  r'  r*  	dataclassrg  rF  r  ry   ri   r  r  r  r  	lru_cacher  r  r  r$  r   r.  rB  rS  rZ  rp  re  rt  ry  rl  finfoiinfor[  r  r}  r  r  r0   r1   r2   <module>r     s  "      	 
   *   $ $ @ @ N F K K 0 5 	 	 	 " & G X X /D	 	x
O x
v	A AHo 
o 
0:7 :7z;? ;(# #
%U Up M M Mz-&? z-z06 0f(_ (HOo HOV)#O )#X &&=? &=R=E_ =E@ Q  ( Q(  (l"O l"`
?? 
?-o -.M? .Mb4 4n
O 
<7+? 7+t!/ !!5 #/ #(0 (	(- 	( II$II$II! 
/ 
"N_ Nb>o >r1   