code: def __init__(self, embedding_size: int, num_heads: int) -> None:
      super(MultiHeadAttention, self).__init__()
      self.num_heads = num_heads
      self.embedding_size = embedding_size
      assert embedding_size % self.num_heads == 0
      self.depth = embedding_size // self.num_heads
      self.wq = tf.keras.layers.Dense(embedding_size)
      self.wk = tf.keras.layers.Dense(embedding_size)
      self.wv = tf.keras.layers.Dense(embedding_size)
      self.dense = tf.keras.layers.Dense(embedding_size)
  
imports: import numpy
  import tensorflow
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: tf.keras.layers with example: obj = MultiHeadAttention(embedding_size, num_heads)