def dropout_mask(x, sz, p: float):
    """
    Return a dropout mask of the same type as `x`, size `sz`, with probability `p` to cancel an element.
    """
    return x.new(*sz).bernoulli_(1 - p).div_(1 - p)