class_docstr: Layers of encoder transformer blocks plus an final layernorm.
  
  Args:
      num_layers (int): number of layers.
      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-6.
      dropout_p (float, optional): Defaults to 0.
docstr: Args:
      hidden-states (:obj:`torch.Tensor` of shape ``(batch, seq_enc, dim_model)``): Input of encoder, might be the embedding of a batch of sequences.
      attention_mask (:obj:`torch.Tensor` of shape ``(batch, seq_enc, seq_enc)``): Avoid invalid areas to participate in the calculation
      position_bias(:obj:`torch.Tensor` of shape ``(num_heads, seq_enc, seq_enc)``) Provides position information to attention mechanism.
  
  Return:
      :obj:`torch.Tensor` of shape ``(batch, seq_enc, dim_model)``: The encoder output.
docstr_tok: ["Args",":","hidden-states","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","seq_enc",",","dim_model",")","``",")",":","Input","of","encoder",",","might","be","the","embedding","of","a","batch","of","sequences",".","attention_mask","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","seq_enc",",","seq_enc",")","``",")",":","Avoid","invalid","areas","to","participate","in","the","calculation","position_bias","(",":","obj",":","`","torch.Tensor","`","of","shape","``","(","num_heads",",","seq_enc",",","seq_enc",")","``",")","Provides","position","information","to","attention","mechanism",".","Return",":",":","obj",":","`","torch.Tensor","`","of","shape","``","(","batch",",","seq_enc",",","dim_model",")","``",":","The","encoder","output","."]
code: def forward(
      self,
      hidden_states: torch.Tensor,
      attention_mask: torch.Tensor,
      position_bias: torch.Tensor,
      use_cache: bool = False,
      past_key_values: Optional[
          List[Tuple[torch.Tensor, torch.Tensor]]
      ] = None,
  ):
      """
      Args:
          hidden-states (:obj:`torch.Tensor` of shape ``(batch, seq_enc, dim_model)``): Input of encoder, might be the embedding of a batch of sequences.
          attention_mask (:obj:`torch.Tensor` of shape ``(batch, seq_enc, seq_enc)``): Avoid invalid areas to participate in the calculation
          position_bias(:obj:`torch.Tensor` of shape ``(num_heads, seq_enc, seq_enc)``) Provides position information to attention mechanism.
  
      Return:
          :obj:`torch.Tensor` of shape ``(batch, seq_enc, dim_model)``: The encoder output.
  
      """  # noqa: E501
      if not use_cache:
          hidden_states = self.layers(
              hidden_states, attention_mask, position_bias
          )
          hidden_states = self.output_layernorm(hidden_states)
          return hidden_states
      else:
          with torch.no_grad():
              current_key_values = []
              for i, module in enumerate(self.layers):
                  hidden_states = module(
                      hidden_states,
                      attention_mask,
                      position_bias,
                      past_key_value=past_key_values[i]
                      if past_key_values
                      else None,
                      use_cache=use_cache,
                  )
                  if use_cache:
                      current_key_values.append(hidden_states[1])
                      hidden_states = hidden_states[0]
              hidden_states = self.output_layernorm(hidden_states)
              if use_cache:
                  return hidden_states, current_key_values
              else:
                  return hidden_states
  
file_path: DeepKE/example/llm/CPM-Bee/src/cpm_live/layers/transformer.py
imports: import torch
  import bmtrain
  from typing import Optional, List, Tuple
  from .blocks import TransformerBlock
  from .layernorm import LayerNorm
label_desc: Use image node_id 2 for calling the Encoder obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, position_bias, use_cache, past_key_values) and returns: hidden_states, hidden_states,  current_key_values, hidden_states