class_docstr: The whole transformer block. A sequence of operation. Consists of self-attention block[, cross-attention block] and feed-forward block.
  
  Args:
      dim_model (int): main dimension of modules in transformer blocks.
      dim_ff (int): dim_ff used in :py:class:`model_center.layer.FeedForward`.
      num_heads (int): num_heads used in :py:class:`model_center.layer.Attention`.
      dim_head (int): dim_head used in :py:class:`model_center.layer.Attention`.
      dtype (optional): Defaults to torch.half.
      eps (float, optional): eps used in :py:class:`model_center.layer.LayerNorm`. Defaults to 1e-5.
      dropout_p (float, optional): Defaults to 0.
class_docstr_tok: ["The","whole","transformer","block",".","A","sequence","of","operation",".","Consists","of","self-attention","block","[",",","cross-attention","block","]","and","feed-forward","block",".","Args",":","dim_model","(","int",")",":","main","dimension","of","modules","in","transformer","blocks",".","dim_ff","(","int",")",":","dim_ff","used","in",":","py",":","class",":","`","model_center.layer.FeedForward","`",".","num_heads","(","int",")",":","num_heads","used","in",":","py",":","class",":","`","model_center.layer.Attention","`",".","dim_head","(","int",")",":","dim_head","used","in",":","py",":","class",":","`","model_center.layer.Attention","`",".","dtype","(","optional",")",":","Defaults","to","torch.half",".","eps","(","float",",","optional",")",":","eps","used","in",":","py",":","class",":","`","model_center.layer.LayerNorm","`",".","Defaults","to","1e-5",".","dropout_p","(","float",",","optional",")",":","Defaults","to","0","."]
docstr: Args:
      self_hidden_states (:obj:`torch.Tensor` of shape ``(batch, seq_self, dim_model)``): Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
      self_attention_mask (:obj:`torch.Tensor` of shape ``(batch, seq_self, seq_self)``): Avoid invalid areas to participate in the calculation of self-attention.
      self_position_bias (:obj:`torch.Tensor` of shape ``(num_heads, seq_self, seq_self)``): Provide positional information to self-attention block.
  
  Return:
      :obj:`torch.Tensor` of shape ``(batch, seq_self, dim_model)``: The output of transformer block.
docstr_tok: ["Args",":","self_hidden_states","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","seq_self",",","dim_model",")","``",")",":","Input","of","transformer","block","(","self-attention","block",")",".","It","can","be","the","raw","embedding","of","a","batch","of","sequences",".","self_attention_mask","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","seq_self",",","seq_self",")","``",")",":","Avoid","invalid","areas","to","participate","in","the","calculation","of","self-attention",".","self_position_bias","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","num_heads",",","seq_self",",","seq_self",")","``",")",":","Provide","positional","information","to","self-attention","block",".","Return",":",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","seq_self",",","dim_model",")","``",":","The","output","of","transformer","block","."]
code: def forward(
      self,
      self_hidden_states: torch.Tensor,
      self_attention_mask: torch.Tensor,
      self_position_bias: Optional[torch.Tensor] = None,
      use_cache: bool = False,
      past_key_value: Optional[
          Tuple[torch.Tensor, torch.Tensor]
      ] = None,
  ):
      """
      Args:
          self_hidden_states (:obj:`torch.Tensor` of shape ``(batch, seq_self, dim_model)``): Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
          self_attention_mask (:obj:`torch.Tensor` of shape ``(batch, seq_self, seq_self)``): Avoid invalid areas to participate in the calculation of self-attention.
          self_position_bias (:obj:`torch.Tensor` of shape ``(num_heads, seq_self, seq_self)``): Provide positional information to self-attention block.
  
      Return:
          :obj:`torch.Tensor` of shape ``(batch, seq_self, dim_model)``: The output of transformer block.
  
      """  # noqa: E501
      # (batch, dim_model, seq_self)
      current_key_value = None
      if not self.mask_att:
          hidden_states = self.self_att(
              self_hidden_states,
              attention_mask=self_attention_mask,
              position_bias=self_position_bias,
              use_cache=use_cache,
              past_key_value=past_key_value,
          )
          if use_cache:
              hidden_states, current_key_value = hidden_states
      else:
          hidden_states = self_hidden_states
  
      # (batch, dim_model, seq_self)
      if not self.mask_ffn:
          hidden_states = self.ffn(hidden_states)
  
      if use_cache:
          return hidden_states, current_key_value
      else:
          return hidden_states
  
file_path: DeepKE/example/llm/CPM-Bee/src/cpm_live/layers/blocks.py
imports: from typing import Optional, Tuple
  import torch
  import bmtrain
  from .layernorm import LayerNorm
  from .attention import Attention
  from .feedforward import FeedForward
label_desc: Use image node_id 2 for calling the TransformerBlock obj's underlying member method code with example usage: obj.forward(_hidden_states, _attention_mask, _position_bias, use_cache, past_key_value) and returns: hidden_states,  current_key_value, hidden_states