docstr: Args:
      image_embedding: image to attend to. Should be shape B x embedding_dim x h x w for any h and w.
      image_pe: the positional encoding to add to the image. Must have the same shape as image_embedding.
      point_embedding: the embedding to add to the query points. Must have shape B x N_points x embedding_dim
      for any N_points.
  
  Returns:
      - the processed point_embedding
      - the processed image_embedding
docstr_tok: ["Args",":","image_embedding",":","image","to","attend","to",".","Should","be","shape","B","x","embedding_dim","x","h","x","w","for","any","h","and","w.","image_pe",":","the","positional","encoding","to","add","to","the","image",".","Must","have","the","same","shape","as","image_embedding",".","point_embedding",":","the","embedding","to","add","to","the","query","points",".","Must","have","shape","B","x","N_points","x","embedding_dim","for","any","N_points",".","Returns",":","-","the","processed","point_embedding","-","the","processed","image_embedding"]
code: def forward(
      self,
      image_embedding: Tensor,
      image_pe: Tensor,
      point_embedding: Tensor,
  ) -> tuple[Tensor, Tensor]:
      """
      Args:
          image_embedding: image to attend to. Should be shape B x embedding_dim x h x w for any h and w.
          image_pe: the positional encoding to add to the image. Must have the same shape as image_embedding.
          point_embedding: the embedding to add to the query points. Must have shape B x N_points x embedding_dim
          for any N_points.
  
      Returns:
          - the processed point_embedding
          - the processed image_embedding
      """
      # BxCxHxW -> BxHWxC == B x N_image_tokens x C
      bs, c, h, w = image_embedding.shape
      image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
      image_pe = image_pe.flatten(2).permute(0, 2, 1)
  
      # Prepare queries
      queries = point_embedding
      keys = image_embedding
  
      # Apply transformer blocks and final layernorm
      for layer in self.layers:
          queries, keys = layer(
              queries=queries,
              keys=keys,
              query_pe=point_embedding,
              key_pe=image_pe,
          )
  
      # Apply the final attenion layer from the points to the image
      q = queries + point_embedding
      k = keys + image_pe
      attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
      queries = queries + attn_out
      queries = self.norm_final_attn(queries)
  
      return queries, keys
  
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 2 for calling the TwoWayTransformer obj's underlying member method code with example usage: obj.forward(image_embedding, image_pe, point_embedding) and returns: queries,  keys