o ( iu@sVddlmZddlmZddlZddlmZddlmZm Z m Z Gdddej Z dS)) annotations)OptionalN)CrossAttentionBlockMLPBlockSABlockcsDeZdZdZ         d!d"fdd Z d#d$dd ZZS)%TransformerBlockz A transformer block, based on: "Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale " FNT hidden_sizeintmlp_dim num_heads dropout_ratefloatqkv_biasbool save_attncausalsequence_length int | Nonewith_cross_attentionuse_flash_attention include_fcuse_combined_linearreturnNonec std|krdkstdtd||dkr!tdt||||_t||_t|||||||| | | d |_ t||_ | |_ t||_ t ||||d| d|_dS) a Args: hidden_size (int): dimension of hidden layer. mlp_dim (int): dimension of feedforward layer. num_heads (int): number of attention heads. dropout_rate (float, optional): fraction of the input units to drop. Defaults to 0.0. qkv_bias(bool, optional): apply bias term for the qkv linear layer. Defaults to False. save_attn (bool, optional): to make accessible the attention matrix. Defaults to False. use_flash_attention: if True, use Pytorch's inbuilt flash attention for a memory efficient attention mechanism (see https://pytorch.org/docs/2.2/generated/torch.nn.functional.scaled_dot_product_attention.html). include_fc: whether to include the final linear layer. Default to True. use_combined_linear: whether to use a single linear layer for qkv projection, default to True. rz'dropout_rate should be between 0 and 1.z-hidden_size should be divisible by num_heads.)rrrrrrrF)r r r rrrN)super__init__ ValueErrorrmlpnn LayerNormnorm1rattnnorm2rnorm_cross_attnr cross_attn) selfr r r r rrrrrrrr __class__h/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/transformerblock.pyrs@      zTransformerBlock.__init__x torch.TensorcontextOptional[torch.Tensor] attn_maskcCsN||j|||d}|jr||j|||d}||||}|S)N)r0)r.)r#r"rr&r%rr$)r'r,r.r0r*r*r+forward]s zTransformerBlock.forward) rFFFNFFTT)r r r r r r r rrrrrrrrrrrrrrrrrrr)NN)r,r-r.r/r0r/rr-)__name__ __module__ __qualname____doc__rr1 __classcell__r*r*r(r+rs Br) __future__rtypingrtorchtorch.nnr Zmonai.networks.blocksrrrModulerr*r*r*r+s