docstr: This model inherits from bmt.DistributedModule.
  Args:
      hidden_q (:obj:`torch.Tensor` of shape ``(batch, len_q, dim_model)``): Indices of input sequence tokens. It will be embedded by model's internal embedding lookup matrix.
      hidden_kv (:obj:`torch.Tensor` of shape ``(batch, len_k, dim_model)``): Length of input sequence before padding.
      attention_mask (:obj:`torch.Tensor` of shape ``(batch, len_q, len_k)``): Used to avoid performing attention on padding token indices.
      position_bias(:obj:`torch.Tensor` of shape ``(num_heads, len_q, len_k)`` or ``(1, num_heads, len_k, len_q)``): Provide positional information about tensor `key_value` and `query`.
  Return:
      out (:obj:`torch.Tensor` of shape ``(batch, len_q, dim_model)``): The attention output.
docstr_tok: ["This","model","inherits","from","bmt.DistributedModule",".","Args",":","hidden_q","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","len_q",",","dim_model",")","``",")",":","Indices","of","input","sequence","tokens",".","It","will","be","embedded","by","model","'s","internal","embedding","lookup","matrix",".","hidden_kv","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","len_k",",","dim_model",")","``",")",":","Length","of","input","sequence","before","padding",".","attention_mask","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","len_q",",","len_k",")","``",")",":","Used","to","avoid","performing","attention","on","padding","token","indices",".","position_bias","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","num_heads",",","len_q",",","len_k",")","``","or","``","(","1",",","num_heads",",","len_k",",","len_q",")","``",")",":","Provide","positional","information","about","tensor","`","key_value","`","and","`","query","`",".","Return",":","out","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","len_q",",","dim_model",")","``",")",":","The","attention","output","."]
code: def forward(
      self,
      hidden_q: torch.Tensor,
      hidden_kv: torch.Tensor,
      attention_mask: torch.BoolTensor,
      position_bias: torch.Tensor,
      use_cache: bool = False,
      past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
  ):
      """This model inherits from bmt.DistributedModule.
      Args:
          hidden_q (:obj:`torch.Tensor` of shape ``(batch, len_q, dim_model)``): Indices of input sequence tokens. It will be embedded by model's internal embedding lookup matrix.
          hidden_kv (:obj:`torch.Tensor` of shape ``(batch, len_k, dim_model)``): Length of input sequence before padding.
          attention_mask (:obj:`torch.Tensor` of shape ``(batch, len_q, len_k)``): Used to avoid performing attention on padding token indices.
          position_bias(:obj:`torch.Tensor` of shape ``(num_heads, len_q, len_k)`` or ``(1, num_heads, len_k, len_q)``): Provide positional information about tensor `key_value` and `query`.
      Return:
          out (:obj:`torch.Tensor` of shape ``(batch, len_q, dim_model)``): The attention output.
      """  # noqa: E501
  
      batch_size = hidden_q.size(0)
      len_q = hidden_q.size(1)
      len_k = hidden_kv.size(1)
  
      h_q = self.project_q(hidden_q) / math.sqrt(
          math.sqrt(self.dim_head)
      )
      h_k = self.project_k(hidden_kv) / math.sqrt(
          math.sqrt(self.dim_head)
      )
      h_v = self.project_v(hidden_kv)
  
      h_q = h_q.view(
          batch_size, len_q, self.num_heads, self.dim_head
      ).permute(0, 2, 1, 3)
      h_k = h_k.view(
          batch_size, len_k, self.num_heads, self.dim_head
      ).permute(0, 2, 1, 3)
      h_v = h_v.view(
          batch_size, len_k, self.num_heads, self.dim_head
      ).permute(0, 2, 1, 3)
  
      if past_kv is not None:
          h_k = torch.cat([past_kv[0], h_k], dim=-2)
          h_v = torch.cat([past_kv[1], h_v], dim=-2)
          len_k = h_k.size(-2)
  
      # (b, n_h, len_q, d_h) @ (b, n_h, d_h, len_k) -> (b, n_h, len_q, len_k)
      score = torch.matmul(
          h_q, h_k.transpose(-1, -2)
      )  # / math.sqrt(self.dim_head) moved to line 75~76
      score = score + position_bias
      score = torch.masked_fill(
          score,
          attention_mask.view(batch_size, 1, len_q, len_k) == False,
          torch.scalar_tensor(
              float("-inf"), device=score.device, dtype=score.dtype
          ),
      )
  
      score = self.softmax(score)
  
      score = torch.masked_fill(
          score,
          attention_mask.view(batch_size, 1, len_q, len_k) == False,
          torch.scalar_tensor(
              0, device=score.device, dtype=score.dtype
          ),
      )
  
      if self.dropout is not None:
          score = self.dropout(score)
  
      # (b, n_h, len_q, len_k) @ (b, n_h, len_k, d_h) -> (b, n_h, len_q, d_h)
      score = torch.matmul(score, h_v)
  
      score = score.view(
          batch_size, self.num_heads, len_q, self.dim_head
      ).permute(0, 2, 1, 3)
      score = score.contiguous().view(
          batch_size, len_q, self.num_heads * self.dim_head
      )
  
      score = self.attention_out(score)
      if use_cache:
          return score, (h_k, h_v)
      else:
          return score
  
file_path: DeepKE/example/llm/CPM-Bee/src/cpm_live/layers/attention.py
imports: from typing import Optional, Tuple
  import torch
  import bmtrain
  import math
  from .linear import Linear
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj.forward(hidden_q, hidden_kv, attention_mask, position_bias, use_cache, past_kv) and returns: score,  h_k,  h_v, score