
    h5              
          S SK JrJr  S SKrS SKrS SKJr  S SKJrJ	r	J
r
  S SKJrJr  S SKJrJrJr  \	" 5       r\R&                  R)                  / SQ/ SQ/ S	Q/S
S9r\R&                  R)                  / SQ/ SQ/ SQ/S
S9rS rS\S\S\S\\\4   4S jr\R2                  R5                  \(       + SS9\R2                  R7                  S/ SQ5      S 5       5       rS rg)    )TuplecastN)assert_allclose)ModelNumpyOps
Softmax_v2)Floats2dInts1d)	has_torchtorch2xpxp2torch)         r   )      r   r   )	      r      fdtype)5w6.?gUכZ?gc<?r   )Y1a ?gM?gˀ珽?r   )gRC?g}Ⱥ?g0?g%7Ŷ?c                     [        SS9n U R                  [        [        5        U " [        SS9u  p[        R
                  " [        SS9   U" [        R                  R                  [        5      5        S S S 5        U " [        SS9u  pU" [        R                  R                  [        5      5      n[        R                  R                  US:H  5      (       d   eg ! , (       d  f       Nn= f)NFnormalize_outputsis_trainzbackprop is not supportedmatchT        )r   
initializeinputsoutputspytestraises
ValueErrorOPSxp
zeros_likeall)model_backpropdXs       Y/home/james-whalen/.local/lib/python3.13/site-packages/thinc/tests/layers/test_softmax.py"test_unnormalized_softmax_backpropr2      s    /E	VW%/KA	z)D	E""7+, 
F .KA	#&&##G,	-B66::bCi     
F	Es   
*C
C*r-   Xtargetsreturnc           
      4   SS K n[        U R                  S5      5      n[        U R                  S5      5      nU R                  S   n[        USS9n[        U5      R	                  5       nXt-  U-   n	X-  n
UR
                  R                  5       nU" X5      nUR                  5         [        [        [        UR
                  R                  R                  U
SS95      5      [        [        [        [        UR                  UR                  5      5      5      4$ )	Nr   Wbsoftmax_temperatureT)requires_grad)dim)torchr   	get_paramattrslongnnCrossEntropyLossbackwardr   r	   r   
functionalsoftmaxTensorgrad)r-   r3   r4   r=   WtbttemperatureXtYt_goldXWbt	XWbt_templossoutputs                r1   torch_softmax_with_temperaturerQ   %   s     	%//#&	'B	%//#&	'B++34K	!4	(Bw$$&GGr>D"I88$$&D)%F
OO(588..66yb6IJHhtELL"'':;<= =    zneeds PyTorch)reasonrJ   )g      ?g      ?g       @c                 x   [        U S S S9n[        R                  R                  SSSSS9R	                  SS	5      n[        R                  S
/5      n[        R                  R                  S	SS9U   nUR                  X$5        [        XU5      u  pVU" USS9u  pxU" Xt-
  5      n	[        XuSS9  [        XSS9  g )Nc                 <    U R                   R                  US   SS9$ )Nr   r   r   )r*   eyeopsshapes     r1   <lambda>*test_softmax_temperature.<locals>.<lambda>B   s    #&&**U1XS*"ArR   c                 6    U R                   R                  USS9$ )Nr   r   )r*   zerosrW   s     r1   rZ   r[   C   s    #&&,,uC,"@rR   )rJ   init_Winit_br;   r   g?r   r   
   r   Tr   g-C6?)atol)
r   r)   r*   arangereshape	asarray1irV   r#   rQ   r   )
rJ   r-   r3   r4   Y_goldYtdXtYr/   r0   s
             r1   test_softmax_temperatureri   =   s     A@E 	b!S,44Q;AmmQC GVVZZ#Z&w/F	Q,Uw?GBD)KA	!*	BA%B$'rR   c                  ^   [         R                  " [        SS9   [        SSS9  S S S 5        [        SS9n SU R                  S'   U R                  [        [        5        [         R                  " [        SS9   U " [        SS9  S S S 5        g ! , (       d  f       Nn= f! , (       d  f       g = f)	Nzsoftmax temperature.*zeror    Fr"   )r   rJ   r   r9   r   )r&   r'   r(   r   r?   r#   r$   r%   )r-   s    r1   !test_reject_incorrect_temperaturerk   T   s    	z)E	FU< 
G /E),EKK%&	VW%	z)E	Ffu% 
G	F 
G	F 
G	Fs   B8B
B
B,)typingr   r   numpyr&   numpy.testingr   	thinc.apir   r   r   thinc.typesr	   r
   
thinc.utilr   r   r   r)   r*   asarrayr$   r%   r2   rQ   markskipifparametrizeri   rk    rR   r1   <module>rw      s       ) 1 1 ( 4 4j	|\B#	N
&&..876
   
!===(.=
8X=0 	M/:8( 9 ;(*&rR   