import itertools,math
def apply(values,p,mask,training=True):
    if not math.isfinite(p) or not 0<=p<=1:
        raise ValueError('probability must be in [0,1]')
    if len(values)!=len(mask) or any(m not in (0,1) for m in mask):
        raise ValueError('binary mask of matching length required')
    if not all(math.isfinite(x) for x in values):
        raise ValueError('finite inputs required')
    if not training or p==0:return list(values)
    if p==1:return [0.0]*len(values)
    return [x*m/(1-p) for x,m in zip(values,mask)]
x=[1.,2.,3.,4.]
y=apply(x,.5,[1,0,1,0])
all_y=[apply(x,.5,m) for m in itertools.product((0,1),repeat=4)]
mean=[sum(v[i] for v in all_y)/len(all_y) for i in range(4)]
print('one mask',y,'mean over 16 masks',mean)
assert y==[2.,0.,6.,0.] and mean==x
assert apply(x,.5,[0]*4,training=False)==x
assert apply(x,1,[1]*4)==[0.]*4