logprob/
softmax.rs

1use super::{Float, FloatIsNanOrPositiveInfinity, LogProb, adding::Ln2};
2use core::iter::Sum;
3
4use alloc::vec::Vec;
5
6///Returns an iterator with the softmax values of a slice of floats.
7///# Errors
8///Returns [`FloatIsNanOrPositiveInfinity`] if any float is NaN or positive infinity.
9#[expect(clippy::missing_panics_doc)]
10pub fn softmax<T: Float + Sum<T> + Ln2>(
11    val: &[T],
12) -> Result<impl Iterator<Item = LogProb<T>> + use<T>, FloatIsNanOrPositiveInfinity> {
13    let v: Vec<_> = val
14        .iter()
15        .map(|x| {
16            if x.is_nan() || (x.is_infinite() && x.is_sign_positive()) {
17                Err(FloatIsNanOrPositiveInfinity)
18            } else {
19                Ok(*x)
20            }
21        })
22        .collect::<Result<Vec<_>, FloatIsNanOrPositiveInfinity>>()?;
23    let max: T = *val
24        .iter()
25        .max_by(|x, y| x.partial_cmp(y).unwrap())
26        .unwrap_or(&T::ZERO);
27    let s: T = v.iter().map(|&x| (x - max).exp()).sum::<T>().ln();
28
29    Ok(v.into_iter().map(move |x| LogProb(x - s - max)))
30}
31
32///This trait allows iterators to have [`softmax`].
33///
34///# Example
35///
36/// ```
37/// use logprob::{LogProb, LogSumExp, Softmax};
38///
39/// # fn main() -> anyhow::Result<()>{
40/// let logits = [-1.0_f64, 0.0, 1.0, 2.0];
41///
42/// let probs: Vec<LogProb<_>> = logits
43///     .iter()
44///     .copied()
45///     .softmax()?
46///     .collect();
47///
48/// let total = probs.iter().log_sum_exp_float();
49/// approx::assert_relative_eq!(total.exp(), 1.0, epsilon = 1e-9);
50/// # Ok(())
51/// # }
52/// ```
53pub trait Softmax: Iterator {
54    ///Gets the softmax from an iterator as another iterator.
55    ///
56    ///# Errors
57    ///Returns [`FloatIsNanOrPositiveInfinity`] if any float is NaN or positive infinity.
58    fn softmax<T: Float + Sum<T> + Ln2>(
59        self,
60    ) -> Result<impl Iterator<Item = LogProb<T>>, FloatIsNanOrPositiveInfinity>
61    where
62        Self: Sized,
63        Self: Iterator<Item = T>,
64    {
65        let v: Vec<_> = self.collect();
66        softmax(&v)
67    }
68}
69
70impl<I: ?Sized> Softmax for I where I: Iterator {}
logprob/softmax.rs

logprob/
softmax.rs
