
    wi 2                       d dl mZ d dlZd dlZd dlZd dlZd dlZd dlZd dl	Zd dl
Z
d dlZd dlmZ ddZddZ e
j        d	d
id	did	did	dig           G d dej                              Z e
j        d	d
id	did	dig           G d dej                              Z G d dej                  Z G d dej                  Zedk    r ej                     dS dS )    )annotationsN)serializationreturnonnx.ModelProtoc                 t    t          j                    } t           j        | _        d| _        d| j        _        | S )Nz	onnx-testtest)onnx
ModelProto
IR_VERSION
ir_versionproducer_namegraphname)models    d/root/.openclaw/workspace/chatterbox_venv_py311/lib/python3.11/site-packages/onnx/test/basic_test.py_simple_modelr      s0    OEE%EEKL    onnx.TensorProtoc            	         t           j                            dt           j        j        dd t          d          D                       S )Nztest-tensor)         c                    g | ]}|d z   S )g      ? ).0xs     r   
<listcomp>z"_simple_tensor.<locals>.<listcomp>!   s    )))!a#g)))r      )r   	data_typedimsvals)r	   helpermake_tensorTensorProtoFLOATranger   r   r   _simple_tensorr'      sG    ;"""())uRyy)))	 #   r   formatprotobuf	textprotojsononnxtxtc                  :    e Zd ZU ded<   d
dZd
dZd
dZd
dZd	S )TestIOstrr(   r   Nonec                    t                      }t          j                            | j                                      |          }t          j        || j                  }|                     ||           d S Nr(   )	r   r   registrygetr(   serialize_protor	   load_model_from_stringassertEqualselfprotoproto_stringloaded_protos       r   #test_load_model_when_input_is_bytesz*TestIO.test_load_model_when_input_is_bytes0   sd    $-11$+>>NNuUU2<TTT-----r   c                   t                      }t          j                            d                              |          }t          j                    }t          j        ||| j	                   t          j
        t          j        |                                          | j	                  }|                     ||           d S )Nr)   r3   )r   r   r4   r5   r6   ioBytesIOr	   
save_modelr(   
load_modelgetvaluer8   )r:   r;   r<   fr=   s        r   5test_save_and_load_model_when_input_has_read_functionz<TestIO.test_save_and_load_model_when_input_has_read_function6   s     %-11*==MMeTTJLLa<<<<rz!**,,'?'?TTT-----r   c                T   t                      }t          j                    5 }t          j                            |d          }t          j        ||| j                   t          j	        || j                  }| 
                    ||           d d d            d S # 1 swxY w Y   d S Nz
model.onnxr3   )r   tempfileTemporaryDirectoryospathjoinr	   rB   r(   rC   r8   r:   r;   temp_dir
model_pathr=   s        r   0test_save_and_load_model_when_input_is_file_namez7TestIO.test_save_and_load_model_when_input_is_file_nameA   s    (** 	2hh==JOE:dkBBBB?:dkJJJLUL111		2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2   A.BB!$B!c                >   t                      }t          j                    5 }t          j        |d          }t          j        ||| j                   t          j        || j                  }| 	                    ||           d d d            d S # 1 swxY w Y   d S rH   )
r   rI   rJ   pathlibPathr	   rB   r(   rC   r8   rN   s        r   /test_save_and_load_model_when_input_is_pathlikez6TestIO.test_save_and_load_model_when_input_is_pathlikeI   s    (** 	2h h==JOE:dkBBBB?:dkJJJLUL111		2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2   A#BBBNr   r0   )__name__
__module____qualname____annotations__r>   rF   rQ   rV   r   r   r   r.   r.   %   sq          KKK. . . .	. 	. 	. 	.2 2 2 22 2 2 2 2 2r   r.   c                  >    e Zd ZU dZded<   ddZddZddZdd	Zd
S )TestIOTensorz'Test loading and saving of TensorProto.r/   r(   r   r0   c                    t                      }t          j                            | j                                      |          }t          j        || j                  }|                     ||           d S r2   )	r'   r   r4   r5   r(   r6   r	   load_tensor_from_stringr8   r9   s       r   $test_load_tensor_when_input_is_bytesz1TestIOTensor.test_load_tensor_when_input_is_bytes_   sf      $-11$+>>NNuUU3LUUU-----r   c                *   t                      }t          j                    }t          j        ||| j                   t          j        t          j        |                                          | j                  }|                     ||           d S r2   )	r'   r@   rA   r	   save_tensorr(   load_tensorrD   r8   )r:   r;   rE   r=   s       r   6test_save_and_load_tensor_when_input_has_read_functionzCTestIOTensor.test_save_and_load_tensor_when_input_has_read_functione   sw      JLL$+6666'
1::<<(@(@UUU-----r   c                T   t                      }t          j                    5 }t          j                            |d          }t          j        ||| j                   t          j	        || j                  }| 
                    ||           d d d            d S # 1 swxY w Y   d S rH   )r'   rI   rJ   rK   rL   rM   r	   rc   r(   rd   r8   rN   s        r   1test_save_and_load_tensor_when_input_is_file_namez>TestIOTensor.test_save_and_load_tensor_when_input_is_file_namem   s      (** 	2hh==JUJt{CCCC+Jt{KKKLUL111		2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2rR   c                >   t                      }t          j                    5 }t          j        |d          }t          j        ||| j                   t          j        || j                  }| 	                    ||           d d d            d S # 1 swxY w Y   d S rH   )
r'   rI   rJ   rT   rU   r	   rc   r(   rd   r8   rN   s        r   0test_save_and_load_tensor_when_input_is_pathlikez=TestIOTensor.test_save_and_load_tensor_when_input_is_pathlikev   s      (** 	2h h==JUJt{CCCC+Jt{KKKLUL111		2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2rW   NrX   )	rY   rZ   r[   __doc__r\   ra   re   rg   ri   r   r   r   r^   r^   R   sw          21KKK. . . .. . . .2 2 2 22 2 2 2 2 2r   r^   c                  >    e Zd Zd
dZd
dZd
dZd
dZd
dZd
dZd	S )TestSaveAndLoadFileExtensionsr   r0   c                <   t                      }t          j                    5 }t          j                            |d          }t          j        ||           t          j        |d          }| 	                    ||           d d d            d S # 1 swxY w Y   d S Nmodel.textprotor*   r3   
r   rI   rJ   rK   rL   rM   r	   rB   rC   r8   rN   s        r   3test_save_model_picks_correct_format_from_extensionzQTestSaveAndLoadFileExtensions.test_save_model_picks_correct_format_from_extension   s    (** 	2hh0ABBJOE:...?:kJJJLUL111	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2   A"BBBc                <   t                      }t          j                    5 }t          j                            |d          }t          j        ||d           t          j        |          }| 	                    ||           d d d            d S # 1 swxY w Y   d S rn   rp   rN   s        r   3test_load_model_picks_correct_format_from_extensionzQTestSaveAndLoadFileExtensions.test_load_model_picks_correct_format_from_extension   s    (** 	2hh0ABBJOE:kBBBB?:66LUL111	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2 	2rr   c                   t                      }t          j                    5 }t          j                            |d          }t          j        ||d           t          j        |d          }| 	                    ||           | 
                    t          j        j        j                  5  t          j        |           d d d            n# 1 swxY w Y   d d d            d S # 1 swxY w Y   d S )Nro   r)   r3   )r   rI   rJ   rK   rL   rM   r	   rB   rC   r8   assertRaisesgoogler)   text_format
ParseErrorrN   s        r   0test_save_model_uses_format_when_it_is_specifiedzNTestSaveAndLoadFileExtensions.test_save_model_uses_format_when_it_is_specified   sS   (** 	,hh0ABBJOE:jAAAA?:jIIILUL111""6?#>#IJJ , ,
+++, , , , , , , , , , , , , , ,	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	,s6   BC(/CC(C	C(C	C((C,/C,c                   t                      }t          j                    5 }t          j                            |d          }t          j        ||           |                     t          j
        j        j                  5  t          j        |d           d d d            n# 1 swxY w Y   t          j        |d          }|                     ||           d d d            d S # 1 swxY w Y   d S )Nzmodel.protobufr*   r3   r)   r   rI   rJ   rK   rL   rM   r	   rB   rv   rw   r)   rx   ry   rC   r8   rN   s        r   0test_load_model_uses_format_when_it_is_specifiedzNTestSaveAndLoadFileExtensions.test_load_model_uses_format_when_it_is_specified   sh   (** 		2hh0@AAJOE:...""6?#>#IJJ @ @ 
;????@ @ @ @ @ @ @ @ @ @ @ @ @ @ @
  ?:jIIILUL111		2 		2 		2 		2 		2 		2 		2 		2 		2 		2 		2 		2 		2 		2 		2 		2 		2 		2s6   AC(B$C($B(	(C(+B(	,/C((C,/C,c                   t                      }t          j                    5 }t          j                            |d          }t          j        ||d           |                     t          j
        j        j                  5  t          j        |           d d d            n# 1 swxY w Y   t          j        |d          }|                     ||           d d d            d S # 1 swxY w Y   d S )Nr   r*   r3   )r   rI   rJ   rK   rL   rM   r	   rB   rv   rw   r)   messageDecodeErrorrC   r8   rN   s        r   Ftest_load_and_save_model_to_path_without_specifying_extension_succeedszdTestSaveAndLoadFileExtensions.test_load_and_save_model_to_path_without_specifying_extension_succeeds   sX    (** 
	2hh88JOE:kBBBB""6?#:#FGG , , 
+++, , , , , , , , , , , , , , ,
  ?:kJJJLUL111
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2 
	2s6   A!C(B$C($B(	(C(+B(	,/C((C,/C,c                >   t                      }t          j                    5 }t          j                            |d          }t          j        ||           |                     t          j
        j        j                  5  t          j        |d           d d d            n# 1 swxY w Y   t          j        |          }|                     ||           t          j        |d          }|                     ||           d d d            d S # 1 swxY w Y   d S )Nr   r*   r3   r)   r|   )r:   r;   rO   rP   r=   #loaded_proto_as_explicitly_protobufs         r   Ttest_load_and_save_model_without_specifying_extension_or_format_defaults_to_protobufzrTestSaveAndLoadFileExtensions.test_load_and_save_model_without_specifying_extension_or_format_defaults_to_protobuf   s    (** 	Ihh88JOE:...""6?#>#IJJ @ @
;????@ @ @ @ @ @ @ @ @ @ @ @ @ @ @  ?:66LUL11126/:3 3 3/ U$GHHH	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	Is7   ADB$D$B(	(D+B(	,ADDDNrX   )	rY   rZ   r[   rq   rt   rz   r}   r   r   r   r   r   rl   rl      s        2 2 2 22 2 2 2
, 
, 
, 
,2 2 2 22 2 2 2 I I I I I Ir   rl   c                  .    e Zd ZddZddZddZddZdS )	TestBasicFunctionsr   r0   c                f    t           j        }t           j        }t           j        }t           j        }d S )N)r	   AttributeProto	NodeProto
GraphProtor
   )r:   _s     r   test_protos_existz$TestBasicFunctions.test_protos_exist   s#    NOOr   c                   t          j                    }|                     |                    d                     t           j        |_        |                                }|                    |           |                     |                    d                     | 	                    |j        t           j                   d S )Nr   )
r	   r
   assertFalseHasFieldr   r   SerializeToStringParseFromString
assertTruer8   )r:   r   model_strings      r   test_version_existsz&TestBasicFunctions.test_version_exists   s    !!55666  ?..00l+++|44555)4?;;;;;r   c                L   t                      }t          |          }|                     |d           d}t          j                            |          }t          |          }|                     |d           t          |j                  }|                     |d           d S )NzNModelProto(ir_version=13, producer_name='onnx-test', graph=GraphProto('test'))a  
           <
             ir_version: 10,
             opset_import: [ "" : 19]
           >
           agraph (float[N] X) => (float[N] C)
           <
             float[1] weight = {1}
           >
           {
              C = Cast<to=1>(X)
           }
        zModelProto(ir_version=10, opset_import={'': 19}, graph=GraphProto('agraph', input=<1 inputs>, output=<1 outputs>, initializer=<1 initializers>, node=<1 nodes>))zhGraphProto('agraph', input=<1 inputs>, output=<1 outputs>, initializer=<1 initializers>, node=<1 nodes>))r   reprr8   r	   parserparse_modelr   )r:   r   
model_repr
text_model
graph_reprs        r   test_model_and_graph_reprz,TestBasicFunctions.test_model_and_graph_repr   s    %[[
\	
 	
 	


 ''
33%[[
 o	
 	
 	

 %+&&
v	
 	
 	
 	
 	
r   c                    d}t           j                            |          }|                     t	          |          d           t	          |j        d                   }|                     |d           d S )Nac  
            <
            ir_version: 9,
            opset_import: [ "" : 15, "custom_domain" : 1],
            producer_name: "FunctionProtoTest",
            producer_version: "1.0",
            model_version: 1,
            doc_string: "A test model for model local functions."
          >
         agraph (float[N] x) => (float[N] out)
         {
            out = custom_domain.Selu<alpha=2.0, gamma=3.0>(x)
         }
         <
         domain: "custom_domain",
         opset_import: [ "" : 15],
         doc_string: "Test function proto"
         >
           Selu
           <alpha: float=1.67326319217681884765625, gamma: float=1.05070102214813232421875>
           (X) => (C)
           {
               constant_alpha = Constant<value_float: float=@alpha>()
               constant_gamma = Constant<value_float: float=@gamma>()
               alpha_x = CastLike(constant_alpha, X)
               gamma_x = CastLike(constant_gamma, X)
               exp_x = Exp(X)
               alpha_x_exp_x = Mul(alpha_x, exp_x)
               alpha_x_exp_x_ = Sub(alpha_x_exp_x, alpha_x)
               neg = Mul(gamma_x, alpha_x_exp_x_)
               pos = Mul(gamma_x, X)
               _zero = Constant<value_float=0.0>()
               zero = CastLike(_zero, X)
               less_eq = LessOrEqual(X, zero)
               C = Where(less_eq, neg, pos)
           }
        zModelProto(ir_version=9, opset_import={'': 15, 'custom_domain': 1}, producer_name='FunctionProtoTest', producer_version='1.0', graph=GraphProto('agraph', input=<1 inputs>, output=<1 outputs>, node=<1 nodes>), functions=<1 functions>)r   z{FunctionProto('Selu', domain='custom_domain', opset_import={'': 15}, input=<1 inputs>, output=<1 outputs>, node=<13 nodes>))r	   r   r   r8   r   	functions)r:   textr   function_reprs       r   test_function_reprz%TestBasicFunctions.test_function_repr  s    $J ''--KK x	
 	
 	
 U_Q/00 J	
 	
 	
 	
 	
r   NrX   )rY   rZ   r[   r   r   r   r   r   r   r   r   r      se           < < < <!
 !
 !
 !
F/
 /
 /
 /
 /
 /
r   r   __main__)r   r   )r   r   )
__future__r   r@   rK   rT   rI   unittestgoogle.protobuf.messagerw   google.protobuf.text_formatparameterizedr	   r   r   r'   parameterized_classTestCaser.   r^   rl   r   rY   mainr   r   r   <module>r      s!   # " " " " " 				 				        " " " "                  #"	:	;	6	9	 "2 "2 "2 "2 "2X "2 "2 "2J #"	:	;	6 #2 #2 #2 #2 #28$ #2 #2 #2LMI MI MI MI MIH$5 MI MI MI`g
 g
 g
 g
 g
* g
 g
 g
T zHMOOOOO r   