import torch


def masked_log_softmax(vector, mask):
    """
    ``torch.nn.functional.log_softmax(vector)`` does not work if some elements of ``vector`` should be
    masked.  This performs a log_softmax on just the non-masked portions of ``vector``.  Passing
    ``None`` in for the mask is also acceptable; you'll just get a regular log_softmax.

    We assume that both ``vector`` and ``mask`` (if given) have shape ``(batch_size, vector_dim)``.

    In the case that the input vector is completely masked, this function returns an array
    of ``0.0``.  You should be masking the result of whatever computation comes out of this in that
    case, anyway, so it shouldn't matter.
    """
    if mask is not None:
        vector = vector + mask.log()
    return torch.nn.functional.log_softmax(vector)