Function softmax_axis

pub fn softmax_axis<A, S, D>(
    args: &ArrayBase<S, D>,
    axis: usize,
) -> ArrayBase<OwnedRepr<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}}$