code: def MHA(self, query, key, value):
      batch_size = key.size(0)
      # b*h, 1, d
      query = (
          query.reshape(batch_size, -1, self.head_num, self.head_dim)
          .transpose(1, 2)
          .reshape(batch_size * self.head_num, -1, self.head_dim)
      )
      # b*h, tn, d
      key = (
          key.reshape(batch_size, -1, self.head_num, self.head_dim)
          .transpose(1, 2)
          .reshape(batch_size * self.head_num, -1, self.head_dim)
      )
      # b*h, tn, pl*d
      value = (
          value.reshape(
              batch_size,
              self.task_num,
              -1,
              self.head_num,
              self.head_dim,
          )
          .permute(0, 3, 1, 2, 4)
          .reshape(batch_size * self.head_num, self.task_num, -1)
      )
  
      attn_weights = (
          torch.bmm(query, key.transpose(1, 2)) * self.scaling
      )  # b*h, 1, tn
      attn_probs = F.dropout(attn_weights, p=0.1)
      attn_probs = F.softmax(attn_probs, dim=-1)  # b*h, 1, tn
      attn_output = torch.bmm(attn_probs, value)  # b*h, 1, pl*d
      prompt_embedding = attn_output.squeeze(1).reshape(
          batch_size, self.head_num, -1, self.head_dim
      )  # b, h, pl, d
      prompt_embedding = prompt_embedding.transpose(1, 2).reshape(
          batch_size, -1, self.embedding_size
      )  # b, pl, e
      prompt_embedding = self.out_proj(prompt_embedding)  # b, pl, e
      return prompt_embedding
  
file_path: TextBox/textbox/model/ptg.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from .pretrained_models import Pretrained_Models
  from transformers import AutoTokenizer, AutoModel
label_desc: Use image node_id 3 for calling the PTG obj's underlying member method code with example usage: obj.MHA(query, key, value) and returns: prompt_embedding