src/transformers/models/layoutlmv3/modeling_tf_layoutlmv3.py/TFLayoutLMv3SelfAttention/call
class TFLayoutLMv3SelfAttention:    def call(
        self,
        hidden_states: tf.Tensor,
        attention_mask: Optional[tf.Tensor],
        head_mask: Optional[tf.Tensor],
        output_attentions: bool,
        rel_pos: Optional[tf.Tensor] = None,
        rel_2d_pos: Optional[tf.Tensor] = None,
        training: bool = False,
    ) -> Union[Tuple[tf.Tensor], Tuple[tf.Tensor, tf.Tensor]]:
        key_layer = self.transpose_for_scores(self.key(hidden_states))
        value_layer = self.transpose_for_scores(self.value(hidden_states))
        query_layer = self.transpose_for_scores(self.query(hidden_states))

        # Take the dot product between "query" and "key" to get the raw attention scores.
        normalised_query_layer = query_layer / self.attention_score_normaliser
        transposed_key_layer = tf.transpose(
            key_layer, perm=[0, 1, 3, 2]
        )  # batch_size, num_heads, attention_head_size, seq_length
        attention_scores = tf.matmul(normalised_query_layer, transposed_key_layer)

        if self.has_relative_attention_bias and self.has_spatial_attention_bias:
            attention_scores += (rel_pos + rel_2d_pos) / self.attention_score_normaliser
        elif self.has_relative_attention_bias:
            attention_scores += rel_pos / self.attention_score_normaliser

        if attention_mask is not None:
            # Apply the attention mask (is precomputed for all layers in TFLayoutLMv3Model call() function)
            attention_scores += attention_mask

        # Normalize the attention scores to probabilities.
        # Use the trick of CogView paper to stabilize training.
        attention_probs = self.cogview_attention(attention_scores)

        attention_probs = self.dropout(attention_probs, training=training)

        # Mask heads if we want to.
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = tf.matmul(attention_probs, value_layer)
        context_layer = tf.transpose(
            context_layer, perm=[0, 2, 1, 3]
        )  # batch_size, seq_length, num_heads, attention_head_size
        shape = tf.shape(context_layer)
        context_layer = tf.reshape(
            context_layer, (shape[0], shape[1], self.all_head_size)
        )  # batch_size, seq_length, num_heads * attention_head_size

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs