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 scaled attention.
  Args:
    v: (batch_size, seq_len_v, hidden_size)
    k: (batch_size, seq_len_k, hidden_size)
    q: (batch_size, seq_len_q, 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","scaled","attention",".","Args",":","v",":","(","batch_size",",","seq_len_v",",","hidden_size",")","k",":","(","batch_size",",","seq_len_k",",","hidden_size",")","q",":","(","batch_size",",","seq_len_q",",","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 scaled_dot_product_attention(q, k, v, mask):
      """
      The implementation of scaled attention.
      Args:
        v: (batch_size, seq_len_v, hidden_size)
        k: (batch_size, seq_len_k, hidden_size)
        q: (batch_size, seq_len_q, 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)
      """
  
      matmul_qk = tf.matmul(
          q, k, transpose_b=True
      )  # (batch_size, seq_len_q, seq_len_k)
  
      # Scaled
      dk = tf.cast(tf.shape(k)[-1], tf.float32)
      scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
  
      # Masked
      if mask is not None:
          scaled_attention_logits += mask * -1e9
  
      # Normalized
      attention_weights = tf.nn.softmax(
          scaled_attention_logits, axis=-1
      )  # (batch_size, seq_len_q, seq_len_k)
  
      # Weighted sum
      output = tf.matmul(
          attention_weights, v
      )  # (batch_size, seq_len_q, depth_v)
  
      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 4 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.scaled_dot_product_attention(q, k, v, mask) and returns: output,  attention_weights