pub fn burn_softmax_last_dim(x: Tensor<2>) -> Tensor<2>
Softmax over the last dim for host-checked logits.