
    Si\#                        S SK Jr  S SKrS SKJr  S SKrS SKJr  S SKJ	r	J
r
Jr   " S S\5      r\S:X  a  \R                  " 5         gg)	    )annotationsN)TestShapeInferenceHelper)TensorProto)	make_nodemake_tensormake_tensor_value_infoc                  x    \ rS rSrSS jrSS jrSS jrSS jrSS jrSS jr	SS jr
SS	 jrSS
 jrSS jrSrg)TestDataPropagation   c           	     6   U R                  S[        R                  S4S[        R                  S4/[        SS/S/5      [        SSS/S/5      // 5      nU R	                  U[        S[        R                  S	5      [        S[        R                  S
5      /SS9  g )Nx)         y)r      r   ShapeshapeExpandzr   )r   r   r   T	data_prop)_make_graphr   INT32r   _assert_inferredr   INT64selfgraphs     ^/var/www/html/ai-image-ml/venv/lib/python3.13/site-packages/onnx/test/data_propagation_test.pytest_expand_symbolic_input.TestDataPropagation.test_expand_symbolic_input   s      ;$$i038I8I92UV'C57)4(S'NSE: 
 	&w0A0A4H&sK,=,=yI  	 	
    c                B   U R                  S[        R                  S4/[        SS/S/5      [        SS/S/[	        S[        R
                  SS	5      S
9// 5      nU R                  U[        S[        R                  S5      [        S[        R
                  S5      /SS9  g )Nr   r   r      r   r   ConstantOfShaper   valuer   r   r)   r   Tr   	r   r   FLOATr   r   r   r   r   r   r   s     r!   (test_constantofshape_with_symbolic_shape<TestDataPropagation.test_constantofshape_with_symbolic_shape#   s      ;$$i01'C57)4%IE%g{/@/@$M	 
 	&w0A0A4H&sK,=,=yI  	 	
r$   c           	         [         R                  R                  S5      nU R                  U[	        S[
        R                  S5      [	        S[
        R                  S5      /SS9  g)	z8Infer the shape of z by propagating the value of xshape.z
            <ir_version: 7, opset_import: [ "" : 18]>
            agraph (float[4, 1, 16] x, float[1, 8, 16] y) => () {
                xshape = Shape (x)
                z = Expand (y, xshape)
            }
        xshaper   r   r      r   Tr   Nonnxparserparse_modelr   r   r   r   r.   r   models     r!   test_model_data_propagation/TestDataPropagation.test_model_data_propagation:   sa    ''
 	&x1B1BDI&sK,=,=zJ  	 	
r$   c           	         [         R                  R                  S5      nU R                  U[	        S[
        R                  S5      [	        S[
        R                  S5      /SS9  g)	zxTest value-propagation through function calls.
Underlying core example is same as previous test_model_data_propagation.
a  
            <ir_version: 7, opset_import: [ "" : 18, "local" : 1 ]>
            agraph (float[4, 1, 16] x, float[1, 8, 16] y) => () {
                xshape = local.GetShape (x)
                z = Expand (y, xshape)
            }
            <domain: "local", opset_import: [ "" : 18 ]>
            GetShape (x) => (shapeval) {
                shapeval = Shape(x)
            }
        r2   r   r   r3   Tr   Nr5   r9   s     r!   test_data_prop_via_function/TestDataPropagation.test_data_prop_via_functionN   sc     ''

 	&x1B1BDI&sK,=,=zJ  	 	
r$   c                4   [         R                  R                  S5      nU R                  U[	        S[
        R                  S5      [	        S[
        R                  S5      [	        S[
        R                  S5      [	        S[
        R                  S5      /S	S
9  g)zTest value-propagation handles multiple calls to same function correctly.
Underlying core example is same as previous test_model_data_propagation.
a  
            <ir_version: 7, opset_import: [ "" : 18, "local" : 1 ]>
            agraph (float[4, 1, 16] x, float[1, 8, 16] y) => () {
                yshape = local.GetShape (y)
                xshape = local.GetShape (x)
                z = Expand (y, xshape)
                w = Expand (y, yshape)
            }
            <domain: "local", opset_import: [ "" : 18 ]>
            GetShape (x) => (shapeval) {
                shapeval = Shape(x)
            }
        yshaper   r2   r   r3   w)r   r4   r   Tr   Nr5   r9   s     r!   test_multiple_calls_to_function3TestDataPropagation.test_multiple_calls_to_functionh   s     ''
 	&x1B1BDI&x1B1BDI&sK,=,=zJ&sK,=,=zJ	  	 		
r$   c                   U R                  S[        R                  S4S[        R                  S4/[        SS/S/5      [        SS/S/5      [        SSS/S	/5      [        S
S	/S/[	        S[        R
                  SS5      S9// 5      nU R                  U[        S[        R                  S5      [        S[        R                  S5      [        S	[        R                  S5      [        S[        R
                  S5      /SS9  g )Nr   r&   r   )r   r   r   r   r2   rA   Addzshaper(   r   r)   r*   r+   r,   r   )r      r4   Tr   r-   r   s     r!   test_shape_arithmetic)TestDataPropagation.test_shape_arithmetic   s     ;$$i038I8I92UV'C58*5'C58*5%(H!5zB%JE%g{/@/@$M		
 
 	&x1B1BDI&x1B1BDI&x1B1BDI&sK,=,=yI	  	 		
r$   c                   U R                  S[        R                  S4S[        R                  S4/[        SS/S/5      [        SS/S/5      [        SSS/S	/5      [        S
S	/S/[	        S[        R
                  SS5      S9// 5      nU R                  U[        S[        R                  S5      [        S[        R                  S5      [        S	[        R                  S5      [        S[        R
                  S5      /SS9  g )Nr   r&   r   r   r   r2   rA   rF   rG   r(   r   r)   r*   r+   r,   )rH      r4   Tr   r-   r   s     r!   $test_shape_arithmetic_with_broadcast8TestDataPropagation.test_shape_arithmetic_with_broadcast   s     ;$$i038I8I42PQ'C58*5'C58*5%(H!5zB%JE%g{/@/@$M		
 
 	&x1B1BDI&x1B1BDI&x1B1BDI&sK,=,=yI	  	 		
r$   c                   U R                  S[        R                  S4S[        R                  S4/[        SS/S/5      [        SS/S/5      [        SSS/S	/5      [        S
S	/S/[	        S[        R
                  SS5      S9// 5      nU R                  U[        S[        R                  S5      [        S[        R                  S5      [        S	[        R                  S5      [        S[        R
                  S5      /SS9  g )Nr    r   r   r   r2   rA   rF   rG   r(   r   r)   r*   r+   r,   )r   Tr   r-   r   s     r!   )test_shape_arithmetic_with_zero_broadcast=TestDataPropagation.test_shape_arithmetic_with_zero_broadcast   s     ;$$b)C1B1BD+IJ'C58*5'C58*5%(H!5zB%JE%g{/@/@$M		
 
 	&x1B1BDI&x1B1BDI&x1B1BDI&sK,=,=rB	  	 		
r$   c                F   [         R                  R                  S5      n[         R                  R	                  USSS5      nUR
                  R                  S   nU R                  UR                  R                  R                  R                  S   R                  S5        g)zZTest that a Concat with an empty tensor as input is handled correctly by data-propagation.z
            <ir_version: 7, opset_import: [ "" : 17]>
            agraph (float[256] y) => (float[N] z)
            <float[0] x = {}>
            {
                z = Concat <axis=0> (x, y)
            }
        Tr      Nr6   r7   r8   shape_inferenceinfer_shapesr    outputassertEqualtypetensor_typer   dim	dim_valuer   r:   inferred_modelrX   s       r!   test_empty_tensor%TestDataPropagation.test_empty_tensor   s    ''	
 --::5$dS%%,,Q/0066::1=GGMr$   c                F   [         R                  R                  S5      n[         R                  R	                  USSS5      nUR
                  R                  S   nU R                  UR                  R                  R                  R                  S   R                  S5        g)z{Test that a Concat with an empty tensor as input is handled correctly by data-propagation.
This time with a negative axis.
z
            <ir_version: 7, opset_import: [ "" : 17]>
            agraph (float[256] y) => (float[N] z)
            <float[0] x = {}>
            {
                z = Concat <axis=-1> (x, y)
            }
        Tr   rT   NrU   r^   s       r!   test_empty_tensor_negative_axis3TestDataPropagation.test_empty_tensor_negative_axis   s     ''	
 --::5$dS%%,,Q/0066::1=GGMr$   rP   N)returnNone)__name__
__module____qualname____firstlineno__r"   r/   r;   r>   rC   rI   rM   rQ   r`   rc   __static_attributes__rP   r$   r!   r
   r
      s6    
$
.
(
4
<
6
6
6N Nr$   r
   __main__)
__future__r   unittestshape_inference_testr   onnx.parserr6   r   onnx.helperr   r   r   r
   rg   mainrP   r$   r!   <module>rs      sF   
 #  9   F FgN2 gNT zMMO r$   