
    /j                    j    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lm	Z	 ddZ
 G d	 d
e          ZdS )    )annotationsN)helper)Base)expect      ?Fc                `   |dk    s|du r(|du r"| t          j        | j        t                    fS | S t           j                            |           t           j                            dd| j                  |k    }dd|z
  z  }|r"|| z  |z  |                    t                    fS || z  |z  S )Nr   FT)dtypeg      ?   )nponesshapeboolrandomseeduniformastype)Xdrop_probabilityr   training_modereturn_maskmaskscales          h/home/longshao/multi-rider-rag/.venv/lib/python3.11/site-packages/onnx/backend/test/case/node/dropout.pydropoutr      s    1 6 6$bgagT22222INN49QQW--1AAD%%&E 3ax%T!2!222!8e    c                  .   e Zd Zedd            Zedd            Zedd            Zedd            Zedd            Zedd            Z	edd	            Z
edd
            Zedd            Zedd            Zedd            Zedd            ZdS )DropoutreturnNonec                 <   t          j        d          } t          j                            ddgdg|           }t           j                            ddd                              t           j                  }t          |          }t          ||g|gd	
           d S )Nr   r   xyinputsoutputsr            test_dropout_defaultr$   r%   namer   int64onnxr   	make_noder   randnr   float32r   r   )r   noder!   r"   s       r   export_defaultzDropout.export_default   s    x{{{$$YuseRV$WWIOOAq!$$++BJ77AJJtQC!3IJJJJJJr   c                 j   t          j        d          } t          j                            dddgdg|           }t          j        d          }t           j                            dd	d
                              t           j                  }t          ||          }t          |||g|gd           d S )Nr   r   r!   rr"   r#   皙?r&   r'   r(   test_dropout_default_ratior*   r   r-   r.   r   r/   r1   r   r0   r   r   r   )r   r2   r5   r!   r"   s        r   export_default_ratiozDropout.export_default_ratio'   s    x{{{$$sCj3%d % 
 
 JsOOIOOAq!$$++BJ77AqMMtQFQC6RSSSSSSr   c                 J   t          j        d          } t          j                            ddgddg|           }t           j                            ddd	                              t           j                  }t          |d
          \  }}t          ||g||gd           d S )Nr   r   r!   r"   zr#   r&   r'   r(   Tr   test_dropout_default_maskr*   r,   )r   r2   r!   r"   r;   s        r   export_default_maskzDropout.export_default_mask3   s    x{{{$$sec3Zd % 
 
 IOOAq!$$++BJ77qd+++1tQC!Q6QRRRRRRr   c                 x   t          j        d          } t          j                            dddgddg|           }t          j        d          }t           j                            d	d
d                              t           j                  }t          ||d          \  }}t          |||g||gd           d S )Nr   r   r!   r5   r"   r;   r#   r6   r&   r'   r(   Tr<   test_dropout_default_mask_ratior*   r8   )r   r2   r5   r!   r"   r;   s         r   export_default_mask_ratioz!Dropout.export_default_mask_ratio>   s    x{{{$$sCj3*4 % 
 
 JsOOIOOAq!$$++BJ77q!...1!Q!Q6W	
 	
 	
 	
 	
 	
r   c                    t          j        d          } t          j                            dg ddg|           }t           j                            ddd                              t           j                  }t          j        d	          }t          j	        d
          }t          |||          }t          ||||g|gd           d S )Nr   r   r!   r5   tr"   r#   r&   r'   r(   r   Tr   test_training_dropout_defaultr*   r   r-   r.   r   r/   r   r0   r   r1   bool_r   r   r   r2   r!   r5   rD   r"   s         r   export_training_defaultzDropout.export_training_defaultN   s    x{{{$$ooou4 % 
 
 IOOAq!$$++BJ77JsOOHTNNAq***!QQC6U	
 	
 	
 	
 	
 	
r   c                    t          j        d          } t          j                            dg dddg|           }t           j                            ddd	                              t           j                  }t          j        d
          }t          j	        d          }t          |||d          \  }}t          ||||g||gd           d S )Nr   r   rC   r"   r;   r#   r&   r'   r(   r   Tr   r   "test_training_dropout_default_maskr*   rG   r   r2   r!   r5   rD   r"   r;   s          r   "export_training_default_ratio_maskz*Dropout.export_training_default_ratio_mask]   s    x{{{$$oooSz % 
 
 IOOAq!$$++BJ77JsOOHTNNq!1$???1q!9F5		
 	
 	
 	
 	
 	
r   c                    t          j        d          } t          j                            dg ddg|           }t           j                            ddd                              t           j                  }t          j        d	          }t          j	        d
          }t          |||          }t          ||||g|gd           d S )Nr   r   rC   r"   r#   r&   r'   r(         ?TrE   test_training_dropoutr*   rG   rI   s         r   export_trainingzDropout.export_trainingo   s    x{{{$$ooou4 % 
 
 IOOAq!$$++BJ77JtHTNNAq***tQ1Is9PQQQQQQr   c                    t          j        d          } t          j                            dg dddg|           }t           j                            ddd	                              t           j                  }t          j        d
          }t          j	        d          }t          |||d          \  }}t          ||||g||gd           d S )Nr   r   rC   r"   r;   r#   r&   r'   r(   rQ   TrL   test_training_dropout_maskr*   rG   rN   s          r   export_training_ratio_maskz"Dropout.export_training_ratio_mask|   s    x{{{$$oooSz % 
 
 IOOAq!$$++BJ77JtHTNNq!1$???1!QQF9U	
 	
 	
 	
 	
 	
r   c                    t          j        d          } t          j                            dg ddg|           }t           j                            ddd                              t           j                  }t          j        d	          }t          j	        d
          }t          |||          }t          ||||g|gd           d S )Nr   r   rC   r"   r#   r&   r'   r(           TrE    test_training_dropout_zero_ratior*   rG   rI   s         r   "export_training_default_zero_ratioz*Dropout.export_training_default_zero_ratio   s    x{{{$$ooou4 % 
 
 IOOAq!$$++BJ77JsOOHTNNAq***!QQC6X	
 	
 	
 	
 	
 	
r   c                    t          j        d          } t          j                            dg dddg|           }t           j                            ddd	                              t           j                  }t          j        d
          }t          j	        d          }t          |||d          \  }}t          ||||g||gd           d S )Nr   r   rC   r"   r;   r#   r&   r'   r(   rX   TrL   %test_training_dropout_zero_ratio_maskr*   rG   rN   s          r   'export_training_default_zero_ratio_maskz/Dropout.export_training_default_zero_ratio_mask   s    x{{{$$oooSz % 
 
 IOOAq!$$++BJ77JsOOHTNNq!1$???1q!9F8		
 	
 	
 	
 	
 	
r   c            
        t           j                            ddgdg          } t          j        g d                              t          j                  }|}t          | |g|gdt          j        dd          g	           d S )
Nr   r!   r"   )r$   r%   )r   r
   test_dropout_default_old    r$   r%   r+   opset_imports)	r.   r   r/   r   arrayr   r1   r   make_opsetidr2   r!   r"   s      r   export_default_oldzDropout.export_default_old   s    {$$5E % 
 
 HZZZ  ''
333C+!.r2667	
 	
 	
 	
 	
 	
r   c            
     $   t           j                            ddgdgd          } t          j                            ddd                              t          j                  }|}t          | |g|gd	t          j	        d
d          g           d S )Nr   r!   r"   g?)r$   r%   ratior&   r'   r(   test_dropout_random_oldra   rb   rc   )
r.   r   r/   r   r   r0   r   r1   r   rf   rg   s      r   export_random_oldzDropout.export_random_old   s    {$$5E	 % 
 
 IOOAq!$$++BJ773C*!.r2667	
 	
 	
 	
 	
 	
r   N)r   r   )__name__
__module____qualname__staticmethodr3   r9   r>   rA   rJ   rO   rS   rV   rZ   r]   rh   rl    r   r   r   r      s       K K K \K 	T 	T 	T \	T S S S \S 
 
 
 \
 
 
 
 \
 
 
 
 \
" 
R 
R 
R \
R 
 
 
 \
 
 
 
 \
 
 
 
 \
& 
 
 
 \
" 
 
 
 \
 
 
r   r   )r   r   FF)
__future__r   numpyr   r.   r   onnx.backend.test.case.baser   onnx.backend.test.case.noder   r   r   rq   r   r   <module>rv      s    # " " " " "            , , , , , , . . . . . .   u
 u
 u
 u
 u
d u
 u
 u
 u
 u
r   