pub fn unsplit_head(d: &Tensor, which: usize) -> Result<Tensor>
split_heads backward for one of q/k/v (which in 0..3): d [h, t, hd] -> [t, 3c] with the other thirds zero.
which