
    7h                        S 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	  SSK
rSSKJrJrJr  SSKJr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(J)r)J*r*J+r+J,r,J-r-J.r.J/r/  SSK0J1r1J2r2J3r3  SSK4J5r5  SSK6J7r7  SSK8J9r9  SSK:J;r;  \	(       a  SSK<J=r=  S S jr>\S\?S\R                  R                  4S j5       rBS rC " S S\35      rD " S S\;5      rE " S S\E5      rF " S S\E5      rGg)!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)contextmanagernullcontext)TYPE_CHECKING   )graph_break_hintstrace_rules	variables)raise_observed_exceptionunimplemented_v2UnspecializeRestartAnalysisUnsupported)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set_example_valueunpatched_nn_module_callunpatched_nn_module_call_impl   )typestrValueMutationNewVariableTracker)invoke_and_store_as_constant)LazyVariableTracker)SliceVariable)UserDefinedObjectVariable)InstructionTranslatorc                   ^ ^ [        US5      (       aj  UU 4S 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 )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_hookc                   > [        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     [/var/www/fran/franai/venv/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r2   r3   r4   s     r8   r9   r:   _   s     Ct44r;   )r   type
isinstancedictitemslisttuplesettorchfxProxyr   node)xkvr7   txs      r8   r7   /initialize_lazy_module.<locals>.convert_to_fakeY   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'N)hasattrr!   r@   _infer_parameters)rK   modargskwargs
proxy_argsproxy_kwargsarg	fake_argsrI   rJ   fake_kwargsr7   s   `          @r8   initialize_lazy_modulerW   O   s     s&''
	 $5T#B 
5?@Zc_S)Z	@9E9K9K9MN9Mq/!,,9MNck:# ( ANs   A?B
module_keyrO   c              #   "  #    UR                  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)Nr   @r%   )name	num_callsget	__class__nn_module_stack)rX   sourcerK   rO   fully_qualified_namer\   s         r8   record_nn_module_stackrb   k   s     !;;=  !5q9I09AJ<q,:J+*>)N:&-6])*z*Bz*s   <B0A= /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   r=   types
MethodTyper   
make_guardr   CLOSURE_MATCH	functoolspartialNOT_PRESENT_IN_GENERIC_DICT)r`   rO   fwdforward_sources       r8   &guard_to_detect_forward_monkeypatchingrr   x   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4U 4S jjrS	 rS
 rS r  SS jrS rSSS\SS4S jrS rS rSS jrS rSS jr      SS jr S     SU 4S jjjrSrU =r$ )NNModuleVariable   module_typerX   valuenn_module_stack_sourcereturnNc                    > [         TU ]  " S0 UD6  Xl        X l        X0l        U R
                  (       d   eU R
                  U l        g )Nr3   )super__init__rw   rX   rx   r`   ry   )selfrw   rX   rx   rQ   r^   s        r8   r}   NNModuleVariable.__init__   s=     	"6"&$
{{{&*kk#rs   c                 @    U R                   =(       d    U R                  $ r2   ry   r`   r~   s    r8   get_nn_module_stack_source+NNModuleVariable.get_nn_module_stack_source       **9dkk9rs   c                     Xl         g r2   ry   r~   r`   s     r8   set_nn_module_stack_source+NNModuleVariable.set_nn_module_stack_source       &,#rs   c                     U R                   $ r2   )rw   r   s    r8   python_typeNNModuleVariable.python_type   s    rs   rK   r-   c                     g r2   r3   )r~   rK   r`   submod	key_extraoptionss         r8   _wrap_submodule NNModuleVariable._wrap_submodule   s     	rs   c                    UR                   R                  U R                  5      n[        U[        R
                  R                  5      (       a  / n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 R                  (       d   e/ n[)        U5       HW  u  puUR                  UR                   R                  UU R                  U[        [        U R                  U5      5      S95        MY     U$ Nr`   )outputget_submodulerX   r>   rD   nn
ModuleDictr@   r	   ConstantVariablecreateregister_attr_or_moduler   r   r`   append
ModuleListParameterList
Sequentialr&   	enumerate)r~   rK   baseresultr[   r   name_varidxs           r8   unpack_var_sequence$NNModuleVariable.unpack_var_sequence   sd   yy&&t7dEHH//00F $

$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  + rs   r[   r(   c                 *   UR                   R                  U R                  5      n[        X25      n[	        [        [        U R                  U5      5      R                  [        R                  5      5        [        R                  R                  U5      $ r2   )r   r   rX   rM   r   r   r   r`   rk   r   HASATTRr	   r   r   )r~   rK   r[   rO   r   s        r8   call_obj_hasattr!NNModuleVariable.call_obj_hasattr   sn     ii%%doo6#:dkk489DD$$	

 ))0088rs   c                 f    UR                   R                  U R                  5      n[        USS5      $ )NtrainingF)r   r   rX   getattrr~   rK   rO   s      r8   is_trainingNNModuleVariable.is_training   s)    ii%%doo6sJ..rs   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   rX   r   tagf_codeco_namemark_class_dynamicr=   r   r   s      r8   convert_to_unspecialized)NNModuleVariable.convert_to_unspecialized   sT    ii%%doo6c" 99
*00c;))rs   c                    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_okrg   )r   r   rX   r    r   r   SUPPORTABLEside_effectshas_pending_mutation_of_attr	load_attrr>   r	   DeletedVariableobject__getattribute__)r~   rK   keyr   mutated_attr	base_dicts         r8   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*=	rs   c                    [        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.r`   __getattr__)r    r   r   r   r   r>   ri   FunctionTyper   r	   UserMethodVariablecall_functionr   r   )r~   r   rK   r[   
obj_source
getattr_fnr   s          r8   _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
 	
rs   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$ 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  [/        U R                   U5        US:X  a&  U(       d  ["        R0                  " UR                  US9$ U(       ag  [2        R4                  " UW	[7        U5      5      n[9        U[:        [<        45      (       a)  UR?                  [        U RA                  5       U5      5        U$ [C        W	[D        5      (       am  U R                   (       a,  [        [        U R                   S5      U5      n[        US5      n["        RF                  " U	RH                  US9RK                  X/0 5      $ [C        U	[L        5      (       a=  ["        RN                  " U	RP                  ["        RR                  " [U        U5      5      US9$ [C        U	[V        5      (       a$  ["        RF                  " U	RY                  U5      US9$ [C        U	[Z        R\                  5      (       a  ["        RN                  " XUS9$ [_        U	5      (       d  [a        U	5      (       a   [2        R4                  " X[7        U5      5      $ [        SS[c        U5       SU S[c        U	5       3S[c        U	5       S3SU S[c        U	5       S[c        U5       S3S/[        Rd                  QS9  ["        R$                  " XUS9$ ! [(         a5    U R+                  XAX R                   S9n
U
b  U
s $ [-        [(        U5         GNf = f)Nrg   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   rK   r[   r   rd   r^   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.)3r`   r   r   r   rX   r   r   rC   inspectgetmror^   updaterg   keysr   r   
DYNAMO_BUGr	   GetAttrVariablegetattr_staticAttributeErrorr   r
   rr   UserDefinedClassVariabler(   buildr   r>   ru   UnspecializedNNModuleVariabler   r   r   propertyUserFunctionVariabler   r   classmethodr   re   r,   r=   staticmethod__get__ri   r   r   r   r&   r   )r~   rK   r[   r`   r   r   object_memberall_class_attribute_namesrH   subobjr   outs               r8   var_getattrNNModuleVariable.var_getattr6  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!''FN64J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CM " 664KK 7  %!M("s   $P* *)Q)Q)(Q)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   e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R,                  b  UR,                  U l        [1        XX#5        UR                   R3                  5       (       a  UR4                  R7                  S5      (       av  UR4                  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                  /[?        X#5      Q76 S
9sS S S 5        $ U R(                  (       d   S5       e[        U[        R@                  RB                  5      (       a#  UR                  n
['        U R(                  S5      nO"URD                  n
['        U R(                  S5      n[G        U
[H        RJ                  5      (       a  U
RL                  n
['        US5      nU /U-   nO![G        U
[H        RN                  5      (       d   eURQ                  [R        RT                  " XS9UU5      sS S S 5        $ ! , (       d  f       g = f)Nz3Expected lazy sequential isn't a valid combination?r   )z	torch.nn.z	torch.ao.ztorch.nn.utils.parametrizeT)check_forward_hookscheck_backward_hooksr%   wrap_fx_proxycall_modulerK   proxyzsMust provide a valid source in order to inline, since inlined function may have default args which must be guarded.rd   
_call_implre   )+r   r   rX   rb   r   r   r>   rD   r   r   r^   rd   r   r   r   r@   r   r   r   r   r`   popcls_to_becomerw   rW   is_root_tracer
__module__
startswithbuilderr   create_proxyr!   rE   GraphModuler   r   ri   rj   re   r   inline_user_function_returnr	   r   )r~   rK   rP   rQ   rO   is_lazyrT   
child_namer   r   fn	fn_sources               r8   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{ "!z*-,,*<*<*>&J$$		99" OO&#1*T[[*2U#V	 :  	 &&(C +? C
 
F $$0'*'8'8D$
 'r= 		((**NN--.HIINN&BB%T 11"52$))00% +48C
 
T {{ Z{ c588#7#788 B *4;;	 BIB *4;; EI"e..//B *9j AI 6D=D!"e&8&8999955222H{
 
 
s   D6NCN1D	N
Nc                 6  >^ ^^^^^^^^^^^ SSK JmJmJm  T R                  mTR
                  R                  T5      mUUUU U4S 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  aT  [        R                  " [        R                  " TR                  R                  5      5      (       a  TR!                  S5      $ TS:X  a  TS   R#                  5       (       d   e[%        TS	   T5      (       d   eTS	   R&                  TS   R(                     n[%        U[*        5      (       a  U$ UR                  mTR
                  R                  T5      nTR
                  R-                  UTT[/        [1        T R2                  T5      5      S
9$ U(       a7  [5        TT5      n	TR                  R6                   ST S3m[9        TU	TTT5      $ UUUUU 4S jmUUUUU4S jn
UUUU U4S jnUUUUU U4S jnS mTS:X  a  TR
                  R:                  R=                  [?        T R2                  S5      5        T(       d  T(       a   e/ nTRA                  5        H  u  mnURC                  U" TU5      5        M      T" U[E        5       S9$ TS:X  a  TR
                  R:                  R=                  [?        T R2                  S5      5        / nTRF                  " S40 U
" SS5      D6 H  u  mnURC                  U" TU5      5        M      T" U[E        5       S9$ TS:X  a  TR
                  R:                  R=                  [?        T R2                  S5      5        / nTRH                  " S40 U
" SSS5      D6 H  u  mnURC                  U" TU5      5        M      T" U[E        5       S9$ TS:X  a  TR
                  R:                  R=                  [?        T R2                  S5      5        / nTRJ                  " S40 U
" SSS5      D6 H  u  mnURC                  U" TU5      5        M      T" U[E        5       S9$ TS:X  a_  TR
                  R:                  R=                  [?        T R2                  S5      5        T(       d  T(       a   eU" TRA                  5       5      $ TS:X  aO  TR
                  R:                  R=                  [?        T R2                  S5      5        U" TRK                  5       5      $ TS :X  aW  TR
                  R:                  R=                  [?        T R2                  S5      5        U" TRF                  " S40 U
" S5      D65      $ TS!:X  aW  TR
                  R:                  R=                  [?        T R2                  S5      5        U" TRH                  " S40 U
" S5      D65      $ TS":X  aX  T(       d  T(       a   e/ nTRM                  5        H#  mURC                  TR!                  T5      5        M%     T" U[E        5       S9$ TS#:X  a&  T(       d  T(       a   eU" TR'                  5       5      $ TS$:X  aS  T(       d  T(       a   e/ nTR'                  5        H  u  mnURC                  U" TU5      5        M      T" U[E        5       S9$ TS%:X  a*  T(       d  T(       a   eTR!                  [O        T5      5      $ TS&:X  a  [%        T[P        RR                  RT                  [P        RR                  RV                  45      (       aN  T(       aG  TS	   R#                  5       (       a/  TR!                  TS	   RY                  5       TRZ                  ;   5      $ TS':X  Ga  T(       d  [O        T5      S:X  d   e[P        RR                  RT                  R\                  [P        RR                  R^                  R\                  [P        RR                  RV                  R\                  [P        RR                  R`                  R\                  [P        RR                  Rb                  R\                  4n[e        T5      R\                  U;  a  [%        TS	   [f        R                  5      (       d   [i        TS	   5      5       eTS	   RY                  5       m[%        T[j        [l        45      (       d   e[5        TT5      Rn                  n	[%        U	[p        Rr                  5      (       d   e[?        [?        T R2                  T5      S(5      nTRu                  [f        Rv                  " U	US
9T /[y        T5      -   T5      $ T R2                  (       d   e[%        TS	   [z        5      (       Ga;  TR
                  R|                  (       Ga  / n[y        [        [O        T5      5      5      TS	   RY                  5          n[        TTS	   RY                  5          5       HT  u  nnUU   m[/        [1        T R2                  T5      5      nURC                  TR
                  R-                  UTUS
95        MV     TTS	   RY                  5          nTR
                  R-                  UT  S)3[/        [1        T R2                  TS	   RY                  5       5      5      S
9nU$ T R                  T5        SS*KBJCn  [%        TS	   U5      (       a  TS	   R                  TR
                  5      mODTS	   R#                  5       (       a  TS	   RY                  5       mO[        S+S,T  S-T S-T S-T 3S./ S/9  TT   nTR
                  R-                  UT R                  T[/        [1        T R2                  T5      5      S
9$ TS0:X  d  [%        T[P        RR                  R                  R                  R                  5      (       a  TS1:X  dC  [%        T[P        RR                  R                  R                  R                  5      (       af  TS2:X  a`  [5        TT5      Rn                  n	[?        [?        T R2                  T5      S(5      nTRu                  [f        Rv                  " U	US
9T /T-   T5      $ TTR                  R                  ;   ai  [        TR                  R                  T   5      (       aB  [        S3 [        R                  " TTR                  5       5       5       5      (       a  U" T5      $ [        TT G]E  TTTT5      $ ! , (       d  f       GNx= f)5Nr%   )r   ListIteratorVariableTupleVariablec           
         > 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_attrr3   r%   r   call_method)rP   rQ   r   )r   r  rX   r"   rG   r!   r  r   )
r[   	mod_proxyrR   rS   r   rP   rQ   moduler~   rK   s
        r8   generic_call_method_helper@NNModuleVariable.call_method.<locals>.generic_call_method_helper  s     		..	I innf5'8v'F$J. ii,,!#1j1'	 -  rs   r   _wrapped_call_implrd   _check_input_dimT_get_item_by_idxr   r   __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r2   )is_python_constantr5   rH   s     r8   r9   UNNModuleVariable.call_method.<locals>.assert_all_args_kwargs_const.<locals>.<genexpr>^  s      0V1$$&&0Vs   z&non-const argument in nn.Module methodcall_method: r   z(Dynamo does not support calling method `z` of ``nn.Module`` z with non-constant arguments.r   )all	itertoolschainvaluesr   )rP   rQ   r  r[   r~   s   r8   assert_all_args_kwargs_constBNNModuleVariable.call_method.<locals>.assert_all_args_kwargs_const]  sy     09fmmo0V   !D+D64&$qI!#f$7x?\!^rs   c                    > 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 r2   )r   r   	signaturebindas_python_constantr@   apply_defaults	arguments)namesr  rH   rI   rJ   
bound_argsrP   r)  rQ   r  r[   s         r8   
get_kwargs0NNModuleVariable.call_method.<locals>.get_kwargsi  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c                    > / nU  HJ  u  p#UR                  TR                  R                  UTU[        T" TR                  U5      5      S95        ML     T" U[        5       S9$ )Nr   mutation_type)r   r   r   r   r`   r'   )	r@   r   r[   r   r  
gen_sourcer   r~   rK   s	       r8   wrap_values1NNModuleVariable.call_method.<locals>.wrap_valuest  sf    F %II55-jd.KL	 6  !& (>N>PQQrs   c                    > 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   r`   )r[   objr   r  r8  r   r~   rK   s     r8   named_embed1NNModuleVariable.call_method.<locals>.named_embed  sV     $++D1II55-jd.KL	 6 
 
rs   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    )splitlenr  r   )r`   r[   
name_splitrH   s       r8   r8  0NNModuleVariable.call_method.<locals>.gen_source  sY    CJ!}"j/A%NN1%#F. j/A% Mrs   named_childrenr   r6  named_parametersr   prefixrecursenamed_buffersr   remove_duplicatenamed_modulesmemochildrenmodules
parametersbuffersr   r(  r@   __len____contains____getitem__re   z.__getitem__(slice))SymNodeVariablez.Unsupported key type for nn.Module.__getitem__r$  r   zEDynamo does not support getitem on `nn.Module` with non-constant key.r   _get_abs_string_index_conv_forward_output_paddingc              3   V   #    U  H  n[        U[        R                  5      v   M!     g 7fr2   )r>   r	   TensorVariabler"  s     r8   r9   /NNModuleVariable.call_method.<locals>.<genexpr>I  s'      ?A 1i6677?s   ')r3   )RrA  r   r  r  rX   r   r   r   rb   r   r   is_torch_inline_allowedr   getfiler^   r  r   r!  r>   r@   rx   r   r   r   r   r`   r   __name__r)   guard_on_key_orderaddr   rF  r   r'   rG  rJ  rL  r   rC  rD   r   r   ParameterDictr.  r   rT  r   r   r   r=   r	   r&   strintre   ri   r   r	  r   rA   r+   exportranger   r   tensorrU  evaluate_exprr   rO  conv_ConvNd_ConvTransposeNdrg   rh   r%  r&  r'  r(  r|   r  ) r~   rK   r[   rP   rQ   constantr  mod_varr   r  r3  r9  r=  r   parambufferbuiltin_supportedsrcr   r   
new_modulenew_module_variablerU  r  r   r  r  r)  r8  r   r  r^   s    `````                   @@@@@@@r8   r  NNModuleVariable.call_method  s@    	LKoo((-	 	2 77 %%b$77Y (!@!@!BB 2$7 
 %%+*M*MOOF,,==>+
 +
 $**400%%7--////d1g}55551gmmDGMM2G'#@AA$$CYY,,S1F9944%mDKK&EF	 5   &B&&//0$w?D/BdFKK
	 
			5 		5	R 	R	 		 ##II((,,ZZ-PQ''F & 5 5 7fk$78 !8'>N>PQQ''II((,,Z]-STF%66  Xy1 e k$67  (>N>PQQ_$II((,,ZZ-PQF & 4 4 !Xy2DE!f k$78! (>N>PQQ_$II((,,ZZ-PQF & 4 4 !VX/AB!f k$78! (>N>PQQZII((,,ZZ-PQ''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^''F.55d;< &'>N>PQQX''v||~..W_''F &fk$78 !/'>N>PQQY''#**3v;77N"6EHH$7$79O9O#PQQQ**,,#**Q**,?  ]"#d)q.00##//##//&&22&&22##//! F|''/@@!$q'9+E+EFFXPTUVPWHXXF1g002!#Sz2222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/$q'?331g++BII6a++--1g002 L+D64&$qI!9 C[F9944%mDKK&EF	 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  F$$---))224899 "v}}?  
 .d337&r4v>>c s   t		
t)rX   rw   ry   rx   rK   r-   rP   list[VariableTracker]rQ   dict[str, VariableTracker]rz   r(   )F)r^  r  __qualname____firstlineno__r(   _nonvar_fieldsr=   rb  rD   r   Moduler}   r   r   r   r   r   r   r   r   r   r   r   r   r  __static_attributes____classcell__r^   s   @r8   ru   ru      s	    	
 
	'	'N22-029>2	2:- )
@
9)
914
9	
9/* .)
VmD^i &i -	i
 
ib @? &	@?
 -@? 
@? @?rs   ru   c                      ^  \ rS rSrSSS1\R
                  kr SU 4S jjrS rS rS r	\
\R                  S	 5       5       rU 4S
 jr        SS jr      SU 4S jjrSS jrSU 4S jjrSS jrSrU =r$ )r   iS  
value_typeis_state_mutatedry   c                 4  > [        U5      [        R                  R                  R                  L a  [        S5      e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 )NzScriptModules aren't supported in UnspecializedNNModuleVariable because their .forward function isn't a static member of their typer  r  rx   Fr3   )r=   rD   jit_scriptRecursiveScriptModuler   r   r|   r}   r  r`   ry   )r~   rx   rQ   lazy_value_to_becomer^   s       r8   r}   &UnspecializedNNModuleVariable.__init__c  s    ;%))++AAAW  6!#*6,+?RV#W E{22
 (,E{|$/u// % '+kk#rs   c                     U$ r2   r3   r~   attr_sources     r8   _wrap_source*UnspecializedNNModuleVariable._wrap_source|      rs   c                 @    U R                   =(       d    U R                  $ r2   r   r   s    r8   r   8UnspecializedNNModuleVariable.get_nn_module_stack_source  r   rs   c                     Xl         g r2   r   r   s     r8   r   8UnspecializedNNModuleVariable.set_nn_module_stack_source  r   rs   c                  |   [         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__)
rD   r   r{  __setattr__r}   rg   r(  rM   idr  )	supportedrH   s     r8   _nn_module_method_ids3UnspecializedNNModuleVariable._nn_module_method_ids  s     XX__00%((//2J2JK	 XX__--446
6q*% *+*< BqzzN6
 	
 
s   =B9B9B9c                   >  [         R                  " U R                  S5      nU[
        R                  R                  R                  [
        R                  R                  R                  [
        R                  R                  R                  4;   a7  UR                  [        R                  " U5      U /0 5      R                  U5      $ [        TU ]9  U5      $ ! [         a  n[        UeS nAff = f)N__iter__)r   r   r  r   NotImplementedErrorrD   r   r   r  r   r   r	  r	   r   r   r|   )r~   rK   r  er^   s       r8   r   1UnspecializedNNModuleVariable.unpack_var_sequence  s    	-''DB HH((HH""++HH((
 
 11..r2  "!"%& w*2..'  	-%1,	-s   !C 
C/#C**C/c                    U R                   n[        U5      (       a*  UR                  b  UR                  U l        [	        XX#5        [        U[        R                  R                  5      (       d6  UR                  R                  [        La  Sn[        U R                  U5      nOSn[        U R                  U5      n[        UR                  [        R                  5      (       Ga4  [        UR                   [        R                  5      (       Ga
  UR                  R                  [        L Ga  UR                   R                  ["        L Ga  SUR$                  ;  Ga  [&        R(                  " US5      n[        U[        R*                  5      (       Ga  UR,                  nU R/                  US5      R1                  5       R3                  5       (       GdF  U R/                  US5      R1                  5       R3                  5       (       Gd  U R/                  US5      R1                  5       R3                  5       (       d  U R/                  US5      R1                  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!  [9        [9        U R6                  S5      U5      n	OS n	[;        U R6                  U5        U R6                  (       a-  [=        [?        [A        U5      5      U RC                  5       X5      O	[E        5       n
U
   [F        RH                  " XiS9RK                  X/[M        U5      -   U5      sS S S 5        $ ! , (       d  f       g = f)N__call__r   rd   _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   )'rx   r   r  r  rW   r>   rD   rE   r  r  re   r#   r   r   ri   rj   r   r$   rg   r   r   r   nn_modules_globals_vtr   realizerC  rd   r`   r   rr   rb   rb  r  r   r   r	   r   r   rA   )r~   rK   rP   rQ   rO   r[   r  forward_method
globals_vtr`   ctxs              r8   r   +UnspecializedNNModuleVariable.call_function  s    jj#  ,"%"3"3"2D9 3 4 455%%-EE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$D00B;;
4;; DdKFF.t{{C@ {{ #BsGd==?  	 11"DRRFT$Z' SSs   01O++
O9c                    > US;   a{  [        U R                  U5      nU R                  (       a!  [        [        U R                  S5      U5      nOS n[        R
                  " XVS9R                  X/[        U5      -   U5      $ U[        U R                  S0 5      ;  Ga   [        R                  " [        U R                  5      U5      n[        U[        5      (       aZ  [        [        [        U R                  S5      U5      S5      nUR                  [        R
                  " UR                   US9U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	[        R4                  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	[        R@                  [        RB                  45      (       d  XR                  R0                  S   ;   a  SU l        U[8        R:                  RD                  RF                  L ao  [        US   [        RH                  5      (       aM  UR                  [        R
                  " [8        R:                  RD                  RJ                  5      XS   /U5      $ [L        T
U ]  XX45      $ ! [         a    S n GNf = f)Nr  r^   r   rg   re   r  z,UnspecializedNNModuleVariable missing methodr$  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  r`   r   r	   r   r   rA   rx   r   r   r=   r   r>   r   r	  re   rM   r  r  r  r   r   	DIFFICULTrg   r.  rZ  r   rD   r   	Parameterr  ru   r   r{  r  r   __delattr__r|   r  )r~   rK   r[   rP   rQ   r  r`   method	attr_namerx   r^   s             r8   r  )UnspecializedNNModuleVariable.call_method  s    77$/B{{#Jt{{K$H$O11"DRRFT$Z'  wtzz:r:: //TZZ0@$G &,//#z$++{CTJJ 55226??6R  
++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
 ui&>&>??))+uxx/A/AA**"5"5m"DD,0D)**"5"5j"AA,0D)%66%CC  !JJ$7$7
$CC -1D)444Q22: :
 5522588??3N3NO7O  w"2T::Y " s    *M- -M=<M=c                     U R                  X5      n[        U[        R                  5      (       a  UR	                  U5      $ g r2   )r   r>   r	   ConstDictVariablemaybe_getitem_const)r~   rK   fieldname_vtdict_vts        r8   getattr_helper,UnspecializedNNModuleVariable.getattr_helperO  s9    ""2-gy::;;..w77rs   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                  (       a  US;   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        U4S jm[        U4S j['        UR)                  5       5       5       5      n[        R*                  " U[-        U5      TS9$ [.        TU ]a  X5      $ )N)r  r  r  r  r   )r  r  c                    > [         R                  R                  U5      n[        TU 5      n[	        TU5      n[
        R                  " X%5      nX64$ r2   )r	   r   r   r   r   r*   )irI   rJ   r   
source_keysource_valuerx   hooks_dict_sources          r8   build_key_valueBUnspecializedNNModuleVariable.var_getattr.<locals>.build_key_value~  sM    0077: 00A1E
01BJO+221Cz!rs   c              3   @   >#    U  H  u  nu  p#T" XU5      v   M     g 7fr2   r3   )r5   r  rI   rJ   r  s       r8   r9   <UnspecializedNNModuleVariable.var_getattr.<locals>.<genexpr>  s%      :WYQa((:Ws   r   )r   r   r   r   rx   r>   r?   rC  r`   r   r   rk   r   EMPTY_NN_MODULE_HOOKS_DICTr	   r  SEQUENCE_LENGTHr_  r`  r   r@   NNModuleHooksDictVariabler=   r|   r   )	r~   rK   r[   
hooks_dicthooks_sourcer   r  r  r^   s	         @@r8   r   )UnspecializedNNModuleVariable.var_getattrU  s    
 
 99))FFtRR$TZZ6
j$//C
Oq4H{{'1$++t'D%(33 , G G
 %66r:: KK
 II**GGSS T2J *4;; =+66|7S7STUII((,,->?	"  :CJDTDTDV:W F 66Z(1B  w"2,,rs   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c  [        [        U5        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   r  r
   r   )r~   rK   r[   r  r   s        r8    manually_trace_nn_module_getattr>UnspecializedNNModuleVariable.manually_trace_nn_module_getattr  sn     ,,T2!!"mW=;%%b*g>C;%%b*g>C;$^R8
rs   )r  ry   r  rz   N)rK   r-   rP   rv  rQ   rw  rz   r(   ru  rt  )r^  r  rx  ry  r,   rz  r}   r  r   r   r   rm   cacher  r   r   r  r  r   r  r|  r}  r~  s   @r8   r   r   S  s      
#	1	1	N22:- __
  
/0A#A &A -	A
 
AFa; &	a;
 -a; 
a;F;-z rs   r   c                       \ rS rSrSrS rSrg)$UnspecializedBuiltinNNModuleVariablei  z_
Differentiates between builtin nn modules (e.g. torch.nn.Linear) and user defined nn modules.
c                     U$ r2   r3   r  s     r8   r  1UnspecializedBuiltinNNModuleVariable._wrap_source  r  rs   r3   N)r^  r  rx  ry  __doc__r  r|  r3   rs   r8   r  r    s    rs   r  c                   6   ^  \ rS rSrSrSU 4S jjrS 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.
c                 n   > UR                  SS 5      nUc   S5       e[        TU ]  " SSU0UD6  X0l        g )Nr`   zMFSDPManagedNNModule depends on having an accurate source to control guarding.rx   r3   )r]   r|   r}   r`   )r~   rx   rQ   r`   r^   s       r8   r}   $FSDPManagedNNModuleVariable.__init__  sE    Hd+! 	
[	
! 	/u//rs   c                     [        U[        [        45      (       d?  [        R                  R
                  R                  (       a  [        U5      $ [        U5      $ U$ r2   )r>   r   r   rD   _dynamoconfigskip_fsdp_guardsr  s     r8   r  (FSDPManagedNNModuleVariable._wrap_source  sL    ,.IJ
 
 }}##44)+662;??rs   r   r  )	r^  r  rx  ry  r  r}   r  r|  r}  r~  s   @r8   r  r    s    	 rs   r  rt  )Hr  rm   r   r&  ri   
contextlibr   r   typingr   torch.nnrD   rA  r   r   r	   excr
   r   r   r   guardsr   r   mutation_guardr   r`   r   r   r   r   r   r   r   utilsr   r   r   r   r   r   r   r   r    r!   r"   r#   r$   r   r&   r'   r(   	functionsr)   lazyr*   listsr+   user_definedr,   torch._dynamo.symbolic_convertr-   rW   rb  r   r{  rb   rr   ru   r   r  r  r3   rs   r8   <module>r     s   .     2    8 8  1 .      = < 3 %   3 D;8 	+s 	+UXX__ 	+ 	+Bw
? w
?tN$= Nb
+H "? rs   