in muse/sampling.py [0:0]
def get_mask_chedule(method, **schedule_kwargs):
if method == "cosine":
return cosine_schedule
elif method == "linear":
return linear_schedule
elif "pow" in method:
return partial(pow, method=method)
elif method == "sigmoid":
return partial(sigmoid_schedule, **schedule_kwargs)
else:
raise ValueError("Unknown schedule method: {}".format(method))