
    '}hns                     T   d dl mZ d dlmZmZmZmZmZmZ d dl	m
c mZ d dlmZ d dlmZ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mZmZm Z m!Z!m"Z"m#Z#m$Z$m%Z% d dl&m'Z' d	d
gZ(ddgZ)g dZ*dgZ+dee    deeee$e f   f   fdZ,dee-   de-fdZ.dedefdZ/dedefdZ0dedeee#   ee   f   fdZ1dedefdZ2de de$dee eeed   f   f   fdZ3dee    deeee"ef   f   ddfdZ4d ee#d!f   d"ee-   de-fd#Z5ded$e-dee-   fd%Z6ed&e!dee-   fd'       Z7ed&e!dee-   fd(       Z8y))    )defaultdict)DictListOptionalSequenceTupleUnionN)	translate)BindingDispatcherSignatureExpr)with_native_function)
AnnotationArgumentBackendIndexBackendMetadataBaseOperatorNameBaseTyBaseTypeDEFAULT_KERNEL_NAMESPACEDeviceCheckTypeDispatchKeyFunctionSchemaNativeFunctionNativeFunctionsGroupOperatorNameReturn
SchemaKindVariant)	concatMapz'adaptive_avg_pool3d_backward.grad_inputz _slow_conv2d_backward.grad_input_cummax_helper_cummin_helper)"_assert_asyncz_assert_async.msg_cslt_sparse_mm_search_assert_scalar_dimI_dimV_has_same_storage_numel_linalg_check_errors_local_scalar_dense%_nested_tensor_from_mask_left_aligned_nnz_use_cudnn_ctc_lossz_use_cudnn_ctc_loss.Tensor#_validate_compressed_sparse_indicesallclose	dense_dimequalis_coalesced	is_pinnedis_same_size	is_set_toq_per_channel_axisq_scaleq_zero_pointqschemerecord_stream
sparse_dimsym_constrain_rangesym_constrain_range_for_size_nested_tensor_storage_offsets'_chunk_grad_outputs_efficient_attention_fused_sdp_choice_print_nested_get_ragged_idx
polygamma_native_functionsreturnc                     t        t              }| D ]Z  }||j                  j                            }|j                  j	                         |vsJ |||j                  j	                         <   \ |S N)r   dictfunc	signaturekind)rD   pre_grouped_native_functionsfds       b/var/www/html/test/engine/venv/lib/python3.12/site-packages/torchgen/native_function_generation.pypre_group_native_functionsrP   e   sm    
 	D !  ()9)9);<vv{{}A%%%!&&++- ('    overload_namec                     | sdS |  dS )Nout_out )rR   s    rO   &get_expected_out_variant_overload_namerW   s   s    %5Am_D+AArQ   rI   c                 :   | j                         t        j                  k(  sJ | j                  j                  J t        | j                  j                         j                  t        | j                  j                              | j                  j                         j                  t        d| j                  j                  j                  j                  d | j                  j                  j                  j                         g      | j"                        S )NrT   nametypedefault
annotationrZ   	argumentsreturns)rK   r   inplacer_   self_argr   rZ   remove_inplacewith_overloadrW   rR   remove_self_annotationwith_out_argsr   argumentr[   r]   r`   )rI   s    rO   self_to_out_signaturerh   |   s    99;*,,,,,>>""...
 YY%%'5524993J3JK
 ..779GG0099>> #~~66??JJ		
  rQ   c                 L   | j                         t        j                  k(  sJ t        |       \  }}t	        | j
                  j                  t        | j
                  j                              | j                  j                         j                  |      t        |            S Nr^   )rK   r   
functionalgenerate_out_args_from_schemar   rZ   rd   rW   rR   r_   rJ   rf   tuplerI   new_returnsnew_out_argss      rO   functional_to_out_signaturerq      s    99;*///// =d CK
 YY$$24993J3JK
 ..**,::
 k" rQ   c           
      J   t        d | j                  D              rJ | j                  D cg c]  }|j                  j                         s|! }}t	        |      dkD  sJ t        d | j                  j                        }dD cg c]	  }||vs| }}t        d | j                  D              }g }g }t        | j                        D ]  \  }	}|j                  j                         rt        t	        | j                        dk(  rdnd|	 |j                  d t        j                  ||	    d      	      }
|j                  |
       |st        d |
j                  |
j                  
      }|j                  |       |j                  |        ||fS c c}w c c}w )Nc              3   h   K   | ]*  }|j                   d uxr |j                   j                   , y wrG   )r]   is_write.0rs     rO   	<genexpr>z0generate_out_args_from_schema.<locals>.<genexpr>   s0      ?@D :Q\\%:%::   02r   c                 J    | j                   g S | j                   j                  S rG   )r]   	alias_set)as    rO   <lambda>z/generate_out_args_from_schema.<locals>.<lambda>   s    ," !,,2H2H rQ   abcdefghijklmnopqrstuvwxyzc              3   h   K   | ]*  }|j                   t        t        j                        k(   , y wrG   )r[   r   r   Tensorru   s     rO   rx   z0generate_out_args_from_schema.<locals>.<genexpr>   s"     WQqvv&--)@@Wry      rT   !rY   )rZ   r[   r]   )anyr`   r[   is_tensor_likelenr    r_   flat_allall	enumerater   r   parseappendr   r]   )rI   rw   tensorlike_retsused_annotationsxvalid_annotationsall_rets_are_tensorsrp   ro   inew_outnew_rets               rO   rl   rl      s   
  DHLL    #',,JQ!&&2G2G2IqJOJ!### H
 01<L3L  W$,,WW#%L !#K$,,' "166  "!$,,/14UCs)VV%++/@/C.DA,FG	G (# !GLLW=O=O ""7+q!#"$ $$I Ks   FF	F F c                 L   | j                         t        j                  k(  sJ t        |       \  }}t	        | j
                  j                         j                  t        | j
                  j                              | j                  j                  |      t        |            S rj   )rK   r   mutablerl   r   rZ   rc   rd   rW   rR   r_   rf   rm   rn   s      rO   mutable_to_out_signaturer      s    99;*,,,,, !>d CKYY%%'5524993J3JK
 ....|<k" rQ   rM   k)r   r   c                    ddl m} |t        j                  k(  r| j                  j                         t        j                  k7  sJ | j                  j                  d      j                  t        t        | j                  j                  j                  j                  d| j                  j                  j                  j                  | j                  j                         t        j                  k(        | j                  j                  j                              }n|t        j                  k(  r| j                  j                         t        j                   k(  rt#        | j                        }n| j                  j                         t        j                  k(  rt%        | j                        }nW| j                  j                         t        j                  k(  rt'        | j                        }nt)        d      t)        d	      |j                         t        j                  k(  r|j                  j+                         n|j                  |      }| j                  j-                         r|d
z  }t.        j0                  |j                  t3        |dt4              ii}dht7        | j8                  h dz        z  }t;        d&i d|d| j<                  dt>        j@                  hdddd dd dd dg di dddddd dd dddtB        jD                  d| jF                  dt7               d| jH                  d dd!dd"dd#dd$|d%| jJ                  |fS )'Nr   )cppT)keep_return_namesF)basera   dunder_methodfunctional_overload)rZ   rR   z^We only bother generating out= functions from either inplace or mutable or functional variantszDWe currently only generate either functional or out= NativeFunctions_symint)kernel
structuredcpp_namespace	generated>   	view_copypt2_compliant_tagnondeterministic_seededrI   !use_const_ref_for_mutable_tensorsvariantsr   structured_delegatestructured_inheritsprecomputedautogenufunc_inner_loopmanual_kernel_registrationmanual_cpp_bindingpython_modulecategory_overridedevice_guarddevice_checkloccpp_no_default_argsis_abstract&has_composite_implicit_autograd_kernel4has_composite_implicit_autograd_nested_tensor_kernel&has_composite_explicit_autograd_kernel5has_composite_explicit_autograd_non_functional_kerneltags	namespacerV   )&torchgen.apir   r   rk   rI   rK   rJ   	with_namer   r   rZ   r   r   r   rR   rT   ra   rh   r   rq   AssertionErrorunambiguous_name
has_symintr   CompositeExplicitAutogradr   r   setr   r   r   r   functionr   NoCheckr   r   r   )rM   r   r   rI   kernel_namebackend_metadatar   s          rO   generate_functionr   
  s4    !J!!!vv{{}
 5 5555 vv$7AA%))..!"#&&++"2"2"@"@()9K9K(K  ffkk77	
 
jnn	 66;;=J...(0DVV[[]j000+AFF3DVV[[]j333.qvv6D p  R
 	
 99;*..( 			""$XXd^ 
 	vvy --II" 60
 =3	NN D
 	 	
	
./.Q.Q	
 &&'		

 	
 !%	
 !%	
 	
 	
  	
 (-	
  %	
 	
 #	
 	
  )00!	
" #	
$ !$%	
& '	
( 49)	
* BG+	
, 48-	
. CH/	
4 5	
6 kk7	
: 	= rQ   rsindicesc                 (   t        |       }|j                         D ]s  }t        j                  |v }t        j                  |v }t        j
                  |v }t        j                  |v }|s|s|s|sUt        d |j                         D              }t        d |j                         D              }	t        d |j                         D              }
|s|	s|
r|rt        |j                               dk(  rjt        |t        j                     j                  j                        t        vr2t        dt        |t        j                     j                               D|r8t        |t        j                     j                  j                        t        v r~|r|t        j                     n<|r|t        j
                     n'|r|t        j                     n|t        j                     }|j                  j!                         t        j                  k(  xs& t        d |j                  j"                  D              }t        d |j$                  D              }| xr |xr |}|sn|slt        |j                  j                        t&        vrGt        |j                  j                        t(        vr"t        dt        |j                         d	      |rWt+        |t        j                        \  }}||t        j                  <   t-        j.                  ||       | j1                  |       |r|s|st+        |t        j                        \  }}||t        j                  <   t-        j.                  ||       | j1                  |       v y )
Nc              3   4   K   | ]  }|j                     y wrG   )r   rv   rM   s     rO   rx   z1add_generated_native_functions.<locals>.<genexpr>  s     FaQ11F   c              3      K   | ]>  }|j                   xr, t        |j                  j                  j                        d k7   @ yw)set_N)
is_view_opstrrI   rZ   r   s     rO   rx   z1add_generated_native_functions.<locals>.<genexpr>  s9      EF@QVV[[%5%5!6&!@@s   AAc              3   4   K   | ]  }|j                     y wrG   )r   r   s     rO   rx   z1add_generated_native_functions.<locals>.<genexpr>  s      )=>88)r   r   zEFound an out= operator that we could not find any other variants of: c              3   P   K   | ]  }|j                   j                            y wrG   )r[   r   ru   s     rO   rx   z1add_generated_native_functions.<locals>.<genexpr>  s$      M,-%%'Ms   $&c              3   6   K   | ]  }d t        |      v   yw)rT   N)r   )rv   op_names     rO   rx   z1add_generated_native_functions.<locals>.<genexpr>  s     QgES\1Qs   zBFound an operator that we could not generate an out= variant for: z.
This type of operators don't have tensor-like return, making it difficult to generate a proper out= variant. If
out= variant is not needed, please add the function name into FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT list.)rP   valuesr   rk   ra   r   rT   r   r   r   r   rI   rZ   &OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLYr   *INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLYrK   r`   r   *MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT-FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANTr   r   
grow_indexr   )r   r   rL   rN   has_functionalhas_inplacehas_mutablehas_out
are_manualhas_view_opsare_composite_implicitbase_fnbase_fn_valid	needs_outgets_out_variantfnmetadatas                    rO   add_generated_native_functionsr   x  s0    $>b#A )002 p#..!3 ((A- ((A-..A% +NF188:FFJ  JK((* L &) )BC((*) &" \-C3qxxz?a/ *..)..334AB )_`cdefpftftdudzdz`{_|}  
 *,,-22778=>   *$$%  z))*  z~~&z,,- " $LL--/:3E3EE  M181E1EM JM QQQI#*{J}J=))*EFGLL--.HI )^_bcjcoco_p^q rr u   0*..IH$&*..!'':		"
 "w2B0*:O:OPH+-*''('':		"aprQ   rets.namesc                     t        |       t        |      k(  sJ t        |       dk(  ryt        |       dk(  r	d|d    dS dt        j                  |       j                          ddj	                  |       dS )	Nr    r   zreturn ;(, );)r   
dispatcherreturns_typecpp_typejoin)r   r   s     rO   
return_strr     sy    t9E
"""
4yA~	Taq
!$$006??AB!DIIeDTCUUWXXrQ   out_varc                     | j                         }g }t        | j                        dkD  }t        |      D ]$  \  }}|	|j	                  |r	d| d| dn|       & |S )Nr   	std::get<>())aliased_return_namesr   r`   r   r   )rI   r   aliased_retsnon_aliased_namesis_out_var_a_tupler   rw   s          rO   gather_nonaliased_inner_retsr     ss    ,,.LT\\*Q.,' 19$$/A)A3b	+w
 rQ   gc                    d| j                   j                  vry | j                  %d| j                  j                  vr| j                  }nY| j                  %d| j                  j                  vr| j                  }n(t	        t        | j                   j                              t        | j                   j                        }t        |j                        }g }g }g }t        t        j                  | j                   j                        t        j                  |j                              D ]  \  }}|j                  |j                  j                  r|j                  d|j                   d|j                   d       |j                  t        |j                   dt        j                   ||j                                     |j                  |j                   d       |j                  t        j"                  |              dj%                  t'        ||j)                               D 	cg c]  }	|	j*                   c}	      }
d	}t-        |j                  j.                        d
kD  rd| dnd}t1        |j                  |      }t3        | j                   j                  j.                  ||z         }dj%                  |      }d|j5                  |j                         | j6                  j                  j9                         rdndz          d| d| d|j                  j                  j;                          d|
 d| dS c c}	w )Nr   zauto z_clone = clone_arg(r   _clone)binds)exprr[   r   outputr   z = r   
r   rZ   z {
  
  z
at::_ops::::call();
  
}
)rk   r   ra   r   r   r   rI   r   zipr   jit_argumentsr]   rt   r   rZ   r   argument_typerg   r   r
   r_   r  r   r`   r   r   defnrT   r   r   )r   target_fsig
target_sigcontextclone_mutable_inputscloned_return_namesa_curra_tgteexprsout_namemaybe_assigninner_return_namesret_strclone_mutable_inputs_strs                   rO   gen_composite_functional_kernelr    s    !,,+++yyAIINN!B99	
	;aiinn#D99 S!2!2344
all//
0C$X]]3J*,G   !2!23  / 8 'E,<,<,E,E ''}$7}BG NN"KK=/#11&L  &&&++f'=>NN:..v67#8$ IIy*:N:N:P'QR!qvvRSEH,/0E0E,F,JU8*C(PRL5hmmXN	!!#58K#KG  $yy)=>sxxz!%%***?*?*AYrJK L M .
8==-->>@A P
) 	  Ss   6Mc                    d| j                   j                  vry t        | j                   j                        }t        | j                  j                        }dj                  t        |j                         |j                               D cg c]  }|j                   c}      }g }d}t        | j                   j                  j                  j                         D ]n  \  }}t        | j                  j                  j                        dk(  r|nd| d| d}	|j                  d|j                   d|	 d	|j                   d|	 d
	       p g }
t        | j                   j                  j                               D ]a  \  }}||
j                  |       t        | j                  j                  j                        dk(  r|nd| d| d}	|
j                  |	       c dj                  |      }d|j                  | j                   j                  j                  j!                         | j                   j                  j#                         rdndz          d| d| j                  j                  j                  j!                          d| d| dt%        | j                   j                  j                  |
       dS c c}w )Nr   r   
tmp_outputr   r   r   r   z  resize_out_helper(z);
  copy_arg(r   r  r   r   r  z
 {
  auto z = at::_ops::r	  r
  r  r  )rT   r   r   rI   rk   r   r
   r_   r  r   r   r`   r   rZ   r   r  r   r   r   )r   r  r  r  r  	copy_outsr  r   out_argfunctional_return_namer   ret_namecopy_outs_strs                rO   gen_composite_out_kernelr'  L  s    !%%**$ aeejj
)C$Q\\%6%67JII"3==?J4H4H4JKLALE IH

 4 4 8 89 


7 1<<$$,,-2 QCr(1- 	
 	\\N"%;$< =LL>2329	


 D !!@!@!BC 	08KK! q||((001Q6  2hZq1 #
 KK./	0 IIi(Mquuzz//1!%%**BWBWBYY_abc d ezq||0055FFHIQVPW X/ aeejj  $'( )	 G 	Ms   K)9collectionsr   typingr   r   r   r   r   r	   torchgen.api.dispatcherapir   torchgen.api.translater
   torchgen.api.typesr   r   r   torchgen.contextr   torchgen.modelr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   torchgen.utilsr    r   r   r   r   rP   r   rW   rh   rq   rl   r   r   r   r   r   r  r'  rV   rQ   rO   <module>r1     s2   # ? ? , , , A A 1    & % . '* & 	. *#1 -V . *
(~.
(	.$z>9:
:;
(B(3- BC B > Bn  *-%
-%
4<h'(-%j> n Fcc$c
>4T2S-T TUUVc\x^x+tL/$ABBCx 
xvYU63;' YS	 Yc Y	~ 	 	S	 	 9'; 9 9 9| 6 4 6# 6 6rQ   