
    /j5                    n    d dl mZ d dlZd dlZd dlmZ d dlmZ d dlm	Z	 d dl
mZ  G d de          ZdS )	    )annotationsN)TensorProto)Base)expect)make_tensorc                  F   e Zd Zedd            Zedd            Zedd            Zedd            Zedd            Zedd            Z	edd	            Z
edd
            Zedd            Zedd            Zedd            Zedd            Zedd            ZdS )QuantizeLinearreturnNonec                    t           j                            dg ddg          } t          j        g d                              t          j                  }t          j        d          }t          j        d          }t          j        g d                              t          j                  }t          | |||g|gd	
           d S )Nr	   xy_scaley_zero_pointyinputsoutputs)r         i  iir      )r               r   test_quantizelinearr   r   name)	onnxhelper	make_nodenparrayastypefloat32uint8r   noder   r   r   r   s        o/home/longshao/multi-rider-rag/.venv/lib/python3.11/site-packages/onnx/backend/test/case/node/quantizelinear.pyexportzQuantizeLinear.export   s    {$$333E % 
 
 H1112299"*EE*Q--x}}H///0077AAw-C&		
 	
 	
 	
 	
 	
    c                 &   t           j                            dg ddg          } t          j        ddgddgd	d
ggddgddgddggddgddgddggggt          j                  }t          j        g dt          j                  }t          j        g dt          j                  }||                    dddd          z  |                    dddd          z                       t          j                  }t          | |||g|gd           d S )Nr	   r   r   r   i^
   i   iir          i   ii@iii*dtype)r         )T         r   r   test_quantizelinear_axisr   )
r   r    r!   r"   r#   r%   r&   reshaper$   r   r'   s        r)   export_axiszQuantizeLinear.export_axis$   sT   {$$333E % 
 
 H BZ$sCj91X3x"c34[4,t= *	
 	
 	
 (999BJ777xRX>>>Aq!,,,|/C/CAq!Q/O/OOWWH
 
 	w-C+		
 	
 	
 	
 	
 	
r+   c                    t           j                            dg ddg          } t          j        g d                              t          j                  }t          j        d          }t          dt          j	        dgd	g          }t          dt          j	        d
gg d          }t          | |||g|gd           d S )Nr	   r   r   r                 ?       @g     j@g      i@r   r   r   r   r6   )r         ?r   i  `   test_quantizelinear_e4m3fnr   )r   r    r!   r"   r#   r$   r%   r   r   FLOAT8E4M3FNr   r'   s        r)   export_e4m3fnzQuantizeLinear.export_e4m3fnC   s    {$$333E % 
 
 H55566==bjII*Q--">;3KaSSTRUVV[5s<P<P<PQQw-C-		
 	
 	
 	
 	
 	
r+   c                    t           j                            dg ddg          } t          j        g d                              t          j                  }t          j        d          }t          dt          j	        dgd	g          }t          dt          j	        d
gg d          }t          | |||g|gd           d S )Nr	   r   r   r   r>   r   r   r   r?   r6   )r   rB   r   i   rC   test_quantizelinear_e5m2r   )r   r    r!   r"   r#   r$   r%   r   r   
FLOAT8E5M2r   r'   s        r)   export_e5m2zQuantizeLinear.export_e5m2W   s    {$$333E % 
 
 H55566==bjII*Q--">;3IA3QTPUVV[3aS:P:P:PQQw-C+		
 	
 	
 	
 	
 	
r+   c                    t           j                            dg ddg          } t          j        g d                              t          j                  }t          j        d          }t          j        d          }t          j        g d                              t          j                  }t          | |||g|gd	
           d S )Nr	   r   r   r   )r?   g      `      @      333333@333333@g      @g         @     rA     )rT   i      i   i  rU   rV     r   rW   r   test_quantizelinear_uint16r   )	r   r    r!   r"   r#   r$   r%   uint16r   r'   s        r)   export_uint16zQuantizeLinear.export_uint16k   s    {$$333E % 
 
 H  
 
 &

 	
  *S//y''H  
 
 &

 	
" 	w-C-		
 	
 	
 	
 	
 	
r+   c                    t           j                            dg ddg          } t          j        g d                              t          j                  }t          j        d          }t          j        d          }t          j        g d                              t          j                  }t          | |||g|gd	
           d S )Nr	   r   r   r   )r?   g     rL   rM   rN   rO   rP   rQ   g    @g    g    @g    g     @g      rR   rS   rA      )r\        i  r   r^   r_   rT   irT    rT   r`   rT   r`   test_quantizelinear_int16r   )	r   r    r!   r"   r#   r$   r%   int16r   r'   s        r)   export_int16zQuantizeLinear.export_int16   s    {$$333E % 
 
 H  
 
& &

' 	
( *S//x}}H  
 
& &

' 	
* 	w-C,		
 	
 	
 	
 	
 	
r+   c                    t           j                            dg ddgd          } t          j        g dg dg dg                              t          j                  }t          j        g d	t          j        
          }t          dt          j
        |j        t          j        |                    }t          dt          j
        |j        g d          }t          | |||g|gd           d S )Nr	   r   r   r   r   r   axisr?         @g333333@g333333!@ir/      	            (   rA   rL         @r3   r   )r   r   r   r6   r   r   r   r5   r5   r6   r6      test_quantizelinear_uint4r   )r   r    r!   r"   r#   r$   r%   asarrayr   r   UINT4shape	ones_liker   r'   s        r)   export_uint4zQuantizeLinear.export_uint4   s   {$$333E	 % 
 
 H$$$      
 
 &

 	
 *___BJ???"K-w}bl7>S>S
 
 "AG-R-R-R
 
 	w-C,		
 	
 	
 	
 	
 	
r+   c                    t           j                            dg ddgd          } t          j        g dg dg dg                              t          j                  }t          j        g d	t          j        
          }t          dt          j
        |j        t          j        |                    }t          dt          j
        |j        g d          }t          | |||g|gd           d S )Nr	   r   r   r   re   rg   ri   rl   rq   r3   r   )r   r   r   r6   r   r5   r5   r6   r6      test_quantizelinear_int4r   )r   r    r!   r"   r#   r$   r%   ru   r   r   INT4rw   rx   r   r'   s        r)   export_int4zQuantizeLinear.export_int4   s   {$$333E	 % 
 
 H$$$      
 
 &

 	
 *___BJ???"K,gmR\'=R=R
 
 !17,R,R,R
 
 	w-C+		
 	
 	
 	
 	
 	
r+   c                    t           j                            dg ddgd          } t          j        g dg dg dgt          j        	          }t          j        g d
t          j        	          }t          dt          j	        |j
        t          j        |                    }t          dt          j	        |j
        g d          }t          | |||g|gd           d S )Nr	   r   r   r   re   rg   )g       g      r@   rL   )rr         @      @g      @r3   rq   r   )r   r   r   r   r   r   r   r   r   r   r   r   test_quantizelinear_uint2r   )r   r    r!   r"   r#   r%   ru   r   r   UINT2rw   
zeros_liker   r'   s        r)   export_uint2zQuantizeLinear.export_uint2  s   {$$333E	 % 
 
 H$$$&&&$$$
 *
 
 
 *___BJ???"K-w}bmG>T>T
 
 "AG-Q-Q-Q
 
 	w-C,		
 	
 	
 	
 	
 	
r+   c                    t           j                            dg ddgd          } t          j        g dg dg dgt          j        	          }t          j        g d
t          j        	          }t          dt          j	        |j
        t          j        |                    }t          dt          j	        |j
        g d          }t          | |||g|gd           d S )Nr	   r   r   r   re   rg   )g      rM   r@   rA   g       g      g333333g333333!r3   rq   r   )r   r   r   r   r]   r]   r   r   r   r]   r]   test_quantizelinear_int2r   )r   r    r!   r"   r#   r%   ru   r   r   INT2rw   r   r   r'   s        r)   export_int2zQuantizeLinear.export_int26  s   {$$333E	 % 
 
 H$$$&&&(((
 *
 
 
 *___BJ???"K,gmR]7=S=S
 
 !17,U,U,U
 
 	w-C+		
 	
 	
 	
 	
 	
r+   c                    t           j                            dg ddgd          } t          j        g dg dg dg                              t          j                  }t          j        g d	t          j        
          }t          dt          j
        |j        t          j        |                    }t          dt          j
        |j        g d          }t          | |||g|gd           d S )Nr	   r   r   r   re   rg   ri   r   rq   r3   r   )r   r   r   r5   r|   r|   r   r   r   g      r]   r   test_quantizelinear_float4e2m1r   )r   r    r!   r"   r#   r$   r%   ru   r   r   
FLOAT4E2M1rw   r   r   r'   s        r)   export_float4e2m1z QuantizeLinear.export_float4e2m1T  s   {$$333E	 % 
 
 H$$$   (((
 
 &

 	
 *___BJ???""MM'""	
 
 "G777	
 
 	w-C1		
 	
 	
 	
 	
 	
r+   c                 d   t           j                            dg ddgdd          } t          j        g dg dg d	gt          j        
          t          j        ddgddgddggt          j        
          t          j        ddgddgddggt          j        
          }j        |j        k    sJ dt          fdt          t          j                            D                       sJ j                 j                 z  dk    sJ j                 j                 z  }t          j        |          }t          j        ||          }t          j        |z  |z                                 t          j                  }t          | |g|gd           d S )Nr	   r   r   r   r   )r   r   rf   
block_size)r   g      (@g      I@r   r@   g       @rr   r   r?   g      4@g      $@rr   r3         ?rh   rL   皙@ffffff@皙@r   r   c              3  \   K   | ]&}|k    j         |         j         |         k    V  'd S Nrw   .0i
block_axisr   r   s     r)   	<genexpr>z;QuantizeLinear.export_blocked_asymmetric.<locals>.<genexpr>  G       
 
J GAJ'-**
 
r+   repeatsrf   &test_quantizelinear_blocked_asymmetricr   )r   r    r!   r"   r#   r%   r&   rw   allrangelenrepeatrintr$   r   )	r(   r   r   y_scale_elementwisey_zero_point_elementwiser   r   r   r   s	         @@@r)   export_blocked_asymmetricz(QuantizeLinear.export_blocked_asymmetricz  s:   {$$333E % 
 
 H&&&$$$&&&
 *
 
 
 (c
c
c

 *
 
 
 xAAA
 (
 
 
 } 22222
 
 
 
 
 
 
3qw<<((
 
 
 
 
 	
 	
 

 wz"W]:%>>!CCCC'*%z)BB !izRRR#%9'
$
 $
 $
  GA++.FFGGNNrxXXw-C9		
 	
 	
 	
 	
 	
r+   c                 :   t           j                            dddgdgddt          j                  } t          j        g dg d	g d
gt
          j                  t          j        ddgddgddggt
          j                  dt          fdt          t          j                            D                       sJ j                 j                 z  dk    sJ j                 j                 z  }t          j        |          }t          j        t          j        |z            dd                              t
          j                  }t#          dt          j        j        |          }t%          | g|gd           d S )Nr	   r   r   r   r   r   )r   r   rf   r   output_dtype)r   r{   ir   r   r   r3   r   rh   rL   r   r   r   c              3  \   K   | ]&}|k    j         |         j         |         k    V  'd S r   r   r   s     r)   r   z:QuantizeLinear.export_blocked_symmetric.<locals>.<genexpr>  r   r+   r   r   r`   rT   )a_mina_max%test_quantizelinear_blocked_symmetricr   )r   r    r!   r   INT16r"   r#   r%   r   r   r   rw   r   clipr   r$   rb   r   r   )r(   r   r   y_valr   r   r   r   s        @@@r)   export_blocked_symmetricz'QuantizeLinear.export_blocked_symmetric  s   {$$#E$* % 
 
 H###$$$&&&
 *
 
 
 (c
c
c

 *
 
 
 
 
 
 
 
 
 
3qw<<((
 
 
 
 
 	
 	
 

 wz"W]:%>>!CCCC'*%z)BB !izRRRGA++,,F%
 
 

&

 	 G	
 
 	w<C8		
 	
 	
 	
 	
 	
r+   N)r
   r   )__name__
__module____qualname__staticmethodr*   r<   rF   rJ   rZ   rc   ry   r   r   r   r   r   r    r+   r)   r	   r	      s       
 
 
 \
& 
 
 
 \
< 
 
 
 \
& 
 
 
 \
& /
 /
 /
 \/
b 7
 7
 7
 \7
r 
 
 
 \
> 
 
 
 \
> 
 
 
 \
< 
 
 
 \
: #
 #
 #
 \#
J ;
 ;
 ;
 \;
z :
 :
 :
 \:
 :
 :
r+   r	   )
__future__r   numpyr"   r   r   onnx.backend.test.case.baser   onnx.backend.test.case.noder   onnx.helperr   r	   r   r+   r)   <module>r      s    # " " " " "            , , , , , , . . . . . . # # # # # #d
 d
 d
 d
 d
T d
 d
 d
 d
 d
r+   