import torch


def make_bow_matrix(word_ids, vocab_size, device=None):
    """
    Create a Bag of Words (BoW) matrix from word_ids.

    Args:
        word_ids (list of lists): a list of lists that contains the ids for each
            word in the sentence. Each id should be between [0, vocab_size).
            The number of lists should be equal to the batch_size, i.e.,
            len(word_ids) = batch_size
        vocab_size (int): the size of the words vocabulary
        device (int or torch.device): if not None, will sent the BoW matrix to
            this device. Default is None.

    Returns:
        torch.FloatTensor of shape (batch_size, vocab_size)
    """
    batch_size = len(word_ids)
    bow = torch.zeros(batch_size, vocab_size).to(device)
    for i, b_ids in enumerate(word_ids):
        for idx in b_ids:
            bow[i, idx] += 1.0
    return bow