def forward(self, hidden_states, attention_mask=None, layer_past=None, head_mask=None, use_cache=False, output_attentions=False):
    query = self.q_proj(hidden_states)
    key = self.k_proj(hidden_states)
    value = self.v_proj(hidden_states)
    query = self._split_heads(query, self.num_heads, self.head_dim)
    key = self._split_heads(key, self.num_heads, self.head_dim)
    value = self._split_heads(value, self.num_heads, self.head_dim)
    if (layer_past is not None):
        past_key = layer_past[0]
        past_value = layer_past[1]
        key = torch.cat((past_key, key), dim=(- 2))
        value = torch.cat((past_value, value), dim=(- 2))
    if (use_cache is True):
        present = (key, value)
    else:
        present = None
    (query_length, key_length) = (query.size((- 2)), key.size((- 2)))
    causal_mask = self.bias[:, :, (key_length - query_length):key_length, :key_length]
    (attn_output, attn_weights) = self._attn(query, key, value, causal_mask, self.masked_bias, self.attn_dropout, attention_mask, head_mask)
    attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
    attn_output = self.out_proj(attn_output)
    attn_output = self.resid_dropout(attn_output)
    outputs = (attn_output, present)
    if output_attentions:
        outputs += (attn_weights,)
    return outputs
def forward(self, hidden_states, attention_mask=None, layer_past=None, head_mask=None, use_cache=False, output_attentions=False):
    query = self.q_proj(hidden_states)
    key = self.k_proj(hidden_states)
    value = self.v_proj(hidden_states)
    query = self._split_heads(query, self.num_heads, self.head_dim)
    key = self._split_heads(key, self.num_heads, self.head_dim)
    value = self._split_heads(value, self.num_heads, self.head_dim)
    if (layer_past is not None):
        past_key = layer_past[0]
        past_value = layer_past[1]
        key = torch.cat((past_key, key), dim=(- 2))
        value = torch.cat((past_value, value), dim=(- 2))
    if (use_cache is True):
        present = (key, value)
    else:
        present = None
    (query_length, key_length) = (query.size((- 2)), key.size((- 2)))
    causal_mask = self.bias[:, :, (key_length - query_length):key_length, :key_length]
    (attn_output, attn_weights) = self._attn(query, key, value, causal_mask, self.masked_bias, self.attn_dropout, attention_mask, head_mask)
    attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
    attn_output = self.out_proj(attn_output)
    attn_output = self.resid_dropout(attn_output)
    outputs = (attn_output, present)
    if output_attentions:
        outputs += (attn_weights,)
    return outputs