class_docstr: Multi-headed attention from "Attention Is All You Need
docstr: Input shape: Batch x Time x Channel
docstr_tok: ["Input","shape",":","Batch","x","Time","x","Channel"]
code: def call(
      self,
      hidden_states: tf.Tensor,
      key_value_states: Optional[tf.Tensor] = None,
      past_key_value: Optional[Tuple[Tuple[tf.Tensor]]] = None,
      attention_mask: Optional[tf.Tensor] = None,
      layer_head_mask: Optional[tf.Tensor] = None,
      training: Optional[bool] = False,
  ) -> Tuple[tf.Tensor, Optional[tf.Tensor]]:
      """Input shape: Batch x Time x Channel"""
  
      # if key_value_states are provided this layer is used as a cross-attention layer
      # for the decoder
      is_cross_attention = key_value_states is not None
      bsz, tgt_len, embed_dim = shape_list(hidden_states)
  
      # get query proj
      query_states = self.q_proj(hidden_states) * self.scaling
      # get key, value proj
      if is_cross_attention and past_key_value is not None:
          # reuse k,v, cross_attentions
          key_states = past_key_value[0]
          value_states = past_key_value[1]
      elif is_cross_attention:
          # cross_attentions
          key_states = self._shape(
              self.k_proj(key_value_states), -1, bsz
          )
          value_states = self._shape(
              self.v_proj(key_value_states), -1, bsz
          )
      elif past_key_value is not None:
          # reuse k, v, self_attention
          key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
          value_states = self._shape(
              self.v_proj(hidden_states), -1, bsz
          )
          key_states = tf.concat(
              [past_key_value[0], key_states], axis=2
          )
          value_states = tf.concat(
              [past_key_value[1], value_states], axis=2
          )
      else:
          # self_attention
          key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
          value_states = self._shape(
              self.v_proj(hidden_states), -1, bsz
          )
  
      if self.is_decoder:
          # if cross_attention save Tuple(tf.Tensor, tf.Tensor) of all cross attention key/value_states.
          # Further calls to cross_attention layer can then reuse all cross-attention
          # key/value_states (first "if" case)
          # if uni-directional self-attention (decoder) save Tuple(tf.Tensor, tf.Tensor) of
          # all previous decoder key/value_states. Further calls to uni-directional self-attention
          # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
          # if encoder bi-directional self-attention `past_key_value` is always `None`
          past_key_value = (key_states, value_states)
  
      proj_shape = (bsz * self.num_heads, -1, self.head_dim)
      query_states = tf.reshape(
          self._shape(query_states, tgt_len, bsz), proj_shape
      )
      key_states = tf.reshape(key_states, proj_shape)
      value_states = tf.reshape(value_states, proj_shape)
  
      src_len = shape_list(key_states)[1]
      attn_weights = tf.matmul(
          query_states, key_states, transpose_b=True
      )
  
      tf.debugging.assert_equal(
          shape_list(attn_weights),
          [bsz * self.num_heads, tgt_len, src_len],
          message=(
              f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
              f" {shape_list(attn_weights)}"
          ),
      )
  
      if attention_mask is not None:
          tf.debugging.assert_equal(
              shape_list(attention_mask),
              [bsz, 1, tgt_len, src_len],
              message=(
                  f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
                  f" {shape_list(attention_mask)}"
              ),
          )
  
          attention_mask = tf.cast(
              attention_mask, dtype=attn_weights.dtype
          )
          attn_weights = (
              tf.reshape(
                  attn_weights, (bsz, self.num_heads, tgt_len, src_len)
              )
              + attention_mask
          )
          attn_weights = tf.reshape(
              attn_weights, (bsz * self.num_heads, tgt_len, src_len)
          )
  
      attn_weights = stable_softmax(attn_weights, axis=-1)
  
      if layer_head_mask is not None:
          tf.debugging.assert_equal(
              shape_list(layer_head_mask),
              [self.num_heads],
              message=(
                  f"Head mask for a single layer should be of size {(self.num_heads)}, but is"
                  f" {shape_list(layer_head_mask)}"
              ),
          )
  
          attn_weights = tf.reshape(
              layer_head_mask, (1, -1, 1, 1)
          ) * tf.reshape(
              attn_weights, (bsz, self.num_heads, tgt_len, src_len)
          )
          attn_weights = tf.reshape(
              attn_weights, (bsz * self.num_heads, tgt_len, src_len)
          )
  
      attn_probs = self.dropout(attn_weights, training=training)
      attn_output = tf.matmul(attn_probs, value_states)
  
      tf.debugging.assert_equal(
          shape_list(attn_output),
          [bsz * self.num_heads, tgt_len, self.head_dim],
          message=(
              f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
              f" {shape_list(attn_output)}"
          ),
      )
  
      attn_output = tf.transpose(
          tf.reshape(
              attn_output, (bsz, self.num_heads, tgt_len, self.head_dim)
          ),
          (0, 2, 1, 3),
      )
      attn_output = tf.reshape(attn_output, (bsz, tgt_len, embed_dim))
  
      attn_output = self.out_proj(attn_output)
      attn_weights: tf.Tensor = tf.reshape(
          attn_weights, (bsz, self.num_heads, tgt_len, src_len)
      )
  
      return attn_output, attn_weights, past_key_value
  
imports: import math
  import random
  from typing import Any, Optional, Tuple, Union
  import numpy
  import tensorflow
  from ...activations_tf import get_tf_activation
  from ...file_utils import DUMMY_INPUTS, add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
  from ...modeling_tf_outputs import TFBaseModelOutputWithPastAndCrossAttentions, TFCausalLMOutputWithCrossAttentions
  from ...modeling_tf_utils import TFCausalLanguageModelingLoss, TFModelInputType, TFPreTrainedModel, TFSharedEmbeddings, get_initializer, keras_serializable, unpack_inputs
  from ...tf_utils import shape_list, stable_softmax
  from ...utils import logging
  from .configuration_xglm import XGLMConfig
label_desc: Use image node_id 3 for calling the TFXGLMAttention obj's underlying member method code with example usage: obj.call(hidden_states, key_value_states, past_key_value, attention_mask, layer_head_mask, training) and returns: attn_output,  attn_weights,  past_key_value