import torch


def dot_attention(query, key, padding_mask):
    """
    Args:
        query(torch.Tensor): shape (batch_size, num_queries, embed_dim)
        key(torch.Tensor): shape (batch_size, num_keys, embed_dim)
        padding_mask(torch.Tensor): shape (batch_size, num_keys)

    Returns:
        attn_output(torch.Tensor): shape (batch_size, num_queries, embed_dim)
        attn_weight(torch.Tensor): shape (batch_size, num_queries, num_keys)
    """
    batch_size, num_queries, num_keys, embed_dim = query.shape[0], query.shape[1], key.shape[1], key.shape[2]
    attn_weight = torch.matmul(query, torch.transpose(key, 1, 2))
    assert attn_weight.shape == torch.Size([batch_size, num_queries, num_keys])
    # mask attention weights
    attn_weight = attn_weight.masked_fill_(padding_mask.unsqueeze(1), float('-inf'))
    attn_weight = torch.softmax(attn_weight, dim=-1)
    # attention pool
    attn_output = torch.matmul(attn_weight, key)
    assert attn_output.shape == torch.Size([batch_size, num_queries, embed_dim])
    return attn_output, attn_weight