torchtune.generation

generate	Generates tokens from a model conditioned on a prompt, and also returns logits for the generations.
sample	Generic sample from a probability distribution.
get_causal_mask_from_padding_mask	Converts a padding mask of shape [bsz, seq_len] to a [bsz, seq_len, seq_len] causal attention mask suitable for consumption by scaled_dot_product_attention~().
get_position_ids_from_padding_mask	Calculates position ids given a padding mask which right-shifts position ids to start from the first valid token.