code: def _make_causal_mask(
      input_ids_shape: torch.Size,
      dtype: torch.dtype,
      device: torch.device,
      past_key_values_length: int = 0,
  ):
      """
      Make causal mask used for bi-directional self-attention.
      """
      bsz, tgt_len = input_ids_shape
      mask = torch.full(
          (tgt_len, tgt_len), torch.finfo(dtype).min, device=device
      )
      mask_cond = torch.arange(mask.size(-1), device=device)
      mask.masked_fill_(
          mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0
      )
      mask = mask.to(dtype)
  
      if past_key_values_length > 0:
          mask = torch.cat(
              [
                  torch.zeros(
                      tgt_len,
                      past_key_values_length,
                      dtype=dtype,
                      device=device,
                  ),
                  mask,
              ],
              dim=-1,
          )
      return mask[None, None, :, :].expand(
          bsz, 1, tgt_len, tgt_len + past_key_values_length
      )
  
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch import nn
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from transformers.modeling_utils import PreTrainedModel
  from transformers.models.llama.modeling_llama import LlamaConfig, LlamaRMSNorm
  from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from colossalai.kernel.cuda_native.mha.flash_attn_2 import HAS_FLASH_ATTN
  from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
  from colossalai.moe.layers import SparseMLP
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.moe.utils import get_activation, set_moe_args
label_desc: Use image node_id 2 for calling a global function with example usage: _make_causal_mask(input_ids_shape, dtype, device, past_key_values_length) and returns: mask