Merge pull request #80 from jim-steele/main

MPS support workaround for complex tensors
2025-09-08 20:01:41 +03:00
parent c59676b008 148f30a467
commit 33cfed3880
1 changed files with 5 additions and 1 deletions
@@ -211,6 +211,10 @@ def apply_rotary_enc(
    # repeat freqs along seq_len dim to match k seq_len
    if repeat_freqs_k:
        r = xk_.shape[-2] // xq_.shape[-2]
-        freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1)
+        if freqs_cis.is_complex() and freqs_cis.device.type == "mps":
+            # MPS doesn't support repeat on complex; cat works fine.
+            freqs_cis = torch.cat([freqs_cis] * r, dim=-2)
+        else:
+            freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)