
    wi                    ^    d dl mZ d dlZd dlmZ  G d de          Z G d de          ZdS )    )annotationsN)OpRunc                  V    e Zd Zd ZddZddZddZddZ	 	 	 	 	 	 	 	 	 	 	 	 	 ddZdS )
CommonLSTMc                r    t          j        | ||           t          |j                  | _        d| _        d S )N   )r   __init__lenoutput	n_outputsn_gatesself	onnx_node
run_paramss      j/root/.openclaw/workspace/chatterbox_venv_py311/lib/python3.11/site-packages/onnx/reference/ops/op_lstm.pyr	   zCommonLSTM.__init__   s3    tY
333Y-..    x
np.ndarrayreturnc                8    ddt          j        |           z   z  S )N   )npexpr   r   s     r   fzCommonLSTM.f   s    Ar

N##r   c                *    t          j        |          S Nr   tanhr   s     r   gzCommonLSTM.g       wqzzr   c                *    t          j        |          S r   r   r   s     r   hzCommonLSTM.h   r"   r   XRBWH_0C_0Pnum_directionsinttuple[np.ndarray, np.ndarray]c	                <   |j         d         }	|j         d         }
|j         d         }t          j        |	|||
g          }g }t          j        |d          \  }}}|}|}t          j        ||j         d         d          D ]1}t          j        |t          j        |                    t          j        |t          j        |                    z   t          j        t          j        |d           z   }t          j        |dd          \  }}}}|                     |||z  z             }|                     |||z  z             }|                     |          }||z  ||z  z   }|                     |||z  z             }|| 	                    |          z  }|
                    |           |}|}3t          j        |          }|dk    r||d d dd d d d f<   | j        dk    r	|d         }n)t          j        |g d          }|d d d d dd d f         }||fS )	Nr   r   r   axis      )r3   r   r   r   )shaper   emptysplitdot	transposeaddr   r!   r$   appendconcatenatelayout)r   r%   r&   r'   r(   r)   r*   r+   r,   
seq_lengthhidden_size
batch_sizeYh_listp_ip_op_fH_tC_tr   gatesior   cCHconcatenatedY_hs                                r   _stepzCommonLSTM._step   s    WQZ
imWQZ
Hj.*kJKK(1a..c3!QWQZa000 	 	Aq",q//**&bl1oo../&"(1a..)* 
 %B//JAq!Qq39}%%Aq39}%%Aq		AC!a%Aq37{##ADFF1IIAMM!CCC~f--Q(AaaaAAAqqqjM;!B%CCQ--AAAAqqq"aaaK.C#vr   Nc           
        d}d}|j         d         }|dk    rBt          j        |d          }t          j        |d          }|?t          |j                   dk    r'|j         d         dk    rt          j        |d          }|?t          |j                   dk    r'|j         d         dk    rt          j        |d          }|?t          |j                   dk    r'|j         d         dk    rt          j        |d          }|?t          |j                   dk    r'|j         d         dk    rt          j        |d          }|?t          |j                   dk    r'|j         d         dk    rt          j        |d          }|j         d         }|j         d         }| j        dk    rt          j        |dd          }|&t          j        d|z  |z  t          j                  }|#t          j        ||z  t          j                  }|"t          j        ||ft          j                  }|"t          j        ||ft          j                  }n t          d	|d
| j	        j
        d          |                     ||||||||          \  }}|                    |j                  }| j        dk    r|fn||                    |j                  fS )Nr4   r   r   r   r1   r0   r3   )dtypezUnsupported value z! for num_directions and operator .)r,   )r5   r   squeezer
   r=   swapaxeszerosfloat32NotImplementedError	__class____name__rP   astyperR   r   )r   r%   r(   r&   r'   sequence_lens	initial_h	initial_cr+   activation_alphaactivation_betaactivationsclip	directionr?   input_forgetr=   r   number_of_peepholesr,   r@   rA   rO   s                          r   _runzCommonLSTM._runL   s   ( Q
11%%%A
11%%%A}QW!1!1agajAooJqq))))+,,q00!'*a// "
=q A A A%	((1,,OA&!++Jyq999	%	((1,,OA&!++Jyq999	}QW!1!1agajAooJqq)))'"+KJ{aK1a((yHQ[;6bjIIIyH0;>bjQQQ Hj+%>bjQQQ	 Hj+%>bjQQQ	%=^ = = $ 7= = =  
 q!Q	9a  
 
3 HHQW~**ttCJJqw4G4G0HHr   )r   r   r   r   )r%   r   r&   r   r'   r   r(   r   r)   r   r*   r   r+   r   r,   r-   r   r.   )NNNNNNNNNNNNN)	rZ   
__module____qualname__r	   r   r!   r$   rP   rf    r   r   r   r      s          
$ $ $ $      0 0 0 0n 
#JI JI JI JI JI JIr   r   c                      e Zd Zd ZdS )LSTMc                >    t                               | ||           d S r   )r   r	   r   s      r   r	   zLSTM.__init__   s     D)Z88888r   N)rZ   rg   rh   r	   ri   r   r   rk   rk      s#        9 9 9 9 9r   rk   )
__future__r   numpyr   onnx.reference.op_runr   r   rk   ri   r   r   <module>rp      s    # " " " " "     ' ' ' ' ' 'KI KI KI KI KI KI KI KI\9 9 9 9 9: 9 9 9 9 9r   