@patrickvonplaten in the documentation of generate 1 function the description of prefix_allowed_tokens_fn argument says

If provided, this function constraints the beam search to allowed tokens only at each step.

Can  this function be applied with sampling as well?