code: def _test_multihead_attn_invalid_shape_impl(self, mha):
      # Batched (3D) query cases
      query = torch.randn(4, 4, 4)
      key = torch.randn(4, 4, 4)
      value = torch.randn(4, 4, 4)
  
      msg = "expected `key` and `value` to be 3-D but found 2-D and 3-D tensors respectively"
      # 3D query, 2D key and 3D value
      with self.assertRaisesRegex(AssertionError, msg):
          mha(query, torch.randn(4, 4), value)
  
      msg = "expected `key` and `value` to be 3-D but found 3-D and 2-D tensors respectively"
      # 3D query, 3D key and 2D value
      with self.assertRaisesRegex(AssertionError, msg):
          mha(query, key, torch.randn(4, 4))
  
      msg = "expected `key_padding_mask` to be `None` or 2-D but found 1-D tensor instead"
      # 3D query, 3D key, 3D value and 1D key_padding_mask
      with self.assertRaisesRegex(AssertionError, msg):
          mha(
              query,
              key,
              value,
              key_padding_mask=torch.tensor(
                  [False, False, True, True], dtype=torch.bool
              ),
          )
  
      msg = "expected `attn_mask` to be `None`, 2-D or 3-D but found 1-D tensor instead"
      # 3D query, 3D key, 3D value and 1D attn_mask
      with self.assertRaisesRegex(AssertionError, msg):
          mha(
              query,
              key,
              value,
              attn_mask=torch.tensor(
                  [False, False, True, True], dtype=torch.bool
              ),
          )
  
      # Unbatched (2D) query cases
      query = torch.randn(4, 4)
      key = torch.randn(4, 4)
      value = torch.randn(4, 4)
  
      msg = "expected `key` and `value` to be 2-D but found 3-D and 2-D tensors respectively"
      # 2D query, 3D key and 2D value
      with self.assertRaisesRegex(AssertionError, msg):
          mha(query, torch.randn(4, 4, 4), value)
  
      msg = "expected `key` and `value` to be 2-D but found 2-D and 3-D tensors respectively"
      # 2D query, 3D key and 2D value
      with self.assertRaisesRegex(AssertionError, msg):
          mha(query, key, torch.randn(4, 4, 4))
  
      msg = "expected `key_padding_mask` to be `None` or 1-D but found 2-D tensor instead"
      # 2D query, 2D key, 2D value and 1D key_padding_mask
      with self.assertRaisesRegex(AssertionError, msg):
          mha(
              query,
              key,
              value,
              key_padding_mask=torch.tensor(
                  [[False, False, True, True] * 2], dtype=torch.bool
              ),
          )
  
      msg = "expected `attn_mask` to be `None`, 2-D or 3-D but found 1-D tensor instead"
      # 2D query, 2D key, 2D value and 1D attn_mask
      with self.assertRaisesRegex(AssertionError, msg):
          mha(
              query,
              key,
              value,
              attn_mask=torch.tensor(
                  [False, False, True, True], dtype=torch.bool
              ),
          )
  
      msg = r"Expected `attn_mask` shape to be \(4, 4, 4\)"
      # 2D query, 2D key, 2D value and 3D incorrect attn_mask
      with self.assertRaisesRegex(AssertionError, msg):
          mha(
              query,
              key,
              value,
              attn_mask=torch.randn(5, 4, 4)
              .bernoulli_()
              .to(torch.bool),
          )
  
imports: import contextlib
  import random
  import unittest
  import unittest.mock
  from torch.nn import MultiheadAttention
  from torch.testing._internal.common_device_type import instantiate_device_type_tests, dtypes, onlyCUDAAndPRIVATEUSE1
  from torch.testing._internal.common_nn import NNTestCase
  from torch.testing._internal.common_utils import run_tests, TEST_NUMPY, TEST_WITH_CROSSREF, parametrize, instantiate_parametrized_tests
  import torch.nn
  import torch
label_desc: Use image node_id 4 for calling the TestMultiheadAttentionNN obj's underlying member method code with example usage: obj._test_multihead_attn_invalid_shape_impl(mha) without return types