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)
      return probs
  
file_path: megatron-lm/megatron/fused_kernels/tests/test_fused_kernels.py
imports: import math
  import torch
  from torch.nn import LayerNorm
  from megatron.model.enums import AttnMaskType
  from megatron.model.fused_layer_norm import MixedFusedLayerNorm
  from megatron.model.fused_softmax import FusedScaleMaskSoftmax
  from megatron.model.utils import attention_mask_func
  from megatron.fused_kernels import load
label_desc: Use image node_id 6 for calling a global function with example usage: forward_torch_softmax(input, mask, scale) and returns: probs