def softmax_attention(self, query, key, padding_mask, causal_mask):
        "Standard softmax self-attention, as in the original Transformer paper"
        seq_len = key.size(2)
        # (sequence_length X sequence_length)
        bias = self.rel_pos_bias(seq_len)
        if seq_len != query.size(2):
            if query.size(2) != 1:
                raise ValueError("Size mismatch between Q and K in softmax attention")
            # (1 X sequence_length)
            bias = bias[-1:]

        # scaled attention
        query = query * self.scaling

        # (batch_size x number of chunks x chunk_size x chunk_size) if chunking
        # (batch_size x 1 x sequence_length x sequence_length) otherwise
        qk = torch.matmul(query, key.transpose(2, 3)) + bias

        # apply causal mask (presumed to be 1/0 for not masked / masked)
        # additive, but convert to 0/-inf (which is not explicitly in the Mega source code)
        if causal_mask is not None:
            additive_causal_mask = torch.zeros_like(causal_mask, dtype=qk.dtype)
            additive_causal_mask = additive_causal_mask.masked_fill((1 - causal_mask).bool(), float("-inf"))
            qk = qk + additive_causal_mask

        if padding_mask is not None:
            # 1 for tokens which are *not masked*
            # 0 for tokens which are *masked*
            # replace masked tokens with -inf to make softmax ignore them
            # need to invert the padding mask to match what mega original did
            padding_mask = 1 - padding_mask
            padding_mask_all = padding_mask.all(dim=-1, keepdim=True)
            padding_mask = torch.logical_and(padding_mask, ~padding_mask_all)
            qk = qk.masked_fill(padding_mask.unsqueeze(2).to(torch.bool), float("-inf"))

        attn_weights = self.softmax(qk).type_as(qk)
        return attn_weights
Standard softmax self-attention, as in the original Transformer paper
dpkd/transformers/src/transformers/models/mega/modeling_mega.py