
    N jR                        S SK r S SKrS SKrS SKrS SKrS SKJr  S SKJrJ	r	J
r
  S SKJrJrJrJrJr  S SKrS SKrS SKrS SKJs  Jr  S SKJrJrJr  S SKJr  S SKJ r J!r!  S SK"J#r#  S S	K$J%r%J&r&  S S
K'J(r(  S SK)J*r*  S SK+J,r,  S SK-J.r.J/r/  S SKJ0r0  S SK1J2r2J3r3J4r4  S SK5J6r6J7r7  S SK8J9r9J:r:  \(       a  S SK;J<r<  \" S5      r=\R|                  " \?5      r@S\2S\S\4   S\S\S\24
S jrA/ SQrBS\CS\C4S jrD SFS\R                  R                  S\GS\R                  R                  4S  jjrHS\R                  R                  SS4S! jrI " S" S#\R                  R                  5      rL\" S$S%S&/5      rM " S' S(\R                  R                  5      rOS)\R                  R                  S*\S\4   S+\S,   S-\S\S\S.\\\P\C\4   \Q\   \R\   4      SS4S/ jrSS0\R                  R                  SS4S1 jrTS2\R                  R                  SS4S3 jrU " S4 S55      rV " S6 S7\W5      rX " S8 S95      rY " S: S;5      rZS<\S=\S\Q\S4   S\P\C\4   S\R                  R                  4
S> jr[ " S? S@\R                  R                  R                  5      r^  SGS=\S\4   SA\\R\/      SB\GS\S\4   4SC jjr_SSSD.S=\S\4   SA\\R\/      S.\\\P\C\4   \Q\   \R\   4      S\S\R                  R                  4   4SE jjr`g)H    N)
namedtuple)CallableIterableSequence)AnyOptionalTYPE_CHECKINGTypeVarUnion)CaptureOutputfullgraph_captureget_traced_fn)disable)argument_namescheck_user_input_output)UserErrorType)dynamo_timedget_metrics_context)_compiling_state_context)TracingContext)_restore_state_dict)_RelaxedConstraint
Constraint)Node)ConstraintViolationError
DimDynamicStatelessSymbolicContext)_PyTreeCodeGen_PyTreeInfo)ArgumentTarget)FakeTensorModeTconstraint_violation_errorfunc.argskwargsreturnc                     SSK JnJn  [        R                  " U5      nU" X#4U5      nU R
                  (       a  U" U R
                  S   U5      4U l        U $ )z
Because we trace a different callable, the sources are all messed up.
Manually patch them so the error message looks correct.
r   )_get_input_paths_replace_sources)torch.export._unliftr*   r+   inspect	signaturer&   )r$   r%   r&   r'   r*   r+   orig_sigflat_input_pathss           p/root/GenerationalWealth/GenerationalWealth/venv/lib/python3.13/site-packages/torch/_dynamo/functional_export.pypost_process_error_msgr2   )   sZ     H  &H'A!&&7<<Q?AQR+
"' &%    ))__export_root__)z_export_root. )z._export_rootr6   textc                 J    U n[          H  u  p#UR                  X#5      nM     U$ )z;Generic utility to clean export_root patterns from strings.)EXPORT_ROOT_REPLACEMENTSreplace)r7   resultpatternreplacements       r1   clean_export_root_stringr>   E   s'    F 85 !9Mr3   graph_moduleis_inline_builtinc                   ^ S[         [        [        [        [        4   4   S[         [        [        [        [        4   4   4U4S jjnS[        [           S[        [           4S jnU R
                  R                   H  nSUR                  ;   a0  U" UR                  S   R                  5       5      UR                  S'   UR                  R                  SS5      nU(       d  Mh  U" UR                  5       5      UR                  S'   M     SU R                  ;   aN  0 nU R                  S   R                  5        H  u  px[        U5      U[        U5      '   M     UU R                  S'   U $ )a  
Clean up nn_module_stack metadata by removing export_root references.

Removes the _export_root module references from nn_module_stack metadata
in graph nodes, which are artifacts from the export process. Fixes two patterns:

1. Keys: Removes "__export_root_" and "__modules['_export_root']_" prefixes
   - Normal case: "L__self____export_root_child" -> "L__self__child"
   - inline_builtin case: Uses numeric ID strings like "140468831433840"

2. Values: Removes "._export_root" and "._modules['_export_root']" from child names
   e.g., "L['self']._export_root.child" -> "L['self'].child"
   e.g., "L['self']._modules['_export_root'].child" -> "L['self'].child"

Also removes the root export entry "L__self____export_root" entirely.

Args:
    graph_module: The GraphModule to clean up
    is_inline_builtin: If True, keys are numeric ID strings and self references
                      (L['self']) are filtered out

Returns:
    The cleaned GraphModule (modified in-place)
nn_module_stackr(   c                    > SU ;   a  U S	 0 nU R                  5        H2  u  nu  p4[        U5      n[        U5      nT(       a  US:X  a  M-  Xd4X'   M4     U$ )NL__self____export_rootz	L['self'])itemsr>   )rB   cleaned_stackkey
child_namechild_class	clean_key
clean_namer@   s          r1   _process_nn_module_stackEclean_nn_module_stack_and_source_fn.<locals>._process_nn_module_stacki   sq     $6 89 .=.C.C.E*C**05I 2*=J !Z;%>(2'@M$ /F r3   source_fn_stackc                    / nU  H  n[        U[        5      (       aZ  [        U5      S:X  aK  Uu  p4[        U[        5      (       a  [	        U5      nUR                  XT45        M_  UR                  U5        Mr  UR                  U5        M     U$ )N   )
isinstancetuplelenstrr>   append)rN   rF   itemnameclsrK   s         r1   _process_source_fn?clean_nn_module_stack_and_source_fn.<locals>._process_source_fn   s    #D$&&3t9> 	dC((!9$!?J!((*):; "((. $$T* $ r3   N dynamo_flat_name_to_original_fqn)dictrT   rR   r#   r   graphnodesmetacopygetrE   r>   )	r?   r@   rL   rY   noderN   clean_name_to_original_fqn	flat_nameoriginal_fqns	    `       r1   #clean_nn_module_stack_and_source_fnrf   M   sS   8c5a=01	c5a= 	!,HQK HQK $ ""((		)+C		+,113,DII'( ))--(94@?+=o>R>R>T+UDII'( ) *\->->>%'"'3'8'8.(

%'(#I )6 ''?	'JK( ' 	<= r3   c                 b   0 nU R                   R                   GH  nUR                  S:X  a  UR                  n[	        U5      nXC:w  a  XBl        [        X5      (       d   e[        R                  R                  R                  X5      n[        R                  R                  R                  XPU5        [        R                  R                  R                  X5        UR                  S:X  d  M  UR                  n[        U[        5      (       d   e[	        U5      n[        U[        5      (       d   e[	        UR                  5      nXC:X  a  GM8  X1;   a  X   Ul        Xbl        GMO  U R                  U5      nU R!                  U5        U R#                  XG5        XBl        Xbl        XAU'   GM     g)z3Remove export_root artifacts from FX graph in-placeget_attrcall_moduleN)r]   r^   optargetr>   hasattrtorchfxr?   	_get_attr_assign_attr	_del_attrrQ   rT   rW   get_submoduledelete_submoduleadd_submodule)r?   clean_named_module_maprb   
old_target
new_targetparamnew_namerk   s           r1   clean_export_rootrz      s]   
 .0 ""((77j J1*=J'(|8888--77Q%%225
S%%//I77m#Jj#....1*=Jj#..../		:H' 34@$	!//
;F))*5&&z:$K I1;:.? )r3   c                   J   ^  \ rS rSrS\S\SS4U 4S jjrS\SS4S	 jrS
rU =r$ )ModuleToTrace   fooin_specr(   Nc                 :   > [         TU ]  5         Xl        X l        g N)super__init___export_rootr   )selfr~   r   	__class__s      r1   r   ModuleToTrace.__init__   s    r3   	flat_argsExportTracerOutputc                     [         R                  " XR                  5      u  p#U R                  " U0 UD6n[         R                  " U5      u  pV[        XV5      $ r   )pytreetree_unflattenr   r   tree_flattenr   )r   r   r&   r'   resout_flatout_specs          r1   forwardModuleToTrace.forward   sK    ,,YE00#005!(55r3   )r   r   )	__name__
__module____qualname____firstlineno__r   r   r   __static_attributes____classcell__r   s   @r1   r|   r|      s6    C # $ 
6# 6*> 6 6r3   r|   r   r   r   c                     ^  \ rS rSrSr SS\R                  R                  S\\	   S\\
\      S\\\4   S\\\\\	4   4   S	\\	   S
S4U 4S jjjrSU 4S jjrSS jrS\S\\S4   S\\\	4   S
\	4U 4S jjrS\S\\	   S\\\	4   S
\	4U 4S jjrS\S
\	4U 4S jjrS
\R                  R                  4U 4S jjrSrU =r$ )DynamoGraphTransformer   zZGraph transformer for dynamo export that flattens inputs/outputs without complex matching.Nmoduleflat_inputsflat_args_dynamic_dimsgraph_input_ordergraph_output_map	fake_moder(   c                   > [         TU ]  U5        [        U5      [        U5      :X  d   eX l        X0l        X@l        XPl        X`l        UR                  R                   Vs/ s H  owR                  S:X  d  M  UPM     snU l        [        S UR                  R                   5       5      U l        0 U l        U R                  5         0 U l        U R#                  5         g s  snf )Nplaceholderc              3   H   #    U  H  oR                   S :X  d  M  Uv   M     g7f)outputN)rj   ).0ns     r1   	<genexpr>2DynamoGraphTransformer.__init__.<locals>.<genexpr>   s     R+=aAQ+=s   "	")r   r   rS   r   r   r   r   r   r]   r^   rj   placeholdersnextoutput_nodenew_input_nodes_create_flattened_inputsold_to_new_mapping_create_placeholder_mapping)	r   r   r   r   r   r   r   r   r   s	           r1   r   DynamoGraphTransformer.__init__   s     	 )*c+.>>>>&&<#!2 0" )/(:(:T(:1ddm>SQ(:TR6<<+=+=RR :<%%' #%((* Us   !C 8C c                   > [        [        U R                  5      5       GH3  n[        TU ]  SU 3S0 5      nXR
                  ;   az  U R
                  U   nU[        U R                  5      :  aR  U R                  U   nUR                  R                  5        H%  u  pVUS:w  d  M  XbR                  R                  U'   M'     U R                  Gb  [        U R                  U   [        R                  5      (       a  U R                  R                  U R                  U   [        [        [        U R                  U   R                   5      5       Vs/ s H5  nXpR"                  U   ;   a  [$        R&                  O[$        R(                  PM7     snS/[        U R                  U   R                   5      -  S9S9UR                  R                  S'   Ou[+        U R                  U   S5      (       a1  U R                  U   R,                  UR                  R                  S'   O&U R                  U   UR                  R                  S'   X R.                  U'   GM6     gs  snf )zKCreate new placeholder nodes for flattened inputs with proper fake tensors.arg_ valN)dynamic_sizesconstraint_sizes)symbolic_context)rangerS   r   r   r   r   r   r_   rE   rb   r   rQ   rm   Tensorfrom_tensorr   shaper   r   DYNAMICSTATICrl   r   r   )	r   ir   graph_placeholder_idxorig_placeholderrG   valuedr   s	           r1   r   /DynamoGraphTransformer._create_flattened_inputs  s   s4++,-A'-QCj"bAK ***(,(>(>q(A%(3t/@/@+AA'+'8'89N'O$&6&;&;&A&A&C
%<9>,,11#6 'D
 ~~)j  #U\\/ / 04~~/I/I$$Q'%= &+3t/?/?/B/H/H+I%J' &K $%(C(CA(F#F !+ 2 2%/%6%6!7 &K' +/#d6F6Fq6I6O6O2P)P
& 0J 0  %%e, ))!,e44/3/?/?/B/F/F  %%e,/3/?/?/B  %%e, '2  #M .*'s   <Ic                     U R                   R                  5        HL  u  pU[        U R                  5      :  d  M   U R                  U   nU R                  U   nX@R
                  U'   MN     g)z1Create mapping from old placeholders to new ones.N)r   rE   rS   r   r   r   )r   user_input_idxr   old_placeholdernew_placeholders        r1   r   2DynamoGraphTransformer._create_placeholder_mapping2  sc     6:5K5K5Q5Q5S1N$s4+<+<'=="&"3"34I"J"&"6"6~"F;J''8	 6Tr3   rk   r&   .r'   c                   > U R                   U R                  ;   a  U R                  U R                      nS HN  nXPR                   R                  ;   d  M  U R                   R                  U   UR                  R                  U'   MP     SU R                   R                  ;   aJ  SUR                  R                  ;  a0  U R                   R                  S   UR                  R                  S'   U$ [        TU ]  XU5      $ )z1Replace old placeholders with new flattened ones.)tensor_dictexample_valueunbacked_bindingsr   )current_noder   r_   rb   r   r   )r   rk   r&   r'   new_argrG   r   s         r1   r   "DynamoGraphTransformer.placeholder<  s    
  7 77--d.?.?@G M++000-1->->-C-CC-HGLL%%c* M
 ))...5@Q@Q3Q+/+<+<+A+A%+H!!%(N 7&vV<<r3   c                   > US   n/ n[        U R                  R                  5       5       Hz  nU R                  U   u  pxUS:X  a  UR                  XH   5        M/  US:X  a,  UR                  n	UR                  U R
                  U	   5        Ma  US:X  d  Mi  UR                  U5        M|     [        T
U ]  U[        U5      40 5      $ )z/Transform output according to graph_output_map.r   	graph_outinputconstant)	sortedr   keysrU   indexr   r   r   rR   )r   rk   r&   r'   original_outputsnew_outputsr   output_typer   	input_idxr   s             r1   r   DynamoGraphTransformer.outputR  s      7 --2245A#44Q7Kk)""#3#89'II	""4#7#7	#BC
*""3' 6 w~fu['9&;R@@r3   r   c                 ~  > Xl         [        TU ]	  U5      n[        US5      (       a  UR                  ULa  S H:  nX1R
                  ;   d  M  UR
                  U   UR                  R
                  U'   M<     UR                  S:w  a6  [        US5      (       a%  UR                  R                  UR                  5        U$ )z.Run node transformation and preserve metadata.rb   )r   r   r   r   rW   )	r   r   run_noderl   rb   r_   rj   _renamerW   )r   r   r;   rG   r   s       r1   r   DynamoGraphTransformer.run_nodeg  s    !!$ 66""v{{!';D&&=,-FF3KFKK$$S) E
 ttxGAv$6$6##AFF+r3   c                 Z  > [         TU ]  5       n[        U R                  S5      (       a  SU R                  R                  ;   a&  U R                  R                  S   UR                  S'   SU R                  R                  ;   a&  U R                  R                  S   UR                  S'   U$ )z:Perform the graph transformation and copy module metadata.r_   r[   dynamo_compile_id)r   	transformrl   r   r_   )r   	result_gmr   s     r1   r    DynamoGraphTransformer.transformx  s    G%'	 4;;''1T[[5E5EEEI[[EUEU 7F	AB #dkk&6&666:kk6F6F (7	23 r3   )
r   r   r   r   r   r   r   r   r   r   r   )r(   N)r   r   r   r   __doc__rm   rn   GraphModulelistr   setintr\   rR   rT   r   r   r   r   r!   r    r   r   r   r   r   r   r   r   r   s   @r1   r   r      s3   d $(+$$+ #Y+ !%SX	+
  S>+ sE#s(O34+ C=+ 
+ +>(2TK==$)(C-$8=BFsCx.=	=,AA$,SMA;?S>A	A*$ 3 "588//  r3   r   module_to_traceorig_callabler   r"   graph_capture_outputdynamic_shapesc           	         S n [        U 5      u  pUR                  R                  UR                  5        [        USS 5      =nGb  UR                  =nGbo  [        U R                  [        R                  R                  [        R                  R                  -  5      (       Gd   UR                  5         UR                  5       nUR                  [         R"                  " U5      UUU5      nU(       a3  UR$                  (       a  UR$                  S   U-   4Ul        O2U4Ul        O)U(       a  [	        U5      nO[&        R)                  SU5        UR*                   He  n[        U[,        R.                  5      (       d  M$  [	        SR1                  [2        R4                  " UR6                  U   5      5       SU S35      nMg     U(       a  [9        XqXE5      nUeg ! [         a  n
U
n S n
A
GNS n
A
ff = f)N	shape_envr   z#Summary of dimension constraints:%sr6   zk
It appears that you're trying to set a constraint on a value which we evaluated to have a static value of z0. Set TORCH_LOGS="+export" for more information.)r   r   build_guards__code__r   getattrdim_constraintsrQ   r   rm   _opsOpOverloadPacket
OpOverloadsolveforced_specializationsprettify_resultsr-   r.   r&   loginfovar_to_rangesympyIntegerjoin	tracebackformat_listvar_to_stackr2   )r   r   r   r   r&   r'   r   r$   fnr5   er   r   r   msgks                   r1   &_suggest_or_raise_constraint_violationr	    s    "&'o.11>>r{{K
 id;	;H ) 9 99_F##JJ''%***?*??
 

 	!0!G!G!I..m,&"	
 &)...33A6<3*/ 47&*/%-Ec-J*9 ''A!U]]++-Ewwy44Y5K5KA5NOPQ RJJK MEE.* ( "%;&t&
" )(	 "_ $ '%&"'s   2G 
G3&G..G3
shuffle_gmc                     U R                   R                  5         U R                  5         [        U R	                  5       5       H  u  p[        X5        [        XU5        M     g r   )r]   eliminate_dead_code	recompiler   named_buffersdelattrsetattr)r
  rW   buffers      r1   _normalize_shuffle_graphr    sL    ((*Z5578
!
&) 9r3   gmc                     U R                   R                   H1  nUR                  S:X  d  M  UR                  S   UR                  S'   M3     g )Nr   r   r   )r]   r^   rj   r_   )r  rb   s     r1   normalize_graph_moduler    s7    77m##yy9DIIe r3   c                   n    \ rS rSrS\S\S\\   SS4S jrS\	\S	4   S\	\	\S	4   \
\\4   4   4S
 jrSrg)InputProcessori  rootnum_argskwarg_namesr(   Nc                 (    Xl         X l        X0l        g r   )r  r  r  )r   r  r  r  s       r1   r   InputProcessor.__init__  s     	 &r3   inputs.c                    Un0 n[        U5      U R                  :  a:  [        [        U R                  X R                  S  5      5      nUS U R                   nU R
                  b  [        U R
                  [        R                  R                  5      (       a  [        U R
                  R                  R                  [        5      (       d   e[        U R
                  R                  R                  R                  S5      (       d   eU R
                  R                  R                  R                  R                  U L d   eU R
                  /UQ7nX#4$ )Ninput_processor)rS   r  r\   zipr  r  rQ   rm   rn   r   r]   _codegen_DynamoBytecodeCodeGenrl   dynamo_bytecode_flattenr  )r   r  r&   r'   s       r1   __call__InputProcessor.__call__  s    t9t}}$#d..]]_0EFGF$--(D99 $))UXX%9%9::!$))//":":<RSSSSIIOO,,DDFW    IIOO,,DDTT II%%D|r3   )r  r  r  )r   r   r   r   objectr   r   rT   r   rR   r\   r$  r   r   r3   r1   r  r    sk    '' ' #Y	'
 
'FCK(	uVS[!4V#44	5r3   r  c                       \ rS rSrSrg)Yieldi  r   N)r   r   r   r   r   r   r3   r1   r(  r(    s    r3   r(  c                   \    \ rS rSrS\S\S\\\4   SS4S jr	\
" SS	9S
\S\4S j5       rSrg)DynamoBytecodeFlatteni
  r  out	f_globalsr(   Nc                 6    Xl         X l        X0l        S U l        g r   )r  r+  r,  	gm_inputsr   r  r+  r,  s       r1   r   DynamoBytecodeFlatten.__init__  s      /"15r3   *do not trace internal dynamo graph capturereasonr  c                   ^  S[         SS 4U 4S jjnT R                  U5      u  p4 T R                  R                  UT R                  S9" U0 UD6  [        e! [
         a    T R                  c   eT R                  s $ f = f)Nexample_inputsr(   c                     > U Tl         [        er   )r.  r(  )r5  r   s    r1   backend_dummy5DynamoBytecodeFlatten.__call__.<locals>.backend_dummy  s    +DNKr3   compiled_fnextra_globals)r&  r  r+  forward_callabler,  r(  r.  RuntimeError)r   r  r7  r&   r'   s   `    r1   r$  DynamoBytecodeFlatten.__call__  s    	6 	d 	 ++F3	"HH%%) &    	">>--->>!	"s   *A %A>=A>)r,  r.  r  r+  )r   r   r   r   r  r   r\   rT   r&  r   dynamo_disabler$  r   r   r3   r1   r*  r*  
  s_    	6'	6 	6 V$		6
 
	6 GH 6  Ir3   r*  c                   p    \ rS rSrS\S\S\\\4   SS4S jr	\
" SS	9S
\\   S\\S4   S\4S j5       rSrg)DynamoBytecodeUnflatteni'  r  r+  r,  r(   Nc                 (    Xl         X l        X0l        g r   )r  r+  r,  r/  s       r1   r    DynamoBytecodeUnflatten.__init__(  s      /"r3   r1  r2  	flat_outsr  .c                 (  ^ S[         S[        [            4U4S jjnU R                  U5      u  pE[        R                  R                  5          U R                  R                  X0R                  S9" U0 UD6nS S S 5        U$ ! , (       d  f       W$ = f)Nr5  r(   c                     > T$ r   r   )r5  rD  s    r1   r7  7DynamoBytecodeUnflatten.__call__.<locals>.backend_dummy6  s	    r3   r9  )	r&  r   r  rm   _C_DisableTorchDispatchr+  r<  r,  )r   rD  r  r7  r&   r'   resultss    `     r1   r$   DynamoBytecodeUnflatten.__call__2  s    	6 	hv6F 	 ++F3XX++-hh//) 0 G . 	 .- s   *B
B)r,  r  r+  )r   r   r   r   r  r   r\   rT   r&  r   r?  r   rR   r$  r   r   r3   r1   rA  rA  '  sw    #'# # V$	#
 
# GH!&)383E	 Ir3   rA  r+  modc                 T
  ^ U R                   c   eU R                   n[        R                  R                  R	                  U5      u  pV[
        R                  " X#45      n[        R                  R                  R                  U[        R                  5        U R                  R                  nUR                  m[        U[        R                  R                   5      (       Ga  UR"                  R%                  5       Tl        UR&                  R%                  5       Tl        [)        U4S jUR*                   5       5      (       d   eTR*                  R-                  UR*                  5        UR.                  R%                  5       Tl        [0        R2                  S:  a.  SS Kn	U	R7                  [        R                  R                   5      n
O%[9        [        R                  R                   SS 5      n
UR:                  R=                  5        H$  u  pU
(       d  M  X;  d  M  UTR:                  U'   M&     UR>                  R%                  5       Tl        UR@                  R%                  5       Tl         URB                  R%                  5       Tl!        URD                  R%                  5       Tl"        TR>                  (       d  TR@                  (       a  [9        TSS 5      c   e[        TRF                  [        RH                  R                  RJ                  5      (       d   eTRF                  RL                  b   eS[N        S[N        S[N        4S	 jnUTRF                  l&        [        U[        R                  R                   5      (       a  TOUn[Q        U[S        U5      [U        URW                  5       5      5      n[Y        XU5      n[[        XU5      n[]        [_        [`        Rb                  " U5      X#5      UU5      TRd                  l3        [i        T5        [k        TS
5      (       a   e[k        TS5      (       a   eUTl6        UTl7        [q        TS5        TRs                  5         U R                  Rt                  Rv                  Rx                  TRz                  S'   U R                   c   eU R                   R|                  TRz                  S'   STRz                  S   l?        [        TRz                  S   5      nU R                   R                  UlA        UTRz                  S'   T$ )Nc              3   F   >#    U  H  n[        TU5      (       + v   M     g 7fr   )rl   )r   mr?   s     r1   r   7create_fx_graph_from_captured_output.<locals>.<genexpr>S  s     GAw|Q///s   !)      r   __annotations___wrapped_callr&   r'   r(   c                   ^^^ SU R                   ;  d   eU R                  mU R                  m[        U 5      R                  mS[
        S[
        S[
        4UUU4S jjn [        R                  " X05      U l        0 U l        0 U l        [        [        U 5      U ]&  " U0 UD6U R                   R                  S5        TU l        TU l        $ ! U R                   R                  S5        TU l        TU l        f = f)Nr   r&   r'   r(   c                 8   > TU l         TU l        T" U /UQ70 UD6$ r   )_forward_hooks_forward_pre_hooks)r   r&   r'   	fwd_hooksfwd_pre_hooksoriginal_forwards      r1   patched_forwardZcreate_fx_graph_from_captured_output.<locals>.dynamo_wrapped_call.<locals>.patched_forwardw  s(    *3D'.;D++DB4B6BBr3   )__dict__rW  rX  typer   r&  types
MethodTyper   r$  pop)r   r&   r'   r\  rY  rZ  r[  s       @@@r1   dynamo_wrapped_callAcreate_fx_graph_from_captured_output.<locals>.dynamo_wrapped_callp  s     555 //	 $ 7 7#':#5#5 C C6 Cf C C
<#(#3#3O#JDL*,D'.0D+ dT;TLVLMM%%i0*3D'.;D+ MM%%i0*3D'.;D+s    AC +C6_dynamo_bytecode_flatten_dynamo_bytecode_unflatten_param_name_to_sourcemodule_call_specsr   Ttracing_context)Bbackend_inputrm   _dynamoconvert_framer   r   tree_leaves
eval_framer   r   INVALID_INPUTr   r,  r?   rQ   nnModule_parametersr`   _buffersall_modulesupdate_non_persistent_buffers_setsysversion_infoannotationlibget_annotationsr   r^  rE   rW  rX  _backward_hooks_backward_pre_hooksrT  rn   _WrappedCallcls_callr&  r  rS   r   r   r*  rA  r"  r   r-   r.   r]   r!  r  rl   re  rf  r  r  output_graphexport_metadatamodule_call_specr_   r   allow_non_fake_inputsr   tensor_to_context)r+  rL  r&   r'   rj  r5   r  flat_real_argsr,  rz  annotationsrW   r   rc  r  r#  dynamo_bytecode_unflattenri  r?   s                     @r1   $create_fx_graph_from_captured_outputr  A  s    (((%%Mmm))77<GA''7N	MM4433 ((22I --L$((#'#3#3#8#8#:  $ 2 2 4GGGGGG$$T]]3,,113 	0 w& '77HK!%((//3DdKK==..0KD{t6.3%%d+ 1 '+&9&9&>&>&@#*.*A*A*F*F*H''+';';'@'@'B$+/+C+C+H+H+J(&&,*I*I
 <$?KKK**EHH,A,A,N,N     --66>>>< <6 <f <4 3FL&&/%dEHHOO<<<$D$T3t9d6;;=6IJO3O)T 7i X"8w((-t<!#L
 <(|%?@@@@|%ABBBB,CL).GL+L12  --==NN )* (((%(%6%6%@%@Lk";?Lk"8$\%6%6{%CDO(+(9(9(K(KO%+:L'(r3   c            	         ^  \ rS rSrS\\   S\S\SS4U 4S jjrS\S\4S	 jr	S
\S\4S jr
SS.S\\   S\S\S\4U 4S jjjrS\\   S\\   S\S\4S jrSS.S\R                  R                   R"                  S\S-  S\4S jjrSrU =r$ )r"  i  orig_arg_namesr#  r  r(   Nc                 b   > [         TU ]  5         Xl        X l        X0l        SU l        S U l        g NF)r   r   r  r#  r  
wrap_tuple_inputs)r   r  r#  r  r   s       r1   r   _DynamoBytecodeCodeGen.__init__  s0     	,'>$)B&/3r3   r  c                 0    Xl         U R                  " U6 nU$ r   )r  r#  )r   r  rJ  s      r1   process_inputs%_DynamoBytecodeCodeGen.process_inputs  s    ..7r3   outputsc                 r    U R                  XR                  5      nU R                  (       a  U4nS U l        U$ r   )r  r  r  )r   r  rJ  s      r1   process_outputs&_DynamoBytecodeCodeGen.process_outputs  s0    00,,G??jGr3   Fexpanded_def	free_varsmaybe_return_annotationr  c                   > U R                   n[        U5      S:  a  US   S:H  OSnU(       a  UR                  SS5        [        TU ]  US S  X#S9n[        U5      S:  a  X`R                  XAU5      -  nU$ )Nr   r   Fr  )r  rS   insertr   
gen_fn_defgen_var_bindings)r   r  r  r  fn_argshas_orig_selffn_definitionr   s          r1   r  !_DynamoBytecodeCodeGen.gen_fn_def  s     %%25g,2Bv-Q'*AJ/ + 
 y>A227|TTMr3   r  c                    U Vs/ s H(  oDR                  S5      S   R                  S5      S   PM*     nn[        U5      S:X  a  SnO*[        U5      S:X  a	  US    S3nOSR                  U5       nSU SSR                  U5       S	3$ s  snf )
N:r   #r6      , z
    _fn_args = (z)
    z,, = self._dynamo_bytecode_flatten(*_fn_args))splitrS   r  )r   r  r  r  xwithout_annotationfn_signatures          r1   r  '_DynamoBytecodeCodeGen.gen_var_bindings  s     FOOYggcl1o33C8;YOw<1L\Q%aj\,L"ii01L 	YY!"##OS 	S Ps   /B
)descsoutput_argsr  c          	          SSR                  U Vs/ s H  n[        U5      PM     sn5       S3nU R                  (       a  SU S3nSU 3$ s  snf )Nz!self._dynamo_bytecode_unflatten((r  z,), _fn_args)(z,)zreturn )r  rT   r  )r   r  r  areturneds        r1   generate_output&_DynamoBytecodeCodeGen.generate_output  s\     7tyyR]A^R]Q#a&R]A^7_6``mn??8*B'H
## B_s   A
)r  r#  r  r  r  )r   r   r   r   r   rT   r   r   r   r  r  boolr  r  rm   rn   rb   r    r&  r  r   r   r   s   @r1   r"  r"    s    4S	4 "*	4 $,4 
4c c 
s s  #9 "%
  
 &SCyS-1#YSFJS	S NR$ 88==11$=Cd]$	$ $r3   r"  constraintsrestore_state_dictc                 D   ^ ^^ S[         S[         S[         4UU U4S jjnU$ )Nr&   r'   r(   c            
        > [         R                  R                  R                  (       a   e[         R                  R                  R	                  SSS9   [        5          [        S5         [        TU UTS9nS S S 5        S S S 5        S S S 5        [        WTX5      nT(       a  [        TU5        U$ ! , (       d  f       N@= f! , (       d  f       NI= f! , (       d  f       NR= f)NFwarn)replay_side_effectsside_effect_replay_policyr   )r  )
rm   rk  configinstall_free_tensorspatchr   r   r   r  r   )r&   r'   r+  r?   r  rL  r  s       r1   inner.dynamo_graph_capture_for_export.<locals>.inner  s    ==''<<<<MM  &&$)V '   !,-#'	C . " <CdS\2 .- "! s<   C C,B39CC3
C=C
C	C
C#)r   )rL  r  r  r  s   ``` r1   dynamo_graph_capture_for_exportr    s)    
S C C  ( Lr3   )r  r   c                t   ^ ^^ UmUmS[         S[         S[        R                  R                  4UUU 4S jjnU$ )a  
Improved dynamo graph capture using transformer approach with proper fake tensor handling.

This function creates a capture instance that handles:
1. PyTree flattening/unflattening with proper input ordering
2. Dynamo graph capture with export-specific context
3. FX graph transformation for export compatibility
4. Proper fake tensor metadata preservation
5. Dynamic dimension constraint handling

Notable improvements over manual approach:
- Uses FX Transformer for cleaner graph manipulation
- Properly handles fake tensor metadata and dynamic dimensions
- Preserves all necessary metadata for export
- More robust error handling and edge case management

TODO:
1. Are we actually gonna run the bytecode?
2. Need to attach guards
r&   r'   r(   c                  	  > [        5          [        R                  " X45      u  p#[        U[        R
                  5        [        TU5      n[        T[        R                  R                  5      (       a  TR                  OTnTnTnSSKJn  U" 5         [        R                  R                  R!                  SSSSSSSS[        R                  R                  R"                  S9	n	[%        5          ['        S5         U	   [)        U[+        U5      TSS9n
U
R,                  R.                  c   e/ nU
R0                  bC  U
R0                  R2                  nU
R0                  R4                  nU
R0                  R6                  nO[        R8                  R;                  [        R                  R                  5       [        R8                  R=                  5       5      nUR>                  RA                  S 5        URC                  5         S n[E        UUUU
U UU5        U
R,                  R.                  RF                  nURH                  nURJ                  nURL                  nURN                  nS S S 5        S S S 5        S S S 5        U VVs/ s H  nU=(       d    S Vs1 s H  nURP                  [S        U5      :X  d  M  [        U[T        5      (       a  M5  URV                  RX                  RZ                  URV                  RX                  R\                  :w  d  My  UR^                  iM     snPM     nnn0 nW HZ  nUU   n[        U[        R                  R`                  Rb                  5      (       d   U5       e[e        U5      UURf                  '   M\     URi                  5        H  u  nnWU   UU'   M     [k        WUUUWW5      Rm                  5       n[o        [q        [s        [t        Rv                  " U5      X5      UW5      5      UR>                  l<        URC                  5         [{        U[        R                  R                  R|                  5        [        U5        WUR                  S	'   UUR                  S
'   UsS S S 5        $ ! , (       d  f       GN$= f! , (       d  f       GN.= f! , (       d  f       GN8= fs  snf s  snnf ! , (       d  f       g = f)Nr  )resetTF)	specialize_intspecialize_floatassume_static_by_defaultautomatic_dynamic_shapes capture_dynamic_output_shape_opscapture_scalar_outputs'constant_fold_autograd_profiler_enabledlog_graph_in_out_metadatar  r   )r   _is_export_deprecated_do_not_user   rh  r   )Ar   r   r   r   r   ro  r|   rQ   rm   rp  rq  r   r6   r  rk  r  r  install_free_tensors_for_exportr   r   r   rR   r   r  rj  r?   r   r5  rn   r   Graphr]   r   r  r	  r  graph_input_idx_to_local_sourceoutput_return_typer   r  t_ididr   constraint_rangevrlowerupperdimsourceGetItemSourcerS   r   rE   r   r   r   r   r   r-   r.   r!  rf   inline_inbuilt_nn_modulesrz   r_   )r&   r'   r   r   r   r   r  r   r  dynamo_config_ctxr+  r5  r]   r   r  graph_inputsr   r   r  r  cr   r   inpr  real_idx	graph_idxtransformed_graph_constraints_dynamic_shapesrL  s                               r1   r  /_dynamo_graph_capture_for_export.<locals>.inner,  s(   %'#)#6#6~#F K#K1L1LM+C9O+5c588??+K+KCKKQTM6BK   G % 4 4 : :#!%)-).15'+8<*. &+]]%9%9%Y%Y !; !$ $%01!'#+& ,59	 //<<HHH,.$$0--::E # 1 1 ; ;I%(%6%6%E%EN!HH001BEHHNNDTUEKK&&t,OO% $I6#!" #&":":"G"G"W"W.NN#2#E#E *33#2#C#C K " 2 &h %& %A */R//"Q%  !+1.@ A	 
 ..11771;M;M;P;P;V;VV AEE/ % # & 13#%c*!&%--*>*>*L*LMMUvUM256G2H!&,,/ $
 (9'>'>'@#)(6y(AH% (A !7&! ! ik  0>"7#4#4]#CTR0##, '')/!5==#7#7#Q#Q /0:J""#672;"";/$} ('F "! 21 &%V&W ('s   C#S/R&;R>E'R	%R-R&5SR=R86R8A R8R8!R=(ES
RR
R#R&&
R5	0S8R==S
S)r   rm   rn   r   )rL  r  r   r  r  r  s   `   @@r1    _dynamo_graph_capture_for_exportr    sE    6 %OL@%S @%C @%EHH,@,@ @% @%D Lr3   )Fr  )ar-   loggingrx  r  r`  collectionsr   collections.abcr   r   r   typingr   r   r	   r
   r   r   rm   torch.fxtorch.utils._pytreeutils_pytreer   torch._dynamo.convert_framer   r   r   torch._dynamo.decoratorsr   r?  torch._dynamo.eval_framer   r   torch._dynamo.excr   torch._dynamo.utilsr   r   torch._export.utilsr   torch._guardsr   torch.exportr   torch.export.dynamic_shapesr   r   r   %torch.fx.experimental.symbolic_shapesr   r   r   torch.fx.graphr   r   torch.fx.noder    r!   torch._subclasses.fake_tensorr"   r#   	getLoggerr   r   r2   r9   rT   r>   rn   r   r  rf   rz   rp  rq  r|   r   Transformerr   r\   rR   r   r	  r  r  r  	Exceptionr(  r*  rA  r  r]   CodeGenr"  r  r  r   r3   r1   <module>r     s     
   " 8 8 ? ?    $ $ W W > L + A 8 ( , F  
 7 * <CL!& 8&
38
& & 	&
 &* 3 3  CH[((&&[;?[
XX[|'<EHH$8$8 '<T '<T
6EHHOO 
6   4{J6OP iUXX11 iXA)XX__A)CH%A) ()A) (	A)
 A) A) U4S>5:tCy#HIJA) 
A)H*)=)= *$ *:uxx33 : : D	I 	 : 4g	g g(-c3hgAEc3hg
XXgTD$UXX^^33 D$R /3$	#s(	$z*+  c3h	> /3MQ	`	#s(	` $z*+` U4S>5:tCy#HIJ	`
 c588'''(`r3   