class_docstr: - Q (query), K (key) and V (value) are split into multiple heads (num_heads)
  - each tuple (q, k, v) are fed to scaled_dot_product_attention
  - all attention outputs are concatenated
docstr: Model forward pass.
  
  Args:
      queries (tf.Tensor): Tensor of queries.
      keys (tf.Tensor): Tensor of keys
  
  Returns:
      tf.Tensor: Output tensor.
docstr_tok: ["Model","forward","pass",".","Args",":","queries","(","tf.Tensor",")",":","Tensor","of","queries",".","keys","(","tf.Tensor",")",":","Tensor","of","keys","Returns",":","tf.Tensor",":","Output","tensor","."]
code: def call(self, queries, keys):
      """Model forward pass.
  
      Args:
          queries (tf.Tensor): Tensor of queries.
          keys (tf.Tensor): Tensor of keys
  
      Returns:
          tf.Tensor: Output tensor.
      """
  
      # Linear projections
      Q = self.Q(queries)  # (N, T_q, C)
      K = self.K(keys)  # (N, T_k, C)
      V = self.V(keys)  # (N, T_k, C)
  
      # --- MULTI HEAD ---
      # Split and concat, Q_, K_ and V_ are all (h*N, T_q, C/h)
      Q_ = tf.concat(tf.split(Q, self.num_heads, axis=2), axis=0)
      K_ = tf.concat(tf.split(K, self.num_heads, axis=2), axis=0)
      V_ = tf.concat(tf.split(V, self.num_heads, axis=2), axis=0)
  
      # --- SCALED DOT PRODUCT ---
      # Multiplication
      outputs = tf.matmul(
          Q_, tf.transpose(K_, [0, 2, 1])
      )  # (h*N, T_q, T_k)
  
      # Scale
      outputs = outputs / (K_.get_shape().as_list()[-1] ** 0.5)
  
      # Key Masking
      key_masks = tf.sign(
          tf.abs(tf.reduce_sum(keys, axis=-1))
      )  # (N, T_k)
      key_masks = tf.tile(key_masks, [self.num_heads, 1])  # (h*N, T_k)
      key_masks = tf.tile(
          tf.expand_dims(key_masks, 1), [1, tf.shape(queries)[1], 1]
      )  # (h*N, T_q, T_k)
  
      paddings = tf.ones_like(outputs) * (-(2**32) + 1)
      # outputs, (h*N, T_q, T_k)
      outputs = tf.where(tf.equal(key_masks, 0), paddings, outputs)
  
      # Future blinding (Causality)
      diag_vals = tf.ones_like(outputs[0, :, :])  # (T_q, T_k)
      tril = tf.linalg.LinearOperatorLowerTriangular(
          diag_vals
      ).to_dense()  # (T_q, T_k)
      masks = tf.tile(
          tf.expand_dims(tril, 0), [tf.shape(outputs)[0], 1, 1]
      )  # (h*N, T_q, T_k)
  
      paddings = tf.ones_like(masks) * (-(2**32) + 1)
      # outputs, (h*N, T_q, T_k)
      outputs = tf.where(tf.equal(masks, 0), paddings, outputs)
  
      # Activation
      outputs = tf.nn.softmax(outputs)  # (h*N, T_q, T_k)
  
      # Query Masking, query_masks (N, T_q)
      query_masks = tf.sign(tf.abs(tf.reduce_sum(queries, axis=-1)))
      query_masks = tf.tile(
          query_masks, [self.num_heads, 1]
      )  # (h*N, T_q)
      query_masks = tf.tile(
          tf.expand_dims(query_masks, -1), [1, 1, tf.shape(keys)[1]]
      )  # (h*N, T_q, T_k)
      outputs *= query_masks  # broadcasting. (N, T_q, C)
  
      # Dropouts
      outputs = self.dropout(outputs)
  
      # Weighted sum
      outputs = tf.matmul(outputs, V_)  # ( h*N, T_q, C/h)
  
      # --- MULTI HEAD ---
      # concat heads
      outputs = tf.concat(
          tf.split(outputs, self.num_heads, axis=0), axis=2
      )  # (N, T_q, C)
  
      # Residual connection
      outputs += queries
  
      return outputs
  
imports: import random
  import numpy
  from tqdm import tqdm
  import tensorflow
  from recommenders.utils.timer import Timer
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.call(queries, keys) and returns: outputs