Fix apply_rope broadcasting dimensions in attention_flax.py

Author: amepas

src/maxdiffusion/models/attention_flax.py

@@ -1697,11 +1697,11 @@ def apply_rope(xq: Array, xk: Array, freqs_cis: Any) -> tuple[Array, Array]:
   if isinstance(freqs_cis, (tuple, list)):
     cos, sin = freqs_cis
     if cos.ndim == 2:
-      cos = cos[None, None, :, :]
-      sin = sin[None, None, :, :]
+      cos = cos[None, :, None, :]
+      sin = sin[None, :, None, :]
     elif cos.ndim == 3 and cos.shape[0] == 1:
-      cos = cos[None, :, :, :]
-      sin = sin[None, :, :, :]
+      cos = cos[:, :, None, :]
+      sin = sin[:, :, None, :]
 
     def _rotate(x):
       x_reshaped = x.reshape(*x.shape[:-1], -1, 2)