o ) i @snddlmZddlmZmZddlZddlmZddlm Z ddl m Z e ddd\Z Z Gd d d ejZdS) ) annotations)OptionalTupleN)get_rel_pos_embedding_layer)optional_importzeinops.layers.torch Rearrange)namecsHeZdZdZ            d$d%fdd Zd&d'd"d#ZZS)(CrossAttentionBlocka A cross-attention block, based on: "Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale " One can setup relative positional embedding as described in NF hidden_sizeint num_heads dropout_ratefloathidden_input_size int | Nonecontext_input_sizedim_headqkv_biasbool save_attncausalsequence_lengthrel_pos_embedding Optional[str] input_sizeOptional[Tuple]attention_dtypeOptional[torch.dtype]use_flash_attentionreturnNonec std|krdkstdtd|r!||}||_n||dkr+td|}|||_| r<| dur