code: def forward_torch_softmax(input, mask, scale):
      input = input * scale
      mask_output = (
          attention_mask_func(input, mask)
          if mask is not None
          else input
      )
      probs = torch.nn.Softmax(dim=-1)(mask_output)
      all_k_masked = mask.all(axis=-1)
      zero_attention_mask = (1.0 - all_k_masked.float())[:, :, :, None]
      probs = probs * zero_attention_mask
      return probs
  
imports: import itertools
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer import AttnMaskType
  from apex.transformer.functional import FusedScaleMaskSoftmax
label_desc: Use image node_id 2 for calling a global function with example usage: forward_torch_softmax(input, mask, scale) and returns: probs