o
    ,h                     @   s  d Z ddlZddlZddlZddlmZmZmZmZ ddl	m
Z
mZmZ ddlmZ ddlmZmZ er:ddlmZ ejejejejejejejejejejejejejejejejejejejeji
Zejejejejejejejejejejejejejejejejejejejeji
Zejejejejejejejejejejejejejejejejejejejeji
Z ejejejejejejejejejejejejejejejejejejejeji
Z!d	d
 Z"ej#ddG dd deZ$ej#ddG dd deZ%ej#ddG dd deZ&ej#ddG dd deZ'ej#ddG dd deZ(ej#ddG dd de
Z)ej#ddG dd de)Z*ej#ddG dd de
Z+ej#ddG dd de
Z,ej#ddG dd  d eZ-ej#ddG d!d" d"e
Z.ej#ddG d#d$ d$e+Z/ej#ddG d%d& d&e+Z0ej#ddG d'd( d(eZ1G d)d* d*ej2Z3ej#ddG d+d, d,e
Z4ej#ddG d-d. d.e
Z5ej#ddG d/d0 d0e
Z6ej#ddG d1d2 d2e
Z7ej#ddG d3d4 d4e
Z8ej#ddG d5d6 d6e
Z9G d7d8 d8e
Z:ej#ddG d9d: d:e
Z;ej#ddG d;d< d<e
Z<ej#ddG d=d> d>e
Z=ej#ddG d?d@ d@e
Z>ej#ddG dAdB dBe
Z?ej#ddG dCdD dDe<Z@ej#ddG dEdF dFe<ZAej#ddG dGdH dHe
ZBej#ddG dIdJ dJe
ZCej#ddG dKdL dLe
ZDej#ddG dMdN dNe
ZEej#ddG dOdP dPeEZFej#ddG dQdR dReFZGej#ddG dSdT dTeEZHej#ddG dUdV dVeZIej#ddG dWdX dXeZJej#ddG dYdZ dZeZKej#ddG d[d\ d\e
ZLej#ddG d]d^ d^e
ZMej#ddG d_d` d`e
ZNej#ddG dadb dbe
ZOej#ddG dcdd ddeZPej#ddG dedf dfeZQdgdhdiedjeeR fdkdlZSdgdhdiefdmdnZTdiedjeUfdodpZVdiedjeeR fdqdrZWdiefdsdtZXdieduefdvdwZYejZdiefdxdyZ[ejZdiefdzd{Z\ejZdiefd|d}Z]ejZdiefd~dZ^ejZdiefddZ_ejZdiefddZ`dS )a  
This module provides Source classes that track the origins of values in PyTorch Dynamo.
Sources represent where values come from (e.g. local variables, globals, attributes) and
are used for guard generation and code reconstruction during compilation.

The module includes specialized sources for:
- Local variables and synthetic locals
- Global variables and constants
- Object attributes and method calls
- NN module specialization (specialized vs unspecialized)
- Random values and tensor properties
- Default argument handling
- FSDP (Fully Sharded Data Parallel) modules

Sources play a key role in Dynamo's guard system by tracking value origins for
guard generation, and in code reconstruction by providing methods to rebuild
the code needed to recreate values.
    N)AnyOptionalTYPE_CHECKINGUnion)ChainedSourceGuardSourceSource   )utils)create_call_functioncreate_instruction)	PyCodegenc                 C   s>   t | trdS z|  tjkrW dS W dS  ty   Y dS w NTF)
isinstanceConstantSourceguard_sourcer   CONSTANTNotImplementedErrorsource r   O/var/www/html/scripts/venv/lib/python3.10/site-packages/torch/_dynamo/source.pyis_constant_sourcec   s   
r   T)frozenc                   @   s\   e Zd ZU eed< dZeed< dZee	e  ed< dZ
eed< dd	d
Zdd Zdd ZdS )LocalSource
local_nameFis_inputNdynamismis_derefed_cell_contentscodegenr   c                 C   s,   | j r|| j d S ||| j d S N)r   
load_derefr   append_outputcreate_loadselfr   r   r   r   reconstruct   s   zLocalSource.reconstructc                 C      t jS r    )r   LOCALr%   r   r   r   r         zLocalSource.guard_sourcec                 C      dt | j dS )NzL[])reprr   r)   r   r   r   name      zLocalSource.namer   r   )__name__
__module____qualname__str__annotations__r   boolr   r   	frozensetr   r&   r   r.   r   r   r   r   r   o   s   
 
r   c                   @   0   e Zd ZU eed< dddZdd Zdd	 Zd
S )SyntheticLocalSourcer   r   r   c                 C      | || j d S r    )r"   r#   r   r$   r   r   r   r&         z SyntheticLocalSource.reconstructc                 C   r'   r    )r   SYNTHETIC_LOCALr)   r   r   r   r      r*   z!SyntheticLocalSource.guard_sourcec                 C   s   d| j dS )NzSYNTHETIC_LOCAL[r,   )r   r)   r   r   r   r.         zSyntheticLocalSource.nameNr0   r1   r2   r3   r4   r5   r&   r   r.   r   r   r   r   r9      
   
 
r9   c                   @   s0   e Zd ZU eed< dd ZdddZdd	 Zd
S )RandomValueSourcerandom_call_indexc                 C   r'   r    )r   RANDOM_VALUEr)   r   r   r   r      r*   zRandomValueSource.guard_sourcer   r   c                 C   s:   | ||jjj | || j | td d S NBINARY_SUBSCR)r"   r#   txoutputrandom_values_varcreate_load_constrA   r   r$   r   r   r   r&      s   zRandomValueSource.reconstructc                 C   s   d| j  S )Nrandom_value_)rA   r)   r   r   r   r.         zRandomValueSource.nameNr0   )r1   r2   r3   intr5   r   r&   r.   r   r   r   r   r@      s
   
 
r@   c                   @   r8   )GlobalSourceglobal_namer   r   c                 C      | |j| jdd d S NTaddr"   create_load_globalrM   r$   r   r   r   r&         zGlobalSource.reconstructc                 C   r'   r    r   GLOBALr)   r   r   r   r      r*   zGlobalSource.guard_sourcec                 C   r+   )NG[r,   r-   rM   r)   r   r   r   r.      r/   zGlobalSource.nameNr0   r>   r   r   r   r   rL      r?   rL   c                   @   r8   )GlobalWeakRefSourcerM   r   r   c                    (      fdd  tdd d S )Nc                      s      jjddS rO   rR   r   r   r%   r   r   <lambda>   s    z1GlobalWeakRefSource.reconstruct.<locals>.<lambda>r   Fadd_push_nullextend_outputr   r$   r   r[   r   r&      s   zGlobalWeakRefSource.reconstructc                 C   r'   r    rU   r)   r   r   r   r      r*   z GlobalWeakRefSource.guard_sourcec                 C   r+   )NrW   z]()rX   r)   r   r   r   r.      r/   zGlobalWeakRefSource.nameNr0   r>   r   r   r   r   rY      s
   
 
rY   c                   @   &   e Zd Zd
ddZdd Zdd Zd	S )WeakRefCallSourcer   r   c                    rZ   )Nc                      s
    j S r    baser   r[   r   r   r\      s   
 z/WeakRefCallSource.reconstruct.<locals>.<lambda>r   Fr]   r$   r   r[   r   r&      s   zWeakRefCallSource.reconstructc                 C   
   | j  S r    rc   r   r)   r   r   r   r         
zWeakRefCallSource.guard_sourcec                 C      | j   dS )Nz()rc   r.   r)   r   r   r   r.         zWeakRefCallSource.nameNr0   r1   r2   r3   r&   r   r.   r   r   r   r   ra      s    
ra   c                   @      e Zd ZdS )CallFunctionNoArgsSourceNr1   r2   r3   r   r   r   r   rl          rl   c                   @   8   e Zd ZU eed< dd ZdddZdd	 Zd
d ZdS )
AttrSourcememberc                 C   b   | j sJ dd| jv r/| jd}t| dt| j d|d d  t| d|d  d S d S Nz9Can't construct an AttrSource without a valid base source.rc   rq   rc   rq   splitobject__setattr__rp   joinr%   member_partsr   r   r   __post_init__      
zAttrSource.__post_init__r   r   c                 C       || j  ||| j d S r    rc   r_   create_load_attrsrq   r$   r   r   r   r&         
zAttrSource.reconstructc                 C   rd   r    re   r)   r   r   r   r      rf   zAttrSource.guard_sourcec                 C   s:   | j  sd| j  d| j dS | j  d| j  S )Nzgetattr(, )rt   )rq   isidentifierrc   r.   r)   r   r   r   r.      s   
zAttrSource.nameNr0   	r1   r2   r3   r4   r5   r}   r&   r   r.   r   r   r   r   rp         
 
	rp   c                   @   ro   )GenericAttrSourcerq   c                 C   rr   rs   rv   r{   r   r   r   r}      r~   zGenericAttrSource.__post_init__r   r   c                 C   r   r    r   r$   r   r   r   r&     r   zGenericAttrSource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   zGenericAttrSource.guard_sourcec                 C      d| j   d| jdS )Nzobject.__getattribute__(r   r   rc   r.   rq   r)   r   r   r   r.     rT   zGenericAttrSource.nameNr0   r   r   r   r   r   r      r   r   c                   @   s$   e Zd ZU dZeed< dddZdS )	LocalCellSourcez
    Conceptually, this class is `LocalSource` for cell objects implicitly
    generated by Python (e.g., captured variables).
    r   r   r   c                 C   r:   r    )r"   create_load_closurer   r$   r   r   r   r&     s   zLocalCellSource.reconstructNr0   )r1   r2   r3   __doc__r4   r5   r&   r   r   r   r   r     s   
 r   c                   @   s4   e Zd ZU dZeed< dddZdd Zd	d
 ZdS )
GradSourcegradrq   r   r   c                 C   r   r    r   r$   r   r   r   r&   *  r   zGradSource.reconstructc                 C   rd   r    re   r)   r   r   r   r   .  rf   zGradSource.guard_sourcec                 C   s   | j   d| j S )Nrt   r   r)   r   r   r   r.   1  r;   zGradSource.nameNr0   )	r1   r2   r3   rq   r4   r5   r&   r   r.   r   r   r   r   r   &  s
   
 
r   c                   @      e Zd Zdd ZdS )ParamBufferSourcec                 C      t | j  S r    #_GUARD_SOURCE_SPECIALIZED_NN_MODULErc   r   r)   r   r   r   r   7  r=   zParamBufferSource.guard_sourceNr1   r2   r3   r   r   r   r   r   r   5      r   c                   @   rk   )UnspecializedParamBufferSourceNrm   r   r   r   r   r   <  rn   r   c                   @   s>   e Zd ZU dZee ed< dd Zdd Zdd Z	d	d
 Z
dS )EphemeralSourceNdescc                 C   r'   r    )r   	EPHEMERALr)   r   r   r   r   N  r*   zEphemeralSource.guard_sourcec                 C   s&   d| j d urd| j   dS d dS )Nz
<ephemeralz:  >)r   r)   r   r   r   r.   Q  s   &zEphemeralSource.namec                 C      t r    r   r%   fnr   r   r   
make_guardT     zEphemeralSource.make_guardc                 C      dS NTr   r)   r   r   r   is_ephemeralW  r   zEphemeralSource.is_ephemeral)r1   r2   r3   r   r   r4   r5   r   r.   r   r   r   r   r   r   r   J  s   
 r   c                   @   s    e Zd ZdZdZdZdd ZdS )TensorPropertyr   r	      c                 C   s.   | t ju rdS | t ju rdS | t ju rdS d S )Nsizestridestorage_offset)r   SIZESTRIDESTORAGE_OFFSETr)   r   r   r   method_name`  s   


zTensorProperty.method_nameN)r1   r2   r3   r   r   r   r   r   r   r   r   r   [  s
    r   c                   @   sH   e Zd ZU eed< dZee ed< dd Zddd	Z	d
d Z
dd ZdS )TensorPropertySourcepropNidxc                 C   s>   | j d usJ | jtju r| jd u sJ d S | jd usJ d S r    )rc   r   r   r   r   r)   r   r   r   r}   n  s   z"TensorPropertySource.__post_init__r   r   c                    s\      fdd  j jd ur  j  tjd ur'dndd d S )Nc                      s     tjdj  S )Ncall_)load_import_fromr
   r1   r   r   r   r[   r   r   r\   w  s    z2TensorPropertySource.reconstruct.<locals>.<lambda>r   r	   F)r^   rc   r   r"   rH   r_   r   r$   r   r[   r   r&   u  s   

z TensorPropertySource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   z!TensorPropertySource.guard_sourcec                 C   s   | j tju r| j  d| j dS | j tju r$| j  d| j dS | j tju r9| jd u s1J | j  dS td| j  )Nz.size()[r,   z
.stride()[z.storage_offset()z
unhandled )	r   r   r   rc   r.   r   r   r   AssertionErrorr)   r   r   r   r.     s   zTensorPropertySource.namer0   )r1   r2   r3   r   r5   r   r   rK   r}   r&   r   r.   r   r   r   r   r   i  s   
 
r   c                   @   ro   )IndexedSourcer   c                 C      | j d usJ d S r    rb   r)   r   r   r   r}     r/   zIndexedSource.__post_init__r   r   c                 C   r   r    r   r$   r   r   r   r&     r   zIndexedSource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   zIndexedSource.guard_sourcec                 C   s   d| j  d| j  dS )N(r   r   )r   rc   r.   r)   r   r   r   r.     rT   zIndexedSource.nameNr0   )	r1   r2   r3   rK   r5   r}   r&   r   r.   r   r   r   r   r     s   
 
r   c                   @   .   e Zd Zdd ZdddZdd Zd	d
 ZdS )NegateSourcec                 C   r   r    rb   r)   r   r   r   r}     r/   zNegateSource.__post_init__r   r   c                 C   r   r    r   r$   r   r   r   r&     r   zNegateSource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   zNegateSource.guard_sourcec                 C   rg   )Nz
.__neg__()rh   r)   r   r   r   r.     s   zNegateSource.nameNr0   r1   r2   r3   r}   r&   r   r.   r   r   r   r   r     
    
r   c                   @   r   )ConvertIntSourcec                 C   r   r    rb   r)   r   r   r   r}     r/   zConvertIntSource.__post_init__r   r   c                 C      || j  d S r    rb   r$   r   r   r   r&     r=   zConvertIntSource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   zConvertIntSource.guard_sourcec                 C      d| j   dS )Nz!cast_symbool_to_symint_guardless(r   rh   r)   r   r   r   r.     r/   zConvertIntSource.nameNr0   r   r   r   r   r   r     r   r   c                   @   r   )FlattenScriptObjectSourcec                 C   r   r    rb   r)   r   r   r   r}     r/   z'FlattenScriptObjectSource.__post_init__r   r   c                 C   r   r    rb   r$   r   r   r   r&     r=   z%FlattenScriptObjectSource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   z&FlattenScriptObjectSource.guard_sourcec                 C   rg   )Nz.__obj_flatten__()rh   r)   r   r   r   r.     ri   zFlattenScriptObjectSource.nameNr0   r   r   r   r   r   r     r   r   c                   @   r   )ScriptObjectQualifiedNameSourcec                 C   r   r    rb   r)   r   r   r   r}     r/   z-ScriptObjectQualifiedNameSource.__post_init__r   r   c                 C   r   r    rb   r$   r   r   r   r&     r=   z+ScriptObjectQualifiedNameSource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   z,ScriptObjectQualifiedNameSource.guard_sourcec                 C   rg   )Nz._type().qualified_name()rh   r)   r   r   r   r.     ri   z$ScriptObjectQualifiedNameSource.nameNr0   r   r   r   r   r   r     r   r   c                   @   r`   )AttrProxySourcer   r   c                 C   r   r    rb   r$   r   r   r   r&     r=   zAttrProxySource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   zAttrProxySource.guard_sourcec                 C   rg   )Nz.get_base()rh   r)   r   r   r   r.     ri   zAttrProxySource.nameNr0   rj   r   r   r   r   r     s    
r   c                   @   s|   e Zd ZU eeef ed< dZeed< e	j
ddddZ
eed< e	j
ddddZeed< dd ZdddZdd Zdd ZdS )DefaultsSourceidx_keyFis_kw)initr-   comparefield_namec              
   C   s   | j sJ d| jr0t| jtsJ t| dd t| d| j   d| j d| j d d S t| jt	s8J t| dd t| d| j   d| j d	| j d
 d S )Nz]Base must be a valid source in order to properly track and guard this Defaults to its origin.r   __kwdefaults__r   rt   z['z']__defaults__[r,   )
rc   r   r   r   r4   rx   ry   r.   r   rK   r)   r   r   r   r}     s   ""zDefaultsSource.__post_init__r   r   c                 C   s@   || j  ||| j ||| j |td d S rC   )rc   r_   r   r   r"   rH   r   r   r$   r   r   r   r&     s   
zDefaultsSource.reconstructc                 C   rd   r    re   r)   r   r   r   r   	  rf   zDefaultsSource.guard_sourcec                 C      | j S r    )r   r)   r   r   r   r.     r*   zDefaultsSource.nameNr0   )r1   r2   r3   r   rK   r4   r5   r   r6   dataclassesr   r   r}   r&   r   r.   r   r   r   r   r     s   
 
r   c                       sT   e Zd ZU eed< dZeed<  fddZddd	Zd
d Z	dd Z
dd Z  ZS )GetItemSourceindexFindex_is_slicec                    sD   | j d usJ t| jtr t d| j  t dd d S d S )Nr   r   T)rc   r   r   slicesuperry   
__reduce__r)   	__class__r   r   r}     s
   zGetItemSource.__post_init__r   r   c                 C   sJ   || j  | jr|||   n	||| j |td d S rC   )rc   r   r"   rH   unpack_slicer   r   r$   r   r   r   r&     s
   
zGetItemSource.reconstructc                 C   rd   r    re   r)   r   r   r   r   $  rf   zGetItemSource.guard_sourcec                 C   s   | j sJ | j\}}|| S r    )r   r   )r%   slice_class
slice_argsr   r   r   r   '  s   

zGetItemSource.unpack_slicec                 C   sH   t | jtrJ | jr| j  d|  dS | j  d| jdS Nr   r,   )r   r   r   r   rc   r.   r   r)   r   r   r   r.   ,  s   zGetItemSource.namer0   )r1   r2   r3   r   r5   r   r6   r}   r&   r   r   r.   __classcell__r   r   r   r   r     s   
 
r   c                   @   ro   )ConstDictKeySourcer   c                 C   rd   r    re   r)   r   r   r   r   ;  rf   zConstDictKeySource.guard_sourcer   r   c                    B      fdd  | j   | j  tdd d S )Nc                           tjdS )Ndict_keys_getitemr   r
   r1   r   r   r   r   r\   @      z0ConstDictKeySource.reconstruct.<locals>.<lambda>r   Fr^   rc   r"   rH   r   r_   r   r$   r   r   r   r&   >     

zConstDictKeySource.reconstructc                 C   r   )Nzlist(dict.keys(z))[r,   rc   r.   r   r)   r   r   r   r.   F  s   zConstDictKeySource.namec                 C   r   r   r   r)   r   r   r   is_dict_keyJ  r   zConstDictKeySource.is_dict_keyNr0   )	r1   r2   r3   r   r5   r   r&   r.   r   r   r   r   r   r   7  s   
 
r   c                   @   8   e Zd ZU eed< dd Zdd Zddd	Zd
d ZdS )DictGetItemSourcer   c                 C   0   ddl m} t| jts|| jsJ d S d S Nr	   )ConstantVariable	variablesr   r   r   r   
is_literalr%   r   r   r   r   r}   V     
zDictGetItemSource.__post_init__c                 C   rd   r    re   r)   r   r   r   r   ]  rf   zDictGetItemSource.guard_sourcer   r   c                 C   sF   || j  t| jtr|| j n	||| j |td d S rC   )rc   r   r   r   r"   rH   r   r$   r   r   r   r&   `  s
   
zDictGetItemSource.reconstructc                 C   s@   t | jtr| j  d| j  dS | j  d| jdS r   r   r   r   rc   r.   r)   r   r   r   r.   k  s   zDictGetItemSource.nameNr0   	r1   r2   r3   r   r5   r}   r   r&   r.   r   r   r   r   r   O  s   
 
r   c                   @   r   )DictSubclassGetItemSourcer   c                 C   r   r   r   r   r   r   r   r}   {  r   z'DictSubclassGetItemSource.__post_init__c                 C   rd   r    re   r)   r   r   r   r     rf   z&DictSubclassGetItemSource.guard_sourcer   r   c                    sZ      fdd  | j t| jtr | j n	  | j  tdd d S )Nc                      r   )Ndict_getitemr   r   r   r   r   r\     r   z7DictSubclassGetItemSource.reconstruct.<locals>.<lambda>r   F)	r^   rc   r   r   r   r"   rH   r_   r   r$   r   r   r   r&     s   

z%DictSubclassGetItemSource.reconstructc                 C   sB   t | jtrd| j  d| j  dS | j  d| jdS )Nzdict.__getitem__(r   r   r   r,   r   r)   r   r   r   r.     s   zDictSubclassGetItemSource.nameNr0   r   r   r   r   r   r   t  s   
 
r   c                   @   s"   e Zd ZdZd	ddZdd ZdS )
ListGetItemSourcezY
    Same as GetItemSource with reconstruct and name overridden to be list specific.
    r   r   c                    sP      fdd  | j | jrtd  | j  tdd d S )Nc                      r   )Nlist_getitemr   r   r   r   r   r\     r   z/ListGetItemSource.reconstruct.<locals>.<lambda>>List[slice] is a temporary object and should not have a sourcer   F)	r^   rc   r   RuntimeErrorr"   rH   r   r_   r   r$   r   r   r   r&     s   

zListGetItemSource.reconstructc                 C   s8   t | jtrJ | jrtdd| j  d| jdS )Nr   zlist.__getitem__(r   r   )r   r   r   r   r   rc   r.   r)   r   r   r   r.     s   zListGetItemSource.nameNr0   )r1   r2   r3   r   r&   r.   r   r   r   r   r     s    
r   c                   @   s   e Zd ZdddZdd ZdS )	TupleIteratorGetItemSourcer   r   c                    r   )Nc                      r   )Ntuple_iterator_getitemr   r   r   r   r   r\     r   z8TupleIteratorGetItemSource.reconstruct.<locals>.<lambda>r   Fr   r$   r   r   r   r&     r   z&TupleIteratorGetItemSource.reconstructc                 C   r   )Nz___tuple_iterator_getitem(r   r   r   r)   r   r   r   r.     rT   zTupleIteratorGetItemSource.nameNr0   )r1   r2   r3   r&   r.   r   r   r   r   r     s    
r   c                   @   r`   )DataclassFieldsSourcer   r   c                    0      fdd  | j  tdd d S )Nc                      r   )Ndataclass_fieldsr   r   r   r   r   r\     r   z3DataclassFieldsSource.reconstruct.<locals>.<lambda>r	   Fr^   rc   r_   r   r$   r   r   r   r&     s
   

z!DataclassFieldsSource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   z"DataclassFieldsSource.guard_sourcec                 C   r   )Nz___dataclass_fields(r   rh   r)   r   r   r   r.     r/   zDataclassFieldsSource.nameNr0   rj   r   r   r   r   r     s    
r   c                   @   r   )
TypeSourcec                 C   r   r    rb   r)   r   r   r   r}     r/   zTypeSource.__post_init__r   r   c                    r   )Nc                           ddS )Nbuiltinstyper   r   r   r   r   r\         z(TypeSource.reconstruct.<locals>.<lambda>r	   Fr   r$   r   r   r   r&        
zTypeSource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   zTypeSource.guard_sourcec                 C   r   )Nztype(r   rh   r)   r   r   r   r.     r/   zTypeSource.nameNr0   r   r   r   r   r   r     s
    
r   c                   @   r`   )OptimizerSourcer   r   c                 C   r   r    rb   r$   r   r   r   r&     r=   zOptimizerSource.reconstructc                 C   rd   r    re   r)   r   r   r   r     rf   zOptimizerSource.guard_sourcec                 C   rd   r    rh   r)   r   r   r   r.     rf   zOptimizerSource.nameNr0   rj   r   r   r   r   r        
r  c                   @   r`   )NNModuleSourcer   r   c                 C   r   r    rb   r$   r   r   r   r&     r=   zNNModuleSource.reconstructc                 C   r   r    r   r)   r   r   r   r     r=   zNNModuleSource.guard_sourcec                 C   rd   r    rh   r)   r   r   r   r.     rf   zNNModuleSource.nameNr0   rj   r   r   r   r   r    r  r  c                   @   r   )UnspecializedNNModuleSourcec                 C   r   r    )%_GUARD_SOURCE_UNSPECIALIZED_NN_MODULErc   r   r)   r   r   r   r     r=   z(UnspecializedNNModuleSource.guard_sourceNr   r   r   r   r   r    r   r  c                   @   r   )"UnspecializedBuiltinNNModuleSourcec                 C   r   r    )-_GUARD_SOURCE_UNSPECIALIZED_BUILTIN_NN_MODULErc   r   r)   r   r   r   r     r=   z/UnspecializedBuiltinNNModuleSource.guard_sourceNr   r   r   r   r   r
    r   r
  c                   @   r   )FSDPNNModuleSourcec                 C   r   r    )_GUARD_SOURCE_FSDP_MODULErc   r   r)   r   r   r   r     r=   zFSDPNNModuleSource.guard_sourceNr   r   r   r   r   r    r   r  c                   @      e Zd Zdd Zdd ZdS )GlobalStateSourcec                 C   r   Nr   r   r)   r   r   r   r.   #  r   zGlobalStateSource.namec                 C   r'   r    rU   r)   r   r   r   r   &  r*   zGlobalStateSource.guard_sourceNr1   r2   r3   r.   r   r   r   r   r   r  !      r  c                   @   r   )TorchFunctionModeStackSourceindc                 C   s   d|    dS )Nz$___get_torch_function_mode_stack_at(r   )
_get_indexr)   r   r   r   r.   .  ri   z!TorchFunctionModeStackSource.namec                 C   s   ddl m} || jS )Nr	   )TorchFunctionModeStackVariable)variables.torch_functionr  get_mode_indexr  )r%   r  r   r   r   r  1  s   z'TorchFunctionModeStackSource._get_indexr   r   c                    s<      fdd   |  g  tdd d S )Nc                      r   )N get_torch_function_mode_stack_atr   r   r   r   r   r\   8  s    z:TorchFunctionModeStackSource.reconstruct.<locals>.<lambda>r	   F)r^   r_   rH   r  r   r$   r   r   r   r&   6  s
   
z(TorchFunctionModeStackSource.reconstructc                 C   r'   r    rU   r)   r   r   r   r   ?  r*   z)TorchFunctionModeStackSource.guard_sourceNr0   )	r1   r2   r3   rK   r5   r.   r  r&   r   r   r   r   r   r  *  s   
 
	r  c                   @   s8   e Zd ZU eed< dddZdd Zdd	 Zd
d ZdS )r   source_namer   r   c                 C   rN   )NFrP   )r"   rS   r  r$   r   r   r   r&   G  rT   zConstantSource.reconstructc                 C   r'   r    )r   r   r)   r   r   r   r   J  r*   zConstantSource.guard_sourcec                 C   r   r    )r  r)   r   r   r   r.   M  r*   zConstantSource.namec                 C   r   r    r   r   r   r   r   r   P  r   zConstantSource.make_guardNr0   )	r1   r2   r3   r4   r5   r&   r   r.   r   r   r   r   r   r   C  s   
 
r   c                   @   s,   e Zd ZdefddZdd Zddd	Zd
S )NumpyTensorSourcereturnc                 C   r   )Nz___from_numpy(r   rh   r)   r   r   r   r.   V  r/   zNumpyTensorSource.namec                 C   rd   r    re   r)   r   r   r   r   Y  rf   zNumpyTensorSource.guard_sourcer   r   c                    r   )Nc                      r   )Ntorch	as_tensorr  r   r   r   r   r\   ]  r  z/NumpyTensorSource.reconstruct.<locals>.<lambda>r	   Fr   r$   r   r   r   r&   \  r  zNumpyTensorSource.reconstructNr0   )r1   r2   r3   r4   r.   r   r&   r   r   r   r   r  T  s    r  c                   @   "   e Zd ZdefddZdd ZdS )SubclassAttrListSourcer  c                 C   rg   )Nz.__tensor_flatten__()[0]rh   r)   r   r   r   r.   d  ri   zSubclassAttrListSource.namec                 C   rd   r    re   r)   r   r   r   r   g  rf   z#SubclassAttrListSource.guard_sourceNr1   r2   r3   r4   r.   r   r   r   r   r   r   b      r   c                   @   r  )FloatTensorSourcer  c                 C   r   )Nz___as_tensor(r   rh   r)   r   r   r   r.   o  r/   zFloatTensorSource.namec                 C   rd   r    re   r)   r   r   r   r   r  rf   zFloatTensorSource.guard_sourceNr!  r   r   r   r   r#  m  r"  r#  c                   @   r  )CallMethodItemSourcer  c                 C   rg   )Nz.item()rh   r)   r   r   r   r.   x  ri   zCallMethodItemSource.namec                 C   rd   r    re   r)   r   r   r   r   {  rf   z!CallMethodItemSource.guard_sourceNr!  r   r   r   r   r$  v  r"  r$  c                   @   r  )ShapeEnvSourcec                 C   r   r  r   r)   r   r   r   r.     r   zShapeEnvSource.namec                 C   r'   r    )r   	SHAPE_ENVr)   r   r   r   r     r*   zShapeEnvSource.guard_sourceNr  r   r   r   r   r%    r  r%  c                   @   r  )BackwardStateSourcec                 C   r   r  r   r)   r   r   r   r.     r   zBackwardStateSource.namec                 C   r'   r    )r   BACKWARD_STATEr)   r   r   r   r     r*   z BackwardStateSource.guard_sourceNr  r   r   r   r   r'    r  r'  Fonly_allow_inputr   r  c                C   s:   t | trt| j|dS t | tsd S |r| jsd S | jS Nr)  )r   r   get_local_source_namerc   r   r   r   r   r*  r   r   r   r,    s   


r,  c                C   s   t | |dd uS r+  )r,  r-  r   r   r   is_from_local_source  ri   r.  c                 C   s   t | d uS r    )get_global_source_namer   r   r   r   is_from_global_source  rJ   r0  c                 C   s(   t | tr
t| jS t | tsd S | jS r    )r   r   r/  rc   rL   rM   r   r   r   r   r/    s
   


r/  c                 C   s,   t | tr
t| jS t | to| jo| j S r    )r   r   is_from_nonlocal_sourcerc   r   r   r   r   r   r   r   r1    s   


r1  targetc                 C   s   t | trt| j|S | |kS r    )r   r   is_from_sourcerc   )r   r2  r   r   r   r3    s   
r3  c                 C   &   t | trdS t | trt| jS dS r   )r   r  r   &is_from_unspecialized_nn_module_sourcerc   r   r   r   r   r5    
   


r5  c                 C   r4  r   )r   r
  r   .is_from_unspecialized_builtin_nn_module_sourcerc   r   r   r   r   r7    r6  r7  c                 C   r4  r   )r   r   r   )is_from_unspecialized_param_buffer_sourcerc   r   r   r   r   r8    r6  r8  c                 C   r4  r   )r   r   r   $is_from_flatten_script_object_sourcerc   r   r   r   r   r9    r6  r9  c                 C   r4  r   )r   r  r   is_from_optimizer_sourcerc   r   r   r   r   r:    r6  r:  c                 C   sr   t | trdS t | trt | jtr| jjdkrdS t | tr-t | jtr-| jjdkr-dS t | tr7t| jS dS )NTr   r   F)	r   r   r   rc   rp   rq   r   r   is_from_defaultsr   r   r   r   r;    s"   




r;  )ar   r   enum	functoolstypingr   r   r   r   torch._guardsr   r   r   r   r
   bytecode_transformationr   r   r   r   r(   LOCAL_SPECIALIZED_NN_MODULErV   GLOBAL_SPECIALIZED_NN_MODULELOCAL_UNSPECIALIZED_NN_MODULEGLOBAL_UNSPECIALIZED_NN_MODULE%LOCAL_UNSPECIALIZED_BUILTIN_NN_MODULE&GLOBAL_UNSPECIALIZED_BUILTIN_NN_MODULELOCAL_FSDP_MODULEGLOBAL_FSDP_MODULEr   r	  r  r  r   	dataclassr   r9   r@   rL   rY   ra   rl   rp   r   r   r   r   r   r   Enumr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r  r  r  r
  r  r  r  r   r  r   r#  r$  r%  r'  r4   r,  r.  r6   r0  r/  r1  r3  	lru_cacher5  r7  r8  r9  r:  r;  r   r   r   r   <module>   sB  















(





$
&

$
*
(



















