Function unsplit_head 

pub fn unsplit_head(d: &Tensor, which: usize) -> Result<Tensor>

Available on crate feature train only.

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