Function softmax_axis 

pub fn softmax_axis<A, S, D>(
    args: &ArrayBase<S, D, A>,
    axis: usize,
) -> Array<A, D>
where
    A: Float + ScalarOperand,
    D: RemoveAxis,
    S: Data<Elem = A>,

Softmax function along a specific axis:

f(x_i) = \frac{e^{x_i}}{\sum_{j} e^{x_j}}