o
    ={c=*                     @   s  d Z ddlZddlZedg dZdd Zdd Zd	d
 Zdd Zdd Z	dd Z
dd Zdd Zdd ZG dd dejZdd Zdd Zdd Zdd  Zd!d" Zdefdefd
efdefde	fde
fdefdefdefdefdefdefd efd"efgZd%d#d$ZdS )&z2Tests for saving/loading function for keras Model.    NModelFn)modelinput_shapeZtarget_shapec                  C   s4   t t jjddddt jjdddg} t| dd	S )
zBasic sequential model.   relur   
activationr      softmaxr	   Nr   Nr
   keras
SequentiallayersDenser   r    r   ?lib/python3.10/site-packages/keras/tests/model_architectures.pybasic_sequential   s   r   c                  C   s2   t t jjdddt jjdddg} t| ddS )z+Sequential model with deferred input shape.r   r   r   r
   r   r   r   r   r   r   r   r   basic_sequential_deferred&   s   r   c                  C   sT   t d} t jdd tdD }|| }t jd|}t | |}t|ddS )zStacked RNN model.r   c                 S   s   g | ]}t jd qS r
   )r   r   ZLSTMCell).0_r   r   r   
<listcomp>4   s    zstacked_rnn.<locals>.<listcomp>r   r
   N   r   r   )r   Inputr   ZRNNranger   Modelr   )inputsZlayerxoutputsr   r   r   r   stacked_rnn1   s   
r%   c                  C   sn   t d} t jjddd| }t jjddd|}t jjddd|}t jd|}t | |}t|dd	S )
zLSTM model.r   r   T)return_sequencesr   r
   Fr   r   r   r   r   LSTMr   r!   r   r"   r#   r$   r   r   r   r   lstm;   s   
r*   c                  C   s   t jddd} t jddd}t jdd| }t jd|}t j||g}t jjd	d
dd|}t jjdddd|}t j| |g||gd}t|ddgddgS )zMulti-input Multi-output model.Nbodyshapenamer   tags
   r      r
   sigmoidpriorityr	   r/   r   r   Z
departmentr"   r$   N   r   r   )	r   r   r   	Embeddingr(   concatenater   r!   r   )Z
body_inputZ
tags_inputr#   Zbody_featuresZpred_1Zpred_2r   r   r   r   multi_input_multi_outputF   s   r;   c                  C   sh   t t jjddddt jjdddg} t jdd}| |}t jjddd|}t ||}t|d	d
S )z0A sequential model nested in a functional model.r   r   r   r   r
   r   r.   r   r   r   )r   r   r   r   r   r!   r   )inner_modelr"   r#   r$   r   r   r   r   nested_sequential_in_functionalX   s   r>   c                  C   s   d} d}d}t jd| fd}t jj|dd}||\}}}||g}t jd|fd}	t jj|ddd}
|
|	|d\}}}t jj|d	d
}||}t ||	g|}t|dd| fdd|fgdd|fS )zSequence to sequence model.r   r
   Nr<   T)return_state)r&   r?   )Zinitial_stater   r   r'   )Znum_encoder_tokensZnum_decoder_tokensZ
latent_dimZencoder_inputsencoderr   Zstate_hZstate_cZencoder_statesZdecoder_inputsZdecoder_lstmZdecoder_outputsZdecoder_denser   r   r   r   
seq_to_seqh   s,   rA   c                  C   s   t jdddd} t jjdddd| }t jd	|}t jjd
ddd|}t jddd}t j||g}t jjddd|}t jjd
ddd|}t j| |g||gd}t|ddgddgS )z#Shared layer in a functional model.)r1   Zint32
main_input)r.   Zdtyper/   r2   r   r1   )Z
output_dim	input_dimZinput_lengthr   r8   r3   Z
aux_outputr5   )r2   Z	aux_inputr-   r
   r   r   main_outputr6   )Nr1   )Nr2   r7   )	r   r   r   r9   r(   r   r:   r!   r   )rB   r#   Zlstm_outZauxiliary_outputZauxiliary_inputrD   r   r   r   r   shared_layer_functional   s.   rE   c                  C   s   t t jjddddt jjddddg} t d}t d}| |}| |}t j||g}t j |}t ||g|}t|ddgdS )z.Shared sequential model in a functional model.r
   r   r   r   )r2   r2   r   )Nr2   r2   r   )Nr   )	r   r   r   ZConv2Dr   r:   ZGlobalAveragePooling2Dr!   r   )r=   Zinputs_1Zinputs_2Zx1Zx2r#   r$   r   r   r   r   shared_sequential   s   

rF   c                       s>   e Zd ZdZd fdd	Zdd Zdd Zed	d
 Z  Z	S )MySubclassModelzA subclass model.r   c                    sZ   t  jdd d|i| _tjjddd| _tjjddd| _tj | _	tj
d	| _d S )
NZmy_subclass_modelr/   rC      r   r   r
   r         ?)super__init___configr   r   r   dense1dense2BatchNormalizationbnDropoutdp)selfrC   	__class__r   r   rL      s   
zMySubclassModel.__init__c                 K   s(   |  |}| |}| |}| |S r+   )rN   rS   rQ   rO   )rT   r"   kwargsr#   r   r   r   call      



zMySubclassModel.callc                 C   s   | j S r+   )rM   rT   r   r   r   
get_config   s   zMySubclassModel.get_configc                 C   s   | di |S )Nr   r   )clsZconfigr   r   r   from_config   s   zMySubclassModel.from_configr   )
__name__
__module____qualname____doc__rL   rX   r[   classmethodr]   __classcell__r   r   rU   r   rG      s    rG   c                  C   s    G dd dt j} t|  ddS )z2A subclass model nested in another subclass model.c                       (   e Zd ZdZ fddZdd Z  ZS )z4nested_subclassed_model.<locals>.NestedSubclassModelzA nested subclass model.c                    sF   t    tjjddd| _tjjddd| _tj | _t	 | _
d S )Nr   r   r   r
   )rK   rL   r   r   r   rN   rO   rP   rQ   rG   inner_subclass_modelrZ   rU   r   r   rL      s
   
z=nested_subclassed_model.<locals>.NestedSubclassModel.__init__c                 S   s(   |  |}| |}| |}| |S r+   )rN   rQ   re   rO   rT   r"   r#   r   r   r   rX      rY   z9nested_subclassed_model.<locals>.NestedSubclassModel.callr^   r_   r`   ra   rL   rX   rc   r   r   rU   r   NestedSubclassModel   s    rh   r   r   r   r!   r   )rh   r   r   r   nested_subclassed_model   s   rj   c                  C   sT   t  } tjdd}| |}tj |}tjjddd|}t||}t|ddS )z.A subclass model nested in a functional model.r   r<   r
   r   r   r   r   )rG   r   r   r   rP   r   r!   r   )re   r"   r#   r$   r   r   r   r   %nested_subclassed_in_functional_model   s   rk   c                     s,   dd  G  fdddt j} t|  ddS )z.A functional model nested in a subclass model.c                  S   sJ   t jdd} t jjddd| }t j |}t jd|}t | |S )N)r   r<   r   r   r   r
   )r   r   r   r   rP   r!   )r"   r#   r$   r   r   r   get_functional_model   s
   zCnested_functional_in_subclassed_model.<locals>.get_functional_modelc                       s*   e Zd ZdZ fddZdd Z  ZS )zNnested_functional_in_subclassed_model.<locals>.NestedFunctionalInSubclassModelz&A functional nested in subclass model.c                    s>   t  jdd tjjddd| _tjjddd| _ | _d S )N%nested_functional_in_subclassed_modelrH   r   r   r   r
   )rK   rL   r   r   r   rN   rO   inner_functional_modelrZ   )rV   rl   r   r   rL      s   zWnested_functional_in_subclassed_model.<locals>.NestedFunctionalInSubclassModel.__init__c                 S   s   |  |}| |}| |S r+   )rN   rn   rO   rf   r   r   r   rX      s   


zSnested_functional_in_subclassed_model.<locals>.NestedFunctionalInSubclassModel.callrg   r   rl   rU   r   NestedFunctionalInSubclassModel   s    rp   r   r   ri   )rp   r   ro   r   rm      s   rm   c                  C   s    G dd dt j} t|  ddS )z!Shared layer in a subclass model.c                       rd   )z?shared_layer_subclassed_model.<locals>.SharedLayerSubclassModelz$A subclass model with shared layers.c                    s>   t  jdd tjjddd| _tjd| _tj | _	d S )NZshared_layer_subclass_modelrH   r   r   r   rJ   )
rK   rL   r   r   r   denserR   rS   rP   rQ   rZ   rU   r   r   rL     s   zHshared_layer_subclassed_model.<locals>.SharedLayerSubclassModel.__init__c                 S   s(   |  |}| |}| |}|  |S r+   )rq   rS   rQ   rf   r   r   r   rX     rY   zDshared_layer_subclassed_model.<locals>.SharedLayerSubclassModel.callrg   r   r   rU   r   SharedLayerSubclassModel  s    rr   r   ri   )rr   r   r   r   shared_layer_subclassed_model   s   rs   c                  C   sX   t jdd} t jd| }t j |}t jd|}t j| |ddd}t|dd	S )
z%A functional model with keyword args.r   r<   r   r
   mF)r/   Z	trainabler   r   )r   r   r   r   rP   r!   r   r)   r   r   r   functional_with_keyword_args  s   ru   c                    s    fddt D }|S )z,Get all models excluding the specified ones.c                    s   g | ]
}|d   vr|qS )r   r   )r   r   exclude_modelsr   r   r   :  s    zget_models.<locals>.<listcomp>)
ALL_MODELS)rw   Zmodelsr   rv   r   
get_models8  s   ry   r+   )ra   collectionsr   
namedtupler   r   r   r%   r*   r;   r>   rA   rE   rF   r!   rG   rj   rk   rm   rs   ru   rx   ry   r   r   r   r   <module>   sR   
