Function softmax 

pub fn softmax(a: &Array, axes: &[isize]) -> Result<Array>

exp(a) / sum(exp(a)) over axes (all axes if empty), with the maximum subtracted first (without gradient), as MLX does.