added rope_type param to attention calls

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -116,6 +116,7 @@ def __init__(
         dtype=dtype,
         mesh=mesh,
         attention_kernel=self.attention_kernel,
+        rope_type=rope_type,
     )
 
     self.audio_norm1 = nnx.RMSNorm(
@@ -138,6 +139,7 @@ def __init__(
         dtype=dtype,
         mesh=mesh,
         attention_kernel=self.attention_kernel,
+        rope_type=rope_type,
     )
 
     # 2. Prompt Cross-Attention
@@ -162,6 +164,7 @@ def __init__(
         dtype=dtype,
         mesh=mesh,
         attention_kernel=self.attention_kernel,
+        rope_type=rope_type,
     )
 
     self.audio_norm2 = nnx.RMSNorm(
@@ -185,6 +188,7 @@ def __init__(
         dtype=dtype,
         mesh=mesh,
         attention_kernel=self.attention_kernel,
+        rope_type=rope_type,
     )
 
     # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
@@ -204,6 +208,7 @@ def __init__(
         dtype=dtype,
         mesh=mesh,
         attention_kernel=self.attention_kernel,
+        rope_type=rope_type,
     )
 
     self.video_to_audio_norm = nnx.RMSNorm(
@@ -227,6 +232,7 @@ def __init__(
         dtype=dtype,
         mesh=mesh,
         attention_kernel=self.attention_kernel,
+        rope_type=rope_type,
     )
 
     # 4. Feed Forward