class_docstr: fused operation: scaling + mask + softmax
  
  Arguments:
      input_in_fp16: flag to indicate if input in fp16 data format.
      input_in_bf16: flag to indicate if input in bf16 data format.
      attn_mask_type: attention mask type (pad or causal)
      scaled_masked_softmax_fusion: flag to indicate user want to use softmax fusion
      mask_func: mask function to be applied.
      softmax_in_fp32: if true, softmax in performed at fp32 precision.
      scale: scaling factor used in input tensor scaling.
class_docstr_tok: ["fused","operation",":","scaling","+","mask","+","softmax","Arguments",":","input_in_fp16",":","flag","to","indicate","if","input","in","fp16","data","format",".","input_in_bf16",":","flag","to","indicate","if","input","in","bf16","data","format",".","attn_mask_type",":","attention","mask","type","(","pad","or","causal",")","scaled_masked_softmax_fusion",":","flag","to","indicate","user","want","to","use","softmax","fusion","mask_func",":","mask","function","to","be","applied",".","softmax_in_fp32",":","if","true",",","softmax","in","performed","at","fp32","precision",".","scale",":","scaling","factor","used","in","input","tensor","scaling","."]
code: def __init__(
      self,
      input_in_fp16,
      input_in_bf16,
      attn_mask_type,
      scaled_masked_softmax_fusion,
      mask_func,
      softmax_in_fp32,
      scale,
  ):
      super().__init__()
      self.input_in_fp16 = input_in_fp16
      self.input_in_bf16 = input_in_bf16
      if self.input_in_fp16 and self.input_in_bf16:
          raise RuntimeError(
              "both fp16 and bf16 flags cannot be active at the same time."
          )
      self.input_in_float16 = self.input_in_fp16 or self.input_in_bf16
      self.attn_mask_type = attn_mask_type
      self.scaled_masked_softmax_fusion = scaled_masked_softmax_fusion
      self.mask_func = mask_func
      self.softmax_in_fp32 = softmax_in_fp32
      self.scale = scale
  
      if not (self.scale is None or softmax_in_fp32):
          raise RuntimeError("softmax should be in fp32 when scaled")
  
      if self.scaled_masked_softmax_fusion:
          if self.attn_mask_type == AttnMaskType.causal:
              self.fused_softmax_func = (
                  scaled_upper_triang_masked_softmax
              )
          elif self.attn_mask_type == AttnMaskType.padding:
              self.fused_softmax_func = scaled_masked_softmax
          else:
              raise ValueError("Invalid attn_mask_type.")
  
imports: import torch
  from apex._autocast_utils import _cast_if_autocast_enabled
  from apex.transformer.enums import AttnMaskType
label_desc: Use image node_id 1 to create a new FusedScaleMaskSoftmax object from inherited base classes: torch.nn with example: obj = FusedScaleMaskSoftmax(input_in_fp16, input_in_bf16, attn_mask_type, scaled_masked_softmax_fusion, mask_func, softmax_in_fp32, scale)