docstr: attn_weights.shape: [4, 12, 24, 24]
  total_mask.shape: [4, 1, 24, 24]
  
  total_mask[0, 0], a 24x24 matrix is like a lower triangular matrix, but
  upper elements are True and lower elements and diagonal are False.
docstr_tok: ["attn_weights.shape",":","[","4",",","12",",","24",",","24","]","total_mask.shape",":","[","4",",","1",",","24",",","24","]","total_mask","[","0",",","0","]",",","a","24x24","matrix","is","like","a","lower","triangular","matrix",",","but","upper","elements","are","True","and","lower","elements","and","diagonal","are","False","."]
code: def test_fused_upper_triangle_mask_softmax(self):
      """
      attn_weights.shape: [4, 12, 24, 24]
      total_mask.shape: [4, 1, 24, 24]
  
      total_mask[0, 0], a 24x24 matrix is like a lower triangular matrix, but
      upper elements are True and lower elements and diagonal are False.
      """
      for dtype, scale, softmax_in_fp32 in itertools.product(
          (torch.half, torch.bfloat16),
          (None, 2.0),
          (False, True),
      ):
          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.causal,
                  )
              return
          fused_fn, torch_fn = self._setup_fused_softmax(
              input_in_fp16,
              input_in_bf16,
              scale,
              softmax_in_fp32,
              AttnMaskType.causal,
          )
  
          attn_weights_0 = (
              torch.randn((4, 12, 24, 24))
              .to(device="cuda", dtype=dtype)
              .requires_grad_(True)
          )
          with torch.no_grad():
              attn_weights_1 = attn_weights_0.clone().requires_grad_(
                  True
              )
          total_mask = (
              ~(torch.tril(torch.randn((24, 24), device="cuda")).bool())
              .unsqueeze(0)
              .unsqueeze(0)
          )
          total_mask = total_mask.repeat((4, 1, 1, 1))
          expected = fused_fn(attn_weights_0, total_mask)
          actual = torch_fn(attn_weights_1, total_mask)
          self.assertEqual(actual, expected, msg=msg)
  
          g0 = torch.randn_like(actual)
          with torch.no_grad():
              g1 = g0.clone()
          actual.backward(g0)
          expected.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 7 for calling the TestFusedScaleMaskSoftmax obj's underlying member method code with example usage: obj.test_fused_upper_triangle_mask_softmax() without return types