pub fn log_softmax(logits: &[f32]) -> Vec<f32>
Compute log(softmax(logits)) in a numerically stable way. Returns a vector of the same length as logits.