Skip to main content

softmax

Function softmax 

pub fn softmax<const D: usize>(
    tensor: Tensor<D>,
    dim: impl AsIndex,
) -> Tensor<D>
Expand description

Applies the softmax function on the input tensor along the given dimension.

softmax(x_i) = exp(x_i) / sum_j(exp(x_j))

§Arguments

  • dim: the dimension along which Softmax will be computed. Negative dimensions are supported and count from the end.

§Panics

  • If dim is outside [-D, D)