code: def attention_mask_func(attention_scores, attention_mask):
      attention_scores.masked_fill_(attention_mask, -10000.0)
      return attention_scores
  
file_path: megatron-lm/megatron/core/transformer/utils.py
imports: from operator import itemgetter
  from typing import Any, Dict, Iterable, Optional, Tuple, Union
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.mapping import ShardedObject, StateDict
  from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_tp_sharded_tensor_for_checkpoint
label_desc: Use image node_id 2 for calling a global function with example usage: attention_mask_func(attention_scores, attention_mask) and returns: attention_scores