o ) ir @sNddlmZddlmZddlZddlmZddlmZGdddej Z dS)) annotations)OptionalN)SABlockcs<eZdZdZ       ddfdd ZdddZZS) SpatialAttentionBlockaZPerform spatial self-attention on the input tensor. The input tensor is reshaped to B x (x_dim * y_dim [ * z_dim]) x C, where C is the number of channels, and then self-attention is performed on the reshaped tensor. The output tensor is reshaped back to the original shape. Args: spatial_dims: number of spatial dimensions, could be 1, 2, or 3. num_channels: number of input channels. Must be divisible by num_head_channels. num_head_channels: number of channels per head. norm_num_groups: Number of groups for the group norm layer. norm_eps: Epsilon for the normalization. attention_dtype: cast attention operations to this dtype. 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 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). N ư>TF spatial_dimsint num_channelsnum_head_channels int | Nonenorm_num_groupsnorm_epsfloatattention_dtypeOptional[torch.dtype] include_fcbooluse_combined_linearuse_flash_attentionreturnNonec spt||_tj|||dd|_|dur ||dkr td|dur(||nd} t|| d|||| d|_dS)NT) num_groupsr epsaffinerz3num_channels must be divisible by num_head_channels) hidden_size num_headsqkv_biasrrrr) super__init__rnn GroupNormnorm ValueErrorrattn) selfrr r r rrrrrr __class__h/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/spatialattention.pyr *s  zSpatialAttentionBlock.__init__x torch.TensorcCs`|}|j}||}|jg|dddRdd}||}|dd|}||}|S)Nr)shaper#reshape transposer%)r&r+residualr/r)r)r*forwardHs $ zSpatialAttentionBlock.forward)NrrNTFF)rr r r r r r r rrrrrrrrrrrr)r+r,)__name__ __module__ __qualname____doc__r r3 __classcell__r)r)r'r*rsr) __future__rtypingrtorchtorch.nnr!Zmonai.networks.blocksrModulerr)r)r)r*s