src/transformers/models/mctct/modeling_mctct.py/MCTCTSelfAttention/forward
class MCTCTSelfAttention:    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        output_attentions=False,
    ):
        mixed_query_layer = self.query(hidden_states)
        mixed_query_layer = mixed_query_layer / math.sqrt(self.attention_head_size)

        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(mixed_query_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        # relative key position embeddings
        positional_embedding = self.distance_embedding.weight
        relative_position_scores = torch.einsum("lh, bche -> bcle", positional_embedding, query_layer.transpose(2, 3))

        relative_position_scores = self.relative_position_embedding_rotate(relative_position_scores)
        attention_scores = attention_scores + relative_position_scores

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

        # Normalize the attention scores to probabilities.
        attention_probs = nn.functional.softmax(attention_scores, dim=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

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

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).flatten(start_dim=-2)

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

        return outputs