def METHOD_NAME(self, query, keys):
    """
    :param query: if batch_first: (b x t_q x n) else: (t_q x b x n)
    :param keys: if batch_first: (b x t_k x n) else (t_k x b x n)
    :returns: (context, scores_normalized)
    context: if batch_first: (b x t_q x n) else (t_q x b x n)
    scores_normalized: if batch_first (b x t_q x t_k) else (t_q x b x t_k)
    """
    # first dim of keys and query has to be 'batch', it's needed for bmm
    if not self.batch_first:
        keys = keys.transpose(0, 1)
        if query.dim() == 3:
            query = query.transpose(0, 1)
    if query.dim() == 2:
        single_query = True
        query = query.unsqueeze(1)
    else:
        single_query = False
    b = query.size(0)
    t_k = keys.size(1)
    t_q = query.size(1)
    # FC layers to transform query and key
    processed_query = self.linear_q(query)
    processed_key = self.linear_k(keys)
    # scores: (b x t_q x t_k)
    scores = self.calc_score(processed_query, processed_key)
    if self.mask is not None:
        mask = self.mask.unsqueeze(1).expand(b, t_q, t_k)
        # I can't use -INF because of overflow check in pytorch
        scores.masked_fill_(mask, -65504.0)
    # Normalize the scores, softmax over t_k
    scores_normalized = F.softmax(scores, dim=-1)
    # Calculate the weighted average of the attention inputs according to
    # the scores
    # context: (b x t_q x n)
    context = torch.bmm(scores_normalized, keys)
    if single_query:
        context = context.squeeze(1)
        scores_normalized = scores_normalized.squeeze(1)
    elif not self.batch_first:
        context = context.transpose(0, 1)
        scores_normalized = scores_normalized.transpose(0, 1)
    return context, scores_normalized