U
    Ph
                     @  sN   d dl mZ d dlmZ d dlmZ d dlmZ ddhZG dd dej	Z
dS )	    )annotationsN)get_act_layer)look_up_optionvitswinc                      s8   e Zd ZdZddddddd	 fd
dZdd Z  ZS )MLPBlockz
    A multi-layer perceptron block, based on: "Dosovitskiy et al.,
    An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale <https://arxiv.org/abs/2010.11929>"
            GELUr   intfloatztuple | strNone)hidden_sizemlp_dimdropout_rateactreturnc                   s   t    d|  krdks(n td|p.|}|dkrDt||nt||d | _t||| _t|| _t	|| _
t|t}|dkrt	|| _n |dkr| j
| _ntdt d	S )
a  
        Args:
            hidden_size: dimension of hidden layer.
            mlp_dim: dimension of feedforward layer. If 0, `hidden_size` will be used.
            dropout_rate: fraction of the input units to drop.
            act: activation type and arguments. Defaults to GELU. Also supports "GEGLU" and others.
            dropout_mode: dropout mode, can be "vit" or "swin".
                "vit" mode uses two dropout instances as implemented in
                https://github.com/google-research/vision_transformer/blob/main/vit_jax/models.py#L87
                "swin" corresponds to one instance as implemented in
                https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_mlp.py#L23


        r      z'dropout_rate should be between 0 and 1.GEGLU   r   r   zdropout_mode should be one of N)super__init__
ValueErrornnLinearlinear1linear2r   fnDropoutdrop1r   SUPPORTED_DROPOUT_MODEdrop2)selfr   r   r   r   Zdropout_modeZdropout_opt	__class__ N/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/blocks/mlp.pyr      s    
&


zMLPBlock.__init__c                 C  s2   |  | |}| |}| |}| |}|S )N)r   r   r   r   r    )r!   xr$   r$   r%   forward?   s
    


zMLPBlock.forward)r   r	   r   )__name__
__module____qualname____doc__r   r'   __classcell__r$   r$   r"   r%   r      s        #r   )
__future__r   torch.nnr   monai.networks.layersr   monai.utilsr   r   Moduler   r$   r$   r$   r%   <module>   s
   