def METHOD_NAME():
    n_h, emb_size = 4, 12
    n_k, n_q, b = 13, 14, 15
    mha = MultiHeadAttention(emb_size, n_h)
    query = torch.ones((b, n_q, emb_size))
    key = torch.ones((b, n_k, emb_size))
    value = torch.ones((b, n_k, emb_size))
    output, attention = mha.forward(query, key, value, n_q, n_k)
    assert output.shape == (b, n_q, emb_size)
    assert attention.shape == (b, n_h, n_q, n_k)