code: def build_attention_mask_3d(source_mask, target_mask, attn_mask_type):
      """
      Returns a 3D attention mask for Megatron given two 2D masks
      :param source_mask - < 0.5 for non-masked, else masked [batch, src length]
      :param target_mask - < 0.5 for non-masked, else masked [batch, tgt length]
      :param attn_mask_type - AttnMaskType enum
      """
      if attn_mask_type == AttnMaskType.padding:
          mask = build_attention_mask_3d_padding(
              source_mask, target_mask
          )
      elif attn_mask_type == AttnMaskType.causal:
          mask = build_attention_mask_3d_causal(
              source_mask, target_mask
          )
      else:
          raise ValueError(
              f"Unsupported attention mask attn_mask_type = {attn_mask_type}"
          )
  
      return mask
  
imports: import itertools
  import math
  from typing import Dict, Iterator, List, Tuple, Union
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 22 for calling a global function with example usage: build_attention_mask_3d(source_mask, target_mask, attn_mask_type) and returns: mask