
    rh                        d Z ddlZddlZddlZddlZddlmZmZ ddlm	Z	 ddl
ZddlmZmZmZ ddlmZmZmZmZ ddlmZmZ dd	lmZ dd
lmZmZmZmZmZm Z m!Z! ddl"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/ ddl0m1Z1m2Z2m3Z3 ddl4m5Z5 ddl6m7Z7 ddl8m9Z9 ddl:m;Z; e	rddl<m=Z= d dZ>ede?dej                  j                  fd       ZBd ZC G d de3      ZD G d de;      ZE G d deE      ZF G d deE      ZGy)!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                     t        |d      rg fdt        ||      \  }}|D cg c]
  } |       }}|j                         D 	ci c]  \  }}	| |	       }
}}	|j                  |||
       yyc c}w c c}	}w )a~  
    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                    t        |       r t        |       fd| D         S t        | t              r*| j	                         D ci c]  \  }}| |       c}}S t        | t
        t        t        f      r t        |       fd| D              S t        | t        j                  j                        rt        | j                        S | S c c}}w )Nc              3   .   K   | ]  } |        y wN .0elemconvert_to_fakes     t/var/www/html/ai-insurance-compliance-backend/venv/lib/python3.12/site-packages/torch/_dynamo/variables/nn_module.py	<genexpr>zBinitialize_lazy_module.<locals>.convert_to_fake.<locals>.<genexpr>[   s      E4!6 E   c              3   .   K   | ]  } |        y wr2   r3   r4   s     r8   r9   zBinitialize_lazy_module.<locals>.convert_to_fake.<locals>.<genexpr>_   s     Ct4Cr:   )r   type
isinstancedictitemslisttuplesettorchfxProxyr   node)xkvr7   txs      r8   r7   z/initialize_lazy_module.<locals>.convert_to_fakeY   s    QtAw E1 EFFAt$:;'')D$!Q?1--DDAeS12tAwCCCCAuxx~~.%affb11 Es   CN)hasattrr!   r?   _infer_parameters)rJ   modargskwargs
proxy_argsproxy_kwargsarg	fake_argsrH   rI   fake_kwargsr7   s   `          @r8   initialize_lazy_modulerU   O   s     s&'
	 $5T6#B 
L5?@c_S)@	@9E9K9K9MNAq/!,,NNc9k:# ( ANs   A7A<
module_keyrM   c              #   "  K   |j                         }|j                  j                  |d      }|dkD  r|  d| n| } 	 ||j                  f|j                  | <   |dz   |j                  |<   d  |j                  | = y # |j                  | = w xY ww)Nr   @r%   )name	num_callsget	__class__nn_module_stack)rV   sourcerJ   rM   fully_qualified_namerZ   s         r8   record_nn_module_stackr`   k   s     !;;=  !5q9I09AJ<q,:J+*>)N:&-6])*z*Bz*s   ;B1A= /B=BBc                    | rd|j                   v rt        |j                   d         rk|j                   d   }t        | d      }t        |      t        j
                  u rt        |d      }t        |j                  t        j                               y t        | j                  t        j                  t        j                  d                   y y )Nforward__func__)attr)__dict__callabler   r<   types
MethodTyper   
make_guardr   CLOSURE_MATCH	functoolspartialNOT_PRESENT_IN_GENERIC_DICT)r^   rM   fwdforward_sources       r8   &guard_to_detect_forward_monkeypatchingrp   x   s     $#,,y2I)J,,y)C'	:NCyE,,,!+NJ!G.33L4N4NOP !!%%$@@y     c                       e Zd Zddddhej                  Zdededej                  j                  ddf fdZ
d Zd	 Zd
 Z	 	 ddZd ZdddeddfdZd Zd ZddZd ZddZ	 	 	 	 	 	 ddZ	 d	 	 	 	 	 d fdZ xZS )NNModuleVariablemodule_typerV   valuenn_module_stack_sourcereturnNc                     t        |   di | || _        || _        || _        | j
                  sJ | j
                  | _        y )Nr3   )super__init__rt   rV   ru   r^   rv   )selfrt   rV   ru   rO   r\   s        r8   rz   zNNModuleVariable.__init__   sC     	"6"&$
{{{&*kk#rq   c                 6    | j                   xs | j                  S r2   rv   r^   r{   s    r8   get_nn_module_stack_sourcez+NNModuleVariable.get_nn_module_stack_source       **9dkk9rq   c                     || _         y r2   rv   r{   r^   s     r8   set_nn_module_stack_sourcez+NNModuleVariable.set_nn_module_stack_source   
    &,#rq   c                     | j                   S r2   )rt   r~   s    r8   python_typezNNModuleVariable.python_type   s    rq   rJ   r-   c                      y r2   r3   )r{   rJ   r^   submod	key_extraoptionss         r8   _wrap_submodulez NNModuleVariable._wrap_submodule   s     	rq   c                 z   |j                   j                  | j                        }t        |t        j
                  j                        rg }|j                         D ]{  \  }}t        j                  j                  |      }|j                   j                  || j                  |t        t        | j                  |                   |j                  |       } |S t        |t        j
                  j                   t        j
                  j"                  t        j
                  j$                  f      sJ t'        |             | j                  sJ g }t)        |      D ]Z  \  }}|j                  |j                   j                  || j                  |t        t        | j                  |                         \ |S Nr^   )outputget_submodulerV   r=   rC   nn
ModuleDictr?   r	   ConstantVariablecreateregister_attr_or_moduler   r   r^   append
ModuleListParameterList
Sequentialr&   	enumerate)r{   rJ   baseresultrY   r   name_varidxs           r8   unpack_var_sequencez$NNModuleVariable.unpack_var_sequence   sj   yy&&t7dEHH//0F $

 (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 	 rq   rY   r(   c                 ,   |j                   j                  | j                        }t        ||      }t	        t        t        | j                  |            j                  t        j                               t        j                  j                  |      S r2   )r   r   rV   rK   r   r   r   r^   ri   r   HASATTRr	   r   r   )r{   rJ   rY   rM   r   s        r8   call_obj_hasattrz!NNModuleVariable.call_obj_hasattr   sp     ii%%doo6d#:dkk489DD$$	

 ))0088rq   c                 f    |j                   j                  | j                        }t        |dd      S )NtrainingF)r   r   rV   getattrr{   rJ   rM   s      r8   is_trainingzNNModuleVariable.is_training   s)    ii%%doo6sJ..rq   c                     |j                   j                  | j                        }t        j                  |       |j
                  j                  dk7  rt        j                  t        |             t        )zIRestart analysis treating this module as an UnspecializedNNModuleVariablerz   )
r   r   rV   r   tagf_codeco_namemark_class_dynamicr<   r   r   s      r8   convert_to_unspecializedz)NNModuleVariable.convert_to_unspecialized   sT    ii%%doo6c" 99
*00c;))rq   c                    |j                   j                  | j                        }t        |      r&t	        dd|  d| ddgt
        j                         |j                   j                  j                  | |      rC|j                   j                  j                  | |d      }t        |t        j                         S t        j                  |d	      }||v S )
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_okre   )r   r   rV   r    r   r   SUPPORTABLEside_effectshas_pending_mutation_of_attr	load_attrr=   r	   DeletedVariableobject__getattribute__)r{   rJ   keyr   mutated_attr	base_dicts         r8   has_key_in_generic_dictz(NNModuleVariable.has_key_in_generic_dict   s    yy&&t7"4(M24&#?5 H&22
 99!!>>tSI9911;;D#RV;WL!,	0I0IJJJ++D*=	irq   c                    t        |      r&t        dd|  d| ddgt        j                         t	        |d      }|y	t        |t        j                        s&t        d
d|  d| ddgt        j                         dt        |d      i}t        j                  || fi |j                  |t        j                  j                  |      gi       S )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=   rg   FunctionTyper   r	   UserMethodVariablecall_functionr   r   )r{   r   rJ   rY   
obj_source
getattr_fnr   s          r8   _custom_getattr_fallbackz)NNModuleVariable._custom_getattr_fallback  s    "4(O&tfAdV45 H&22
 (tL
*e&8&89P&tfAdV4-A '22	$ Z
MBC++JHHVV++224892
 	
rq   c                 r   | j                   xr t        | j                   |      }|j                  j                  | j                        }t
        j                  |d      }d}t               }t        j                  |j                        D ]+  }|j                  |j                  j                                - | j                   s%t        dd|  d| dg t        j                          |dk(  rt#        j$                  | ||      S ||v r||   }	nYd	|v r||d	   v r||vr	|d	   |   }	nAd
|v r||d
   v r	|d
   |   }	n-d|v r||d   v r	|d   |   }	n	 t        j&                  ||      }	d}|dk(  rt/        | j                   |       |dk(  r#|s!t#        j0                  |j                  |      S |rat3        j4                  |	t7        |            }t9        |t:        t<        f      r)|j?                  t        | jA                         |             |S tC        	tD              rk| j                   r,t        t        | j                   d      |      }t        |d      }t#        jF                  |	jH                  |      jK                  || gi       S tC        |	tL              r>t#        jN                  |	jP                  t#        jR                  tU        |            |      S tC        |	tV              r&t#        jF                  |	jY                  |      |      S tC        |	tZ        j\                        rt#        jN                  |	| |      S t_        |	      sta        |	      r t3        j4                  ||	t7        |            S t        ddtc        |       d| dtc        |	       dtc        |	       dd| dtc        |	       dtc        |       ddgt        jd                         t#        j$                  | ||      S # t(        $ r9 | j+                  |||| j                         }
|
|
cY S t-        t(        |       Y w xY w)Nre   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   rJ   rY   r   rb   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   rV   r   r   rB   inspectgetmror\   updatere   keysr   r   
DYNAMO_BUGr	   GetAttrVariablegetattr_staticAttributeErrorr   r
   rp   UserDefinedClassVariabler(   buildr   r=   rs   UnspecializedNNModuleVariabler   r   r   propertyUserFunctionVariabler   r   classmethodr   rc   r,   r<   staticmethod__get__rg   r   r   r   r&   r   )r{   rJ   rY   r^   r   r   object_memberall_class_attribute_namesrG   subobjr   outs               r8   var_getattrzNNModuleVariable.var_getattr6  s$   >DKK!>yy&&t7++D*=	$'E!/ 	@A%,,QZZ__->?	@ {{0&tfAdV47 6)445 :,,T4GG9t_F)#	*--55z*40Fi'DIm4L,L}-d3F9$:1F)Fz*40F //d; % 924;;E;}55dnnVTT!''FN64JKC# 02OPQ ..t>>@$G J fh';;'
4;;(LdSF'7F 55KK!  -dHb12 , 33OO77T
C! 
 - 55NN4(   2 23 33FDPP!&)Xf-=&,,R9OPP B274=/$Oabijpbqars"`ahioap`qqr s6tfIgfoEVVqryz~r  rA  AB  C> +66		
 ((tFCCM " 66"4DKK 7  %!M("s   O4 4,P6"P65P6c                    |j                   j                  | j                        }t        | j                  | j	                         ||      5  t        |      }t        |t        j                  j                        r|j                  j                  t        j                  j                  j                  u rt        |      r| j                  |       |rJ d       |rJ |\  }|j                  j                         D ]m  \  }}|j!                  |j                   j#                  || j                  |t%        t'        | j(                  |                  |gi        |j+                         }o |cd d d        S |r+|j,                  |j,                  | _        t1        ||||       |j                   j3                         r|j4                  j7                  d      rw|j4                  dk7  rht        |dd      r| j                  |       ddlm}	  |	| |j                   j<                  d	| j                  gt?        ||       
      cd d d        S | j(                  sJ d       t        |t        j@                  jB                        r#|j                  }
t'        | j(                  d      }n"|jD                  }
t'        | j(                  d      }tG        |
tH        jJ                        r|
jL                  }
t'        |d      }| g|z   }ntG        |
tH        jN                        sJ |jQ                  tS        jT                  |
|      ||      cd d d        S # 1 sw Y   y xY w)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rJ   proxyzsMust provide a valid source in order to inline, since inlined function may have default args which must be guarded.rb   
_call_implrc   )+r   r   rV   r`   r   r   r=   rC   r   r   r\   rb   r   r   r   r?   r   r   r   r   r^   popcls_to_becomert   rU   is_root_tracer
__module__
startswithbuilderr   create_proxyr!   rD   GraphModuler   r   rg   rh   rc   r   inline_user_function_returnr	   r   )r{   rJ   rN   rO   rM   is_lazyrR   
child_namer   r   fn	fn_sources               r8   r   zNNModuleVariable.call_function  s    ii%%doo6#OOT<<>C
 a	 %S)G3 3 34MM))UXX-@-@-H-HH%c* 11"5 # I{ "!z*-,,*<*<*> #&J$$		99" OO&#1*T[[*2U#V	 :  	 &&(C# Ca	 a	F $$0'*'8'8D$
 'r3f= 		((*NN--.HINN&BB%T 11"52$0"))00% +48Ca	 a	T {{ Z{ c588#7#78 B *4;;	 BIB *4;; EI"e../B *9j AI 6D=D!"e&8&899955222iH{a	 a	 a	s   D"M8CMC8MMc                     ddl mmm  j                  j
                  j                         fd}dv r j                        S dk(  r8t         j                   j                               5   |      cd d d        S dk(  rMt        j                  t        j                  j                  j                              rj!                  d      S dk(  rd   j#                         sJ t%        d	         sJ d	   j&                  d   j(                     }t%        |t*              r|S |j                  j
                  j                        }j
                  j-                  |t/        t1         j2                              
      S |r7t5              }	j                  j6                   d dt9        |	      S  fdfd}
 fd} fd}d dk(  rj
                  j:                  j=                  t?         j2                  d             srJ g }jA                         D ]  \  }|jC                   ||               |tE                     S dk(  rj
                  j:                  j=                  t?         j2                  d             g } jF                  d4i  |
dd      D ]  \  }|jC                   ||               |tE                     S dk(  rj
                  j:                  j=                  t?         j2                  d             g } jH                  d4i  |
ddd      D ]  \  }|jC                   ||               |tE                     S dk(  rj
                  j:                  j=                  t?         j2                  d             g } jJ                  d4i  |
ddd      D ]  \  }|jC                   ||               |tE                     S dk(  rUj
                  j:                  j=                  t?         j2                  d             srJ  |jA                               S dk(  rOj
                  j:                  j=                  t?         j2                  d              |jK                               S d k(  rWj
                  j:                  j=                  t?         j2                  d              | jF                  d4i  |
d            S d!k(  rWj
                  j:                  j=                  t?         j2                  d              | jH                  d4i  |
d            S d"k(  rOsrJ g }jM                         D ]"  |jC                  j!                               $  |tE                     S d#k(  rsrJ  |j'                               S d$k(  rJsrJ g }j'                         D ]  \  }|jC                   ||               |tE                     S d%k(  r srJ j!                  tO                    S d&k(  rt%        tP        jR                  jT                  tP        jR                  jV                  f      rCrAd	   j#                         r.j!                  d	   jY                         jZ                  v       S d'k(  rىstO              dk(  sJ tP        jR                  jT                  j\                  tP        jR                  j^                  j\                  tP        jR                  jV                  j\                  tP        jR                  j`                  j\                  tP        jR                  jb                  j\                  f}te              j\                  |vrt%        d	   tf        j                        sJ ti        d	                d	   jY                         t%        tj        tl        f      sJ t5              jn                  }	t%        |	tp        jr                        sJ t?        t?         j2                        d(      }ju                  tg        jv                  |	|
       gty              z         S  j2                  sJ t%        d	   tz              r:j
                  j|                  rg }ty        t        tO                          d	   jY                            }t        d	   jY                                  D ]V  \  }}||   t/        t1         j2                              }|jC                  j
                  j-                  ||
             X d	   jY                            }j
                  j-                  |  d)t/        t1         j2                  d	   jY                                     
      }|S  j                         dd*lBmC} t%        d	   |      rd	   j                  j
                        nBd	   j#                         rd	   jY                         nt        d+d,  d- d- d- d.g /          }j
                  j-                  | j                  t/        t1         j2                              
      S d0k(  szt%        tP        jR                  j                  j                  j                        rd1k(  s=t%        tP        jR                  j                  j                  j                        rgd2k(  rbt5              jn                  }	t?        t?         j2                        d(      }ju                  tg        jv                  |	|
       gz         S j                  j                  v r^t        j                  j                           r<t        d3 t        j                  j                               D              r |      S t         E        S # 1 sw Y   xY w)5Nr%   )r   ListIteratorVariableTupleVariablec           	          	j                   j                  dj                  di       }t        |j                         t              \  }}ddlm}  |		j                   j                  d| |g||            S )Nget_attrr3   r%   r   call_method)rN   rO   r   )r   r   rV   r"   rF   r!   r   r   )
rY   	mod_proxyrP   rQ   r   rN   rO   moduler{   rJ   s
        r8   generic_call_method_helperz@NNModuleVariable.call_method.<locals>.generic_call_method_helper  s     		..	I innf5'8v'F$J. ii,,!#1j1'	 -  rq   r   _wrapped_call_implrb   _check_input_dimT_get_item_by_idxr   r   __resultc                      t        d t        j                   j                               D              s#t	        dd d d  d d d dg        y y )	Nc              3   <   K   | ]  }|j                           y wr2   )is_python_constantr5   rG   s     r8   r9   zUNNModuleVariable.call_method.<locals>.assert_all_args_kwargs_const.<locals>.<genexpr>^  s      +,$$&s   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   )rN   rO   r  rY   r{   s   r8   assert_all_args_kwargs_constzBNNModuleVariable.call_method.<locals>.assert_all_args_kwargs_const]  st     09fmmo0V  !D+D64&$qI!#f$7x?\!^rq   c                              t        	
      } t        j                  |      j                  D cg c]  }|j	                          c}i j                         D ci c]  \  }}||j	                          c}}}|j                          |j                  }| D ci c]  }|||   
 c}S c c}w c c}}w c c}w r2   )r   r   	signaturebindas_python_constantr?   apply_defaults	arguments)namesr   rG   rH   rI   
bound_argsrN   r  rO   r  rY   s         r8   
get_kwargsz0NNModuleVariable.call_method.<locals>.get_kwargsi  s    (*&B3**2.33378a1'')89?HA1a**,,HJ %%'#--J.34Az!}$44 9H 5s   B6%B;&Cc                     g }| D ]M  \  }}|j                  j                  j                  ||t         j                  |                         O  |t                     S )Nr   mutation_type)r   r   r   r   r^   r'   )	r?   r   rY   r   r   
gen_sourcer   r{   rJ   s	       r8   wrap_valuesz1NNModuleVariable.call_method.<locals>.wrap_valuest  sm    F % fII55-jd.KL	 6  (>N>PQQrq   c                      j                  |       j                  j                  || t         j                  |                   g      S r   )r   r   r   r   r^   )rY   objr   r  r$  r   r{   rJ   s     r8   named_embedz1NNModuleVariable.call_method.<locals>.named_embed  sV     $++D1II55-jd.KL	 6 
 
rq   c                     |j                  d      }|d   dk(  r| S t        |      dkD  r,|j                  d      }t        | |      } t        |      dkD  r,| S )N.r    )splitlenr   r   )r^   rY   
name_splitrG   s       r8   r$  z0NNModuleVariable.call_method.<locals>.gen_source  s[    CJ!}"j/A%NN1%#FA. j/A% Mrq   named_childrenr   r"  named_parametersr   prefixrecursenamed_buffersr   remove_duplicatenamed_modulesmemochildrenmodules
parametersbuffersr   r  r?   __len____contains____getitem__rc   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   P   K   | ]  }t        |t        j                           y wr2   )r=   r	   TensorVariabler  s     r8   r9   z/NNModuleVariable.call_method.<locals>.<genexpr>I  s%       1i667s   $&r3   )Rr+  r   r   r  rV   r   r   r   r`   r   r   is_torch_inline_allowedr   getfiler\   r
  r   r  r=   r?   ru   r   r   r   r   r^   r   __name__r)   guard_on_key_orderaddr   r/  r   r'   r0  r3  r5  r   r-  rC   r   r   ParameterDictr  r   r=  r   r   r   r<   r	   r&   strintrc   rg   r   r   r   r@   r+   exportranger   r   tensorr>  evaluate_exprr   r8  conv_ConvNd_ConvTransposeNdre   rf   r  r  r  r  ry   r  ) r{   rJ   rY   rN   rO   constantr  mod_varr   r   r   r%  r(  r   parambufferbuiltin_supportedsrcr   r   
new_modulenew_module_variabler>  r   r   r   r  r  r$  r   r  r\   s    `````                   @@@@@@@r8   r  zNNModuleVariable.call_method  s<    	LKoo((-	 	2 77 %%b$77Y (!@!@!BB 8 2$78 8
 %%+*M*MOOF,,==>+
 $**400%%7--///d1g}5551gmmDGMM2G'#@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 9fk$789'>N>PQQ''II((,,Z]-STF6v66  Xy1  8e k$678 (>N>PQQ_$II((,,ZZ-PQF 4 4 4 !Xy2DE! 9f k$789 (>N>PQQ_$II((,,ZZ-PQF 4 4 4 !VX/AB! 9f k$789 (>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6v66OI9NOPPYII((,,ZZ-PQ3v33Lj6KLMMV^''F =.55d;<='>N>PQQX''v||~..W_''F & 9fk$789'>N>PQQY''#**3v;77N"6EHH$7$79O9O#PQQ**,#**Q**,?  ]"#d)q.00##//##//&&22&&22##//! F|''/@@!$q'9+E+EFXPTUVPWHXXF1g002!#Sz222VT*33!"e&8&8999 DKK!>
K55222cBFT$Z'  ;;;$q'=1 99##F  c&k 23DG4N4N4PQD'0Q8R8R8T1U'V 	V"3i,]4;;-LMII== & #'* > 	 "(Q(B(B(D!EJ*,))*K*K"& 34-)$++tAw7Q7Q7ST  +L +' /. 11"5/$q'?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O+ 6588#3#3#8#8#I#IJ-- &//B":dkk4#@*MI11..r)D  F$$---))22489 "v}}? 
 .d337&r4v>>c8 8s   
qq!rJ   r-   rN   list[VariableTracker]rO   dict[str, VariableTracker]rw   r(   )F)rF  r   __qualname__r(   _nonvar_fieldsr<   rJ  rC   r   Modulerz   r   r   r   r   r   r   r   r   r   r   r   r   r  __classcell__r\   s   @r8   rs   rs      s	    	
 
	'	'N22-029>2	2:- )
@
9)
914
9	
9/* .)
VmD^i &i -	i
 
ib @? &	@?
 -@? 
@? @?rq   rs   c                        e Zd Zdddhej                  Z	 d fdZd Zd Zd Ze	e
j                  d               Z fd	Z	 	 	 	 	 	 	 	 dd
Z	 	 	 	 	 	 d fdZddZd fdZddZ xZS )r   
value_typeis_state_mutatedrv   c                 *   t        |      t        j                  j                  j                  u rt        d      d|v r+t        |d   dd       }t        |      |u rt        |      |d<   t        |    dd|i| d| _	        | j                  | _        y )NzScriptModules aren't supported in UnspecializedNNModuleVariable because their .forward function isn't a static member of their typere  r   ru   Fr3   )r<   rC   jit_scriptRecursiveScriptModuler   r   ry   rz   rf  r^   rv   )r{   ru   rO   lazy_value_to_becomer\   s       r8   rz   z&UnspecializedNNModuleVariable.__init__c  s    ;%))++AAAW  6!#*6,+?RV#W E{22
 (,E{|$/u// % '+kk#rq   c                     |S r2   r3   r{   attr_sources     r8   _wrap_sourcez*UnspecializedNNModuleVariable._wrap_source|      rq   c                 6    | j                   xs | j                  S r2   r}   r~   s    r8   r   z8UnspecializedNNModuleVariable.get_nn_module_stack_source  r   rq   c                     || _         y r2   r   r   s     r8   r   z8UnspecializedNNModuleVariable.set_nn_module_stack_source  r   rq   c                  b   t         j                  j                  j                  t         j                  j                  j                  h} t         j                  j                  j
                  j                         D ch c]'  }t        |d      r|| vrt        |j                        ) c}S c c}w )N__code__)
rC   r   ra  __setattr__rz   re   r  rK   idrt  )	supportedrG   s     r8   _nn_module_method_idsz3UnspecializedNNModuleVariable._nn_module_method_ids  s|     XX__00%((//2J2JK	 XX__--446
q*%!9*< qzzN
 	
 
s   =,B,c                    	 t        j                  | j                  d      }|t
        j                  j                  j                  t
        j                  j                  j                  t
        j                  j                  j                  fv r6|j                  t        j                  |      | gi       j                  |      S t        | 9  |      S # t        $ r}t        |d }~ww xY w)N__iter__)r   r   re  r   NotImplementedErrorrC   r   r   rz  r   r   r   r	   r   r   ry   )r{   rJ   r   er\   s       r8   r   z1UnspecializedNNModuleVariable.unpack_var_sequence  s    	-''DB HH((HH""++HH((
 
 11..r2  "!"%& w*2..'  	-%1,	-s    C 	C*C%%C*c                 V   | j                   }t        |      r+|j                  |j                  | _        t	        ||||       t        |t        j                  j                        s5|j                  j                  t        urd}t        | j                  |      }nd}t        | j                  |      }t        |j                  t        j                        rt        |j                   t        j                        r|j                  j                  t        u r|j                   j                  t"        u rd|j$                  vrt'        j(                  |d      }t        |t        j*                        r]|j,                  }| j/                  |d      j1                         j3                         s"| j/                  |d      j1                         j3                         s| j/                  |d      j1                         j3                         s| j/                  |d      j1                         j3                         s|j/                  |d      j3                         sx|j/                  |d	      j3                         sX|j/                  |d
      j3                         s8|j/                  |d      j3                         sd}| j                  j4                  }| j6                  r!t9        t9        | j6                  d      |      }	nd }	t;        | j6                  |       | j6                  r.t=        t?        tA        |            | jC                         ||      n	tE               }
|
5  tG        jH                  ||	      jK                  || gtM        |      z   |      cd d d        S # 1 sw Y   y xY w)N__call__r   rb   _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   )'ru   r   r   re  rU   r=   rC   rD   r   r~  rc   r#   r   r   rg   rh   r   r$   re   r   r   r   nn_modules_globals_vtr   realizer-  rb   r^   r   rp   r`   rJ  rv  r   r   r	   r   r   r@   )r{   rJ   rN   rO   rM   rY   r   forward_method
globals_vtr^   ctxs              r8   r   z+UnspecializedNNModuleVariable.call_function  s    jj#  ,"%"3"3"2sD&9 3 4 45%%-EED$/BD$/B
 3<<!1!12s~~u'7'78%%)AA''+HH-$33CCN.%*<*<=55
$$R):;CCEIIK'',ABJJLPPR'',<=EEGKKM'',@AIIKOOQ!--b2NOSSU!--b2JKOOQ!--b2IJNNP!--b2MNRRT$D00B;;
4;; DdKFF.t{{C@ {{ #BsGd==?S  	  	11"VDRRTFT$Z'	 	 	s    5NN(c                    |dv rzt        | j                  |      }| j                  r!t        t        | j                  d      |      }nd }t	        j
                  ||      j                  || gt        |      z   |      S |t        | j                  di       vr	 t        j                  t        | j                        |      }t        |t              r\t        t        t        | j                  d      |      d      }|j                  t	        j
                  |j                   |      ||      S t#        |d      rat%        |j&                        | j)                         v r<t+        dd|  d	| d	| d	| d
| d| j                   dgt,        j.                         |dk(  rd| j                  j0                  v r|d   j3                         }|d   }	t        |	t        j4                        r*|	j7                         t8        j:                  j<                  u s|| j                  j0                  d   v rd| _        no|| j                  j0                  d   v rd| _        nLt        |	t        j@                  t        jB                  f      s|| j                  j0                  d   v rd| _        |t8        j:                  jD                  jF                  u rjt        |d   t        jH                        rM|j                  t	        j
                  t8        j:                  jD                  jJ                        | |d   g|      S tL        
|   ||||      S # t        $ r d }Y w xY w)Nr  r\   r   re   rc   rt  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   ru  r   r   r%   Tr   r   )(r   re  r^   r   r	   r   r   r@   ru   r   r   r<   r   r=   r   r   rc   rK   rv  rt  rx  r   r   	DIFFICULTre   r  rC  r   rC   r   	Parameterrf  rs   r   ra  ru  r   __delattr__ry   r  )r{   rJ   rY   rN   rO   r   r^   method	attr_nameru   r\   s             r8   r  z)UnspecializedNNModuleVariable.call_method  s    77$/B{{#Jt{{K$H$O11"VDRRTFT$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 !DJJ$7$7
$CC -1D)444Q22:
 5522588??3N3NO47O  w"2tT6::Y " s   )M MMc                 ~    | j                  ||      }t        |t        j                        r|j	                  |      S y r2   )r   r=   r	   ConstDictVariablemaybe_getitem_const)r{   rJ   fieldname_vtdict_vts        r8   getattr_helperz,UnspecializedNNModuleVariable.getattr_helperO  s8    ""2u-gy::;..w77rq   c                    |dv r|j                   j                  j                  | |      st        | j                  |      }t        |t              rmt        |      dk(  r_| j                  r>t        | j                  |      }t        |j                  t        j                               t        j                  i       S | j                  r|dv r|j                   j                  j                  | |      st        | j                  |      }t        | j                  |      t        j                  t        j                                |j                   j"                  j%                         fdt        fdt'        |j)                               D              }t        j*                  |t-        |            S t.        | a  ||      S )N)r  r  r  r  r   )r  r  c                     t         j                  j                  |      }t        |       }t	        |      }t        j                  ||      }||fS r2   )r	   r   r   r   r   r*   )irH   rI   r   
source_keysource_valueru   hooks_dict_sources          r8   build_key_valuezBUnspecializedNNModuleVariable.var_getattr.<locals>.build_key_value~  sQ    0077: 00A1E
01BJO+221lCEz!rq   c              3   >   K   | ]  \  }\  }} |||        y wr2   r3   )r5   r  rH   rI   r  s       r8   r9   z<UnspecializedNNModuleVariable.var_getattr.<locals>.<genexpr>  s'      -6QA1a(s   r   )r   r   r   r   ru   r=   r>   r-  r^   r   r   ri   r   EMPTY_NN_MODULE_HOOKS_DICTr	   r  SEQUENCE_LENGTHrG  rH  r   r?   NNModuleHooksDictVariabler<   ry   r   )	r{   rJ   rY   
hooks_dicthooks_sourcer   r  r  r\   s	         @@r8   r   z)UnspecializedNNModuleVariable.var_getattrU  s    
 
 99))FFtTR$TZZ6
j$/C
Oq4H{{'1$++t'D%(33 , G G
 %66r:: KK
 II**GGdS T2J *4;; =+66|7S7STUII((,,->?	"  :CJDTDTDV:W F 66Z(1B  w"2t,,rq   c                     t        j                  |      }| j                  |d|      }|| j                  |d|      }|| j                  |d|      }|t        t        |       |S )a  
        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{   rJ   rY   r  r   s        r8    manually_trace_nn_module_getattrz>UnspecializedNNModuleVariable.manually_trace_nn_module_getattr  sn     ,,T2!!"mW=;%%b*g>C;%%b*g>C;$^R8
rq   rw   N)rJ   r-   rN   r]  rO   r^  rw   r(   r\  r[  )rF  r   r_  r,   r`  rz   ro  r   r   r   rk   cacherx  r   r   r  r  r   r  rb  rc  s   @r8   r   r   S  s      
#	1	1	N22:- __
  
/0A#A &A -	A
 
AFa; &	a;
 -a; 
a;F;-zrq   r   c                       e Zd ZdZd Zy)$UnspecializedBuiltinNNModuleVariablezg
    Differentiates between builtin nn modules (e.g. torch.nn.Linear) and user defined nn modules.
    c                     |S r2   r3   rm  s     r8   ro  z1UnspecializedBuiltinNNModuleVariable._wrap_source  rp  rq   N)rF  r   r_  __doc__ro  r3   rq   r8   r  r    s    rq   r  c                   *     e Zd ZdZd fdZd Z xZS )FSDPManagedNNModuleVariablea  
    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                 l    |j                  dd       }|J d       t        |   dd|i| || _        y )Nr^   zMFSDPManagedNNModule depends on having an accurate source to control guarding.ru   r3   )r[   ry   rz   r^   )r{   ru   rO   r^   r\   s       r8   rz   z$FSDPManagedNNModuleVariable.__init__  sG    Hd+! 	
[	
! 	/u//rq   c                     t        |t        t        f      s:t        j                  j
                  j                  rt        |      S t        |      S |S r2   )r=   r   r   rC   _dynamoconfigskip_fsdp_guardsrm  s     r8   ro  z(FSDPManagedNNModuleVariable._wrap_source  sG    ,.IJ
 }}##44)+662;??rq   r  )rF  r   r_  r  rz   ro  rb  rc  s   @r8   r  r    s    	rq   r  r[  )Hr  rk   r   r  rg   
contextlibr   r   typingr   torch.nnrC   r+  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-   rU   rJ  r   ra  r`   rp   rs   r   r  r  r3   rq   r8   <module>r     s   .     2    8 8  1 .      = < 3 %   3 D;8 	+s 	+UXX__ 	+ 	+Bw
? w
?tN$= Nb
+H "? rq   