class_docstr: Multi-Head Attention
class_docstr_tok: ["Multi-Head","Attention"]
code: def call(
      self,
      query: tf.Tensor,
      key: tf.Tensor,
      value: tf.Tensor,
      mask: tf.Tensor = None,
      **kwargs: Any,
  ) -> tf.Tensor:
      batch_size = query.shape[0]
  
      # linear projections of Q, K, V
      query, key, value = [
          tf.transpose(
              tf.reshape(
                  linear(x, **kwargs),
                  shape=[batch_size, -1, self.num_heads, self.d_k],
              ),
              perm=[0, 2, 1, 3],
          )
          for linear, x in zip(self.linear_layers, (query, key, value))
      ]
  
      # apply attention on all the projected vectors in batch
      x, attn = scaled_dot_product_attention(
          query, key, value, mask=mask
      )
  
      # Concat attention heads
      x = tf.transpose(x, perm=[0, 2, 1, 3])
      x = tf.reshape(
          x, shape=[batch_size, -1, self.num_heads * self.d_k]
      )
  
      return self.output_linear(x, **kwargs)
  
imports: import math
  from typing import Any, Callable, Optional, Tuple
  import tensorflow
  from tensorflow.keras import layers
  from doctr.utils.repr import NestedObject
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.call(query, key, value, mask) and returns: self