Struct Softmax 

pub struct Softmax<'src, 'dst> { /* private fields */ }

Computes the softmax function over a slice of floats.

The implementation uses a three-pass approach for numerical stability. See https://ogunlao.github.io/2020/04/26/you_dont_really_know_softmax.html. and https://arxiv.org/abs/2001.04438.

Implementations

impl<'src, 'dst> Softmax<'src, 'dst>

pub fn new(input: &'src [f32], output: &'dst mut [MaybeUninit<f32>]) -> Self

Construct a softmax operation which reads input and writes to to output.

pub fn new_mut(input: &'dst mut [f32]) -> Self

where 'dst: 'src,

Construct a softmax operation which updates input in place.

pub fn flush_nans_to_zero(self, flush: bool) -> Self

Replace NaN values in the output with zeros.

This option exists to changing handling of the case where the input values are all negative infinity. In that case the normal output would be NaN.

In the context of attention operations which use negative infinity to represent masked token positions, it is preferable to produce zeros as the output if all input positions are masked. See https://github.com/pytorch/pytorch/issues/41508.

Trait Implementations

impl<'dst> SimdOp for Softmax<'_, 'dst>

type Output = &'dst mut [f32]

The normalized elements.

fn eval<I: Isa>(self, isa: I) -> Self::Output

Evaluate the operation using the given instruction set.

fn dispatch(self) -> Self::Output

where Self: Sized,

Dispatch this operation using the preferred ISA for the current platform.