docstr: 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
docstr_tok: ["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"]
code: 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:
      """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
      """
      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
  
imports: from __future__ import annotations
  import math
  from torch import nn
  from kornia.contrib.models.sam.architecture.common import MLPBlock
  from kornia.core import Module, Tensor
  from kornia.core.check import KORNIA_CHECK
label_desc: Use image node_id 1 to create a new TwoWayAttentionBlock object from inherited base classes: Module with example: obj = TwoWayAttentionBlock(embedding_dim, num_heads, mlp_dim, activation, attention_downsample_rate, skip_first_layer_pe)