U
    Mh                  	   @   s  d dl mZmZ d dlm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mZmZmZmZmZ edede
eededeejddddd	Zed
ed
e
eededeejddddd	Zededeededeejddddd	Zededee	dedeejddddd	ZeddZededeededeejddddd	Zddeeeee edddZeeedddZ eeee dddZ!eedddZ"eee dddZ#eee dd d!Z$eeee dd"d#Z%dS )$    )ListOptional)
dispatcher)	BaseCppType	BaseCTypeBindingboolTConstRefCTypeCTypelongT
NamedCTypetensorT)ArgumentBaseTyBaseTypeFunctionSchemaNativeFunctionNativeFunctionsViewGroupbase)nametypeN)r   r   default
annotation)r   Znctypeargumentr   Zmutated_viewZmutated_view_idxreapply_viewszat::functionalizationZInverseReturnModeZinverse_return_mode)r   )g
is_reverseinclude_namespacer   returnc                C   sb   |d kr|st |r t| j|S |s(t | jd k	s6t |rH| jjj n| jjj }d| dS )Nz
at::_ops::z::call)AssertionErrorreverse_nameviewZ	view_copyfuncr   unambiguous_name)r   r   r   r   api_name r%   P/var/www/html/venv/lib/python3.8/site-packages/torchgen/api/functionalization.pyr   [   s    r   )fr   r   c                 C   s*   | j j }|rd| dS | dS d S )Nz+at::functionalization::FunctionalInverses::Z_inverse)r"   r   r#   )r'   r   r$   r%   r%   r&   r    s   s    r    )r"   r   r   c                C   sV   | j j}|d jttjks t|dd  }dd |D }|rBtntg}|	| |S )Nr      c                 S   s   g | ]}t j|d dqS )T)Zremove_non_owning_ref_typesr   r   .0ar%   r%   r&   
<listcomp>   s    z%capture_arguments.<locals>.<listcomp>)
	argumentsflat_allr   r   r   Tensorr   inverse_return_mode_bindingreapply_views_bindingextend)r"   r   argsnon_self_argsZnon_self_value_bindingsZall_bindingsr%   r%   r&   capture_arguments   s    

r6   )r"   r   c                 C   s4   t | jdkst| jD ]}|j stqttS )Nr(   )lenreturnsr   r   Zis_tensor_liker   r   )r"   retr%   r%   r&   returns_type   s    
r:   )r   r   c                 C   s   | rt ttgS t tgS d S )N)base_bindingmutated_view_bindingmutated_view_idx_binding)r   r%   r%   r&   outer_arguments   s    
r>   c                 C   s4   t | jdks,t | jdkr0| jd j r0tS d S )Nr(   r   )r7   r8   r   Zis_list_liker=   )r"   r%   r%   r&   inner_call_index   s    r?   c                 C   sz   | j j}|d jttjks t|dd  }dd |D }|sHtg| S t| }|d k	rhtt	t
|g| S tt	t
g| S d S )Nr   r(   c                 S   s   g | ]}t |qS r%   r)   r*   r%   r%   r&   r-      s     z#inner_arguments.<locals>.<listcomp>)r.   r/   r   r   r   r0   r   r;   r?   r<   r1   )r"   r   r4   r5   Znon_self_bindingsZindex_bindingr%   r%   r&   inner_arguments   s*    
r@   )&typingr   r   Ztorchgen.apir   Ztorchgen.api.typesr   r   r   r   r	   r
   r   r   r   Ztorchgen.modelr   r   r   r   r   r   r0   r;   r<   r=   boolr2   ZInverseReturnModeTr1   strr   r    r6   r:   r>   r?   r@   r%   r%   r%   r&   <module>   s   ,             	


