
    ɯwgz                     (   d dl Z d dlZd dlZ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 d dlZd dlmZ d dlmZ d dlZd dlmZmZmZ d dlmZmZ d dlmZ d dlmZ d dlm Z  d dl!m"Z" d d	lm#Z#m$Z$m%Z% d d
l&m'Z'm(Z( ddl)m*Z*m+Z+ ddlm,Z,m-Z-m.Z.  ej^                  e0      Z1ejd                  jf                  Z3ejd                  jh                  Z4ejd                  jj                  Z5ejd                  jl                  Z6ejd                  jn                  Z7 eg e3jp                  e3jr                  e3jt                  e3jv                  e3jx                  e3jz                  e3j|                  e3j~                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e e3j                  e3j                  e3j                  e3j                  j                  e5j                  e6j                        Z^i  e       e^Z_e3j                  e3j                  e3j                  e3j                  j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  e3j                  j                  e3j                  e3j                  e3j                  gZp ee_ep       deeej                  j                  ej                  j                  f      dedef   fdZt ete3j                  j                  g      dej                  dewddfd       Zx ete3j                  j                  g      dej                  dewddfd       Zz ete3j                  j                  g      ddddej                  de	ej                  j                     de	ej                  j                     ddfd       Z{ ete3j                  g      e	 	 ddej                  de	ej                  j                     de	ej                  j                     dej                  fd              Z ete3j                   g      deeeej                  f      dej                  j                  dedej                  fd        Z ete3j                  j                  g      deeeej                  f      d!ee   dedej                  fd"       Z ete3j                  g      d#ej                  d$ej                  d%ej                  d&ee   d'eeee   f   d(eeee   f   d)eeee   f   d*ed+ee   d,ed-ee   de
ej                  ej                  ej                  f   fd.       Z ete3j
                  j                  g      ddej                  d/edej                  fd0       Z ete3j                  g      ed1ej                  d2ej                  dej                  fd3              Z ete3j                  g      e	 	 dd1ej                  d4ej                  d5ej                  d6ej                  j                  d7ej                  j                  dej                  fd8              Z ete3j                  g      ed1ej                  d9ej                  dej                  fd:              Z ete3j                  j                  g      	 dd;eej                     d<edej                  fd=       Z ete3j                  g      dej                  dej                  fd>       Z ete3j                  g      dd?dej                  d@ej                  d7e	ej                  j                     dej                  fdA       Z ete3j                  g      d1ej                  dej                  fdB       Z ete3j                  e3j                   g      d1ej                  dej                  fdC       Z ete3j"                  j                  g      ddDd1ej                  dEe	ej$                     dej                  fdF       Z ete3j&                  e4j&                  g      d1ej                  dGej                  dej                  fdH       Z ete3j(                  e4j(                  g      d1ej                  dGej                  dej                  fdI       Z ete3j*                        	 	 dd1ej                  d<e	e   dKedej                  fdL       Z ete3j,                        	 	 dd1ej                  d<e	e   dKedej                  fdM       Z ete3j.                  g      d1ej                  d<edNedOedej                  f
dP       Z ete3j0                  j                  g      d1ej                  deeeej                  f      dej                  fdQ       Z ete3j0                  j4                  g      d1ej                  dRej4                  dej                  fdS       Z	 ddej                  dTe	ej8                     dej8                  fdUZ ete3j<                        ddddVd1ej                  dRe	ej4                     dWe	ej>                     dTe	ej8                     dedej                  fdX       Z ete3j@                        ddddVd1ej                  dRe	ej4                     dWe	ej>                     dTe	ej8                     dedej                  fdY       Z ete3jB                        ddddJdJejD                  dZd1ej                  deeef   dRe	ej4                     d[e	ejH                     dWe	ej>                     d\ed]edTej8                  dej                  fd^       Z ete3jJ                  j                        ddddVd1ej                  d_edRe	ej4                     dWe	ej>                     dTe	ej8                     dedej                  fd`       Z ete3jJ                  jL                        ddddVd1ej                  daed_edRe	ej4                     dWe	ej>                     dTe	ej8                     dedej                  fdb       Z ete3jP                  j                        d_edeeeej                  f      dedej                  fdc       Z ete5j                  j                        d$ej                  d%ej                  ddej                  dej                  fde       Z\ ete6j                  j                        d$ej                  dfej                  dgej                  d%ej                  dhej                  diej                  ddej                  djej                  dkej                  dledej                  fdm       Z] etejd                  jj                  jR                        dnej                  dej                  fdo       Z ete3jV                  g      e	 	 	 ddpej                  dqej                  dredsedtedej                  fdu              Z ete3jX                  jZ                        	 dd1eej                     dveej                     dweej                     dxedeej                     f
dy       Z ete3j^                  jZ                        	 dd1eej                     dveej                     dweej                     dxedeej                     f
dz       Z ete3jb                  jZ                        d{eej                     d|eej                     d%ej                  j                  deej                     fd}       Ze3jf                  j                  ji                  ejj                  jl                  jn                         ete3jf                        d$ej                  d%ej                  ddej                  ej                     d~ej                  ej                     dej                  ej                     dededede
ej                  ej                  ej                  f   fd              Z e jp                  d      deeedef   f   fd       Zdeeedef   f   fdZ ete3jv                        d1ej                  dej                  dej                  dej                  fd       Z ete7jx                  jz                        d$ej                  dedededRej4                  de
ej                  ej                  f   fd       Z ete3j~                        	 dd1ej                  dej                  dej                  dedej                  f
d       Z ete3j                        	 dd1ej                  dej                  dej                  dedej                  f
d       Z ete3j                  j                        ed#ej                  dej                  d<edej4                  dej                  f
d              Ze ete3j                        ddd1ej                  d<edej                  dej                  dewdedej                  fd       Z ete3j                        	 	 	 	 ddej                  dee   d'e	eeee   f      d(eeee   f   d)eeee   f   dede
ej                  ej                  f   fd       Zy)    N)AnyCallableDictListOptionalTupleUnion)core_aten_decompositionsget_decompositionsremove_decompositions)_grid_sampler_2dpw_cast_for_opmath)extra_random_decomps)counters)	out_dtype)pad_listlike)elementwise_dtypesELEMENTWISE_TYPE_PROMOTION_KINDtype_to_dtype)definitely_trueguard_size_oblivious   )configinductor_prims)is_gpu,needs_fallback_due_to_atomic_add_limitationsuse_scatter_fallbackopsreturn.c                     t        |       r| gn| D ]!  }|t        v st        j                  d|        # t	        j
                  | t              S )Nzduplicate decomp: %s)callabledecompositionslogwarningdecompregister_decomposition)r   ops     b/home/mcse/projects/flask/flask-venv/lib/python3.12/site-packages/torch/_inductor/decomposition.pyr&   r&   r   sI      }se# 5KK.45 ((n==    tensormsgc                      y N r*   r+   s     r(   assert_async_msg_decompr0   }       
r)   c                      y r-   r.   r/   s     r(   "functional_assert_async_msg_decompr3      r1   r)   )minmaxsymbolr4   r5   c                     y r-   r.   )r6   r4   r5   s      r(   sym_constrain_range_for_sizer8      s     r)   xc                 R    || j                  |      } || j                  |      } | S r-   )	clamp_min	clamp_max)r9   r4   r5   s      r(   clampr=      s/     KK
KKHr)   size
fill_valuekwargsc                     |j                  d      }|.t        t        |            |d<   t        j                  | |fi |S t
        S )Ndtype)getr   typetorchfullNotImplemented)r>   r?   r@   rB   s       r(   rF   rF      sF     JJwE}'Z(89wzz$
5f55r)   physical_layoutc                     dgt        |       z  }t        |      D ]
  \  }}|||<    t        j                  |D cg c]  }| |   	 c}fi |j	                  |      S c c}w Nr   )len	enumeraterE   emptypermute)r>   rH   r@   permpls         r(   empty_permutedrR      se     3T?D/* 1Q;;9AQ9DVDLLTRR9s   Agrad_outputinputweight
bias_sizesstridepaddingdilation
transposedoutput_paddinggroupsoutput_maskc                 .   |
d   rt        | j                  j                        st        S t        j                  | dgt        t        d| j                                     z         }t        j                  | |||||||||	|
d   |
d   dg      \  }}}|||fS )N   r   r   F)
r   devicerD   rG   atensumlistrangedimconvolution_backward)rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r]   	grad_biasgrad_inpgrad_weight_s                  r(   rf   rf      s     q>(:(:(?(?!@qcDq+//:K1L,M&MNI#88	QQ/ Hk1 k9--r)   decimalsc                 H    d|z  }t         j                  | |z        d|z  z  S )Ng      $@g      ?)ra   round)r9   rk   ten_pow_decimalss      r(   	round_decro      s,    X~::a**+s5E/EFFr)   selfbatch2c                 b   t         j                  rkt        | j                  d   dk(        st        |j                  d   dk(        r5| j	                  d      |j	                  d      z  j                  d      }|S | j                  j                  dk(  rt        | j                  d      dk(        ryt        |j                  d      dk(        r\t        d   dxx   dz  cc<   t        j
                  | j                  d      |j                  d      z  dd	      j	                  d      S t        S )
Nr   r_   re   cpuinductordecompose_bmmTre   keepdim)r   coordinate_descent_tuningr   shape	unsqueezerb   r`   rD   r>   r   rE   squeezerG   )rp   rq   outs      r(   bmmr      s    ''

1 237KLLOq 8
 >>"%(8(8(;;@@Q@GCJ{{5 		! 127KKKOq 8
 Z 1Q6199Q&.."44!Til r)   mat1mat2betaalphac                    | j                   j                  dk(  r7t        |j                  d      dk(        rt        |j                  d      dk(        rgt        d   dxx   dz  cc<   t        j                  |j                  d      |j                  d      z  dd      j                  d      }||z  || z  z   S t        |j                  d      dk(        ryt        |j                  d      d	k        r\t        |j                  d      d	k        r?t        d   dxx   dz  cc<   |j                  |z  j                  dd      }||z  || z  z   S t        S )
Nru   r   r   rs   rv   decompose_addmmTrx      )r`   rD   r   r>   r   rE   rb   r}   r|   r   TrG   )rp   r   r   r   r   r~   s         r(   addmmr      s2    {{5 		! 127KIIbMQ8
 Z !23q83))Q$,,r"224il  3;,, 1!23		! 23		! 23Z !23q8366D=%%!T%:C3;,,r)   input2c           	      j   t         j                  rit        | j                  d   dk(        st        |j                  d   dk(        r3| j	                  d      |j	                  d      z  j                  d      S | j                  j                  dk(  rt        | j                  d      dk(        rt        | j                  d      dkD        rt        |j                  d      dk(        r| j                  |j                  k(  rt        t        j                  |       t        j                  |      z   dk        rYt        d   d	xx   dz  cc<   t        j                  t        | j                  d            D cg c]  }| |d d f   |z   c}      S t        | j                  d      dk(        ryt        |j                  d      dk(        r\t        d   d	xx   dz  cc<   t        j
                  | j!                  d      |j!                  d      z  dd
      j	                  d      S t"        S c c}w )Nr   r   r_   rt   ru   rs       rv   decompose_mmTrx   )r   rz   r   r{   r|   rb   r`   rD   r>   rB   r   rE   numelr   catrd   r}   rG   )rp   r   is      r(   mmr     s    ''

1 237KLLOq 8
 NN1%(8(8(;;@@Q@GG{{5  2!!34$TYYq\A%56$V[[^q%89v||+T!2U[[5H!HR OPZ 0A509951;NOad1a4j61OPP		! 127KKKOq 8
 Z 0A5099Q&.."44!Til  Ps   :H0tensorsre   c                    ddl m dt        j                  dt        ffd}t        t        ||             t              dk(  rd   j                         S dt              cxk  rt        |       k  r#n n t        j                  j                        S t              dkD  rt        fdD              rd   }t        |j                        }dk  rt        |j                        z   n|j                  t                      |j                        j                   | j#                  dz         j                         S t$        S )Nr   )r   r9   r   c                     t        | j                        dk(  r | j                  d   dk(        ryt        | j                        k  r | j                     dk(        ryy)Nr   r   FT)rK   r{   )r9   re   r   s    r(   non_empty_tensorzcat.<locals>.non_empty_tensorD  sW    " qww<1!5aggajAo!FQWW"6qwws|q7H"Ir)   r   c              3   ,   K   | ]  }|d    u   yw)r   Nr.   ).0tfiltered_tensorss     r(   	<genexpr>zcat.<locals>.<genexpr>f  s!      )%&a  )s   )%torch.fx.experimental.symbolic_shapesr   rE   Tensorboolrc   filterrK   clonera   r   defaultallr{   insertr|   expandflattenrG   )r   re   r   inpr{   r   r   s    `   @@r(   r   r   =  s/   
 KELL T 2 F#3W=>
!"((**	
S!"	1S\	1xx 0#66 q S )*:) & q!SYY&)AgcC		N"3S#./0(s}}S!((%088cAgFLLNN r)   c                 "   | j                         ret        j                  t        j                  | j                        t        d      t        j                  | j                  | j                              S t        | t        j                        \  }}t        j                  t        j                  || j                        }t        j                  | dk  |d      }t        j                  t        j                  |       t        d      |      S )Nnan)type_promotion_kindrB   r`   r   g        )
is_complexrE   whereisnanrealfloatatan2imagr   r   INT_TO_FLOATscalar_tensormathpir`   )r9   rj   rB   r   rets        r(   angler   s  s    ||~{{KKuu{{166166/J
 	
 "	;HHHAu 
		TWWE!((	CB
++a!eR
%C;;u{{1~uU|S99r)   r   yc                b   t        j                  |       xr | j                         }t        j                  |      xr |j                         }|r|st        S |}|||z  }t        j                  | j
                  |j
                        }dt         j                  dt         j                  fd} || j                  | j                  j
                              } ||j                  |j                  j
                              }	t        j                  ||	z   d      j                  |      }
|
S )Nr*   r   c                     | j                   ^ }}|dz  dk7  rt        d      g ||dz  d}| j                  |      }|S )zNReshape tensor from [*initial_dims, last_dim] to *initial_dims, last_dim/2, 2]r_   r   zQThe size of the last dimension must be even to reshape it to [..., last_dim/2, 2])r{   AssertionErrorview)r*   initial_dimslast_dim	new_shapereshaped_tensors        r(   reshape_tensor_complexz#add.<locals>.reshape_tensor_complex  s^     #),,x a<1 c 
 6l5HM515	 ++i0r)   )	start_dim)
rE   	is_tensorr   rG   promote_typesrB   r   r   r   r   )r9   r   r   x_is_complex_tensory_is_complex_tensorzcomplex_typer   
x_reshaped
z_reshapedresults              r(   addr     s      //!,?//!,?&9	AAI&&qww8Lu||  " (qvv||(<=J'qvv||(<=J]]:
2bAFF|TFMr)   c                 4    | j                         rJ d       | S )NzTODO: implement this)r   rp   s    r(   conj_physicalr     s     8"88 Kr)   c                     | S r-   r.   r   s    r(   liftr     s    Kr)   )	generatorr   c                    |J t        j                  | t         j                        | k  j                  | j                        S )N)rB   )rE   	rand_likefloat32torB   )rp   r   s     r(   	bernoullir     s9     OOD6=AA$**MMr)   otherc                 b    t        j                  t        j                  |      || kD  z  | |      S r-   rE   r   r   rp   r   s     r(   fminr     '    ;;u{{5)UT\:D%HHr)   c                 b    t        j                  t        j                  |      || k  z  | |      S r-   r   r   s     r(   fmaxr     r   r)   Fry   c                 x    | j                   t        j                  k(  rt        j                  | ||      S t        S Nrx   )rB   rE   r   anyrG   rp   re   ry   s      r(   amaxr     -     zzUZZyy388r)   c                 x    | j                   t        j                  k(  rt        j                  | ||      S t        S r   )rB   rE   r   r   rG   r   s      r(   aminr     r   r)   startlengthc                 N    t        j                  | |||      j                         S r-   )rE   narrowr   )rp   re   r   r   s       r(   narrow_copyr     s"     <<c5&17799r)   c                 J    t         j                  | |      j                         S r-   )ra   r   r   )rp   r>   s     r(   view_copy_defaultr     s    
 99T4 &&((r)   rB   c                 @    | j                  |      j                         S r-   )r   r   )rp   rB   s     r(   view_copy_dtyper     s    
 775>!!r)   memory_formatc                 X    |t         j                  u s|t        j                  |       S |S r-   )rE   preserve_formatutilssuggest_memory_format)r*   r   s     r(   get_like_layoutr      s.    
 ---1F**622r)   )rB   r`   r   r`   c                    t        j                  g | j                         f|xs | j                  |xs | j                  d|j                  t        | |            S Nr   r   )rE   randr>   rB   r`   r   r   rp   rB   r`   r   r@   s        r(   r   r     s^     ::$))+!tzz$ 	
 	bt];b<=r)   c                    t        j                  g | j                         f|xs | j                  |xs | j                  d|j                  t        | |            S r   )rE   randnr>   rB   r`   r   r   r   s        r(   
randn_liker     s^     ;;$))+!tzz$ 	
 	bt];b<=r)   )rB   layoutr`   
pin_memoryrequires_gradr   r   r   r  c                    t        j                  g | j                         ||xs | j                  |xs | j                  |xs | j
                  |      j                  t        | |            S )N)rB   r   r`   r  r   )rE   rF   r>   rB   r   r`   r   r   )rp   r?   rB   r   r`   r   r  r   s           r(   	full_liker  -  sa     ::$))+!tzz$$# 	bt];b<=r)   highc                    t        j                  j                  d|g | j                         f|xs | j                  |xs | j
                  d|j                  t        | |            S )Nr   r   r   ra   randintlowr>   rB   r`   r   r   )rp   r  rB   r`   r   r@   s         r(   randint_liker	  C  sl     <<	$))+ !tzz$  	bt];b<=r)   r  c                    t        j                  j                  ||g | j                         f|xs | j                  |xs | j
                  d|j                  t        | |            S r   r  )rp   r  r  rB   r`   r   r@   s          r(   randint_like_lowr  W  sl     <<$))+ !tzz$  	bt];b<=r)   c                 F    t        j                  j                  d| |fi |S rJ   )ra   r  r  )r  r>   r@   s      r(   r  r  l  s"     <<AtT4V44r)   biasc                     t         j                  j                  j                  |      }t         j                  j                  j	                  | |||j                         d         S rJ   )rE   r   
_quantized$wrapped_fbgemm_pack_gemm_matrix_fp16!wrapped_fbgemm_linear_fp16_weightr>   )rT   rU   r  packed_weights       r(   #linear_dynamic_fp16_unpacked_weightr  u  sO     II((MMfUM99AA}dFKKM!$4 r)   input_scaleinput_zero_pointweight_scaleweight_zero_point	out_scaleout_zero_pointout_channelc
           	          t         j                  j                  j                  ||||      }
t         j                  j                  j	                  | |||
|||	      S r-   )rE   r   r  _wrapped_linear_prepack#_wrapped_quantized_linear_prepacked)rT   r  r  rU   r  r  r  r  r  r  r  s              r(   wrapped_quantized_linearr    s\     II((@@/M 99CC r)   packedc                     dt         j                  dt         j                  fd} || dddf         } || ddd f         }| dd df   j                  t         j                        |z  |z   S )Nu8r   c                      fddD        \  }}}}t         j                  dk(  r4||dz  z   |dz  z   |dz  z   j                  t        j                        d   S |dz  |dz  z   |dz  z   |z   j                  t        j                        d   S )Nc              3   f   K   | ](  }d |f   j                  t        j                         * yw.N)r   rE   int32)r   nr!  s     r(   r   zPq_embedding_bag_byte_unpack_decomp.<locals>.bitcast_u8_to_f32.<locals>.<genexpr>  s&     GQbajmmEKK0Gs   .1)r   r   r_      little   r      r$  )sys	byteorderr   rE   r   )r!  r9   r   r   ws   `    r(   bitcast_u8_to_f32z=q_embedding_bag_byte_unpack_decomp.<locals>.bitcast_u8_to_f32  s    G,G
1a==H$aLAG,R8>>u}}MiXX"Wb)Q!V4q8>>u}}MiXXr)   .i)rE   r   r   r   )r  r.  scalesoffsetss       r(   "q_embedding_bag_byte_unpack_decompr2    sz    Yell Yu|| Y vc2b5j12FsBCx 01G#ss(u}}-6@@r)   agridinterpolation_modepadding_modealign_cornersc                     | j                   t        j                   d      k(  xr' |dk(  xr  | j                  t        j                         }t	        | |||||      }|S )Nru   r   r   )r4  r5  r6  r7  _expand_grid)r`   rE   is_contiguouscontiguous_formatdecomp_grid_sampler_2d)r3  r4  r5  r6  r7  r9  outputs          r(   grid_sampler_2dr>    sm     	
ELL'' 	C!#	COO%*A*AOBL $	-!#!F Mr)   left_tensorsright_tensorsscalarc                     t         j                  j                  | t         j                  j                  ||      |      S Nr   )ra   _foreach_addr   _foreach_mulrp   r?  r@  rA  s       r(   _foreach_addcmul_scalarrG    <     !!d$$\=A "  r)   c                     t         j                  j                  | t         j                  j                  ||      |      S rC  )ra   rD  r   _foreach_divrF  s       r(   _foreach_addcdiv_scalarrK    rH  r)   start_tensorsend_tensorsc           	          t         j                  j                  | t         j                  j	                  t         j
                  j                  ||       |            S r-   )ra   rD  r   rE  Scalar_foreach_sub)rL  rM  rU   s      r(   _foreach_lerp_scalarrQ    sL     !!  "";>	
 r)   running_meanrunning_vartrainingexponential_average_factorepsilonc           
          t         j                  | |||||||      \  }}	}
|r||	|
fS ||j                  d      |j                  d      fS )Nr   )ra   native_batch_norm	new_zeros)rT   rU   r  rR  rS  rT  rU  rV  r3  bcs              r(   miopen_batch_normr]    sl     $$"	GAq! 1ay	 r)   c                      i t         t        S r-   )r"   r   r.   r)   r(   fast_random_decompsr_    s    5n5 455r)   c                  B    t         j                  rt        S t               S )z"decomps can change based on config)r   fallback_randomr"   r_  r.   r)   r(   select_decomp_tablerb    s      r)   masksourcec                    ddl m}m}  || j                  |j                        rt
        j                  | |g      \  } }|j                  d      j                  d      dz
  }d | ||fD        \  }}}t
        j                  |||gd      }	t        j                  ||	|      j                  | j                        S t        S )Nr   )BackendFeaturehas_backend_featurers   r   c              3   <   K   | ]  }|j                           y wr-   )r   )r   r9   s     r(   r   z!masked_scatter.<locals>.<genexpr>*  s     ,WQQYY[,Ws   )codegen.commonrf  rg  r`   MASKED_SCATTER_WITH_INDEXra   broadcast_tensorsreshapecumsum_unsafe_masked_indexrE   r   r   r{   rG   )
rp   rc  rd  rf  rg  
source_idx	self_flat	mask_flatsource_flatr   s
             r(   masked_scatterrs    s     D4;;(P(PQ ++T4L9
d\\"%,,Q/!3
,W4vBV,W)	9k**;	J<QRS{{9fi8==djjIIr)   	quant_min	quant_maxepsc                    t        j                  |       \  }}||z
  t        ||z
        z  }t        j                  |t        j                  |g            }|t        j
                  ||z        j                  t         j                        z
  }t        j                  |||      }|j                  t         j                        |j                  t         j                        fS r-   )rE   aminmaxr   r5   r   rm   r   intr=   float64int64)	rT   rt  ru  rv  rB   min_valmax_valscale
zero_points	            r(   choose_qparams_tensorr  0  s     }}U+GWw%	I(=">>EIIeU\\3%01EU[[59<<UYYGGJZI>J88EMM"JMM%++$>>>r)   index
accumulatec                     | j                         }t        j                  ||g|j                  |j                        |      }|j                  | j                        S r-   )r   rE   	index_putrl  r{   )rp   r  rd  r  	flatteneds        r(   putr  @  sL     IE7FNN5;;7I TZZ((r)   c                 V    t         j                  | |||      }| j                  |      S )N)r  )ra   r  copy_)rp   r  rd  r  r~   s        r(   put_r  N  s(     ((4:(
>C::c?r)   r=  input_dtypec                     | |z  }t        j                  ||d      }t        j                  | ||      }| j                  |k7  r|j                  |      }|j                         S )NTrx   )rE   rb   r   fmarB   r   
contiguous)rS   r=  re   r  new_grad_outputsum_new_grad
grad_inputs          r(   _softmax_backward_datar  Y  sc     "F*O99_#tDL##VG\?KJ K']];/
  ""r)   Tinclude_selfsrcreduction_typer  c                   |dk(  rt        | j                        s| j                  j                  xs | j                  j                  }t	        j
                  |      }|r)| }t	        j
                  |       j                  |||      }	nN| j                  ||d      }t	        j                  |       j                  |||      }	|	j                  |	dk  d      }	|j                  |||      }|r||	z  S ||	z  S t        t        j                  j                  || j                  |j                  |j                  j                  d      rt         S | j"                  |dz   d  j%                         | j"                  d | j%                         z  }
|j%                         g| j"                  |dz   d  | j"                  d | }g t'        | j(                  |z
  | j(                        dt'        d| j(                  |z
        }|j+                  t        j,                        j/                  |
      j1                  |      j3                  |      }| j5                  |||||      S )Nmeanr   r   Tr  )r   rB   is_floating_pointr   rE   	ones_like	index_add
index_fill
zeros_likemasked_fillr   ra   scatter_reduce_twor`   rD   rG   r{   r   rd   ndimr   r{  repeat_interleaverl  rN   scatter_reduce)rp   re   r  r  r  r  true_divisiononesr~   countsrepeatsindex_shaperO   scatter_indexs                 r(   index_reducer  o  s    (T

) 

44M

8M8Ms#C__T*44S%FF//#ua0C%%d+55c5$GF''
A6FmmC,,sV|?#-?  

		

 jjq#))+djj#.>.D.D.FFG;;=L4::cAgi#8L4::ds;KLKNU499s?DII.NNE!TYY_4MND		7	#				  !   r)   kernel_size	ceil_modec                    |dk(  rddg}|dk(  rddg}|s|}t        |d      }t        |d      }t        |d      }t        |d      }|d   |d   z  }t        j                  j                  j	                  ||      s0|t        j
                  t        j                        j                  kD  rt        S t        j                  | |||||      \  }}t        j                  ||d   | j                  d      ||      }	||	fS )Nr   r   r_   rs   )r   rE   	_inductorlowering'should_fallback_max_pool2d_with_indicesiinfoint8r5   rG   prims#_low_memory_max_pool2d_with_offsets)_low_memory_max_pool2d_offsets_to_indicesr>   )
r9   r  rW   rX   rY   r  window_sizevalsr1  indicess
             r(   max_pool2d_with_indicesr    s    1}q6!|a&{A.KHa(H7A&G&!$Fa.;q>1K 	  HH	
 UZZ0444==	MD' ==A	r
G =r)   )NNrX  )r   r   )NFr-   )r   r   F)r   )F)Nr   r   F)	functoolsloggingr   r+  typingr   r   r   r   r   r   r	   rE   torch._decomp_decompr%   torch._prims_common_prims_commonr   $torch.ao.quantization.fx._decomposedr
   r   r   torch._decomp.decompositionsr   r<  r   $torch._decomp.decompositions_for_rngr   torch._dynamo.utilsr   !torch._higher_order_ops.out_dtyper   torch._inductor.utilsr   r   r   r   r   r   r    r   r   r   r   r   	getLogger__name__r#   r   ra   r  	quantizedr  quantized_decomposed_adaptive_avg_pool2d_backwardaddmvarangebitwise_and_bitwise_or_
clamp_min_dist
empty_likeflipgeluhardtanhindex_selectlcm
leaky_relulinalg_vector_norm_log_softmax max_pool2d_with_indices_backward_native_batch_norm_legit#_native_batch_norm_legit_functional$_native_batch_norm_legit_no_training_batch_norm_with_update"_batch_norm_with_update_functional_batch_norm_no_updatebatch_norm_backwardrY  native_group_normnative_layer_normnll_loss2d_backward_softmaxsin_sqrt__to_copytril_indicestriu_indicesupsample_bilinear2dvecr  r  inductor_decompositionsr"   _unsafe_indexrn  #_unsafe_masked_index_put_accumulate+_scaled_dot_product_flash_attention_for_cpur   r  r<   r;   gluselect_scatterslice_scattersplitr   r}   rb   unbinddecomps_to_exclude_opsOperatorBaseOpOverloadPacketr&   _assert_asyncr+   strr0   _functional_assert_asyncr3   r8   SymInttypesNumberr=   rF   ry  rR   rf   r   rm   rk   ro   r   r   r   r   r   r   r   r   detach_r   	Generatorr   r   r   r   r   	view_copyr   rB   r   r   r   r   r`   r   r  r   r   r   r	  	low_dtyper  r  embedding_bag_byte_unpackr2  r>  _foreach_addcmulrO  rG  _foreach_addcdivrK  _foreach_lerprQ  r]  py_impl_CDispatchKeyAutograd	lru_cacher_  rb  rs  choose_qparamsr*   r  r  r  r  r  r.   r)   r(   <module>r     s      
  D D D   # + 
 F ( 7 . 
 X $  g!yy~~		II	YY!!
yy55 ,'**'

' 	' 		'
 	' 	' 			' 	' 			' 			' 	' 	' 	' 	' 	'  	!'" 	--#'$ 	%%%'& 	00''( 	11)'* 	$$+', 	//-'. 	""/'0 	  1'2 	3'4 	5'6 	7'8 	  9': 	;'< 			='> 	

?'@ 	A'B 	C'D 	E'F 	G'H 	  $$I'J 	55K'L 	++M') T K,.J2IJ
 	,,44<<NNNNHHJJLLHHKK " n&8 9>	eEJJ++UZZ-H-HHI	J>c3h> ++//01ELL s t  2
 66::;<u|| # $  = ::BBCD )-(,	LL 
%++$$	% 
%++$$	%	
 
 E % )-(,	||		%++$$	%	 
%++$$	%	 \\		  &	 $	
uS%,,&'
(	""	 	 \\		 %	  ,,4456S
uS%,,&'
(S#YS S \\	S 7S 2234..<<. LL. S		.
 #tCy.!. 3S	>". CcN#. . I. . d. 5<<u||34. 5.@ ,,-.G G GU\\ G /G
 
#
,,LL \\  $* %
  ! !
,,
,, ,, ++

	
 ;; \\  &6 	"
,,LL \\  #F ))*+ 2%,,2	2 \\2 ,2j %:U\\ :ell : &:& 
#
 +/	&||&||& EKK&&'	&
 \\& $&R ++,-   .
 DLL12u||   3 //01 ,0N
,,N (N \\	N 2N EJJ/0Iu|| IELL IU\\ I 1I EJJ/0Iu|| IELL IU\\ I 1I 		" 
,,	#  \\	 # 		" 
,,	#  \\	 # ))*+:
,,:	: : 	:
 \\: ,: //01)
,,)
uS%,,&'
() \\) 2) --./"
,,";;" \\" 0" 48LLE//0  ' $(%)37=
,,= EKK = U\\"	=
 E//0= = \\= (=  ( $(%)37=
,,= EKK = U\\"	=
 E//0= = \\= )=  '
 $(%)%)).)>)>=
,,=c5j!= EKK 	=
 U\\"= U\\"= = = &&= \\= (=* ))112
 $(%)37=
,,=
= EKK 	=
 U\\"= E//0= = \\= 3=& ))334 $(%)37=
,,=	= =
 EKK = U\\"= E//0= = \\= 5=( ,,-5
5
uS%,,&'
(5 5 \\	5 .5 	EEMMN<<LL ,, \\	 O 
;;CCD<< ll LL	
 ,, || ,, || LL  \\ E4 		++EEF
Au|| 
A 
A G
A --./  ||
,,  	
  \\  0: --445
 	
u||
u||$ % 	
 
%,, 6 --445
 	
u||
u||$ % 	
 
%,, 6 **112
%
ell#
 KK
 
%,,	
 3
 ''(<(<(E(EF../<<LL //%,,
' //%,,/	
 .  !&  5<<u||34 0 G< T6T#xS'9"9: 6 6!T#xS'9"9: ! ++,
,,
,, LL \\	 -$ ,;;BBC?<<?? ? 
	?
 ;;? 5<<%&? D? !
 	
)
,,
)<<
) LL
) 	
)
 \\
) "
) 		"
 	
,,<< LL 	
 \\ # 33;;<##LL# 
# 	#
 \\#  =#( ))* 1
,,1	1 <<1 
	1
 1 1 \\1 +1h 445 /3%&&'/||/c/ U3S	>*+/ 3S	>"	/
 CcN#/ / 5<<%&/ 6/r)   