o
    h>                     @   s  d dl Z d dlZd dlZd dl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 ddlmZ 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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iZdd	 Zd
d Zded f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%G d"d# d#ej&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 d0d1 d1eZ.ejddG d2d3 d3eZ/ejddG d4d5 d5eZ0ejddG d6d7 d7e0Z1ejddG d8d9 d9e0Z2ejddG d:d; d;e
Z3ejddG d<d= d=e
Z4ejddG d>d? d?eZ5ejddG d@dA dAe
Z6ddBde
fdCdDZ7dS )E    N)AnyOptionalUnion)ChainedSourceGuardSourceSource   )utils)create_call_functioncreate_instruction)	enum_reprc                 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   J/var/www/html/ai/venv/lib/python3.10/site-packages/torch/_dynamo/source.pyis_constant_source(   s   
r   c                 C   s$   |   tjtjtjtjtjtjfv S N)r   r   LOCALGLOBALLOCAL_NN_MODULEGLOBAL_NN_MODULELOCAL_FSDP_MODULEGLOBAL_FSDP_MODULEr   r   r   r   is_input_source4   s   r   r   )GetItemSourceODictGetItemSourcec                 C   sn   | j |}t| jtr|| j| |S |r,t| ts J ||| 	  |S ||| j |S r   )
basereconstructr   indexr   extendr   appendcreate_load_constunpack_slice)r   codegenindex_is_sliceinstrsr   r   r   reconstruct_getitem?   s   r+   T)frozenc                   @   s:   e Zd ZU eed< dZeed< dd Zdd Zdd	 Z	d
S )LocalSource
local_nameFcell_or_freevarc                 C   s   | | jgS r   )create_loadr.   selfr(   r   r   r   r"   U      zLocalSource.reconstructc                 C      t jS r   )r   r   r2   r   r   r   r   X      zLocalSource.guard_sourcec                 C      dt | j dS )NL[])reprr.   r5   r   r   r   name[      zLocalSource.nameN)
__name__
__module____qualname__str__annotations__r/   boolr"   r   r;   r   r   r   r   r-   P   s   
 r-   c                   @   .   e Zd ZU eed< dd Zdd Zdd ZdS )	RandomValueSourcerandom_call_indexc                 C   r4   r   r   RANDOM_VALUEr5   r   r   r   r   c   r6   zRandomValueSource.guard_sourcec                 C   s"   | |jjj|| jtdgS NBINARY_SUBSCR)r0   txoutputrandom_values_varr&   rE   r   r1   r   r   r   r"   f   s   
zRandomValueSource.reconstructc                 C   s   d| j  S )Nrandom_value_)rE   r5   r   r   r   r;   m      zRandomValueSource.nameN)r=   r>   r?   intrA   r   r"   r;   r   r   r   r   rD   _   s
   
 rD   c                   @   s6   e Zd ZU eed< eed< dd Zdd Zdd Zd	S )
GeneratorStateSourcedeviceinitial_seedc                 C   r4   r   rF   r5   r   r   r   r   v   r6   z!GeneratorStateSource.guard_sourcec                 C      t  r   r   r1   r   r   r   r"   y   s   z GeneratorStateSource.reconstructc                 C   s    d| j  d| j }d| dS )Ngenerator_state__r8   r9   )rQ   rR   )r2   r;   r   r   r   r;   }   s   zGeneratorStateSource.nameN)	r=   r>   r?   r@   rA   rO   r   r"   r;   r   r   r   r   rP   q   s   
 rP   c                   @   rC   )	GlobalSourceglobal_namec                 C   s   |j | jdddgS )NFTadd)create_load_globalrX   r1   r   r   r   r"         zGlobalSource.reconstructc                 C   r4   r   r   r   r5   r   r   r   r      r6   zGlobalSource.guard_sourcec                 C   r7   )NG[r9   r:   rX   r5   r   r   r   r;      r<   zGlobalSource.nameNr=   r>   r?   r@   rA   r"   r   r;   r   r   r   r   rW      s
   
 rW   c                   @   $   e Zd Zdd Zdd Zdd ZdS )DummyGlobalSourcec                 C   rS   r   rT   r1   r   r   r   r"      r6   zDummyGlobalSource.reconstructc                 C   r4   r   r]   r5   r   r   r   r      r6   zDummyGlobalSource.guard_sourcec                 C      dS N r   r5   r   r   r   r;         zDummyGlobalSource.nameNr=   r>   r?   r"   r   r;   r   r   r   r   rb          rb   c                   @   rC   )	GlobalWeakRefSourcerX   c                 C   s   |j | jdddgtddS )NTrY   r   F)r[   rX   r
   r1   r   r   r   r"      s   zGlobalWeakRefSource.reconstructc                 C   r4   r   r]   r5   r   r   r   r      r6   z GlobalWeakRefSource.guard_sourcec                 C   r7   )Nr^   z]()r_   r5   r   r   r   r;      r<   zGlobalWeakRefSource.nameNr`   r   r   r   r   ri      s
   
 ri   c                   @   6   e Zd ZU eed< dd Zdd Zdd Zdd	 Zd
S )
AttrSourcememberc                 C   sb   | 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.r!   rl   )r!   rl   splitobject__setattr__rk   join)r2   member_partsr   r   r   __post_init__   s   
zAttrSource.__post_init__c                 C   s   | j ||| j S r   )r!   r"   create_load_attrsrl   r1   r   r   r   r"      s   zAttrSource.reconstructc                 C   
   | j  S r   r!   r   r5   r   r   r   r         
zAttrSource.guard_sourcec                 C   s:   | j  sd| j  d| j dS | j  d| j  S )Nzgetattr(, )rm   )rl   isidentifierr!   r;   r5   r   r   r   r;      s   
zAttrSource.nameN)	r=   r>   r?   r@   rA   rt   r"   r   r;   r   r   r   r   rk      s   
 	rk   c                   @      e Zd Zdd ZdS )ParamBufferSourcec                 C      t | j  S r   _GUARD_SOURCE_NN_MODULEr!   r   r5   r   r   r   r      r3   zParamBufferSource.guard_sourceNr=   r>   r?   r   r   r   r   r   r}          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_OFFSETr5   r   r   r   method_name   s   


zTensorProperty.method_nameN)r=   r>   r?   r   r   r   r   r   r   r   r   r      s
    r   c                   @   sF   e Zd ZU eed< dZee ed< dd Z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   )r!   r   r   r   r   r5   r   r   r   rt      s   z"TensorPropertySource.__post_init__c                 C   s^   g | j ||| j }| jd ur||| j |t	| jd ur(dndd |S )Nr   r   T)
r!   r"   create_load_attrr   r   r   r%   r&   r$   r
   )r2   r(   instructionsr   r   r   r"      s   

z TensorPropertySource.reconstructc                 C   rv   r   rw   r5   r   r   r   r      rx   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()[r9   z
.stride()[z.storage_offset()z
unhandled )	r   r   r   r!   r;   r   r   r   AssertionErrorr5   r   r   r   r;      s   zTensorPropertySource.name)r=   r>   r?   r   rA   r   r   rO   rt   r"   r   r;   r   r   r   r   r      s   
 r   c                   @   ,   e Zd Zdd Zdd Zdd Zdd Zd	S )
NegateSourcec                 C      | j d usJ d S r   r!   r5   r   r   r   rt     r<   zNegateSource.__post_init__c                 C   rS   r   rT   r1   r   r   r   r"     r6   zNegateSource.reconstructc                 C   rv   r   rw   r5   r   r   r   r   	  rx   zNegateSource.guard_sourcec                 C   s   | j   dS )Nz
.__neg__()r!   r;   r5   r   r   r   r;     s   zNegateSource.nameNr=   r>   r?   rt   r"   r   r;   r   r   r   r   r     s
    r   c                   @   sz   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
 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   rm   z['z']__defaults__[r9   )
r!   r   r   r   r@   rp   rq   r;   r   rO   r5   r   r   r   rt     s   ""zDefaultsSource.__post_init__c                 C   s<   | j |}||| j ||| jtdg |S rH   )r!   r"   r$   ru   r   r&   r   r   )r2   r(   r*   r   r   r   r"   )  s   
zDefaultsSource.reconstructc                 C   rv   r   rw   r5   r   r   r   r   4  rx   zDefaultsSource.guard_sourcec                 C      | j S r   )r   r5   r   r   r   r;   7  r6   zDefaultsSource.nameN)r=   r>   r?   r   rO   r@   rA   r   rB   dataclassesr   r   rt   r"   r   r;   r   r   r   r   r     s   
 r   c                       sR   e Zd ZU eed< dZeed<  fddZdd Zdd	 Z	d
d Z
dd Z  ZS )r   r#   Fr)   c                    sD   | j d usJ t| jtr t d| j  t dd d S d S )Nr#   r)   T)r!   r   r#   slicesuperrq   
__reduce__r5   	__class__r   r   rt   @  s
   zGetItemSource.__post_init__c                 C   s   g t | || jdtdS )Nr)   rI   )r+   r)   r   r1   r   r   r   r"   G  s
   zGetItemSource.reconstructc                 C   rv   r   rw   r5   r   r   r   r   M  rx   zGetItemSource.guard_sourcec                 C   s   | j sJ | j\}}|| S r   )r)   r#   )r2   slice_class
slice_argsr   r   r   r'   P  s   

zGetItemSource.unpack_slicec                 C   s   t | jtr| j  d| j  dS | jr$| j  d|  dS t | jtjr>| j  dt	| j| 
   dS | j  d| jdS )Nr   r9   )r   r#   r   r!   r;   r)   r'   enumEnumr   r   is_localr5   r   r   r   r;   U  s   &zGetItemSource.name)r=   r>   r?   r   rA   r)   rB   rt   r"   r   r'   r;   __classcell__r   r   r   r   r   ;  s   
 r   c                   @      e Zd Zdd Zdd ZdS )TupleIteratorGetItemSourcec                 C   s4   | tjd g | j||| jtddS )Ntuple_iterator_getitemr   T)load_import_fromr	   r=   r!   r"   r&   r#   r
   r1   r   r   r   r"   c  s   

z&TupleIteratorGetItemSource.reconstructc                 C   s   d| j   d| jdS )Nz___tuple_iterator_getitem(ry   rz   )r!   r;   r#   r5   r   r   r   r;   k  s   zTupleIteratorGetItemSource.nameN)r=   r>   r?   r"   r;   r   r   r   r   r   a  s    r   c                   @   r   )

TypeSourcec                 C   r   r   r   r5   r   r   r   rt   q  r<   zTypeSource.__post_init__c                 C   "   | dd | j|tdd S )Nbuiltinstyper   Tr   r!   r"   r
   r1   r   r   r   r"   t     zTypeSource.reconstructc                 C   rv   r   rw   r5   r   r   r   r   x  rx   zTypeSource.guard_sourcec                 C      d| j   dS )Nztype(rz   r   r5   r   r   r   r;   {  r<   zTypeSource.nameNr   r   r   r   r   r   o  s
    r   c                   @   rj   )SuperSourcer   c                 C   s    | j d usJ | jd usJ d S r   )r   r!   r5   r   r   r   rt     s   zSuperSource.__post_init__c                 C   s.   | dd | j|| j| tdd S )Nr   r   r   T)r   r   r"   r!   r
   r1   r   r   r   r"     s   

zSuperSource.reconstructc                 C   rv   r   rw   r5   r   r   r   r     rx   zSuperSource.guard_sourcec                 C   s   d| j   d| j  dS )Nzsuper(ry   rz   )r   r;   r!   r5   r   r   r   r;     s   zSuperSource.nameN)	r=   r>   r?   r   rA   rt   r"   r   r;   r   r   r   r   r     s   
 r   c                   @   rj   )r    r#   c                 C   r   r   r   r5   r   r   r   rt     r<   z ODictGetItemSource.__post_init__c                 C   s(   | tjjgt| |ddtddS )NFr   r   T)_create_load_constcollectionsOrderedDict__getitem__r+   r
   r1   r   r   r   r"     s   zODictGetItemSource.reconstructc                 C   rv   r   rw   r5   r   r   r   r     rx   zODictGetItemSource.guard_sourcec                 C   s   t | jtrd| jj d| jj }d| j  d| dS t | jtr3d| j  d| j  dS d| j  d| jdS )Nz__load_module("z").z___odict_getitem(ry   rz   )r   r#   r   r>   r?   r!   r;   r   )r2   repr   r   r   r;     s   zODictGetItemSource.nameN)	r=   r>   r?   r   rA   rt   r"   r   r;   r   r   r   r   r      s   
 r    c                   @   ra   )NNModuleSourcec                 C   s   | j |S r   )r!   r"   r1   r   r   r   r"     rN   zNNModuleSource.reconstructc                 C   r~   r   r   r5   r   r   r   r     r3   zNNModuleSource.guard_sourcec                 C   rv   r   r   r5   r   r   r   r;     rx   zNNModuleSource.nameNrg   r   r   r   r   r     rh   r   c                   @   r|   )NotNNModuleSourcec                 C   r~   r   )_GUARD_SOURCE_NOT_NN_MODULEr!   r   r5   r   r   r   r     r3   zNotNNModuleSource.guard_sourceNr   r   r   r   r   r     r   r   c                   @   r|   )FSDPNNModuleSourcec                 C   r~   r   )_GUARD_SOURCE_FSDP_MODULEr!   r   r5   r   r   r   r     r3   zFSDPNNModuleSource.guard_sourceNr   r   r   r   r   r     r   r   c                   @   r   )GlobalStateSourcec                 C   rc   rd   r   r5   r   r   r   r;     rf   zGlobalStateSource.namec                 C   r4   r   r]   r5   r   r   r   r     r6   zGlobalStateSource.guard_sourceNr=   r>   r?   r;   r   r   r   r   r   r         r   c                   @   s8   e Zd ZU eed< dd Zdd Zdd Zdd	d
ZdS )r   source_namec                 C   s   |j | jdddgS )NFrY   )r[   r   r1   r   r   r   r"     r\   zConstantSource.reconstructc                 C   r4   r   )r   r   r5   r   r   r   r     r6   zConstantSource.guard_sourcec                 C   r   r   )r   r5   r   r   r   r;     r6   zConstantSource.nameFc                 C   rS   r   rT   )r2   fnis_volatiler   r   r   
make_guard  r6   zConstantSource.make_guardN)F)	r=   r>   r?   r@   rA   r"   r   r;   r   r   r   r   r   r     s   
 r   c                   @   s*   e Zd ZdefddZdd Zdd ZdS )	NumpyTensorSourcereturnc                 C   r   )Nz__as_tensor(rz   r   r5   r   r   r   r;     r<   zNumpyTensorSource.namec                 C   rv   r   rw   r5   r   r   r   r     rx   zNumpyTensorSource.guard_sourcec                 C   r   )Ntorch	as_tensorr   Tr   r1   r   r   r   r"     r   zNumpyTensorSource.reconstructN)r=   r>   r?   r@   r;   r   r"   r   r   r   r   r     s    r   c                   @   r   )ShapeEnvSourcec                 C   rc   rd   r   r5   r   r   r   r;     rf   zShapeEnvSource.namec                 C   r4   r   )r   	SHAPE_ENVr5   r   r   r   r     r6   zShapeEnvSource.guard_sourceNr   r   r   r   r   r     r   r   allow_cell_or_freevarc                C   s8   t | trt| j|dS t | tsdS |s| jrdS dS )Nr   FT)r   r   is_from_local_sourcer!   r-   r/   )r   r   r   r   r   r     s   


r   )8r   r   r   typingr   r   r   torch._guardsr   r   r   re   r	   bytecode_transformationr
   r   r   r   r   r   r   r   r   r   r   r   r   r   r+   	dataclassr-   rD   rP   rW   rb   ri   rk   r}   r   r   r   r   r   r   r   r   r   r    r   r   r   r   r   r   r   r   r   r   r   r   <module>   s    












&

)
%









