def __init__(
        self,
        embedding_dim: int,
        num_heads: int,
        mlp_dim: int = 2048,
        activation: type[Module] = nn.ReLU,
        attention_downsample_rate: int = 2,
        skip_first_layer_pe: bool = False,
    ) -> None:
        """Construct a transformer block with four layers.

        (1) self-attention of sparse inputs, (2) cross attention of
        sparse inputs to dense inputs, (3) mlp block on sparse inputs, and (4) cross attention of dense inputs to sparse
        inputs.

        Args:
            embedding_dim: the channel dimension of the embeddings
            num_heads: the number of heads in the attention layers
            mlp_dim: the hidden dimension of the mlp block
            activation: the activation of the mlp block
            skip_first_layer_pe: skip the PE on the first layer
            attention_downsample_rate: downsampling rate from embedding dimension

        """
        super().__init__()
        self.self_attn = Attention(embedding_dim, num_heads)
        self.norm1 = nn.LayerNorm(embedding_dim)

        self.cross_attn_token_to_image = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
        self.norm2 = nn.LayerNorm(embedding_dim)

        self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
        self.norm3 = nn.LayerNorm(embedding_dim)

        self.norm4 = nn.LayerNorm(embedding_dim)
        self.cross_attn_image_to_token = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)

        self.skip_first_layer_pe = skip_first_layer_pe
Construct a transformer block with four layers.

        (1) self-attention of sparse inputs, (2) cross attention of
        sparse inputs to dense inputs, (3) mlp block on sparse inputs, and (4) cross attention of dense inputs to sparse
        inputs.

        Args:
            embedding_dim: the channel dimension of the embeddings
            num_heads: the number of heads in the attention layers
            mlp_dim: the hidden dimension of the mlp block
            activation: the activation of the mlp block
            skip_first_layer_pe: skip the PE on the first layer
            attention_downsample_rate: downsampling rate from embedding dimension

        
kornia/contrib/models/sam/architecture/transformer.py