+
    Pji                        ^ RI t ^ RIt^ RIt^ RIHt ^ RIHtHtH	t	H
t
 ^ RIHt ^ RIHtHt ]! RRR7      t]! RRR7      t]P$                  P'                  4       t. R8Ot. R9Ot. R:OtR	 tR
 tR;R ltR tR tR tR tR tR t R t!R t"R t#R t$R t%R t&R t'R t(R])! 4       3R])! ^^^^R7      3R])! ^R7      3R])! ^^^R 7      3R!])! ^R7      3R"])! ^^R#7      3R$])! 4       3R%])! ^^R#7      3R&])! ^^R#7      3.	t*R' t+R( t,R) t-R* t.R+ t/R,])! 4       3R-])! ^^^R.7      3R/])! ^R7      3R0])! 4       3R1])! ^^R#7      3R2])! 4       3.t0R3 t1R4 t2R5 t3R6 t4R7 t5R# )<    N)EinopsError)_enumerate_directions	rearrangereducerepeat)FLOAT_REDUCTIONS)collect_test_backendsis_backend_testedFsymboliclayersTc                     \         P                  ! . RO4      p \        V R4       \        P                  ! \
        4      ;_uu_ 4        \        V R4       RRR4       \        V R4       \        P                  ! \
        4      ;_uu_ 4        \        V R4       RRR4       R#   + '       g   i     LU; i  + '       g   i     R# ; i)   za b c d ... ->  a b c ... dza b c d (...) ->  a b c ... dN... ->  (...)z(...) -> (...)r   r   r   r   r   )npzerosr   pytestraisesr   xs    n/Users/ahmed/devFolder/Ultron/claude-voice/gateway/.venv/lib/python3.14/site-packages/einops/tests/test_ops.py"test_collapsed_ellipsis_errors_outr   ,   s}    
!Aa./	{	#	#!45 
$ a!	{	#	#!%& 
$	#	 
$	# 
$	#	#s   B!
B4!B1	4C	c                  .   \         P                  ! R4      P                  . RO4      p \         F0  p\         P                  ! V \        W4      4      '       d   K+  Q V4       h	  \         F6  w  r#\         P                  ! \        W4      \        W4      4      '       d   K6  Q h	  R FG  p\         F:  w  r#\         P                  ! \        WVR7      \        WVR7      4      '       d   K:  Q h	  KI  	  . \        Op\         F  pVP                  V4       K  	  R# )   	reductionN  r               minmaxsum)
r   arangereshapeidentity_patternsarray_equalr   equivalent_rearrange_patternsequivalent_reduction_patternsr   extend)r   patternpattern1pattern2r   all_rearrange_patternspattern_pairss          r   test_ellipsis_ops_numpyr4   7   s    
		#$,,_=A$~~a1!677@@7 % <~~i4i6LMMMM < +	"?H>>&	"JFSTjsLtuuuu #@ +
 2016%%m4 7    c                  aaa VVV3R lpV! V4      p\         P                  pRp	V'       df   VP                   U
u. uF   p
\        P	                  4       V	8  d   T
MRNK"  	  pp
V P                  VR7      pV! V4      pV P                  WV3.4      pM(V! V P                  V4      4      pV P                  V4      pV! W~4       R# u up
i )z
Helper to test result of operation (rearrange or transpose) against numpy
if reduction == 'rearrange', rearrange op is tested, otherwise reduce
c                 H   < SR 8X  d   \        V S3/ SB # \        V SS3/ SB # )r   )r   r   )r   axes_lengthsr/   r   s   &r   	operation)check_op_against_numpy.<locals>.operationO   s1    #Q8<88!Wi@<@@r5   g      ?N)shape)	r   r+   r;   rngrandomcreate_symboleval_symbol
from_numpyto_numpy)backendnumpy_inputr/   r8   r   is_symbolicr9   numpy_resultcheck_equalp_none_dimensiondsymbol_shapesymbolresult_symbolbackend_results   &&fff&         r   check_op_against_numpyrM   I   s    A [)L..KQ\QbQbcQbASZZ\-==4GQbc&&\&:!&) ,,]k=R<ST"7#5#5k#BC )).9- ds   &Cc                 ~   \         P                  ! R4      P                  . RO4      p R F  p\        VRR7       F{  p\        \        \        P                  ! \         4      ,            F  p\        W V/ RVR7       K  	  R	 F0  p\        P                  ! \          F  p\        W V/ WAR7       K  	  K2  	  K}  	  K  	  R# )
z'Checking various patterns against numpyFr   r   )r8   r   rD   Nr   r   )TFr$   )r   r(   r)   r	   r*   list	itertoolschainr,   rM   r-   )r   rD   rB   r/   r   s        r   test_ellipsis_ops_imperativerR   d   s    
		#$,,_=A$,k%PG,tIOOEb4c/ddd&bK]h e
 3	(0MNG*G"	  O 3 Q %r5   c                     ^ RI p ^ RIHp V P                  R8  d   \        P
                  ! 4        \        P                  ! R4      P                  . RO4      p\        \        \        P                  ! \         4      ,            Fd  p\        W#4      pVP                  V P                  V4      V4      p\        P                   ! VP#                  V^ ,           4      V4      '       d   Kd  Q h	  R# )r   N	array_api2.0.0r   r   )numpyeinopsrU   __version__r   skipr   r(   r)   r*   rO   rP   rQ   r,   r   from_dlpackr+   asnumpy)xpAAr   r/   expectedresults         r   test_rearrange_array_apira   u   s    &	~~
		#$,,_=A$tIOO=Z,['\\\Q(bnnQ/9~~bjj!4h???? ]r5   c            
         ^ RI p ^ RIHp V P                  R8  d   \        P
                  ! 4        \        P                  ! R4      P                  . RO4      p\        P                  ! \          F  pR F{  p\        W#VR7      pVP                  V P                  V4      W4R7      p\        P                  ! VP                  \        P                   ! V^ ,           4      4      V4      '       d   K{  Q h	  K  	  R# )r   NrT   rV   r   r   r   r$   )rW   rX   rU   rY   r   rZ   r   r(   r)   rP   rQ   r-   r   r[   r+   r\   asarray)r]   r^   r   r/   r   r_   r`   s          r   test_reduce_array_apird      s    &	~~
		#$,,_=A??$AB.IaI>HYYr~~a0'YOF>>"**RZZ
-C"DhOOOO / Cr5   c                  F   . ROp \         P                  ! \         P                  ! V 4      4      P                  V 4      pR FQ  p\	        W4      p\        \         P                  ! W4      4      ^ 8X  g   Q hVP                  VP                  8X  d   KQ  Q h	  \	        VR4      p\         P                  ! VP                  4       VP                  4       4      '       g   Q h\	        VR4      p\         P                  ! W4      '       g   Q h\	        VR4      p\	        VR4      p\         P                  ! WE4      '       g   Q h\	        \	        VR4      R^^R7      p\         P                  ! W4      '       g   Q h\        \        RV 4      4      p\	        VR3/ VB p\	        VR3/ VB p\         P                  ! W4      '       g   Q h\         P                  ! ^4      P                  . RO4      p\	        VR	4      pVR,          VR,          8X  g   Q hVR,          VR,          8X  g   Q hR
# )r   a b c d e f -> f e d c b aza b c d e f -> a (b) (c d e) fz0a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11zf e d c b a -> a b c d e fza b c d e f -> (f d) c (e b) az(f d) c (e b) a -> a b c d e f)brH   abcdefza b c -> b c aN)r   r   r    r"         )za b c d e f -> a b c d e fzb a c d e f -> a b d e f crf   za b c d e f -> (f e) d (c b a)za b c d e f -> (f e d c b a)r   r    r!   )r   r   r    )r   r    r   )r   r   r   )r   r   r   )r   r(   prodr)   r   len	setdiff1ddtyper+   flattendictzip)	r;   r   r/   r`   result1result2sizestempx2s	            r    test_rearrange_consistency_numpyrx      s   E
		"''%.!))%0A 1&2<<*+q000||qww&&& q:;F>>!))+v~~'78888qLMF>>!$$$$78G78G>>'++++y$DEGgklpqrF>>!$$$$Xu%&EQ8BEBDt=GGF>>!$$$$	9		%	%i	0Br+,Fg;&/)))g;&/)))r5   c                     \        ^^
4       F  p \        P                  ! ^V ,          4      P                  ^.V ,          4      p\        P                  V 4      pRP                  R \        V 4       4       4      pRP                  R V 4       4      pVR,           V,           p\        W4      p\        P                  ^ ^^
V .4       F2  pV\        V4      ,          V\        Wr,          4      ,          8X  d   K2  Q h	  K  	  \        ^^
4       EF7  p \        P                  ! ^V ,          4      P                  ^.V ,          4      p\        P                  V 4      pRP                  R \        V 4      RRR1,           4       4      pRP                  R VRRR1,           4       4      pVR,           V,           p\        W4      pVP                  VP                  8X  g   Q h\        P                  ! V4      p\        V4       F"  w  rWV	,	          ^,          V
,          ,          pK$  	  \        P                  ! Wh4      '       d   EK8  Q h	  R# )	r    c              3   F   "   T F  pR \        V4      ,           x  K  	  R# 5iiNstr.0axiss   & r   	<genexpr>4test_rearrange_permutations_numpy.<locals>.<genexpr>   s     "M}t3T??}   !c              3   F   "   T F  pR \        V4      ,           x  K  	  R# 5ir|   r~   r   s   & r   r   r      s     #LC#d)OOr    -> c              3   F   "   T F  pR \        V4      ,           x  K  	  R# 5ir|   r~   r   s   & r   r   r      s     "S?Rt3T???Rr   Nc              3   F   "   T F  pR \        V4      ,           x  K  	  R# 5ir|   r~   r   s   & r   r   r      s     #R@QC#d)OO@Qr   )ranger   r(   r)   r<   permutationjoinr   integerstupler;   
zeros_like	enumerater+   )n_axesinputr   left_expressionright_expression
expressionr`   pickexpected_resultoriginal_axisresult_axiss              r   !test_rearrange_permutations_numpyr      s   2,		!V)$,,aS6\:oof-(("MuV}"MM88#L#LL$v-0@@
5-LLAF|4Dt%d6G0H)IIII 5  2,		!V)$,,aS6\:oof-(("SuV}TrT?R"SS88#RDbD@Q#RR$v-0@@
5-||u{{***--.*3K*@&M-!71 <LLO +A ~~f6666 r5   c                  F   \          EF  p \        R V P                  4       \         EF  p\        P
                  ! RRR7      P                  . RO4      pVR9   d&   W"P                  R4      P                  4       ,          pR/ \        W!4      ! 4       .R/ \        W!4      ! 4       .R\        ^^R7      \        W!4      ! 4       .R/ \        W!4      ! RR	7      P                  ^^^ 4      P                  R^.4      .R
/ \        W!4      ! RR	7      P                  ^^^ 4      P                  R^.4      .R/ \        W!4      ! RR	7      P                  ^^^ 4      P                  R^.4      .R/ \        W!4      ! RR	7      P                  ^^^ 4      P                  R.4      .R\        ^R7      V..pV Fi  w  rEp\        V P                  VP                  4       4      V3RV/VB pV P                  V4      p\        P                   ! Wv4      '       d   Ka  Q RV 24       h	  EK  	  EK  	  R# )Reduction tests for int64ro   float64a b c d e -> 	a ... -> z(a1 a2) ... (e1 e2) -> )a1e2a b c d e -> (e c) ar   a ... c d e -> (e c) aa b c d e ... -> (e c) aa b c d e -> (e c a)(a a2) ... -> (a2 a) ...a2r   z
Failed at Nr   r   meanrl   r   r    r   )imp_op_backendsprintframework_name
REDUCTIONSr   r(   r)   astyper   getattrrq   	transposer   r@   copyrA   allclose)rB   r   r   
test_casesr/   r8   r   r`   s           r   test_reduction_imperativesr      s   "?$g&<&<=#III/w?GGXE,,Y 7 < < >> "ge&?&ABb'%";"=>*DA!,<ge>W>YZ*E-6:DDQ1MUUWY[\V]^ -E-6:DDQ1MUUWY[\V]^ /E-6:DDQ1MUUWY[\V]^
 (WU-FF-S-]-]^_abde-f-n-npros-tu+TQZ?)J, ;E6 2 25::< @'oU^obno ))&1{{6;;Sz'=SS; ;E7 $ #r5   c                     \          EF  p \        R V P                  4       \         EF  p\        P
                  ! RRR7      P                  . RO4      pW"P                  R4      P                  4       ,          pR/ \        W!4      ! 4       .R/ \        W!4      ! 4       .R\        ^^R7      \        W!4      ! 4       .R/ \        W!4      ! RR	7      P                  ^^^ 4      P                  R^.4      .R
/ \        W!4      ! RR	7      P                  ^^^ 4      P                  R^.4      .R/ \        W!4      ! RR	7      P                  ^^^ 4      P                  R^.4      .R/ \        W!4      ! RR	7      P                  ^^^ 4      P                  R.4      .R\        ^R7      V..pV EF.  w  rEpVP                  VP                   Uu. uF  pRNK  	  up.pV FT  p	V P                  V	4      p
\        W3RV/VB pV P                  WV3.4      p\        P                   ! W4      '       d   KT  Q h	   . p	/ VCp\#        RVP                  4       F3  w  rW9   d   V	P%                  R4       WV&   K"  V	P%                  V4       K5  	  V P                  V	4      p
\        W3RV/VB pV P                  WV3.4      p\        P                   ! W4      '       d   EK/  Q h	  EK  	  EK  	  R# u upi )r   r   r   r   r   r   z(a a2) ... (e e2) -> )r   r   r   r   r   r   r   r   r   Nr   abcder   r   r   r   )sym_op_backendsr   r   r   r   r(   r)   r   r   r   rq   r   r;   r>   r   r?   r   rr   append)rB   r   r   r   r/   r8   expected_numpy_result_shapesr;   sym
result_symr`   _axes_lengthsr   lengths                   r   test_reduction_symbolicr      s   "?$g&<&<=#III/w?GGXELL388::E !"ge&?&ABb'%";"=>($!*:GE<U<WX*E-6:DDQ1MUUWY[\V]^ -E-6:DDQ1MUUWY[\V]^ /E-6:DDQ1MUUWY[\V]^
 (WU-FF-S-]-]^_abde-f-n-npros-tu+TQZ?)J, AK<'<++ekk'Bkk'BC#E!//6C!'!Z	!Z\!ZJ$00El^LF;;vEEEE	 $ E$4|$4M(+GU[[(A?!LL.28$/!LL0 )B "//6C!'![	![]![J$00El^LF;;vEEEE+ AK5 $ #: (Cs   K"c                     \          EF  p \        R V P                  4       . \        ORN EF  pRp\        P
                  pVR9   d   Rp\        P                  p^pRV P                  9   d   ^pRV P                  9   d   ^	p\        V4       EFi  p\        P                  ^^VR7      p\        P                  V4      pVR8X  d   ^ M\        P                  V^,           4      pRP                  R	 \        V4       4       4      p	RP                  R
 WxR  4       4      p
V	R,           V
,           p\        P                  ! ^^\        P                  ! V4      ,           VR7      P                  V4      pVR8X  d   WP                  4       ,          p\!        WVR7      pVP#                  V4      pV^ 8  d%   \%        W4      ! \'        \        V4      4      R7      pV! W4      '       g   Q h\)        WW/ RR7       EKl  	  EK  	  EK  	  R# )zStress-testing reduction for r   r   rl   r   oneflowpaddlesizerz   c              3   F   "   T F  pR \        V4      ,           x  K  	  R# 5ir   Nr~   r   r}   s   & r   r   4test_reduction_stress_imperatives.<locals>.<genexpr>A  s     Dmc!fmr   c              3   F   "   T F  pR \        V4      ,           x  K  	  R# 5ir   r~   r   s   & r   r   r   B  s      M7L!s1v7Lr   N->r   r   r   F)r   r8   rD   r   )r   r   r   r   r   r+   r   r   r<   r   r   r   r(   rl   r)   r   r   r   r   r   rM   )rB   r   ro   coincidemax_dimr   r;   r   skippedleftrightr/   r   rs   rt   s                  r   !test_reduction_stress_imperativesr   /  s   "?-w/E/EF3:3{3IE~~H,,!;;GG2227111.Q7!oof5(K7!S\\&ST*=UxxDeFmDD M{87L MM+-IIaRWWU^!35AII%P&MA yA++k2Q;%g9uU7^?TUG1111&w7^`nstt ) 4 #r5   c                     \         P                  ! R
4      P                  . RO4      P                  R4      p W P	                  4       ,          p R pR pR pR p^ RIHpHpHpH	p VP                  VVP                  VVP                  VVP                  V/p	\         F  p
V
P                  V	9  d   K  WP                  ,          pV
P                  V 4      p\         Fd  w  r\        RV
P                  W4       \        WVR7      p\        WVR7      p\         P                   ! VV
P#                  V4      4      '       d   Kd  Q h	  K  	  R	# )r   float32c                 $    V P                  V4      # )N)	logsumexp)r   tuple_of_axess   &&r   logsumexp_torchAtest_reduction_with_callable_imperatives.<locals>.logsumexp_torchS  s    {{=))r5   c                 ,    ^ RI pVP                  W4      # r   N)
tensorflowreduce_logsumexp)r   r   tfs   && r   logsumexp_tf>test_reduction_with_callable_imperatives.<locals>.logsumexp_tfV  s    ""144r5   c                 6    ^ RI Hu Hp VP                  W4      # r   )tensorflow.keras.backendkerasrB   r   )r   r   ks   && r   logsumexp_kerasAtest_reduction_with_callable_imperatives.<locals>.logsumexp_keras[  s    ,,{{1,,r5   c                     V P                  V4      pW P                  VR R7      ,
          p\        P                  ! V4      p\        P                  ! W1R7      p\        P                  ! V4      V,           # )T)keepdimsr   )r&   r   expr'   log)r   r   minusedys   &&  r   logsumexp_numpyAtest_reduction_with_callable_imperatives.<locals>.logsumexp_numpy`  sS    %%&md33FF1IFF1)vvay7""r5   )NumpyBackendTensorflowBackendTFKerasBackendTorchBackendz!Test reduction with callable for r   Nr   r   )r   r(   r)   r   r&   einops._backendsr   r   r   r   r   r   r@   r-   r   r   r   rA   )x_numpyr   r   r   r   r   r   r   r   backend2callbackrB   backend_callback	x_backendr0   r1   output_numpyoutput_backends                    r   (test_reduction_with_callable_imperativesr   O  s+   ii)*22?CJJ9UG{{}G*5
-
# _^ 	##_((,%%##_	 #!!)99+,B,BC&&w/	"?H5w7M7Mxb!'OL#ICSTN;;  0   	 #@ #r5   c                  v   \          EF(  p \        R V P                  4       . ^.. RO. RO3 F  p\        P                  ! \        P
                  ! V4      4      P                  V4      p\        V4      p\        V P                  V4      4      p\        V4      \        V4      u;8X  d   \        V4      8X  g   Q h Q hV Uu. uF  qPP                  V4      NK  	  pp\        W44       F@  w  rgVP                  VP                  8X  g   Q h\        P                  ! Wg4      '       d   K@  Q h	  EK  	  EK+  	  R# u upi )ztesting directions forN)r   r   r   r   r    r"   ri   )r   r   r   r   r(   rl   r)   r   r@   rm   rA   rr   r;   r   )rB   r;   r   axes1axes2axax1ax2s           r   test_enumerating_directionsr    s    "?&(>(>?1#y,7E		"''%.)11%8A)!,E)'*<*<Q*?@Eu:U9s5z9999949:Eb%%b)EE:-yyCII---{{3,,,, . 8 # ;s    D6c                  "   \          EFy  p \        R V P                  4       R EFX  p. ^.^^.. RO^.^,          .pV EF9  p\        V4       Uu. uFC  p\        P
                  ! WD\        P                  ! V4      ,           4      P                  V4      NKE  	  ppV Uu. uF  q`P                  V4      NK  	  pp\        P                  ! V4      p\        VR4      p	\        VR4      p
\        P                  ! W4      '       g   Q h\        P                  ! WP                  V
4      4      '       g   Q h\        VR4      p	\        VR4      p
\        P                  ! WP                  V
4      4      '       d   EK:  Q h	  EK[  	  EK|  	  R# u upi u upi )ztesting shapes for ...->...zb ... -> ... bN)r   r   r"   r  )r   r   r   r   r   r(   rl   r)   r@   rc   r   r+   rA   )rB   n_arraysr   r;   r}   arrays1arrayarrays2result0rs   rt   s              r    test_concatenations_and_stackingr    sG   "?#W%;%;<!H1#1v|aS1W=FTYZbTcdTcq299QBGGEN(:;CCEJTcdBIJ'--e4'J**W-#GZ8#GZ8~~g7777~~g/?/?/HIIII#G-=>#G-=>~~g/?/?/HIIII   " #
 eJs   A	F
F
c            	         \          EF[  p V R9   d   K  \        P                  ! ^^4      P                  . RO4      P	                  R4      p/ p\
         F  pVP                  V4      p\        VR4      '       g   K(  \        VRV R7      p\        VRV R7      p\        VRV ^R7      p\        VRV R7      pVP                  4        VP                  VP                  4      p	WVP                  &   K  	  \        R	VP                  4       4       VP                  4        FC  w  rVP                  4        F*  w  r\        P                   ! W4      '       d   K#  Q WR
.4       h	  KE  	  EK^  	  R# )anyr   gradza b c -> c ar   z
c a -> a cza (c1 c2) -> a)r   c1z... -> zcomparing gradients forzprovided different gradientsN)r  allrk   )r   r   r(   r)   r   r   r@   hasattrr   backwardrA   r  r   r   keysitemsr   )r   r   resultsrB   y0y1y2y3y4r  name1grad1name2grad2s                 r   test_gradients_imperativesr#    s+   Z	&IIa'//	:AA)L&G##A&B2v&&Ni@BLI>B,	aHBI;BKKM##BGG,D.2G**+ ' 	'8#MMOLE '{{500`5A_2``0 !0 ,)  r5   c                     \          F  p \        R V P                  4       \        P                  ! ^RR7      P                  . RO4      p. ROpV Fj  p\        P                  ! W4      pV P                  V4      pV P                  WS4      pV P                  V4      p\        P                  ! Wt4      '       d   Kj  Q h	  K  	  R# zTiling tests for r   r   N)r   r   r    r   r"   )r   )r   r   r   r    r   )r    r   r   r!   r   )
r   r   r   r   r(   r)   tiler@   rA   r+   )rB   r   r   repeatsr_   	convertedrepeatedr`   s           r   test_tiling_imperativesr*    s    "!7#9#9:		)73;;OL


 "Gwwu.H**51I||I7H%%h/F>>&3333 " #r5   c                  z   \          EF/  p \        R V P                  4       \        P                  ! ^RR7      P                  . RO4      p. ROpV F  p\        P                  ! W4      pV P                  VP                  4      pV P                  V P                  WS4      WQ..4      p\        P                  ! Wd4      '       g   Q hV P                  R.\        VP                  4      ,          4      pV P                  V P                  WS4      WQ..4      p\        P                  ! Wd4      '       d   K  Q h	  EK2  	  R# r%  )r   r   r   r   r(   r)   r&  r>   r;   r?   r+   rm   )rB   r   r   r'  r_   r   r`   s          r   test_tiling_symbolicr,    s    "?!7#9#9:		)73;;OL


 "Gwwu.H''4C((c)Csl^TF>>&3333''U[[1A(ABC((c)Csl^TF>>&3333 " #r5   za b c -> c a bza b c -> (c copy a b))r   arg   cza b c -> (a copy) b c r   za b c -> (c a) (copy1 b copy2))r-  copy1copy2za ...  -> a ... copyz... c -> ... (copy1 c copy2))r0  r1  z...  -> ... z ...  -> copy1 ... copy2 za b c  -> copy1 a copy2 b c () c                   VP                  R4      w  r4VR,           V,           p\        W3/ VB p\        We3RR/VB p\        We3RR/VB p\        P                  ! W4      '       g   Q h\        P                  ! W4      '       g   Q hR# )z*Checks repeat pattern by running reductionr   r   r%   r&   N)splitr   r   r   r+   )	r   repeat_patternru   r   r   reduce_patternr)  reduced_minreduced_maxs	   &&,      r   check_reversionr8    s     &&t,KDT\D(Na151HLULeLKLULeLK>>!))))>>!))))r5   c                      \         P                  ! ^4      P                  . RO4      p \        V R^R7      p\         P                  ! V R,          V4      '       g   Q h\
         F  w  r#\        W3/ VB  K  	  R# )r   za b c -> copy a b c r/  Nr   r    r"   )r   r(   r)   r   r+   repeat_test_casesr8  )r   x1r/   axis_dimensionss       r   test_repeat_numpyr>    s_    
		)$$Y/A	)	2B>>!D'2&&&&$5 6o6 %6r5   c                  h   \         P                  ! ^4      P                  . RO4      p \         F  p\	        RVP
                  4       \         F]  w  r#\        W3/ VB pVP                  V 4      p\        WR3/ VB pVP                  V4      p\         P                  ! Wt4      '       d   K]  Q h	  K  	  R# r   zRepeat tests for Nr:  )r   r(   r)   r   r   r   r;  r   r@   rA   r+   )r   rB   r/   r=  r_   r(  r)  r`   s           r   test_repeat_imperativesrA    s    
		)$$Y/A"!7#9#9:(9$Ga<O<H**1-IiDODH%%h/F>>&3333 ): #r5   c            	      ~   \         P                  ! ^4      P                  . RO4      p \         F  p\	        RVP
                  4       \         Fh  w  r#\        W3/ VB pVP                  V P                  4      pVP                  \        WR3/ VB WP..4      p\         P                  ! Wd4      '       d   Kh  Q h	  K  	  R# r@  )r   r(   r)   r   r   r   r;  r   r>   r;   r?   r+   )r   rB   r/   r=  r_   r   r`   s          r   test_repeat_symbolicrC    s    
		)$$Y/A"!7#9#9:(9$Ga<O<H''0C(()P)PTWS[R\]F>>&3333 ): #r5   c                     ^ RI p ^ RIHp V P                  R8  d   \        P
                  ! 4        \        P                  ! ^4      P                  . RO4      p\         Fi  w  r4\        W#3/ VB pVP                  ! V P                  V4      V3/ VB p\        P                  ! VP                  V^ ,           4      V4      '       d   Ki  Q h	  R# )r   NrT   rV   r:  )rW   rX   rU   rY   r   rZ   r   r(   r)   r;  r   r[   r+   r\   )r]   r^   r   r/   r=  r_   r`   s          r   test_repeat_array_apirE     s    &	~~
		)$$Y/A$5 !882>>!,gII~~bjj!4h????	 %6r5   za b c d -> c a d bza b c d -> (c 2 d a b))r-  r.  rH   z1 b c d -> (d copy 1) 3 b c z1 ...  -> 3 ... z"() ... d -> 1 (copy1 d copy2) ... z$1 b c d -> (1 1) (1 b) 2 c 3 d (1 1)c                      \         P                  ! ^04      P                  . RO4      p \         F  w  r\	        W3/ VB  K  	  R# )r   N)r   r   r!   r#   )r   r(   r)   test_cases_repeat_anonymousr8  )r   r/   r=  s      r   test_anonymous_axesrH  <  s6    
		- ((6A$? 6o6 %@r5   c            	         \         P                  ! R4      P                  . RO4      p \         P                  ! \	        \        V 4      R4      \	        V R4      4      '       g   Q h\         P                  ! \        \        V 4      RR4      \        V RR4      4      '       g   Q h\         P                  ! \        \        V 4      R^R7      \        V R^R7      4      '       g   Q hR# )	r   z... -> (...)za ... e -> (...)r%   z...  -> b (...))rg   Nr   r   )r   r(   r)   r+   r   rO   r   r   r   s    r   test_list_inputsrJ  B  s    
		#$,,_=A>>$q'>*!^$    >>tAw*E2q$e,    >>tAw)Q/q#q)   r5   c                  v   \        R 4      '       g   \        P                  ! 4        ^ RIp RV P                  P
                  n        R pV P                  VRRR7      pV P                  VRRR7      pV P                  . R	OR7      pV P                  V! V4      V! V4      RR7      '       g   Q hV P                  V! V4      V! V4      RR7      '       g   Q hV P                  . R
OR7      pV P                  V! V4      V! V4      RR7      '       g   Q hV P                  V! V4      V! V4      RR7      '       g   Q hR# )torchNTc                 j    V P                   vrr4\        V R ^R7      p \        V RRV^,          R7      p V # )z"(a a2) b c ... -> b (c a2) (a ...)r   zb ca2 A -> b Ar'   )ca2)r;   r   r   )r   r-  rg   r.  others   &    r   func14test_torch_compile_with_dynamic_shape.<locals>.func1[  s8    ''aa=!D1&1q59r5   F)dynamic	fullgraphr   gh㈵>)atol)r!   r"   r#   r    )r#   r    r!   r   r    )
r
   r   rZ   rL  _dynamoconfigverbosecompilerandnr   )rL  rP  func1_compiled_staticfunc1_compiled_dynamicr   s        r   %test_torch_compile_with_dynamic_shaper\  S  s   W%% $(EMM  "MM%$MO"]]5$$]O&A>>/2E!H4>HHHH>>03U1XD>IIII)A>>/2E!H4>HHHH>>03U1XD>IIIIr5   c                 @   a  \        V 3R  l\        ^4       4       4      # )c              3   D   <"   T F  pSV,	          ^,          x  K  	  R# 5i)r   N )r   r}   r   s   & r   r   bit_count.<locals>.<genexpr>q  s     /YQ!||Ys    )r'   r   r   s   fr   	bit_countra  p  s    /U2Y///r5   c                 .   \         P                  ! \        ^@4       U u. uF  p \        V 4      ^,          ^ 8H  NK  	  up 4      P	                  ^.^,          4      p\
         EF+  p\        RVP                  4       \        ^4       EF9  p\         P                  ! WRR7      p\         P                  ! WRR7      p\         P                  ! WE4      '       d   Q h\        R4      p\        V4      p\        V4      pRW&   RP                  V4      R,           RP                  V4      ,           p	\        VP                  V4      V	RR	7      p
\        VP                  V4      V	R
R	7      p\         P                  ! WBP                  V
4      4      '       g   Q h\         P                  ! WRP                  V4      4      '       d   EK:  Q h	  \         P                  ! VRRR7      p\         P                  ! VRRR7      pRp	\        VP                  V4      V	RR	7      p
\        VP                  V4      V	R
R	7      p\         P                  ! WBP                  V
4      4      '       g   Q h\         P                  ! WRP                  V4      4      '       d   EK,  Q h	  R# u up i )z5Checks that any/all reduction works in all frameworkszReduction any/all tests for T)r   r   rh   1rz   r   r  r   r  za b ... -> 1 1 ...N)r   r   )r   rc   r   ra  r)   r   r   r   r  r  r+   rO   r   r   r@   rA   )r   x_nprB   r   expected_result_anyexpected_result_allaxesaxes_inaxes_outr/   res_anyres_alls               r   #test_reduction_imperatives_booleansrl  t  s   ::dD1	!q(Q.DEMMqcTUgVD"?,g.D.DE!HD"$&&4"H"$&&4"H~~&9OOOO>D4jGDzH HNxx(F2chhx6HIGW//5w%PGW//5w%PG>>"57G7G7PQQQQ>>"57G7G7PQQQQ $ !ffTF ffTF&++D17eL++D17eL~~13C3CG3LMMMM~~13C3CG3LMMMM7 # Es   J)r	  za b c d e-> a b c d eza b c d e ...-> ... a b c d eza b c d e ...-> a ... b c d ez... a b c d e -> ... a b c d eza ... e-> a ... eza ... -> a ... za ... c d e -> a (...) c d e))za b c d e -> (a b) c d eza b ... -> (a b) ... )za b c d e -> a b (c d) ez... c d e -> ... (c d) e)za b c d e -> a b c d ez... -> ... )za b c d e -> (a b c d e)r   )za b c d e -> b (c d e) aza b ... -> b (...) a)za b c d e -> b (a c d) eza b ... e -> b (a ...) e))r   z	 ... ->  )za b c d e -> (e a)za ... e -> (e a))za b c d e -> d (a e)z a b c d e ... -> d (a e) )za b c d e -> (a b)z ... c d e  -> (...) )r   F)6rP   rW   r   r   rX   r   einops.einopsr   r   r   r   einops.testsr   r   r	   r
   r   r   r=   default_rngr<   r*   r,   r-   r   r4   rM   rR   ra   rd   rx   r   r   r   r   r   r  r  r#  r*  r,  rq   r;  r8  r>  rA  rC  rE  rG  rH  rJ  r\  ra  rl  r_  r5   r   <module>rp     s       J J 7 A'uE'eDii	 ! ! '5$.6"@P #*L78!TH2Fju@0f-J&a64"4, tvdQ!q9:t|,%taq'BCTq\*#T%;<TV $Qa"89&1A(>? *7
44@& 46"ta123#Tq\2 )4aq+AB+TV4 7"J:0Nr5   