@triton.jit
def _drop_and_scale(SEEDS, row, p, offsets, x):
    seed = SEEDS + row
    random = tl.rand(seed.to(tl.int32), offsets)
    x_keep = random > p
    zero = 0.0
    zero = zero.to(x.dtype)
    return tl.where(x_keep, (x / (1 - p)).to(x.dtype), zero)
