docstr: attention_scores.shape = [4, 12, 24, 24]
  mask = None
docstr_tok: ["attention_scores.shape","=","[","4",",","12",",","24",",","24","]","mask","=","None"]
code: def test_fused_scale_softmax(self):
      """
      attention_scores.shape = [4, 12, 24, 24]
      mask = None
      """
      for dtype, scale, softmax_in_fp32, shape in itertools.product(
          (torch.half, torch.bfloat16),
          (None, 2.0),
          (False, True),
          ((4, 12, 24, 24), (32, 12, 4, 214)),
      ):
          msg = f"{dtype}-{scale}-{softmax_in_fp32}"
          input_in_fp16 = dtype == torch.half
          input_in_bf16 = dtype == torch.bfloat16
          if not (scale is None or softmax_in_fp32):
              with self.assertRaises(RuntimeError, msg=msg):
                  self._setup_fused_softmax(
                      input_in_fp16,
                      input_in_bf16,
                      scale,
                      softmax_in_fp32,
                      AttnMaskType.padding,
                  )
              return
          fused_fn, torch_fn = self._setup_fused_softmax(
              input_in_fp16,
              input_in_bf16,
              scale,
              softmax_in_fp32,
              AttnMaskType.padding,
          )
  
          attention_scores_0 = (
              torch.randn(shape)
              .to(device="cuda", dtype=dtype)
              .requires_grad_(True)
          )
          with torch.no_grad():
              attention_scores_1 = (
                  attention_scores_0.clone().requires_grad_(True)
              )
          mask = None
  
          expected = fused_fn(attention_scores_0, mask)
          actual = torch_fn(attention_scores_1, mask)
          self.assertEqual(actual, expected, msg=msg)
  
          g0 = torch.rand_like(actual)
          with torch.no_grad():
              g1 = g0.clone()
          expected.backward(g0)
          actual.backward(g1)
  
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 5 for calling the TestFusedScaleMaskSoftmax obj's underlying member method code with example usage: obj.test_fused_scale_softmax() without return types