class_docstr: Multi-Head Attention block, used in Transformer, Set Transformer and so on
  
  Parameters
  ----------
  d_model : int
      The feature size (input and output) in Multi-Head Attention layer.
  num_heads : int
      The number of heads.
  d_head : int
      The hidden size per head.
  d_ff : int
      The inner hidden size in the Feed-Forward Neural Network.
  dropouth : float
      The dropout rate of each sublayer.
  dropouta : float
      The dropout rate of attention heads.
  
  Notes
  -----
  This module was used in SetTransformer layer.
class_docstr_tok: ["Multi-Head","Attention","block",",","used","in","Transformer",",","Set","Transformer","and","so","on","Parameters","--","--","--","--","--","d_model",":","int","The","feature","size","(","input","and","output",")","in","Multi-Head","Attention","layer",".","num_heads",":","int","The","number","of","heads",".","d_head",":","int","The","hidden","size","per","head",".","d_ff",":","int","The","inner","hidden","size","in","the","Feed-Forward","Neural","Network",".","dropouth",":","float","The","dropout","rate","of","each","sublayer",".","dropouta",":","float","The","dropout","rate","of","attention","heads",".","Notes","--","--","-","This","module","was","used","in","SetTransformer","layer","."]
docstr: Compute multi-head self-attention.
  
  Parameters
  ----------
  x : torch.Tensor
      The input tensor used to compute queries.
  mem : torch.Tensor
      The memory tensor used to compute keys and values.
  lengths_x : list
      The array of node numbers, used to segment x.
  lengths_mem : list
      The array of node numbers, used to segment mem.
docstr_tok: ["Compute","multi-head","self-attention",".","Parameters","--","--","--","--","--","x",":","torch.Tensor","The","input","tensor","used","to","compute","queries",".","mem",":","torch.Tensor","The","memory","tensor","used","to","compute","keys","and","values",".","lengths_x",":","list","The","array","of","node","numbers",",","used","to","segment","x.","lengths_mem",":","list","The","array","of","node","numbers",",","used","to","segment","mem","."]
code: def forward(self, x, mem, lengths_x, lengths_mem):
      """
      Compute multi-head self-attention.
  
      Parameters
      ----------
      x : torch.Tensor
          The input tensor used to compute queries.
      mem : torch.Tensor
          The memory tensor used to compute keys and values.
      lengths_x : list
          The array of node numbers, used to segment x.
      lengths_mem : list
          The array of node numbers, used to segment mem.
      """
      batch_size = len(lengths_x)
      max_len_x = max(lengths_x)
      max_len_mem = max(lengths_mem)
      device = x.device
      lengths_x = th.as_tensor(lengths_x, dtype=th.int64, device=device)
      lengths_mem = th.as_tensor(
          lengths_mem, dtype=th.int64, device=device
      )
  
      queries = self.proj_q(x).view(-1, self.num_heads, self.d_head)
      keys = self.proj_k(mem).view(-1, self.num_heads, self.d_head)
      values = self.proj_v(mem).view(-1, self.num_heads, self.d_head)
  
      # padding to (B, max_len_x/mem, num_heads, d_head)
      queries = F.pad_packed_tensor(queries, lengths_x, 0)
      keys = F.pad_packed_tensor(keys, lengths_mem, 0)
      values = F.pad_packed_tensor(values, lengths_mem, 0)
  
      # attention score with shape (B, num_heads, max_len_x, max_len_mem)
      e = th.einsum("bxhd,byhd->bhxy", queries, keys)
      # normalize
      e = e / np.sqrt(self.d_head)
  
      # generate mask
      mask = _gen_mask(lengths_x, lengths_mem, max_len_x, max_len_mem)
      e = e.masked_fill(mask == 0, -float("inf"))
  
      # apply softmax
      alpha = th.softmax(e, dim=-1)
      # the following line addresses the NaN issue, see
      # https://github.com/dmlc/dgl/issues/2657
      alpha = alpha.masked_fill(mask == 0, 0.0)
  
      # sum of value weighted by alpha
      out = th.einsum("bhxy,byhd->bxhd", alpha, values)
      # project to output
      out = self.proj_o(
          out.contiguous().view(
              batch_size, max_len_x, self.num_heads * self.d_head
          )
      )
      # pack tensor
      out = F.pack_padded_tensor(out, lengths_x)
  
      # intra norm
      x = self.norm_in(x + out)
  
      # inter norm
      x = self.norm_inter(x + self.ffn(x))
  
      return x
  
imports: import numpy
  import torch
  import torch.nn
  from ...backend import pytorch
  from ...base import dgl_warning
  from ...readout import broadcast_nodes, max_nodes, mean_nodes, softmax_nodes, sum_nodes, topk_nodes
label_desc: Use image node_id 3 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(x, mem, lengths_x, lengths_mem) and returns: x