def forward(self, image_embedding: Tensor, image_pe: Tensor, point_embedding: Tensor) -> tuple[Tensor, Tensor]:
        """Run forward.

        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
Run forward.

        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

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