code: def _setup_fused_softmax(
      self,
      input_in_fp16,
      input_in_bf16,
      scale=None,
      softmax_in_fp32=False,
      attn_mask_type=AttnMaskType.padding,
  ):
      fused_fn = FusedScaleMaskSoftmax(
          input_in_fp16=input_in_fp16,
          input_in_bf16=input_in_bf16,
          mask_func=attention_mask_func,
          scale=scale,
          softmax_in_fp32=softmax_in_fp32,
          attn_mask_type=attn_mask_type,
          scaled_masked_softmax_fusion=True,
      )
      torch_fn = FusedScaleMaskSoftmax(
          input_in_fp16=input_in_fp16,
          input_in_bf16=input_in_bf16,
          mask_func=attention_mask_func,
          scale=scale,
          softmax_in_fp32=softmax_in_fp32,
          attn_mask_type=attn_mask_type,
          scaled_masked_softmax_fusion=False,
      )
      return fused_fn, torch_fn
  
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 1 for calling the TestFusedScaleMaskSoftmax obj's underlying member method code with example usage: obj._setup_fused_softmax(input_in_fp16, input_in_bf16, scale, softmax_in_fp32, attn_mask_type) and returns: fused_fn,  torch_fn