
    ȅi~                        S r SSKrSSKrSSKrSSKrSSKrSSKJrJr  SSK	J
r
Jr  SSKJrJrJr  SSKrSSKJr  SSKJrJrJr  SS	KJrJrJr  SS
KJrJr  SSKJ r   SSK!J"r"J#r#J$r$J%r%J&r&J'r'J(r(  SSK)J*r*J+r+J,r,J-r-J.r.J/r/J0r0J1r1J2r2J3r3J4r4J5r5J6r6J7r7  SSK8J9r9J:r:J;r;J<r<  SSK=J>r>  SSK?J@r@  SSKAJBrB  SSKCJDrD  \(       a  SSKEJFrF  SSKGJHrH  SSS\R                  R                  S\\<   S\K\L\<4   SS4
S jrM\
S\LS\SSS\R                  R                  S\4
S j5       rNS\\   S\R                  R                  SS4S  jrO " S! S"\<5      rP " S# S$\D5      rQ " S% S&\Q5      rR " S' S(\Q5      rSg))aQ  
This module implements variable tracking for PyTorch nn.Module instances during Dynamo tracing.

It provides specialized handling for different types of nn.Module instances through several key classes:

- NNModuleVariable: Handles instance-specific module tracing, specializing on module id() and placing
  parameters directly on the torch.fx.GraphModule. This creates one graph per module instance.

- UnspecializedNNModuleVariable: Provides class-level module tracing, treating nn.Modules like other
  user-defined objects and passing parameters as inputs to the FX graph. This creates one graph per
  module class.

- UnspecializedBuiltinNNModuleVariable: Specifically handles built-in PyTorch modules (e.g. nn.Linear)
  with appropriate optimizations.

- FSDPManagedNNModuleVariable: Special handling for FSDP-wrapped modules with modified guarding behavior
  and parameter handling.

The module integrates with Dynamo's broader tracing functionality to handle module method calls,
parameter access, hooks, and other nn.Module behaviors while maintaining proper scoping and guarding
of module state.
    N)IterableSequence)contextmanagernullcontext)AnyOptionalTYPE_CHECKING)Source   )graph_break_hintstrace_rules	variables)raise_observed_exceptionunimplementedUnspecializeRestartAnalysis)GuardBuilderinstall_guard)GenerationTracker)
AttrSourceConstDictKeySourceDictGetItemSourceFSDPNNModuleSourceGetItemSourceNNModuleSourceUnspecializedNNModuleSource)get_custom_getattrget_fake_valueis_lazy_moduleis_namedtupleis_safe_constantistensoristypennmodule_has_hooksobject_has_getattributeproxy_args_kwargsraise_args_mismatchset_example_valueunpatched_nn_module_callunpatched_nn_module_call_impl   )raise_type_error_exctypestrValueMutationNewVariableTracker)invoke_and_store_as_constant)LazyVariableTracker)SliceVariable)UserDefinedObjectVariable)InstructionTranslator)ConstantVariabletxr3   modargskwargsreturnc                   ^ ^ [        US5      (       ay  S[        S[        4UU 4S jjm[        X#5      u  pEU Vs/ s H  nT" U5      PM     nnUR                  5        VV	s0 s H  u  pUT" U	5      _M     n
nn	 UR	                  XU
5        ggs  snf s  sn	nf ! [
         a5  n[        [
        T [        U5      (       a  [        U5      OS/S9   SnAgSnAff = f)aj  
Fairly coupled helper used by NNModuleVariable and UnspecializedNNModuleVariable.

Used to cause lazy module to be initialized (and delete its init hook) before tracing. Especially
useful now that 'allowed' modules graph-break on hooks, calling this first ensures there is no hook
by the time we trace __call__ and thus no graph-break for lazy allowed modules.
_initialize_hookxr9   c                   > [        U 5      (       a  [        U 5      " U4S jU  5       6 $ [        U [        5      (       a,  U R	                  5        VVs0 s H  u  pUT" U5      _M     snn$ [        U [
        [        [        45      (       a  [        U 5      " U4S jU  5       5      $ [        U [        R                  R                  5      (       a  [        U R                  T5      $ U $ s  snnf )Nc              3   4   >#    U  H  nT" U5      v   M     g 7fN .0elemconvert_to_fakes     [/home/james-whalen/.local/lib/python3.13/site-packages/torch/_dynamo/variables/nn_module.py	<genexpr>Binitialize_lazy_module.<locals>.convert_to_fake.<locals>.<genexpr>_   s      E14!6!61   c              3   4   >#    U  H  nT" U5      v   M     g 7fr?   r@   rA   s     rE   rF   rG   c   s     Ct44rH   )r   type
isinstancedictitemslisttuplesettorchfxProxyr   node)r<   kvrD   r5   s      rE   rD   /initialize_lazy_module.<locals>.convert_to_fake]   s    QAw E1 EFFAt$$:;'')D)$!?1--)DDAeS122AwCCCCAuxx~~..%affb11 Es   C'z0AttributeError during lazy module initializationr7   N)hasattrr   r%   rM   _infer_parametersAttributeErrorr   str)r5   r6   r7   r8   
proxy_argsproxy_kwargsarg	fake_argsrU   rV   fake_kwargserD   s   `           @rE   initialize_lazy_modulerc   N   s     s&''
	s 
	s 
	 
	 $5T#B 
5?@Zc_S)Z	@9E9K9K9MN9Mq/!,,9MN	!!#+>% ( AN  
	$ 1vv FK
	s#   B!B:B 
C#+CC
module_keysourcec              #   J  #    UR                   n[        R                  " SSU5      nUR                  R	                  US5      nUS:  a  U  SU 3OU n  XCR
                  4UR                  U '   US-   UR                  U'   S v   UR                  U 	 g ! UR                  U 	 f = f7f)Nz9\._(?:modules|parameters|buffers)\[(['\"])([^'\"\]]+)\1\]z.\2r   @r*   )nameresub	num_callsget	__class__nn_module_stack)rd   re   r5   r6   fully_qualified_namerk   s         rE   record_nn_module_stackrp   {   s      ";;66D
   !5q9I09AJ<q,:J+*>)N:&-6])*z*Bz*s   AB#0B B#B  B#c           	         U (       a  SUR                   ;   a  [        UR                   S   5      (       al  UR                   S   n[        U S5      n[        U5      [        R
                  L a  [        US5      n[        UR                  [        R                  5      5        g [        U R                  [        R                  " [        R                  SS95      5        g g )Nforward__func__)attr)__dict__callabler   rJ   types
MethodTyper   
make_guardr   CLOSURE_MATCH	functoolspartialNOT_PRESENT_IN_GENERIC_DICT)re   r6   fwdforward_sources       rE   &guard_to_detect_forward_monkeypatchingr      s      $#,,y2I)J)J,,y)C'	:NCyE,,,!+NJ!G.33L4N4NOP !!%%$@@y     c                   .  ^  \ rS rSrSSSS1\R
                  krS\S\S\R                  R                  S\SS4
U 4S	 jjrS\4S
 jrS\SS4S jrS\4S jrSSS\S\R                  R                  S\S\SS4S jrSSS\\   4S jrSSS\SS4S jrSSS\4S jrS%S jrSSS\S\4S jrS\R                  R                  SSS\S\S\\   4
S jrSSS\S\4S jrSSS \\   S\\\4   S\4S! jr S&SSS\S \\   S\\\4   S"\S\4U 4S# jjjrS$r U =r!$ )'NNModuleVariable   module_typerd   valuenn_module_stack_sourcer8   r9   Nc                    > [         TU ]  " S0 UD6  Xl        X l        X0l        U R
                  U l        U R
                  U l        g )Nr@   )super__init__r   rd   r   re   r   )selfr   rd   r   r8   rm   s        rE   r   NNModuleVariable.__init__   s>     	"6"&$
 #kk&*kk#r   c                 V    U R                   =(       d    U R                  nU(       d   eU$ r?   r   re   r   ress     rE   get_nn_module_stack_source+NNModuleVariable.get_nn_module_stack_source   "    ))8T[[
s
r   re   c                     Xl         g r?   r   r   re   s     rE   set_nn_module_stack_source+NNModuleVariable.set_nn_module_stack_source       &,#r   c                     U R                   $ r?   )r   )r   s    rE   python_typeNNModuleVariable.python_type   s    r   r5   r3   submod	key_extraoptionsc                     g r?   r@   )r   r5   re   r   r   r   s         rE   _wrap_submodule NNModuleVariable._wrap_submodule   s     	r   c                 f   UR                   R                  U R                  5      n/ n[        U[        R
                  R                  5      (       a  UR                  5        Hx  u  pE[        R                  R                  U5      nUR                   R                  UU R                  U[        [        U R                  U5      5      S9  UR                  U5        Mz     U$ [        U[        R
                  R                   [        R
                  R"                  [        R
                  R$                  45      (       d   ['        U5      5       e[)        U5       HW  u  puUR                  UR                   R                  UU R                  U[        [        U R                  U5      5      S95        MY     U$ Nre   )outputget_submodulerd   rK   rQ   nn
ModuleDictrM   r   r4   createregister_attr_or_moduler   r   re   append
ModuleListParameterList
Sequentialr,   	enumerate)r   r5   baseresultrh   r   name_varidxs           rE   unpack_var_sequence$NNModuleVariable.unpack_var_sequence   sS   yy&&t7(*dEHH//00 $

$55<<TB		11OO)-T*JK	 2  h' !- M588&&(>(>@S@ST
 
 	4=	 
 %T?KCMM		11OO)-S*IJ	 2  + r   rh   r4   c                 *   UR                   R                  U R                  5      n[        X25      n[	        [        [        U R                  U5      5      R                  [        R                  5      5        [        R                  R                  U5      $ r?   )r   r   rd   rY   r   r   r   re   ry   r   HASATTRr   r4   r   )r   r5   rh   r6   r   s        rE   call_obj_hasattr!NNModuleVariable.call_obj_hasattr   sn     ii%%doo6#:dkk489DD$$	

 ))0088r   c                 f    UR                   R                  U R                  5      n[        USS5      $ )NtrainingF)r   r   rd   getattrr   r5   r6   s      rE   is_trainingNNModuleVariable.is_training  s)    ii%%doo6sJ..r   c                     UR                   R                  U R                  5      n[        R                  " U5        UR
                  R                  S:w  a  [        R                  " [        U5      5        [        e)zIRestart analysis treating this module as an UnspecializedNNModuleVariabler   )
r   r   rd   r   tagf_codeco_namemark_class_dynamicrJ   r   r   s      rE   convert_to_unspecialized)NNModuleVariable.convert_to_unspecialized  sT    ii%%doo6c" 99
*00c;))r   keyc                    UR                   R                  U R                  5      n[        U5      (       a#  [	        SSU  SU 3SS/[
        R                  QS9  UR                   R                  R                  X5      (       aC  UR                   R                  R                  XSS9n[        U[        R                  5      (       + $ [        R                  US	5      nX%;   $ )
Nz3Custom __getattribute__ in nn.Module dict key checkzhas_key_in_generic_dict  }Dynamo does not support checking key existence on `nn.Module` instances that have a custom `__getattribute__` method defined.1Avoid defining `__getattribute__` in your module.gb_typecontextexplanationhintsT)
deleted_okru   )r   r   rd   r$   r   r   SUPPORTABLEside_effectshas_pending_mutation_of_attr	load_attrrK   r   DeletedVariableobject__getattribute__)r   r5   r   r   mutated_attr	base_dicts         rE   has_key_in_generic_dict(NNModuleVariable.has_key_in_generic_dict  s    yy&&t7"4((M24&#?5 H&22
 99!!>>tII9911;;DRV;WL!,	0I0IJJJ++D*=	r   r   
obj_sourcec                    [        U5      (       a#  [        SSU  SU 3SS/[        R                  QS9  [	        USS9nUc  g	[        U[        R                  5      (       d#  [        S
SU  SU 3SS/[        R                  QS9  S[        US5      0n[        R                  " XP40 UD6R                  U[        R                  R                  U5      /0 5      $ )zDCheck for a __getattr__ and handle it specially if it is implementedz5Custom __getattribute__ in nn.Module attribute accessvar_getattr r   r   r   r   T)ignore_nn_module_getattrNz6torch.nn.Module with a non-function custom __getattr__zDynamo detected a nn.Module object with a custom `__getattr__` method, but this method is not a standard Python function (e.g., it might be implemented in C/C++). Dynamo cannot currently trace into such non-standard `__getattr__` methods.zAvoid using objects with non-standard __getattr__ methods within the compiled region. If possible, implement __getattr__ as a standard Python function.re   __getattr__)r$   r   r   r   r   rK   rw   FunctionTyper   r   UserMethodVariablecall_functionr4   r   )r   r   r5   rh   r   
getattr_fnr   s          rE   _custom_getattr_fallback)NNModuleVariable._custom_getattr_fallback-  s     #4((O&tfAdV45 H&22
 (tL
*e&8&899P&tfAdV4-A '22	$ Z
MBC++JHHVV++224892
 	
r   c                 	   U R                   =(       a    [        U R                   U5      nUR                  R                  U R                  5      n[
        R                  US5      nSn[        5       n[        R                  " UR                  5       H,  nUR                  UR                  R                  5       5        M.     U R                   (       d"  [        SSU  SU 3S/ [        R                   QS9  US:X  a  ["        R$                  " XUS9$ S n	X%;   a  XR   n	OcS	U;   a  X%S	   ;   a  X';  a	  US	   U   n	OGS
U;   a  X%S
   ;   a	  US
   U   n	O0SU;   a  X%S   ;   a	  US   U   n	O [        R&                  " XB5      n	SnUS:X  a  [3        U R                   U5        US:X  a&  U(       d  ["        R4                  " UR                  US9$ U(       af  [6        R8                  " X[;        U5      5      n[=        U[>        [@        45      (       a)  URC                  [        U RE                  5       U5      5        U$ [G        U	[H        5      (       am  U R                   (       a,  [        [        U R                   S5      U5      n[        US5      n["        RJ                  " U	RL                  US9RO                  X/0 5      $ [G        U	[P        5      (       a=  ["        RR                  " U	RT                  ["        RV                  " [/        U5      5      US9$ [G        U	[X        5      (       a$  ["        RJ                  " U	R[                  U5      US9$ [G        U	[\        R^                  5      (       a  ["        RR                  " XUS9$ [a        U	5      (       d  [c        U	5      (       a   [6        R8                  " X[;        U5      5      $ [        SS[e        U5       SU S[e        U	5       3S[e        U	5       S3SU S[e        U	5       S[e        U5       S3S/[        Rf                  QS9  ["        R$                  " XUS9$ ! [(         aO    U R+                  XAX R                   S9n
U
b  U
s $ [-        [(        US[/        U5      R0                   SU S3/S9   GNf = f)Nru   Tzgetattr with no sourcer   r   zDynamo does not know how to access an attribute on an `nn.Module` instance that lacks a source. This is usually an internal error in Dynamo.r   r   _modules_parameters_buffersF)r   r5   rh   r   '' object has no attribute 'rX   rr   rm   fgetz$Unsupported nn.Module attribute typeznn.Module subclass: z, name: z, attribute type: z>Dynamo does not support tracing nn.Module attributes of type ``zRefactor your code so that `z	` (type `z`) is not an attribute of `zqCurrently supported attribute types are methods, classmethods, staticmethods, properties, constants, and tensors.)4re   r   r   r   rd   r   r   rP   inspectgetmrorm   updateru   keysr   r   
DYNAMO_BUGr   GetAttrVariablegetattr_staticr[   r   r   rJ   __name__r   UserDefinedClassVariabler.   buildr   rK   r   UnspecializedNNModuleVariabler   r   r"   propertyUserFunctionVariabler   r   classmethodr   rs   r2   staticmethod__get__rw   r   r    r!   r,   r   )r   r5   rh   re   r   r   object_memberall_class_attribute_namesr<   subobjr   outs               rE   var_getattrNNModuleVariable.var_getattr_  s;   >DKK!>yy&&t7++D*=	$'E!/A%,,QZZ__->? 0 {{0&tfAdV47 6)445 :,,TGG_F)#*--5z*40Fi'Dm4L,L}-d3F9$:1F)Fz*40F //; % 924;;E;}55dnnVTT!''N64JKC# 02OPQQ ..t>>@$G J fh'';;'
4;;(LdSF'7F 55KK!  -Hb12 ,, 33OO77T
C! 
 -- 55NN4(!   2 233 33FPP!&))Xf-=-=&,,R9OPPB274=/$Oabijpbqars"`ahioap`qqr s6tfIgfoEVVqryz~r  rA  AB  C> +66		
 ((FCCQ " 664KK 7  %!M("d4j1122MdVSTUVs   &P+ +)R*RRr7   c                 <   UR                   R                  U R                  5      n[        U R                  U R	                  5       X5         [        U5      n[        U[        R                  R                  5      (       Ga"  UR                  R                  [        R                  R                  R                  L a  [        U5      (       a  U R                  U5        U(       a   S5       eU(       a  [        USS[        U5       S35        Uu  nUR                   R#                  5        Hj  u  pxUR%                  UR                   R'                  UU R                  U[)        [+        U R,                  U5      5      S9U/0 5        UR/                  5       nMl     UsS S S 5        $ U(       a*  UR0                  b  UR0                  U l        [5        XX#5        UR                   R7                  5       (       a  UR8                  R;                  S5      (       av  UR8                  S:w  af  [        USSS	9(       a  U R                  U5        S
SKJn	  U	" UUR                   R@                  " SU R                  /[C        X#5      Q76 S9sS S S 5        $ [        U[        RD                  RF                  5      (       a#  UR                  n
[+        U R,                  S5      nO"URH                  n
[+        U R,                  S5      n[K        U
[L        RN                  5      (       a(  U
RP                  n
[+        US5      nU /[S        U5      -   nO![K        U
[L        RT                  5      (       d   eURW                  [X        RZ                  " XS9UU5      sS S S 5        $ ! , (       d  f       g = f)Nz3Expected lazy sequential isn't a valid combination?ztorch.nn.Module.Sequentialz0 kwargs kwargsr   )z	torch.nn.z	torch.ao.ztorch.nn.utils.parametrizeT)check_forward_hookscheck_backward_hooksr*   wrap_fx_proxycall_moduler5   proxyrr   
_call_implrs   ).r   r   rd   rp   r   r   rK   rQ   r   r   rm   rr   r#   r   r&   lenr   rM   r   r   r   r   re   popcls_to_becomer   rc   is_root_tracer
__module__
startswithbuilderr  create_proxyr%   rR   GraphModuler  r"   rw   rx   rs   rN   r   inline_user_function_returnr   r   )r   r5   r7   r8   r6   is_lazyr_   
child_namer   r  fn	fn_sources               rE   r   NNModuleVariable.call_function  s    ii%%doo6#OOT<<>
 %S)G3 3 344MM))UXX-@-@-H-HH%c** 11"5 # I{ '4"v;-w/	 *-,,*<*<*>&J$$		99" OO&#1*T[[*2U#V	 :  	 &&(C +? O
 
R $$0'*'8'8D$
 'r= 		((**NN--.HIINN&BB%T 11"52$))00% +48O
 
` c588#7#788 B *4;;	 BIB *4;; EI"e..//B *9j AI 6DJ.D!"e&8&8999955222H
 
 
s   EN#CN	C:N
Nconstantc                   >^ ^^^^^^^ ^!^"^#^$^% SSK JmJmJm   T R                  m#TR
                  R                  T#5      m$S[        S[        4UUU$U U4S jjnTS;   a  T R                  TTT5      $ TS:X  a8  [        T R                  T R                  5       TT$5         U" T5      sS S S 5        $ TS:X  aU  [        R                  " [        R                  " T$R                   R"                  5      5      (       a  TR$                  " S	5      $ TS
:X  a  TS   R'                  5       (       d  [)        TST$ ST S35        [+        TS   T 5      (       d  [)        TST$ ST S35        TS   R,                  TS   R.                     n[+        U[0        5      (       a  U$ UR                  m#TR
                  R                  T#5      nTR
                  R3                  UT#T#[5        [7        T R8                  T#5      5      S9$ U(       a7  [;        T$T5      n	T$R                   R<                   ST S3m[?        TU	TTT5      $ SIUUU$UU 4S jjm!S[        S[@        [        [B        4   4UU!UU$U4S jjn
S[D        [F        [B        [B        4      SS4UU"U#U%U U4S jjnS[        S[B        SS4UU U"U#U U4S jjnS[H        S[        S[H        4S jm"TS:X  a  TR
                  RJ                  RM                  [O        T R8                  S5      5        T(       d  T(       a&  [Q        TTS [S        T5       S![S        T5       S"35        / m%T$RU                  5        H  u  mnT%RW                  U" TU5      5        M      T" T%[Y        5       S#9$ TS$:X  a  TR
                  RJ                  RM                  [O        T R8                  S%5      5        / nT$RZ                  " SJ0 U
" S&S'5      D6 H  u  mnURW                  U" TU5      5        M      T" U[Y        5       S#9$ TS(:X  a  TR
                  RJ                  RM                  [O        T R8                  S)5      5        / nT$R\                  " SJ0 U
" S&S'S*5      D6 H  u  mnURW                  U" TU5      5        M      T" U[Y        5       S#9$ TS+:X  a  TR
                  RJ                  RM                  [O        T R8                  S5      5        / nT$R^                  " SJ0 U
" S,S&S*5      D6 H  u  mnURW                  U" TU5      5        M      T" U[Y        5       S#9$ TS-:X  a  TR
                  RJ                  RM                  [O        T R8                  S5      5        T(       d  T(       a&  [Q        TTS [S        T5       S![S        T5       S"35        U" T$RU                  5       5      $ TS.:X  aO  TR
                  RJ                  RM                  [O        T R8                  S5      5        U" T$R_                  5       5      $ TS/:X  aW  TR
                  RJ                  RM                  [O        T R8                  S%5      5        U" T$RZ                  " SJ0 U
" S'5      D65      $ TS0:X  aW  TR
                  RJ                  RM                  [O        T R8                  S)5      5        U" T$R\                  " SJ0 U
" S'5      D65      $ TS1:X  ao  T(       d  T(       a&  [Q        TTS [S        T5       S![S        T5       S"35        / nT$ H$  nURW                  TR$                  " U5      5        M&     T" U[Y        5       S#9$ TS2:X  aJ  T(       d  T(       a&  [Q        TTS [S        T5       S![S        T5       S"35        U" T$R-                  5       5      $ TS:X  aw  T(       d  T(       a&  [Q        TTS [S        T5       S![S        T5       S"35        / nT$R-                  5        H  u  mnURW                  U" TU5      5        M      T" U[Y        5       S#9$ TS3:X  aO  T(       d  T(       a&  [Q        TTS [S        T5       S![S        T5       S"35        TR$                  " [S        T$5      5      $ TS4:X  a  T" T Ra                  T5      [Y        5       S#9$ TS5:X  a  [+        T$[b        Rd                  Rf                  [b        Rd                  Rh                  45      (       aO  T(       aH  TS   R'                  5       (       a0  TR$                  " TS   Rk                  5       T$Rl                  ;   5      $ TS6:X  Ga  T(       d  [S        T5      S:w  a&  [Q        TTS7[S        T5       S![S        T5       S"35        [b        Rd                  Rf                  Rn                  [b        Rd                  Rp                  Rn                  [b        Rd                  Rh                  Rn                  [b        Rd                  Rr                  Rn                  [b        Rd                  Rt                  Rn                  4n[w        T$5      Rn                  U;  a  TS   R'                  5       (       a,  [+        TS   Rk                  5       [        [x        45      (       d   [{        S8S9T  S:T S:T S:T 3S;T S<T$ S=3S>/S?9  [;        T$T5      R|                  n	[+        U	[~        R                  5      (       d   e[O        [O        T R8                  T5      S@5      nTR                  [        R                  " U	US9T /[        T5      -   T5      $ [+        TS   [        5      (       Ga;  TR
                  R                  (       Ga  / n[        [        [S        T$5      5      5      TS   Rk                  5          n[        T$TS   Rk                  5          5       HT  u  nnUU   m#[5        [7        T R8                  T#5      5      nURW                  TR
                  R3                  UT#US95        MV     T$TS   Rk                  5          nTR
                  R3                  UT  SA3[5        [7        T R8                  TS   Rk                  5       5      5      S9nU$ T R                  T5        SSBKJJKn  Sn[+        TS   U5      (       a  TS   R                  TR
                  5      nODTS   R'                  5       (       a  TS   Rk                  5       nO[{        SCS9T  S:T S:T S:T 3SD/ S?9  T$U   nTR
                  R3                  UT R                  U[5        [7        T R8                  U5      5      S9$ TSE:X  d  [+        T$[b        Rd                  R                  R                  R                  5      (       a  TSF:X  dC  [+        T$[b        Rd                  R                  R                  R                  5      (       ao  TSG:X  ai  [;        T$T5      R|                  n	[O        [O        T R8                  T5      S@5      nTR                  [        R                  " U	US9T /[        T5      -   T5      $ TT$R                   R                  ;   ai  [        T$R                   R                  T   5      (       aB  [        SH [        R                  " TTR                  5       5       5       5      (       a  U" T5      $ [        T&T G]a  TT[        T5      T5      $ ! , (       d  f       GN= f)KNr*   )r4   ListIteratorVariableTupleVariablerh   r9   c           
         > T	R                   R                  STR                  S0 5      n[        UR                  T5        [        TT5      u  p#SSKJn  U" T	T	R                   R                  SU U/UQ7US9S9$ )Nget_attrr@   r*   r  call_method)r7   r8   r  )r   r  rd   r'   rT   r%   r  r  )
rh   	mod_proxyr]   r^   r  r7   r8   moduler   r5   s
        rE   generic_call_method_helper@NNModuleVariable.call_method.<locals>.generic_call_method_helperK  s     		..	I innf5'8v'F$J. ii,,!#1j1'	 -  r   r  _wrapped_call_implrr   _check_input_dimT_get_item_by_idxz``nn.Module`` z's call method z# requires a constant index argumentr   z# requires a tuple as first argumentr   __resultc                     > [        S [        R                  " T TR                  5       5       5       5      (       d   [	        SST ST ST  ST 3ST ST S3/ S9  g g )	Nc              3   @   #    U  H  oR                  5       v   M     g 7fr?   )is_python_constantrB   r<   s     rE   rF   UNNModuleVariable.call_method.<locals>.assert_all_args_kwargs_const.<locals>.<genexpr>  s      0V1$$&&0V   z&non-const argument in nn.Module methodcall_method: r   (Dynamo does not support calling method `` of ``nn.Module`` z with non-constant arguments.r   )all	itertoolschainvaluesr   )r7   r8   r(  rh   r   s   rE   assert_all_args_kwargs_constBNNModuleVariable.call_method.<locals>.assert_all_args_kwargs_const  sy     09fmmo0V   D+D64&$qI!#f$7x?\!^r   namesc                    > T" 5         [        T	T
5      n[        R                  " U5      R                  " T Vs/ s H  o"R	                  5       PM     sn0 TR                  5        VVs0 s H  u  p4X4R	                  5       _M     snnD6nUR                  5         UR                  nU  Vs0 s H  o3XS   _M	     sn$ s  snf s  snnf s  snf r?   )r   r   	signaturebindas_python_constantrM   apply_defaults	arguments)r@  r  r<   rU   rV   
bound_argsr7   r>  r8   r(  rh   s         rE   
get_kwargs0NNModuleVariable.call_method.<locals>.get_kwargs  s    (*&B **2.333784a'')489?H1**,,HJ %%'#--J.34ez}$e44 9H 5s   B;(C *CrM   zvariables.ListIteratorVariablec                    > / nU  HJ  u  p#UR                  T	R                  R                  UTU[        T" TR                  U5      5      S95        ML     T" T[        5       S9$ )Nr   mutation_type)r   r   r   r   re   r-   )
rM   r   rh   r   r"  
gen_sourcer   named_childrenr   r5   s
       rE   wrap_values1NNModuleVariable.call_method.<locals>.wrap_values  sm     F %II55-jd.KL	 6  !& (.>.@ r   objzvariables.TupleVariablec                    > T" TR                   " U 5      TR                  R                  UTU [        T" TR                  U 5      5      S9/5      $ r   )r   r   r   r   re   )rh   rQ  r4   r#  rM  r   r   r5   s     rE   named_embed1NNModuleVariable.call_method.<locals>.named_embed  sV     $++D1II55-jd.KL	 6 
 
r   re   c                     UR                  S5      nUS   S:X  a  U $ [        U5      S:  a-  UR                  S5      n[        X5      n [        U5      S:  a  M-  U $ )N.r    )splitr  r  r   )re   rh   
name_splitr<   s       rE   rM  0NNModuleVariable.call_method.<locals>.gen_source  sY    CJ!}"j/A%NN1%#F. j/A% Mr   rN  r   z0 args and 0 kwargsz
 args and r  rK  named_parametersr   prefixrecursenamed_buffersr   remove_duplicatenamed_modulesmemochildrenmodules
parametersbuffersr   r=  __len____iter____contains____getitem__z1 args and 0 kwargsz6Invalid or non-const argument in nn.Module __getitem__r7  r   r8  r9  z+ with a non-constant or non-(str, int) key.z9Use constant arguments of type str or int for __getitem__r   rs   z.__getitem__(slice))SymNodeVariablez.Unsupported key type for nn.Module.__getitem__zEDynamo does not support getitem on `nn.Module` with non-constant key._get_abs_string_index_conv_forward_output_paddingc              3   @   #    U  H  oR                  5       v   M     g 7fr?   )	is_tensorr4  s     rE   rF   /NNModuleVariable.call_method.<locals>.<genexpr>  s     R+QaKKMM+Qr6  )r9   Nr@   )YrW  r4   r"  r#  rd   r   r   r\   r.   r   rp   r   r   is_torch_inline_allowedr   getfilerm   r-  r   r3  r+   rK   rM   r   r   r   r   r   re   r   r   r/   rL   r   r   rO   r
   guard_on_key_orderaddr   r&   r  rN  r   r-   r[  r^  r`  r   rQ   r   r   ParameterDictrD  r   ri  r   r   r   rJ   intr   rs   rw   r   r  r   r   rN   r1   exportranger   r   tensorrj  evaluate_exprrc  conv_ConvNd_ConvTransposeNdru   rv   r:  r;  r<  r=  r   r&  )'r   r5   rh   r7   r8   r   r)  mod_varr   r  rH  rO  rS  r[  paramr^  buffernamed_modules_listr   tmpitems_resultbuiltin_supportedsrcr   r   
new_modulenew_module_variablerj  	key_valuer  r4   r"  r#  r>  rM  r   r(  rN  rm   s'   `````                         @@@@@@@@rE   r&  NNModuleVariable.call_method>  sG    	LKoo((-	S 	_ 	 	2 77 %%b$77Y (!@!@!BB 2$7 
 %%+*M*MOOF,,==>+
 +
 $**400%%7--//$$VHOD6Ade d1g}55$$VHOD6Ade 1gmmDGMM2G'#@AA$$CYY,,S1F9944%mDKK&EF	 5   &B&&//0$w?D/BdFKK
	 
			5s 		5tCH~ 		5 		5	E#s(O,	-	 	"	c 	 	0I 	 		v 	S 	V 	 ##II((,,ZZ-PQv#)4ykCK=@	 57N & 5 5 7f%%k$&?@ !8'.>.@  ''II((,,Z]-ST68%66  Xy1 e !''D%(@A  ( 0@0B  _$II((,,ZZ-PQ35M & 4 4 !Xy2DE!f $$[v%>?! (EUEWXX_$II((,,ZZ-PQ8: & 4 4 !VX/AB!f #))+dF*CD! ("2B2D  ZII((,,ZZ-PQv#)4ykCK=@	 v44677YII((,,ZZ-PQv33566\!II((,,Z]-STv66OI9NOPPYII((,,ZZ-PQv33Lj6KLMMV^v#)4ykCK=@	 F.55c:; '>N>PQQXv#)4ykCK=@	 v||~..W_v#)4ykCK=@	 35L &f##Kf$=> !/'DTDVWWYv#)4ykCK=@	 $**3v;77Z'((,<L<N  N"6EHH$7$79O9O#PQQQ**,,#**Q**,?  ]"Ta#)4ykCK=@	 ##//##//&&22&&22##//! F|''/@@G..00"47#=#=#?#sLL! X"/vQtfAdV1VH M%##'&(;F8Cn%p X VT*33!"e&8&89999 DKK!>
K55222cBFT$Z'  $q'=11 99###F  c&k 23DG4N4N4PQD'0Q8R8R8T1U'VV"3i,]4;;-LMII== & #'* >  (W "(Q(B(B(D!EJ*,))*K*K"& 34-)$++tAw7Q7Q7ST  +L +' /. 11"5/I$q'?33 G11"))<	a++-- G668	L+D64&$qI!9 I&F9944%mDKK&KL	 5   ++6588#3#3#8#8#@#@AAO+ 6588#3#3#8#8#I#IJJ-- &//B":dkk4#@*MI11..r)Dd#  F$$---))224899R9??4+QRRR-d337&r4dVDDq
 s   z00
z?)rd   r   r   re   r   )r5   r3   r9   N)F)"r   r  __qualname____firstlineno__r.   _nonvar_fieldsrJ   r\   rQ   r   Moduler   r   r
   r   r   r   r   rN   r   r   boolr   r   r   r   r   r  r   rL   r   r&  __static_attributes____classcell__rm   s   @rE   r   r      sG    	
 
	'	'N22-029>2TW2	2F 
- -D - T  #  	
   
&= $BW <
9)
914
9	
9/5 /$ /* *A    PT  .0
hhoo0
 $0
 	0

 0
 
/	"0
dpD5 pDS pD_ pDdk#k 'k S/)*	k
 
kf GE#GE GE '	GE
 S/)*GE GE 
GE GEr   r   c            
         ^  \ rS rSrSSS1\R
                  kr S\R                  R                  S\	SS4U 4S	 jjr
S
\S\4S jrS\4S jrS\SS4S jr\\R"                  S\\   4S j5       5       rSSS\\   4U 4S jjrSSS\\   S\\\4   S\4S jrSSS\S\\   S\\\4   S\4
U 4S jjrSSS\S\S\\   4S jrSSS\S\4U 4S jjrSSS\S\4S jr Sr!U =r"$ )r   i  
value_typeis_state_mutatedr   r   r8   r9   Nc                 h  > [        U5      [        R                  R                  R                  L a%  [        S[        U5      S/ [        R                  QS9  SU;   a,  [        US   SS 5      n[        U5      UL a  [        U5      US'   [        TU ],  " SSU0UD6  SU l        U R                  U l        g )	NzFUnspecializedNNModuleVariable wrapped around ScriptModules unsupportedzScriptModules aren't supported in UnspecializedNNModuleVariable because their .forward function isn't a static member of their type.r   r  r  r   Fr@   )rJ   rQ   jit_scriptRecursiveScriptModuler   r\   r   	DIFFICULTr   r   r   r  re   r   )r   r   r8   lazy_value_to_becomerm   s       rE   r   &UnspecializedNNModuleVariable.__init__  s    ;%))++AAA`E
X&00 6!#*6,+?RV#W E{22
 (,E{|$/u// % '+kk#r   attr_sourcec                     U$ r?   r@   r   r  s     rE   _wrap_source*UnspecializedNNModuleVariable._wrap_source      r   c                 V    U R                   =(       d    U R                  nU(       d   eU$ r?   r   r   s     rE   r   8UnspecializedNNModuleVariable.get_nn_module_stack_source  r   r   re   c                     Xl         g r?   r   r   s     rE   r   8UnspecializedNNModuleVariable.set_nn_module_stack_source  r   r   c                     [         R                  R                  R                  [         R                  R                  R                  [         R                  R                  R
                  1n [         R                  R                  R                  R                  5        Vs1 s H2  n[        US5      (       d  M  X;  d  M  [        UR                  5      iM4     sn$ s  snf )N__code__)rQ   r   r  __setattr__r   __delattr__ru   r=  rY   idr  )	supportedr<   s     rE   _nn_module_method_ids3UnspecializedNNModuleVariable._nn_module_method_ids  s    
 HHOO''HHOO$$HHOO''
	 XX__--446
6q*% *+*< BqzzN6
 	
 
s    C9C Cr5   r3   c                   >  [         R                  " U R                  S5      nU[
        R                  R                  R                  [
        R                  R                  R                  [
        R                  R                  R                  4;   a7  UR                  [        R                  " X5      U /0 5      R                  U5      $ [        TU ]9  U5      $ ! [         a  n[        UeS nAff = f)Nrg  )r   r   r  r[   NotImplementedErrorrQ   r   r   rg  r   r   r  r.   r   r   r   )r   r5   r  rb   rm   s       rE   r   1UnspecializedNNModuleVariable.unpack_var_sequence  s    	-''DB HH((HH""++HH((
 
 11%%b-  "!"%& w*2..'  	-%1,	-s   !C 
C/#C**C/r7   c                 	   U R                   n[        U5      (       a*  UR                  b  UR                  U l        [	        XX#5        [        U[        R                  R                  5      (       d  Sn[        U R                  U5      nOSn[        U R                  U5      n[        UR                  [        R                  5      (       Ga  [        UR                  [        R                  5      (       Ga  UR                  R                  [         L Ga  UR                  R                  ["        L Gae  SUR$                  ;  GaT  [&        R(                  " US5      n[        U[        R*                  5      (       Ga  UR,                  nU R/                  US5      R1                  5       R3                  5       (       Gd  U R/                  US5      R1                  5       R3                  5       (       Gd  U R/                  US5      R1                  5       R3                  5       (       Gdu  U R/                  US5      R1                  5       R3                  5       (       GdA  UR/                  US5      R3                  5       (       Gd  UR/                  US	5      R3                  5       (       d  UR/                  US
5      R3                  5       (       d  UR/                  US5      R3                  5       (       d  UR/                  US5      R3                  5       (       d  UR/                  US	5      R3                  5       (       db  UR/                  US
5      R3                  5       (       d=  UR/                  US5      R3                  5       (       d  SnU R                  R4                  nU R6                  (       a  U R9                  U5      n	OS n	[;        U R6                  U5        U R6                  (       a.  [=        [?        [A        U5      5      U RC                  5       UU5      O	[E        5       n
U
   [        U[        R*                  [        RF                  RH                  45      (       d=  [J        RL                  " XU	S9nURO                  X/[Q        U5      -   U5      sS S S 5        $ [R        RT                  " XiS9RO                  X/[Q        U5      -   U5      sS S S 5        $ ! , (       d  f       g = f)N__call__r  rr   _backward_hooks_backward_pre_hooks_forward_hooks_forward_pre_hooks_global_backward_pre_hooks_global_backward_hooks_global_forward_hooks_global_forward_pre_hooksr   )+r   r   r  r  rc   rK   rQ   rR   r  r   r"   r  rw   rx   r  rs   r(   r)   ru   r   r   r   nn_modules_globals_vtr  realizer  rr   re   get_source_by_walking_mror   rp   r\   r  r   r   r  ScriptFunctionr.   r   r   rN   r   r   )r   r5   r7   r8   r6   rh   r  forward_method
globals_vtre   ctxfn_vts               rE   r   +UnspecializedNNModuleVariable.call_function)  s    jj#  ,"%"3"3"2D9#uxx3344D$/BD$/B
 3<<!1!122s~~u'7'788%%)AA''+HH-$33CCN.%*<*<==55
$$R):;CCEIIKK'',ABJJLPPRR'',<=EEGKKMM'',@AIIKOOQQ!--b2NOSSUU!--b2JKOOQQ!--b2IJNNPP!--b2MNRRTT!--b2NOSSUU!--b2JKOOQQ!--b2IJNNPP!--b2MNRRTT$D00B;;33D9FF.t{{C@ {{ #BsG//1	  	 b5#5#5uyy7O7O"PQQ'--bVD**2vT
/BFK S !55bHVVd+V SSs   A-S0S
Srh   c                    > US;   ao  [        U R                  U5      nU R                  (       a  U R                  U5      nOS n[        R
                  " XUS9nUR                  X/[        U5      -   U5      $ U[        U R                  S0 5      ;  Ga   [        R                  " [        U R                  5      U5      n[        U[        5      (       aL  [        U R                  U5      S5      n[        R
                  " XR                   US9nUR                  XU5      $ [#        US5      (       a`  [%        UR&                  5      U R)                  5       ;   a9  [+        SSU  SU SU SU 3S	U S
U R                   3S/[,        R.                  QS9  US:X  Ga  SU R                  R0                  ;   a  US   R3                  5       n	US   n
U
R5                  5       (       a+  U
R7                  5       [8        R:                  R<                  L d  XR                  R0                  S   ;   a  SU l        OvXR                  R0                  S   ;   a  SU l        OR[        U
[@        RB                  [@        RD                  45      (       d  XR                  R0                  S   ;   a  SU l        U[8        R:                  RF                  RH                  L a"  [        US   [@        RJ                  5      (       d'  U[8        R:                  RF                  RL                  L aP  [        R
                  " U[8        R:                  RF                  RL                  5      nUR                  XUS   /U5      $ [N        TU ]  X[        U5      U5      $ ! [         a    S n GNf = f)Nr+  r   ru   rs   r  z,UnspecializedNNModuleVariable missing methodr7  r   z'Dynamo does not support tracing method z of nn.Module z@Dynamo does not really define unspecialized nn.Module very well.r   r  r   r   r*   Tr   r   ))r   r  re   r  r.   r   r   rN   r   r   r   rJ   r[   rK   r   r   rs   rY   r  r  r  r   r   r  ru   rD  ro  r   rQ   r   	Parameterr  r   r   r   r  r  r   r  r   r&  )r   r5   rh   r7   r8   r  re   r  method	attr_namer   rm   s              rE   r&  )UnspecializedNNModuleVariable.call_methodx  s     77$/B{{77=#))"@E&&r6DJ+>GGwtzz:r:: //TZZ0@$G &,//#D$B$B4$H*U'--b//&Q**2V<< 
++v'4+E+E+GGJ+D64&$qI"I$~^b^h^h]i jZ*44	 }$$**:M:M)M !G668	Q
 OO%%%*;*;*=ASAS*S**"5"5m"DD,0D)**"5"5j"AA,0D)%66%CC  !JJ$7$7
$CC -1D) %((//555tAw	(A(ABB588??666 (--b%((//2M2MN**2d1gGGw"2T$Z@@I " s   *M= =NNfieldname_vtc                     U R                  X5      n[        U[        R                  5      (       a  UR	                  U5      $ g r?   )r  rK   r   ConstDictVariablemaybe_getitem_const)r   r5   r  r  dict_vts        rE   getattr_helper,UnspecializedNNModuleVariable.getattr_helper  s;     ""2-gy::;;..w77r   c           
      ,  >^^ US;   a  UR                   R                  R                  X5      (       d  [        U R                  U5      n[        U[        5      (       at  [        U5      S:X  ae  U R                  (       a>  [        U R                  U5      n[        UR                  [        R                  5      5        [        R                  " 0 5      $ U R                  (       Ga"  US;   Ga  UR                   R                  R                  X5      (       d  [        U R                  U5      n[        U R                  U5      m[        TR                  [        R                   5      5        UR                   R"                  R%                  T5        S[&        S[(        S[(        S[*        [,        [,        4   4U4S jjm[        U4S	 j[/        UR1                  5       5       5       5      n[        R2                  " U[5        U5      TS
9$ [6        TU ]q  X5      $ )N)r  r  r  r  r   )r  r  irU   rV   r9   c                    > [         R                  R                  U5      n[        TU 5      n[	        TU5      n[
        R                  " X%5      nX64$ r?   )r   r4   r   r   r   r0   )r  rU   rV   r   
source_keysource_valuer   hooks_dict_sources          rE   build_key_valueBUnspecializedNNModuleVariable.var_getattr.<locals>.build_key_value  sO      0077: 00A1E
01BJO+221Cz!r   c              3   @   >#    U  H  u  nu  p#T" XU5      v   M     g 7fr?   r@   )rB   r  rU   rV   r  s       rE   rF   <UnspecializedNNModuleVariable.var_getattr.<locals>.<genexpr>  s%      :WYQa((:Ws   r   )r   r   r   r   r   rK   rL   r  re   r   r   ry   r   EMPTY_NN_MODULE_HOOKS_DICTr   r  SEQUENCE_LENGTHrs  rt  rv  r   rO   r.   r   rM   NNModuleHooksDictVariablerJ   r   r  )	r   r5   rh   
hooks_dicthooks_sourcer   r  r  rm   s	         @@rE   r  )UnspecializedNNModuleVariable.var_getattr  s    
 
 99))FFtRR$TZZ6
j$//C
Oq4H{{'1$++t'D%(33 , G G
 %66r:: KKK
 II**GGSS T2J *4;; =+66|7S7STUII((,,->?"""#&"78"  :CJDTDTDV:W F 66Z(1B  w"2,,r   c           	      (   [         R                  " U5      nU R                  USU5      nUc  U R                  USU5      nUc  U R                  USU5      nUc4  [        [        US[        U R                  5      R                   SU S3/S9  Uc   eU$ )z
Dynamo tracing of nn.Module __getattr__ can be expensive if the model
has deep submodule hierarchy. Since the __getattr__ is stable, we can
directly look into the underlying datastructures. This saves a lot of
compilation time.
r   r   r   r   r   rX   )r   r4   r  r   r[   rJ   r   r   )r   r5   rh   r  r  s        rE    manually_trace_nn_module_getattr>UnspecializedNNModuleVariable.manually_trace_nn_module_getattr  s     ,,T2!!"mW=;%%b*g>C;%%b*g>C;$TZZ(1122MdVSTU 
r   )r  r   r  )#r   r  r  r  r2   r  rQ   r   r  r   r   r
   r  r   r   r   r{   cacherP   rv  r  rN   r.   r   r   rL   r\   r   r&  r   r  r  r  r  r  r  s   @rE   r   r     s     
#	1	1	N2ehhoo 2 2 2< 6 F 
- -D - __
3s8 
  
/&= /$BW /0M#M 'M S/)*	M
 
M^XA#XA XA '	XA
 S/)*XA 
XAt)25@O	/	"=-5 =-S =-_ =-~)14	 r   r   c                   *    \ rS rSrSrS\S\4S jrSrg)$UnspecializedBuiltinNNModuleVariablei4  z_
Differentiates between builtin nn modules (e.g. torch.nn.Linear) and user defined nn modules.
r  r9   c                     U$ r?   r@   r  s     rE   r  1UnspecializedBuiltinNNModuleVariable._wrap_source9  r  r   r@   N)r   r  r  r  __doc__r
   r  r  r@   r   rE   r  r  4  s     6 r   r  c                   v   ^  \ rS rSrSrS\R                  R                  S\SS4U 4S jjr	S\S\4S	 jr
S
rU =r$ )FSDPManagedNNModuleVariablei>  a  
Tracing behavior: trace into submodules and treat them as Unspecialized, do not
register parameters to the top-level, treat them as function inputs.

Guards behavior: if 'skip_fsdp_guards', many guards that would be installed
by a vanilla UnspecializedNNModuleVariable are simply dropped, on the basis
that a user wrapping their model in FSDP(model) is already opting into a
requirement to not modify internal model state, which would already break FSDP without
compilation.
r   r8   r9   Nc                 l   > UR                  S5      nUc   S5       e[        TU ]  " SSU0UD6  X0l        g )Nre   zMFSDPManagedNNModule depends on having an accurate source to control guarding.r   r@   )rl   r   r   re   )r   r   r8   re   rm   s       rE   r   $FSDPManagedNNModuleVariable.__init__J  sC    H%! 	
[	
! 	/u//r   r  c                     [        U[        [        45      (       d?  [        R                  R
                  R                  (       a  [        U5      $ [        U5      $ U$ r?   )rK   r   r   rQ   _dynamoconfigskip_fsdp_guardsr  s     rE   r  (FSDPManagedNNModuleVariable._wrap_sourceS  sL    ,.IJ
 
 }}##44)+662;??r   r   )r   r  r  r  r  rQ   r   r  r   r   r  r  r  r  s   @rE   r  r  >  sB    	ehhoo      r   r  )Tr  r{   r   r;  ri   rw   collections.abcr   r   
contextlibr   r   typingr   r   r	   torch.nnrQ   torch._guardsr
   rW  r   r   r   excr   r   r   guardsr   r   mutation_guardr   re   r   r   r   r   r   r   r   utilsr   r   r   r   r    r!   r"   r#   r$   r%   r&   r'   r(   r)   r   r+   r,   r-   r.   	functionsr/   lazyr0   listsr1   user_definedr2   torch._dynamo.symbolic_convertr3   r   r4   r   r  rL   r\   rc   rp   r   r   r   r  r  r@   r   rE   <module>r     sw  .    	  . 2 / /    8 8 V V 0 .       S R 3 %   3 D***	* ?
#* o%&	*
 
*Z ++#+)@+GLxx++ +( V #(88?? 	 FRE REji$= iX+H "? r   