code: def forward(
      self, query: torch.Tensor, key_soft: torch.Tensor
  ) -> torch.Tensor:
      values = self.W_value(key_soft)
      split_size = self.num_units // self.num_heads
      values = torch.stack(
          torch.split(values, split_size, dim=2), dim=0
      )
  
      out_soft = scores_soft = None
      querys = self.W_query(query)  # [N, T_q, num_units]
      keys = self.W_key(key_soft)  # [N, T_k, num_units]
  
      # [h, N, T_q, num_units/h]
      querys = torch.stack(
          torch.split(querys, split_size, dim=2), dim=0
      )
      # [h, N, T_k, num_units/h]
      keys = torch.stack(torch.split(keys, split_size, dim=2), dim=0)
      # [h, N, T_k, num_units/h]
  
      # score = softmax(QK^T / (d_k ** 0.5))
      scores_soft = torch.matmul(
          querys, keys.transpose(2, 3)
      )  # [h, N, T_q, T_k]
      scores_soft = scores_soft / (self.key_dim**0.5)
      scores_soft = F.softmax(scores_soft, dim=3)
  
      # out = score * V
      # [h, N, T_q, num_units/h]
      out_soft = torch.matmul(scores_soft, values)
      out_soft = torch.cat(
          torch.split(out_soft, 1, dim=0), dim=3
      ).squeeze(
          0
      )  # [N, T_q, num_units]
  
      return out_soft
  
imports: import math
  from typing import Tuple
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from TTS.tts.layers.delightful_tts.conv_layers import ConvNorm
label_desc: Use image node_id 2 for calling the StyleEmbedAttention obj's underlying member method code with example usage: obj.forward(query, key_soft) and returns: out_soft