Function softmax 

pub fn softmax(
    arr: &ArrayBase<OwnedRepr<f64>, Dim<[usize; 1]>>,
) -> ArrayBase<OwnedRepr<f64>, Dim<[usize; 1]>>

Softmax function for probability distribution