code: def __init__(self, args, mode="normal"):
      super(MSA, self).__init__()
      if mode == "copy":
          nhead, head_dim = 1, args.nhid
          qninp, kninp = args.dec_ninp, args.nhid
      if mode == "normal":
          nhead, head_dim = args.nhead, args.head_dim
          qninp, kninp = args.nhid, args.nhid
      self.attn_drop = nn.Dropout(0.1)
      self.WQ = nn.Linear(
          qninp,
          nhead * head_dim,
          bias=True if mode == "copy" else False,
      )
      if mode != "copy":
          self.WK = nn.Linear(kninp, nhead * head_dim, bias=False)
          self.WV = nn.Linear(kninp, nhead * head_dim, bias=False)
      self.args, self.nhead, self.head_dim, self.mode = (
          args,
          nhead,
          head_dim,
          mode,
      )
  
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 1 to create a new MSA object from inherited base classes: nn with example: obj = MSA(args, mode)