Function repeat_kv 

pub fn repeat_kv(xs: Tensor, n_rep: usize) -> Result<Tensor>

Repeats a key or value tensor for grouped query attention The input tensor should have a shape (batch, num_kv_heads, seq_len, head_dim),