code: def test_self_multihead_attn_pad_mask(self):
      grads = torch.randn_like(self.tst_inputs)
      pad_mask_byte = torch.tril(
          torch.ones(
              self.tst_inputs.size(1),
              self.tst_inputs.size(0),
              device=torch.device("cuda"),
              dtype=torch.uint8,
          ),
          1,
      )
      pad_mask_bool = pad_mask_byte.to(torch.bool)
  
      ref_outputs, _ = self.ref_layer.forward(
          self.ref_inputs,
          self.ref_inputs,
          self.ref_inputs,
          key_padding_mask=pad_mask_bool,
          need_weights=False,
          attn_mask=None,
          is_training=True,
      )
  
      tst_outputs, _ = self.tst_layer.forward(
          self.tst_inputs,
          self.tst_inputs,
          self.tst_inputs,
          key_padding_mask=pad_mask_byte,
          need_weights=False,
          attn_mask=None,
          is_training=True,
      )
  
      self.ref_inputs.backward(grads)
      self.tst_inputs.backward(grads)
  
      torch.testing.assert_close(
          self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5
      )
      torch.testing.assert_close(
          ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3
      )
      torch.testing.assert_close(
          self.ref_inputs.grad,
          self.tst_inputs.grad,
          atol=1e-3,
          rtol=1e-3,
      )
  
file_path: apex/apex/contrib/test/multihead_attn/test_self_multihead_attn.py
imports: import unittest
  import torch
label_desc: Use image node_id 4 for calling the SelfMultiheadAttnTest obj's underlying member method code with example usage: obj.test_self_multihead_attn_pad_mask() without return types