class_docstr: Implementation of Multi-Head Attention
  
  The MH Block in paper: Accurate Learning of Graph Representations
  with Graph Multiset Pooling
  
  Args:
      dim_Q: The dimension of query tensor.
      dim_K: The dimension of key tensor.
      dim_V: The dimension of value tensor.
      num_heads: The number of attention heads.
      conv: The Graph Convolution Block for GMH. Default: None
      layer_norm: Whether to use LayerNorm.
docstr: Args:           
      Q: Query tensor with shape [-1, dim_Q]
      K: Key tensor with shape [-1, dim_K]
      graph: "Pgl.Graph" instance. Default: None
      mask: A tensor indicating the position of dummy nodes
  Return:
      A tensor with shape (-1, Q.shape[0], dim_V)
docstr_tok: ["Args",":","Q",":","Query","tensor","with","shape","[","-1",",","dim_Q","]","K",":","Key","tensor","with","shape","[","-1",",","dim_K","]","graph",":","``","Pgl.Graph","''","instance",".","Default",":","None","mask",":","A","tensor","indicating","the","position","of","dummy","nodes","Return",":","A","tensor","with","shape","(","-1",",","Q.shape","[","0","]",",","dim_V",")"]
code: def forward(self, Q, K, graph=None, mask=None):
      """
      Args:
          Q: Query tensor with shape [-1, dim_Q]
          K: Key tensor with shape [-1, dim_K]
          graph: "Pgl.Graph" instance. Default: None
          mask: A tensor indicating the position of dummy nodes
      Return:
          A tensor with shape (-1, Q.shape[0], dim_V)
      """
      Q = self.proj_q(Q)
      if graph is not None:
          graph, x = graph
          K, V = self.layer_k(graph, x), self.layer_v(graph, x)
          K, _ = to_dense_batch(K, graph)
          V, _ = to_dense_batch(V, graph)
      else:
          K, V = self.layer_k(K), self.layer_v(K)
  
      # h_k = self.dim_V // self.num_heads
      Q_ = paddle.concat(Q.split(self.num_heads, 2), axis=0)
      K_ = paddle.concat(K.split(self.num_heads, 2), axis=0)
      V_ = paddle.concat(V.split(self.num_heads, 2), axis=0)
      if mask is not None:
          # mask [batch_size, seq_len]
          mask = paddle.concat(
              [mask for _ in range(self.num_heads)], axis=0
          )
          mask = mask.reshape((K_.shape[0], K_.shape[1]))
          mask = mask.unsqueeze(1)
          att_score = Q_.bmm(K_.transpose([0, 2, 1]))
          att_score = att_score / math.sqrt(self.dim_V)
          A = F.softmax(mask + att_score, 1)
      else:
          A = F.softmax(
              Q_.bmm(K_.transpose([0, 2, 1])) / math.sqrt(self.dim_V), 1
          )
      output = paddle.concat(
          (Q_ + A.bmm(V_)).split(self.num_heads, 0), 2
      )
  
      if self.layer_norm:
          output = self.ln0(output)
      output = output + F.relu(self.proj_o(output))
      if self.layer_norm:
          output = self.ln0(output)
  
      return output
  
imports: import paddle
  import paddle.nn
  import math
  import paddle.nn.functional
  import pgl
  from pgl.utils.transform import to_dense_batch
label_desc: Use image node_id 2 for calling the MAB obj's underlying member method code with example usage: obj.forward(Q, K, graph, mask) and returns: output