easydel.inference.sampling_funcs

easydel.inference.sampling_funcs.dynamic_sample_tokens(tokens: Array, length: Array, logits: Array, top_p: Array, temperature: Array, random_sampling: Array, presence_penalty: Array, frequency_penalty: Array, repetition_penalty: Array, rngs: Array) → Array

Vectorized version of dynamic_sample_tokens. Takes similar arguments as dynamic_sample_tokens but with additional array axes over which dynamic_sample_tokens is mapped.

easydel.inference.sampling_funcs.sample_top_p_efficient(logits: Array, top_p: Array, temperature: Array, rng: PRNGKey, top_k_for_computation: int = 64) → Array

easydel.inference.sampling_funcs.vmaped_sample_top_p_efficient(logits: Array, top_p: Array, temperature: Array, rng: PRNGKey, top_k_for_computation: int = 64) → Array

Vectorized version of sample_top_p_efficient. Takes similar arguments as sample_top_p_efficient but with additional array axes over which sample_top_p_efficient is mapped.