
    wiE                       d dl mZ d dlZd dlZd dlmZ d dlZd dlZd dlmZ erd dl	m
Z
  G d dej                  Z G d d	ej                  Z G d
 dej                  Z G d dej                  Z G d dej                  Z ej        dddg dddddg ddg           G d dej                              Zedk    r ej                     dS dS )    )annotationsN)TYPE_CHECKING)defs)Sequencec                  6    e Zd Zd	dZd	dZd	dZd	dZd	dZdS )

TestSchemareturnNonec                    t          j        d          }|                     |j        t           j        j        j                   d S )NRelu)r   
get_schemaassertEqualnode_determinismOpSchemaNodeDeterminismDeterministic)selfrelu_schemas     e/root/.openclaw/workspace/chatterbox_venv_py311/lib/python3.11/site-packages/onnx/test/schema_test.pytest_get_schemazTestSchema.test_get_schema   sC    of--($-*G*U	
 	
 	
 	
 	
    c                .    t          j        d           d S )NConv)r   r   r   s    r   test_typecheckzTestSchema.test_typecheck   s    r   c                    t          j        d          j        d         j        }|                     t          |          t          j                   |                     |j        t          j        j                   d S )NBatchNormalizationepsilon)	r   r   
attributesdefault_valuer   typeonnxAttributeProtoFLOAT)r   vs     r   test_attr_default_valuez"TestSchema.test_attr_default_value   s`    O011<YGUa$"5666!4!:;;;;;r   c                    t          j        d          }|                     t          |j                  t
          j                   |                     |j        t           j        j	        j
                   d S )NSelu)r   r   r   r!   function_bodyr"   FunctionProtor   r   r   r   )r   selu_schemas     r   test_function_bodyzTestSchema.test_function_body"   sg    of--k788$:LMMM($-*G*U	
 	
 	
 	
 	
r   c                   t          j        d          }|                     |j        t           j        j        j                   |                     |j                   t          j        d          }|                     |j        t           j        j        j	                   | 
                    |j                   t          j        d          }|                     |j        t           j        j        j	                   | 
                    |j                   t          j        d          }|                     |j        t           j        j        j                   |                     |j                   d S )NRandomNormalLiker   CastLikeIf)r   r   r   r   r   r   NonDeterministic
assertTruenon_deterministicr   assertFalse)r   rand_schema	bn_schemacast_like_schema	if_schemas        r   test_node_determinismz TestSchema.test_node_determinism)   sJ   o&899($-*G*X	
 	
 	
 	5666O$899	&(E(S	
 	
 	
 	4555?:66-M)7	
 	
 	
 	);<<<OD))	&M):	
 	
 	
 		344444r   Nr	   r
   )__name__
__module____qualname__r   r   r&   r,   r9    r   r   r   r      sx        
 
 
 
       < < < <

 
 
 
5 5 5 5 5 5r   r   c                  P    e Zd Zd ZddZddZddZd Zd Zd	 Z	d
 Z
d Zd ZdS )TestOpSchemac                    t          j        ddd          }|                     |t           j                   |                     |j        t           j        j        j                   d S Ntest_optest_domain   )r   r   assertIsInstancer   r   r   r   r   schemas     r   	test_initzTestOpSchema.test_initC   s_    y-;;fdm444#T]%B%P	
 	
 	
 	
 	
r   r	   r
   c           	        t          j        dddt           j                            dd          gddgdfg          }|                     |j        d           |                     |j        d           |                     |j        d           |                     t          |j                  d           |                     |j        d	         j        d           |                     |j        d	         j	        d           |                     t          |j
                  d           |                     |j
        d	         j        d           |                     |j
        d	         j        dg           d S )
NrC   rD   rE   input1Ttensor(int64) inputstype_constraintsr   )r   r   FormalParameterr   namedomainsince_versionlenrP   type_strrQ   type_param_strallowed_type_strsr   	op_schemas     r   test_init_with_inputsz"TestOpSchema.test_init_with_inputsK   so   MM11(C@@A"_$5r:;
 
 
	 	333)=9990!444Y-..222)!,18<<<)!,5s;;;Y788!<<<3A6EsKKK&q);o=N	
 	
 	
 	
 	
r   c                   t          j        dddt           j                            dd          t           j                            dd          gt           j                            dd          t           j                            dd          gdd	gd
fgt           j                            dt           j        j        j        d          g          }|                     t          |j                  d           |                     |j        d         j	        d           |                     |j        d         j
        d           |                     |j        d         j	        d           |                     |j        d         j
        d           |                     t          |j                  d           |                     |j        d         j	        d           |                     |j        d         j
        d           |                     |j        d         j	        d           |                     |j        d         j
        d           |                     t          |j                  d           |                     |j        d         j        d           |                     |j        d         j        d	g           |                     t          |j                  d           |                     |j        d         j	        d           |                     |j        d         j        t           j        j        j                   |                     |j        d         j        d           d S )NrC   rD   rE   rK   rL   input2output1output2rM   rN   attr1attr1 descriptionrP   outputsrQ   r      r   )r   r   rR   	AttributeAttrTypeINTSr   rV   rP   rS   rW   rd   rQ   rX   rY   r   r!   descriptionrZ   s     r   +test_init_creates_multi_input_output_schemaz8TestOpSchema.test_init_creates_multi_input_output_schema_   s   M--h<<--h<<
 --i==--i==  #_$5r:;''T]38:M 
 
 
	& 	Y-..222)!,18<<<)!,5s;;;)!,18<<<)!,5s;;;Y.//333*1-2I>>>*1-6<<<*1-2I>>>*1-6<<<Y788!<<<3A6EsKKK&q);o=N	
 	
 	
 	Y122A666-g6;WEEE ).0F0K	
 	
 	
 	-g6BDWXXXXXr   c                   t          j        ddd          }|                     |j        d           |                     |j        d           |                     |j        d           |                     t          |j                  d           |                     t          |j                  d           |                     t          |j	                  d           d S )NrC   rD   rE   r   )
r   r   r   rS   rT   rU   rV   rP   rd   rQ   rZ   s     r   $test_init_without_optional_argumentsz1TestOpSchema.test_init_without_optional_arguments   s    M)]A>>	333)=9990!444Y-..222Y.//333Y788!<<<<<r   c                   |                      t                    5  t          j        dd           d d d            n# 1 swxY w Y   |                      t                    5  t          j        ddd           d d d            n# 1 swxY w Y   t          j        ddd          }|                     |j        d           d S )NrD   rE   )rT   rU   {   rC   )assertRaises	TypeErrorr   r   r   rS   rG   s     r   	test_namezTestOpSchema.test_name   sA   y)) 	A 	AMa@@@@	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	Ay)) 	1 	1M#}a000	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 y-;;i00000!   >AA#BB
B
c                   |                      t                    5  t          j        dd           d d d            n# 1 swxY w Y   |                      t                    5  t          j        ddd           d d d            n# 1 swxY w Y   t          j        ddd          }|                     |j        d           d S )NrC   rE   )rS   rU   rn   rD   )ro   rp   r   r   r   rT   rG   s     r   rD   zTestOpSchema.test_domain   s0   y)) 	; 	;My::::	; 	; 	; 	; 	; 	; 	; 	; 	; 	; 	; 	; 	; 	; 	;y)) 	- 	-M)S!,,,	- 	- 	- 	- 	- 	- 	- 	- 	- 	- 	- 	- 	- 	- 	- y-;;66666rr   c                    |                      t                    5  t          j        dd           d d d            n# 1 swxY w Y   t          j        ddd          }|                     |j        d           d S rB   )ro   rp   r   r   r   rU   rG   s     r   test_since_versionzTestOpSchema.test_since_version   s    y)) 	4 	4M)]333	4 	4 	4 	4 	4 	4 	4 	4 	4 	4 	4 	4 	4 	4 	4 y-;;-q11111s   =AAc                l    t          j        dddd          }|                     |j        d           d S )NrC   rD   rE   test_doc)doc)r   r   r   rx   rG   s     r   rw   zTestOpSchema.test_doc   s8    y-
KKKZ00000r   c           	        t           j                            ddd          g}t          j        ddd|ddgd	fg
          }|                     t	          |j                  d           |                     |j        d         j        d           |                     |j        d         j        d           |                     |j        d         j        d           d S )NrK   rL   The first input.rS   rW   ri   rC   rD   rE   rM   rN   rO   r   )	r   r   rR   r   rV   rP   rS   rW   ri   )r   rP   rH   s      r   test_inputszTestOpSchema.test_inputs   s     M))9K *  

 "_$5r:;
 
 
 	V]++Q///q).999q)2C888q)57IJJJJJr   c           	        t           j                            ddd          g}t          j        ddd|ddgd	fg
          }|                     t	          |j                  d           |                     |j        d         j        d           |                     |j        d         j        d           |                     |j        d         j        d           d S )Nr_   rL   zThe first output.r{   rC   rD   rE   rM   rN   )rd   rQ   r   )	r   r   rR   r   rV   rd   rS   rW   ri   )r   rd   rH   s      r   test_outputszTestOpSchema.test_outputs   s     M)):M *  
 "_$5r:;
 
 
 	V^,,a000*/;;;*3S999*68KLLLLLr   Nr:   )r;   r<   r=   rI   r\   rj   rl   rq   rD   ru   rw   r|   r~   r>   r   r   r@   r@   B   s        
 
 

 
 
 
((Y (Y (Y (YT= = = =1 1 17 7 72 2 21 1 1K K K(M M M M Mr   r@   c                      e Zd Zd ZdS )TestFormalParameterc           	     8   d}d}d}t           j        j        j        }d}d}t           j        j        j        }t           j                            |||||||          }|                     |j        |           |                     |j	        |           |                     |j
        |           |                     |j        |           |                     |j        |           |                     |j        |           |                     |j        |           d S )NrK   tensor(float)rz   TrE   )param_optionis_homogeneous	min_aritydifferentiation_category)r   r   FormalParameterOptionSingleDifferentiationCategoryUnknownrR   r   rS   rW   ri   optionr   r   r   )	r   rS   rW   ri   r   r   r   r   formal_parameters	            r   rI   zTestFormalParameter.test_init   s1   "(}:A	#'=#H#P =88%)%= 9 
 
 	).555)2H===)5{CCC)0,???)8.III)3Y???57O	
 	
 	
 	
 	
r   N)r;   r<   r=   rI   r>   r   r   r   r      s#        
 
 
 
 
r   r   c                  r    e Zd Zej                            dddgdfddddgdfdg          dd            ZdS )TestTypeConstraintParamsingle_typerL   r   Test descriptiondouble_typesrM   )tuplerL   )r   rM   r   _strrX   allowed_typesSequence[str]ri   r	   r
   c                   t           j                            |||          }|                     |j        |           |                     |j        t          |                     |                     |j        |           d S N)r   r   TypeConstraintParamr   ri   rY   listrX   )r   r   rX   r   ri   type_constraints         r   rI   z!TestTypeConstraintParam.test_init   s}    & -;;M;
 
 	4kBBB:D<O<OPPP7HHHHHr   N)
r   r   rX   r   r   r   ri   r   r	   r
   )r;   r<   r=   parameterizedexpandrI   r>   r   r   r   r      s         ''C/!24FG /2"	 S		
 I I I I I Ir   r   c                      e Zd Zd Zd ZdS )TestAttributec                   d}t           j        j        j        }d}t           j                            |||          }|                     |j        |           |                     |j        |           |                     |j        |           d S )N	test_attrzTest attribute)	r   r   rg   STRINGSrf   r   rS   r!   ri   )r   rS   type_ri   	attributes        r   rI   zTestAttribute.test_init  s    &.&M++D%EE	...///.<<<<<r   c                r   t          j        d          j        d         j        }|                     |t
          j                   t           j                            d|d          }| 	                    ||j                   | 	                    d|j
                   | 	                    d|j                   d S )Nr   r   ra   rb   )r   r   r   r    rF   r"   r#   r   rf   r   rS   ri   )r   r    r   s      r   test_init_with_default_valuez*TestAttribute.test_init_with_default_value   s    O011<YGU 	 	mT-@AAAM++G]DWXX		(?@@@).111,i.CDDDDDr   N)r;   r<   r=   rI   r   r>   r   r   r   r     s7        = = =E E E E Er   r   CustomOp   rN   )rE   re         )op_type
op_version	op_domaintrap_op_versiontestc                  l    e Zd ZU ded<   ded<   ded<   ded<   dd
ZddZd Zd Zd Zd Z	d Z
d ZdS )TestOpSchemaRegisterr   r   intr   r   z	list[int]r   r	   r
   c                    |                      t          j                            | j        | j                             d S r   )r4   r"   r   hasr   r   r   s    r   setUpzTestOpSchemaRegister.setUpD  s0    t|T^DDEEEEEr   c                    g | j         | j        D ]h}t          j        t          j        j                  5  t          j                            | j        || j	                   d d d            n# 1 swxY w Y   id S r   )
r   r   
contextlibsuppressr"   r   SchemaErrorderegister_schemar   r   )r   versions     r   tearDownzTestOpSchemaRegister.tearDownH  s    ?-?t? 	S 	SG$TY%:;; S S	++DL'4>RRRS S S S S S S S S S S S S S S	S 	Ss   ,A..A2	5A2	c                B   g | j         | j        D ]}t          j        | j        | j        |          }t          j                            |           |                     t          j        	                    | j        || j                             g | j         | j        D ]}}t          j        
                    |j        ||j                  }t          j        | j        | j        |          }|                     t          |          t          |                     ~d S r   )r   r   r   r   r   r   r"   register_schemar2   r   r   rS   rT   r   r   )r   r   r[   registered_ops       r   test_register_multi_schemaz/TestOpSchemaRegister.test_register_multi_schemaN  s   ?-?t? 	R 	RG I
 I%%i000OODIMM$,PPQQQQ?-?t? 
	A 
	AG I00)9 M  I
 S//Y@@@@
	A 
	Ar   c                   d| j          d| j         d| j          d| j         d	}t          j                            |          }t          j        | j        | j         | j        t          j                            dd          t          j                            dd	          gt          j                            d
d          gddgdfgt          j        	                    dt          j        j
        j        d          g          }|                     t          j        j                  5  t          j                            |d           d d d            n# 1 swxY w Y   t          j                            |           | j        D ]`}t          j                            t          j        | j        | j         |t          j                            d
d	          g                     at          j                            |d           d S )Nzc
            <
                ir_version: 7,
                opset_import: [
                    "z" : z
                ]
            >
            agraph (float[N, 128] X, int32 Y) => (float[N] Z)
            {
                Z = .z-<attr1=[1,2]>(X, Y)
            }
           rK   rL   r^   int32r_   r   rN   ra   rb   rc   T)check_custom_domain)rd   )r   r   r   r"   parserparse_modelr   r   rR   rf   rg   rh   ro   checkerValidationErrorcheck_modelr   r   )r   inputmodelr[   r   s        r   .test_using_the_specified_version_in_onnx_checkzCTestOpSchemaRegister.test_using_the_specified_version_in_onnx_checkc  sV    n	  +//	  ^  '+l   ''..MLNO--h<<--h@@
 --i==  #_$5r:;''T]38:M 
 
 
	$ t|;<< 	F 	FL$$U$EEE	F 	F 	F 	F 	F 	F 	F 	F 	F 	F 	F 	F 	F 	F 	F	!!),,,+ 
	 
	GI%%LN55iII	  	 	 	 	 	  D AAAAAs   '"EEEc                F   t          j        | j        | j        | j                  }t
          j                             |           |                     t
          j         j                  5  t
          j                             |           d d d            d S # 1 swxY w Y   d S r   )	r   r   r   r   r   r"   r   ro   r   rZ   s     r   Atest_register_schema_raises_error_when_registering_a_schema_twicezVTestOpSchemaRegister.test_register_schema_raises_error_when_registering_a_schema_twice  s    MLNO
 
	
 		!!),,,ty455 	1 	1I%%i000	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1s   ) BBBc                   g | j         | j        D ]}t          j        | j        | j        |          }t          j                            |           |                     t          j        	                    |j
        ||j                             t          j                            |j
        | j        |j                   | j         D ]@}|                     t          j        	                    |j
        ||j                             At          j        	                    |j
        | j        |j                  rRt          j                            |j
        | j        |j                  }|                     |j        | j                   d S d S r   )r   r   r   r   r   r   r"   r   r2   r   rS   rT   r   r   
assertLessrU   )r   r   r[   rH   s       r   $test_deregister_the_specified_schemaz9TestOpSchemaRegister.test_deregister_the_specified_schema  sY   ?-?t? 	V 	VG I
 I%%i000OODIMM).'9CSTTUUUU	##INDOYEUVVV+ 	V 	VGOODIMM).'9CSTTUUUU9==):JKK 	CY))1A F OOF0$/BBBBB		C 	Cr   c                    |                      t          j        j                  5  t          j                            | j        | j        | j                   d d d            d S # 1 swxY w Y   d S r   )ro   r"   r   r   r   r   r   r   r   s    r   @test_deregister_schema_raises_error_when_opschema_does_not_existzUTestOpSchemaRegister.test_deregister_schema_raises_error_when_opschema_does_not_exist  s    ty455 	W 	WI''dot~VVV	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	Ws   1A##A'*A'c                   t          j        | j        | j        | j                  t
          j                                        t
          j                             j        j	        j
                  }t
          j                             j        j
                  }fd} |t
          j                                                   } |t
          j                                                   }|                     t          |          d           |                     t          |          d           |                     t          |          t                               |                     t          |          t                               |                     t          |d                   t                               |                     t          |d                   t                               d S )Nc                     fd| D             S )Nc                4    g | ]}|j         j         k    |S r>   )rS   ).0opr[   s     r   
<listcomp>znTestOpSchemaRegister.test_legacy_schema_accessible_after_deregister.<locals>.filter_schema.<locals>.<listcomp>  s'    FFF2BGy~,E,EB,E,E,Er   r>   )schemasr[   s    r   filter_schemazZTestOpSchemaRegister.test_legacy_schema_accessible_after_deregister.<locals>.filter_schema  s    FFFFFFFFr   rE   r   )r   r   r   r   r   r"   r   r   rS   rU   rT   get_all_schemasget_all_schemas_with_historyr   rV   r   )r   schema_aschema_br   schema_cschema_dr[   s         @r   .test_legacy_schema_accessible_after_deregisterzCTestOpSchemaRegister.test_legacy_schema_accessible_after_deregister  s   MLNO
 
	
 		!!),,,9''NI3Y5E
 
 9''		8HII	G 	G 	G 	G 	G !=!:!:!<!<== =!G!G!I!IJJX***X***XI777XI777Xa[))3y>>:::Xa[))3y>>:::::r   Nr:   )r;   r<   r=   __annotations__r   r   r   r   r   r   r   r   r>   r   r   r   r   +  s         & LLLOOONNNF F F FS S S SA A A*/B /B /Bb1 1 1C C C&W W W; ; ; ; ;r   r   __main__)
__future__r   r   unittesttypingr   r   r"   r   collections.abcr   TestCaser   r@   r   r   r   parameterized_classr   r;   mainr>   r   r   <module>r      s;   # " " " " "                             )((((((,5 ,5 ,5 ,5 ,5" ,5 ,5 ,5^VM VM VM VM VM8$ VM VM VMr
 
 
 
 
(+ 
 
 
<I I I I Ih/ I I I8E E E E EH% E E E, #" "+||		
 	
 "+||		
 	
 $O; O; O; O; O;8, O; O;% $O;d zHMOOOOO r   