o
    h                  	   @   s  d dl Z d dlZd dlmZmZmZmZmZmZm	Z	m
Z
mZ d dlmZ d dlmZmZmZmZmZ d dlmZmZmZ e
deeeeeef eeef Ze
deeee eeZe
de	eef ee Ze jdeeeef d	ed fd
dZdeegef d	eegef fddZdeeegef d	eeegef fddZ deeegef d	eeegef fddZ!deeegef d	eeegef fddZ"deeegef d	eeegef fddZ#deeeeef gef d	eeeeef gef fddZ$dS )    N)	AnyCallableDictIteratorListOptionalTupleTypeVarUnion)BackendIndexDispatchKeyNativeFunctionNativeFunctionsGroupNativeFunctionsViewGroup)contextSTFF2F3greturnc              	   #   s    t | tr
| j nt | tr| j n|  t fdd- tj j j	d d V  W d    n1 s5w   Y  W d    d S W d    d S 1 sMw   Y  d S )Nc                      s   d j  d j S )Nzin native_functions.yaml line z:
  )locfunc fr   F/var/www/html/ai/venv/lib/python3.10/site-packages/torchgen/context.py<lambda>6   s    z)native_function_manager.<locals>.<lambda>)!use_const_ref_for_mutable_tensorsuse_ilistref_for_tensor_lists)

isinstancer   outr   viewr   localparametrizer   part_of_structured_group)r   r   r   r   native_function_manager'   s    

"r'   r   c                    s$   t  dtdtf fdd}|S )Nr   r   c                    s4   t |   | W  d    S 1 sw   Y  d S Nr'   r   r   r   r   wrapperD   s   
$z%with_native_function.<locals>.wrapper)	functoolswrapsr   r   r   r+   r   r*   r   with_native_functionC   s   r/   c                    (   t  dtdtdtf fdd}|S )Nr   f2r   c                    6   t |   | |W  d    S 1 sw   Y  d S r(   r)   )r   r1   r*   r   r   r+   M   s   
$z)with_native_function_and.<locals>.wrapper)r,   r-   r   r   r   r.   r   r*   r   with_native_function_andL   s   r3   c                    r0   )Nslfr   r   c                    s6   t |  | |W  d    S 1 sw   Y  d S r(   r)   r4   r   r*   r   r   r+   W      
$z,method_with_native_function.<locals>.wrapper)r,   r-   r   r   r   r.   r   r*   r   method_with_native_functionV   s   r7   c                    r0   )Nr4   r   r   c                    s:   t |d   | |W  d    S 1 sw   Y  d S )Nr   r)   r5   r*   r   r   r+   b   s   $z3method_with_nested_native_function.<locals>.wrapper)r,   r-   r   r   r   r.   r   r*   r   "method_with_nested_native_function_      r8   c                    r0   )Nr   backend_indexr   c                    r2   r(   r)   )r   r:   r*   r   r   r+   o   r6   z/with_native_function_and_index.<locals>.wrapper)r,   r-   r   r   r   r.   r   r*   r   with_native_function_and_indexl   r9   r;   c                    s0   t  dtdtttf dtf fdd}|S )Nr   backend_indicesr   c                    r2   r(   r)   )r   r<   r*   r   r   r+   {   r6   z1with_native_function_and_indices.<locals>.wrapper)r,   r-   r   r   r   r   r   r.   r   r*   r    with_native_function_and_indicesx   s   $r=   )%
contextlibr,   typingr   r   r   r   r   r   r   r	   r
   torchgen.localr$   torchgen.modelr   r   r   r   r   torchgen.utilsr   r   r   r   boolstrr   r   contextmanagerr'   r/   r3   r7   r8   r;   r=   r   r   r   r   <module>   s^    ,

		&*	*
	

