code: def chatglm_attention_forward(
      self,
      hidden_states: torch.Tensor,
      position_ids,
      attention_mask: torch.Tensor,
      layer_id,
      layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
      use_cache: bool = False,
      output_attentions: bool = False,
  ):
      """
      hidden_states: [seq_len, batch, hidden_size]
      attention_mask: [(1, 1), seq_len, seq_len]
      """
  
      # [seq_len, batch, 3 * hidden_size]
      mixed_raw_layer = self.query_key_value(hidden_states)
  
      # [seq_len, batch, 3 * hidden_size] -->
      #   [seq_len, batch, num_attention_heads, 3 * hidden_size_per_attention_head]
      new_tensor_shape = mixed_raw_layer.size()[:-1] + (
          self.num_attention_heads_per_partition,
          3 * self.hidden_size_per_attention_head,
      )
      mixed_raw_layer = mixed_raw_layer.view(*new_tensor_shape)
      # [seq_len, batch, num_attention_heads, hidden_size_per_attention_head]
      (
          query_layer,
          key_layer,
          value_layer,
      ) = self.split_tensor_along_last_dim(mixed_raw_layer, 3)
  
      if self.position_encoding_2d:
          q1, q2 = query_layer.chunk(2, dim=(query_layer.ndim - 1))
          k1, k2 = key_layer.chunk(2, dim=(key_layer.ndim - 1))
          cos, sin = self.rotary_emb(q1, seq_len=position_ids.max() + 1)
          position_ids, block_position_ids = (
              position_ids[:, 0, :].transpose(0, 1).contiguous(),
              position_ids[:, 1, :].transpose(0, 1).contiguous(),
          )
          q1, k1 = apply_rotary_pos_emb_index(
              q1, k1, cos, sin, position_ids
          )
          q2, k2 = apply_rotary_pos_emb_index(
              q2, k2, cos, sin, block_position_ids
          )
          query_layer = torch.concat([q1, q2], dim=(q1.ndim - 1))
          key_layer = torch.concat([k1, k2], dim=(k1.ndim - 1))
      else:
          position_ids = position_ids.transpose(0, 1)
          cos, sin = self.rotary_emb(
              value_layer, seq_len=position_ids.max() + 1
          )
          # [seq_len, batch, num_attention_heads, hidden_size_per_attention_head]
          query_layer, key_layer = apply_rotary_pos_emb_index(
              query_layer, key_layer, cos, sin, position_ids
          )
  
      # [seq_len, batch, hidden_size]
      context_layer, present, attention_probs = attention_fn(
          self=self,
          query_layer=query_layer,
          key_layer=key_layer,
          value_layer=value_layer,
          attention_mask=attention_mask,
          hidden_size_per_partition=self.hidden_size_per_partition,
          layer_id=layer_id,
          layer_past=layer_past,
          use_cache=use_cache,
      )
  
      output = self.dense(context_layer)
  
      outputs = (output, present)
  
      if output_attentions:
          outputs += (attention_probs,)
  
      return outputs
  
file_path: BigDL/python/llm/src/bigdl/llm/transformers/models/chatglm.py
imports: import math
  import torch
  import torch.utils.checkpoint
  import torch.nn.functional
  from typing import Optional, Tuple
  from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache
label_desc: Use image node_id 4 for calling a global function with example usage: chatglm_attention_forward(hidden_states, position_ids, attention_mask, layer_id, layer_past, use_cache, output_attentions) and returns: outputs