rope arg changed

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -676,7 +676,7 @@ def __init__(
 
         cross_attn_pos_embed_max_pos = max(self.pos_embed_max_pos, self.audio_pos_embed_max_pos)
         self.cross_attn_rope = LTX2RotaryPosEmbed(
-            dim=self.audio_cross_attention_dim,
+            dim=inner_dim,
             patch_size=self.patch_size,
             patch_size_t=self.patch_size_t,
             base_num_frames=cross_attn_pos_embed_max_pos,
@@ -690,7 +690,7 @@ def __init__(
             num_attention_heads=self.num_attention_heads,
         )
         self.cross_attn_audio_rope = LTX2RotaryPosEmbed(
-            dim=self.audio_cross_attention_dim,
+            dim=audio_inner_dim,
             patch_size=self.audio_patch_size,
             patch_size_t=self.audio_patch_size_t,
             base_num_frames=cross_attn_pos_embed_max_pos,

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -279,7 +279,6 @@ def test_transformer_3d_model_dot_product_attention(self):
             audio_attention_head_dim=128,
             audio_cross_attention_dim=32,
             num_layers=1, # Reduced layers for speed
-            config=self.config,
             scan_layers=False,
             mesh=self.mesh,
             attention_kernel="dot_product"