U Ph@snddlmZddlZddlmZddlmZmZddl Z ddl m Z dgZ ddZ dd d d d d dddZdS)) annotationsN)repeat)ListUnionbuild_sincos_position_embeddingcsfdd}|S)Ncs.t|tjjr t|ts t|Stt|S)N) isinstance collectionsabcIterablestrtupler)xnZ/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/blocks/pos_embed_utils.pyparsesz_ntuple..parser)rrrrr_ntuples r@zUnion[int, List[int]]intfloatztorch.nn.Parameter) grid_size embed_dim spatial_dims temperaturereturnc CsP|dkrtd}||}|\}}tj|tjd}tj|tjd} tj|| dd\}} |ddkrhtd|d} tj| tjd| } d|| } td || g} td | | g} tjt | t | t | t | gd d d d d d d f}n@|d kr4td }||}|\}}}tj|tjd}tj|tjd} tj|tjd}tj|| |dd\}} }|ddkr~td|d} tj| tjd| } d|| } td || g} td | | g} td || g}tjt | t | t | t | t |t |gd d d d d d d f}nt dt |}d|_|S)a Builds a sin-cos position embedding based on the given grid size, embed dimension, spatial dimensions, and temperature. Reference: https://github.com/cvlab-stonybrook/SelfMedMAE/blob/68d191dfcc1c7d0145db93a6a570362de29e3b30/lib/models/mae3d.py Args: grid_size (List[int]): The size of the grid in each spatial dimension. embed_dim (int): The dimension of the embedding. spatial_dims (int): The number of spatial dimensions (2 for 2D, 3 for 3D). temperature (float): The temperature for the sin-cos position embedding. Returns: pos_embed (nn.Parameter): The sin-cos position embedding as a fixed parameter. )dtypeij)indexingrzHEmbed dimension must be divisible by 4 for 2D sin-cos position embeddingg?zm,d->md)dimNrzHEmbed dimension must be divisible by 6 for 3D sin-cos position embeddingz6Spatial Dimension Size {spatial_dims} Not Implemented!F)rtorcharangefloat32meshgridAssertionErroreinsumflattencatsincosNotImplementedErrornn Parameter requires_grad)rrrrZ to_2tupleZ grid_size_thwZgrid_hZgrid_wZpos_dimomegaout_hout_wZpos_embZ to_3tupledZgrid_dout_d pos_embedrrrr#s^  D      )rr) __future__rcollections.abcr itertoolsrtypingrrr%torch.nnr0__all__rrrrrr s