add voxcpm with some bug

src/position_embed/rope.rs

@@ -64,6 +64,7 @@ pub fn apply_rotary_pos_emb_vision(
     // cos, sin -> (seq_len, head_dim) -> (seq_len, 1, head_dim)
     let cos = cos.unsqueeze(D::Minus2)?;
     let sin = sin.unsqueeze(D::Minus2)?;
+    
     let q_embed = q
         .broadcast_mul(&cos)?
         .add(&rotate_half(q)?.broadcast_mul(&sin)?)?;
@@ -78,15 +79,42 @@ pub fn apply_rotary_pos_emb(
     k: &Tensor,
     cos: &Tensor,
     sin: &Tensor,
+    tof32: bool,
 ) -> Result<(Tensor, Tensor)> {
-    // sin/cos: (bs, 1, seq_len, head_dim)
+    // sin/cos: to (bs, 1, seq_len, head_dim)
     // q/k: (bs, n_head, seq_len, head_dim)
+    let mut cos = cos.clone();
+    let mut sin = sin.clone();
+    if cos.rank() == 2 {
+        // (seq_len, head_dim) -> (1, 1, seq_len, head_dim)
+        cos = cos.unsqueeze(0)?.unsqueeze(0)?;
+        sin = sin.unsqueeze(0)?.unsqueeze(0)?;
+    }
+    if cos.rank() == 3 {
+        // (bs, seq_len, head_dim) -> (bs, 1, seq_len, head_dim)
+        cos = cos.unsqueeze(1)?;
+        sin = sin.unsqueeze(1)?;
+    }
+    let orig_dtype = q.dtype();
+    let q = if tof32 {
+        &q.to_dtype(DType::F32)?
+    } else {
+        q
+    };
+    let k = if tof32 {
+        &k.to_dtype(DType::F32)?
+    } else {
+        k
+    };
+    let cos = cos.to_dtype(q.dtype())?;
+    let sin = sin.to_dtype(q.dtype())?;
+
     let q_embed = q
         .broadcast_mul(&cos)?
-        .add(&rotate_half(q)?.broadcast_mul(&sin)?)?;
+        .add(&rotate_half(q)?.broadcast_mul(&sin)?)?.to_dtype(orig_dtype)?;
     let k_embed = k
         .broadcast_mul(&cos)?
-        .add(&rotate_half(k)?.broadcast_mul(&sin)?)?;
+        .add(&rotate_half(k)?.broadcast_mul(&sin)?)?.to_dtype(orig_dtype)?;
     Ok((q_embed, k_embed))
 }