Function dfdx::tensor_ops::softmax

pub fn softmax<T: Reduce<Axes>, Axes>(t: T) -> T

Computes the softmax function across Axes.

Equivalent to exp(log_softmax(t)).

Pytorch equivalent: t.softmax(Axes)

Related functions: logsumexp(), log_softmax()

Example:

let t: Tensor3D<2, 3, 5> = TensorCreator::zeros();
let _ = t.softmax::<Axis<2>>();

Using multi axis softmax:

let _ = t.softmax::<Axes2<1, 2>>();