class_docstr:  Multi-headed attention is based on "Attention
  is all you Need" (https://arxiv.org/pdf/1706.03762.pdf).
class_docstr_tok: ["Multi-headed","attention","is","based","on","``","Attention","is","all","you","Need","''","(","https",":","\/\/arxiv.org\/pdf\/1706.03762.pdf",")","."]
docstr: The implementation of Multi-headed attention.
  Args:
    inputs = (v, k, q)
    q: (batch_size, seq_len_q, hidden_size)
    k: (batch_size, seq_len_k, hidden_size)
    v: (batch_size, seq_len_v, hidden_size)
    mask: (batch_size, seq_len_q, seq_len_k)
  
  Returns:
    output: (batch_size, seq_len_q, hidden_size)
    attention_weights: (batch_size, num_heads, seq_len_q, seq_len_k)
docstr_tok: ["The","implementation","of","Multi-headed","attention",".","Args",":","inputs","=","(","v",",","k",",","q",")","q",":","(","batch_size",",","seq_len_q",",","hidden_size",")","k",":","(","batch_size",",","seq_len_k",",","hidden_size",")","v",":","(","batch_size",",","seq_len_v",",","hidden_size",")","mask",":","(","batch_size",",","seq_len_q",",","seq_len_k",")","Returns",":","output",":","(","batch_size",",","seq_len_q",",","hidden_size",")","attention_weights",":","(","batch_size",",","num_heads",",","seq_len_q",",","seq_len_k",")"]
code: def call(self, inputs, training=None, mask=None):
      """
      The implementation of Multi-headed attention.
      Args:
        inputs = (v, k, q)
        q: (batch_size, seq_len_q, hidden_size)
        k: (batch_size, seq_len_k, hidden_size)
        v: (batch_size, seq_len_v, hidden_size)
        mask: (batch_size, seq_len_q, seq_len_k)
  
      Returns:
        output: (batch_size, seq_len_q, hidden_size)
        attention_weights: (batch_size, num_heads, seq_len_q, seq_len_k)
      """
      q, k, v = inputs
      batch_size = tf.shape(q)[0]
  
      q = self.wq(q)  # (batch_size, seq_len_q, hidden_size)
      k = self.wk(k)  # (batch_size, seq_len_k, hidden_size)
      v = self.wv(v)  # (batch_size, seq_len_v, hidden_size)
  
      q = self.split_heads(
          q, batch_size
      )  # (batch_size, num_heads, seq_len_q, depth)
      k = self.split_heads(
          k, batch_size
      )  # (batch_size, num_heads, seq_len_k, depth)
      v = self.split_heads(
          v, batch_size
      )  # (batch_size, num_heads, seq_len_v, depth)
  
      # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
      # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
      (
          scaled_attention,
          attention_weights,
      ) = self.scaled_dot_product_attention(q, k, v, mask)
  
      scaled_attention = tf.transpose(
          scaled_attention, perm=[0, 2, 1, 3]
      )  # (batch_size, seq_len_q, num_heads, depth)
  
      concat_attention = tf.reshape(
          scaled_attention, (batch_size, -1, self.hidden_size)
      )  # (batch_size, seq_len_q, hidden_size)
  
      output = self.dense(
          concat_attention
      )  # (batch_size, seq_len_q, hidden_size)
  
      return output, attention_weights
  
file_path: delta/delta/layers/sub_tf.py
imports: from absl import logging
  import delta.compat
  import numpy
  from delta.layers.base_layer import Layer
  from delta.layers.utils_tf import shape_list
label_desc: Use image node_id 3 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.call(inputs, training, mask) and returns: output,  attention_weights