class_docstr: Multi-headed attention block.
  
  Args:
      num_heads (int): number of attention heads
      dim_model (int): input/output dimensionality
      dropout (float): the dropout probability
      mask (torch.Tensor): the mask tensor
      device: (Optional[torch.device]).
  
  Example::
  
      self.attention = MultiHeadedAttention(
          num_heads=attn_heads, dim_model=hidden, dropout=dropout, device=device
      )
      self.attention.forward(query, key, value, mask=mask)
class_docstr_tok: ["Multi-headed","attention","block",".","Args",":","num_heads","(","int",")",":","number","of","attention","heads","dim_model","(","int",")",":","input\/output","dimensionality","dropout","(","float",")",":","the","dropout","probability","mask","(","torch.Tensor",")",":","the","mask","tensor","device",":","(","Optional","[","torch.device","]",")",".","Example",":",":","self.attention","=","MultiHeadedAttention","(","num_heads=attn_heads",",","dim_model=hidden",",","dropout=dropout",",","device=device",")","self.attention.forward","(","query",",","key",",","value",",","mask=mask",")"]
docstr: forward function
  
  Args:
      query (torch.Tensor): the query tensor
      key (torch.Tensor): the key tensor
      value (torch.Tensor): the value tensor
      mask (torch.Tensor): the mask tensor
  
  Returns:
      torch.Tensor.
docstr_tok: ["forward","function","Args",":","query","(","torch.Tensor",")",":","the","query","tensor","key","(","torch.Tensor",")",":","the","key","tensor","value","(","torch.Tensor",")",":","the","value","tensor","mask","(","torch.Tensor",")",":","the","mask","tensor","Returns",":","torch.Tensor","."]
code: def forward(
      self,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
  ) -> torch.Tensor:
      """
      forward function
  
      Args:
          query (torch.Tensor): the query tensor
          key (torch.Tensor): the key tensor
          value (torch.Tensor): the value tensor
          mask (torch.Tensor): the mask tensor
  
      Returns:
          torch.Tensor.
      """
      batch_size = query.size(0)
  
      # 1) Do all the linear projections in batch from dim_model => num_heads x d_k
      # self.linear[0, 1, 2] is query weight matrix, key weight matrix, and
      # value weight matrix, respectively.
      # l(x) represents the transformed query matrix, key matrix and value matrix
      # l(x) has shape (batch_size, seq_len, dim_model). You can think l(x) as
      # the matrices from a one-head attention; or you can think
      # l(x).view(...).transpose(...) as the matrices of num_heads attentions,
      # each attention has d_k dimension.
      query, key, value = [
          linearLayer(x)
          .view(batch_size, -1, self.num_heads, self.d_k)
          .transpose(1, 2)
          for linearLayer, x in zip(
              self.linear_layers, (query, key, value)
          )
      ]
  
      # 2) Apply attention on all the projected vectors in batch.
      # x shape: batch_size, num_heads, seq_len, d_k
      x, attn = self.attention(
          query, key, value, mask=mask, dropout=self.dropout
      )
  
      # 3) "Concat" using a view and apply a final linear.
      # each attention's output is d_k dimension. Concat num_heads attention's outputs
      # x shape: batch_size, seq_len, dim_model
      x = (
          x.transpose(1, 2)
          .contiguous()
          .view(batch_size, -1, self.num_heads * self.d_k)
      )
      return self.output_linear(x)
  
imports: import copy
  import math
  from typing import Callable, Optional, Tuple
  import torch
  import torch.nn
  from torchrec.modules.embedding_configs import EmbeddingConfig
  from torchrec.modules.embedding_modules import EmbeddingCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 2 for calling the MultiHeadedAttention obj's underlying member method code with example usage: obj.forward(query, key, value, mask) and returns: self