o
    hT                     @   s  d dl Z d dlZd dlZd dlZd dlmZ d dlmZmZm	Z	m
Z
mZmZ d dlmZmZ d dlZd dlm  mZ d dlm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" d d
l#m$Z$ edej%ej&Z'de(de)fddZ*e j+G dd dZ,G dd dej-Z.G dd dZ/e	de0Z1G dd dej2Z3G dd dZ4G dd dZ5G dd dZ6G dd  d Z7G d!d" d"Z8G d#d$ d$Z9G d%d& d&Z:G d'd( d(e:Z;G d)d* d*e:Z<G d+d, d,e:Z=G d-d. d.e:Z>G d/d0 d0e:Z?G d1d2 d2Z@G d3d4 d4ZAdS )5    N)abstractmethod)AnyCallableNewTypeOptionalTypeVarUnion)overrideSelf)TracingContext)
FakeTensorFakeTensorModeTensor)MetaConverterMetaTensorDescMetaTensorDescriberSymNode)ShapeEnv)no_dispatch	_SymNodeTnamereturnc                 C   s
   |  dS )z
    An ops filter which allows pickle-safe ops. Pickle-safe ops are built-in
    ones where it will be possible to unpickle on any machine which has PyTorch.
    )ztorch.ops.atenztorch.ops.fbgemm)
startswithr    r   ]/home/www/facesmatcher.com/frenv_anti/lib/python3.10/site-packages/torch/fx/_graph_pickler.py_ops_filter_safe   s   r   c                   @   s(   e Zd ZU eZeeegef  e	d< dS )Options
ops_filterN)
__name__
__module____qualname__r   r   r   r   strbool__annotations__r   r   r   r   r   (   s   
 r   c                	       s   e Zd ZdZddejdee ddf fddZe	de
deed	ef eed	f f fd
dZe	de
dee fddZedde
dee defddZededede
fddZ  ZS )GraphPicklerzb
    GraphPickler is a Pickler which helps pickling fx graph - in particular
    GraphModule.
    Nfileoptionsr   c                    s4   t  | |p
t | _tt | _tdd| _d S )NF)Z	copy_data)	super__init__r   r(   _UnpickleStateTokenobject_unpickle_stater   _meta_tensor_describer)selfr'   r(   	__class__r   r   r*   5   s   zGraphPickler.__init__obj.c                 C   s   t |trt| |S t |tjjrt| |S t |tjj	tjj
fr)t| |S t |tr4t| |S t |tjr@t| |S t |tjjrMt| |S t |tjjrVJ t| | }r`|S tS N)
isinstancer   _TensorPickleDatareduce_helpertorchfxGraphModule_GraphModulePickleData_opsZOperatorBaseOpOverloadPacket_OpPickleDatar   _ShapeEnvPickleDataSymInt_SymNodePickleData_guardsr   _TracingContextPickleDataNode_TorchNumpyPickleDataNotImplemented)r/   r2   reducer   r   r   reducer_overrideC   s    

zGraphPickler.reducer_overridec                 C   s   || j u rdS d S )Nunpickle_state)r-   )r/   r2   r   r   r   persistent_idn   s   
zGraphPickler.persistent_idc                 C   sH   t  }| ||}|| | W  d   S 1 sw   Y  dS )z#
        Pickle an object.
        N)ioBytesIOdumpgetvalue)clsr2   r(   streampicklerr   r   r   dumpsu   s
   


$zGraphPickler.dumpsdata	fake_modec                 C   sH   t |}t| }t||}| W  d   S 1 sw   Y  dS )z%
        Unpickle an object.
        N)_UnpickleStaterJ   rK   _GraphUnpicklerload)rR   rS   staterO   Z	unpicklerr   r   r   loads   s
   
$zGraphPickler.loadsr3   )r    r!   r"   __doc__rJ   rK   r   r   r*   r	   r,   tupler   r   rG   r#   rI   classmethodbytesrQ   staticmethodr   rX   __classcell__r   r   r0   r   r&   /   s    "*	 r&   c                   @   s   e Zd ZdeddfddZdS )rT   rS   r   Nc                 C   s   || _ t | _d S r3   )rS   r   meta_converter)r/   rS   r   r   r   r*      s   z_UnpickleState.__init__)r    r!   r"   r   r*   r   r   r   r   rT      s    rT   r+   c                       sB   e Zd Zdejdeddf fddZededefdd	Z	  Z
S )
rU   rO   rH   r   Nc                    s   t  | || _d S r3   )r)   r*   r-   )r/   rO   rH   r0   r   r   r*      s   
z_GraphUnpickler.__init__pidc                 C   s   |dkr| j S td)NrH   zInvalid persistent ID)r-   pickleUnpicklingError)r/   r`   r   r   r   persistent_load   s   
z_GraphUnpickler.persistent_load)r    r!   r"   rJ   rK   rT   r*   r	   r,   rc   r^   r   r   r0   r   rU      s    rU   c                   @   sx   e Zd ZU eeef ed< edede	de
eeege	f e
eef f fddZde	ddfd	d
Zdede	fddZdS )r>   rR   rP   r2   r   c                 C      | j | ||jffS r3   unpickler-   rN   rP   r2   r   r   r   r6      s   z!_ShapeEnvPickleData.reduce_helperenvNc                 C   s*   |j rJ |j | _| jd= | jd= d S )NZtracked_fakesZfake_tensor_cache)Z_translation_validation_enabled__dict__copyrR   )r/   rh   r   r   r   r*      s   
z_ShapeEnvPickleData.__init__rH   c                 C   sB   |j sJ |j jsJ | j D ]\}}t|j j|| q|j jS r3   )rS   	shape_envrR   itemssetattr)r/   rH   kvr   r   r   rf      s
   
z_ShapeEnvPickleData.unpickle)r    r!   r"   dictr#   r,   r%   r[   r&   r   rZ   r   r
   rT   r+   r6   r*   rf   r   r   r   r   r>      s   
 	r>   c                   @   sv   e Zd Zedededeeee	gef eee
f f fddZdeddfdd	Zdefd
dZde	dejfddZdS )r@   rP   r2   r   c                 C   s8   | |j |jf}t|tjrtj|fS tdt| )NzUnhandled SymNode type )	noder-   r4   r7   r?   r@   unpickle_sym_intNotImplementedErrortype)rN   rP   r2   argsr   r   r   r6      s   
z _SymNodePickleData.reduce_helperrq   Nc                 C   s$   |j | _|j| _|j| _|j| _d S r3   )Z_exprexprrk   pytypeZ_hinthint)r/   rq   r   r   r   r*      s   z_SymNodePickleData.__init__c                 C   s0   ddl m} | jd usJ || j| j| j| jS )Nr   r   )torch.fx.experimental.sym_noder   rk   rv   rw   rx   )r/   r   r   r   r   _to_sym_node   s   z_SymNodePickleData._to_sym_noderH   c                 C   s   t |  S r3   )r7   r?   rz   r/   rH   r   r   r   rr      s   z#_SymNodePickleData.unpickle_sym_int)r    r!   r"   r[   r&   r   rZ   r   r
   rT   r+   r6   r   r*   rz   r7   r?   rr   r   r   r   r   r@      s    r@   c                   @   sx   e Zd ZU ee ed< edededee	e
egef ee
ef f fddZdededd	fd
dZdedefddZd	S )r5   metadatarP   r2   r   c                 C   s   | j | |j||jffS r3   )rf   r.   r-   rg   r   r   r   r6      s   
z_TensorPickleData.reduce_helper	describertNc                 C   s|   | |}|jrt|jtjjjsJ tj|d d| _	t
jD ]}|dv r%qt| j	|d u s;J d| dt| j	| qd S )NrS   )rS   	view_funcz
not None: z: )Zdescribe_tensorr   r4   r7   Z_subclassesZ
meta_utilsZ_FakeTensorViewFuncdataclassesreplacer|   r   Z_UNSERIALIZABLEgetattr)r/   r}   r~   r|   rn   r   r   r   r*      s   

z_TensorPickleData.__init__rH   c                    sX   t j| j jd}dtg tjf dttjt	f dt
f fdd} j| jj|d d S )Nr   make_meta_tdevicer   c                    s:   t   t j|  |W  d    S 1 sw   Y  d S r3   )r   r   rS   )r   r   rH   r   r   	with_fake  s   $z-_TensorPickleData.unpickle.<locals>.with_fake)r   r   r|   rS   r   r7   r   r   r   r#   r   r_   Zmeta_tensorrk   )r/   rH   r|   r   r   r   r   rf     s$   
z_TensorPickleData.unpickle)r    r!   r"   r   r   r%   r[   r&   rZ   r   r
   rT   r+   r6   r   r   r*   rf   r   r   r   r   r5      s   
 
r5   c                   @   s   e Zd Zedededeeee	e
gef ee	ef f  fddZdededdfd	d
Zde
dedef fddZededee	 fddZdS )rD   rP   r2   r   c                 C   s"   |  | }r| j||jffS d S r3   )from_objectrf   r-   )rN   rP   r2   rR   r   r   r   r6     s   z#_TorchNumpyPickleData.reduce_helpermodr   Nc                 C      || _ || _d S r3   )r   r   )r/   r   r   r   r   r   r*   ,     
z_TorchNumpyPickleData.__init__rH   .c                 C   s&   t t| j| j}tjjj	 | S r3   )
r   	importlibimport_moduler   r   r7   _dynamo	variablesmiscZget_np_to_tnp_map)r/   rH   npr   r   r   rf   0  s   z_TorchNumpyPickleData.unpickletnpc                 C   s   t |sd S tjjj }z|| }sW d S W n
 ty#   Y d S w t|dd  }s.d}t|dd  }s8d S |tt	
||ksDJ | ||S )Nr!   numpyr    )callabler7   r   r   r   Zget_tnp_to_np_mapget	TypeErrorr   r   r   )rN   r   Z	tnp_to_npr   r   r   r   r   r   r   4  s    
z!_TorchNumpyPickleData.from_object)r    r!   r"   r[   r&   r,   r   rZ   r   r
   rT   r+   r6   r#   r*   rf   r   r   r   r   r   rD     s"    rD   c                   @   sz   e Zd Zededejjdee	e
egejjf ee
ef f fddZdejjdeddfd	d
ZdedejjfddZdS )r:   rP   r2   r   c                 C   s   | j | ||j|jffS r3   )rf   r(   r-   rg   r   r   r   r6   K  s   
z$_GraphModulePickleData.reduce_helpergmr(   Nc                 C   sH   t |tjjjr| }n| }|j | _	| j	d= t
|j|| _d S )N_graph)r4   r7   r8   Z_lazy_graph_moduleZ_LazyGraphModuleZ_real_recompileZ	recompileri   rj   gm_dict_GraphPickleDatar   graph)r/   r   r(   Z_python_coder   r   r   r*   W  s   
z_GraphModulePickleData.__init__rH   c                 C   s.   t jjt jj}| j|_| j|||_|S r3   )	r7   r8   r9   __new__r   ri   r   rf   r   )r/   rH   r   r   r   r   rf   a  s   z_GraphModulePickleData.unpickle)r    r!   r"   r[   r&   r7   r8   r9   rZ   r   r
   rT   r+   r6   r   r*   rf   r   r   r   r   r:   J  s    

r:   c                   @   sd   e Zd Zdejjdeejjd f deddfddZdejj	ded ejjf d	e
dejjfd
dZdS )_NodePickleDatarq   mappingr(   r   Nc                    sp   t tjj fdd|j| _t tjj fdd|j| _|j| _|j| _t	
|j|| _|j| _|j| _d S )Nc                        |  S r3   r   nr   r   r   <lambda>o      z*_NodePickleData.__init__.<locals>.<lambda>c                    r   r3   r   r   r   r   r   r   q  r   )pytreetree_map_onlyr7   r8   rC   ru   kwargsr   opr=   ra   targetrt   meta)r/   rq   r   r(   r   r   r   r*   i  s   z_NodePickleData.__init__r   rH   c                    sx   t t fdd| j}t t fdd| j}| j|}t|s)t|t	s)J |
| j|||| j| j}| j|_|S )Nc                    r   r3   r   r   r   r   r   r     r   z*_NodePickleData.unpickle.<locals>.<lambda>c                    r   r3   r   r   r   r   r   r     r   )r   r   r   ru   r   r   rf   r   r4   r#   Zcreate_noder   r   rt   r   )r/   r   r   rH   ru   r   r   rq   r   r   r   rf     s   z_NodePickleData.unpickle)r    r!   r"   r7   r8   rC   rp   r   r*   GraphrT   rf   r   r   r   r   r   h  s&    
r   c                
   @   s   e Zd Zedededeeegef ee	 f fddZ
edededd fddZed	ed
eed ed f dedd fddZededefddZed	edefddZeded	edefddZdS )r=   rP   r   r   c                 C   s   |  ||j}|j|jffS r3   )ra   r(   rf   r-   )rN   rP   r   resultr   r   r   r6     s   z_OpPickleData.reduce_helperr(   c                 C   s   t |tr	t|S tjj|}t |tjjr| 	|t
|S t |tjjr,| 	|t|S |dr>|dd\}}t||S |drO|dd\}}t|S tdt| d| d| )N)z	builtins.zmath.ztorch..   z	operator.zTARGET:  )r4   r#   _OpStrPickleDatar7   r8   rC   Z_pretty_print_targetr;   
OpOverload
_pickle_op_OpOverloadPickleDatar<   _OpOverloadPacketPickleDatar   split_OpBuiltinPickleData_OpOperatorPickleDatars   rt   )rN   r   r(   r   rootdetail_r   r   r   ra     s   



z_OpPickleData.pickler   dataclsr   r   c                 C   s4   |j  }r|| sddlm} |d|  || S )Nr   )BypassFxGraphCachez"Unable to pickle non-standard op: )r   Ztorch._inductor.codecacher   )r   r   r(   r   r   r   r   r   r     s   z_OpPickleData._pickle_oprH   c                 C   s   d S r3   r   r{   r   r   r   rf     s   z_OpPickleData.unpicklec                 C   s8   d|v r| dd\}}t | }| ||S t | S )zC
        Like `globals()[name]` but supports dotted names.
        r   r   )r   globals_getattr_by_name)rN   r   r   restr   r   r   r   _lookup_global_by_name  s
   

z$_OpPickleData._lookup_global_by_namer   c                 C   s4   d|v r| dd\}}t| |} d|v st| |S )zG
        Like `getattr(root, name)` but supports dotted names.
        r   r   )r   r   )r   r   r   r   r   r   r     s
   

z_OpPickleData._getattr_by_nameN)r    r!   r"   r[   r&   r,   rZ   r   rT   r+   r6   r   ra   r]   r#   r   rt   r   r   rf   r   r   r   r   r   r   r=     s:    r=   c                   @   s0   e Zd ZdeddfddZdedefddZdS )	r   r   r   Nc                 C   
   || _ d S r3   r   r/   r   r   r   r   r*        
z_OpStrPickleData.__init__rH   c                 C   s   | j S r3   r   r{   r   r   r   rf     s   z_OpStrPickleData.unpickle)r    r!   r"   r#   r*   rT   rf   r   r   r   r   r         r   c                   @   4   e Zd ZdeddfddZdedejjfddZ	dS )	r   r   r   Nc                 C   r   r3   r   r   r   r   r   r*     r   z_OpOverloadPickleData.__init__rH   c                 C   "   |  | j}t|tjjsJ |S r3   )r   r   r4   r7   r;   r   r/   rH   r2   r   r   r   rf        z_OpOverloadPickleData.unpickle)
r    r!   r"   r#   r*   rT   r7   r;   r   rf   r   r   r   r   r         r   c                   @   r   )	r   r   r   Nc                 C   r   r3   r   r   r   r   r   r*     r   z$_OpOverloadPacketPickleData.__init__rH   c                 C   r   r3   )r   r   r4   r7   r;   r<   r   r   r   r   rf     r   z$_OpOverloadPacketPickleData.unpickle)
r    r!   r"   r#   r*   rT   r7   r;   r<   rf   r   r   r   r   r     r   r   c                   @   s4   e Zd ZdededdfddZdedefdd	ZdS )
r   r   r   r   Nc                 C   r   r3   )r   r   )r/   r   r   r   r   r   r*     r   z_OpBuiltinPickleData.__init__rH   c                 C   sR   | j dkrt| jS | j dkrdd l}| || jS | j dkr'| t| jS t)Nbuiltinsmathr   r7   )r   __builtins__r   r   r   r   r7   rs   )r/   rH   r   r   r   r   rf     s   


z_OpBuiltinPickleData.unpickler    r!   r"   r#   r*   rT   r,   rf   r   r   r   r   r     s    r   c                   @   s0   e Zd ZdeddfddZdedefddZdS )	r   r   r   Nc                 C   r   r3   r   r   r   r   r   r*     r   z_OpOperatorPickleData.__init__rH   c                 C   s   dd l }| || jS )Nr   )operatorr   r   )r/   rH   r   r   r   r   rf     s   z_OpOperatorPickleData.unpickler   r   r   r   r   r     r   r   c                   @   sD   e Zd ZdejjdeddfddZdejjde	dejjfd	d
Z
dS )r   r   r(   r   Nc                 C   sB   |j | _|j| _i }|jD ]
}t|||||< qt| | _d S r3   )Z_tracer_cls
tracer_clsZ_tracer_extrastracer_extrasnodesr   rZ   values)r/   r   r(   r   rq   r   r   r   r*     s   
z_GraphPickleData.__init__r   rH   c                 C   s:   t j|| j| j}i }| jD ]}||||||< q|S r3   )r7   r8   r   r   r   r   rf   )r/   r   rH   r   r   ndr   r   r   rf   !  s
   
z_GraphPickleData.unpickle)r    r!   r"   r7   r8   r   r   r*   r9   rT   rf   r   r   r   r   r     s    r   c                   @   sn   e Zd Zededejjdee	e
egejjf ee
ef f fddZdeddfdd	Zd
edefddZdS )rB   rP   r2   r   c                 C   rd   r3   re   rg   r   r   r   r6   .  s
   z'_TracingContextPickleData.reduce_helpercontextNc                 C   sL   |j | _ |j| _|j| _|j| _|j| _|j| _|j| _|j| _|j| _d S r3   )	module_contextframe_summary_stackloc_in_frameaot_graph_nameparams_flatparams_flat_unwrap_subclassesparams_unwrapped_to_flat_indexoutput_strides#force_unspec_int_unbacked_size_like)r/   r   r   r   r   r*   =  s   z"_TracingContextPickleData.__init__rH   c                 C   sV   t |j}| j|_| j|_| j|_| j|_| j|_| j|_| j|_| j	|_	| j
|_
|S r3   )r   rS   r   r   r   r   r   r   r   r   r   )r/   rH   r   r   r   r   rf   S  s   
z"_TracingContextPickleData.unpickle)r    r!   r"   r[   r&   r7   rA   r   rZ   r   r
   rT   r+   r6   r*   rf   r   r   r   r   rB   -  s    
rB   )Br   r   rJ   ra   abcr   typingr   r   r   r   r   r   Ztyping_extensionsr	   r
   r7   Ztorch.utils._pytreeutilsZ_pytreer   Ztorch._guardsr   Ztorch._subclasses.fake_tensorr   r   r   Ztorch._subclasses.meta_utilsr   r   r   ry   r   Z%torch.fx.experimental.symbolic_shapesr   Ztorch.utils._mode_utilsr   r?   ZSymFloatr   r#   r$   r   	dataclassr   Picklerr&   rT   r,   r+   	UnpicklerrU   r>   r@   r5   rD   r:   r   r=   r   r   r   r   r   r   rB   r   r   r   r   <module>   sJ     [
>,(E


 