class_docstr: Compute 'Scaled Dot Product Attention
  
  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
      dropout (nn.Dropout): the dropout layer
  
  Example::
  
      self.attention = Attention()
      x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout)
class_docstr_tok: ["Compute","'Scaled","Dot","Product","Attention","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","dropout","(","nn.Dropout",")",":","the","dropout","layer","Example",":",":","self.attention","=","Attention","(",")","x",",","attn","=","self.attention","(","query",",","key",",","value",",","mask=mask",",","dropout=self.dropout",")"]
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
      dropout (nn.Dropout): the dropout layer
  
  Returns:
      Tuple[torch.Tensor, 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","dropout","(","nn.Dropout",")",":","the","dropout","layer","Returns",":","Tuple","[","torch.Tensor",",","torch.Tensor","]",":"]
code: def forward(
      self,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
      dropout: Optional[nn.Dropout] = None,
  ) -> Tuple[torch.Tensor, 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
          dropout (nn.Dropout): the dropout layer
  
      Returns:
          Tuple[torch.Tensor, torch.Tensor]:
      """
      scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(
          query.size(-1)
      )
  
      if mask is not None:
          scores = scores.masked_fill(mask == 0, -1e9)
  
      p_attn = nn.functional.softmax(scores, dim=-1)
  
      if dropout is not None:
          p_attn = dropout(p_attn)
  
      return torch.matmul(p_attn, value), p_attn
  
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 1 for calling the Attention obj's underlying member method code with example usage: obj.forward(query, key, value, mask, dropout) and returns: torch,  p_attn