code: def forward(self, inp1, inp2, mask=None):
      B, L2, H = inp2.shape
      NH, HD = self.nhead, self.head_dim
      if self.mode == "copy":
          q, k, v = self.WQ(inp1), inp2, inp2
      else:
          q, k, v = self.WQ(inp1), self.WK(inp2), self.WV(inp2)
      L1 = 1 if inp1.ndim == 2 else inp1.shape[1]
      if self.mode != "copy":
          q = q / math.sqrt(H)
      q = q.view(B, L1, NH, HD).permute(0, 2, 1, 3)
      k = k.view(B, L2, NH, HD).permute(0, 2, 3, 1)
      v = v.view(B, L2, NH, HD).permute(0, 2, 1, 3)
      pre_attn = torch.matmul(q, k)
      if mask is not None:
          pre_attn = pre_attn.masked_fill(mask[:, None, None, :], -1e8)
      if self.mode == "copy":
          return pre_attn.squeeze(1)
      else:
          alpha = self.attn_drop(torch.softmax(pre_attn, -1))
          attn = (
              torch.matmul(alpha, v)
              .permute(0, 2, 1, 3)
              .contiguous()
              .view(B, L1, NH * HD)
          )
          ret = attn
          if inp1.ndim == 2:
              return ret.squeeze(1)
          else:
              return ret
  
imports: import math
  import torch
  import torch.nn.functional
  from torch import nn
  from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
  from utlis import *
  import dgl.function
  from dgl.nn.functional import edge_softmax
label_desc: Use image node_id 2 for calling the MSA obj's underlying member method code with example usage: obj.forward(inp1, inp2, mask) and returns: pre_attn, ret, ret