
     Ti'                     b    S r SSKrSSKrSSKJr  SSKJr  SSK	J
r
  S rS r " S S	5      rS
 rg)zw
A one-layer Whisper encoder model test case, with inputs: audio_features.
This is an onnxscript version of the model.
    N)script)opset18)FLOATc                    ^ ^^^^^^^^^	^
^^^^^^^^^^^ [        5       S[        S   S[        S   4UUUUU UUUUUUUUU	UUUUUUUU
4S jj5       nUR                  5       nU$ )Naudio_features)   P     return)r       c                 
  > [         R                  " U SSS9n[         R                  " U T<T;SSS/SS/S/S9n[         R                  " US5      n[         R                  " U5      n[         R
                  " US5      n[         R                  " SU5      n[         R                  " X&5      n[         R                  " UT>T=SSS/SS	/S/S9n[         R                  " US5      n	[         R                  " U	5      n
[         R
                  " U
S5      n[         R                  " SU5      n[         R                  " X5      n[         R                  " U/ S
QS9n[         R
                  " UT?5      n[         R                  " UTJTISSSS9n[         R                  " TNSS/S9n[         R                  " UU5      n[         R
                  " UTM5      n[         R                  " US5      n[         R                  " US/S/S/SS9n[         R                  " UUSS9n[         R                  " U/ SQS9n[         R                  " THSS/S9n[         R                  " UU5      n[         R                  " US/S/S/SS9n[         R                  " UUSS9n[         R                  " TPSS/S9n[         R                  " UU5      n[         R
                  " UTO5      n[         R                  " US/S/S/SS9n[         R                  " UUSS9n [         R                  " U / SQS9n![         R                  " U/ SQS9n"[         R                  " UU"5      n#[         R                  " U#SS9n$[         R                  " U$U!5      n%[         R                  " U%/ SQS9n&[         R                  " US/S/SS9n'[         R                  " U&U'SS9n([         R                  " TLSS/S9n)[         R                  " U(U)5      n*[         R
                  " U*TK5      n+[         R
                  " UU+5      n,[         R                  " U,TGTFSSSS9n-[         R                  " TCSS/S9n.[         R                  " U-U.5      n/[         R
                  " U/TB5      n0[         R                  " U0S5      n1[         R                  " U15      n2[         R
                  " U2S5      n3[         R                  " SU35      n4[         R                  " U0U45      n5[         R                  " TESS/S9n6[         R                  " U5U65      n7[         R
                  " U7TD5      n8[         R
                  " U,U85      n9[         R                  " U9TAT@SSSS9n:U:$ )Nr   r   )endstartNOTSET)grouppadsauto_padstrides	dilationsgU?g      ?g      ?   )r   r   r   )permg   >)
stash_typeepsilonaxisg      ?r      @   )r   )	allowzero)r   r   r      )r   r   r    r   r   )r   ShapeConvDivErfAddMul	TransposeLayerNormalizationMatMulConcatReshapeSoftmax)Qr   val_0conv1dval_2val_3val_5val_7geluconv1d_1val_9val_10val_12val_14gelu_1permuteadd_20
layer_normval_17val_18linearmul_18val_25view	transposeval_27linear_1val_31view_1val_33val_34linear_2val_37view_2transpose_2transpose_3matmulsoftmaxmatmul_1transpose_4val_42_unsafe_viewval_44val_45linear_3add_141layer_norm_1val_48val_49linear_4val_51val_52val_54val_56gelu_2val_57val_58linear_5add_170layer_norm_2encoder_encoder_conv1_biasencoder_encoder_conv1_weightencoder_encoder_conv2_biasencoder_encoder_conv2_weight&encoder_encoder_embed_positions_weightencoder_encoder_layer_norm_bias!encoder_encoder_layer_norm_weight!encoder_encoder_layers_0_fc1_bias#encoder_encoder_layers_0_fc1_weight!encoder_encoder_layers_0_fc2_bias#encoder_encoder_layers_0_fc2_weight.encoder_encoder_layers_0_final_layer_norm_bias0encoder_encoder_layers_0_final_layer_norm_weight0encoder_encoder_layers_0_self_attn_k_proj_weight2encoder_encoder_layers_0_self_attn_layer_norm_bias4encoder_encoder_layers_0_self_attn_layer_norm_weight0encoder_encoder_layers_0_self_attn_out_proj_bias2encoder_encoder_layers_0_self_attn_out_proj_weight.encoder_encoder_layers_0_self_attn_q_proj_bias0encoder_encoder_layers_0_self_attn_q_proj_weight.encoder_encoder_layers_0_self_attn_v_proj_bias0encoder_encoder_layers_0_self_attn_v_proj_weightsQ                                                              j/var/www/html/ai-image-ml/venv/lib/python3.13/site-packages/onnxscript/rewriter/models/_whisper_encoder.py
main_graphmake_model.<locals>.main_graph)   s    n!1=(&QCc	
 FI.E"E3'C'{{6)<<(&QCc	
 Hi0U#VS)S&)X.##F;W&LM//@>)

 ""<Aq6
 
F3V%STVU+vsRDqAvv;%%d>	""<Aq6
 >>*f5taS2$Q?6Q?""<Aq6
 
F3;;v'UVtaS2$Q?6Q?''\B''\B	;7//&r2>>';7''|Dvu1={FaH"">aV
 f5;;v'WX++fh/11<:)
 ""#FaQRVTf5;;v'HIXy1V$VS)S&)Xv.""#FaQRVT/;;v'HI++gx011-+)
     )r   r   to_model_proto)rk   rh   rg   rj   ri   rv   ru   rz   ry   rt   r|   r{   rx   rw   rs   rr   ro   rn   rq   rp   rm   rl   r~   models   ``````````````````````  r}   
make_modelr      sa    0 Xlk*l	|	l l l l l\ %%'ELr   c                  
   [         R                  R                  S5        [         R                  R                  SS5      R	                  [         R
                  5      n [         R                  R                  SSS5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  SSS5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  SS5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  SS5      R	                  [         R
                  5      n	[         R                  R                  SS5      R	                  [         R
                  5      n
[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  SS5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  SS5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  SS5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[         R                  R                  S5      R	                  [         R
                  5      n[        U UUUUUUUUU	U
UUUUUUUUUUU5      nU$ )N
   r   r   r	   r    i   )nprandomseedrandastypefloat32r   )rk   rh   rg   rj   ri   rv   ru   rz   ry   rt   r|   r{   rx   rw   rs   rr   ro   rn   rq   rp   rm   rl   r   s                          r}   make_model_with_random_weightsr      s   IINN2-/YY^^D#-F-M-Mbjj-Y*#%99>>#r1#=#D#DRZZ#P !#!4!;!;BJJ!G#%99>>#sA#>#E#Ebjj#Q !#!4!;!;BJJ!G;=99>>#;N;U;U


<8 :<9L9S9STVT^T^9_679yy~~c37O7V7V


84 68YY^^C5H5O5OPRPZPZ5[279yy~~c37O7V7V


84 8:yy~~c37O7V7V


84 68YY^^C5H5O5OPRPZPZ5[29;S9Q9X9X


:6 8:yy~~c7J7Q7QRTR\R\7]479yy~~c7J7Q7QRTR\R\7]457YY^^C5H5O5OPRPZPZ5[2*,))..s*C*J*J2::*V'(*		t(<(C(CBJJ(O%*,))..d*C*J*J2::*V'(*		s(;(B(B2::(N%(*		s(;(B(B2::(N%&(iinnS&9&@&@&L#.$"$"<:86886:886+)+))'-E0 Lr   c                        \ rS rSrS rS rSrg)_WhisperEncoderTest   c                     [        U S5      (       d/  [        5       n[        R                  R	                  U5      nX l        U R
                  $ )N_onnx_model)hasattrr   irserdedeserialize_modelr   )selfmodel_protor   s      r}   get_onnx_model"_WhisperEncoderTest.get_onnx_model   s=    t]++8:KHH..{;E$r   c                    [        U S5      (       de  [        R                  R                  S5        S[        R                  R	                  SSS5      R                  [        R                  5      0nXl        U R                  $ )N_ort_inputsr   r   r   r	   r
   )r   r   r   r   r   r   r   r   )r   inputss     r}   get_ort_inputs"_WhisperEncoderTest.get_ort_inputs   s_    t]++IINN2 "))..B"="D"DRZZ"PF  &r   )r   r   N)__name__
__module____qualname____firstlineno__r   r   __static_attributes__ r   r}   r   r      s      r   r   c                      [        5       $ )N)r   r   r   r}   whisper_encoder_testr      s      r   )__doc__numpyr   onnx_irr   
onnxscriptr   onnxscript.onnx_opsetr   onnxscript.onnx_typesr   r   r   r   r   r   r   r}   <module>r      s7   
    ) 'HV:z   $!r   