import torch


def pack_batch_masked(padded_batch, mask):
    """
    Packs a padded batch of size BxLxD, according to mask of size BxL,
    to size SxD where S is the sum of unpadded lengths of all sequences in batch.
    :param padded_batch: BxLxD
    :param mask: BxL
    :return: packed batch (tensor of shape SxD), list with lengths of packed sequences (summing to S).
    """
    # Compute number of desired outputs (entity mentions) in each chunk
    mask_lengths = list(mask.long().sum(dim=1))
    # Flatten batch and chunk size dimensions, to create one long tensor (lstm_out and mask).
    padded_batch = padded_batch.contiguous().view(-1, padded_batch.shape[2])
    mask = mask.view(-1)
    # use only the rows in lstm_out according to the mask:
    mask_indices = mask.nonzero().squeeze()
    # The following check is necessary to avoid error, because Variables cannot be indexed with empty tensor :(
    if len(mask_indices) == 0:
        output = torch.FloatTensor()
        if padded_batch.is_cuda: output = output.cuda()
    else:
        output = padded_batch[mask_indices]
    return output, mask_lengths