mistralrs_core/
layers_utils.rs

1use candle_core::{Result, Tensor};
2
3pub fn repeat_kv(x: Tensor, n_rep: usize) -> Result<Tensor> {
4    if n_rep == 1 {
5        Ok(x)
6    } else {
7        let (b_sz, n_kv_head, seq_len, head_dim) = x.dims4()?;
8        Tensor::cat(&vec![&x; n_rep], 2)?.reshape((b_sz, n_kv_head * n_rep, seq_len, head_dim))
9    }
10}