class_docstr: input:
      query --- [N, T_q, query_dim]
      key --- [N, T_k, key_dim]
  output:
      out --- [N, T_q, num_units]
class_docstr_tok: ["input",":","query","--","-","[","N",",","T_q",",","query_dim","]","key","--","-","[","N",",","T_k",",","key_dim","]","output",":","out","--","-","[","N",",","T_q",",","num_units","]"]
code: def forward(self, query, key):
      queries = self.W_query(query)  # [N, T_q, num_units]
      keys = self.W_key(key)  # [N, T_k, num_units]
      values = self.W_value(key)
  
      split_size = self.num_units // self.num_heads
      queries = torch.stack(
          torch.split(queries, split_size, dim=2), dim=0
      )  # [h, N, T_q, num_units/h]
      keys = torch.stack(
          torch.split(keys, split_size, dim=2), dim=0
      )  # [h, N, T_k, num_units/h]
      values = torch.stack(
          torch.split(values, split_size, dim=2), dim=0
      )  # [h, N, T_k, num_units/h]
  
      # score = softmax(QK^T / (d_k ** 0.5))
      scores = torch.matmul(
          queries, keys.transpose(2, 3)
      )  # [h, N, T_q, T_k]
      scores = scores / (self.key_dim**0.5)
      scores = F.softmax(scores, dim=3)
  
      # out = score * V
      out = torch.matmul(scores, values)  # [h, N, T_q, num_units/h]
      out = torch.cat(torch.split(out, 1, dim=0), dim=3).squeeze(
          0
      )  # [N, T_q, num_units]
  
      return out
  
file_path: TTS/TTS/tts/layers/gst_layers.py
imports: import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(query, key) and returns: out