
    ЏkhpI                        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
 d dlmZmZmZ d dlmZmZmZ d dlmZ g dZddgZd	 Zej0                  j3                  d
      ej0                  j3                  d      ej0                  j3                  d       ee       G d d                                  Z ee       G d d             Z ee       G d d             Zy)    N)is_array_api_strictmake_xp_test_casexp_default_dtype	xp_device)xp_assert_equalxp_assert_closexp_assert_less)log_softmax	logsumexpsoftmax)_wrap_radians)float32float64int32int64	complex64
complex128r   r   c                    | j                  t        j                   dz
  t        j                   dddddt        j                  t        j                  dz   g	      }| j                  t        j                  dz
  t        j                  dddddt        j                  t        j                   dz   g	      }t        ||       }t	        ||d       y )N   gYnr   gYn)xp)atol)asarraymathpir   r   )r   xrefress       ^/var/www/teggl/fontify/venv/lib/python3.12/site-packages/scipy/special/tests/test_logsumexp.pytest_wrap_radiansr       s    


TWWHQJ"gvq$''477196 	7A
**dggai"gvq$''DGG8A:7 8C
b
!CC1%    z/ignore:invalid value encountered:RuntimeWarningz0ignore:divide by zero encountered:RuntimeWarningz*ignore:overflow encountered:RuntimeWarningc                   >   e Zd Zd Zd Zd Zd Zd Zd Zd Z	d Z
d	 Zej                  j                  d
dg df      d        Zd Zej                  j                  de      d        Zej                  j                  de      ej                  j                  de      d               Zd Zej                  j                  dddg      d        Zd Zej                  j                  d ej0                  ej4                   ej4                   eej4                   d       eej4                   d       eej4                   ej4                         eej4                   ej4                          eej4                  d       eej4                  d       eej4                  ej4                         eej4                  ej4                          eej4                   d       eej4                   d       eej4                   d       eej4                   d       eej4                  d       eej4                  d       eej4                  d       eej4                  d      gd !            d"        Zd# Zej                  j                  d$d%d&g g      d'        Zd( Zy))*TestLogSumExpc                 *   |j                  g       }|j                  |j                         }t        t        |      |       |j	                  d|j
                        }|j                  |j                  |j                  |                  }t        t        |      |       |j                  ddg      }|j                  dt        j                  d      z         }t        t        |      |       d}|j                  |fd      }|j                  dt        j                  |      z         }t        t        |      |       |j                  dgd	z        }|j                  |      }|j                  ||g      }|j                  ||g      }	t        |j                  t        |	            |j                  |             t        |j                  t        |	d
            |j                  |d
             t        |j                  t        |	d            |j                  |d             |j                  |j                  g      }
|j                  |j                  g      }t        t        |
      |
d
          t        t        |
       |
d
           t        t        |      |d
          t        t        |j                  |j                   |j                   g            |
d
           |j                  ddgdt        j                   gg      }|j                  ddg      }t        t        |d      |       |j                  |d      }t        t        |dd      |       t        t        |d      |j                  d             y )N      i@dtype     @@       @  '  g     @Ww'&l7i@B r   axisr   g    _Bg|=g    _r   T)r.   keepdims)r   )r   infr   r   aranger   logsumexpr   r   fullstacknannpexpand_dims)selfr   adesiredbnr   logxXlogXr1   r8   r   s                r   test_logsumexpzTestLogSumExp.test_logsumexp%   s   JJrN**bffW%	!g. IId"**I-&&q	*+	!g. JJu~&**Vdhhsm34	!g.GGQD% **Wtxx{23	!g.JJw()vvayHHaVxxt%y/;yA679JKyA679JK jj"&&"jj"&&"	#A/	3$#a&1	#A/	"**rvvgw-?"@ACF7K JJu(* +jj$'	!"-s3 nnSrn*	!"t<cB 		!(3RZZ5EFr!   c           	         |j                  d|j                        }|j                  ddd      }|j                  |j                  ||j	                  |      z              }t        t        ||      |       |j                  ddg      }|j                  ddg      }|j                  dt        j                  d      z         }t        t        ||      |       |j                  d	gd
z        }|j                  ddd
      }|j                  |      }|j                  ||f      }|j                  ||f      }|j                  ||f      }	t        |j	                  t        ||	            |j                  |	|z               t        |j	                  t        ||	d            |j                  |	|z  d             t        |j	                  t        ||	d            |j                  |	|z  d             y )Nr%   r&                 r>   r*   g333333?g333333@r,   i r   r   )r>   r.   r-   )r2   r   r3   r4   r5   r   r   r   r   linspacer7   )
r;   r   r<   r>   r=   r   r@   rA   rB   Bs
             r   test_logsumexp_bzTestLogSumExp.test_logsumexp_bW   s   IId"**I-IIdB$&&"&&),-	!q)73JJd|$JJSz"**TDHHW$556	!q)73JJw'(KK4(vvayHHaVxxt%HHaVy34bffQUmDy;<bffQUQRf>STy;<bffQUQRf>STr!   c                     |j                  g d      }|j                  g d      }t        ||d      \  }}t        ||j                  d             t        ||j                  d             y )N)r   r   r   )r   r   r   Tr>   return_sign      ?rF   )r   r   r   r   r;   r   r<   r>   rss         r   test_logsumexp_signz!TestLogSumExp.test_logsumexp_signl   sT    JJy!JJ{#aT212::b>*2::c?+r!   c                     |j                  ddg      }|j                  ddg      }t        ||d      \  }}|j                  |      rJ |j                  |      rJ |dk  sJ |dk(  sJ y )Nr   r   TrL   r   )r   r   isfiniteisnanrO   s         r   test_logsumexp_sign_zeroz&TestLogSumExp.test_logsumexp_sign_zerot   so    JJ1vJJ2waT21;;q>!!88A;1uuAvvr!   c                     |j                  d      }|j                  |      }t        |d|d      \  }}|j                  |j                  cxk(  rdk(  sJ  J t        |d|d      \  }}|j                  |j                  cxk(  rdk(  sJ  J y )Nr            rY   T)r.   r>   rM   r   rY   r[   r   rZ   ones	ones_liker   shaperO   s         r   test_logsumexp_sign_shapez'TestLogSumExp.test_logsumexp_sign_shape~   s    GGL!LLOaT:1ww!''.Y.....1$?1ww!''*U*****r!   c                    |j                  g d      }t        |d      \  }}|j                  |j                  |            }||j	                  |      z  }t        ||       t        ||j                  |      z  |       y )N)y      ?      ?y       @      y             @T)rM   )r   r   r4   r5   absr   )r;   r   r<   rP   rQ   expected_sumexpexpected_signs          r   test_logsumexp_complex_signz)TestLogSumExp.test_logsumexp_complex_sign   sj    JJ01-1&&+'"&&*AA=)BFF1I7r!   c                     |j                  d      }|j                  |      }t        |d|      }|j                  dk(  sJ t        |d|      }|j                  dk(  sJ y )NrX   rY   )r.   r>   r\   r]   r^   )r;   r   r<   r>   rP   s        r   test_logsumexp_shapez"TestLogSumExp.test_logsumexp_shape   s]    GGL!LLOaa1%ww)###af*ww&   r!   c                     |j                  ddg      }|j                  ddg      }t        t        ||      |j                  d             y )Nr   r+   r   rG   rN   )r   r   r   r;   r   r<   r>   s       r   test_logsumexp_b_zeroz#TestLogSumExp.test_logsumexp_b_zero   s?    JJ5z"JJ1v	!q)2::b>:r!   c                 b    |j                  d      }|j                  d      }t        ||       y )N)r[   r   rY   r   )rZ   r      rG   )zerosr_   r   rk   s       r   test_logsumexp_b_shapez$TestLogSumExp.test_logsumexp_b_shape   s(    HH\"GGI!qr!   argr   )r   rY   rZ   c                     t        |      t        t        j                  t        j                  |                  k(  sJ y N)r   r9   r   
atleast_1d)r;   rq   s     r   test_xp_invalid_inputz#TestLogSumExp.test_xp_invalid_input   s+    ~2::bmmC6H+I!JJJJr!   c                     ddg}t        j                  dt        j                  d      z         }t	        t        |      |       y )Nr*   r(   r)   )r9   r   r   r3   r   r   )r;   r<   r=   s      r   test_array_likezTestLogSumExp.test_array_like   s4    4L**Vdhhsm34	!g.r!   r'   c                     t        ||      }|j                  ddg|      }|j                  |d      r|j                  d      j                  n|}|j                  dt	        j
                  d      z   |      }t        t        |      |       y )Nr(   r&   integralrN   r)   )getattrr   isdtyper'   r   r3   r   r   )r;   r'   r   r<   desired_dtyper=   s         r   test_dtypes_azTestLogSumExp.test_dtypes_a   sw    E"JJu~UJ313E:1NB--# 	**Vdhhsm3=*I	!g.r!   dtype_adtype_bc                    t        ||      }t        ||      }|j                  ddg|      }|j                  ddg|      }t        |      rk||fD cg c]  }|j                  |d      s| }	}t	        |	      dk  r*|	j                  |j                  d      j                          |j                  |	 }
n/|j                  ||      }
|j                  |
d      rt        |      }
|j                  t        j                  t        j                  d      t        j                  d      z
        |
      }t        t        ||      |       y c c}w )NrY   r   r&   r   ry   rN   rG   )rz   r   r   r{   lenappendr'   result_typer   r   r3   r5   r   r   )r;   r~   r   r   
xp_dtype_a
xp_dtype_br<   r>   r'   xp_float_dtypesr|   r=   s               r   test_dtypes_abzTestLogSumExp.test_dtypes_ab   s,    R)
R)
JJ1vZJ0JJ2wjJ1r" 4>z2J E&(jj
&C  % EO E?#a'&&rzz"~';';<*BNNO<MNN:zBMzz-4/3}**TXXdhhqkDHHQK&?@*V	!q)73Es   Ec                     |j                  ddg      }t        |      }|j                  |d   |d         }t        ||       y )NrE   g      Dr   r   )r   r   	logaddexpr   )r;   r   r<   r   r   s        r   test_gh18295zTestLogSumExp.test_gh18295   sA     JJU|$lll1Q41&S!r!   r   r   c                 `   t         j                  j                  d      }t        ||      }d}|j	                  dd|      d|j	                  dd|      z  z   }|j                  ||      }t        |d      }|j                  |j                  |j                  |      d            }|j                  |j                  |      |j                        }t        |j                  |j                  |            |       t        ||       t        |dd	      \  }	}
|j                  |j                  |      d      }t        |j                  |j                  |
            |       t        |	|j!                  |j                  |                   t        |
||j                  |      z         y )
Nl   &DN$)
   d   r   (                 ?r&   r-   T)rM   r.   )r9   randomdefault_rngrz   uniformr   r   r3   r4   r5   	full_likeimagr   r	   rd   r   real)r;   r   r'   rngra   r   r   r   maxoutsgns              r   test_gh21610zTestLogSumExp.test_gh21610   sM    ii##$67E"KK2u%ckk!R.G(GGJJqJ&"ffRVVBFF1IAV./ll2773</rvvbggcl+S1S!QDq9SffRVVAYQf'rvvbggcl+S1RWWRVVC[12S_-r!   c                    |j                  ddg      }t        |      }|j                  d      }t        |j                  |      |j                  |             t        |j	                  |      |j	                  |      dd       y )Nr   y        y<y9B.?i<gV瞯<)r   rtol)r   r   r   r   r   )r;   r   r   r   r   s        r   test_gh21709_small_imaginaryz*TestLogSumExp.test_gh21709_small_imaginary   sj     JJ567l
 jjCDbggcl3bggclGr!   zx,yrE   g       g%u?g}b@g"~j@gH@rY   )repeatc           
          t        |j                  ||g            }|j                  |j                  |j	                  |j                  ||g                        }t        ||       y rs   )r   r   r3   r4   r5   r   )r;   r   yr   r   r   s         r   test_gh22601_infinite_elementsz,TestLogSumExp.test_gh22601_infinite_elements  sR    4 

Aq6*+ffRVVBFF2::q!f#5678S!r!   c                     |j                  ddg      }|j                  ddg      }t        |       t        ||       t        ||j                  ddg             t        ||j                  ddg             y)zFTest that logsumexp doesn't accidentally write back to its parameters.g      @g      @g      @r)   rG   N)r   r   r   rk   s       r   test_no_writebackzTestLogSumExp.test_no_writeback  se    JJBx JJBx !!q2::r2h/02::r2h/0r!   x_rawrN   r   c                     |D ][  }|j                  ||      }t        t        |            t        |      k(  sJ t        t        ||            t        |      k(  r[J  y)z(Test input device propagation to output.)devicerG   N)r   r   r   )r;   r   r   devicesdr   s         r   test_devicezTestLogSumExp.test_device(  s`      	@A

5
+AYq\*il:::YqA./9Q<???	@r!   c                 "   |j                  g d      }|j                  g d      }t        t        ||      t        ||d      d          |j                  g d      }t        t        ||      |j                  |j                               y )N)grC@g X?g	ml@)_+J?eY@g|H4rG   TrL   r   )r   r   gR%)r   r   r   r8   rk   s       r   test_gh22903zTestLogSumExp.test_gh229030  sp     JJ;<JJ<=	!q)9Q!+Nq+QR JJ<=	!q)2::bff+=>r!   N) __name__
__module____qualname__rC   rJ   rR   rV   rb   rg   ri   rl   rp   pytestmarkparametrizeru   rw   dtypesr}   r   r   r   r   itproductr9   r1   complexr   r   r   r    r!   r   r#   r#       s   
0GdU*,+
8!; [[UQ	N3K 4K/
 [[Wf-/ ./ [[Y/[[Y/4 0 04(" [[W{L&AB. C.0H [[UJBJJVVGFFRVVGR RVVGS!RVVGRVV$RVVGbffW%BFFBBFFC BFFBFF#BFFRVVG$RVVGV$RVVGV$RVVGV$RVVGV$BFFF#BFFF#BFFF#BFFF#'	
( +% ."/."1 [[WsD"o6@ 7@?r!   r#   c                   j    e Zd Zd Zd Zej                  j                  dd      d        Zd Z	d Z
y	)
TestSoftmaxc                    t        t        |j                  g d            |j                  g d      d       t        t        |j                  ddg            |j                  ddg      d       t        t        |j                  ddg            |j                  dt        j                  g      dt        j                  z   z  d       |j                  d	|j                  
      }|j                  g d|j                  
      }t        t        |      |d       t        t        |dz         |d       t        t        |j                  |d            |j                  |d      d       y )N)r(   rE   rE   rE   rN   rE   rE   rE   vIh%<=r   rN         ?rE   r   r[   r&   )g*lI9i?g{O?gW-R?gI?r   rY   rY   )r   r   r   r9   er2   r   reshape)r;   r   r   expecteds       r   test_softmax_fixturesz!TestSoftmax.test_softmax_fixturesA  s)   

+> ?@

#345	B

B8 45

B8,5	:

B8 45

B:.BDD9"	$ IIarzzI*:: 3 ;=**  F
 	
H59 	C((? 	

1f 56

8V45	Br!   c           	      n   t        t        |j                  ddgddgg      d      |j                  ddgddgg      d       t        t        |j                  ddgddgg      d      |j                  d	dgd	dgg      d       |j                  g d
g dg      }|j                  g dg dg      }t        t        |d      |d       t        t        |j                  d      |j                  d       |j	                  |d      }t        t        |d      |j	                  |d      d       y )Nr(   rE   r   r-   r   r   r   r   rN   )g      9rE   g      9@g      I@)rN   g     Pt@g     h@g     p@)g	+9g.K|T}%m;g]7=g\?)rE   gpk&gTV6?gd?rY   rY   rY   r   rY   )r   r   r   Tr   )r;   r   r   r   x3ds        r   test_softmax_multi_axesz#TestSoftmax.test_softmax_multi_axes]  s$   

UBK%+E FQO

RHr2h#78u	F

UBK%+E FQO

RHr2h#78u	F
 JJ002 3::  3 4	5 6 	*H5A!,hjjuE jjI&&1

8Y7e	Er!   array_api_strictint->float promotionreasonc                 v    t        t        |j                  g d            |j                  g d      d       y N)r*   r   r   r   r   r   r   )r   r   r   r;   r   s     r   test_softmax_int_arrayz"TestSoftmax.test_softmax_int_arrayw  s*    

? ;<

#345	Br!   c                 X    t        t        d      t        j                  d      d       y )Nr*   rN   r   r   r   r   r9   r   r;   s    r   test_softmax_scalarzTestSoftmax.test_softmax_scalar|  s    rzz"~EBr!   c                 `    t        t        g d      t        j                  g d      d       y r   r   r   s    r   test_softmax_array_likez#TestSoftmax.test_softmax_array_like  s!    0

#345	Br!   N)r   r   r   r   r   r   r   xfail_xp_backendsr   r   r   r   r!   r   r   r   ?  sG    B8E4 [[""#5>T"UB VBCBr!   r   c                      e Zd Zd Zej
                  j                  dd      d        Zd Zd Z	e
d        Ze
d	        Zej
                  j                  d
ddg      d        Zd Zej
                  j                  dd ej"                  d       ej$                  d      z  fdddgddggfg      d        Zd Zd Zd Zy)TestLogSoftmaxc                 v    t        t        |j                  ddg            |j                  ddg      d       y )Nr(   rN   rE        8r   r   r   r
   r   r   s     r   test_log_softmax_basicz%TestLogSoftmax.test_log_softmax_basic  s0    BJJr{$;<

B;/e	=r!   r   r   r   c                 v    t        t        |j                  ddg            |j                  ddg      d       y Nr*   r   rE   r   r   r   r   r   s     r   test_log_softmax_int_arrayz)TestLogSoftmax.test_log_softmax_int_array  s0    BJJay$9:

B;/e	=r!   c                 2    t        t        d      dd       y )NrN   rE   r   r   )r   r
   r   s    r   test_log_softmax_scalarz&TestLogSoftmax.test_log_softmax_scalar  s    C(#E:r!   c                 `    t        t        ddg      t        j                  ddg      d       y r   )r   r
   r9   r   r   s    r   test_log_softmax_array_likez*TestLogSoftmax.test_log_softmax_array_like  s'    T1I.

B;/e	=r!   c                     | j                  d| j                        }g d}|| j                  || j                        fS )Nr[   r&   g:-g:-guZgi,ܿ)r2   r   r   r   r   expects      r   data_1dzTestLogSoftmax.data_1d  s=    IIarzzI*( "**V2::*666r!   c                     | j                  | j                  d| j                        d      }g dg dg}|| j                  || j                        fS )N   r&   )rY   r[   r   )r   r2   r   r   r   s      r   data_2dzTestLogSoftmax.data_2d  sR    JJryy"**y5v>))	* "**V2::*666r!   offsetr   r   c                 d    | j                  |      \  }}||z  }t        t        |      |d       y )Nr   r   )r   r   r
   )r;   r   r   r   r   s        r   test_log_softmax_translationz+TestLogSoftmax.test_log_softmax_translation  s/     LL$	6	VAU;r!   c                     | j                  |      \  }}|j                  |d      }|j                  |d      }t        t        |      |d       y )Nr   r   r   )r   r   r   r
   r;   r   r   r   s       r   test_log_softmax_noneaxisz(TestLogSoftmax.test_log_softmax_noneaxis  sF     LL$	6JJq&!FF+AU;r!   zaxis_2d, expected_2dr   r   r   rE   r   c                     |j                  ddgddgg      }t        t        ||      |j                  ||j                        d       y )Nr(   rN   r-   r&   r   r   )r   r   r
   r'   )r;   axis_2dexpected_2dr   r   s        r   	test_axeszTestLogSoftmax.test_axes  sG    
 JJeR[12AG4

;agg
>U	Lr!   c                 ^    | j                  |      \  }}t        t        |d      |d       y )Nr   r-   r   r   )r   r   r
   r   s       r   test_log_softmax_2d_axis1z(TestLogSoftmax.test_log_softmax_2d_axis1  s'    LL$	6AA.UCr!   c                     | j                  |      \  }}t        t        |j                  d      |j                  d       y )Nr   r-   r   r   )r   r   r
   r   r   s       r   test_log_softmax_2d_axis0z(TestLogSoftmax.test_log_softmax_2d_axis0  s/    LL$	6ACCa0&((Gr!   c                     | j                  |      \  }}|j                  |d      }|j                  |d      }t        t        |d      |d       y )Nr   r   r-   r   r   )r   r   r   r
   r   s       r   test_log_softmax_3dz"TestLogSoftmax.test_log_softmax_3d  sG    LL$	6JJq)$FI.AF3V%Hr!   N)r   r   r   r   r   r   r   r   r   r   staticmethodr   r   r   r   r   r9   r3   r_   r   r   r   r   r   r!   r   r   r     s   = [[""#5>T"U= V=;= 7 7 7 7 [[X3x0< 1<< [[3	
FBFF3K'"''&/)*	
b%[2u+&'6 L	L
DHIr!   r   )	itertoolsr   r   r   numpyr9   scipy._lib._array_apir   r   r   r   scipy._lib._array_api_no_0dr   r   r	   scipy.specialr
   r   r   scipy.special._logsumexpr   r   integral_dtypesr    r   filterwarningsr#   r   r   r   r!   r   <module>r     s       @ @9 9 : 9 2 
MG$& MNNOHI9X? X?  J P OX?v 7AB AB ABH ;PI PI  PIr!   