import torch


def dot_product(query, key):
    """Luong dot-product attention.

    Args:
        query (Tensor[bsz, tgt_len, hsz])
        key (Tensor[bsz, src_len, hsz])

    Returns:
        Tensor[bsz, tgt_len, src_len]
    """
    return torch.softmax(query @ key.transpose(1, 2), dim=-1)