
    ΅i                         S SK r S SKrS SKJs  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\4S	 jrS
\S\S\S\\\	4   4S jrS\S\S\S\S\S\	SS4S jrg)    N)ExportedProgram)
AOTI_FILESpackage_pt2)FileLike   )LoweredBackendModuleepreturnc                 l    " S S[         R                  R                  5      nU" U R                  5       U R                  R
                  U R                  R                  5      nU R                  u  p4[        R                  " X445      u  pV[         R                  R                  U[        U5      5      nU$ )Nc                      ^  \ rS rSrS\R
                  R                  S\R                  S\R                  SS4U 4S jjr	S r
S	rU =r$ )
<get_new_ep_with_flat_inputs_outputs.<locals>.FlattenedModule   original_modulein_specout_specr
   Nc                 F   > [         TU ]  5         Xl        X l        X0l        g N)super__init__r   r   r   )selfr   r   r   	__class__s       ^/home/james-whalen/.local/lib/python3.13/site-packages/torch/nativert/backends/_lower_utils.pyr   Eget_new_ep_with_flat_inputs_outputs.<locals>.FlattenedModule.__init__   s      G#2 "L$M    c                     [         R                  " XR                  5      nUu  p4U R                  " U0 UD6n[         R                  " U5      u  pg[        U5      $ r   )pytreetree_unflattenr   r   tree_flattentuple)r   flat_inputsinputsargskwargsoutputsflat_outputs_s           r   forwardDget_new_ep_with_flat_inputs_outputs.<locals>.FlattenedModule.forward   sO    **;EF!LD**D;F;G$11':OL&&r   )r   r   r   )__name__
__module____qualname____firstlineno__torchfxGraphModuler   TreeSpecr   r'   __static_attributes____classcell__)r   s   @r   FlattenedModuler      sJ    		%"XX11		% __		% oo			%
 		%	' 	'r   r3   )r-   nnModulemodule	call_specr   r   example_inputsr   r   exportr   )r	   r3   flattened_moduler"   r#   r    r&   flat_eps           r   #get_new_ep_with_flat_inputs_outputsr<      s    '%((// '* '
		R\\))2<<+@+@ $$LD(($8NKll!!"2E+4FGGNr   exported_program
model_name
backend_idc                   ^ U R                   u  p4U R                  R                  m[        U 5      n[        R
                  " X445      u  pg[        R                  R                  UR                  5       [        U5      SS0S9n[        U[        5      (       d!  [        S[        U5      R                   35      e[!        XRUS9n	U4S jn
["        R$                  " X5      U	l        [        R(                  R)                  XU5      nX4$ )zx
Lower an exported program to AOTInductor and return a delegate ExportedProgram
with the `executorch_call_delegate` HOP
zaot_inductor.packageT)optionszaoti_files must be a list, got )module_namec                    > [         R                  " X45      u  p4[        R                  R                  " U /UQ76 nTb  Ub  [         R
                  " UT5      $ U$ r   )r   r   r-   _higher_order_opsexecutorch_call_delegater   )r   r"   r#   r    r&   r%   r   s         r   patched_forward/lower_exported_program.<locals>.patched_forwardD   s^    ,,d^<..GG

 L$<((x@@r   )r8   r7   r   r<   r   r   r-   	_inductoraot_compiler6   r   
isinstancelistAssertionErrortyper)   r   types
MethodTyper'   r9   )r=   r>   r?   r"   r#   r;   r    r&   
aoti_fileslowered_aoti_modulerF   aoti_delegate_epr   s               @r   lower_exported_programrS   ,   s     $22LD))22H12BCG(($8NK,,%,7Mt6T - J j$''-d:.>.G.G-HI
 	
 /  #("2"2?"X||**+>fM''r   foriginal_epdelegate_epdelegate_filesc                 8    [        U XU SU 3U0U SU 3U0S9  g)zP
Package a pt2 archive file that can be consumed by NativeRT with AOTI Delegate
-)exported_programsrP   N)r   )rT   r>   r?   rU   rV   rW   s         r   #package_nativert_with_aoti_delegater[   U   sA     	l!J<(+
 "l!J<0.A r   )rN   r-   torch.utils._pytreeutils_pytreer   torch.exportr   !torch.export.pt2_archive._packager   r   torch.typesr   _lowered_aoti_moduler   r<   strr   rS   r[    r   r   <module>re      s      $ $ ( E   6O  @&(%&(36&(DG&(
?J&'&(R  !	
 !  
r   