code: def scaled_masked_softmax(inputs, mask, scale):
      # input is 4D tensor (b, np, sq, sk)
      if mask is not None:
          args = _cast_if_autocast_enabled(inputs, mask, scale)
          with torch.cuda.amp.autocast(enabled=False):
              return ScaledMaskedSoftmax.apply(*args)
      else:
          args = _cast_if_autocast_enabled(inputs, scale)
          with torch.cuda.amp.autocast(enabled=False):
              return ScaledSoftmax.apply(*args)
  
imports: import torch
  from apex._autocast_utils import _cast_if_autocast_enabled
  from apex.transformer.enums import AttnMaskType
label_desc: Use image node_id 2 for calling a global function with example usage: scaled_masked_softmax(inputs, mask, scale) and returns: ScaledMaskedSoftmax, ScaledSoftmax