use superfluid_abi::SamplingParams;
pub fn argmax(row: &[f32]) -> u32 {
if row.first().is_none_or(|v| v.is_nan()) {
return 0;
}
let best = max_of(row);
const BLOCK: usize = 64;
for (b, block) in row.chunks(BLOCK).enumerate() {
if max_of(block) == best {
let at = block.iter().position(|&v| v == best).unwrap_or(0);
return (b * BLOCK + at) as u32;
}
}
0
}
fn max_of(row: &[f32]) -> f32 {
#[cfg(target_arch = "aarch64")]
{
use std::arch::aarch64::{vdupq_n_f32, vld1q_f32, vmaxnmq_f32, vmaxnmvq_f32};
let (chunks, rest) = row.as_chunks::<16>();
let mut best = unsafe {
let low = vdupq_n_f32(f32::NEG_INFINITY);
let (mut a, mut b, mut c, mut d) = (low, low, low, low);
for chunk in chunks {
let p = chunk.as_ptr();
a = vmaxnmq_f32(a, vld1q_f32(p));
b = vmaxnmq_f32(b, vld1q_f32(p.add(4)));
c = vmaxnmq_f32(c, vld1q_f32(p.add(8)));
d = vmaxnmq_f32(d, vld1q_f32(p.add(12)));
}
vmaxnmvq_f32(vmaxnmq_f32(vmaxnmq_f32(a, b), vmaxnmq_f32(c, d)))
};
for &v in rest {
if v > best {
best = v;
}
}
best
}
#[cfg(not(target_arch = "aarch64"))]
{
let mut best = f32::NEG_INFINITY;
for &v in row {
if v > best {
best = v;
}
}
best
}
}
pub fn penalizes(p: &SamplingParams) -> bool {
let repeat = if p.repeat_penalty == 0.0 { 1.0 } else { p.repeat_penalty };
repeat != 1.0 || p.presence_penalty != 0.0 || p.freq_penalty != 0.0
}
pub fn apply_penalties(row: &mut [f32], recent: &[u32], p: &SamplingParams, exempt: &std::collections::HashSet<u32>) {
if !penalizes(p) {
return;
}
let repeat = if p.repeat_penalty == 0.0 { 1.0 } else { p.repeat_penalty };
let mut counts: std::collections::HashMap<u32, u32> = std::collections::HashMap::new();
for &t in recent {
if !exempt.contains(&t) {
*counts.entry(t).or_insert(0) += 1;
}
}
for (t, n) in counts {
let Some(l) = row.get_mut(t as usize) else { continue };
if repeat != 1.0 {
*l = if *l > 0.0 { *l / repeat } else { *l * repeat };
}
*l -= p.presence_penalty + p.freq_penalty * n as f32;
}
}
pub fn apply_logit_bias(row: &mut [f32], bias: &[(u32, f32)]) {
for &(t, b) in bias {
if let Some(l) = row.get_mut(t as usize) {
*l += b;
}
}
}
pub fn uniform(position: u64) -> f64 {
let mut z = position.wrapping_add(0x9E37_79B9_7F4A_7C15);
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
z ^= z >> 31;
(z >> 11) as f64 / (1u64 << 53) as f64
}
fn by_prob_desc(a: &(u32, f32), b: &(u32, f32)) -> std::cmp::Ordering {
b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal).then(a.0.cmp(&b.0))
}
fn sort_head(cands: &mut [(u32, f32)], head: usize) {
if head < cands.len() {
cands.select_nth_unstable_by(head - 1, by_prob_desc);
}
cands[..head].sort_unstable_by(by_prob_desc);
}
const INITIAL_HEAD: usize = 128;
fn grow(head: usize, n: usize) -> usize {
head.saturating_mul(8).clamp(head + 1, n)
}
const DEAD_BELOW: f32 = -36.736_8;
#[inline(always)]
fn exp_neg(x: f32) -> f32 {
const LOG2E: f32 = std::f32::consts::LOG2_E;
const LN2_HI: f32 = 0.693_359_4;
const LN2_LO: f32 = -2.121_944_4e-4;
let live = x > DEAD_BELOW;
let x = if live { x } else { 0.0 };
let n = (x * LOG2E - 0.5) as i32;
let nf = n as f32;
let r = (x - nf * LN2_HI) - nf * LN2_LO;
let poly = ((((1.987_569_1e-4 * r + 1.398_2e-3) * r + 8.333_452e-3) * r + 4.166_579_6e-2) * r + 1.666_666_5e-1) * r + 0.5;
let y = poly * (r * r) + r + 1.0;
let scale = f32::from_bits(((n + 127) as u32) << 23);
if live {
y * scale
} else {
0.0
}
}
fn softmax_weights(row: &mut [f32], inv_t: f32) -> f32 {
for l in row.iter_mut() {
*l *= inv_t;
}
let max = max_of(row);
let mut acc = [0.0f32; 16];
let (chunks, rest) = row.as_chunks_mut::<16>();
#[cfg(target_arch = "aarch64")]
unsafe {
use std::arch::aarch64::*;
let sums = acc.as_mut_ptr();
let top = vdupq_n_f32(max);
for chunk in chunks.iter_mut() {
let p = chunk.as_mut_ptr();
let x = [
vsubq_f32(vld1q_f32(p), top),
vsubq_f32(vld1q_f32(p.add(4)), top),
vsubq_f32(vld1q_f32(p.add(8)), top),
vsubq_f32(vld1q_f32(p.add(12)), top),
];
if vmaxnmvq_f32(vmaxnmq_f32(vmaxnmq_f32(x[0], x[1]), vmaxnmq_f32(x[2], x[3]))) <= DEAD_BELOW {
for (q, _) in x.iter().enumerate() {
vst1q_f32(p.add(4 * q), vdupq_n_f32(0.0));
}
continue;
}
for (q, xq) in x.iter().enumerate() {
let e = exp_neg_x4(*xq);
vst1q_f32(p.add(4 * q), e);
vst1q_f32(sums.add(4 * q), vaddq_f32(vld1q_f32(sums.add(4 * q)), e));
}
}
}
#[cfg(not(target_arch = "aarch64"))]
for chunk in chunks.iter_mut() {
for j in 0..16 {
let e = exp_neg(chunk[j] - max);
chunk[j] = e;
acc[j] += e;
}
}
let quarter = |q: usize| (acc[q] + acc[q + 4]) + (acc[q + 8] + acc[q + 12]);
let mut sum = (quarter(0) + quarter(1)) + (quarter(2) + quarter(3));
for l in rest {
let e = exp_neg(*l - max);
*l = e;
sum += e;
}
sum
}
#[cfg(target_arch = "aarch64")]
#[inline(always)]
unsafe fn exp_neg_x4(x: std::arch::aarch64::float32x4_t) -> std::arch::aarch64::float32x4_t {
use std::arch::aarch64::*;
let zero = vdupq_n_f32(0.0);
let live = vcgtq_f32(x, vdupq_n_f32(DEAD_BELOW));
let x = vbslq_f32(live, x, zero);
let n = vcvtq_s32_f32(vsubq_f32(vmulq_f32(x, vdupq_n_f32(std::f32::consts::LOG2_E)), vdupq_n_f32(0.5)));
let nf = vcvtq_f32_s32(n);
let r = vsubq_f32(vsubq_f32(x, vmulq_f32(nf, vdupq_n_f32(0.693_359_4))), vmulq_f32(nf, vdupq_n_f32(-2.121_944_4e-4)));
let mut poly = vaddq_f32(vmulq_f32(vdupq_n_f32(1.987_569_1e-4), r), vdupq_n_f32(1.398_2e-3));
poly = vaddq_f32(vmulq_f32(poly, r), vdupq_n_f32(8.333_452e-3));
poly = vaddq_f32(vmulq_f32(poly, r), vdupq_n_f32(4.166_579_6e-2));
poly = vaddq_f32(vmulq_f32(poly, r), vdupq_n_f32(1.666_666_5e-1));
poly = vaddq_f32(vmulq_f32(poly, r), vdupq_n_f32(0.5));
let y = vaddq_f32(vaddq_f32(vmulq_f32(poly, vmulq_f32(r, r)), r), vdupq_n_f32(1.0));
let scale = vreinterpretq_f32_s32(vshlq_n_s32::<23>(vaddq_s32(n, vdupq_n_s32(127))));
vbslq_f32(live, vmulq_f32(y, scale), zero)
}
fn top_head(probs: &[f32], head: usize, out: &mut Vec<(u32, f32)>) {
const BLOCK: usize = 16;
out.clear();
let n = probs.len();
if head >= n / 4 {
out.extend(probs.iter().enumerate().map(|(i, &p)| (i as u32, p)));
sort_head(out, head);
return;
}
let mut floor = f32::NEG_INFINITY;
for (b, block) in probs.chunks(BLOCK).enumerate() {
if max_of(block) <= floor {
continue;
}
for (j, &p) in block.iter().enumerate() {
if p > floor {
out.push(((b * BLOCK + j) as u32, p));
if out.len() == 2 * head {
out.select_nth_unstable_by(head - 1, by_prob_desc);
out.truncate(head);
floor = out[head - 1].1;
}
}
}
}
sort_head(out, head);
}
pub fn sample(row: &mut [f32], p: &SamplingParams, u: f64) -> u32 {
if p.temperature <= 0.0 {
return argmax(row);
}
if row.is_empty() {
return 0;
}
let sum = softmax_weights(row, 1.0 / p.temperature);
if sum > 0.0 {
for w in row.iter_mut() {
*w /= sum;
}
}
let full_total: f32 = if sum > 0.0 { 1.0 } else { 0.0 };
let n = row.len();
let mut cands: Vec<(u32, f32)> = Vec::new();
let k_cap = if p.top_k == 0 { n } else { (p.top_k as usize).clamp(1, n) };
let mut head = if p.top_k == 0 { INITIAL_HEAD.min(n) } else { k_cap };
loop {
top_head(row, head, &mut cands);
let front = &cands[..head];
let mut keep = head.min(k_cap);
let mut cut = p.top_k != 0;
if p.min_p > 0.0 {
let floor = p.min_p * front[0].1;
let c = front[..keep].iter().take_while(|c| c.1 >= floor).count();
if c < keep {
keep = c.max(1);
cut = true;
} else if !cut && head < n {
head = grow(head, n);
continue;
}
}
if p.top_p < 1.0 {
let mut acc = 0.0f32;
let mut found = None;
for (i, c) in front[..keep].iter().enumerate() {
acc += c.1;
if acc >= p.top_p {
found = Some(i + 1);
break;
}
}
match found {
Some(k) => {
keep = k.max(1);
cut = true;
}
None if !cut && head < n => {
head = grow(head, n);
continue;
}
None => {}
}
}
let total: f32 = if cut { front[..keep].iter().map(|c| c.1).sum() } else { full_total };
let mut target = u as f32 * total;
let mut last_with_mass = front[0].0;
for c in &front[..keep] {
if target < c.1 {
return c.0;
}
target -= c.1;
if c.1 > 0.0 {
last_with_mass = c.0;
}
}
if cut || head == n {
return last_with_mass;
}
head = grow(head, n);
}
}
pub fn logprobs(row: &[f32], chosen: u32, k: usize) -> (f32, Vec<(u32, f32)>) {
let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let lse = max + row.iter().map(|&l| (l - max).exp()).sum::<f32>().ln();
let lp = row.get(chosen as usize).map(|&l| l - lse).unwrap_or(f32::NEG_INFINITY);
let k = k.min(row.len());
if k == 0 {
return (lp, Vec::new());
}
let mut top: Vec<(u32, f32)> = row.iter().enumerate().map(|(i, &l)| (i as u32, l - lse)).collect();
sort_head(&mut top, k);
top.truncate(k);
(lp, top)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn uniform_is_a_pure_function_of_position() {
assert_eq!(uniform(7), uniform(7));
assert_ne!(uniform(7), uniform(8));
assert!((0.0..1.0).contains(&uniform(u64::MAX)));
}
#[test]
fn greedy_is_lowest_index_argmax() {
assert_eq!(argmax(&[1.0, 3.0, 3.0, 2.0]), 1);
assert_eq!(argmax(&[]), 0);
let plain = |row: &[f32]| -> u32 {
let mut best = 0usize;
for (i, &v) in row.iter().enumerate() {
if v > row[best] {
best = i;
}
}
best as u32
};
let mut st = 0xA11CE_u64;
for n in 1..40usize {
for case in 0..40 {
let mut row: Vec<f32> = (0..n).map(|_| (lcg(&mut st) * 6.0).floor()).collect();
match case % 5 {
1 => row[n / 2] = f32::NAN,
2 => row[0] = f32::NAN,
3 => row[n - 1] = f32::INFINITY,
4 => row.iter_mut().for_each(|v| *v = f32::NEG_INFINITY),
_ => {}
}
assert_eq!(argmax(&row), plain(&row), "{row:?}");
}
}
}
#[test]
fn top_k_one_is_greedy_at_any_temperature() {
let p = SamplingParams { temperature: 1.0, top_k: 1, top_p: 1.0, ..Default::default() };
let mut row = vec![0.1, 0.5, 0.2];
assert_eq!(sample(&mut row, &p, 0.99), 1);
}
#[test]
fn top_p_cuts_on_the_original_mass_after_top_k() {
let row: Vec<f32> = [0.4f32, 0.3, 0.3].iter().map(|p| p.ln()).collect();
let p = SamplingParams { temperature: 1.0, top_k: 2, top_p: 0.5, ..Default::default() };
assert_eq!(sample(&mut row.clone(), &p, 0.95), 1, "the second candidate survives the cut");
assert_eq!(sample(&mut row.clone(), &p, 0.1), 0);
let p1 = SamplingParams { temperature: 1.0, top_k: 2, top_p: 0.39, ..Default::default() };
assert_eq!(sample(&mut row.clone(), &p1, 0.95), 0, "a threshold under the best probability keeps only it");
}
#[test]
fn a_draw_at_the_top_of_the_range_stays_on_a_token_with_mass() {
let top = 1.0 - f64::EPSILON / 2.0;
assert_eq!(top as f32, 1.0, "the draw this is about");
let mut row = vec![f32::NEG_INFINITY; 64];
(row[5], row[9]) = (1.0, 1.0);
let p = SamplingParams { temperature: 1.0, top_k: 40, top_p: 1.0, ..Default::default() };
assert_eq!(sample(&mut row.clone(), &p, top), 9, "the last token with any mass");
let p = SamplingParams { temperature: 1.0, top_k: 0, top_p: 1.0, ..Default::default() };
assert_eq!(sample(&mut row.clone(), &p, top), 9);
assert_eq!(sample(&mut row.clone(), &p, 0.0), 5);
}
#[test]
fn top_p_zero_keeps_the_top_token() {
let row: Vec<f32> = [0.4f32, 0.3, 0.3].iter().map(|p| p.ln()).collect();
for top_p in [0.0, -1.0] {
let p = SamplingParams { temperature: 1.0, top_p, ..Default::default() };
for u in [0.0, 0.5, 0.99] {
assert_eq!(sample(&mut row.clone(), &p, u), 0, "top_p {top_p}, draw {u}");
}
}
let off = SamplingParams { temperature: 1.0, top_p: 1.0, ..Default::default() };
assert_eq!(sample(&mut row.clone(), &off, 0.99), 2, "1.0 is the whole distribution");
}
#[test]
fn repeat_penalty_demotes_a_seen_token() {
let p = SamplingParams { repeat_penalty: 100.0, ..Default::default() };
let mut row = vec![0.5, 10.0, 0.7];
apply_penalties(&mut row, &[1], &p, &Default::default());
assert_eq!(argmax(&row), 2);
}
fn reference_sample(row: &[f32], p: &SamplingParams, u: f64) -> u32 {
let mut weights = row.to_vec();
let sum = softmax_weights(&mut weights, 1.0 / p.temperature);
let mut cands: Vec<(u32, f32)> = weights.iter().enumerate().map(|(i, &w)| (i as u32, w / sum)).collect();
let full_total: f32 = 1.0;
cands.sort_by(by_prob_desc);
let mut bounded = false;
if p.top_k != 0 {
cands.truncate((p.top_k as usize).clamp(1, cands.len()));
bounded = true;
}
if p.min_p > 0.0 {
let floor = p.min_p * cands[0].1;
let keep = cands.iter().take_while(|c| c.1 >= floor).count().max(1);
if keep < cands.len() {
bounded = true;
}
cands.truncate(keep);
}
if p.top_p < 1.0 {
let mut acc = 0.0f32;
let mut keep = 0usize;
for c in &cands {
acc += c.1;
keep += 1;
if acc >= p.top_p {
bounded = true;
break;
}
}
cands.truncate(keep.max(1));
}
let total: f32 = if bounded { cands.iter().map(|c| c.1).sum() } else { full_total };
let mut target = u as f32 * total;
for c in &cands {
if target < c.1 {
return c.0;
}
target -= c.1;
}
cands.last().map(|c| c.0).unwrap_or(0)
}
fn lcg(state: &mut u64) -> f32 {
*state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
((*state >> 40) as f32) / ((1u64 << 24) as f32)
}
#[test]
fn head_sorting_matches_the_full_sort_reference() {
let n = 3001usize;
let mut st = 0x5EED_u64;
for shape in 0..3 {
let mut row: Vec<f32> = (0..n).map(|_| lcg(&mut st) * [12.0, 3.0, 0.02][shape]).collect();
row[17] += [6.0, 1.0, 0.0][shape];
for &top_k in &[0u32, 1, 7, 50, 5000] {
for &top_p in &[1.0f32, 0.9, 0.3] {
for &min_p in &[0.0f32, 0.1] {
for &temperature in &[0.7f32, 1.5] {
let p = SamplingParams { temperature, top_k, top_p, min_p, ..Default::default() };
let mut us = vec![0.0, 0.001, 0.25, 0.5, 0.75, 0.999, 0.999_999];
us.extend((0..8).map(|_| lcg(&mut st) as f64));
for u in us {
let want = reference_sample(&row, &p, u);
let got = sample(&mut row.clone(), &p, u);
assert_eq!(got, want, "shape {shape} top_k {top_k} top_p {top_p} min_p {min_p} t {temperature} u {u}");
}
}
}
}
}
}
}
#[test]
fn exp_neg_is_the_exponential_to_two_ulp() {
let mut worst = 0.0f64;
let mut x = 0.0f32;
while x > DEAD_BELOW {
let (got, want) = (exp_neg(x) as f64, (x as f64).exp());
worst = worst.max(((got - want) / want).abs());
x -= 0.000_37;
}
assert!(worst < 2.5e-7, "relative error {worst:e}");
assert_eq!(exp_neg(0.0), 1.0);
for dead in [DEAD_BELOW, -37.0, -1000.0, f32::NEG_INFINITY, f32::NAN] {
assert_eq!(exp_neg(dead), 0.0, "{dead}");
}
}
#[test]
fn the_row_pass_is_exp_neg_element_for_element() {
let mut st = 0xE4B_u64;
for n in [1usize, 15, 16, 17, 31, 32, 33, 100, 1000] {
let mut row: Vec<f32> = (0..n).map(|_| lcg(&mut st) * 120.0 - 60.0).collect();
if n > 20 {
(row[3], row[18]) = (f32::NAN, f32::NEG_INFINITY);
}
let inv_t = 1.3f32;
let scaled: Vec<f32> = row.iter().map(|l| l * inv_t).collect();
let max = scaled.iter().copied().filter(|v| !v.is_nan()).fold(f32::NEG_INFINITY, f32::max);
let want: Vec<f32> = scaled.iter().map(|l| exp_neg(l - max)).collect();
let mut acc = [0.0f32; 16];
let whole = n / 16 * 16;
for (i, e) in want[..whole].iter().enumerate() {
acc[i % 16] += e;
}
let quarter = |q: usize| (acc[q] + acc[q + 4]) + (acc[q + 8] + acc[q + 12]);
let mut want_sum = (quarter(0) + quarter(1)) + (quarter(2) + quarter(3));
for e in &want[whole..] {
want_sum += e;
}
let sum = softmax_weights(&mut row, inv_t);
assert_eq!(row.iter().map(|v| v.to_bits()).collect::<Vec<_>>(), want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(), "n {n}");
assert_eq!(sum.to_bits(), want_sum.to_bits(), "n {n}");
}
}
#[test]
fn top_head_is_the_head_of_the_sorted_row() {
let mut st = 0xBEEF_u64;
for &n in &[1usize, 15, 16, 17, 400, 5000] {
for &levels in &[2.0f32, 37.0, 1.0e6] {
let row: Vec<f32> = (0..n).map(|_| (lcg(&mut st) * levels).floor() / levels).collect();
let mut whole: Vec<(u32, f32)> = row.iter().enumerate().map(|(i, &p)| (i as u32, p)).collect();
whole.sort_by(by_prob_desc);
for &head in &[1usize, 2, 7, 128, 1000] {
let head = head.min(n);
let mut got = Vec::new();
top_head(&row, head, &mut got);
assert_eq!(&got[..head], &whole[..head], "n {n} levels {levels} head {head}");
}
}
}
}
#[test]
fn a_row_with_no_mass_answers_the_first_token() {
let p = SamplingParams { temperature: 0.8, top_k: 20, top_p: 0.9, ..Default::default() };
for u in [0.0, 0.5, 0.999] {
assert_eq!(sample(&mut vec![f32::NEG_INFINITY; 500], &p, u), 0);
assert_eq!(sample(&mut vec![f32::NAN; 500], &p, u), 0);
}
}
#[test]
#[ignore]
fn sampler_cost() {
let n = 151_936usize;
let mut st = 0xC057_u64;
let rows: Vec<Vec<f32>> = (0..3)
.map(|shape| {
let mut row: Vec<f32> = (0..n).map(|_| lcg(&mut st) * [30.0, 12.0, 45.0][shape]).collect();
row[4242] += 8.0;
row
})
.collect();
let time = |what: &str, f: &mut dyn FnMut(&[f32]) -> u32| {
let mut sink = 0u32;
let t = std::time::Instant::now();
for i in 0..300 {
sink ^= f(&rows[i % 3]);
}
eprintln!("{what}: {:.0} us ({sink})", t.elapsed().as_secs_f64() * 1e6 / 300.0);
};
let p = SamplingParams { temperature: 0.7, top_k: 20, top_p: 0.8, ..Default::default() };
let open = SamplingParams { temperature: 1.0, top_k: 0, top_p: 0.95, ..Default::default() };
time("copy the row", &mut |r| r.to_vec().len() as u32);
time("argmax", &mut |r| argmax(r));
time("sample top-k 20 top-p 0.8 (with the copy)", &mut |r| sample(&mut r.to_vec(), &p, 0.37));
time("sample top-p 0.95, no top-k (with the copy)", &mut |r| sample(&mut r.to_vec(), &open, 0.37));
time("reference: exp per element and a full sort (with the copy)", &mut |r| reference_sample(r, &p, 0.37));
let mut w = rows[0].clone();
let t = std::time::Instant::now();
for _ in 0..300 {
w.copy_from_slice(&rows[0]);
std::hint::black_box(softmax_weights(&mut w, 1.0 / 0.7));
}
eprintln!("softmax_weights alone (with the copy): {:.0} us", t.elapsed().as_secs_f64() * 1e6 / 300.0);
if let Some(path) = std::env::var_os("SUPERFLUID_SAMPLER_ROWS") {
let bytes = std::fs::read(path).unwrap();
let all: Vec<f32> = bytes.as_chunks::<4>().0.iter().map(|b| f32::from_ne_bytes(*b)).collect();
let real: Vec<&[f32]> = all.chunks_exact(n).collect();
let mut w = vec![0.0f32; n];
let mut cands = Vec::new();
let (mut a, mut b, mut c) = (0.0, 0.0, 0.0);
for i in 0..300 {
w.copy_from_slice(real[i % real.len()]);
let t = std::time::Instant::now();
let sum = softmax_weights(&mut w, 1.0 / 0.7);
a += t.elapsed().as_secs_f64();
let t = std::time::Instant::now();
for x in w.iter_mut() {
*x /= sum;
}
b += t.elapsed().as_secs_f64();
let t = std::time::Instant::now();
top_head(&w, 20, &mut cands);
c += t.elapsed().as_secs_f64();
}
eprintln!("real rows: weights {:.0} us, division {:.0} us, head {:.0} us", a * 1e6 / 300.0, b * 1e6 / 300.0, c * 1e6 / 300.0);
}
}
#[test]
fn flat_row_draws_deep_in_the_id_order() {
let row = vec![0.0f32; 4000];
let p = SamplingParams { temperature: 1.0, top_k: 0, top_p: 1.0, ..Default::default() };
let got = sample(&mut row.clone(), &p, 0.9999);
assert_eq!(got, reference_sample(&row, &p, 0.9999));
assert!(got >= 3990, "equal probabilities walk the id order: {got}");
}
#[test]
fn logprobs_top_k_is_the_sorted_head() {
let row = vec![1.0f32, 5.0, 3.0, 5.0, 0.5];
let (_, top) = logprobs(&row, 0, 3);
assert_eq!(top.iter().map(|t| t.0).collect::<Vec<_>>(), vec![1, 3, 2]);
assert!(top[0].1 > top[2].1);
}
#[test]
fn logprobs_are_normalized() {
let (lp, top) = logprobs(&[0.0, 0.0], 0, 2);
assert!((lp - (0.5f32).ln()).abs() < 1e-6);
assert_eq!(top.len(), 2);
}
}