o
    i eh                     @   s8  d Z ddlZddlZddlZddlZddlZddlm  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mZ dd	lmZ G d
d deZeddedg dG dd deZeddedg dG dd deZeddedg dG dd deZdd Zdd Zdd Zdd ZdS )z!Module implementing RNN wrappers.    N)lstm)AbstractRNNCell)serialization_lib)generic_utils)
tf_inspect)	tf_export)
deprecatedc                       s~   e Zd ZdZ fddZdd Zdd Zdd	 Zed
d Z	edd Z
edd Zdd Z fddZedddZ  ZS )_RNNCellWrapperzBase class for cells wrappers V2 compatibility.

    This class along with `rnn_cell_impl._RNNCellWrapperV1` allows to define
    wrappers that are compatible with V1 and V2, and defines helper methods for
    this purpose.
    c                    s@   t  j|i | || _t|j}d|jv p|jd u| j_	d S )Ntraining)
super__init__cellr   getfullargspeccallargsvarkw
_call_specexpects_training_arg)selfr   r   kwargscell_call_spec	__class__ Q/var/www/myenv/lib/python3.10/site-packages/keras/src/layers/rnn/cell_wrappers.pyr   3   s   
z_RNNCellWrapper.__init__c                 K   s   t )a  Calls the wrapped cell and performs the wrapping logic.

        This method is called from the wrapper's `call` or `__call__` methods.

        Args:
          inputs: A tensor with wrapped cell's input.
          state: A tensor or tuple of tensors with wrapped cell's state.
          cell_call_fn: Wrapped cell's method to use for step computation
            (cell's `__call__` or 'call' method).
          **kwargs: Additional arguments.

        Returns:
          A pair containing:
          - Output: A tensor with cell's output.
          - New state: A tensor or tuple of tensors with new wrapped cell's
            state.
        )NotImplementedErrorr   inputsstatecell_call_fnr   r   r   r   _call_wrapped_cell;   s   z"_RNNCellWrapper._call_wrapped_cellc                 K   s   | j ||fd| jji|S )a  Runs the RNN cell step computation.

        When `call` is being used, we assume that the wrapper object has been
        built, and therefore the wrapped cells has been built via its `build`
        method and its `call` method can be used directly.

        This allows to use the wrapped cell and the non-wrapped cell
        equivalently when using `call` and `build`.

        Args:
          inputs: A tensor with wrapped cell's input.
          state: A tensor or tuple of tensors with wrapped cell's state.
          **kwargs: Additional arguments passed to the wrapped cell's `call`.

        Returns:
          A pair containing:

          - Output: A tensor with cell's output.
          - New state: A tensor or tuple of tensors with new wrapped cell's
            state.
        r   )r    r   r   )r   r   r   r   r   r   r   r   O   s   z_RNNCellWrapper.callc                 C   s   | j | d| _dS )zBuilds the wrapped cell.TN)r   buildbuilt)r   inputs_shaper   r   r   r!   i   s   
z_RNNCellWrapper.buildc                 C   s   | j S N)r   r   r   r   r   wrapped_celln   s   z_RNNCellWrapper.wrapped_cellc                 C      | j jS r$   )r   
state_sizer%   r   r   r   r(   r      z_RNNCellWrapper.state_sizec                 C   r'   r$   )r   output_sizer%   r   r   r   r*   v   r)   z_RNNCellWrapper.output_sizec                 C   sF   t t| jd  | j||W  d    S 1 sw   Y  d S N	ZeroState)tf
name_scopetype__name__r   
zero_stater   
batch_sizedtyper   r   r   r1   z   s   $z_RNNCellWrapper.zero_statec                    s@   d| j jj| j  di}t  }tt| t|  S )Nr   )
class_nameconfig)r   r   r0   
get_configr   dictlistitemsr   r6   base_configr   r   r   r7   ~   s   
z_RNNCellWrapper.get_configNc                 C   s6   |  }ddlm} ||d|d}| |fi |S )Nr   )deserializer   custom_objects)copykeras.src.layers.serializationr=   pop)clsr6   r?   deserialize_layerr   r   r   r   from_config   s   
z_RNNCellWrapper.from_configr$   )r0   
__module____qualname____doc__r   r    r   r!   propertyr&   r(   r*   r1   r7   classmethodrE   __classcell__r   r   r   r   r	   +   s     



r	   z'Please use tf.keras.layers.RNN instead.znn.RNNCellDropoutWrapper)v1c                       st   e Zd ZdZ								d fdd	Zdd Zd	d
 Z	dddZdd Z fddZ	e
d fdd	Z  ZS )DropoutWrapperz@Operator adding dropout to inputs and outputs of the given cell.      ?FNc
                    s  t |tjr
tdt j|fdi|
 |	dur%t|	s%td|	 |	p(t_	t
dM dd }|df|d	f|d
ffD ]5\}}||\}}|durk|dk sU|dkr_td| d| td| t| q?td| | q?W d   n1 sw   Y  |_|_|_d_d_d_|rdu rtddd fdd t jtjrjdk r|du rtdt| fdd|_t|j fdd|j_t|j fdd|j_dS dS )a  Create a cell with added input, state, and/or output dropout.

        If `variational_recurrent` is set to `True` (**NOT** the default
        behavior), then the same dropout mask is applied at every step, as
        described in: [A Theoretically Grounded Application of Dropout in
        Recurrent Neural Networks. Y. Gal, Z.
        Ghahramani](https://arxiv.org/abs/1512.05287).

        Otherwise a different dropout mask is applied at every time step.

        Note, by default (unless a custom `dropout_state_filter` is provided),
        the memory state (`c` component of any `LSTMStateTuple`) passing through
        a `DropoutWrapper` is never modified.  This behavior is described in the
        above article.

        Args:
          cell: an RNNCell, a projection to output_size is added to it.
          input_keep_prob: unit Tensor or float between 0 and 1, input keep
            probability; if it is constant and 1, no input dropout will be
            added.
          output_keep_prob: unit Tensor or float between 0 and 1, output keep
            probability; if it is constant and 1, no output dropout will be
            added.
          state_keep_prob: unit Tensor or float between 0 and 1, output keep
            probability; if it is constant and 1, no output dropout will be
            added.  State dropout is performed on the outgoing states of the
            cell. **Note** the state components to which dropout is applied when
            `state_keep_prob` is in `(0, 1)` are also determined by the argument
            `dropout_state_filter_visitor` (e.g. by default dropout is never
            applied to the `c` component of an `LSTMStateTuple`).
          variational_recurrent: Python bool.  If `True`, then the same dropout
            pattern is applied across all time steps per run call. If this
            parameter is set, `input_size` **must** be provided.
          input_size: (optional) (possibly nested tuple of) `TensorShape`
            objects containing the depth(s) of the input tensors expected to be
            passed in to the `DropoutWrapper`.  Required and used **iff**
            `variational_recurrent = True` and `input_keep_prob < 1`.
          dtype: (optional) The `dtype` of the input, state, and output tensors.
            Required and used **iff** `variational_recurrent = True`.
          seed: (optional) integer, the randomness seed.
          dropout_state_filter_visitor: (optional), default: (see below).
            Function that takes any hierarchical level of the state and returns
            a scalar or depth=1 structure of Python booleans describing which
            terms in the state should be dropped out.  In addition, if the
            function returns `True`, dropout is applied across this sublevel.
            If the function returns `False`, dropout is not applied across this
            entire sublevel.  Default behavior: perform dropout on all terms
            except the memory (`c`) state of `LSTMCellState` objects, and don't
            try to apply dropout to
            `TensorArray` objects:
            ```
            def dropout_state_filter_visitor(s):
              # Never perform dropout on the c state.
              if isinstance(s, LSTMCellState):
                return LSTMCellState(c=False, h=True)
              elif isinstance(s, TensorArray):
                return False
              return True
            ```
          **kwargs: dict of keyword arguments for base layer.

        Raises:
          TypeError: if `cell` is not an `RNNCell`, or `keep_state_fn` is
            provided but not `callable`.
          ValueError: if any of the keep_probs are not between 0 and 1.
        zokeras LSTM cell does not work with DropoutWrapper. Please use LSTMCell(dropout=x, recurrent_dropout=y) instead.r4   Nz9dropout_state_filter_visitor must be callable. Received: DropoutWrapperInitc                 S   s   t | }t |}||fS r$   )r-   convert_to_tensorget_static_value)vtensor_valueconst_valuer   r   r   tensor_and_const_value   s   

z7DropoutWrapper.__init__.<locals>.tensor_and_const_valueinput_keep_probstate_keep_proboutput_keep_probr      z
Parameter z# must be between 0 and 1. Received _z7When variational_recurrent=True, dtype must be providedc                 S   s   t dgt |  fdS )NrY   r   )r-   concatTensorShapeas_list)sr   r   r   convert_to_batch_shape   s   z7DropoutWrapper.__init__.<locals>.convert_to_batch_shapec                    s    | }t jj||dS )N)seedr4   )r-   randomuniform)r^   
inner_seedshape)r_   r4   r   r   batch_noise&  s   z,DropoutWrapper.__init__.<locals>.batch_noiserN   zdWhen variational_recurrent=True and input_keep_prob < 1.0 or is unknown, input_size must be providedc                        | d| dS )Ninputrc   	_gen_seedir^   re   r   r   r   <lambda>5      z)DropoutWrapper.__init__.<locals>.<lambda>c                    rf   )Nr   rh   ri   rk   rm   r   r   rn   <  ro   c                    rf   )Noutputrh   ri   rk   rm   r   r   rn   C  ro   )
isinstancer   LSTMCell
ValueErrorr   r   callable	TypeError%_default_dropout_state_filter_visitor_dropout_state_filterr-   r.   setattrfloat_variational_recurrent_input_size_seed_recurrent_input_noise_recurrent_state_noise_recurrent_output_noise_input_keep_probnumbersReal_enumerated_map_structure_up_tor(   r*   )r   r   rV   rX   rW   variational_recurrent
input_sizer4   r`   dropout_state_filter_visitorr   rU   probattrtensor_prob
const_probr   )re   r_   r4   r   r   r      s   O


zDropoutWrapper.__init__c                 C   sN   | j d u rd S d||f }t| j | d}tt| d d dd@ S )Nz%s_%dzutf-8      i)r|   strencodeinthashlibmd5	hexdigest)r   salt_prefixindexsaltstringr   r   r   rj   I  s
   
 zDropoutWrapper._gen_seedc                 C   s4   || }t |}t ||| }||  |S )z7Performs dropout given the pre-calculated noise tensor.)r-   floordivide	set_shape	get_shape)r   unused_indexvaluenoise	keep_probrandom_tensorbinary_tensorretr   r   r   $_variational_recurrent_dropout_valueP  s
   
z3DropoutWrapper._variational_recurrent_dropout_valuec                    s^   |du r|}j s fdd}t||g||gR  S  fdd}t||g|||gR  S )zADecides whether to perform standard dropout or recurrent dropout.Nc                    s0   t |tr|rtjj|d  | dS |S )NrN   )rater`   )rq   boolr-   nndropoutrj   )rl   
do_dropoutrR   r   r   r   r   r   r   n  s   
z(DropoutWrapper._dropout.<locals>.dropoutc                    s"   t |tr|r| || S |S r$   )rq   r   r   )rl   r   rR   n)r   r   r   r   r     s
   )rz   r   )r   valuesr   recurrent_noiser   shallow_filtered_substructurer   r   r   r   _dropout]  s"   

zDropoutWrapper._dropoutc           	      K   s   dd }|| j r| |d| j| j }|||fi |\}}|| jr7tjj| j|}| |d| j	| j|}|| j
rF| |d| j| j
}||fS )a'  Runs the wrapped cell and applies dropout.

        Args:
          inputs: A tensor with wrapped cell's input.
          state: A tensor or tuple of tensors with wrapped cell's state.
          cell_call_fn: Wrapped cell's method to use for step computation
            (cell's `__call__` or 'call' method).
          **kwargs: Additional arguments.

        Returns:
          A pair containing:

          - Output: A tensor with cell's output.
          - New state: A tensor or tuple of tensors with new wrapped cell's
            state.
        c                 S   s   t | t p	| dk S )NrY   )rq   ry   )pr   r   r   _should_dropout  s   z:DropoutWrapper._call_wrapped_cell.<locals>._should_dropoutrg   r   rp   )r   r   r}   _state_keep_probr-   __internal__nestget_traverse_shallow_structurerw   r~   _output_keep_probr   )	r   r   r   r   r   r   rp   	new_stater   r   r   r   r      s:   


z!DropoutWrapper._call_wrapped_cellc                    sp   | j | j| j| j| j| jd}| jtkr%t| j\}}}|	|||d t
  }tt| t|  S )z*Returns the config of the dropout wrapper.)rV   rX   rW   r   r   r`   )
dropout_fndropout_fn_typedropout_fn_module)r   r   r   rz   r{   r|   rw   rv   _serialize_function_to_configupdater   r7   r8   r9   r:   )r   r6   functionfunction_typefunction_moduler<   r   r   r   r7     s*   

zDropoutWrapper.get_configc                    sF   d|v r|  }t||ddd}|d ||d< tt| j||dS )Nr   r   r   r   r>   )r@   _parse_config_to_functionrB   r   rM   rE   )rC   r6   r?   dropout_state_filterr   r   r   rE     s   

zDropoutWrapper.from_config)rN   rN   rN   FNNNNr$   )r0   rF   rG   rH   r   rj   r   r   r    r7   rJ   rE   rK   r   r   r   r   rM      s(     2
05rM   znn.RNNCellResidualWrapperc                       sH   e Zd ZdZd fdd	Zdd Z fddZed fd	d
	Z  Z	S )ResidualWrapperzBRNNCell wrapper that ensures cell inputs are added to the outputs.Nc                       t  j|fi | || _dS )a  Constructs a `ResidualWrapper` for `cell`.

        Args:
          cell: An instance of `RNNCell`.
          residual_fn: (Optional) The function to map raw cell inputs and raw
            cell outputs to the actual cell outputs of the residual network.
            Defaults to calling nest.map_structure on (lambda i, o: i + o),
            inputs and outputs.
          **kwargs: dict of keyword arguments for base layer.
        N)r   r   _residual_fn)r   r   residual_fnr   r   r   r   r     s   
zResidualWrapper.__init__c           	         sB   |||fi |\}}dd   fdd}| j p|||}||fS )aG  Run the cell and apply the residual_fn.

        Args:
          inputs: cell inputs.
          state: cell state.
          cell_call_fn: Wrapped cell's method to use for step computation
            (cell's `__call__` or 'call' method).
          **kwargs: Additional arguments passed to the wrapped cell's `call`.

        Returns:
          Tuple of cell outputs and new state.

        Raises:
          TypeError: If cell inputs and outputs have different structure (type).
          ValueError: If cell inputs and outputs have different structure
            (value).
        c                 S   s   |   |   d S r$   )r   assert_is_compatible_withinpoutr   r   r   assert_shape_match  s   z>ResidualWrapper._call_wrapped_cell.<locals>.assert_shape_matchc                    s2   t j| | t j | | t jdd | |S )Nc                 S   s   | | S r$   r   r   r   r   r   rn     s    zQResidualWrapper._call_wrapped_cell.<locals>.default_residual_fn.<locals>.<lambda>)r-   r   assert_same_structuremap_structure)r   outputsr   r   r   default_residual_fn  s
   
z?ResidualWrapper._call_wrapped_cell.<locals>.default_residual_fn)r   )	r   r   r   r   r   r   r   r   res_outputsr   r   r   r      s   z"ResidualWrapper._call_wrapped_cellc                    sR   | j durt| j \}}}|||d}ni }t  }tt| t|  S )z+Returns the config of the residual wrapper.N)r   residual_fn_typeresidual_fn_module)r   r   r   r7   r8   r9   r:   )r   r   r   r   r6   r<   r   r   r   r7   %  s   

zResidualWrapper.get_configc                    s<   d|v r|  }t||ddd}||d< tt| j||dS )Nr   r   r   r>   )r@   r   r   r   rE   )rC   r6   r?   residual_functionr   r   r   rE   7  s   
zResidualWrapper.from_configr$   )
r0   rF   rG   rH   r   r    r7   rJ   rE   rK   r   r   r   r   r     s    $r   znn.RNNCellDeviceWrapperc                       s<   e Zd ZdZ fddZdd Zdd Z fdd	Z  ZS )
DeviceWrapperz=Operator that ensures an RNNCell runs on a particular device.c                    r   )aO  Construct a `DeviceWrapper` for `cell` with device `device`.

        Ensures the wrapped `cell` is called with `tf.device(device)`.

        Args:
          cell: An instance of `RNNCell`.
          device: A device string or function, for passing to `tf.device`.
          **kwargs: dict of keyword arguments for base layer.
        N)r   r   _device)r   r   devicer   r   r   r   r   M  s   

zDeviceWrapper.__init__c              	   C   s   t t| jd / t jj| j | j	||W  d    W  d    S 1 s,w   Y  W d    d S 1 s<w   Y  d S r+   )
r-   r.   r/   r0   compatrL   r   r   r   r1   r2   r   r   r   r1   Z  s   "zDeviceWrapper.zero_statec                 K   sF   t jj| j |||fi |W  d   S 1 sw   Y  dS )z!Run the cell on specified device.N)r-   r   rL   r   r   r   r   r   r   r    _  s   $z DeviceWrapper._call_wrapped_cellc                    s0   d| j i}t  }tt| t|  S )Nr   )r   r   r7   r8   r9   r:   r;   r   r   r   r7   d  s   

zDeviceWrapper.get_config)	r0   rF   rG   rH   r   r1   r    r7   rK   r   r   r   r   r   H  s    r   c                 C   sX   t | tjrt| }d}| j}nt| r| j}d}| j}n	tdt	|  |||fS )z(Serialize the function for get_config().lambdar   z&Unrecognized function type for input: )
rq   python_types
LambdaTyper   	func_dumprF   rt   r0   rs   r/   )r   rp   output_typemoduler   r   r   r   j  s   

r   c           	      C   s   t  }| |d}|tjv r|tj| j n|dur'tjd|t	dd |r.|| | |}|dkrCt
j| | |dd}|S |dkrZt
 rOtd	tj| | |d
}|S td| d)z)Reconstruct the function from the config.Nz;{} is not loaded, but a layer uses it. It may cause errors.   )
stacklevelr   zfunction in wrapper)r?   printable_module_namer   a1  Requested the deserialization of a layer with a Python `lambda` inside it. This carries a potential risk of arbitrary code execution and thus it is disallowed by default. If you trust the source of the saved model, you can pass `safe_mode=False` to the loading function in order to allow `lambda` loading.)globsz Unknown function type received: z+. Expected types are ['function', 'lambda'])globalsrB   sysmodulesr   __dict__warningswarnformatUserWarningr   deserialize_keras_objectin_safe_moders   r   	func_loadru   )	r6   r?   func_attr_namefunc_type_attr_namemodule_attr_namer   r   r   r   r   r   r   r   |  s>   




r   c                 C   s   t | tj S r$   )rq   r-   TensorArray)substater   r   r   rv     s   rv   c                    s2   dg  fdd}t jjj| |g|R i |S )Nr   c                     s.    d g| R i |} d  d7  < |S )Nr   rY   r   )
inner_argsinner_kwargsrixmap_fnr   r   enumerated_fn  s   z6_enumerated_map_structure_up_to.<locals>.enumerated_fn)r-   r   r   map_structure_up_to)shallow_structurer   r   r   r   r   r   r   r     s   r   ) rH   r   r   r   typesr   r   tensorflow.compat.v2r   v2r-   keras.src.layers.rnnr   &keras.src.layers.rnn.abstract_rnn_cellr   keras.src.savingr   keras.src.utilsr   r    tensorflow.python.util.tf_exportr   "tensorflow.python.util.deprecationr   r	   rM   r   r   r   r   rv   r   r   r   r   r   <module>   s<   	h
  [
X
 3