code: def expand_mask(
      mask: torch.Tensor, tgt_length: int
  ) -> torch.BoolTensor:
      """
      Expands attention_mask from `[batch_size, src_length]` to `[batch_size, 1, tgt_length, src_length]`.
      """
      batch_size, src_length = mask.shape
      tgt_length = tgt_length if tgt_length is not None else src_length
  
      expanded_mask = ~(mask[:, None, :].to(torch.bool))
      return expanded_mask.expand(batch_size, tgt_length, src_length)
  
imports: from typing import Optional, Tuple, Union
  import os
  import torch
  import torch.distributed
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.activations import ACT2FN
  from transformers.file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast, TokenClassifierOutput
  from transformers.modeling_utils import PreTrainedModel
  from transformers import GPTNeoXConfig
  from loguru import logger
  from text_generation_server.utils.layers import TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelRowLinear, TensorParallelHead
label_desc: Use image node_id 2 for calling a global function with example usage: expand_mask(mask, tgt_length) and returns: expanded_mask