class_docstr: Fixed (1 line):
      Add `n_head //= dist.get_world_size(self.mp_group)`
class_docstr_tok: ["Fixed","(","1","line",")",":","Add","`","n_head","\/\/=","dist.get_world_size","(","self.mp_group",")","`"]
code: def forward(
      self, query, key, value, attention_inputs, output_attentions=False
  ):
      # query has shape batch_size x seq_len x d_model
      # key and value have shapes batch_size x context_len x d_model
      (
          position_embeds,
          token_type_mat,
          attention_mask,
          cls_mask,
      ) = attention_inputs
  
      batch_size, seq_len, _ = query.shape
      context_len = key.shape[1]
      n_head, d_head = self.config.n_head, self.config.d_head
      n_head //= dist.get_world_size(self.mp_group)
  
      # Shape batch_size x seq_len x n_head x d_head
      q_head = self.q_head(query).view(
          batch_size, seq_len, n_head, d_head
      )
      # Shapes batch_size x context_len x n_head x d_head
      k_head = self.k_head(key).view(
          batch_size, context_len, n_head, d_head
      )
      v_head = self.v_head(value).view(
          batch_size, context_len, n_head, d_head
      )
  
      q_head = q_head * self.scale
      # Shape n_head x d_head
      r_w_bias = self.r_w_bias * self.scale
      # Shapes batch_size x n_head x seq_len x context_len
      content_score = torch.einsum(
          "bind,bjnd->bnij", q_head + r_w_bias, k_head
      )
      positional_attn = self.relative_positional_attention(
          position_embeds, q_head, context_len, cls_mask
      )
      token_type_attn = self.relative_token_type_attention(
          token_type_mat, q_head, cls_mask
      )
  
      # merge attention scores
      attn_score = content_score + positional_attn + token_type_attn
  
      # precision safe in case of mixed precision training
      dtype = attn_score.dtype
      attn_score = attn_score.float()
      # perform masking
      if attention_mask is not None:
          attn_score = attn_score - INF * (
              1 - attention_mask[:, None, None].float()
          )
      # attention probability
      attn_prob = torch.softmax(attn_score, dim=-1, dtype=dtype)
      attn_prob = self.attention_dropout(attn_prob)
  
      # attention output, shape batch_size x seq_len x n_head x d_head
      attn_vec = torch.einsum("bnij,bjnd->bind", attn_prob, v_head)
  
      # Shape shape batch_size x seq_len x d_model
      attn_out = self.post_proj(
          attn_vec.reshape(batch_size, seq_len, n_head * d_head)
      )
      attn_out = self.hidden_dropout(attn_out)
  
      output = self.layer_norm(query + attn_out)
      return (output, attn_prob) if output_attentions else (output,)
  
imports: import torch
  import torch.distributed
  from transformers.models.funnel.modeling_funnel import INF, FunnelRelMultiheadAttention
label_desc: Use image node_id 1 for calling the FunnelRelMultiheadAttention_ obj's underlying member method code with example usage: obj.forward(query, key, value, attention_inputs, output_attentions) and returns: unknown