import torch


def length_penalty(sequence_lengths, penalty_factor):
    """ 
    Calculate the length penalty according to 
    https://arxiv.org/abs/1609.08144
    lp(Y) =(5 +|Y|)^α / (5 + 1)^α

    Input:
    sequence_lenthgs: the sequences length of all hypotheses of size [batch size x beam size x vocab size]
    penalty_factor: A scalar that weights the length penalty.
    
    Returns:
    The length penalty factor, a tensor fo shape [batch size x beam size].
    """
    return torch.div((5. + sequence_lengths)**penalty_factor, (5. + 1.)
                **penalty_factor)