pub fn softmax_last_dim(input: &Tensor) -> Result<Tensor, KernelError>
Softmax along the last tensor dimension.