code: def get_masks(self, input_ids, past_length, padding_mask=None):
      batch_size, seq_length = input_ids.shape
      full_attention_mask = torch.ones(
          batch_size, seq_length, seq_length, device=input_ids.device
      )
      full_attention_mask.tril_()
      if past_length:
          full_attention_mask = torch.cat(
              (
                  torch.ones(
                      batch_size,
                      seq_length,
                      past_length,
                      device=input_ids.device,
                  ),
                  full_attention_mask,
              ),
              dim=-1,
          )
  
      if padding_mask is not None:
          full_attention_mask = (
              full_attention_mask * padding_mask.unsqueeze(1)
          )
      if not past_length and padding_mask is not None:
          full_attention_mask -= padding_mask.unsqueeze(-1) - 1
      full_attention_mask = (full_attention_mask < 0.5).bool()
      full_attention_mask.unsqueeze_(1)
      return full_attention_mask
  
file_path: ColossalAI/colossalai/inference/engine/modeling/chatglm2.py
imports: from typing import List, Optional, Tuple
  import torch
  from transformers.utils import logging
  from colossalai.inference.kv_cache import BatchInferState
  from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer import ShardConfig
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel, GLMBlock, GLMTransformer, SelfAttention, split_tensor_along_last_dim
  from ._utils import copy_kv_to_mem_cache
label_desc: Use image node_id 1 for calling a global function with example usage: get_masks(input_ids, past_length, padding_mask) and returns: full_attention_mask