def forward(self, hidden_states, attention_mask, head_mask=None, output_attentions=False, keep_multihead_output=False):
    '\n        Input:\n            hidden_states[batch, seq_len, hidden_dim]\n            attention_mask[batch,  1, 1, seq_len]\n        Output:\n            context_states[batch, seq_len, hidden_dim]\n            attention_probs[seq_len, hidden_dim]\n        '
    query = self.query(hidden_states)
    key = self.key(hidden_states)
    value = self.value(hidden_states)
    multi_query = self.transpose_(query)
    multi_key = self.transpose_(key)
    multi_value = self.transpose_(value)
    attention_scores = torch.matmul(multi_query, multi_key.transpose((- 1), (- 2)))
    attention_scores = (attention_scores / self.scale_factor)
    attention_scores = (attention_scores + attention_mask)
    attention_probs = F.softmax(attention_scores, dim=(- 1))
    attention_probs = self.dropout(attention_probs)
    if (head_mask is not None):
        attention_probs = (attention_probs * head_mask)
    context_states = torch.matmul(attention_probs, multi_value)
    if keep_multihead_output:
        self.multihead_output = context_states
        self.multihead_output.retain_grad()
    context_states = context_states.permute(0, 2, 1, 3)
    context_states = context_states.contiguous().view((context_states.size()[:(- 2)] + ((- 1),)))
    if output_attentions:
        return (context_states, attention_probs)
    return context_states
Input:
    hidden_states[batch, seq_len, hidden_dim]
    attention_mask[batch,  1, 1, seq_len]
Output:
    context_states[batch, seq_len, hidden_dim]
    attention_probs[seq_len, hidden_dim]
def forward(self, hidden_states, attention_mask, head_mask=None, output_attentions=False, keep_multihead_output=False):
    '\n        Input:\n            hidden_states[batch, seq_len, hidden_dim]\n            attention_mask[batch,  1, 1, seq_len]\n        Output:\n            context_states[batch, seq_len, hidden_dim]\n            attention_probs[seq_len, hidden_dim]\n        '
    query = self.query(hidden_states)
    key = self.key(hidden_states)
    value = self.value(hidden_states)
    multi_query = self.transpose_(query)
    multi_key = self.transpose_(key)
    multi_value = self.transpose_(value)
    attention_scores = torch.matmul(multi_query, multi_key.transpose((- 1), (- 2)))
    attention_scores = (attention_scores / self.scale_factor)
    attention_scores = (attention_scores + attention_mask)
    attention_probs = F.softmax(attention_scores, dim=(- 1))
    attention_probs = self.dropout(attention_probs)
    if (head_mask is not None):
        attention_probs = (attention_probs * head_mask)
    context_states = torch.matmul(attention_probs, multi_value)
    if keep_multihead_output:
        self.multihead_output = context_states
        self.multihead_output.retain_grad()
    context_states = context_states.permute(0, 2, 1, 3)
    context_states = context_states.contiguous().view((context_states.size()[:(- 2)] + ((- 1),)))
    if output_attentions:
        return (context_states, attention_probs)
    return context_states