U
    qhg                    @  s|  d Z ddlmZ ddlZddl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 ddlmZ dd	lmZ ddlZddlZddlZdd
lmZmZ ddlmZ ddlmZmZm Z m!Z!m"Z"m#Z#m$Z$ ddl%m&Z& ddl'm(Z( ddl)m*Z* ddl+m,Z,m-Z-m.Z.m/Z/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5m6Z6m7Z7m8Z8m9Z9m:Z: ddl;m<Z<m=Z=m>Z> ddl?m@Z@ ddlAmBZBmCZC ddlDmEZE ddlFmGZGmHZH ddlImJZJmKZK ddlLmMZM ddlNmOZO ddlPmQZQmRZR ddlSmTZTmUZU ddlVmTZW ddlXmYZY eZeddZ[eZeddZ\d d! Z]d"d# Z^d$d% Z_d&d' Z`d(d) Zad*d+ Zbd,d- Zcd.d/d0d0d0d.d1d2d3ZdG d4d5 d5ejeZfegd6krxejhd7d8 dS )9zYou can run a specific test by using the following syntax.

::

    python onnx/test/reference_evaluator_test.py TestReferenceEvaluator.test_function_attribute_nested_graph
    )annotationsN)redirect_stdoutwraps)StringIO)getenv)dedent)Sequence)assert_allcloseassert_almost_equal)AttributeProtoFunctionProto
ModelProtoTensorProtocheckerparsersubbyte)get_roi_align_input_values)check_model)onnx_opset_version)float32_to_bfloat16float32_to_float8e4m3float32_to_float8e5m2make_function
make_graph
make_modelmake_model_gen_version	make_nodemake_operatorsetidmake_opsetidmake_sequence_type_protomake_tensormake_tensor_sequence_value_infomake_tensor_value_infomake_value_info)float8e4m3_to_float32float8e5m2_to_float32
from_array)ReferenceEvaluator)OpRunOpRunExpand)load_op)_get_indices_is_out)Cast_19Celu)Adam)_vcelu1)_col2im_naive_implementation_2dcol2im_naive_implementation)Conv_conv_implementation)r4   )_conv_implementation_im2colORT_MAX_IR_SUPPORTED_VERSION8$ORT_MAX_ONNX_OPSET_SUPPORTED_VERSIONZ18c                   s   t   fdd}|S )Nc                    s>   zdd l }~W n  tk
r.   tdd Y nX  | | d S )Nr   zonnxruntime not installed)onnxruntimeImportErrorunittestSkipTest)argskwargsr:   fn F/tmp/pip-unpacked-wheel-xnis5xre/onnx/test/reference_evaluator_test.pywrapperV   s    z'skip_if_no_onnxruntime.<locals>.wrapperr   rA   rD   rB   r@   rC   skip_if_no_onnxruntimeU   s    	rF   c                   s   t   fdd}|S )Nc                    s>   zdd l }~W n  tk
r.   tdd Y nX  | | d S )Nr   ztorch not installed)torchr;   r<   r=   )r>   r?   rG   r@   rB   rC   rD   d   s    z!skip_if_no_torch.<locals>.wrapperr   rE   rB   r@   rC   skip_if_no_torchc   s    	rH   c                   s   t   fdd}|S )Nc                    s>   zdd l }~W n  tk
r.   tdd Y nX  | | d S )Nr   ztorchvision not installed)torchvisionr;   r<   r=   )r>   r?   rI   r@   rB   rC   rD   r   s    z'skip_if_no_torchvision.<locals>.wrapperr   rE   rB   r@   rC   skip_if_no_torchvisionq   s    	rJ   c                   s   t   fdd}|S )Nc                    s>   zdd l }~W n  tk
r.   tdd Y nX  | | d S )Nr   zml-dtypes not installed)	ml_dtypesr;   r<   r=   )r>   r?   rK   r@   rB   rC   rD      s    z%skip_if_no_ml_dtypes.<locals>.wrapperr   rE   rB   r@   rC   skip_if_no_ml_dtypes   s    	rL   c                 C  s*   t |trt| ||S t|}t| ||S N)
isinstanceintr"   r    r$   )nameZ	elem_typeshapeZs_typerB   rB   rC   make_sequence_value_info   s    
rR   c                 C  sV   dd l }t}| jD ]}|jdkr|j} q,q| jtks>|tkrBd S |j|  dgdS )Nr   ) zai.onnxZCPUExecutionProvider)Z	providers)	r:   r9   opset_importdomainversion
ir_versionr7   ZInferenceSessionSerializeToString)
onnx_modelZortZonnx_domain_opsetopsetrB   rB   rC   run_ort_inference   s    

 r[   c                   s  t |tstdt|dt| jt|krHtd| jd|dtd  t fddt	 D }t
| j| }t	 D ]`}|| || d	 || d	   }t|| | || d
  || d	  ||  d	 }	|	||< qt|}
tj|
| jd}t|}t|jd   }t	|D ]}t||jd   }t|}t	|D ]p}t||}t|}|| |ddd
f  ||  }t|}t|| jrd
||| < n| t| ||| < qTq*|S )zNaive implementation for `im2col`.

    Args:
        data: image (float)
        kernel_shape: kernel shape
        dilations: dilations
        pads: pads
        strides: strides

    Returns:
        result
    zUnexpected type z for kernel_shape.zShape mismatch z and .   c                   s    g | ]}| |   fqS rB   rB   .0iZn_dimspadsrB   rC   
<listcomp>   s     z/im2col_naive_implementation.<locals>.<listcomp>   r   dtypeN)rN   tuple	TypeErrortypelenrQ   
ValueErrornparrayrangelistrO   zerosrf   prodr,   r-   )datakernel_shape	dilationsrb   stridesZnew_padsZlist_output_shapedZkdndZoutput_shapereskernel_sizeZres_sizer`   Zi_resZt_resjZi_kernelZt_kernelZi_imgZt_imgrB   ra   rC   im2col_naive_implementation   s>    
 *


 r{   z
np.ndarrayztuple[int, ...]zSequence[int])imgrs   rt   rb   ru   returnc           
      C  s   d }t | jd D ]l}t | jd D ]X}t| ||df ||||}|d krn| jd d |j }	tj|	| jd}||||df< q$q|jd t|  d }	||	S )Nr   rd   .r]   re   ))rn   rQ   r{   rl   emptyrf   rj   reshape)
r|   rs   rt   rb   ru   rx   ncoutZ	new_shaperB   rB   rC   im2col   s         r   c                   @  sJ  e Zd ZdZedddddZedddZdd Zdd Zdd Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd  Zd!d" Zd#d$ Zd%d& Zd'd( Zd)d* Zd+d, Zd-d. Zd/d0 Zd1d2 Zd3d4 Zd5d6 Zd7d8 Zd9d: Zd;d< Zd=d> Zd?d@ Z dAdB Z!dCdD Z"dEdF Z#dGdH Z$dIdJ Z%dKdL Z&dMdN Z'dOdP Z(dQdR Z)dSdT Z*dUdV Z+dWdX Z,dYdZ Z-d[d\ Z.d]d^ Z/e0d_d` Z1e0dadb Z2dcdd Z3dedf Z4dgdh Z5didj Z6dkdl Z7dmdn Z8dodp Z9e:dqdr Z;dsdt Z<e:dudv Z=e:dwdx Z>dydz Z?d{d| Z@d}d~ ZAdd ZBdd ZCdd ZDdd ZEdd ZFdd ZGdd ZHeIJeKLdddd ZMeIJeKLdddd ZNdd ZOdd ZPe0dd ZQdd ZRe:eSdd ZTdd ZUdd ZVdd ZWdd ZXdd ZYdd ZZdd Z[dd Z\edddZ]e^j^_e`adebjcebjdebjdgddggebjedebjcddgddggebjfdgfdebjcddgddggebjedebjcddgddggebjfdgfdebjcddgddggebjedebjcddgddggebjfdgfdebjcddgddggebjedebjcddgddggebjfdgfdebjcddgddggebjedebjcddgddggebjfdgfdebjcddgddggebjedebjcddgddggebjfdgfdebjcddgddggebjedebjcddgddggebjfdgfdebjcddgddggebjedebjcddgddggebjfdgfdebjcebjdebjdgddggebjedebjcddgddggebjfdgfg	ddgddddZge^j^_dddgddddddZhdd Zidd Zjdd Zkdd Zldd Zmdd Zndd Zodd Zpdd Zqdd Zrdd Zsd d Ztdd Zudd Zvdd Zwdd	 Zxd
d Zydd Zze^j^_deb{d|dd eb{dd|ddeb}dddddddgddddgddddggfdeb{d|dd eb{dd|ddeb~dddddddgddddgddddggfeb{d|dddddgdddgdddggeb}dddddddddddgddddddddgd d!d"d#dd$dd%ggfeb{dddgd&d'gddd&dddddgfeb~d(eb~d)eb}d)dddfeb~d*eb~d)eb}d)dddfgd+d, Ze^j^_eb{d|ddeb{dd|ddeb}dddddddgdddd-gdd.d/d0ggfeb{d|ddeb{dd|ddeb~dddd&dddgd1dd2dgd3dd4d/ggfebeb{d|ddgd ebebcddgddggebcddgddgggeb}d5ddddddgddddggddddgd1d1d6d6gggfeb{d|dddddgdddgdddggeb}ddddddddddd6gd7dd2dŐddd$d.gdd8dd9d2d6d0d:ggfeb{dddgddgddd'ddddd1gfeb~d(eb~d)eb}d)dddfeb~d*eb~d)eb}d)dddfgd;d< Zd=d> Zd?d@ ZdAdB ZdCdD ZdEdF ZdGdH ZdIdJ ZdKdL ZdMdN Ze^j^_dOdPdQdRdSdTdUgdVdW Ze^j^_dXdYdZd[d\d]gd^d_ Zd`da Zdbdc Zddde Ze^j^_dfdggdhdigdjdkgdYfdldmdngdogdpdqdrgdZfdsdtdudvgdtdudvgdwdxdygdZfgdzd{ Ze^j^_d|d}gd~ddddgdddggddgfdddgd~dddgddgddggdddgfddd}gd~dddgddgddggdddgfddd}gddd}ggd~ddddgdddgdddggdddgdddgdddgggdddgdddggfdddgdddddgdddgdddggdddgfdddgdddddgdddgdddggdddgfddgddddddddgddddddggddgfg ddebcg |dg fgdd Zdd Ze^j^_dddgddddgdZfdddgdddggddddgdddggdfddddgdddddgd[fgeIJejdkddd ZeIJejdkddd Ze^j^_ejd&dddddd2dgddddddd2dgfejd&ddddddgdddddddgfejdgdgfejdddddddd2gddddddddgfejdddddddgdddddddgfejdgdgfgeIJeKLdddd Ze^j^_e`aejejfejejfdd Ze^j^_e`aejejfejejfdd ZdÐdĄ ZdŐdƄ ZdǐdȄ Ze^j^_dɐdʐdːd̐d͐dΐdϐdАdѐdg
edӐdԄ Ze^j^_dՐd֐dאdgedِdڄ Zdېd܄ ZdS (  TestReferenceEvaluatora6  
        <
            ir_version: 7,
            opset_import: [ "": 10, "com.microsoft": 1]
        >
        agraph (float[N, M] B01, float[N, M] B11, float[N, M] B21) => (float[N, M] D0)
        {
            C0 = Add(B01, B11)
            C1 = Sub(B11, B21)
            D0 = Mul(C0, C1)
        }
        strr   )m_defr}   c                 C  s   t | }t| |S )zjParses a model from a string representation, including checking
        the model for correctness
        )r   parse_modelr   r   )r   mrB   rB   rC   _load_model   s    

z"TestReferenceEvaluator._load_modelFN            ?c              
     sZ  t dtjd d g}t dtjd d g}t dtjd d g}t dtjd g}tdddgdg}| rtdddgdg}	|d k	r|d	k rr rtd
dgdg d}
qtd
dgdgd}
n* rtd
dgdg d}
ntd
dgdg}
t||	|
gd|||g|g}nrttjgtjdddnd } rBttj gtjdddnd }d|rRdnd|r^dndg}td
|dg}
dd ||fD }t||	|
gd|||g|g|d} fdd}n4tdddgdg}	t||	gd|||g|g}dd }|d krt	|}nt	|t
d|gd}zt| W n4 tk
rP } ztd| |W 5 d }~X Y nX ||fS )NXABYMatMulXAAddZY_clip   ZClip)minmax)r   )r   lrre   mirP   marS   c                 S  s   g | ]}|r|qS rB   rB   )r_   _rB   rB   rC   rc   '  s      z=TestReferenceEvaluator._linear_regression.<locals>.<listcomp>initializerc                   s   t || |  S rM   )rl   clipxab	max_value	min_valuerB   rC   f,  s    z4TestReferenceEvaluator._linear_regression.<locals>.fc                 S  s   || | S rM   rB   r   rB   rB   rC   <lambda>2      z;TestReferenceEvaluator._linear_regression.<locals>.<lambda>opset_importszchecker fails for
)r#   r   FLOATr   r   r'   rl   rm   float32r   r   r   	ExceptionAssertionError)r   rZ   r   r   r   r   r   r   node1node2Znode3graphr   r   inputsr   r   rY   erB   r   rC   _linear_regression  sh            

"z)TestReferenceEvaluator._linear_regressionc              	   C  s4   t dtjd d g}| t t| W 5 Q R X d S )Nr   )r#   r   r   assertRaisesrh   r(   )selfr   rB   rB   rC   #test_reference_evaluator_exceptions=  s    z:TestReferenceEvaluator.test_reference_evaluator_exceptionsc                 C  s   t t j}t| t|}| |jdddg | |jdg | |j	ddd t
jddgd	d
ggt
jd}t
jddgddggt
jd}t
jddgddggt
jd}|d |||dd }|| ||  }t|| d S NB01B11B21D0
   rd   rS   zcom.microsoftr   r]      re               r   r   r   )r   r   m2_defr   r   r(   assertEqualinput_namesoutput_namesopsetsrl   rm   r   runr
   r   r   sessr   yzrx   expectedrB   rB   rC   %test_reference_evaluator_no_attributeB  s    
z<TestReferenceEvaluator.test_reference_evaluator_no_attributec                 C  s   t t j}t| t|}| |jdddg | |jdg | |j	ddd t
jddgd	d
ggt
jd}t
jddgddggt
jd}t
jddgddggt
jd}|jd |||ddd}| |t || ||  }t||d  d S )Nr   r   r   r   r   rd   r   r   r]   r   re   r   r   r   r   r   r   r   r   r   T)Zintermediate)r   r   r   r   r   r(   r   r   r   r   rl   rm   r   r   assertIsInstancedictr
   r   rB   rB   rC   2test_reference_evaluator_no_attribute_intermediateP  s    
zITestReferenceEvaluator.test_reference_evaluator_no_attribute_intermediatec                 C  s   t t j}t| t| }| |jdddg | |j	dg | |j
ddd tjddgd	d
ggtjd}tjddgddggtjd}tjddgddggtjd}|d |||dd }|| ||  }t|| d S r   )r   r   r   r   r   r(   rX   r   r   r   r   rl   rm   r   r   r
   r   rB   rB   rC   +test_reference_evaluator_no_attribute_bytes_  s    
zBTestReferenceEvaluator.test_reference_evaluator_no_attribute_bytesc           	   
   C  s*  t t j}tjddgddggtjd}tjddgdd	ggtjd}tjd
dgddggtjd}| jddX t|dd}t }t	| |
d |||d W 5 Q R X | }d}| || W 5 Q R X | jddb t|dd}t }t	| |
d |||d W 5 Q R X | }tdd}| || W 5 Q R X | jddb t|dd}t }t	| |
d |||d W 5 Q R X | }tdd}| || W 5 Q R X | jddb t|dd}t }t	| |
d |||d W 5 Q R X | }tdd}| || W 5 Q R X d S )Nr   rd   r]   r   re   r   r   r   r   r   r   r   r   )levelverboser   z:Add(B01, B11) -> C0
Sub(B11, B21) -> C1
Mul(C0, C1) -> D0
a  
                 +I B01: float32:(2, 2) in [0.0, 3.0]
                 +I B11: float32:(2, 2) in [4.0, 7.0]
                 +I B21: float32:(2, 2) in [-7.0, -4.0]
                Add(B01, B11) -> C0
                 + C0: float32:(2, 2) in [4.0, 10.0]
                Sub(B11, B21) -> C1
                 + C1: float32:(2, 2) in [8.0, 14.0]
                Mul(C0, C1) -> D0
                 + D0: float32:(2, 2) in [32.0, 140.0]
                
a  
                 +I B01: float32:(2, 2):[0.0, 1.0, 2.0, 3.0]
                 +I B11: float32:(2, 2):[4.0, 5.0, 6.0, 7.0]
                 +I B21: float32:(2, 2):[-4.0, -5.0, -6.0, -7.0]
                Add(B01, B11) -> C0
                 + C0: float32:(2, 2):[4.0, 6.0, 8.0, 10.0]
                Sub(B11, B21) -> C1
                 + C1: float32:(2, 2):[8.0, 10.0, 12.0, 14.0]
                Mul(C0, C1) -> D0
                 + D0: float32:(2, 2):[32.0, 60.0, 96.0, 140.0]
                   a  
                 +I B01: float32:(2, 2):[0.0, 1.0, 2.0, 3.0]
                 +I B11: float32:(2, 2):[4.0, 5.0, 6.0, 7.0]
                 +I B21: float32:(2, 2):[-4.0, -5.0, -6.0, -7.0]
                Add(B01, B11) -> C0
                -- begin Add.run(2 inputs)
                -- done Add.run -> 1 outputs
                 + C0: float32:(2, 2):[4.0, 6.0, 8.0, 10.0]
                Sub(B11, B21) -> C1
                -- begin Sub.run(2 inputs)
                -- done Sub.run -> 1 outputs
                 + C1: float32:(2, 2):[8.0, 10.0, 12.0, 14.0]
                Mul(C0, C1) -> D0
                -- begin Mul.run(2 inputs)
                -- done Mul.run -> 1 outputs
                 + D0: float32:(2, 2):[32.0, 60.0, 96.0, 140.0]
                )r   r   r   rl   rm   r   subTestr(   r   r   r   getvaluer   r   lstrip)	r   r   r   r   r   r   stdoutr   logrB   rB   rC   -test_reference_evaluator_no_attribute_verbosem  s`    



zDTestReferenceEvaluator.test_reference_evaluator_no_attribute_verbosec           	      C  s   t  \}}tjddgddggtjd}tjddgtjd}tjdgtjd}||||}t|}|d |||dd }t|| d S )Nr   rd   r]   r   re   r   r   r   r   )r   r   rl   rm   r   r(   r   r
   )	r   r   r   r   r   r   r   r   gotrB   rB   rC   test_reference_evaluator_lr  s    z2TestReferenceEvaluator.test_reference_evaluator_lrc           
   	   C  s6  | j dd tjdd\}}tjddgddggtjd	}tjddgtjd	}tjd
gtjd	}||||}t|}|jd }| |j	j
d |d |||dd }	t||	 W 5 Q R X | j dd tjdd d\}}tjddgddggtjd	}tjddgtjd	}tjd
gtjd	}||||}t|}|jd }| |j	j
d |d |||dd }	t||	 W 5 Q R X | j dd tjdd d\}}tjddgddggtjd	}tjddgtjd	}tjd
gtjd	}||||}t|}|jd }| |j	j
d |d |||dd }	t||	 W 5 Q R X d S )Nmin+maxoptT)r   r   rd   r]   r   re   r   r~   ZClip_11r   r   )r   r   r   )r   r   )r   r   r   rl   rm   r   r(   	rt_nodes_r   	__class____name__r   r
   
r   r   r   r   r   r   r   r   Z	last_noder   rB   rB   rC    test_reference_evaluator_lr_clip  sB    


z7TestReferenceEvaluator.test_reference_evaluator_lr_clipc           
   	   C  s  | j dd tjddd\}}tjddgdd	ggtjd
}tjddgtjd
}tjdgtjd
}||||}t|}|jd }| |j	j
d | |jd | |jd |d |||dd }	t||	 W 5 Q R X | j dd tjddd d\}}tjddgdd	ggtjd
}tjddgtjd
}tjdgtjd
}||||}t|}|jd }| |j	j
d | |jd | |jd |d |||dd }	t||	 W 5 Q R X | j dd tjddd d\}}tjddgdd	ggtjd
}tjddgtjd
}tjdgtjd
}||||}t|}|jd }| |j	j
d | |jd | |jd |d |||dd }	t||	 W 5 Q R X d S )Nr   r   Tr   )r   rZ   r   rd   r]   r   re   r   r~   ZClip_6r   r   )r   rZ   r   g   r   )r   rZ   r   g   G)r   r   r   rl   rm   r   r(   r   r   r   r   r   r   r   r
   r   rB   rB   rC   "test_reference_evaluator_lr_clip_6  s^    
  

  

z9TestReferenceEvaluator.test_reference_evaluator_lr_clip_6c                 C  sT   t d}t|}tjdddgtjdd}|d d|id }|}t|| d S )Naf  
            <
              ir_version: 8,
              opset_import: [ "" : 14, "local" : 1],
              producer_name: "test",
              producer_version: "1.0",
              model_version: 1,
              doc_string: "Test preprocessing model"
            >
            agraph (uint8[H, W, C] x) => (uint8[H, W, C] x_processed)
            {
                x_processed = local.func(x)
            }

            <
              opset_import: [ "" : 14 ],
              domain: "local",
              doc_string: "function 1"
            >
            f1 (x) => (y) {
                y = Identity(x)
            }

            <
              opset_import: [ "" : 14 ],
              domain: "local",
              doc_string: "function 2"
            >
            f2 (x) => (y) {
                y = Identity(x)
            }

            <
              opset_import: [ "" : 14, "local" : 1 ],
              domain: "local",
              doc_string: "Preprocessing function."
            >
            func (x) => (y) {
                x1 = local.f1(x)
                y = local.f2(x1)
            }
        r   rd   r   re   rd   rd   r   r   )	r   r   r(   rl   rm   uint8r   r   r
   )r   r   r   r   resultr   rB   rB   rC   test_nested_local_functions  s    -z2TestReferenceEvaluator.test_nested_local_functionsc           
      C  s   t dtjd d g}t dtjd g}tddgdgdgdd}t|gd|g|g}t|tddgd	}t| t	d

dtj}|jddd}t|}|d d|id }	t||	 d S )Nr   r   	ReduceSumrd   axeskeepdimsrsrS   r   r   <   r   r   r   axisr   r   r#   r   r   r   r   r   r   r   rl   aranger   astyper   sumr(   r   r
   
r   r   r   r   r   rY   r   r   r   r   rB   rB   rC   test_reduce_sum_11S  s    z)TestReferenceEvaluator.test_reduce_sum_11c           
      C  s   t dtjd d g}t dtjd g}tddgdgdgdd}t|gd|g|g}t|tddgd	}t| t	d

dtj}|| jddd}t|}|d d|id }	t||	 d S )Nr   r   ReduceSumSquarerd   r   r   rS   r   r   r   r   r   r   r   r  rB   rB   rC   test_reduce_sum_square_11`  s    z0TestReferenceEvaluator.test_reduce_sum_square_11c                 C  s   t dtjd d g}t dtjd d g}t dtjd g}tdddgdgdd}t|gd||g|g}t|tdd	gd
}t| t	
ddt	j}t	jdgt	jd}|jddd}	t|}
|
d ||dd }t|	| d S )Nr   r   r   r   rd   r   r   rS      r   r   r   re   r   r   r   r   r#   r   r   INT64r   r   r   r   r   rl   r   r   r   r   rm   int64r   r(   r   r
   r   r   r   r   r   r   rY   r   r   r   r   r   rB   rB   rC   test_reduce_sum_13m  s    z)TestReferenceEvaluator.test_reduce_sum_13c                 C  sD  t  }d}td|t|dg}tdddgdg}t }d|_d|_tj|_|j	| t
|d	ddgdg|g|dg}tdtjd d g}tdtjd g}tdtjd g}	ttd	ddgdg|dd
gd||g|	g}
t|
||gd}t|}tddtj}tjdgtjd}|d ||dd }|jddd}t|| d S )NcustomrS   rd   r   r   r   r   r   ZMyReduceSum)rU   r   exampler   	functionsr   r   r]   r~   re   )r   r   r   r   )r   r   r   r   rP   ref_attr_nameINTri   	attributeappendr   r#   r   r   r	  r   r   r(   rl   r   r   r   r   rm   r
  r   r   r
   )r   rZ   
new_domainr   nodeattZmy_reduce_sumr   r   r   r   rY   r   r   r   r   r   rB   rB   rC   test_reduce_sum_attribute|  sZ    
	  z0TestReferenceEvaluator.test_reduce_sum_attributec                 C  s   t dtjd d g}t dtjd d g}t dtjd g}tdddgdgdd}t|gd||g|g}t|tdd	gd
}t| t	
ddt	j}t	jdgt	jd}|| jddd}	t|}
|
d ||dd }t|	| d S )Nr   r   r   r  rd   r  r   rS      r   r   r   re   r   r  r   r  r  rB   rB   rC   test_reduce_sum_square_18  s    z0TestReferenceEvaluator.test_reduce_sum_square_18c                 C  s   t dtjd d g}t dtjd d g}t dtjd g}tdddgdgdd}t|gd||g|g}t|tdd	gd
}t| t	
ddt	j}t	jg t	jd}|jdd}	t|}
|
d ||dd }t|	| d S )Nr   r   r   r   rd   r  r   rS   r  r   r   r   re   r  r   r  r  rB   rB   rC   test_reduce_sum_13_empty_axes  s    z4TestReferenceEvaluator.test_reduce_sum_13_empty_axesc                 C  s   t dtjd d g}t dtjd d g}t dtjd g}tdddgdgdd}t|gd||g|g}t|tdd	gd
}t| t	
ddt	j}t	jg t	jd}|| jdd}	t|}
|
d ||dd }t|	| d S )Nr   r   r   r  rd   r  r   rS   r  r   r   r   re   r  r   r  r  rB   rB   rC   $test_reduce_sum_square_18_empty_axes  s    z;TestReferenceEvaluator.test_reduce_sum_square_18_empty_axesc           	      C  s   t dtjd d g}t dtjd g}tddgdgddd}t|gd|g|g}t|tddgd	}t| t	d

dtj}t|}|d d|id }t|| d S )Nr   r   r   rd   r   Znoop_with_empty_axesr   rS   r  r   r   r   r   r#   r   r   r   r   r   r   r   rl   r   r   r   r   r(   r   r
   	r   r   r   r   r   rY   r   r   r   rB   rB   rC   "test_reduce_sum_13_empty_axes_noop  s    z9TestReferenceEvaluator.test_reduce_sum_13_empty_axes_noopc           	      C  s   t dtjd d g}t dtjd g}tddgdgddd}t|gd|g|g}t|tddgd	}t| t	d

dtj}t|}|d d|id }t|| | d S )Nr   r   r  rd   r  r   rS   r  r   r   r   r   r  r   rB   rB   rC   )test_reduce_sum_square_18_empty_axes_noop  s         z@TestReferenceEvaluator.test_reduce_sum_square_18_empty_axes_noopc                 C  s   t dtjd d g}t dtjd g}t dtjd g}tdddgdg}t|gd||g|g}t|tddgd}t| t	d	
d
tj}tjdgtjd}||k}	t|}
|
d ||dd }t|	| d S )Nr   r   ZGreatergrS   r  r   r   r]   r]   r]   re   r   r   r   r#   r   r   r   r   r   r   r   rl   r   r   r   r   rm   r(   r   r
   r   r   r   r#  r   r   rY   r   r   r   r   r   rB   rB   rC   test_greater  s    z#TestReferenceEvaluator.test_greaterc                 C  sp   t dddgdg}tddtj}tjdgtjd}||k}t|}|d ||d	d
 }t	|| d S )Nr$  r   r   r#  r   r&  r]   re   r'  r   )
r   rl   r   r   r   r   rm   r(   r   r
   )r   r   r   r   r   r   r   rB   rB   rC   test_node_proto  s    z&TestReferenceEvaluator.test_node_protoc                 C  s   t dtjd d g}t dtjd g}t dtjd g}tdddgdg}t|gd||g|g}t|tddgd}t| t	d	
d
tj}tjdgtjd}||k}	t|}
|
d ||dd }t|	| d S )Nr   r   r#  ZGreaterOrEqualr%  rS   r  r   r   r&  r]   re   r'  r   r(  r)  rB   rB   rC   test_greater_or_equal  s    z,TestReferenceEvaluator.test_greater_or_equalc              	   C  s  t dtjd g}tdg dgttjdgtjdd}t|gdg |g}t dtjd g}tdg dgttjdgtjdd}t|gdg |g}ttjdgtjdd	d
}tddd	gdg}tddgdg||d}t dtjd d g}	t dtjd g}
t||gd|	g|
g|gd}t	|}t
|}| t|d tjdgtjd}|d d|id }ttjdgtjd| tjdgtjd}|d d|id }ttjdgtjd| d S )NCConstantrd   re   )value_floatsgthenr   gelsezeror   r$  r   GIfr#  )then_branchelse_branchr%  r   zReferenceEvaluator(X) -> Zr~   )r#   r   r   r   r'   rl   rm   r   r   r   r(   r   r   r   r
   )r   r-  ZbthenZ
bthen_bodyZbelseZ
belse_bodyr2  ZgreaterZnode_ifr   r#  r   	model_defr   r   r   rB   rB   rC   test_if  sJ    zTestReferenceEvaluator.test_ifc                 C  s  t dtjdg}t dtjdg}tdddddgtj}tdddddgtj}tdg dgt|d	}tdg dgt|d	}t	|gd
g |g}t	|gdg |g}tddgdg||d}	t
 }
d|
_d|
_|
jdg |
jdg |
j|	g t }|
jtd|g t	tdddgdgdgdt dtjg gt dtjdggd}t|dtd|tddgd}|j|
g t|}|d dtdi}tjdddddgtjd}t||d  d S )Nthen_outr   else_outrd   r]   r   r   r.  r   outputsvalue	then_body	else_bodyr4  f_condZf_resr   r<  r5  r6  r  rA   rS   condrx   )rU   r   r<  r   )nodesrP   r   r<  test)Zproducer_namer   Tre   r   )r#   r   r   rl   rm   r   r   r   r'   r   r   rU   rP   inputextendoutputr  r   rT   r   BOOLr   r  r(   r   r
   )r   r9  r:  r   r   Zthen_const_nodeZelse_const_noder>  r?  if_noder   rZ   r   r   r   r   r   rB   rB   rC   test_if_functionI  sb          z'TestReferenceEvaluator.test_if_functionc                 C  s  t  }d}td|t|dg}tdg dg}t }d|_d|_tj|_|j	| tdd	d
gdg}tdddgdg}t
|dd	d
gdg|||g|dg}td	tjd d g}	td
tjd d g}
tdtjd g}ttdd	d
gdg|tdtjdgdgdtddgdggd|	|
g|g}t|||gd}t|}tddtj}tjddgtjd}|d ||dd }t|| d }t|| d S )Nr  rS   rd   r.  r   r=  biasr   r   r   r   r   r   LinearRegressionY1Zformer_Bq=
ףp?)rU   rK  Absr  r  r   r  r~   re   r  r   )r   r   r   r   rP   r  TENSORri   r  r  r   r#   r   r   r   r!   r   r(   rl   r   r   r   r   rm   r   absr
   )r   rZ   r  r   cstr  r   r   linear_regressionr   r   r   r   rY   r   r   r   r   r   rB   rB   rC   test_function_attribute|  s`    

  z.TestReferenceEvaluator.test_function_attributec                 C  sH  t  }d}td|t|dg}tdg dg}t }d|_d|_tj|_|j	| tdg dg}t }d|_d	|_tj|_|j	| t
dtjd g}t
dtjd g}t|gd
g |g}	t|gdg |g}
tdg dgttjdgtjdd}tddgdg}tdddgdg}tddgdg|	|
d}tdddgdg}tdddgdg}t|dddgdg||||||g|dd	g}t
dtjd d g}t
dtjd d g}t
dtjd g}ttdddgdg|tdtjdgd gtd!tjdgd"gd#td$dgdggd%||g|g}t|||gd&}t| t|}| |jd jjd' | |jd jjjd( |jd j}d)}|jD ]6}|jjdrf|jstd*t| d+d,}qf|std-td. d/!tj}tjdd0gtjd}|"d |d |d1d }t#|| d  }t$|| |"d |d2 |d1d }t#|| d" }t$|| d S )3Nr  rS   rd   r.  B1r=  bias1B2bias2r>  r?  r2  r   re   r;  	ReduceMinr   Xminr$  r@  r4  r   rA  r   r   r   r   r   rL  rM  	former_B1rN  	former_B2	  )rU   rV  rX  rO  r  r  Z
OpFunctionr(   FzNested node z2 declares no linked attribute but a subgraph does.TzBNo node 'If' was found, has_linked_attribute could not be checked.r   r  r~   r  r   )%r   r   r   r   rP   r  rP  ri   r  r  r#   r   r   r   r'   rl   rm   r   r   r!   r   r   r(   r   r   r   r   Zimpl_
startswithZhas_linked_attributer   r   r   r   r   rQ  r
   )r   rZ   r  r   cst1r  cst2r9  r:  r>  r?  r2  minir@  rI  r   r   rS  r   r   r   r   rY   r   fctcheckedr  r   r   r   r   rB   rB   rC   $test_function_attribute_nested_graph  s    
   

z;TestReferenceEvaluator.test_function_attribute_nested_graphc                  C  s  t  }d}td|t|dg}tdg dg}t }d|_d|_tj|_|j	| tdg dg}t }d|_d	|_tj|_|j	| t
dtjd g}t
dtjd g}t|gd
g |g}	t|gdg |g}
tdg dgttjdgtjdd}tdddgdg}tddgdg|	|
d}tdg dg}t }d|_d|_tj|_|j	| t
dtjd g}t|gd
g |g}t
dtjd g}t|||gdg |g}tdg dgttjdgtjdd}tddgdg}tdddgdg}tddgdg||d}tddd gd!g}td"d!dgd#g}t|d$dd gd#g||||||g|dd	dg}t
dtjd d g}t
d tjd d g}t
d#tjd g}ttd$dd gd%g|td&tjdgd'gtd(tjdgd)gtd*tjdgd+gd,td-d%gd#ggd.||g|g}t|||gd/}t| t|}td0d1tj}tjdd2gtjd}|d |d |d3d }t|| d) }t|| |d |d4 |d3d }t|| d5 }t|| |d |d6 |d3d }t|| d' }t|| d S )7Nr  rS   rd   r.  rU  r=  rV  rW  rX  r>  r?  c100d   re   r;  r$  rZ  Z
f_cond_100r4  ZB4rA  ZB3bias3r2  r   rY  r   ZLessZf_cond_zeror   r   r   r   r   r   rL  rM  r[  rN  r\  r]  Z	former_B3i)rU   rV  rX  rg  rO  r  r  r   r  r~   r  r   ix  i  )r   r   r   r   rP   r  rP  ri   r  r  r#   r   r   r   r'   rl   rm   r   r   r!   r   r   r(   r   r   r   r   rQ  r
   ) r   rZ   r  r   r_  r  r`  r9  r:  Z
then_body1Z
else_body1re  r@  rI  Zcst3Z
then_body2Z
else_body2r2  ra  r   r   rS  r   r   r   r   rY   r   r   r   r   r   rB   rB   rC   +test_function_attribute_nested_nested_graph!  s    	
	  

zBTestReferenceEvaluator.test_function_attribute_nested_nested_graphc              	   C  s  G dd d}G dd dt }G dd dt }tdtjd d g}tdtjd g}tddgdgd	d
d}t|gd|g|g}t|td
dgd}t	d
dtjd }	| t t| W 5 Q R X tddgdgd	d
d}t|gd|g|g}t|td
dgd}| t t||gd W 5 Q R X tddgdgd	d
d}t|gd|g|g}t|td
dgd}| t t||gd W 5 Q R X tddgdgd	d
d}t|gd|g|g}t|td
dgd}t|||gd}
|
d d|	id }d|	d	  }t|| d S )Nc                   @  s    e Zd ZdZdd Zdd ZdS )z:TestReferenceEvaluator.test_custom_node.<locals>._InvAlphar  c                 S  s   || _ || _d S rM   )	onnx_node
run_params)r   ri  rj  rB   rB   rC   __init__  s    zCTestReferenceEvaluator.test_custom_node.<locals>._InvAlpha.__init__c                 S  s   d|| j   fS Nrd   alphar   r   rB   rB   rC   _run  s    z?TestReferenceEvaluator.test_custom_node.<locals>._InvAlpha._runN)r   
__module____qualname__	op_domainrk  rp  rB   rB   rB   rC   	_InvAlpha  s   rt  c                   @  s   e Zd Zdd ZdS )z:TestReferenceEvaluator.test_custom_node.<locals>.InvAlpha2c                 S  s   d|| j   fS rl  rm  ro  rB   rB   rC   rp    s    z?TestReferenceEvaluator.test_custom_node.<locals>.InvAlpha2._runN)r   rq  rr  rp  rB   rB   rB   rC   	InvAlpha2  s   ru  c                   @  s   e Zd ZdZdddZdS )z9TestReferenceEvaluator.test_custom_node.<locals>.InvAlphar  Nc                 S  s   |p| j }d||  fS rl  rm  r   r   rn  rB   rB   rC   rp    s    
z>TestReferenceEvaluator.test_custom_node.<locals>.InvAlpha._run)Nr   rq  rr  rs  rp  rB   rB   rB   rC   InvAlpha  s   rx  r   r         ?r  rn  rU   r   rd   r   r   r   new_opsr   )r)   r#   r   r   r   r   r   r   rl   r   r   r   r   r   NotImplementedErrorr(   rh   r   r
   )r   rt  ru  rx  r   r   r   r   rY   r   r   r   r   rB   rB   rC   test_custom_node  s8    
z'TestReferenceEvaluator.test_custom_nodec           	   	   C  s   G dd dt }tdtjd d g}tdtjd g}tddgdgddd}t|gd|g|g}t|tdd	gd
}t	d
dtjd	 }t||gd}| t |d d|i W 5 Q R X d S )Nc                   @  s   e Zd ZdZdddZdS )zDTestReferenceEvaluator.test_custom_no_output_tuple.<locals>.InvAlphar  Nc                 S  s   |p| j }d||  S rl  rm  rv  rB   rB   rC   rp    s    
zITestReferenceEvaluator.test_custom_no_output_tuple.<locals>.InvAlpha._run)Nrw  rB   rB   rB   rC   rx    s   rx  r   r   ry  r  rz  r   rd   r   r   r   r{  r)   r#   r   r   r   r   r   r   rl   r   r   r   r   r(   r   rh   r   	r   rx  r   r   r   r   rY   r   refrB   rB   rC   test_custom_no_output_tuple  s    z2TestReferenceEvaluator.test_custom_no_output_tuplec           	   	   C  s   G dd dt }tdtjd d g}tdtjd g}tddgdgddd}t|gd|g|g}t|tdd	gd
}t	d
dtjd	 }t||gd}| t |d d|i W 5 Q R X d S )Nc                   @  s   e Zd ZdZdddZdS )zATestReferenceEvaluator.test_custom_empty_output.<locals>.InvAlphar  Nc                 S  s   t  S rM   )rg   rv  rB   rB   rC   rp    s    zFTestReferenceEvaluator.test_custom_empty_output.<locals>.InvAlpha._run)Nrw  rB   rB   rB   rC   rx    s   rx  r   r   ry  r  rz  r   rd   r   r   r   r{  )r)   r#   r   r   r   r   r   r   rl   r   r   r   r   r(   r   rk   r   r  rB   rB   rC   test_custom_empty_output  s    z/TestReferenceEvaluator.test_custom_empty_outputc           	   	   C  s   G dd dt }tdtjd d g}tdtjd g}tddgdgddd}t|gd|g|g}t|tdd	gd
}t	d
dtjd	 }t||gd}| t |d d|i W 5 Q R X d S )Nc                   @  s   e Zd ZdZdddZdS )z@TestReferenceEvaluator.test_custom_tuple_tuple.<locals>.InvAlphar  Nc                 S  sF   |p| j }ttd||  gg}t|ts0tt|d tsBt|S )Nrd   r   )rn  rg   rN   r   r   r   rn  rx   rB   rB   rC   rp    s
    
zETestReferenceEvaluator.test_custom_tuple_tuple.<locals>.InvAlpha._run)Nrw  rB   rB   rB   rC   rx    s   rx  r   r   ry  r  rz  r   rd   r   r   r   r{  r  r  rB   rB   rC   test_custom_tuple_tuple  s    
z.TestReferenceEvaluator.test_custom_tuple_tuplec           	   	     s   G dd d G  fdddt }tdtjd d g}tdtjd g}tddgdgddd	}t|gd
|g|g}t|tddgd}t	d
dtjd }t||gd}| t |d d|i W 5 Q R X d S )Nc                   @  s   e Zd ZdS )zLTestReferenceEvaluator.test_custom_tuple_unexpected_type.<locals>.CustomTypeNr   rq  rr  rB   rB   rB   rC   
CustomType
  s   r  c                      s   e Zd ZdZd fdd	ZdS )zJTestReferenceEvaluator.test_custom_tuple_unexpected_type.<locals>.InvAlphar  Nc                   s0   t   g}t|t stt|d  s,t|S )Nr   )rg   rN   r   r  r  rB   rC   rp    s    zOTestReferenceEvaluator.test_custom_tuple_unexpected_type.<locals>.InvAlpha._run)Nrw  rB   r  rB   rC   rx    s   rx  r   r   ry  r  rz  r   rd   r   r   r   r{  r  r  rB   r  rC   !test_custom_tuple_unexpected_type	  s    	z8TestReferenceEvaluator.test_custom_tuple_unexpected_typec                 C  s  t dtjg }t dtjg }t dtjg }tdtjd }tdtjd }tdddd	d
gtj	}t
dg dgtdtj|j| tdd}t
dg dgtdtjddgdd}t
dg dgtdtjddgdd}	t
dg dgtdtjddgdd}
t
dddgdgd}t
dddgdgd}t
ddddgdgd}t
d ddgdgd}t
d!dgdgd}t||||	||
|||g	d"|||g||g}t
d#d$d%d&gd'g|d(}t
d)d'gd*gddd+}td
tj}g }tdtj}ttd,t d$tj|jt d%tj|jtd&tjg gt d*tjd g||gd-d.}tjd/d/d0d/d0d1d/d0d1d2d/d0d1d2d3gtj	d4}t|}|||d5}|d |}t||d  d S )6Ncond_incond_out
iter_countseq_inseq_outrd   r]   r   r   r   r.  r   Zconst_tensor_x)rP   Z	data_typeZdimsvalsr;  oneZconst_tensor_onerB   Zslice_startZconst_tensor_zerord   r   r   Zconst_tensor_axesr   endr   r<  Z	UnsqueezeZ	slice_endSliceZ	slice_outZSequenceInsertIdentity	loop_bodyZLoop
trip_countrB  	seq_emptyZseq_res)r   r<  bodyZConcatFromSequencerx   )r   r<  r   Znew_axisZ	loop_test)rP   r   r<  rC  )r   r          @      @      @      @re   )r  rB  r  )r#   r   rH  r	  r"   r   rl   rm   r   r   r   r!   rQ   flattenfloatr   r
  Zbool_r   rR   r(   r   r
   )r   r  r  r  r  r  r   Zx_const_nodeZone_const_nodeZzero_const_nodeZ	axes_nodeadd_nodeZend_unsqueeze_nodeZ
slice_nodeZinsert_nodeZidentity_noder  r  Znode_concatr  r  rB  r7  r   oinfr   r   rB   rB   rC   	test_loop   s             z TestReferenceEvaluator.test_loopc                 C  s   t dtjd g}t dtjd g}tddgdgdd}t|gd|g|g}t|}t| t|}|d dt	j
dt	jdd	 id
 }| |jd | |jt	j | | d | | d d S )Nr   r   Z	Bernoulli        seedr%  r]   r   re   ry  r   gh㈵grZ|
 ?r#   r   r   r   r   r   r   r(   r   rl   rp   r   r   rQ   rf   assertGreaterr   
assertLessr   r   r   r   r   r   rY   r   r   rB   rB   rC   test_onnxt_runtime_bernoulli  s    $z3TestReferenceEvaluator.test_onnxt_runtime_bernoullic                 C  s   t dtjd g}tdg dgdddgd}t|gdg |g}t|}t| t|}|d i d }| 	|j
d	 | 	|jtj | | d | | d
 d S )Nr   ZRandomUniformr  r]   r   r  rQ   r%  r   r  rd   )r#   r   r   r   r   r   r   r(   r   r   rQ   rf   rl   r   r  r   r  r   r   r   r   r   rY   r   r   rB   rB   rC   !test_onnxt_runtime_random_uniform  s    z8TestReferenceEvaluator.test_onnxt_runtime_random_uniformc                 C  s   t dtjd g}t dtjd g}tddgdgdd}t|gd|g|g}t|}t| t|}|d dt	j
dt	jdid	 }| |jd | |jt	j | | d	 | | d
 d S )Nr   r   ZRandomUniformLiker  r  r%  r  re   r   rd   r  r  rB   rB   rC   &test_onnxt_runtime_random_uniform_like  s     z=TestReferenceEvaluator.test_onnxt_runtime_random_uniform_likec                 C  s   t dtjd g}tdg dgdddgd}t|gdg |g}t|}t| t|}|d i d }| 	|j
d	 | 	|jtj d S )
Nr   ZRandomNormalr  r]   r   r  r%  r   r  )r#   r   r   r   r   r   r   r(   r   r   rQ   rf   rl   r   r  rB   rB   rC    test_onnxt_runtime_random_normal  s    z7TestReferenceEvaluator.test_onnxt_runtime_random_normalc                 C  s   t dtjd g}t dtjd g}tddgdgdd}t|gd|g|g}t|}t| t|}|d dt	j
dt	jdid	 }| |jd | |jt	j d S )
Nr   r   ZRandomNormalLiker  r  r%  r  re   r   )r#   r   r   r   r   r   r   r(   r   rl   rp   r   r   rQ   rf   r  rB   rB   rC   %test_onnxt_runtime_random_normal_like  s     z<TestReferenceEvaluator.test_onnxt_runtime_random_normal_likec                 C  s^   t jdd}| |jd tjddgddggtjd}t j|dd}t|dd}t	|| d S )Nry  rm  r   rd   r~   r]   re   )
r/   creater   rn  rl   rm   r   evalr1   r
   )r   instr   r   r   rB   rB   rC   test_eval_celu  s    z%TestReferenceEvaluator.test_eval_celuc                 C  sN   t jddgddggt jd}tj|tjd}tj|tjd}|}t|| d S )Nr   rd   r~   r]   re   to)	rl   rm   r   r.   r  r   FLOAT8E4M3FNr   r
   )r   r   r   Zdyr   rB   rB   rC   test_eval_cast  s
    z%TestReferenceEvaluator.test_eval_castc                 C  sv   t dd}| |jd |jdd}| |jd tjddgddggtjd	}|j|dd}t	|dd}t
|| d S )
NrS   r/   ry  rm  r   rd   r~   r]   re   )r+   r   rs  r  rn  rl   rm   r   r  r1   r
   )r   Zcelur  r   r   r   rB   rB   rC   test_eval_celu_load_op  s    
z-TestReferenceEvaluator.test_eval_celu_load_opc                 C  s   t jdd}| |jd d S )Nry  rm  )r0   r  r   rn  )r   r  rB   rB   rC   test_create_adam  s    z'TestReferenceEvaluator.test_create_adamc              	   C  sP  t dtjd d d d g}t dtjd d d d g}t dtjd d d d g}t dtjd d d d g}tddddgdgddddgddgddgd}t|gd	|||g|g}t|td
dgd}t|}|d krd S t|dd}	| 	|	j
d t t|tgdd}
| 	|
j
d t t|dd}| 	|j
d t d\}}t|D ]}t|D ]}tjdd||ftjd}d|dd||f< tjdtjd}tdtdd d|ddd d d d f< tjdggggtjd}|d |||dd }|	d |||dd }t|| |
d |||dd }t|| |d |||dd }t|| q@q2d S )Nr   r   r   Wr4   rd   r]   )rb   rt   ru   r%  rS      r   F)	optimizedr   )r|  r  T)r   r   re   r   rd   rd   r   r   	   r   r~      r   r  r   )r#   r   r   r   r   r   r   r[   r(   r   r   r4   ConvOptimizedrn   rl   rp   r   minimumr   r   rm   r   r
   )r   r   r   r   r  r  r   rY   sess1sess2Zsess3Zsess4sHsWr`   rz   r   r   Zgot3Zgot4rB   rB   rC   	test_conv  sL    
0

z TestReferenceEvaluator.test_convc                 C  sP  t dtjd d d d g}t dtjd d d d g}t dtjd d d d g}t dtjd g}t dtjd g}t dtjd g}t dtjd g}t dtjd g}t d	tjd g}	td
dddddddd	gdg}
t|
gd||||||||	g|g}t|tddgd}t|}|d krd S t	|}d\}}t
|D ]}t
|D ]}tjdd||ftjd}d|dd||f< | jd||d tjdtjd}d|ddd d d d f< |tjdgtjdtjdgtjd|tjdgtjdtjdgtjdtjdgtjdtdgtjd}|d |d }|d |d }t|| W 5 Q R X | jd||d tjdtjd}tdtdd d|ddd d d d f< |tjdgtjdtjdgtjd|tjdgtjdtjdgtjdtjdgtjdtdgtjd}|d |d }|d |d }t|| W 5 Q R X | jd||d tjdtjd}d|ddd d d d f< |tjdgtjdtjdgtjd|tjd gtjdtjd!gtjdtjd"gtjdtdgtjd}|d |d }|d |d }t|| W 5 Q R X q2q$tjd!d#d$d%d&d'd(gd)d*d+d,d-dd.gd/d0d1d2d'd3d4gd5d6d7d8d+d$d+gd9d:d;d<d=d>d?gd@dAdBdCdDdEdFgd!dGdHdIdJdKdLggtjddM}tjdgtjd}tjdgtjd}tjdgtjdd}tjdNgtjd}tjd!gtjd}tjd"gtjd}tjdOgtjd}	||||||||	d}|d |d }|d |d }t|| d S )PNr   wr   x_scalew_scaley_scalex_zero_pointw_zero_pointy_zero_pointZQLinearConvr%  rS   r  r   r   r   rd   re   r   r   Z1x1)r  r`   rz   )rd   rd   rd   rd   )r   r  r  r  r  r  r  r  Z3x3r  r]   r  r     g칟>n?   gO Y@   g{XZ?               :   r   ;      _      @   8                                 '   |   M   P   f   +            S   )   (         \      *         rd   rd   r   r   gO\?{   )r#   r   UINT8r   r   r   r   r   r[   r(   rn   rl   rp   r   r   rm   r   r   r
   r  r   r   )r   r   r  r   r  r  r  r  r  r  r  r   rY   r  r  r  r  r`   rz   feedsr   r   rB   rB   rC   test_qlinearconv1  s    


0

	
z'TestReferenceEvaluator.test_qlinearconvc                 C  s  t dtjd d d d g}t dtjd d d d g}t dtjd d d d g}t dtjd d d d g}tdddgdg|||d}	tddgdg}
td	ddgd
g|||dd}tddgdg}tddd
gdg}t|	|
|||gd||g||g}t|tddtddgd}t|	gd||g|g}t|tddgd}t|}zt	|}|d krHW d S W n t
k
rd   d }Y nX d\}}t|}t|D ]N}t|D ]<}tjdd||ftjd}d|dd||f< tjd#|tjd}tdt||d df d|ddd d d d f< |d ||d}|d k	rP|d ||dd }t|d  |  zt|d  |d   W nZ tk
r } z:td| d| d| d | d!|d  d"|d  |W 5 d }~X Y nX qqd S )$Nr   rM  Y2r  r4   rb   ru   rt   ShaperQ   ZIm2ColZximZexperimental)rb   ru   rt   rU   ZFlattenZwflatr   r%  rS   r  rd   r   )r   r   re   r   r   r]   r~   r  )r   r  zDiscrepancies: pads=z, dilations=z
, strides=z, kernel_shape=r   
!=
)rd   rd   )r#   r   r   r   r   r   r   r   r(   r[   r;   rl   rq   rn   rp   r   r  r   r   r   r
   Zravelr   )r   rs   rb   ru   rt   r   rM  r  r  r  Z
node_shapeZnode_imZ	node_flatZnode_gemr   rY   Z
graph_convZonnx_model_convr   Z	sess_convr  r  nkerr`   rz   r   Zort_resr   rB   rB   rC   common_test_im2col  s         	  




 
.z)TestReferenceEvaluator.common_test_im2colc                 C  s&   | j dddddgddgddgd d S )Nrd   rd   rd   r]   r  r	  r   rB   rB   rC   test_im2col_1x1  s     
  z&TestReferenceEvaluator.test_im2col_1x1c                 C  s&   | j dddddgddgddgd d S )Nr&  rd   r]   r  r  r  rB   rB   rC   test_im2col_2x2  s     
  z&TestReferenceEvaluator.test_im2col_2x2c                 C  s&   | j dddddgddgddgd d S )Nr  rd   r]   r  r  r  rB   rB   rC   test_im2col_3x3  s     
  z&TestReferenceEvaluator.test_im2col_3x3c                 C  s&   | j dddddgddgddgd d S )Nr  r   rd   r]   r   r  r  r  rB   rB   rC   test_im2col_3x3_pads  s     
  z+TestReferenceEvaluator.test_im2col_3x3_padsc                 C  s&   | j dddddgddgddgd d S )Nr  r   rd   r]   r  r  r  rB   rB   rC   test_im2col_3x3_strides	  s     
  z.TestReferenceEvaluator.test_im2col_3x3_stridesc                 C  s&   | j dddddgddgddgd d S )Nr   r   rd   r]   r  r  r  rB   rB   rC   test_im2col_5x5  s     
  z&TestReferenceEvaluator.test_im2col_5x5c                 C  sx  dd l }tdtjd d d g}tdtjd d d g}tdtjd g}tdtjd g}tddddgdgddddgddgddgd}t|gd	|||g|g}t|td
dgd}t	|}	t
dddddgdddddgdddddgddddd gd!d"d#d$d%gggt
j}t
d&d&gt
j}
t
dd&gt
j}|jjt|
|d'}|	d |||
d(}||| }t||d  d S ))Nr   r   r   Ir   Col2Imrd   r  r%  rS   r  r   r         @      &@      0@      5@r        @      (@      1@      6@r         @      *@g      2@      7@r        "@      ,@      3@      8@r  r        .@g      4@g      9@r   )output_sizery   r   r   r  )rG   r#   r   r   r	  r   r   r   r   r(   rl   rm   r   r   r
  nnFoldrg   r   
from_numpynumpyr
   )r   rG   r   r   ISBSr  r   rY   r   image_shapeblock_shapefoldr   rG  rB   rB   rC   test_col2im  sB    

z"TestReferenceEvaluator.test_col2imc              	   C  sl  dd l }tdtjd d d g}tdtjd d d g}	tdtjd g}
tdtjd g}tddddgdg|||d}t|gd||
|g|	g}t|td	d
gd}t	|}|j
jt|t|t|t|t|d}t|}t|D ]}t|D ]|}tjd||ftjd}d|d||f< tj|tjd}tj|tjd}||| }|d |||d}t||d  qqd S )Nr   r   r   r  r   r  r  r%  rS   r  r   )r&  ry   Zdilationpaddingstriderd   re   r   r'  )rG   r#   r   r   r	  r   r   r   r   r(   r(  r)  rg   r   rl   rq   rn   rp   r   rm   r
  r*  r+  r   r
   )r   sizer.  r/  rb   ru   rt   rG   r   r   r,  r-  r  r   rY   r   r0  r  r`   rz   Zi_shapeZb_shaperG  r   rB   rB   rC   common_test_col2im;  sB    
z)TestReferenceEvaluator.common_test_col2imc                 C  s*   | j dddddddgddgddgd d S )Nr   r   r   r]   r   r   rd   r  r5  r  rB   rB   rC   test_col2im_2x3f  s       
  z&TestReferenceEvaluator.test_col2im_2x3c                 C  s*   | j dddddddgddgddgd d S )N   r6  r7  rd   r  r8  r  rB   rB   rC   test_col2im_2x3_padsl  s       
  z+TestReferenceEvaluator.test_col2im_2x3_padsc           	   	   C  s   t jddgt jd}d|d d< t ddgdt ddgt ddddgt ddgf\}}}}}t||||||}t||||||}t|| d S )	Nr   r:  re   r   r   r   r7  rd   )rl   rp   r   rm   r2   r3   r
   )	r   rr   r.  rs   rt   rb   r3  Zr1Zr2rB   rB   rC   test_col2im_2dr  s2              z%TestReferenceEvaluator.test_col2im_2dc                 C  sb  t dtjddddg}t dtjddddg}t dtjdg}t d	tjddddg}td
dddgd	gdddgddgddddgddgd}t|gd|||g|g}t|tddgd}td	d
tjtjdddgdddgdddgggdddgddd gd!d"d#gggd$d%d&gd'd(d)gd*d+d,gggd-d.d/gd0d1d2gd3d4d5ggggtjd6tjd7d8d9d:gtjd6d;}d<|d d d < |d }|d }|d }tjd|jd6}t|jd< D ]}	tdD ]~}
||	|	d |
|
d f }||
 }t|d=ddgddddgddg}t|| }|||
  }|||	|
d d d d f< qqt|}|d |}t||d< d>d? d S )@Nr   r]   r   r   r  rd   r   r   r   r4   grouprt   rs   rb   ru   r%  rS   r  r   i   )r]   r   r   r   g    ޚg   ]?g    ~ng   `cKg   @A?g   ,?g   $w?g   ʩ?g   @g   ug    ?g   g    ?g    ?g   g   
?g   
g   @V?g    ?g   Jg     g   ¿g   i_g   ?g    c?g   ?g   rg    sg   g    N{g   @ʲ?g   m?g    @?g   X?g    Ug   `Ζre   g   @[¿g   &g   (g   `䠾r  r   r  h㈵>atol)r#   r   r   r   r   r   r   rl   r   r   r   r   rm   r   rf   rn   rQ   r   matmulr  r(   r   r
   )r   r   r  r   r   r  r   rY   r  r   r%  r   r  c2mulref1got1rB   rB   rC   test_conv_im2col_group4  s    

NQ]  z.TestReferenceEvaluator.test_conv_im2col_group4c                 C  sf  t dtjddddg}t dtjddddg}t dtjdg}t dtjd d d d g}td	dddgdgdddgddgddddgddgd
}t|gd|||g|g}t|tddgd}td	d
tjd tjdtjdtjdtjdd}d|d d< t|}	|	d |}
tjdddgdddgdddggdddgdddgdddggggtjd}t||
d  d S )Nr   rd   r   r   r  r]   r   r   r4   r=  r%  rS   r  r   l   rd   r   r   r   r]   r   r   r   re   r]   r  )r   r   r   rd   r  r  r!  r  r#  r  r   r   )r#   r   r   r   r   r   r   rl   r   r   r   r   rp   r(   r   rm   r
   r   r   r  r   r   r  r   rY   r  rE  rF  r   rB   rB   rC   test_conv_strides
  s>    


z(TestReferenceEvaluator.test_conv_stridesc           
   	   C  s  t dtjd d d d g}t dtjd d d d g}tddgdgddgddddgddgd}t|gd|g|g}t|td	d
gd}dtdd d d 	d
tji}tjddddgddddgddddgddddggggtjd}t|}|d |}	t||	d   d S )!Nr   r   MaxPoolr   rd   r]   rs   rb   ru   r%  rS   r  r   1   r~   r  g      H@g     G@g     F@g     E@g     D@g      D@g      C@g      B@g      ;@g      :@r$  r  r  r        $@r  re   r   r#   r   r   r   r   r   r   rl   r   r   r   r   rm   r(   r   r
   
r   r   r   r  r   rY   r  r   rE  rF  rB   rB   rC   test_max_pool_2d_12  s4    
&




z)TestReferenceEvaluator.test_max_pool_2d_1c           
      C  sf  t dtjd d d d g}t dtjd d d d g}tddgdgddgddddgddgd}t|gd|g|g}t|td	d
gd}dtjdddddddgdddddddgddddddd gd!d"d#d$d%d&d'gd(d)d*d+d,d-d.gd/d0d+d1d2d3d4gd5d6d7d8d/d9d:ggggtj	d;i}tjd<d=d>d>gd?d@dAdBgdCdDdBdBgdCdEdCdFggggtj	d;}t
|}|d |}	t||	dG  d S )HNr   r   rN  r   rd   r]   rO  r%  rS   r  r   i  if  i  iB  i  i     i  i  i  r      r  iF  iB  B   i  ix  iT  i  i  o   i  is  iL  i  iJ  i  i  4   i  Q   i'  i;  ih  i  i>  i  i  i  i  i|  i  K   iA  i  i  re   g     X@g     @g     @g     8@g     x@g     @g     P@g     `@g     @g     @g      @r   )r#   r   r   r   r   r   r   rl   rm   r   r(   r   r
   rS  rB   rB   rC   test_max_pool_2d_2S  sP    
 




z)TestReferenceEvaluator.test_max_pool_2d_2c                 C  s  t dtjd d g}t dtjd d g}t dtjd d g}t dtjd d g}tddddgdgddd}t|gd	|||g|g}t|td
dgd}tj	dddddggtj
dt	ddggtj	ddggtj
dd}t|}	|	d |}
tj	dddddggtj
d}t||
d  d S )Nr   r  Ur   ScatterElementsrd   r   r   Z	reductionr%  rS   r  r   r   r  r  r  r  re   g?g @r   r  r]  r   r#   r   r   r	  r   r   r   r   rl   rm   r   r(   r   r
   r   r   ZIndr]  r   r  r   rY   r  rE  rF  r   rB   rB   rC   test_scatter_elements  s*    z,TestReferenceEvaluator.test_scatter_elementsc                 C  s   t dtjd d g}t dtjd d g}t dtjd g}t dtjd d g}tddddgdg}t|gd|||g|g}t|tddgd	}tj	d
dggtj
dt	ddggtj	dgtj
dd}t|}	|	d |}
tj	ddggtj
d}t||
d  d S )Nr   r  r]  r   Z	ScatterNDr%  rS   r  r   r   r  re   r   r  r`  ra  rb  rB   rB   rC   test_scatternd  s&    z%TestReferenceEvaluator.test_scatterndc              	     s&  ddd d fdd	}t dtjd d d g}t dtjd g}t d	tjd g}t d
tjd d d d g}tdddd	gd
gddddgd}t|gd|||g|g}t|tddgd}t	d
tjdtddgtddgd}	t|}
|
d |	}||	d dddddddgd}t||d  d S )Nrd   c                 S  sN  | \}}}}~||d  |d  | | dks2t ||d  |d  | | dksVt ||d  |d  | | d }	||d  |d  | | d }
tt||}t||}|tt|	|
 }tt||| }|tt|
|	 }|dd|dd }|dd|dd }tt||| dd}|||fS )Nr   r]   rd   r   r~   )r   rl   repeatr   Ztiler   )x_shapefield_heightfield_widthr2  r3  Nr-  Hr  Z
out_heightZ	out_widthZi0i1Zj0Zj1r`   rz   krB   rB   rC   get_im2col_indices  s    $$  zCTestReferenceEvaluator.test_col2im_impl.<locals>.get_im2col_indicesr   c                   sd  |\}}}}	||d  |d  |	|d  |d   }
}t j|||
|f| jd} |||||\}}}| || | d|}|ddd}t j|td |||f| | }|d dkr|d  |j	d 7  < n|d dkr|d  d9  < |d dkr|d  |j	d 7  < n|d dkr0|d  d9  < |d d d d |d |d |d |d f }|S )Nr   r]   rd   r   re   r~   )
rl   rp   rf   r   Z	transposeaddatslicecopyrQ   )colsrf  rg  rh  r2  r3  ri  r-  rj  r  ZH_paddedZW_paddedZx_paddedrl  r`   rz   Zcols_reshapedrx   rm  rB   rC   col2im_indices  s4        
0z?TestReferenceEvaluator.test_col2im_impl.<locals>.col2im_indicesr   r,  r-  r   r  r   )rb   r%  rS   r  r   r[  )rd   r   r   r   )r   r,  r-  )rd   rd   r   r   )rg  rh  r2  )Nrd   )r   r   Nrd   )r#   r   r   r	  r   r   r   r   rl   r   r   r   r   rm   r(   r   r
   )r   rt  r   r,  r-  r   r  r   rY   r  rE  rF  r   rB   rs  rC   test_col2im_impl  s:       
        
z'TestReferenceEvaluator.test_col2im_implc                 C  s  t dtjd d d d g}t dtjd d d d g}t dtjd g}t dtjd d d d g}tddddgdgddgddgddgddddgddgd	}t|gd
|||g|g}t|tddgd}td	d
tjtd	d
tjtjddddgtjdd}t|}	|	d |}
tjddddgddddgddddgd d!d"d#gd$d%d&d'ggggtjd}t||
d  |d  d9  < d|d d(< t|}	|	d |}
tjd)d*ddgd+d,ddgddddgddddgddddggggtjd}t||
d  d S )-Nr   r  r   r   ConvTransposerd   r   r   )rt   rs   Zoutput_paddingrb   ru   r%  rS   r  r   r   )rd   r   r   r      )r   rd   r   r   re   r  i  i  i  iA  id  i.  id/  i!  i   i3  i4  i$  il$  i8  i\:  i(  i  i+  i,  i#  )r   r   r   r   r   r   r      rR  rL  rB   rB   rC   test_conv_transpose_2d	  sd    











z-TestReferenceEvaluator.test_conv_transpose_2dc                 C  s  t dtjd d d d g}t dtjd d d d g}t dtjd g}t dtjd d d d g}tddddgdgdddgd}t|gd	|||g|g}t|td
dgd}td	d
tjtd	d
tjtjddddgtjdd}tjddddddgddddddgddddddgddddddgddddd d!gdd"d#dd$d%ggdddd&d'dgdddd(d)d*gdd+d,d-d.d/gd0d1d2d3d4d5gd6d7d8d9d:d;gd<d=d>d?d@dAggggtjd}	t|}
|
d |}t|	|d  d S )BNr   r  r   r   rv  Z
SAME_UPPERr]   )auto_padru   r%  rS   r  r   r  r  r  )rd   r]   r   r   r   re   r  rd   r   r   r   rx  r  r:     rw  #      r   L   r     3   r      r      r  &         T   r     R   $   r  Z   rY  t   A   c   rH  rV  r  i$     H   N   r  [      h   rR  )r   r   r  r   r   r  r   rY   r  r   rE  rF  rB   rB   rC   test_conv_transpose_2d_upperI	  sN    	z3TestReferenceEvaluator.test_conv_transpose_2d_upperz1.21.5z+op_dft and op_stft requires numpy >= 1.21.5c                 C  s  t dtjd d d g}t dtjd g}t dtjd g}t dtjd d d d g}tdddddgdg}t|gd|||g|g}t|tddgd	}t	d

dtjtjdtjdtjdtjdd}|d }t|d }t|d }|d d }	|d jd | | d }
tjd|
|	dgtjd}t|
D ]^}|| }|| | }tj|d||df }|d|	 }tj|j|jfdd|d|f< q*t|}|d |}t||d  d S )Nsignal
frame_stepframe_lengthr   STFTrS   r%     r   r  rd   r  rd   rx  re   r  )r  r  r  r]   rd   r   r   )r#   r   r   r	  r   r   r   r   rl   r   r   r   r   rm   r
  rO   rQ   r   rn   fftstackrealimagr(   r   r
   )r   r  r  r  r   r  r   rY   r  onesided_lengthnstftsr   r`   startstopcomplex_outc_outrE  rF  rB   rB   rC   	test_stft	  s<    
"z TestReferenceEvaluator.test_stftc                 C  s  t dtjd d d g}t dtjd g}t dtjd g}t dtjd g}t dtjd d d d g}tdddddgdg}t|gd||||g|g}t|tdd	gd
}t	d
dtjtjdtjdddtdtj tj	dddtjd d   tjdtjdd}	|	d }t|	d }|	d }t|	d }|d d }
d|jd |jd  d  }tjd||
dgtjd}t|D ]j}|| }|| | }tj|d||df | d|
 }|d|
 }tj|j|jfdd|d|f< qtt|}|d |	}t||d  d S )Nr  r  windowr  r   r  r%  rS   r  r   r  r  rx  re   ry  r]   r   r  rd   )r  r  r  r  r  )r#   r   r   r	  r   r   r   r   rl   r   r   r   r   rm   r
  cospirO   rQ   r   rn   r  r  r  r  r(   r   r
   )r   r  r  r  r  r   r  r   rY   r  r  r  r   r`   r  r  r  r  rE  rF  rB   rB   rC   test_stft_with_window	  sL    
* "z,TestReferenceEvaluator.test_stft_with_windowc                 C  s   t dtjd d d d g}t dtjd d g}t dtjd d d d g}t dtjd g}tddddgdgddddd	|d
	}t|gd|||g|g}t|tddgdS )Nr   roisr   r  ZRoiAlignr   r]   r   Zoutput_half_pixel)Zoutput_heightZoutput_widthsampling_ratiospatial_scaleZcoordinate_transformation_modemoder%  rS   r  r   )r#   r   r   r	  r   r   r   r   )r   r  r   r  r   r,  r  r   rB   rB   rC   get_roi_align_model	  s"    z*TestReferenceEvaluator.get_roi_align_modelc                 C  sp   |  |}t \}}}|||d}t|}|d kr6d S |d |}t|}	|	d |}
t|d |
d dd d S )Nr   r  r  r   r?  r@  )r  r   r[   r   r(   r
   )r   r  rY   r   batch_indicesr  r  r   r   r  r   rB   rB   rC   common_test_roi_align 
  s    
z,TestReferenceEvaluator.common_test_roi_alignc              	   C  sH   | j dd | d W 5 Q R X | j dd | d W 5 Q R X d S )Navgr  r   )r   r  r  rB   rB   rC   test_roi_align
  s    z%TestReferenceEvaluator.test_roi_alignc                 C  s   dd l }ddlm} | |}t|}t \}}}|d |||d}	|dddd}
|
||||g}t||	d dd	 d S )
Nr   )RoIAlignr  r  r   r]   )r  r  r?  r@  )	rG   Ztorchvision.opsr  r  r(   r   r   r*  r
   )r   r  rG   r  rY   r   r   r  r  r   r   r   rB   rB   rC   common_test_roi_align_torch
  s    
z2TestReferenceEvaluator.common_test_roi_align_torchc              	   C  s&   | j dd | d W 5 Q R X d S )Nr  r  )r   r  r  rB   rB   rC   test_roi_align_torch!
  s    z+TestReferenceEvaluator.test_roi_align_torchc                 C  s8  t dtjd g}t dtjd g}t dtjd g}t dtjd g}t dtjd g}tddgddddgdd}t|gd	|g||||g}t|td
dgd}dtd	tj
i}	tjdddgtj
dtjdddgtj
dtjdddgtj
dtjdgtj
dg}
t|}|d |	}tdD ]}t|
| ||  qd S Nr   rM  r  Y3Y4Splitr   )Znum_outputsr%  rS   r  r   r   r   rd   r]   re   r   r   r   r   rx  r  r#   r   r   r   r   r   r   rl   r   r   r   rm   r(   r   rn   r
   r   r   rM  r  r  r  r  r   rY   r  r   rE  rF  r`   rB   rB   rC   
test_split*
  s$    z!TestReferenceEvaluator.test_splitc                 C  sL  t dtjd g}t dtjd g}t dtjd g}t dtjd g}t dtjd g}tdddgddddg}t|gd|g||||g}t|td	d
gd}td	tj
tjddddgtjdd}	tjdddgtj
dtjdddgtj
dtjddgtj
dtjddgtj
dg}
t|}|d |	}tdD ]}t|
| ||  q.d S )Nr   rM  r  r  r  r  splitr%  rS   r  r   r   r   r]   re   )r   r  r   rd   r   r   r   r   rx  r  )r#   r   r   r   r   r   r   rl   r   r   r   rm   r
  r(   r   rn   r
   r  rB   rB   rC   test_split_2B
  s(    z#TestReferenceEvaluator.test_split_2c                 C  s  t dtjd g}t dtjd g}t dtjd g}t dtjd g}t dtjd g}tddgddddgdd}t|gd	|g||||g}t|td
dgd}dtd	tj
i}	tjdddgtj
dtjdddgtj
dtjdddgtj
dtjdgtj
dg}
t|}|d |	}tdD ]}t|
| ||  qdtd	tj
i}	tjdddgtj
dtjdddgtj
dtjdddgtj
dtjg tj
dg}
t|}|d |	}tdD ]}t|
| ||  qd S r  r  r  rB   rB   rC   test_split_num_outputs_4]
  s8    z/TestReferenceEvaluator.test_split_num_outputs_4c           
      C  s   t dtjd d g}t dtjd g}tddgdgdd}t|gd|g|g}t|tddgd	}dt	d

dtji}t|}|d |}tjdddgtjd
d}	| |	 |d   d S )Nr   r   ZArgMinrd   r  r%  rS   r  r   r  r   r   r   re   r~   rd   r#   r   r   r	  r   r   r   r   rl   r   r   r   r   r(   r   rm   r
  r   tolist
r   r   r   r  r   rY   r  rE  rF  r   rB   rB   rC   test_argmin
  s    z"TestReferenceEvaluator.test_argminc           
      C  s   t dtjd d g}t dtjd g}tddgdgdd}t|gd|g|g}t|tddgd	}dt	d

dtji}t|}|d |}tjdddgtjd
d}	| |	 |d   d S )Nr   r   ZArgMaxrd   r  r%  rS   r  r   r  r  r   re   r  r   r  r  rB   rB   rC   test_argmax
  s    z"TestReferenceEvaluator.test_argmaxc                 C  s0  t dtjd d g}t dtjd g}t dtjd g}t dtjd g}t dtjd g}tdddddgdgtdddgdgg}t|d	||||g|g}t|td
dgd}tj	dggtj
dtj	dgtj
dtj	dgtj
dtj	dgtj
dd}	t|}
|
d |	}tj	dgtj
d}| | |d   d S )Nr   startsendsr   r   r  TSqueezer%  rS   r  r   r   re   rd   )r   r  r  r   r#   r   r   r	  r   r   r   r   rl   rm   r
  r(   r   r   r  )r   r   r  r  r   r   rC  r   rY   r  rE  rF  r   rB   rB   rC   test_slice_squeeze
  s&    z)TestReferenceEvaluator.test_slice_squeezec           
      C  s   t dtjd d g}t dtjd g}tddgdgdgdgdgdtddgdgdgd	g}t|d
|g|g}t|tddgd}dtj	dggtj
di}t|}|d |}tj	dgtj
d}	| |	 |d   d S )Nr   r   r  r  r   rd   )r   r  r  r  )r   r%  rS   r   r   re   r  )
r   r   r   rC  r   rY   r  rE  rF  r   rB   rB   rC   test_slice_squeeze_6
  s    z+TestReferenceEvaluator.test_slice_squeeze_6c           	      C  sD  t dtjd d g}t dtjd d g}tddgdg}t|gd|g|g}t|tddgd}t| t|}|j	d }| 
|jjd	 |d dtjd
tjdid }| 
|jd | 
|d d t|tddgd}t| t|}|j	d }| 
|jjd |d dtjd
tjdid }| 
|jd | 
|d d d S )Nr   r   
ReduceMeanr%  rS   r  r   r   ZReduceMean_1r  re   r
  r   r   rd   r  ZReduceMean_18)r#   r   r   r   r   r   r   r   r(   r   r   r   r   r   rl   onesr   rQ   )	r   r   r   r   r   rY   r   clsr   rB   rB   rC   test_onnxrt_reduce_mean
  s(    
 
 z.TestReferenceEvaluator.test_onnxrt_reduce_meanr  c           
   P   C  s  g }t dtjd dgt dtjd gg}t dtjd d gt dtjd gg}| dkr|ttjdgtjdd	d
 t|dd	gdgddd}nt|dgdgddgdd}tddgdgdd
tdddgdgdd
|tddgdgdd
g}t	|d|||}g }dddddddddd d!d"d#d$d%d&d'd(d)d*d+d,d-d.d/d0d1d2d3d4d5d6d7d8d9d:d;d<d=d>d?d@dAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd[d\d]d^d_d`dadbdcdddedfgP}|ttj|tj
ddgdhd
 |ttjdigtjddjd
 t dktjd dgg}t dltjd digt dmtjd digg}tdndkdhgdodpgdq|ddrtdsdpgdtgduddgdvtdwdtgdxgdyd
tdzdxdjgdldmgd{ddd|g}t	|d}|||}t|td~| gd}	|	S )NZnext_inr   nextZnext_outZscan_outr  rd   re   Zaxis_redr   Zcdistdf_17_C0Zcdistdf_17_reduced0Zcdistdf_17_ReduceSumSquarer   )rP   r   )rP   r   r   r  Zcdistd_17_IdentityZSubZcdistdf_17_SubZcdistdf_17_IdentityZOnnxIdentityg   @:?g   g   g    ?g   `?g   9?g   `Ƕٿg   Jg    Og   3?g   h	ܿg    sg   g   @`6?g   dg    Lg   ?g   @_ÿg   @Eg   g?g   lg   пg   q?g    v?g   C?g   3?g   `~?g   @Q?g   ˄g   ڞg   @׿g   `Kg    g    ?g   g    zӿg   @E?g   Xӿg   <?g   `ؿg    u(@g    ǿg    Eg    m?g   g    ?g   @?g   T?g   ?g   Tg   @Bʿg    M	?g    9?g   `Y@g    -?g   `Q?g   E?g   &7׿g    [?g   ]Lg    Xg    i?g   g   `ۜg   g   ?g   ;˿g   g   kZ?g   Z?g    &?g   `Ah?g    lg   @D?g    ?G?g    p?g   hg   @>?g   Wg   `1Dֿ)r}  r   Z
Sc_Scancstr]   Z
To_TopKcstrE  valuesindicesZScanZUU032UUZUU033UUZSc_Scan)rP   r  Znum_scan_inputsZ	TransposeZTr_transposed0ZTr_Transpose)rP   permZSqrtZSq_Y0ZSq_SqrtZTopKZTo_TopK)rP   largestsorteddummyrS   r   )r#   r   r   r  r'   rl   rm   r
  r   r   r   r   r	  r   r   )
rZ   	reduce_opZinitializersr   r<  Znode_reducerC  r   Z
list_valuerY   rB   rB   rC   _cdist_model
  sJ   
   Rz#TestReferenceEvaluator._cdist_modelrY  g]l,@g.s,&2@re   r]   r   r   r   ReduceL1gQ@gK-J3@g%s}N@gZOc10@r  r   r  ReduceL2g&?gq#i?giE@g=@r   ZReduceLogSumgd?gUU?g(?gQOIR?ZReduceLogSumExpg?g?g9z%4~@g	Y;@	ReduceMaxgS?gsn`{?go@g@r  gύ]`?gJ?g%s}N @gZOc10@gNs	@ghi
@gXE!@gHg`y#@
ReduceProdr  rO   )rZ   c                   sr  |\ }t ddt j}i }| | }t|}|d d|i}||d|f<  fdd|jd j	jD }	|	d j
}
t |	d jfd|
i}t||gd	}|d d|i}||d
|f< d}| D ]\}}tt|t|D ]\}}|j|jkr$td d| d|j d| d|j dt ||  }|dkrtd| d d| d| d| d| qqd S )Nrx  )r~   r   rE  r  c                   s   g | ]}|j j r|qS rB   )r   r   r^  )r_   r   r  rB   rC   rc     s   z9TestReferenceEvaluator.test_op_reduce.<locals>.<listcomp>r   Z	op_schemar{  Zref_clconstantzShape mismatch for z, :z != r\   gư>zDiscrepancies (max=z) for r   r  )rl   r   r   r   r   r  r(   r   r   r  Z_schemari   r   itemszipreversedrQ   r   rQ  r   )r   Zreduce_op_expectedrZ   r   r   resultsmodelr   r   ZclZschemaZnew_clZbaselinerl  vr   r   ZdiffrB   r  rC   test_op_reduce  s6    `

&&z%TestReferenceEvaluator.test_op_reduce)r  )r  )r  )rZ   	ref_opsetc                 C  s   t dtjd d d d g}t dtjd d d d g}tddgdgg}t|d|g|g}tjddddtj	}t
|td|gd}t|}	|	d d|id	 }
t
|td|gd}t|}|d d|id	 }| |j|
j t||
 d S )
Nr   r   ZMeanVarianceNormalizationr%  r   rd   rS   r   r   )r#   r   r   r   r   rl   randomrandr   r   r   r   r(   r   r   rQ   r
   )r   rZ   r  r   r   rC  r   r   rY   r  r   Zref_onnx_modelZref_expectedr   rB   rB   rC   test_mvn  s    zTestReferenceEvaluator.test_mvnc           
      C  s   dd }| }t jddgddggt jd}t jdd	ggt jd}t jd
dggt jd}t |||g}t|}|||d}|d |}	t||	d  d S )Nc                  S  s<  g } g }g }g }t  dd}g }tdddgdgddd	}|| d
d | D }tddddgdg||}|| |tdtjg  |tdtjg  |tdtjg  |tdtjg  tdddgdgddd	}| | tdddgdgddd	}| | dd | D }	t| d||}
t	|
|	|d}|S )Nrd   )rS   custom_domainZConcatzx:0zx:1Zr__0r   rS   )r   rU   c                 S  s&   g | ]\}}t ||d krdn|qS rl  r   r_   rU   rV   rB   rB   rC   rc   >  s   zZTestReferenceEvaluator.test_concat_in_a_function.<locals>.create_model.<locals>.<listcomp>r  Zconcat_2I__0I__1I__2Zr__4Zr__3c                 S  s&   g | ]\}}t ||d krdn|qS rl  r  r  rB   rB   rC   rc   Y  s   Znumpyxr  )
r   r   r  r  r   r#   r   DOUBLEr   r   )rC  r   r<  r  r   Z	nodes_fctr  Zopset_imports_fctrb  r   r   rY   rB   rB   rC   create_model3  sd    

    
    
  zFTestReferenceEvaluator.test_concat_in_a_function.<locals>.create_modelr   r   r   r   re   rd   r]   r~   )r  r  r  r   )rl   rm   Zfloat64Zvstackr(   r   r
   )
r   r  rY   x1Zx2Zx3r   r  r  r   rB   rB   rC   test_concat_in_a_function2  s    2z0TestReferenceEvaluator.test_concat_in_a_functionc              	   C  s   t dtjd g}t dtjd g}tttddgdgtjdgd|g|g}t|}t	dddtj
g}|d d|id	 }| |t	dddtj
gtjk  d S )
Nr   r   Castr  r%  Clp?Q}>ʈÿr  r   )r#   r   r   STRINGr   r   r   r(   rl   rm   nanr   
assertTruer   Zstr_allr   r   r   r  r  rr   r   rB   rB   rC   test_cast_float_to_stringo  s"    
"z0TestReferenceEvaluator.test_cast_float_to_stringc              	   C  s   t dtjd g}t dtjd g}tttddgdgtjdtddgdgtjdgd|g|g}t|}t	ddd	tj
g}|d d|id
 }t|t	ddd	tj
g d S )Nr   r   r  r#  r  r%  r  r  r  r   )r#   r   r   r   r   r   r  r(   rl   rm   r   r   r
   r  rB   rB   rC   "test_cast_float_to_string_and_back  s     z9TestReferenceEvaluator.test_cast_float_to_string_and_backc                 C  sP  t dtjd }t dtjd }t dtjd }tdddgdgddg}tt|d||g|g}t|}t	
dd	t	j}t	jdt	jd
}|d ||d}	t	jddgddgddgggt	jd
t	jddgddgddgggt	jd
t	jddgddgddgggt	jd
gg}
| t|
d t|	d  t|
d |	d D ]\}}t|| q6d S Nr   r   r#  SplitToSequencer]   r  r%  r  rd   r   r   re   r'  r  r   r  r  r  r  r  r  r  r!  r"  r%  r  r  rQ  r  r  r  r   r#   r   r   r	  	UNDEFINEDr   r   r   r(   rl   r   r   r   r   rm   r
  r   r   rj   r  r
   r   r   r   r#  rC  r  r  rr   r  r   r   r   r   rB   rB   rC   test_split_to_sequence  s"    """z-TestReferenceEvaluator.test_split_to_sequencec                 C  sV  t dtjd }t dtjd }t dtjd }tdddgdgddg}tt|d||g|g}t|}t	
dd	t	j}t	jdddgt	jd
}|d ||d}	t	jddgddgddgggt	jd
t	jddgddgddgggt	jd
t	jddgddgddgggt	jd
gg}
| t|
d t|	d  t|
d |	d D ]\}}t|| q<d S r  r	  r  rB   rB   rC   test_split_to_sequence_1d  s"    """z0TestReferenceEvaluator.test_split_to_sequence_1dc           
        s   t dtjd }t dtjd }tddgdgdddg}tt|d|g|g}t|}t	d
d	tj |d d i} fd
dt jd D g}| t|d t|d  t|d |d D ]\}}	t||	 qd S )Nr   r#  r  r]   r   r   r%  r  r  c                   s"   g | ]} d d d d |f qS rM   rB   r^   rr   rB   rC   rc     s     zTTestReferenceEvaluator.test_split_to_sequence_nokeepdims_noinput.<locals>.<listcomp>)r#   r   r   r
  r   r   r   r(   rl   r   r   r   r   r   rn   rQ   r   rj   r  r
   )
r   r   r#  rC  r  r  r   r   r   r   rB   r  rC   )test_split_to_sequence_nokeepdims_noinput  s    z@TestReferenceEvaluator.test_split_to_sequence_nokeepdims_noinputc                 C  s  t dtjd g}t dtjd g}t dtjd g}t dtjd g}t dtjd g}tttddgdgtjdtddgd	gtjdtd
g dgtdtjdgdddddgdtd
g dgtdtjdgdddddgdtddgdgtjdtdd	gdgtjdtddgdgtjdtddgdgtjdgd|g||||g}t	|}t
jdddddgt
jd}tdd |D  t
dd |D }	t
dd |D }
|d d|i}t|	|d  t|
|d  t|	|d  t|
|d  d S )Nr   F1F2F3F4r  Zf81r  Zf82r.  ZC1r   r   rd   r]   皙?rU  r=  ZC2r%  re   c                 S  s   g | ]}t |qS rB   r   r_   r   rB   rB   rC   rc     s     z;TestReferenceEvaluator.test_cast_float8.<locals>.<listcomp>c                 S  s   g | ]}t t|qS rB   r%   r   r  rB   rB   rC   rc     s     c                 S  s   g | ]}t t|qS rB   )r&   r   r  rB   rB   rC   rc     s     r   )r#   r   r   r   r   r   r  
FLOAT8E5M2r!   r(   rl   rm   r   printr   r
   )r   r   r  r  r  r  r  r  rr   	expected1	expected2r   rB   rB   rC   test_cast_float8  sn          
z'TestReferenceEvaluator.test_cast_float8c           	   
   C  s   t dtjd g}t dtjd g}tttddgdgtjdtdddgdgdd	tddgdgtjdgd
|g|g}tjddgtj	d}tdd |D }t
|}|d d|i}t|d | G dd dt}t
||gd}|d d|i}t|d | d S )Nr   r   r  Zf8r  CastLikeZf32r   )saturater%  g    cAre   c                 S  s"   g | ]}t t|d ddd dqS )TF)uzr  )r   r  r  rB   rB   rC   rc     s
    z@TestReferenceEvaluator.test_cast_like_float8.<locals>.<listcomp>c                   @  s   e Zd ZdZdS )z>TestReferenceEvaluator.test_cast_like_float8.<locals>.CastLikerS   N)r   rq  rr  rs  rB   rB   rB   rC   r     s   r{  )r#   r   r   r   r   r   FLOAT8E4M3FNUZrl   rm   r   r(   r   r
   r*   )	r   r   r   r  rr   r   r  r   r  rB   rB   rC   test_cast_like_float8  s4    z,TestReferenceEvaluator.test_cast_like_float8c           
   	   C  s   t dtjd g}t dtjd g}t dtjd g}tttddgdgtjdtddgdgtjdgd|g||g}t|}t	j
ddd	d
dgt	jd}t	
dd |D }t	
dd |D }|d d|i}	| | |	d   | | |	d   d S )Nr   r  r  r  r  r%  r   rd   r]   r  rU  re   c                 S  s   g | ]}t |qS rB   r  r  rB   rB   rC   rc   8  s     zBTestReferenceEvaluator.test_cast_float8_output.<locals>.<listcomp>c                 S  s   g | ]}t |qS rB   )r   r  rB   rB   rC   rc   9  s     )r#   r   r   r  r  r   r   r   r(   rl   rm   r   r   r   r  )
r   r   r  r  r  r  rr   r  r  r   rB   rB   rC   test_cast_float8_output'  s(    z.TestReferenceEvaluator.test_cast_float8_outputc           	      C  s  t jddddddddd	d
dddt jt j t jgt jd}tjt jdddddddddddddddt jgt jdtjt jdddddddddddddddt jgt jdtjt jdddddddddddddddt jgt jdtj	t jddddddddd ddddddt jgt jdi}d!d" }|
 D ]p\}}| j|d#T ||}t|}|d d$|id% }t|| | |j|j | |j|j W 5 Q R X qd S )&Ng	?i`  i  iP  i0  i  iig-C6?g{Gz?i       j@re   g      ?g      v@g      z@g      t@g      p@g      pg      Xr  g      ?g      |@g      |g      n@g      ng      ?g      x@g      ?g      @g      g?c                 S  st   t dtjd g}t dtjd g}tddgdg| d}tddgdgtjd}t||gd|g|g}t|}t| |S )Nr   r   r  r  r  r   )r#   r   r   r   r   r   r   )r  r   r   r   r   r   rY   rB   rB   rC   model_cast_cast  s    zCTestReferenceEvaluator.test_float8_4_types.<locals>.model_cast_castr  r   r   )rl   rm   infr   r   r   r  r!  r  FLOAT8E5M2FNUZr  r   r(   r   r
   r   rQ   rf   )	r   r   r   r%  r  expectrY   r  r   rB   rB   rC   test_float8_4_types>  s        W

z*TestReferenceEvaluator.test_float8_4_typesc                 C  s   t dtjd g}t dtjd g}tttddgdgtjdgd|g|g}t|}tj	dddd	d
gtj
d}t	dd |D }|d d|i}| | |d   d S )Nr   r   r  r  r%  r   rd   r]   r$  rU  re   c                 S  s   g | ]}t |qS rB   )r   r  rB   rB   rC   rc     s     zDTestReferenceEvaluator.test_cast_bfloat16_output.<locals>.<listcomp>)r#   r   r   BFLOAT16r   r   r   r(   rl   rm   r   r   r   r  )r   r   r   r  r  rr   r  r   rB   rB   rC   test_cast_bfloat16_output  s     
z0TestReferenceEvaluator.test_cast_bfloat16_outputc                 C  s   t dtjd g}t dtjd g}tttdg dgtdtjdgdgdtdg dgtdtjdgd	gdtd
dddgdgtdddgdgddgd|g|g}t|}t	j
dddddgt	jd}t	j
dddddgt	jd}|d d|i}t||d  d S )Nr   r   r.  scalerd   r  r  r2  r  QuantizeLinearr  DequantizeLinearr   r  r%  r]   r$  rU  re        r#   r   r   r   r   r   r!   r  r(   rl   rm   r   r   r
   r   r   r   r  r  rr   r   r   rB   rB   rC   test_quantize_linear_e4m3  s:    z0TestReferenceEvaluator.test_quantize_linear_e4m3c                 C  s   t dtjd g}t dtjd g}tttddddgdgtdddgdgdd	gd
|g|gtdtjdgdgtdtjdgdgg}t|}t	j
dddddgt	jd}t	j
dddddgt	jd}|d d|i}t||d  d S )Nr   r   r-  r,  r2  r  r.  r   r  r%  rd   r  r  r]   r$  rU  re   r/  r0  r1  r2  rB   rB   rC   %test_quantize_linear_e4m3_initializer  s(    z<TestReferenceEvaluator.test_quantize_linear_e4m3_initializerc                 C  s   t dtjd g}t dtjd g}tttdg dgtdtjdgdgdtdg dgtdtjdgd	gdtd
dddgdgtdddgdgddgd|g|g}t|}t	j
dddddgt	jd}t	j
dddddgt	jd}|d d|i}t||d  d S )Nr   r   r.  r,  rd   r  r  r2  r  r-  r  r.  r   r  r%  r]   r$  rU  re   i  r0  )r#   r   r   r   r   r   r!   r  r(   rl   rm   r   r   r
   r2  rB   rB   rC   test_quantize_linear_e5m2	  s:    z0TestReferenceEvaluator.test_quantize_linear_e5m2c                 C  s   t dtjd g}t dtjd g}tttddddgdggd|g|gtdtjdgdgtdtjdgd	gg}t|}t	j
d
dddddddddddgt	jd}t	j
d	dddddddddddgt	jd}|d d|i}t||d  d S )Nr   r   r-  r,  r2  r%  rd   r    r  g      `r        333333@333333@g      @g         @     re   i  i  i     i  i  r   )r#   r   r   UINT16r   r   r   r!   r(   rl   rm   r   uint16r   r
   r2  rB   rB   rC   test_quantize_linear_uint16)  sb    z2TestReferenceEvaluator.test_quantize_linear_uint16c                 C  s   t dtjd g}t dtjd g}tttddddgdggd|g|gtdtjdgdgtdtjdgd	gg}t|}t	j
d
dddddddddddddddgt	jd}t	j
d	dddddddd d!d d"d d"d d"gt	jd}|d d|i}t||d#  d S )$Nr   r   r-  r,  r2  r%  rd   r  r  r  g     r  r7  r8  r9  r:  r;  g    @g    g    @g    g     @g      r<  r=  re   r~   i     i  r  r6  ii r   )r#   r   r   INT16r   r   r   r!   r(   rl   rm   r   int16r   r
   r2  rB   rB   rC   test_quantize_linear_int16e  sr    z1TestReferenceEvaluator.test_quantize_linear_int16c                 C  s   t dtjd g}t dtjd g}tttddddgdgddgd|g|gtdtjd	gd
gtdtjd	gdgg}t|}t	j
ddddgt	jd}t	j
ddd
dgt	jd}|d d|i}t||d  d S )Nr   r   r.  r,  r2  r   r  r%  rd   r  r6  i0u  iy  r>  i  re   g     g     g      }@)r#   r   r?  r   r   r   r   r!   r(   rl   rm   r@  r   r   r
   r2  rB   rB   rC   test_dequantize_linear_uint16  s0       z4TestReferenceEvaluator.test_dequantize_linear_uint16c                 C  s   t dtjd g}t dtjd g}tttddddgdgddgd|g|gtdtjd	gd
gtdtjd	gdgg}t|}t	j
ddddgt	jd}t	j
ddddgt	jd}|d d|i}t||d  d S )Nr   r   r.  r,  r2  r   r  r%  rd   r  i iiii  re   g     @g     @g       g     @)r#   r   rC  r   r   r   r   r!   r(   rl   rm   rD  r   r   r
   r2  rB   rB   rC   test_dequantize_linear_int16  s0       z3TestReferenceEvaluator.test_dequantize_linear_int16r  r   rd   r   r  r   r   rx  r  g      ?ry  r  r  r  r     r     r  D   r  r  r     r~   r  )r  r  r  )r   r   r   r   c                 C  s
  t dtjd g}t dtjd g}tj|tjd}	tj|tjd}
tt	t
ddddgdg||dgd|g|gtdtj|	j|	tdtj|	j|
g}t|}tj|tjd}|d k	rtj|tjd}|d d|i}t||d	  n&| t |d d|i W 5 Q R X d S )
Nr   r   re   r-  r,  r2  r   
block_sizer%  r   )r#   r   r   INT8rl   rm   r   int8r   r   r   r!   rQ   r(   r   r
   r   rk   r   r   r,  Z
zero_pointr   rM  r   r   r   Z
scale_dataZzp_datar  r  rr   r   rB   rB   rC   test_blocked_quantize_linear  sF    ;	   z3TestReferenceEvaluator.test_blocked_quantize_linearr:  -   r   rW  r  r}  r|  6   )r]   r]   r]   r  r  "   r{  E   c                 C  s
  t dtjd g}t dtjd g}tj|tjd}	tj|tjd}
tt	t
ddddgdg||dgd|g|gtdtj|	j|	tdtj|	j|
g}t|}tj|tjd}|d k	rtj|tjd}|d d|i}t||d	  n&| t |d d|i W 5 Q R X d S )
Nr   r   re   r.  r,  r2  rL  r%  r   )r#   r   rN  r   rl   rm   r   rO  r   r   r   r!   rQ   r(   r   r
   r   rk   rP  rB   rB   rC   test_blocked_dequantize_linear=  sF    E	   z5TestReferenceEvaluator.test_blocked_dequantize_linearc              	   C  s   dd }d}d}d}d}t dtjddd	d	g}t d
tjd }tddgd
g||||dg}tt|d|g|g}	t|	}
tj	
ddddtj}|
d d|i}||||||}| t|t|d  d S )Nc                 S  s   t dt j}t | jD ]p\}}}}	t| |td|tt	
|d d  td|tt	|d d  d ||	f d |||||	f< q| ||| |  |  }
|
S )N)r   r   r   r   r   rd   r]   r   )rl   rp   r   r   ZndindexrQ   r   r   rO   mathfloorr   ceil)r   rn  betarK  r4  Z
square_sumr   r   hr  r   rB   rB   rC   	_expected  s&      z2TestReferenceEvaluator.test_lrn.<locals>._expectedg-C6*?ry  r  r   r   r   2   r#  ZLRN)rn  rZ  rK  r4  r%  r   )r#   r   r   r
  r   r   r   r(   rl   r  r  r   r   r   r   rj   )r   r\  rn  rZ  rK  r4  r   r#  rC  r  r  rr   r   r   rB   rB   rC   test_lrn  s    zTestReferenceEvaluator.test_lrnc                 C  s   t ddt jd t ddt jt jdt jdd}tddgdgddgdgd	d
}tf ||}tf ||}t	|| d S )Nr   )rd   rd   r   rd   r   r   r  re   r  NOTSETr>  rt   rs   rb   ru   rz  
rl   r   r   r   r   rp   r   r5   r6   r
   r   r  r?   r   r   rB   rB   rC   test_conv_im2col_1d  s    z*TestReferenceEvaluator.test_conv_im2col_1dc                 C  s   t ddt jd t ddt jt jdt jdd}tddgdgd	d	gdgd
d}tf ||}tf ||}t	|| d S )Nr  )r]   r   r~   rd   )r~   r   r   r  re   r  r   r   r_  r`  ra  rb  rB   rB   rC   test_conv_im2col_1d_pad0  s    z/TestReferenceEvaluator.test_conv_im2col_1d_pad0c                 C  s   t ddt jd t ddt jt jdt jdd}tdddgd	d	gddddgddgd
d}tf ||}tf ||}t	|| d S )N   )rd   rd   r   rK  rd   r  r  r  re   r  r   r_  r`  ra  rb  rB   rB   rC   test_conv_im2col_2d  s"    
z*TestReferenceEvaluator.test_conv_im2col_2dc                 C  s   t ddt jd dt ddt j t jdt jdd	}tdddgddgd
d
d
d
gddgdd}tf ||}tf ||}t	|| d S )Nr   )r]   r   r   r~   rd   r]   r  )r~   r   rd   r]   r  re   r  r   r_  r`  ra  rb  rB   rB   rC   test_conv_im2col_2d_pad0  s"    
z/TestReferenceEvaluator.test_conv_im2col_2d_pad0c                 C  s   t ddt jd dt ddt j t jdt jdd	}tdddgd
d
gddgd dd}tf ||}tf ||}t	|| d S )Nr  )rd   rd   r   r~   rd   r]   r  r  r  re   r  r   Z
SAME_LOWER)r>  rt   rs   ru   rb   rz  ra  rb  rB   rB   rC   test_conv_im2col_2d_autopad  s    z2TestReferenceEvaluator.test_conv_im2col_2d_autopadc              
   C  s   t ddt jd t ddt jt jdt jdd}tddddgd	d	d	gddddddgdddgd
d}tf ||}tf ||}t	|| d S )Ni  )rd   rd   r   r   r  rd   rw  )rd   rd   r   r   r   r  re   r  r   r_  r`  ra  rb  rB   rB   rC   test_conv_im2col_3d'  s*    
z*TestReferenceEvaluator.test_conv_im2col_3dc                 C  s   t ddt jd t ddt jt jdt jdd}tdddgd	d	gddddgd
d
gdd}tf ||}tf ||}t	|| d S )NrH  rI  rd   rS  rJ  rK  re   r  r   r]   r_  r`  ra  rb  rB   rB   rC   test_conv_im2col_2d_strides<  s    
z2TestReferenceEvaluator.test_conv_im2col_2d_stridesc                 C  s   t ddt jd t ddt jt jdt jdd}tdd	dgd
d
gddddgd	d	gdd}tf ||}tf ||}t	|| d S )NrH  rI  rd   rS  rJ  rK  re   r  r]   r   r_  r`  ra  rb  rB   rB   rC   test_conv_im2col_2d_dilationsN  s    
z4TestReferenceEvaluator.test_conv_im2col_2d_dilations)r   )r  )r  )rY  )r  )r  )r  c           
      C  s   t dtjd }t dtjd }tddtj}t|dgdgddg}t	t
|d|g|g}t|}|d d|i}|d }	| |	tj | |	jd d S )	Nr   r   r   )rd   r   r]   r   r  r%  rB   )r#   r   r   rl   r   r   r   r   r   r   r   r(   r   r   Zndarrayr   rQ   )
r   opr   r   rr   rC  r  r  r   rrB   rB   rC   test_reduce_op_no_axis`  s    z-TestReferenceEvaluator.test_reduce_op_no_axisr  rK  r   )r   )r   )r   c                 C  sJ  t dtjd }t dtjd }t dtjd }t dtjd }tjdgtjd}tddddgdgdd	}tt	|gd
|||g|g}t
|}	tjdgtjdd| }
tjddg| tjd}|	d |
||dd }| |jd|  | |jtj tddg|tj}|	d |
||dd }| |jd|  | |jtj d S )Nr   PVr   r   re   ZPadr  )r   r<  r  r%  rd   r  )r   rp  rq  r   ro  r   r   )r   )r#   r   r   r	  rl   rm   r   r   r   r   r(   r   r
  r   r   rQ   rf   re  r   )r   Zdimr   rp  rq  r   r=  r  r  r  r   pr   rB   rB   rC   test_padw  s"    zTestReferenceEvaluator.test_padc              
   C  s   t dtjd }t dtjd }tddgdgdtddgdgtdtjdgdgd	g}tt|d
|g|g}t|}t	j
dt	jd}|d d|id }| |jt  | |jt	j tt	j
dt	jd| d S )Nr   r   r  rQ   r  ConstantOfShaper=  rd   r;  r%  re   r   )r#   r   r   r   r!   r?  r   r   r(   rl   rm   r   r   r   rQ   rg   rf   r@  r
   r   r   r   rC  r  r  r   r   rB   rB   rC   test_constant_of_shape  s"    	z-TestReferenceEvaluator.test_constant_of_shapec                 C  s   t dtjd }t dtjd }tdg dgtdtjdgdgdtd	dgd
gdtdd
gdgtdtjdgdgdtdddgdgg}tt|d|g|g}t	|}t
jdt
jd}|d d|id }| |jt  | |jt
j tt
jdt
jd| d S )Nr   r   r.  liker   rd   r]   r  r  rQ   r  rt  rR  r=  r;  r  r%  re   r   )r#   r   r   r   r!   r?  r	  r   r   r(   rl   rm   r   r   r   rQ   rg   rf   r@  r
   ru  rB   rB   rC   test_constant_of_shape_castlike  s0    z6TestReferenceEvaluator.test_constant_of_shape_castlikec              j   C  s  dt jdddddddd	d
dddddddddddddddddddddd d!d"d#d$d%d&d'd(d)d*d+d,d-d.d/d0d1d2d3d4d5d6d7d8d9d:d;d<d=d>d?d@dAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd[d\d]d^d_d`dadbdcdddegddfdgdhdidjdkdldmdndodpdqdrdsdtdudvdwdxdydzd{d|d}d~dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddgdddddddddddddddddddddddddddddddddddddddddddddddddddddddd ddddddddd	d
dddddddddddddddddddddd d!d"d#d$d%d&d'd(d)d*d+d,d-gdd.d/d0d1d2d3d4d5d6d7d8d9d:d;d<d=d>d?d@dAdBdCdDdEdFdGdHdIdJdKdLdMdNdOdPdQdRdSdTdUdVdWdXdYdZd[d\d]d^d_d`dadbdcdddedfdgdhdidjdkdldmdndodpdqdrdsdtdudvdwdxdydzd{d|d}d~dddddddddddddddddddgdddddddddddddddddddddddddddddddddddddddddddddddddddÐdĐdŐdƐdǐdȐdɐdʐdːd̐d͐dΐdϐdАdѐdҐdӐdԐdՐd֐dאdؐdِdڐdېdܐdݐdސdߐddddddddddddddddddddddgdgt jdi}t jdddddddddd ddddddddddd	d
ddddddddddddddddddddddddddddd d!d"d#ddd$d%d#dd&d'd(d)d*d+d,d-d.d/d0d1d dd2dd3dd4d5d6d
d7dd8d9dd:dd;d<d7dd=dd>dd?dd@gddAdBdCdDdd-d5ddEdFd!dGdHddIddJddddKd2ddLdMd!dFd>ddNd'dOd.ddFddPdHdQddRdSdJd"dddTddUdVdGdWdd)d&dXdYdZddId[dDdHd\d]d^d_dLddRdRdddddddd`dddFd3ddadadbdddcd6ddJd\dFddded/d!dgdd3dfddd7ddTd
dFd^ddgdGdd"dWdFd3d)dPdhdidjdSd\dkd@dlddLd@dmdndRd=djddd(ddddoddpdqd]drdd!dsdtdddjdRdudHd2d:ddd`d/d7dd	dkdbd0d6dvd3ddwdJd drdddxdyd8dzddJd{d^d_d'dd|d=dOd@ddndwdxddzdngddd}d
d/drdddd1dCddd dddpdrd@ddd/d:dDdwd+dd-ddjd~dddd=dnddydHd-djd)d
ddd
d<d	dvdOd2dd4ddEdddLdOdd"dd$ddddNd)ddDdedddQd@dMdd7djdJddhddndWddzd"dEddddd^ddddnddbddXd2gdd1dgdd-dWddd#dd+dMdVddyd"ddddWd*dddwd=dddddd:d_dd^d"ddddddd3ddtdodkddddPd_ddd4dXd$dLddTdd"d'dddud@ddKdd	d0dd]dd!djdddd-ddd3ddEddddldDdddd.dddddd]dudHdWgdgt jdt jdt jdt jd1t jdg}tdtjd }tdtjd }tdtjd }tdtjd }tddgdddgg}t	t
|d|g|||gtdt d gd}t|}	|	d |}
| t|
d tdddD ]}t|| |
|  qd S (  Nr   g6`gQs[ؿgq2 d!?ggc%^?gQF
?g2o?g?gܝ>G?gG@b?gpA?g/ʿg`Y`?g??g@@?gߥ.gt?>P?g<0?ge~% 9g"?gf<߉xgs nY?gA?gOd* ǼgɀqgZ9?g} ixg`?g«kgY Ȧgd?JKg>Y?g LQ?gu@gc0@"?gZֈɾ?g.ykj?g}Ͽh?g2_ѿgg``m?gVHϿgO?g ?g2 vE?g~ۺg?Ͽg:h T?g(ۿ㣿gp}L?g& ׿gӿg֟?gj@@[?g(SHcÿgR,^Rgͼӟ(ɷ?g'U?gם՟P?g2ÿgk{?g; ȿgW^U0?gn ޒgGgĿg3ſ2ѿgT) 8?g^ҿgSQg=vpp?gS_ſgxпgK/Zeg)4$w?g_!@r?g8z?g0(Ngグg|p=fugoҿgN ?gPL2 ?gɗ֓?gtY?g؁?g"cٿF?ge] d`?gut~g/!?g:'?gS*gH?g=T@Au1?gH{?g_@?ggi- BI?gU?q?gHT`EϿgi4-ٿg" ?gʃg9Ͽg3ֿg3d?gP>g@Ӫg- ,Ggk?g d;mgKT?g8?giz_?gဉDοg0+?g ¬?gO?gy_?g%D?g&j'?g+[gu?gv `˿g	?m?g{A>?gggCv?g4.&?g#?g>ġ@?gTg xg"Ug@F?gZ 7?g$s4`?g?u~ؿg~`ӽ?g?sZgs#|?gZ&!?g|Zzɿg(Ϳ?g^stS$ſg+q9gL>?gUo?gR^Igɿg?V??g蒷?g8-?g_ֿg@}\οgcg)?gĿgOɿg9Z?gWK v]?g>@{?gsVg14 @[n?gEc@ѿgVAg;?r?gj &;g#% ÿg:`i?gB@oƿgf_Xjg=?gHCſgk @]ſgG5?g	. ?gr{?g*8⮤g3Z  ,Zgä?g{`b?gjxvԿg^?g1?g?g 7gzPl߅h?g|?gdy1?g$ 9¿gtՉ?g0_?ge߿gx?X?gI??gG]t?gߛgfϾ"Z?g&@?g̅vʿgҿ@ſg?	?gyW?g?FgEFI@3Կg?gZa}gѠ 
:?g5 g)?g~ ?g69q?g?g6XCg|8ҿg?Ԫοg  [0?g.,?g_gԻ?g՟E:g@Ŀg@T?g|iͿgf k̿gE& ?g'@gu^ eŮgP{?gSh?5og=6% ?gg]~:gv!g MgB	 |?g6`r?g]ſgB7?g%Q?g7ف?gg g))L?g*- Jӿg9;ґ?gCr?g>Yͱ?g(?KĿgu gSV?gޟNZ?gE?g?_Sݿg|k`HѿgRS?gۨe?g#@ſg4׶gդϟ?g!5n?g]y ?g-?g}a ?g7_1ԿgrXƿg(;?gQ㿠Kg[#g9*X_4V?go@¿g<~?	ѿgVH?g`2J¿g}3?gڒ ?g.K?ÓgGd?gӕ`?g#K@+g1 ߷?g!?gב ǿg\!>?gmCl7\?gCn~W?g9hgW?2L?gM?gXp?go, 
ǿgs̠?gFė 1?g?ݿg?gTԄ?g	 %ag"d?gg=!p?g?qgAEȿg
?g9c"Vպ?g{wXk?gƣg-@ɧӿgˬ?0)?gcſgc fgR˷䤿g	 ?gu9e?glû?g!H?g1?j?gDXԿg7Z??g߄ gs?g%`g%(hg f?g)1ƿg0c@?gU
g`?֠9g=?Nn?gle lgᩱg0@}gK>?gw4򟌆?gWkByEgt??gCU?g2 v?g邀?g; ۿgW\md?g^?ս?gmg,@?gT }ĿgN?gGb Dտgu?gg4\m?gVKן?g7$@g?¿gE?gfFS?g~P?gBMӿgTzտg(?g?Ԏ?g2~>͌?g}8_g$9 ?g?gﲿf?gÑ!`zg@\ F?g~@eǿg8EU?gsC Y?gg ĿgP?go	 
geʿg_ ?g\u8hg Cɿgz]ڈgu% ?gyd*?gf* xq?gZ߿ǥ?g1_9?g^?gQÄɰͿg!`Tg?gԓͧgs)@ig?grJ?gB?gIzgq"?gt:꾿gJd`k?gu=>R?g gr?gu]?gc?gTyf?gH*?gf?g1g&o}пgG`?gxKu?gu@?gnO~ҿgJ_!g9 kg +g-nt?g[?g _ÿg2Z?n?g_S`?gA 0cu?gwsRֿgܣgaS˭?gI?g-a`bпg(???g]["=gLֿ#~ѿgZG`?gvϿVڿge?ꖿg|?g2MԿg^;?gX_?g3o-߱?gFV߼?g9ǿgqM&gj,??g:Y+?g?g&@.g(䟑w?gl_{gΧr ?g-f??g g] (?g]v	 Oݿgj`PKӿg$ߏ?gyB@BX?g*gS>ϿgY?geƿg4@?g,P &ݿgz`ҿg	~?gf?	?g?1gq^{пg?g)`$?gHo}?gqyF?g> G?g87@Eɿgɷg	_Ng^E`g)˿gkd`?gjF_g}пgzη?gL@ĿgSؿg>&`e<?gB`h٤?g> gWmN?gb_gtߏvgA gb9?տg7B篿g?g '-\?g>+Ʃ?gzi_`?gbP?gHWo?g!/gw?gEf@
?ge$5g@?gW ?gh<g; ?g g?gbj@ÿg#@`5g%`|?gR`Rre   r  r  r  r     r     r  r     i   r                    w         z            r  y         r         ]      a   r  r     r     J   U         s         k         q   r     r        rX  ^      r     X      r0              r        r  F   r  b   r~  r  x         r           g   r     }         j   p   r   r     r  r        r  r  r  r      n   r     u   r     r  r  Y   rf  r                 W               r  ~   r     rH  r  r  r  r  V         >   rZ  r     m      v   r  r  r  `   C   r        r        rV  9   rU  G      rP  g    v?r   r,  ZzpZDynamicQuantizeLinearr%  rS   rd   r   r   r]   r~   )rl   rm   r   r   r#   r   r   r  r   r   r   r   r   r(   r   r   rj   rn   r
   )r   r  r   r   r   ZScaleZZprC  r  r  r   r`   rB   rB   rC   test_dynamic_quantize_linear  sp    gggg               gggg                       z3TestReferenceEvaluator.test_dynamic_quantize_linearabcdefz.comz.netzabc.comzdef.netcatZdogZsnakescatsZdogsZsnakes)r  r  r  rB   r      ßr   Zaau   ßßyyc                 C  s   t dtjd }t dtjd }t dtjd }tdddgdgd}tt|gd||g|g}	t|	}
|
d t	|t	|d^}}tj
|| | |jjdd	h | |j| d S )
Nr   r   r   ZStringConcatr  r%  )r   r   Or]  )r#   r   r  r   r   r   r(   r   rl   rm   testingassert_array_equalZassertInrf   kindr   rQ   )r   r   r   r   expected_shaper   r   r   r  r  r  r   r   rB   rB   rC   test_string_concat  s    	"z)TestReferenceEvaluator.test_string_concatz1,2,3z4,5,6,123456z1,z4,6rS   z5,6zhello world !z  hello   world !z hello world   ! ZhelloZworld!z	o-n-n--x-z	o-n----nx-or   r   Znx r  c                 C  s   t dtjd }t dtjd }t dtjd }tddgddg||d}	tt|	gd|g||g}
t|
}tj	|t
d}|d d|i^}}}tj|tj	|t
d tj|tj	|tjd d S )Nr   Splits	MaxSplitsZStringSplit)r   r<  	delimitermaxsplitr%  re   )r#   r   r  ZINT32r   r   r   r(   rl   rm   objectr   r  r  r
  )r   r   r  r  Zexpected_splitZexpected_num_splitsr   r  r  r  r  r  r   Z
num_splitsr   rB   rB   rC   test_string_split  s&    O z(TestReferenceEvaluator.test_string_splitc                 C  s   t ddddddddd	gd
gd}t|gdtdtjd d gtdtjdgtdtjdgtdtjd d gtdtjdgtdtjdgtdtjdgtd	tjdggtd
tjd d gg}t|tddgdd}t|}t	
ddddgddddgg}|d8 }|t	j}t	j
dgt	jd}t	j
dgt	jd}t	
ddd gd!d"dgddd#gdd$d%gg}|d8 }|t	j}t	j
d&gt	jd}	t	j
d'gt	jd}
t	j
d(gt	jd}t	j
d)gt	jd}|d t|||||	|
||d*}t	jt	j
d+d,d)gdd-dggt	jd|d  d S ).NZQLinearMatMulr   a_scalea_zero_pointr   b_scaleb_zero_pointr  r  r   r  r%  rd   rS   r}  r  r   rW   r  r  r      r   r  r  r  r  gF%u{?re   r  r     r   r     rB  r  g!u|?ig'?i)r   r  r  r   r  r  r  r  r  ii)r   r   r#   r   r   rN  r   r   r(   rl   rm   r   rO  r   r   r   r  r  )r   r  r   rY   r   r   r  r  r   r  r  r  r  r   rB   rB   rC   test_qlinearconv_int8Z  s|    

 
 * z,TestReferenceEvaluator.test_qlinearconv_int8zwww.google.comzwww.facebook.comzwww.bbc.co.ukzwww\.[\w.-]+\.\bcom\bTZOnnxZ
tensorflowZNumpyZPytorchZCythonZnumbaz^[A-Z][a-z]*$r7  zaccount@gmail.comzaccount@hotmail.comz	not emailzaccount2@yahoo.comz,(\W|^)[\w.\-]{0,25}@(yahoo|gmail)\.com(\W|$)win32z)google-re2 package is not built for win32c                 C  s   t dtjd }t dtjd }tddgdg|d}tt|gd|g|g}t|}	|	d dt	
|i^}
}t	j|
| | |
jjd | |
j| d S )Nr   r   RegexFullMatchr   r<  patternr%  r   )r#   r   r  rH  r   r   r   r(   r   rl   rm   r  r  r   rf   r  rQ   )r   r   r  r   r  r   r   r  r  r  r   r   rB   rB   rC   test_regex_full_match  s    z,TestReferenceEvaluator.test_regex_full_matchc              	   C  s   t dtjd }t dtjd }tddgdgdd}tt|gd|g|g}t|}| t	 |
d dtdgi W 5 Q R X d S )Nr   r   r  zx)r  r%  r   )r#   r   r  rH  r   r   r   r(   r   rk   r   rl   rm   )r   r   r   r  r  r  rB   rB   rC   test_regex_invalid_pattern  s    z1TestReferenceEvaluator.test_regex_invalid_patterng      ?gffffff
@r  ig      -ir  z1.22.0z'The test requires numpy 1.22.0 or laterc           	      C  s   t dtjd g}t dtjd g}tttdg dgtdtjdgdgdtdg dgtd|dgd	gdtd
dddgdgtdddgdgd	dgd|g|g}t|}|d dt	
|i}t||d	  d S )Nr   r   r.  r,  rd   r  r  r2  r   r-  r  r.  r  r%  )r#   r   r   r   r   r   r!   r(   r   rl   Zasarrayr
   )	r   Zqtyperr   r   r   r   r  r  r   rB   rB   rC   test_quantize_linear_int4  s6    z0TestReferenceEvaluator.test_quantize_linear_int4c           
        s   t d|d g}t d|d g}tttddgdg|dgd|g|g}t|}tjdddd	d
dgtjd}|tj	k t fdd|D }|
d d|i}	| | |	d   d S )Nr   r   r  r  r%  r   rd   g333333@g@r   r   re   c                   s   g | ]}t j| d qS )signedr   Zfloat32_to_4bit_unpackedr  r  rB   rC   rc   $  s     z@TestReferenceEvaluator.test_cast_int4_output.<locals>.<listcomp>)r#   r   r   r   r(   rl   rm   r   r   INT4r   r   r  
r   Z	cast_fromZcast_tor   r   r  r  rr   r  r   rB   r  rC   test_cast_int4_output  s&    

z,TestReferenceEvaluator.test_cast_int4_outputc           
        s   t d|d g}t d|d g}tttddgdgtjdgd|g|g}t|}tjt	dtj
d}|tjkrrtjntj t fdd	|D }|d d|i}	| | |	d
   d S )Nr   r   r  r  r%  r   re   c                   s   g | ]}t | qS rB   r  r  Zcast_from_nprB   rC   rc   @  s     z?TestReferenceEvaluator.test_cast_int4_input.<locals>.<listcomp>r   )r#   r   r   r   r   r   r(   rl   rm   rn   r   UINT4r  Zuint4Zint4r   r   r  r  rB   r  rC   test_cast_int4_input)  s&    
z+TestReferenceEvaluator.test_cast_int4_inputc           
      C  s  t dtjdg}t dtjdg}t dtjdg}tdddgdgtdg dgd	gd
dtdddgdgddgtddgd}tdddgdgtdg dgd	gddtdddgdgtddgdgttddgdgddgdg |gttdddgdgddgdg |gddgtddtddgd}tttddgdgddgd|g|gdtddtddg||gd }dtj	d!gtj
d"i}t|}|d |}	| |	d# tj	d$gtj
d" d S )%Nr   ri  rG  thisfctaddinput2r.  r  r   CC0r/  rP   r   A1r   rS   r   r   rb  rE  CCr$  rB  r4  rU   r0  r1  :IFr5  r6  rP   rd   rD  r   rW   r   r  r   re   r   r   r#   r   r   r   r   r   r   r   rl   rm   r   r(   r   r   
r   r   rG  r#  func_def_addfunc_defr7  r  r  r   rB   rB   rC   'test_a_function_calling_a_function_onceE  sz    
!z>TestReferenceEvaluator.test_a_function_calling_a_function_oncec           
      C  s  t dtjdg}t dtjdg}t dtjdg}tdddgdgtdg dgd	gd
dtdddgdgddgtddgd}tdddgdgtdg dgd	gddtdddgdgtddgdgttddgdgddgdg |gttdddgdgddgdg |gddgtddtddgd}tttddgdgddtddgdgddgd|g|gd tddtddg||gd!}dtj	d"gtj
d#i}t|}|d |}	| |	d$ tj	d%gtj
d# d S )&Nr   ri  rG  r  r  r  r.  r  r   r  r   r   r!  r   rS   r   r   rb  rE  r"  r$  rB  r4  r#  r0  r1  r$  r%  rd   ztmprD  r   r&  r   re   r   r'  r(  rB   rB   rC   )test_a_function_calling_a_function_double  s|    
!	z@TestReferenceEvaluator.test_a_function_calling_a_function_doublec           
      C  s  t dtjdg}t dtjdg}t dtjdg}tdddgdgtdg dgd	gd
dtdddgdgddgtddgd}tdddgdgtdg dgd	gddtdddgdgtddgdgttddgdgddgdg |gttdddgdgddgdg |gddgtddtddgd}tttddgdgddtddgdgddgd|g|gd tddtddg||gd!}G d"d# d#t}||}|j	
 D ]}	| |	| qd S )$Nr   ri  rG  r  r  r  r.  r  r   r  r   r   r!  r   rS   r   r   rb  rE  r"  r$  rB  r4  r#  r0  r1  r$  r%  rd   r,  rD  r   r&  c                   @  s   e Zd ZdS )z[TestReferenceEvaluator.test_overload_reference_implementation.<locals>.MyReferenceEvaluatorNr  rB   rB   rB   rC   MyReferenceEvaluator&  s   r/  )r#   r   r   r   r   r   r   r   r(   Z
functions_r  r   )
r   r   rG  r#  r)  r*  r7  r/  r  r  rB   rB   rC   &test_overload_reference_implementation  s|    
!	z=TestReferenceEvaluator.test_overload_reference_implementation)r  zG?)r  r1  )r  q=
ףp?)r!  r2  )r  333333?)r'  r3  )r  r   )r   r   )FLOAT16gMb`?)r*  g{Gz?c                 C  s"  t t|}tttddgdg|dtddgdg|dtdddgdgtddgd	gtjdgd
tdtjd d d gtdtjd d d ggtd	tjd d d ggtddgdd}t|dd}t	
dd dt	j}t	
dd dt	j}t||d}|| }	|d |d }
t|	|
|d d S )Nr  r   Xcr  r   Ycr   Zcr#  rw   rS   r  r  r  r   r   r   r]   r   r   r'  r@  )getattrr   r   r   r   r   r#   r   r(   rl   r   r   r   r   r   r   r
   )r   styperA  ityper  r  r   r   r  r   r   rB   rB   rC   test_add_custom_dtype-  s0    

z,TestReferenceEvaluator.test_add_custom_dtype)r  )r   )r4  )r*  c           
      C  s  t t|}tttddgdg|dtddgdg|dtdddgdgtddgd	gtjdgd
tdtjd d d gtdtjd d d ggtd	tjd d d ggtddgdd}t	|}t
dd dt
j}t
dd d dt
j}t||d}||k}|d |d }	t||	 d S )Nr  r   r5  r  r   r6  r$  r7  r#  rw   rS   r  r  r  r8  r'  r   )r9  r   r   r   r   rH  r#   r   r   r(   rl   r   r   r   r   r   r   r   )
r   r:  r;  r  r  r   r   r  r   r   rB   rB   rC   test_cmp_custom_dtypeZ  s0    

 z,TestReferenceEvaluator.test_cmp_custom_dtypec                 C  s  t ttddddgdgdddgd	tdtjd tdtjd tdtjd gtdtjd gtd
dgd}tj	dtj
dd}tjdggggtjd}tjdggggtj
d}tjddddddddddddddddgtj
dd}t|}|d |||d}t||d  d S )Nr^  rr   r  updatesr#  r   rn  r_  rP   rS   r  r   r  re   )r]   r]   r]   r]   r   rd   )rr   r  r>  )r   r   r   r#   r   r   r	  r   rl   rp   r   r   rm   r
  r(   r   r
   )r   r  rr   r  r>  r   r  r   rB   rB   rC   test_scatter_elements_4d  s>    	
" z/TestReferenceEvaluator.test_scatter_elements_4d)FNr   r   )r  )r  )r   rq  rr  r   staticmethodr   r   r   r   r   r   r   r   r   r   r   r  r  r  r  r  r  r  r!  r"  r*  r+  r,  r8  rJ  rT  rd  rh  r~  r  r  r  r  r  r  r  r  r  r  r  r  r  r  rF   r  r  r	  r  r  r  r  r  r  rH   r1  r5  r9  r;  r<  rG  rM  rT  r\  rc  rd  ru  ry  r  r<   ZskipIfversion_utilsZnumpy_older_thanr  r  r  r  r  r  rJ   r  r  r  r  r  r  r  r  r  r  parameterizedexpand	itertoolsproductrl   rm   r   r   r
  r  r  r  r  r  r  r  r  r  r"  r#  r)  r+  r3  r4  r5  rA  rE  rF  rG  r   r   rp   r  rQ  ZdstackrV  r^  rc  rd  rf  rg  rh  ri  rj  rk  rn  rs  rv  rx  r  r  r  r  sysplatformr  r  r   r  r  r  r   r4  r  r  r+  r.  r0  rL   r<  r=  r?  rB   rB   rB   rC   r      s  8T
%/43	,37n}3 


*
~E
'+

 (!2ME6
1
5
' ; 
 Z_#=4%   <D
.	
6	,
 "

	


	


8*
2	
4	:
J	&
"


	


B)#"        -$(0

 
 
,,
6GF
 
 	6




NOKr   __main__r]   )	verbosity)i__doc__
__future__r   rD  rW  rF  r<   
contextlibr   	functoolsr   ior   osr   textwrapr   typingr	   r+  rl   rB  rA  Znumpy.testingr
   r   Zonnx._custom_element_typesZ_custom_element_typesr  Zonnxr   r   r   r   r   r   r   Z$onnx.backend.test.case.node.roialignr   Zonnx.checkerr   Z	onnx.defsr   Zonnx.helperr   r   r   r   r   r   r   r   r   r   r    r!   r"   r#   r$   Zonnx.numpy_helperr%   r&   r'   Zonnx.referencer(   Zonnx.reference.op_runr)   r*   Zonnx.reference.opsr+   Z%onnx.reference.ops._op_common_indicesr,   r-   Zonnx.reference.ops._op_listr.   r/   Z3onnx.reference.ops.aionnx_preview_training._op_listr0   Zonnx.reference.ops.op_celur1   Zonnx.reference.ops.op_col2imr2   r3   Zonnx.reference.ops.op_convr4   r5   Zonnx.reference.ops_optimizedr  Z.onnx.reference.ops_optimized.op_conv_optimizedr6   rO   r7   r9   rF   rH   rJ   rL   rR   r[   r{   r   ZTestCaser   r   mainrB   rB   rB   rC   <module>   s   $	D0                                             i
