U
    Ç9%eV  ã                   @   st   d dl Z d dlZd dlZd dlmZ ddlmZmZ G dd„ dƒZdd„ Z	d	d
„ Z
G dd„ dƒZdd„ Zdd„ ZdS )é    N)ÚModuleé   )ÚExactWeakKeyDictionaryÚis_lazy_modulec                   @   s*   e Zd Zeƒ Zdd„ Zdd„ Zdd„ ZdS )ÚMutationTrackerc                 C   s   d| _ g | _d S )Nr   )Úmutation_countÚwatchers)Úself© r
   ú[/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/torch/_dynamo/mutation_guard.pyÚ__init__   s    zMutationTracker.__init__c                 C   s@   |  j d7  _ | j}g | _|D ]}|ƒ }|d k	r| |¡ qd S )Nr   )r   r   Z
invalidate)r	   ÚnameÚtmpÚrefZguardedr
   r
   r   Úon_mutation   s    zMutationTracker.on_mutationc                 C   s   | j  t |¡¡ d S ©N)r   ÚappendÚweakrefr   )r	   Úguarded_coder
   r
   r   Útrack   s    zMutationTracker.trackN)Ú__name__Ú
__module__Ú__qualname__r   Údbr   r   r   r
   r
   r
   r   r   
   s   	r   c                 C   s:   t t| ƒƒ | tjkr"tƒ tj| < tj|  }| |¡ dS )z+invalidate guarded_code when obj is mutatedN)Úensure_patchedÚtyper   r   r   )Úobjr   Útrackerr
   r
   r   Úwatch   s
    

r   c                    s8   t | ddƒr4d| _| j‰ t ˆ ¡‡ fdd„ƒ}|| _d S )NÚ___needs_mutation_patchTFc                    s6   zt j|   |¡ W n tk
r(   Y nX ˆ | ||ƒS r   )r   r   r   ÚKeyError)r	   ÚkeyÚvalue©Zoriginal_setattrr
   r   Úcustom_setattr-   s
    z&ensure_patched.<locals>.custom_setattr)Úgetattrr   Ú__setattr__Ú	functoolsÚwraps)Úclsr$   r
   r#   r   r   (   s    r   c                   @   sL   e Zd ZdZeƒ Zeƒ Zedd„ ƒZe	dd„ ƒZ
edd„ ƒZedd	„ ƒZd
S )ÚGenerationTrackerr   c                 C   s   | j | j|< d S r   )Ú
generationÚgeneration_values©r)   r   r
   r
   r   Útag=   s    zGenerationTracker.tagc                 C   s    t | tjjƒst‚dtj| < d S )NT)Ú
issubclassÚtorchÚnnr   ÚAssertionErrorr*   Údynamic_classes)r)   r
   r
   r   Úmark_class_dynamicA   s    z$GenerationTracker.mark_class_dynamicc                 C   s   || j krdS | j | S )Néÿÿÿÿ)r,   r-   r
   r
   r   Úget_generation_valueF   s    
z&GenerationTracker.get_generation_valuec                 C   s   || j ko| j | | jkS r   )r,   r+   r-   r
   r
   r   ÚcheckL   s    
þzGenerationTracker.checkN)r   r   r   r+   r   r3   r,   Úclassmethodr.   Ústaticmethodr4   r6   r7   r
   r
   r
   r   r*   8   s   


r*   c                 C   sV   t | tjjƒrd| jkrdS t| dƒr,| jS t| ƒr8dS tj	 
t| ƒ¡pPt | ¡}|S )z5Check for nn.Modules() created dynamically or mutatedÚforwardTÚtorchdynamo_force_dynamicF)Ú
isinstancer0   r1   r   Ú__dict__Úhasattrr;   r   r*   r3   Úgetr   r7   )r   Zdynr
   r
   r   Úis_dynamic_nn_moduleT   s    
ÿr@   c                     sT   t tddƒrBtj‰ ‡ fdd„} | t_tj‰‡fdd„}|t_dt_t jd7  _d	S )
z¥
    Monkey patch torch.nn.Module.__init__ and torch.nn.Module.__setstate__
    so we can detect nn.Module instances created dynamically inside forward methods.
    Ú___needs_generation_tag_patchTc                    s   ˆ | f|ž|Ž t  | ¡ d S r   ©r*   r.   )r	   ÚargsÚkwargs)Úinitr
   r   Úpatched_initl   s    z5install_generation_tagging_init.<locals>.patched_initc                    s   ˆ | |ƒ t  | ¡ d S r   rB   )r	   Ústate)Úsetstater
   r   Úpatched_setstatet   s    
z9install_generation_tagging_init.<locals>.patched_setstateFr   N)r%   r   r   Ú__setstate__rA   r*   r+   )rF   rI   r
   )rE   rH   r   Úinstall_generation_tagging_initc   s    rK   )r'   r   Ztorch.nnr0   r   Úutilsr   r   r   r   r   r*   r@   rK   r
   r
   r
   r   Ú<module>   s   
