import torch


def self_attention(query: torch.Tensor,
                   key: torch.Tensor,
                   value: torch.Tensor
                   ) -> torch.Tensor:
    """ A naive O(n^2) complexity implementation of self-attention

    Args:
        query: query of shape BxHxNxD
        key: key of shape BxHxN'xD
        value: value of shape BxHxN'xD

        where B is the batch size, H is the number of heads, N is the sequence length of the query,
        N' is the sequence length of the key and value (can be N), and D is the feature size.

    Returns: output of self-attention of shape BxHxNxD

    """

    score = torch.einsum("bhqd,bhkd->bhqk", query, key)
    score = score.softmax(dim=-1)
    return torch.einsum("bhqv,bhvd->bhqd", score, value)