code: def bert_extended_attention_mask(attention_mask):
      # We create a 3D attention mask from a 2D tensor mask.
      # [b, 1, s]
      attention_mask_b1s = attention_mask.unsqueeze(1)
      # [b, s, 1]
      attention_mask_bs1 = attention_mask.unsqueeze(2)
      # [b, s, s]
      attention_mask_bss = attention_mask_b1s * attention_mask_bs1
      # [b, 1, s, s]
      extended_attention_mask = attention_mask_bss.unsqueeze(1)
  
      # Convert attention mask to binary:
      extended_attention_mask = extended_attention_mask < 0.5
  
      return extended_attention_mask
  
imports: import torch
  from megatron import get_args
  from megatron.core import tensor_parallel
  from megatron.model.enums import AttnMaskType
  from megatron.model.language_model import parallel_lm_logits
  from megatron.model.language_model import get_language_model
  from megatron.model.utils import get_norm
  from megatron.model.utils import openai_gelu, erf_gelu
  from megatron.model.utils import get_linear_layer
  from megatron.model.utils import init_method_normal
  from megatron.model.utils import scaled_init_method_normal
  from .module import MegatronModule
label_desc: Use image node_id 1 for calling a global function with example usage: bert_extended_attention_mask(attention_mask) and returns: extended_attention_mask