Function log_softmax_3d 

pub fn log_softmax_3d(logits: &[Vec<Vec<f64>>]) -> Vec<Vec<Vec<f64>>>

Numerically stable log-softmax over the last axis. Input / output shape: [B][M][|V|].