Fix typical sampling.

Ai00-X · May 24, 2024 · 39fd373 · 39fd373
1 parent 3391f3f
commit 39fd373
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/crates/ai00-core/src/sampler/typical.rs b/crates/ai00-core/src/sampler/typical.rs
@@ -69,10 +69,11 @@ impl Sampler for TypicalSampler {
     fn sample(&mut self, probs: &[f32]) -> u16 {
         let TypicalSampler { params, state } = self;
 
-        let entropy = -probs.iter().map(|x| x * x.ln()).sum::<f32>();
+        let probs = probs.iter().map(|&x| (x, x.ln())).collect_vec();
+        let entropy = -probs.iter().map(|(x, y)| x * y).sum::<f32>();
         let sorted = probs
-            .iter()
-            .map(|&x| (x, (x - entropy).abs()))
+            .into_iter()
+            .map(|(x, y)| (x, (-y - entropy).abs()))
             .enumerate()
             .sorted_unstable_by(|(_, (_, x)), (_, (_, y))| x.total_cmp(y))
             .map(|(id, (x, _))| (id, x))