pub fn log_softmax(logits: &[f64]) -> Vec<f64>
Compute log-softmax: log(softmax(x)) with the log-sum-exp trick.