code: def get_jit_fused_glm_block_forward():
      from .chatglm2_6b.modeling_chatglm import GLMBlock
  
      def forward(
          self: GLMBlock,
          hidden_states,
          attention_mask,
          rotary_pos_emb,
          kv_cache=None,
          use_cache=True,
      ):
          # hidden_states: [s, b, h]
          # Layer norm at the beginning of the transformer layer.
          layernorm_output = self.input_layernorm(hidden_states)
          # Self attention.
          attention_output, kv_cache = self.self_attention(
              layernorm_output,
              attention_mask,
              rotary_pos_emb,
              kv_cache=kv_cache,
              use_cache=use_cache,
          )
  
          # Residual connection.
          if self.apply_residual_connection_post_layernorm:
              residual = layernorm_output
          else:
              residual = hidden_states
  
          layernorm_input = self.dropout_add(
              attention_output,
              residual,
              self.hidden_dropout,
              self.training,
          )
  
          # Layer norm post the self attention.
          layernorm_output = self.post_attention_layernorm(
              layernorm_input
          )
  
          # MLP.
          mlp_output = self.mlp(layernorm_output)
  
          # Second residual connection.
          if self.apply_residual_connection_post_layernorm:
              residual = layernorm_output
          else:
              residual = layernorm_input
  
          output = self.dropout_add(
              mlp_output, residual, self.hidden_dropout, self.training
          )
  
          return output, kv_cache
  
      return forward
  
file_path: ColossalAI/colossalai/shardformer/modeling/chatglm2.py
imports: from typing import List, Optional, Tuple
  import torch
  import torch.utils.checkpoint
  from torch.nn import CrossEntropyLoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer import ShardConfig
  from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel
label_desc: Use image node_id 2 for calling a global function with example usage: get_jit_fused_glm_block_forward() and returns: forward, output,  kv_cache