def __init__(
        self,
        depth: int,
        embedding_dim: int,
        num_heads: int,
        mlp_dim: int,
        activation: type[Module] = nn.ReLU,
        attention_downsample_rate: int = 2,
    ) -> None:
        """Construct a transformer decoder that attends to an input image using queries whose positional embedding is
        supplied.

        Args:
            depth: number of layers in the transformer
            embedding_dim: the channel dimension for the input embeddings
            num_heads: the number of heads for multihead attention. Must divide embedding_dim
            mlp_dim: the channel dimension internal to the MLP block
            activation: the activation to use in the MLP block
            attention_downsample_rate: downsampling rate from embedding dimension

        """  # noqa: D205
        super().__init__()
        self.depth = depth
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        self.mlp_dim = mlp_dim
        self.layers = nn.ModuleList()

        for i in range(depth):
            self.layers.append(
                TwoWayAttentionBlock(
                    embedding_dim=embedding_dim,
                    num_heads=num_heads,
                    mlp_dim=mlp_dim,
                    activation=activation,
                    attention_downsample_rate=attention_downsample_rate,
                    skip_first_layer_pe=(i == 0),
                )
            )

        self.final_attn_token_to_image = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
        self.norm_final_attn = nn.LayerNorm(embedding_dim)
Construct a transformer decoder that attends to an input image using queries whose positional embedding is
        supplied.

        Args:
            depth: number of layers in the transformer
            embedding_dim: the channel dimension for the input embeddings
            num_heads: the number of heads for multihead attention. Must divide embedding_dim
            mlp_dim: the channel dimension internal to the MLP block
            activation: the activation to use in the MLP block
            attention_downsample_rate: downsampling rate from embedding dimension

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