Skip to main content

accept_sampled

Function accept_sampled 

Source
pub fn accept_sampled(
    probs: &Array,
    draft: &Array,
    draft_probs: Option<&Array>,
    eos: &Array,
    key: &Array,
) -> Array
Expand description

Lossless sampled acceptance (Leviathan et al., Chen et al.): probs [N + 1, vocab] is the target’s sampling distribution at each position (after temperature and filters), draft_probs [N, vocab] the drafter’s (None: deterministic drafts, a point mass on each draft token). Draft token i is accepted when u_i < min(1, p_i(d_i) / q_i(d_i)); after the first rejection at position a, next is sampled from norm(max(p_a − q_a, 0)) (p_a itself when that has no mass), and after a fully accepted draft from p_N. The output’s distribution is the target’s.

key (a threefry key, crate::random) is split into two: the first draws the uniforms u ([N]), the second next; the same key gives the same result.