class_docstr: Multi-Head Attention
class_docstr_tok: ["Multi-Head","Attention"]
code: def forward(
      self,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      mask=None,
  ) -> torch.Tensor:
      batch_size = query.size(0)
  
      # linear projections of Q, K, V
      query, key, value = [
          linear(x)
          .view(batch_size, -1, self.num_heads, self.d_k)
          .transpose(1, 2)
          for linear, x in zip(self.linear_layers, (query, key, value))
      ]
  
      # apply attention on all the projected vectors in batch
      x, attn = scaled_dot_product_attention(
          query, key, value, mask=mask
      )
  
      # Concat attention heads
      x = (
          x.transpose(1, 2)
          .contiguous()
          .view(batch_size, -1, self.num_heads * self.d_k)
      )
  
      return self.output_linear(x)
  
imports: import math
  from typing import Any, Callable, Optional, Tuple
  import torch
  from torch import nn
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(query, key, value, mask) and returns: self