code: def forward(self, query, key, value, mask):
      if not (query is key and key is value):
          raise NotImplementedError(
              "query, key and value must be the same Tensor for now."
          )
      if mask is not None:
          raise NotImplementedError("mask is currently not supported.")
  
      query_projected = torch.nn.functional.linear(
          query, self.in_proj_weight, self.in_proj_bias
      )
  
      batch_size = query_projected.size(0)
      embed_dim = query_projected.size(2)
      head_dim = embed_dim // (self.num_heads * 3)
  
      query, key, value = query_projected.chunk(3, -1)
  
      query = query.view(
          batch_size, -1, self.num_heads, head_dim
      ).transpose(1, 2)
      key = key.view(
          batch_size, -1, self.num_heads, head_dim
      ).transpose(1, 2)
      value = value.view(
          batch_size, -1, self.num_heads, head_dim
      ).transpose(1, 2)
  
      # the output of sdp = (batch, num_heads, seq_len, head_dim)
      attn = torch.nn.functional.scaled_dot_product_attention(
          query,
          key,
          value,
          attn_mask=None,
          dropout_p=0.0,
          is_causal=False,
      )
  
      attn = attn.transpose(1, 2).reshape(
          batch_size, -1, self.num_heads * head_dim
      )
      # Match return signature of nn.MHA
      return self.out_proj(attn), None
  
imports: import argparse
  import itertools
  import random
  import warnings
  from dataclasses import dataclass
  from pathlib import Path
  from pprint import pprint
  from typing import List, Optional
  import numpy
  import torch
  import torch.utils.benchmark
  from prettytable import PrettyTable
  from torch.backends.cuda import sdp_kernel
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling the CompositeMHA obj's underlying member method code with example usage: obj.forward(query, key, value, mask) and returns: self,  None