U
    rh	                     @  sd   d dl mZ d dlmZ d dlZd dlZd dlmZ d dl	m
Z
 dddd	d
ZG dd deZdS )    )annotations)AnyN)Base)expectz
Any | Noner   )optionalreturnc                 C  s   | d k	st | S )N)AssertionError)r    r	   R/tmp/pip-unpacked-wheel-xnis5xre/onnx/backend/test/case/node/optionalgetelement.py-optional_get_element_reference_implementation   s    r   c                   @  s0   e Zd ZeddddZeddddZdS )OptionalHasElementNone)r   c                  C  s   t ddddgt j} tjjtjjdgd}tj	|}tjj
ddgdgd	}t| }t|| g|g|gd
d t|| g|g|gdd d S )N            Z	elem_typeshapeOptionalGetElementoptional_inputoutputinputsoutputsZ)test_optional_get_element_optional_tensorr   r   Zinput_type_protosnameZ test_optional_get_element_tensor)nparrayastypeZfloat32onnxhelpermake_tensor_type_protoTensorProtoFLOATmake_optional_type_proto	make_noder   r   )r   tensor_type_protooptional_type_protonoder   r	   r	   r
   export_get_element_tensor   s6      z,OptionalHasElement.export_get_element_tensorc                  C  s   t ddddgt jg} tjjtjjdgd}tj	|}tj
|}tjjddgdgd	}t| }t|| g|g|gd
d t|| g|g|gdd d S )Nr   r   r   r   r   r   r   r   r   Z+test_optional_get_element_optional_sequencer   Z"test_optional_get_element_sequence)r   r   r   Zint32r   r    r!   r"   ZINT32Zmake_sequence_type_protor$   r%   r   r   )r   r&   Zseq_type_protor'   r(   r   r	   r	   r
   export_get_element_sequence3   s8      z.OptionalHasElement.export_get_element_sequenceN)__name__
__module____qualname__staticmethodr)   r*   r	   r	   r	   r
   r      s   r   )
__future__r   typingr   Znumpyr   r   Zonnx.backend.test.case.baser   Zonnx.backend.test.case.noder   r   r   r	   r	   r	   r
   <module>   s   