def dropout_mask(x, size, prob):
    """
    We pass size in so that we get broadcasting along the sequence dimension
    in RNNDropout.
    """
    return x.new(*size).bernoulli_(1 - prob).div_(1 - prob)
